capnhook-ml 0.1.0__cp312-cp312-macosx_11_0_arm64.whl

simd/binary.hpp

@@ -0,0 +1,84 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdlib>
+#include <stdexcept>
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+#include <hwy/highway.h>
+
+#include "../alloc.hpp"
+
+namespace nb = nanobind;
+
+HWY_BEFORE_NAMESPACE();
+namespace hwy {
+namespace HWY_NAMESPACE {
+namespace capnhook {
+
+template <typename T, typename Op>
+nb::ndarray<nb::numpy, T, nb::ndim<1>> binary(nb::ndarray<T, nb::c_contig> a,
+                     nb::ndarray<T, nb::c_contig> b) {
+    const size_t N = a.shape(0);
+    if (b.shape(0) != N) throw std::runtime_error("shape mismatch");
+    const T* A = a.data();
+    const T* B = b.data();
+
+    const size_t bytes = N * sizeof(T);
+    void* raw = aligned_alloc64(bytes);
+    T* C = static_cast<T*>(raw);
+
+#if defined(_MSC_VER)
+    nb::capsule deleter(C, [](void* p) noexcept { _aligned_free(p); });
+#else
+    nb::capsule deleter(C, [](void* p) noexcept { free(p); });
+#endif
+
+    const ScalableTag<T> d;
+    const size_t L = Lanes(d);
+    Op op;
+    size_t i = 0;
+
+    for (; i + L <= N; i += L) {
+        auto va = Load(d, A + i);
+        auto vb = Load(d, B + i);
+        auto vc = op(va, vb);
+        Store(vc, d, C + i);
+    }
+    for (; i < N; ++i) {
+        C[i] = op(A[i], B[i]);
+    }
+
+    return nb::ndarray<nb::numpy, T, nb::ndim<1>>(C, { N }, deleter);
+}
+
+#define DEFINE_SIMD_BINARY_OP(Symbol, expr_scalar, expr_simd)                      \
+struct Symbol##Op {                                                                  \
+    template <typename V> HWY_INLINE V operator()(V a, V b) const {                  \
+        return (expr_simd);                                                          \
+    }                                                                                \
+    HWY_INLINE float  operator()(float  a, float  b) const { return (expr_scalar); }  \
+    HWY_INLINE double operator()(double a, double b) const { return (expr_scalar); }  \
+};                                                                                   \
+inline nb::ndarray<nb::numpy, float, nb::ndim<1>>                                    \
+Symbol(nb::ndarray<float,  nb::c_contig> a,                                          \
+       nb::ndarray<float,  nb::c_contig> b) {                                        \
+    return binary<float,  Symbol##Op>(a, b);                                          \
+}                                                                                    \
+inline nb::ndarray<nb::numpy, double, nb::ndim<1>>                                   \
+Symbol(nb::ndarray<double, nb::c_contig> a,                                          \
+       nb::ndarray<double, nb::c_contig> b) {                                        \
+    return binary<double, Symbol##Op>(a, b);                                          \
+}
+
+DEFINE_SIMD_BINARY_OP(add, a + b, Add(a, b))
+DEFINE_SIMD_BINARY_OP(sub, a - b, Sub(a, b))
+DEFINE_SIMD_BINARY_OP(mul, a * b, Mul(a, b))
+DEFINE_SIMD_BINARY_OP(div, a / b, Div(a, b))
+
+} // capnhook
+} // HWY_NAMESPACE
+} // hwy
+HWY_AFTER_NAMESPACE();
+
+namespace capnhook = hwy::HWY_NAMESPACE::capnhook;

simd/linalg.hpp

