Add autograd support for simple binary ops

src/core/binary_ops.cpp

@@ -9,16 +9,39 @@ 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;
+        }
+        if (inputs[1].requires_grad()) {
+            grad_inputs[1] = grad_output;
+        }
+        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

@@ -15,9 +15,16 @@ Tensor empty_like_reduced(const Tensor &self, int dim, ScalarType dtype);
 Tensor zeros(std::vector<int64_t> shape, ScalarType dtype, int device = 0);
 std::ostream &operator<<(std::ostream &os, const Tensor &t);
 
+class GradFunction : public intrusive_ptr_target {
+public:
+    virtual std::vector<Tensor> backward(Tensor grad_output) = 0;
+    std::vector<Tensor> inputs;
+};
+
 class Tensor {
 private:
     intrusive_ptr<TensorImpl> impl_;
+    intrusive_ptr<GradFunction> grad_fn_;
     friend Tensor empty(std::vector<int64_t> shape, ScalarType dtype, int device);
     friend Tensor empty(int64_t *shape, int ndim, ScalarType dtype, int device, bool inverse);
     friend Tensor empty_like(const Tensor &self);
@@ -84,9 +91,15 @@ class Tensor {
     intrusive_ptr<TensorStorage> storage() const {
         return impl_.get()->storage();
     }
+    intrusive_ptr<GradFunction> grad_fn() const {
+        return grad_fn_;
+    }
     bool defined() const {
         return impl_.get() && impl_.get()->defined();
     }
+    bool has_grad_fn() const {
+        return grad_fn_.get() != nullptr;
+    }
     int device() const {
         return impl_.get()->device();
     }
@@ -101,7 +114,19 @@ class Tensor {
     bool is_contiguous() const {
         return impl_.get()->is_contiguous();
     }
+    bool requires_grad() const {
+        return impl_.get()->requires_grad();
+    }
+    void set_requires_grad(bool flag) {
+        return impl_.get()->set_requires_grad(flag);
+    }
+    Tensor *grad() {
+        return impl_.get()->grad_.get();
+    }
 
+    void set_grad_fn(GradFunction *fn);
+    void update_grad(Tensor grad);
+    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/include/tensor_impl.h

@@ -91,6 +91,11 @@ class TensorStorage : public intrusive_ptr_target {
     }
 };
 
+class Tensor;
+using TensorDeleter = void (*)(Tensor *);
+inline void delete_nothing(Tensor *) {
+}
+
 class TensorImpl : public intrusive_ptr_target {
 private:
     int dim_;
@@ -101,19 +106,22 @@ class TensorImpl : public intrusive_ptr_target {
     intrusive_ptr<TensorStorage> storage_;
     int64_t storage_offset_ = 0;
     bool is_contiguous_ = true;
+    bool requires_grad_ = false;
 
 public:
     TensorImpl(std::vector<int64_t> &shape, ScalarType dtype);
     TensorImpl(std::vector<int64_t> &shape, std::vector<int64_t> &strides, ScalarType dtype);
     TensorImpl(int64_t *shape, int ndim, ScalarType dtype, bool inverse);
     TensorImpl() :
-        storage_() {
+        storage_(), grad_(nullptr, &delete_nothing) {
     }
     TensorImpl(const TensorImpl &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_(nullptr, &delete_nothing),
+        storage_offset_(other.storage_offset_),
+        is_contiguous_(other.is_contiguous_),
+        requires_grad_(other.requires_grad_) {
     }
     TensorImpl &operator=(const TensorImpl &other) {
         dim_ = other.dim_;
@@ -122,8 +130,10 @@ class TensorImpl : public intrusive_ptr_target {
         dtype_ = other.dtype_;
         numel_ = other.numel_;
         storage_ = other.storage_;
+        grad_ = nullptr;
         storage_offset_ = other.storage_offset_;
         is_contiguous_ = other.is_contiguous_;
+        requires_grad_ = other.requires_grad_;
         return *this;
     }
     TensorImpl(TensorImpl &&other) = default;
@@ -190,6 +200,15 @@ class TensorImpl : public intrusive_ptr_target {
     bool is_contiguous() const {
         return is_contiguous_;
     }
+    bool requires_grad() const {
+        return requires_grad_;
+    }
+    void set_requires_grad(bool flag) {
+        requires_grad_ = flag;
+    }
     void new_storage_(int device);
     void as_strided_(std::vector<int64_t> sizes, std::vector<int64_t> strides, int64_t storage_offset = 0);
+
+public:
+    std::unique_ptr<Tensor, TensorDeleter> grad_;
 };

src/core/tensor.cpp

@@ -1,4 +1,6 @@
+#include <queue>
 #include <iomanip>
+#include <unordered_map>
 
 #include "tensor.h"
 #include "tensor_shape.h"
@@ -66,6 +68,63 @@ Tensor zeros(std::vector<int64_t> shape, ScalarType dtype, int device) {
     return output;
 }
 
+void Tensor::set_grad_fn(GradFunction *fn) {
+    grad_fn_.unsafe_set_ptr(fn);
+}
+
+void Tensor::update_grad(Tensor grad) {
+    auto impl = this->impl_.get();
+    if (impl->grad_.get()) {
+        *impl->grad_ += grad;
+    } else {
+        Tensor grad_ = empty_like(grad);
+        grad_.copy_(grad);
+        impl->grad_ = std::unique_ptr<Tensor, TensorDeleter>(new Tensor(grad_), [](Tensor *t) { delete t; });
+    }
+}
+
+void Tensor::backward(Tensor grad_output) {
+    std::queue<Tensor *> ready;
+    std::unordered_map<TensorImpl *, int> needed;
+    std::unordered_map<TensorImpl *, Tensor> grad_acc;
+    // build needed counts
+    ready.push(this);
+    while (!ready.empty()) {
+        auto t = ready.front();
+        ready.pop();
+        if (t->has_grad_fn()) {
+            for (auto &input : t->grad_fn_.get()->inputs) {
+                if (!input.requires_grad()) continue;
+                needed[input.impl()] += 1;
+                ready.push(&input);
+            }
+        }
+    }
+    // calculate gradients
+    grad_acc[this->impl()] = grad_output;
+    ready.push(this);
+    while (!ready.empty()) {
+        auto t = ready.front();
+        ready.pop();
+        auto go = grad_acc[t->impl()];
+        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()) continue;
+                auto &acc = grad_acc[inputs[i].impl()];
+                acc = acc.defined() ? (acc + gis[i]) : gis[i];
+                if (--needed[inputs[i].impl()] == 0) {
+                    ready.push(&inputs[i]);
+                }
+            }
+        } else if (t->requires_grad()) {
+            t->update_grad(go);
+        }
+    }
+}
+
 void Tensor::copy_from_cpu_ptr(void *ptr) {
     dmemcpy_h2d(data_ptr(), ptr, storage_bytes());
 }
@@ -297,6 +356,10 @@ void print_tensor_(std::ostream &os, const Tensor &t, std::vector<int64_t> indic
 }
 
 std::ostream &operator<<(std::ostream &os, const Tensor &t) {
+    if (!t.defined()) {
+        os << "Tensor(Undefined)";
+        return os;
+    }
     t.data_ptr();
     os << "tensor(shape=[";
     for (int i = 0; i < t.dim(); ++i)

src/core/tensor_impl.cpp

@@ -8,7 +8,8 @@ std::ostream &operator<<(std::ostream &os, const dim_t &d) {
     return os;
 }
 
-TensorImpl::TensorImpl(std::vector<int64_t> &shape, ScalarType dtype) {
+TensorImpl::TensorImpl(std::vector<int64_t> &shape, ScalarType dtype) :
+    grad_(nullptr, &delete_nothing) {
     CHECK_FAIL(shape.size() <= MAX_TENSOR_DIMS);
     dtype_ = dtype;
     dim_ = shape.size();
@@ -20,7 +21,8 @@ TensorImpl::TensorImpl(std::vector<int64_t> &shape, ScalarType dtype) {
     }
 }
 
-TensorImpl::TensorImpl(std::vector<int64_t> &shape, std::vector<int64_t> &strides, ScalarType dtype) {
+TensorImpl::TensorImpl(std::vector<int64_t> &shape, std::vector<int64_t> &strides, ScalarType dtype) :
+    grad_(nullptr, &delete_nothing) {
     CHECK_FAIL(shape.size() <= MAX_TENSOR_DIMS);
     CHECK_FAIL(strides.size() <= MAX_TENSOR_DIMS);
     dtype_ = dtype;
@@ -34,7 +36,8 @@ TensorImpl::TensorImpl(std::vector<int64_t> &shape, std::vector<int64_t> &stride
     }
 }
 
-TensorImpl::TensorImpl(int64_t *shape, int ndim, ScalarType dtype, bool inverse) {
+TensorImpl::TensorImpl(int64_t *shape, int ndim, ScalarType dtype, bool inverse) :
+    grad_(nullptr, &delete_nothing) {
     CHECK_FAIL(ndim <= MAX_TENSOR_DIMS);
     dtype_ = dtype;
     dim_ = ndim;

src/register.cpp

@@ -211,5 +211,15 @@ 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)
+        .def("grad", [](Tensor &self) {
+            if (self.grad() && self.grad()->defined()) {
+                return *self.grad();
+            } else {
+                return Tensor();
+            }
+        });
 }

test/test_tensor.py

@@ -282,6 +282,31 @@ 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
+        grad_ = np.random.uniform(-10, 10, size=(2, 3)).astype(np.float32)
+        grad = kfunca.from_numpy(grad_, 0)
+        # a
+        a_ = np.random.uniform(-10, 10, size=(2, 3)).astype(np.float32)
+        a = kfunca.from_numpy(a_, 0)
+        a.set_requires_grad(True)
+        # b
+        b_ = np.random.uniform(-10, 10, size=(2, 3)).astype(np.float32)
+        b = kfunca.from_numpy(b_, 0)
+        b.set_requires_grad(True)
+        # c
+        c_ = np.random.uniform(-10, 10, size=(2, 3)).astype(np.float32)
+        c = kfunca.from_numpy(c_, 0)
+        # cal
+        ca = c + a
+        ab = a + b
+        accb = ca + ab
+        accba = accb + a
+        # backward
+        accba.backward(grad)
+        assert_allclose(a.grad(), grad * 3)
+        assert_allclose(b.grad(), grad)
 
 
 if __name__ == '__main__':