[#60]: Improve performance of strided kernels

cmake/cuda.cmake

@@ -3,6 +3,7 @@ find_package(CUDAToolkit REQUIRED)
 add_library(${CUDA_LIBNAME}
   STATIC
     src/kernels.cu
+    src/operators.cu
     src/sum.cu
     src/welford.cu
     src/prng.cu

example/main.cpp

@@ -5,26 +5,25 @@
 using namespace smollnet;
 
 int main() {
-  constexpr int batch_size = 1024;
-  constexpr int num_features = 2048;
+  constexpr int input_size = 10;
 
   manual_seed(1234);
-  Tensor input = rand({batch_size, num_features}, DataType::f32, Device::CUDA);
- // input.print_elms();
-
-  Tensor result = mse(input,input);
-  // auto targets_h = targets.cpu();
-
-  // auto net = Dense(Linear(num_features, 64), LayerNorm(), GeLU(), Linear(64, 1));
-
-  // for (int epoch = 0; epoch < 64; ++epoch) {
-  //   auto res = net.forward(input);
-  //   auto loss = mse(res, targets);
-  //   fmt::print("epoch[{}]: Loss={}\n", epoch, static_cast<float*>(loss.cpu().data())[0]);
-  //   loss.backward();
-
-  //   auto optim = SGD(net.parameters(), 0.005f);
-  //   optim.step();
-  //   optim.zero_grad();
-  // }
+
+  Tensor input = rand({input_size, 128}, DataType::f32, Device::CUDA);
+  Tensor targets = rand({input_size, 1}, DataType::f32, Device::CUDA);
+  auto targets_h = targets.cpu();
+
+  auto net = Dense(Linear(128, 64), LayerNorm(), GeLU(), Linear(64, 1));
+
+  for (int epoch = 0; epoch < 64; ++epoch) {
+    auto res = net.forward(input);
+    auto loss = mse(res, targets);
+    fmt::print("epoch[{}]: Loss={}\n", epoch, static_cast<float*>(loss.cpu().data())[0]);
+    loss.backward();
+
+    auto optim = SGD(net.parameters(), 0.005f);
+    optim.step();
+    optim.zero_grad();
+  }
 }
+

python/bindings.cpp

@@ -122,11 +122,27 @@ PYBIND11_MODULE(smollnet, m) {
       .def("print", &smollnet::Tensor::print)
       .def("print_elms", &smollnet::Tensor::print_elms)
 
-      .def("add", &smollnet::Tensor::add)
-      .def("sub", &smollnet::Tensor::sub)
+      .def("neg", &smollnet::Tensor::neg)
+      .def("add", py::overload_cast<const smollnet::Tensor &>(
+                      &smollnet::Tensor::add, py::const_))
+      .def("add",
+           py::overload_cast<float>(&smollnet::Tensor::add, py::const_))
+      .def("sub", py::overload_cast<const smollnet::Tensor &>(
+                      &smollnet::Tensor::sub, py::const_))
+      .def("sub",
+           py::overload_cast<float>(&smollnet::Tensor::sub, py::const_))
+      .def("rsub", &smollnet::Tensor::rsub)
       .def("sum", &smollnet::Tensor::sum, py::arg("dim"),
            py::arg("keep_dim") = false)
-      .def("mul", &smollnet::Tensor::mul)
+      .def("mul", py::overload_cast<const smollnet::Tensor &>(
+                      &smollnet::Tensor::mul, py::const_))
+      .def("mul",
+           py::overload_cast<float>(&smollnet::Tensor::mul, py::const_))
+      .def("div", py::overload_cast<const smollnet::Tensor &>(
+                      &smollnet::Tensor::div, py::const_))
+      .def("div",
+           py::overload_cast<float>(&smollnet::Tensor::div, py::const_))
+      .def("rdiv", &smollnet::Tensor::rdiv)
       .def("matmul", &smollnet::Tensor::matmul)
 
       .def("transpose", &smollnet::Tensor::transpose)
@@ -136,12 +152,35 @@ PYBIND11_MODULE(smollnet, m) {
       .def("cpu", &smollnet::Tensor::cpu)
       .def("copy", &smollnet::Tensor::copy)
 
+      .def(-py::self)
       .def(py::self + py::self)
       .def(py::self - py::self)
       .def(py::self * py::self)
+      .def(py::self / py::self)
+      .def(py::self + float())
+      .def(float() + py::self)
+      .def(py::self - float())
+      .def(float() - py::self)
+      .def(py::self * float())
+      .def(float() * py::self)
+      .def(py::self / float())
+      .def(float() / py::self)
       .def(py::self += py::self)
       .def(py::self -= py::self)
-      .def(py::self *= py::self);
+      .def(py::self *= py::self)
+      .def(py::self /= py::self)
+      .def(py::self += float())
+      .def(py::self -= float())
+      .def(py::self *= float())
+      .def(py::self /= float())
+      .def("__matmul__",
+           [](const smollnet::Tensor &l, const smollnet::Tensor &r) {
+             return l.matmul(r);
+           })
+      .def("__rmatmul__",
+           [](const smollnet::Tensor &r, const smollnet::Tensor &l) {
+             return l.matmul(r);
+           });
 
   py::enum_<smollnet::DataType>(m, "DataType")
       .value("f32", smollnet::DataType::f32)
@@ -162,7 +201,11 @@ PYBIND11_MODULE(smollnet, m) {
   m.def("sigmoid", &smollnet::sigmoid);
 
   m.def("matmul", &smollnet::matmul);
+  m.def("neg", &smollnet::neg);
+  m.def("add", &smollnet::add);
+  m.def("sub", &smollnet::sub);
   m.def("mul", &smollnet::mul);
+  m.def("div", &smollnet::div);
   m.def("sum", &smollnet::sum, py::arg("tensor"), py::arg("dim"),
         py::arg("keep_dim") = false);
 

src/autograd.cpp

@@ -21,6 +21,17 @@ Tensor create_grad_tensor(const Tensor &tensor) {
                tensor.device());
 }
 
+Tensor reduce_broadcast_gradient(Tensor grad, const Tensor &input) {
+  auto sizes = input.dims();
+  for (int dim = 0; dim < input.ndims(); ++dim) {
+    if (sizes[dim] == 1 and grad.size(dim) > 1) {
+      grad = sum(grad, dim, true);
+    }
+  }
+
+  return grad;
+}
+
 // AddFunction implementation
 AddFunction::AddFunction(const Tensor &lhs, const Tensor &rhs) {
   inputs = {lhs, rhs};
@@ -155,6 +166,34 @@ MulFunction::backward(const std::vector<Tensor> &grad_outputs) {
   return grad_inputs;
 }
 
+DivFunction::DivFunction(const Tensor &lhs, const Tensor &rhs) {
+  inputs = {lhs, rhs};
+  needs_input_grad = {lhs.initialized() && lhs.requires_grad(),
+                      rhs.initialized() && rhs.requires_grad()};
+}
+
+std::vector<Tensor>
+DivFunction::backward(const std::vector<Tensor> &grad_outputs) {
+  ASSERT(grad_outputs.size() == 1,
+         "DivFunction expects exactly one gradient output");
+
+  std::vector<Tensor> grad_inputs(2);
+
+  if (needs_input_grad[0]) {
+    Tensor grad = grad_outputs[0] / inputs[1];
+    grad_inputs[0] = reduce_broadcast_gradient(grad, inputs[0]);
+  }
+
+  if (needs_input_grad[1]) {
+    Tensor grad = (grad_outputs[0] * inputs[0]) / (inputs[1] * inputs[1]);
+    grad = reduce_broadcast_gradient(grad, inputs[1]);
+    launch_negative(grad.data(), grad.numel());
+    grad_inputs[1] = grad;
+  }
+
+  return grad_inputs;
+}
+
 // MatmulFunction implementation
 MatmulFunction::MatmulFunction(const Tensor &lhs, const Tensor &rhs)
     : lhs_shape(lhs.dims()), rhs_shape(rhs.dims()) {

src/autograd.hpp

@@ -49,6 +49,13 @@ struct MulFunction : Function {
   void print() const override { printf("MulFunction\n"); }
 };
 
+struct DivFunction : Function {
+  DivFunction(const Tensor &lhs, const Tensor &rhs);
+  std::vector<Tensor>
+  backward(const std::vector<Tensor> &grad_outputs) override;
+  void print() const override { printf("DivFunction\n"); }
+};
+
 struct MatmulFunction : Function {
   MatmulFunction(const Tensor &lhs, const Tensor &rhs);
   std::vector<Tensor>