Add autograd support for simple binary ops

CMakeLists.txt

@@ -26,6 +26,7 @@ include_directories(${ROOTDIR}/src/core)
 include_directories(${ROOTDIR}/src/core/include)
 include_directories(${ROOTDIR}/src/core/utils)
 include_directories(${ROOTDIR}/src/core/utils/memory)
+# set(CMAKE_CXX_FLAGS  "${CMAKE_CXX_FLAGS} -g")
 file(GLOB CORE_SRCS "src/core/*.cpp")
 add_library(kfunca_core SHARED ${CORE_SRCS})
 

src/core/binary_ops.cpp

@@ -9,16 +9,41 @@ Tensor &add_out(Tensor &out, const Tensor &left, const Tensor &right) {
     return out;
 }
 
+Tensor &add_(Tensor &self, const Tensor &other) {
+    return add_out(self, self, other);
+}
+
+class AddGradFunction : public GradFunction {
+public:
+    std::vector<Tensor> backward(Tensor grad_output) override {
+        std::vector<Tensor> grad_inputs;
+        grad_inputs.resize(2);
+        if (inputs[0].requires_grad()) {
+            grad_inputs[0] = grad_output;
+            std::cout << "call AddGradFunction to left\n";
+        }
+        if (inputs[1].requires_grad()) {
+            grad_inputs[1] = grad_output;
+            std::cout << "call AddGradFunction to right\n";
+        }
+        return grad_inputs;
+    }
+    AddGradFunction(const Tensor &left, const Tensor &right) {
+        inputs.push_back(left);
+        inputs.push_back(right);
+    }
+};
+
 Tensor add(const Tensor &left, const Tensor &right) {
     Tensor out;
     out = add_out(out, left, right);
+    out.set_requires_grad((left.requires_grad() || right.requires_grad()));
+    if (out.requires_grad()) {
+        out.set_grad_fn(new AddGradFunction(left, right));
+    }
     return out;
 }
 
-Tensor &add_(Tensor &self, const Tensor &other) {
-    return add_out(self, self, other);
-}
-
 Tensor &sub_out(Tensor &out, const Tensor &left, const Tensor &right) {
     auto iter = TensorIterator().add_output(out).add_input(left).add_input(right).build_for_loops();
     sub_kernel(iter);

src/core/include/tensor.h

@@ -103,15 +103,24 @@ class TensorStorage : public intrusive_ptr_target {
     }
 };
 