@@ -0,0 +1,109 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdlib>
+#include <cstring>
+#include <stdexcept>
+#include <vector>
+#include <cmath>
+#include <algorithm>
+#include <type_traits>
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+
+#include "../alloc.hpp"
+
+#ifdef USE_ACCELERATE
+  #include <Accelerate/Accelerate.h>   
+#else
+  #include <cblas.h>                  
+#endif
+
+namespace nb = nanobind;
+
+HWY_BEFORE_NAMESPACE();
+namespace hwy {
+namespace HWY_NAMESPACE {
+namespace capnhook {
+
+template <typename T>
+T dot(nb::ndarray<T, nb::c_contig> a, nb::ndarray<T, nb::c_contig> b) {
+    const size_t N = a.shape(0);
+    
+    if (b.shape(0) != N) {
+        throw std::runtime_error("dot: vectors must have the same length");
+    }
+    
+    const T* A = a.data();
+    const T* B = b.data();
+    
+    T result = 0;
+    
+    if constexpr (std::is_same_v<T, float>) {
+        result = cblas_sdot(N, A, 1, B, 1);
+    } else if constexpr (std::is_same_v<T, double>) {
+        result = cblas_ddot(N, A, 1, B, 1);
+    }
+    
+    return result;
+}
+
+template <typename T>
+nb::ndarray<nb::numpy, T, nb::ndim<2>> matmul(nb::ndarray<T, nb::c_contig, nb::ndim<2>> A,
+                      nb::ndarray<T, nb::c_contig, nb::ndim<2>> B) {
+    size_t M = A.shape(0), K = A.shape(1),
+           K2 = B.shape(0), N = B.shape(1);
+    if (K2 != K) throw std::runtime_error("matmul: inner dims must match");
+
+    size_t bytes = M * N * sizeof(T);
+    void* raw = aligned_alloc64(bytes);
+    T* C   = static_cast<T*>(raw);
+
+    T alpha = T(1), beta = T(0);
+    // row‑major
+    if constexpr (std::is_same_v<T, float>) {
+        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
+                    M, N, K, alpha,
+                    A.data(), K,
+                    B.data(), N,
+                    beta,
+                    C, N);
+    } else {
+        cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
+                    M, N, K, alpha,
+                    A.data(), K,
+                    B.data(), N,
+                    beta,
+                    C, N);
+    }
+
+#if defined(_MSC_VER)
+    nb::capsule deleter(C, [](void* p) noexcept { _aligned_free(p); });
+#else
+    nb::capsule deleter(C, [](void* p) noexcept { free(p); });
+#endif
+    return { C, { M, N }, deleter };
+}
+
+template <typename T>
+T trace(nb::ndarray<T, nb::c_contig, nb::ndim<2>> A) {
+    size_t M = A.shape(0), N = A.shape(1);
+    size_t n = std::min(M, N);
+    const T* data = A.data();
+    T sum = T(0);
+    for (size_t i = 0; i < n; ++i)
+        sum += data[i * N + i];
+    return sum;
+}
+
+template <typename T>
+T norm(nb::ndarray<T, nb::c_contig> a) {
+    return std::sqrt(dot(a, a));
+}
+
+} // capnhook
+} // HWY_NAMESPACE
+} // hwy
+HWY_AFTER_NAMESPACE();
+
+namespace capnhook = hwy::HWY_NAMESPACE::capnhook;

simd/reduce.hpp

