From b4bda95077a469b590c74e127053953db7b64823 Mon Sep 17 00:00:00 2001 From: xytpai Date: Fri, 15 Aug 2025 03:43:39 +0800 Subject: [PATCH] add autograd binary add --- src/core/binary_ops.cpp | 31 ++++++++++++++--- src/core/include/tensor.h | 25 ++++++++++++++ src/core/include/tensor_impl.h | 25 ++++++++++++-- src/core/tensor.cpp | 63 ++++++++++++++++++++++++++++++++++ src/core/tensor_impl.cpp | 9 +++-- src/register.cpp | 12 ++++++- test/test_tensor.py | 25 ++++++++++++++ 7 files changed, 179 insertions(+), 11 deletions(-) diff --git a/src/core/binary_ops.cpp b/src/core/binary_ops.cpp index bad796c..8dee420 100644 --- a/src/core/binary_ops.cpp +++ b/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 backward(Tensor grad_output) override { + std::vector 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); diff --git a/src/core/include/tensor.h b/src/core/include/tensor.h index 40a098b..0d0f202 100644 --- a/src/core/include/tensor.h +++ b/src/core/include/tensor.h @@ -15,9 +15,16 @@ Tensor empty_like_reduced(const Tensor &self, int dim, ScalarType dtype); Tensor zeros(std::vector 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 backward(Tensor grad_output) = 0; + std::vector inputs; +}; + class Tensor { private: intrusive_ptr impl_; + intrusive_ptr grad_fn_; friend Tensor empty(std::vector 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 storage() const { return impl_.get()->storage(); } + intrusive_ptr 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 &indices) const; diff --git a/src/core/include/tensor_impl.h b/src/core/include/tensor_impl.h index f0bc229..9d7a022 100644 --- a/src/core/include/tensor_impl.h +++ b/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 storage_; int64_t storage_offset_ = 0; bool is_contiguous_ = true; + bool requires_grad_ = false; public: TensorImpl(std::vector &shape, ScalarType dtype); TensorImpl(std::vector &shape, std::vector &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 sizes, std::vector strides, int64_t storage_offset = 0); + +public: + std::unique_ptr grad_; }; diff --git a/src/core/tensor.cpp b/src/core/tensor.cpp index aaad7d8..f5a365c 100644 --- a/src/core/tensor.cpp +++ b/src/core/tensor.cpp @@ -1,4 +1,6 @@ +#include #include +#include #include "tensor.h" #include "tensor_shape.h" @@ -66,6 +68,63 @@ Tensor zeros(std::vector 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(new Tensor(grad_), [](Tensor *t) { delete t; }); + } +} + +void Tensor::backward(Tensor grad_output) { + std::queue ready; + std::unordered_map needed; + std::unordered_map 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 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) diff --git a/src/core/tensor_impl.cpp b/src/core/tensor_impl.cpp index 9223c3d..2d08658 100644 --- a/src/core/tensor_impl.cpp +++ b/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 &shape, ScalarType dtype) { +TensorImpl::TensorImpl(std::vector &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 &shape, ScalarType dtype) { } } -TensorImpl::TensorImpl(std::vector &shape, std::vector &strides, ScalarType dtype) { +TensorImpl::TensorImpl(std::vector &shape, std::vector &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 &shape, std::vector &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; diff --git a/src/register.cpp b/src/register.cpp index e3a41cc..1dd7fa4 100644 --- a/src/register.cpp +++ b/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(); + } + }); } diff --git a/test/test_tensor.py b/test/test_tensor.py index c43a517..e5063c1 100644 --- a/test/test_tensor.py +++ b/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__':