+class GradFunction : public intrusive_ptr_target {
+public:
+    virtual std::vector<Tensor> backward(Tensor grad_output) = 0;
+    std::vector<Tensor> inputs;
+};
+
 class Tensor {
     int dim_;
     dim_t shape_;
     dim_t stride_;
     ScalarType dtype_;
     int64_t numel_;
     intrusive_ptr<TensorStorage> storage_;
+    intrusive_ptr<TensorStorage> grad_storage_;
+    intrusive_ptr<GradFunction> grad_fn_;
     int64_t storage_offset_ = 0;
     bool is_contiguous_ = true;
+    bool requires_grad_ = false;
 
     void new_storage_(int device);
 
@@ -128,13 +137,17 @@ class Tensor {
 
 public:
     Tensor() :
-        storage_() {
+        storage_(), grad_storage_(), grad_fn_() {
     }
     Tensor(const Tensor &other) :
         dim_(other.dim_), shape_(other.shape_), stride_(other.stride_),
         dtype_(other.dtype_), numel_(other.numel_),
-        storage_(other.storage_), storage_offset_(other.storage_offset_),
-        is_contiguous_(other.is_contiguous_) {
+        storage_(other.storage_),
+        grad_storage_(other.grad_storage_),
+        grad_fn_(other.grad_fn_),
+        storage_offset_(other.storage_offset_),
+        is_contiguous_(other.is_contiguous_),
+        requires_grad_(other.requires_grad_) {
     }
     Tensor &operator=(const Tensor &other) {
         dim_ = other.dim_;
@@ -143,8 +156,11 @@ class Tensor {
         dtype_ = other.dtype_;
         numel_ = other.numel_;
         storage_ = other.storage_;
+        grad_storage_ = other.grad_storage_;
+        grad_fn_ = other.grad_fn_;
         storage_offset_ = other.storage_offset_;
         is_contiguous_ = other.is_contiguous_;
+        requires_grad_ = other.requires_grad_;
         return *this;
     }
     Tensor(Tensor &&other) = default;
@@ -199,9 +215,18 @@ class Tensor {
     intrusive_ptr<TensorStorage> storage() const {
         return storage_;
     }
+    intrusive_ptr<TensorStorage> grad_storage() const {
+        return grad_storage_;
+    }
+    intrusive_ptr<GradFunction> grad_fn() const {
+        return grad_fn_;
+    }
     bool defined() const {
         return storage_.get() != nullptr;
     }
+    bool has_grad_fn() const {
+        return grad_fn_.get() != nullptr;
+    }
     int device() const {
         return storage_.get()->device();
     }
@@ -216,7 +241,15 @@ class Tensor {
     bool is_contiguous() const {
         return is_contiguous_;
     }
+    bool requires_grad() const {
+        return requires_grad_;
+    }
+    void set_requires_grad(bool flag) {
+        requires_grad_ = flag;
+    }
 
+    void set_grad_fn(GradFunction *fn);
+    void backward(Tensor grad_output);
     void copy_from_cpu_ptr(void *ptr);
     void copy_to_cpu_ptr(void *ptr) const;
     any_t item(const std::vector<int64_t> &indices) const;

src/core/tensor.cpp

@@ -1,6 +1,8 @@
 #include <iostream>
 #include <iomanip>
 #include <vector>
+#include <tuple>
+#include <queue>
 
 #include "tensor.h"
 #include "tensor_shape.h"
@@ -79,6 +81,30 @@ Tensor::Tensor(int64_t *shape, int ndim, ScalarType dtype, bool inverse) {
     }
 }
 
+void Tensor::set_grad_fn(GradFunction *fn) {
+    grad_fn_.unsafe_set_ptr(fn);
+}
+
+void Tensor::backward(Tensor grad_output) {
+    std::queue<std::tuple<Tensor *, Tensor>> q;
+    q.push({this, grad_output});
+    while (!q.empty()) {
+        auto [t, go] = q.front();
+        q.pop();
+        if (t->has_grad_fn()) {
+            auto grad_fn = t->grad_fn_.get();
+            auto gis = grad_fn->backward(go);
+            auto &inputs = grad_fn->inputs;
+            for (int i = 0; i < inputs.size(); ++i) {
+                if (inputs[i].requires_grad()) {
+                    // update grad input
+                    q.push({&inputs[i], gis[i]});
+                }
+            }
+        }
+    }
+}
+
 void Tensor::copy_from_cpu_ptr(void *ptr) {
     dmemcpy_h2d(data_ptr(), ptr, storage_bytes());
 }

src/register.cpp

@@ -210,5 +210,8 @@ PYBIND11_MODULE(kfunca, m) {
         .def("index_put_", &Tensor::index_put_)
         .def("half", &Tensor::_half)
         .def("bfloat16", &Tensor::_bfloat16)
-        .def("float", &Tensor::_float);
+        .def("float", &Tensor::_float)
+        .def("requires_grad", &Tensor::requires_grad)
+        .def("set_requires_grad", &Tensor::set_requires_grad)
+        .def("backward", &Tensor::backward);
 }

test/test_tensor.py

@@ -281,6 +281,21 @@ def test_index_put(self):
         values_pt = torch.from_numpy(values.numpy())
         arr_gpu_pt.index_put_(indices_t, values_pt)
         assert_allclose(arr_gpu, arr_gpu_pt)
+
+    def test_basic_backward(self):
+        grad = np.random.uniform(-10, 10, size=(2, 3)).astype(np.float32)
+        grad_gpu = kfunca.from_numpy(grad, 0)
+
+        arr1 = np.random.uniform(-10, 10, size=(2, 3)).astype(np.float32)
+        arr_gpu1 = kfunca.from_numpy(arr1, 0)
+        arr_gpu1.set_requires_grad(True)
+
+        arr2 = np.random.uniform(-10, 10, size=(2, 3)).astype(np.float32)
+        arr_gpu2 = kfunca.from_numpy(arr2, 0)
+        arr_gpu2.set_requires_grad(False)
+
+        arr_gpu3 = arr_gpu1 + arr_gpu2
+        arr_gpu3.backward(grad_gpu)
 
 
 if __name__ == '__main__':