@@ -0,0 +1,261 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdlib>
+#include <algorithm>
+#include <cmath>
+#include <stdexcept>
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+#include <hwy/highway.h>
+
+#include "../alloc.hpp"
+#include "binary.hpp"
+
+namespace nb = nanobind;
+
+HWY_BEFORE_NAMESPACE();
+namespace hwy {
+namespace HWY_NAMESPACE {
+namespace capnhook {
+
+template <typename T>
+T reduce_sum(nb::ndarray<T, nb::c_contig> a) {
+    const T* A = a.data();
+    size_t N = a.shape(0);
+    
+    if (N == 0) throw std::runtime_error("reduce_sum: zero-length input");
+    if (N == 1) return A[0];
+    
+    const ScalableTag<T> d;
+    auto acc = Zero(d);
+    size_t i = 0, L = Lanes(d);
+    
+    if (N < L) {
+        T sum = T(0);
+        for (size_t j = 0; j < N; j++) {
+            sum += A[j];
+        }
+        return sum;
+    }
+    
+    for (; i + L <= N; i += L) {
+        acc = Add(acc, Load(d, A + i));
+    }
+    
+    T total = GetLane(SumOfLanes(d, acc));
+    for (; i < N; ++i) total += A[i];
+    return total;
+}
+
+template <typename T>
+T reduce_min(nb::ndarray<T, nb::c_contig> a) {
+    const T* A = a.data();
+    size_t N = a.shape(0);
+    
+    if (N == 0) throw std::runtime_error("reduce_min: zero-length input");
+    if (N == 1) return A[0];
+    
+    const ScalableTag<T> d;
+    size_t L = Lanes(d);
+    
+    if (N < L) {
+        T min_val = A[0];
+        for (size_t j = 1; j < N; j++) {
+            min_val = std::min(min_val, A[j]);
+        }
+        return min_val;
+    }
+    
+    auto acc = Load(d, A);
+    size_t i = L; 
+    
+    for (; i + L <= N; i += L) {
+        acc = Min(acc, Load(d, A + i));
+    }
+    
+    T m = GetLane(MinOfLanes(d, acc));
+    for (; i < N; ++i) m = std::min(m, A[i]);
+    return m;
+}
+
+template <typename T>
+T reduce_max(nb::ndarray<T, nb::c_contig> a) {
+    const T* A = a.data();
+    size_t N = a.shape(0);
+    
+    if (N == 0) throw std::runtime_error("reduce_max: zero-length input");
+    if (N == 1) return A[0]; 
+    
+    const ScalableTag<T> d;
+    size_t L = Lanes(d);
+    
+    if (N < L) {
+        T max_val = A[0];
+        for (size_t j = 1; j < N; j++) {
+            max_val = std::max(max_val, A[j]);
+        }
+        return max_val;
+    }
+    
+    auto acc = Load(d, A); 
+    size_t i = L; 
+    
+    for (; i + L <= N; i += L) {
+        acc = Max(acc, Load(d, A + i));
+    }
+    
+    T m = GetLane(MaxOfLanes(d, acc));
+    for (; i < N; ++i) m = std::max(m, A[i]);
+    return m;
+}
+
+template <typename T>
+T reduce_prod(nb::ndarray<T, nb::c_contig> a) {
+    const T* A = a.data();
+    size_t N = a.shape(0);
+    
+    if (N == 0) throw std::runtime_error("reduce_prod: zero-length input");
+    if (N == 1) return A[0]; 
+    
+    const ScalableTag<T> d;
+    size_t L = Lanes(d);
+    
+    if (N < L) {
+        T prod = T(1);
+        for (size_t j = 0; j < N; j++) {
+            prod *= A[j];
+        }
+        return prod;
+    }
+    
+    auto acc = Set(d, T(1));
+    size_t i = 0;
+    
+    for (; i + L <= N; i += L) {
+        acc = Mul(acc, Load(d, A + i));
+    }
+    
+    T product = GetLane(SumOfLanes(d, acc));
+    for (; i < N; ++i) product *= A[i];
+    return product;
+}
+
+
+template <typename T>
+T reduce_mean(nb::ndarray<T, nb::c_contig> a) {
+    size_t N = a.shape(0);
+    return reduce_sum<T>(a) / T(N);
+}
+
+template <typename T>
+T reduce_var(nb::ndarray<T, nb::c_contig> a) {
+    size_t N = a.shape(0);
+    T mu = reduce_mean<T>(a);
+    const T* A = a.data();
+    T var = T(0);
+    for (size_t i = 0; i < N; ++i) {
+        T d = A[i] - mu;
+        var += d * d;
+    }
+    return var / T(N);
+}
+
+template <typename T>
+T reduce_std(nb::ndarray<T, nb::c_contig> a) {
+    return std::sqrt(reduce_var<T>(a));
+}
+
+
+template <typename T>
+bool reduce_any(nb::ndarray<T, nb::c_contig> a) {
+    const T* A = a.data();
+    size_t N = a.shape(0);
+    for (size_t i = 0; i < N; ++i) if (A[i] != T(0)) return true;
+    return false;
+}
+
+template <typename T>
+bool reduce_all(nb::ndarray<T, nb::c_contig> a) {
+    const T* A = a.data();
+    size_t N = a.shape(0);
+    for (size_t i = 0; i < N; ++i) if (A[i] == T(0)) return false;
+    return true;
+}
+
+
+template <typename T>
+size_t argmax(nb::ndarray<T, nb::c_contig> a) {
+    const T* A = a.data();
+    size_t N = a.shape(0);
+    if (N == 0) throw std::runtime_error("argmax: zero-length input");
+    size_t idx = 0;
+    T best = A[0];
+    for (size_t i = 1; i < N; ++i) {
+        if (A[i] > best) { best = A[i]; idx = i; }
+    }
+    return idx;
+}
+
+template <typename T>
+size_t argmin(nb::ndarray<T, nb::c_contig> a) {
+    const T* A = a.data();
+    size_t N = a.shape(0);
+    if (N == 0) throw std::runtime_error("argmin: zero-length input");
+    size_t idx = 0;
+    T best = A[0];
+    for (size_t i = 1; i < N; ++i) {
+        if (A[i] < best) { best = A[i]; idx = i; }
+    }
+    return idx;
+}
+
+
+template <typename T>
+nb::ndarray<nb::numpy, T, nb::ndim<1>>
+cumsum(nb::ndarray<T, nb::c_contig> a) {
+    size_t N = a.shape(0);
+    const T* A = a.data();
+    size_t bytes = N * sizeof(T);
+    void* raw = aligned_alloc64(bytes);
+    T* C = static_cast<T*>(raw);
+#if defined(_MSC_VER)
+    nb::capsule deleter(C, [](void* p) noexcept { _aligned_free(p); });
+#else
+    nb::capsule deleter(C, [](void* p) noexcept { free(p); });
+#endif
+    T acc = T(0);
+    for (size_t i = 0; i < N; ++i) {
+        acc += A[i];
+        C[i] = acc;
+    }
+    return { C, { N }, deleter };
+}
+
+template <typename T>
+nb::ndarray<nb::numpy, T, nb::ndim<1>>
+cumprod(nb::ndarray<T, nb::c_contig> a) {
+    size_t N = a.shape(0);
+    const T* A = a.data();
+    size_t bytes = N * sizeof(T);
+    void* raw = aligned_alloc64(bytes);
+    T* C = static_cast<T*>(raw);
+#if defined(_MSC_VER)
+    nb::capsule deleter(C, [](void* p) noexcept { _aligned_free(p); });
+#else
+    nb::capsule deleter(C, [](void* p) noexcept { free(p); });
+#endif
+    T acc = T(1);
+    for (size_t i = 0; i < N; ++i) {
+        acc *= A[i];
+        C[i] = acc;
+    }
+    return { C, { N }, deleter };
+}
+
+} // capnhook
+} // HWY_NAMESPACE
+} // hwy
+HWY_AFTER_NAMESPACE();
+
+namespace capnhook = hwy::HWY_NAMESPACE::capnhook;

