[#80]: Support larger tensors

.gitignore

@@ -10,3 +10,4 @@ __pycache__/
 *.pyo
 venv/
 .env/
+.codex

python/bindings.cpp

@@ -4,13 +4,30 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
+#include <algorithm>
+#include <stdexcept>
+
 using namespace pybind11::literals;
 namespace py = pybind11;
 
 // Generic templated function to bind tensor creation functions
 template <typename FuncType>
 void bind_tensor_creation_overloads(pybind11::module &m, const char *func_name,
                                     FuncType &&func) {
+  m.def(
+      func_name,
+      [func](const std::vector<int64_t> &dims,
+             smollnet::DataType dtype = smollnet::DataType::f32,
+             smollnet::Device device = smollnet::Device::CUDA,
+             bool requires_grad = false) {
+        if (dims.size() > smollnet::kMaxTensorDims) {
+          throw std::invalid_argument("Tensor rank exceeds kMaxTensorDims");
+        }
+        return func(dims.data(), dims.size(), dtype, device, requires_grad);
+      },
+      "dims"_a, "dtype"_a = smollnet::DataType::f32,
+      "device"_a = smollnet::Device::CUDA, "requires_grad"_a = false);
+
   // 1D version
   m.def(
       func_name,
@@ -55,6 +72,11 @@ void bind_tensor_creation_overloads(pybind11::module &m, const char *func_name,
 
 // Function objects for each tensor creation function
 struct RandFunctor {
+  auto operator()(const int64_t *dims, size_t rank, smollnet::DataType dtype,
+                  smollnet::Device device, bool requires_grad) const {
+    return smollnet::rand(dims, rank, dtype, device, requires_grad);
+  }
+
   template <size_t N>
   auto operator()(const int64_t (&dims)[N], smollnet::DataType dtype,
                   smollnet::Device device, bool requires_grad) const {
@@ -63,6 +85,11 @@ struct RandFunctor {
 };
 
 struct ZerosFunctor {
+  auto operator()(const int64_t *dims, size_t rank, smollnet::DataType dtype,
+                  smollnet::Device device, bool requires_grad) const {
+    return smollnet::zeros(dims, rank, dtype, device, requires_grad);
+  }
+
   template <size_t N>
   auto operator()(const int64_t (&dims)[N], smollnet::DataType dtype,
                   smollnet::Device device, bool requires_grad) const {
@@ -71,6 +98,11 @@ struct ZerosFunctor {
 };
 
 struct OnesFunctor {
+  auto operator()(const int64_t *dims, size_t rank, smollnet::DataType dtype,
+                  smollnet::Device device, bool requires_grad) const {
+    return smollnet::ones(dims, rank, dtype, device, requires_grad);
+  }
+
   template <size_t N>
   auto operator()(const int64_t (&dims)[N], smollnet::DataType dtype,
                   smollnet::Device device, bool requires_grad) const {
@@ -79,6 +111,11 @@ struct OnesFunctor {
 };
 
 struct EmptyFunctor {
+  auto operator()(const int64_t *dims, size_t rank, smollnet::DataType dtype,
+                  smollnet::Device device, bool requires_grad) const {
+    return smollnet::empty(dims, rank, dtype, device, requires_grad);
+  }
+
   template <size_t N>
   auto operator()(const int64_t (&dims)[N], smollnet::DataType dtype,
                   smollnet::Device device, bool requires_grad) const {
@@ -155,7 +192,15 @@ PYBIND11_MODULE(smollnet, m) {
       .def("matmul", &smollnet::Tensor::matmul)
 
       .def("transpose", &smollnet::Tensor::transpose)
-      .def("expand", &smollnet::Tensor::expand)
+      .def("expand",
+           [](const smollnet::Tensor &tensor, const std::vector<int64_t> &dims) {
+             if (dims.size() > smollnet::kMaxTensorDims) {
+               throw std::invalid_argument("Tensor rank exceeds kMaxTensorDims");
+             }
+             smollnet::TensorShape shape{};
+             std::copy(dims.begin(), dims.end(), shape.begin());
+             return tensor.expand(shape);
+           })
 
       .def("cuda", &smollnet::Tensor::cuda)
       .def("cpu", &smollnet::Tensor::cpu)

src/autograd.hpp

@@ -63,8 +63,8 @@ struct MatmulFunction : Function {
   void print() const override { printf("MatmulFunction\n"); }
 
 private:
-  std::array<int64_t, 3> lhs_shape;
-  std::array<int64_t, 3> rhs_shape;
+  TensorShape lhs_shape;
+  TensorShape rhs_shape;
 };
 
 struct ReLUFunction : Function {
@@ -103,7 +103,7 @@ struct SumFunction : Function {
 
 private:
   int64_t dim_;
-  std::array<int64_t, 3> input_shape_;
+  TensorShape input_shape_;
 };
 
 struct MseFunction : Function {

src/kernels.cuh

@@ -4,6 +4,8 @@
 #include <cstddef>
 #include <cstdint>
 
+#include "types.hpp"
+
 namespace smollnet {
 
 constexpr int32_t ROW_MAJOR = 0;
@@ -12,24 +14,24 @@ constexpr int32_t DEPTH_MAJOR = 2;
 
 
 struct StrideAndSize {
-  std::array<int64_t, 3> stride;
+  int64_t stride[kMaxTensorDims] = {};
 
   int64_t rank;
-  std::array<int64_t, 3> size;
+  int64_t size[kMaxTensorDims] = {};
 };
 
 struct StrideInfo {
   // size of the output operation
-  int64_t output_size[3];
+  int64_t output_size[kMaxTensorDims] = {};
 
-  int64_t a_stride[3];
-  int64_t b_stride[3];
+  int64_t a_stride[kMaxTensorDims] = {};
+  int64_t b_stride[kMaxTensorDims] = {};
   int64_t rank;
 };
 
 struct SizeInfo {
-  int64_t a_size[3];
-  int64_t b_size[3];
+  int64_t a_size[kMaxTensorDims] = {};
+  int64_t b_size[kMaxTensorDims] = {};
 };
 
 enum class WelfordType : uint8_t{
@@ -58,9 +60,8 @@ void launch_div(float *out, float *a, float *b, size_t numElems);
 void launch_div_strided(void *dst, void *a, void *b, const StrideInfo &s,
                         size_t total);
 
-void launch_sum_dim0(void *out, void *in, const StrideAndSize& s_input, const StrideAndSize& s_output);
-void launch_sum_dim1(void *out, void *in, const StrideAndSize& s_input, const StrideAndSize& s_output);
-void launch_sum_dim2(void *out, void *in, const StrideAndSize& s_input, const StrideAndSize& s_output);
+void launch_sum_dim(void *out, void *in, const StrideAndSize &s_input,
+                    const StrideAndSize &s_output, int64_t dim);
 
 void launch_matmul(void *out, void *left, void *right,
                    const StrideInfo &strides, const SizeInfo &sizes,

src/operators.cu

@@ -5,23 +5,19 @@
 
 namespace smollnet {
 
-__device__ __forceinline__ void compute_dimensions(int (&dims)[3], size_t idx,
-                                                   const StrideInfo &s) {
-
-  if (s.rank == 3) {
-    int64_t rest = s.output_size[1] * s.output_size[2];
-    dims[0] = idx / rest;
-    int64_t rem = idx % rest;
-    dims[1] = rem / s.output_size[2];
-    dims[2] = rem % s.output_size[2];
-  } else if (s.rank == 2) {
-    dims[0] = idx / s.output_size[1];
-    dims[1] = idx % s.output_size[1];
-    dims[2] = 0;
-  } else { // rank == 1
-    dims[0] = idx;
-    dims[1] = 0;
-    dims[2] = 0;
+__device__ __forceinline__ void compute_strided_offsets(size_t idx,
+                                                        const StrideInfo &s,
+                                                        int64_t &offA,
+                                                        int64_t &offB) {
+  offA = 0;
+  offB = 0;
+
+  for (int64_t dim = s.rank - 1; dim >= 0; --dim) {
+    const int64_t coord = idx % s.output_size[dim];
+    idx /= s.output_size[dim];
+
+    offA += coord * s.a_stride[dim];
+    offB += coord * s.b_stride[dim];
   }
 }
 
@@ -73,13 +69,9 @@ __global__ void add_strided_kernel(float *__restrict__ out,
   if (idx >= total)
     return;
 
-  int dims[3] = {0, 0, 0};
-  compute_dimensions(dims, idx, s);
-
-  int64_t offA = dims[0] * s.a_stride[0] + dims[1] * s.a_stride[1] +
-                 dims[2] * s.a_stride[2];
-  int64_t offB = dims[0] * s.b_stride[0] + dims[1] * s.b_stride[1] +
-                 dims[2] * s.b_stride[2];
+  int64_t offA = 0;
+  int64_t offB = 0;
+  compute_strided_offsets(idx, s, offA, offB);
 
   out[idx] = a[offA] + b[offB];
 }
@@ -124,13 +116,9 @@ __global__ void mul_strided_kernel(float *__restrict__ out,
   if (idx >= total)
     return;
 
-  int dims[3] = {0, 0, 0};
-  compute_dimensions(dims, idx, s);
-
-  int64_t offA = dims[0] * s.a_stride[0] + dims[1] * s.a_stride[1] +
-                 dims[2] * s.a_stride[2];
-  int64_t offB = dims[0] * s.b_stride[0] + dims[1] * s.b_stride[1] +
-                 dims[2] * s.b_stride[2];
+  int64_t offA = 0;
+  int64_t offB = 0;
+  compute_strided_offsets(idx, s, offA, offB);
 
   out[idx] = a[offA] * b[offB];
 }
@@ -167,13 +155,9 @@ __global__ void sub_strided_kernel(float *out, float *a, float *b, StrideInfo s,
   if (idx >= total)
     return;
 
-  int dims[3] = {0, 0, 0};
-  compute_dimensions(dims, idx, s);
-
-  int64_t offA = dims[0] * s.a_stride[0] + dims[1] * s.a_stride[1] +
-                 dims[2] * s.a_stride[2];
-  int64_t offB = dims[0] * s.b_stride[0] + dims[1] * s.b_stride[1] +
-                 dims[2] * s.b_stride[2];
+  int64_t offA = 0;
+  int64_t offB = 0;
+  compute_strided_offsets(idx, s, offA, offB);
 
   out[idx] = a[offA] - b[offB];
 }
@@ -210,13 +194,9 @@ __global__ void div_strided_kernel(float *__restrict__ out,
   if (idx >= total)
     return;
 
-  int dims[3] = {0, 0, 0};
-  compute_dimensions(dims, idx, s);
-
-  int64_t offA = dims[0] * s.a_stride[0] + dims[1] * s.a_stride[1] +
-                 dims[2] * s.a_stride[2];
-  int64_t offB = dims[0] * s.b_stride[0] + dims[1] * s.b_stride[1] +
-                 dims[2] * s.b_stride[2];
+  int64_t offA = 0;
+  int64_t offB = 0;
+  compute_strided_offsets(idx, s, offA, offB);
 
   out[idx] = a[offA] / b[offB];
 }

src/sgd.cpp

@@ -8,17 +8,16 @@ namespace smollnet {
 
 void SGD::step() const {
   for (const auto &p : params_) {
-    ASSERT(
-        p.size(0) == p.grad().size(0),
-        fmt::format("Size 0 mismatch!: {} vs {}", p.size(0), p.grad().size(0)));
-    ASSERT(
-        p.size(1) == p.grad().size(1),
-        fmt::format("Size 1 mismatch!: {} vs {}", p.size(1), p.grad().size(1)));
-    ASSERT(
-        p.size(2) == p.grad().size(2),
-        fmt::format("Size 2 mismatch!: {} vs {}", p.size(2), p.grad().size(2)));
+    Tensor grad = p.grad();
+    ASSERT(p.ndims() == grad.ndims(),
+           fmt::format("Rank mismatch!: {} vs {}", p.ndims(), grad.ndims()));
+    for (int64_t dim = 0; dim < p.ndims(); ++dim) {
+      ASSERT(p.size(dim) == grad.size(dim),
+             fmt::format("Size {} mismatch!: {} vs {}", dim, p.size(dim),
+                         grad.size(dim)));
+    }
 
-    launch_sgd_update(p.data(), p.grad().data(), lr_, p.numel());
+    launch_sgd_update(p.data(), grad.data(), lr_, p.numel());
   }
 }