simd/unary.hpp

@@ -0,0 +1,83 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdlib>
+#include <cmath>
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+#include <hwy/highway.h>
+#include <hwy/contrib/math/math-inl.h>
+
+#include "../alloc.hpp"
+
+namespace nb = nanobind;
+
+HWY_BEFORE_NAMESPACE();
+namespace hwy {
+namespace HWY_NAMESPACE {
+namespace capnhook {
+
+template <typename T, typename Op>
+nb::ndarray<nb::numpy, T, nb::ndim<1>> unary(nb::ndarray<T, nb::c_contig> a) {
+    const size_t N = a.shape(0);
+    const T* A = a.data();
+
+    const size_t bytes = N * sizeof(T);
+    void* raw = aligned_alloc64(bytes);
+    T* C = static_cast<T*>(raw);
+
+#if defined(_MSC_VER)
+    nb::capsule deleter(C, [](void* p) noexcept { _aligned_free(p); });
+#else
+    nb::capsule deleter(C, [](void* p) noexcept { free(p); });
+#endif
+
+    const ScalableTag<T> d;
+    Op op;
+    size_t i = 0;
+    const size_t L = Lanes(d);
+
+    for (; i + L <= N; i += L) {
+        auto v = Load(d, A + i);
+        Store(op(d, v), d, C + i);
+    }
+
+    for (; i < N; ++i) {
+        C[i] = op(A[i]);
+    }
+
+    return { C, { N }, deleter };
+}
+
+#define DEFINE_SIMD_UNARY_OP(Symbol, expr_scalar, expr_simd)        \
+struct Symbol##Op {                                                  \
+    template <class D, class V>                                     \
+    HWY_INLINE V operator()(D d, V v) const {                       \
+        return expr_simd;                                           \
+    }                                                                \
+    HWY_INLINE float operator()(float x) const { return (expr_scalar); } \
+    HWY_INLINE double operator()(double x) const { return (expr_scalar); } \
+};                                                                   \
+inline nb::ndarray<nb::numpy, float, nb::ndim<1>>                     \
+Symbol(nb::ndarray<float, nb::c_contig> a) {                         \
+    return unary<float, Symbol##Op>(a);                              \
+}                                                                    \
+inline nb::ndarray<nb::numpy, double, nb::ndim<1>>                    \
+Symbol(nb::ndarray<double, nb::c_contig> a) {                        \
+    return unary<double, Symbol##Op>(a);                             \
+}
+
+DEFINE_SIMD_UNARY_OP(exp, std::exp(x), hwy::HWY_NAMESPACE::Exp(d, v))
+DEFINE_SIMD_UNARY_OP(log, std::log(x), hwy::HWY_NAMESPACE::Log(d, v))
+DEFINE_SIMD_UNARY_OP(sqrt, std::sqrt(x), hwy::HWY_NAMESPACE::Sqrt(v))
+DEFINE_SIMD_UNARY_OP(sin, std::sin(x), hwy::HWY_NAMESPACE::Sin(d, v))
+DEFINE_SIMD_UNARY_OP(cos, std::cos(x), hwy::HWY_NAMESPACE::Cos(d, v))
+DEFINE_SIMD_UNARY_OP(asin, std::asin(x), hwy::HWY_NAMESPACE::Asin(d, v))
+DEFINE_SIMD_UNARY_OP(acos, std::acos(x), hwy::HWY_NAMESPACE::Acos(d, v))
+
+}  // capnhook
+}  // HWY_NAMESPACE
+}  // hwy
+HWY_AFTER_NAMESPACE();
+
+namespace capnhook = hwy::HWY_NAMESPACE::capnhook;