From c6cb39c16580bfb963330eb7c3a78f3740a53c62 Mon Sep 17 00:00:00 2001 From: xytpai Date: Tue, 5 Aug 2025 16:29:07 +0800 Subject: [PATCH 1/4] add func interface --- src/core/binary_ops.cpp | 21 +++++++++++++++++---- src/core/include/tensor.h | 8 ++++++++ 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/src/core/binary_ops.cpp b/src/core/binary_ops.cpp index bad796c..c4f206e 100644 --- a/src/core/binary_ops.cpp +++ b/src/core/binary_ops.cpp @@ -3,22 +3,35 @@ namespace gpu { + Tensor &add_out(Tensor &out, const Tensor &left, const Tensor &right) { auto iter = TensorIterator().add_output(out).add_input(left).add_input(right).build_for_loops(); add_kernel(iter); return out; } +Tensor &add_(Tensor &self, const Tensor &other) { + return add_out(self, self, other); +} + +class AddFunction : public Function { +public: + Tensor forward(std::vector inputs) override { + CHECK_FAIL(inputs.size() == 2, "AddFunction requires exactly two inputs."); + return add(inputs[0], inputs[1]); + } +private: + Tensor left_; + Tensor right_; + +} + Tensor add(const Tensor &left, const Tensor &right) { Tensor out; out = add_out(out, 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 8c1282a..71e5baa 100644 --- a/src/core/include/tensor.h +++ b/src/core/include/tensor.h @@ -114,6 +114,12 @@ class TensorStorage : public intrusive_ptr_target { } }; +class Function : public intrusive_ptr_target { +public: + virtual std::vector forward() = 0; + virtual void backward(std::vector grad_output) = 0; +}; + class Tensor { int dim_; dim_t shape_; @@ -121,6 +127,8 @@ class Tensor { ScalarType dtype_; int64_t numel_; intrusive_ptr storage_; + intrusive_ptr grad_storage_; + intrusive_ptr grad_fn_; int64_t storage_offset_ = 0; bool is_contiguous_ = true; From f569eca68b72ed6460b2689147ea4840134fc9fe Mon Sep 17 00:00:00 2001 From: xytpai Date: Fri, 8 Aug 2025 12:30:14 +0800 Subject: [PATCH 2/4] refine structure --- src/core/binary_ops.cpp | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/src/core/binary_ops.cpp b/src/core/binary_ops.cpp index c4f206e..a197890 100644 --- a/src/core/binary_ops.cpp +++ b/src/core/binary_ops.cpp @@ -3,7 +3,6 @@ namespace gpu { - Tensor &add_out(Tensor &out, const Tensor &left, const Tensor &right) { auto iter = TensorIterator().add_output(out).add_input(left).add_input(right).build_for_loops(); add_kernel(iter); @@ -16,20 +15,24 @@ Tensor &add_(Tensor &self, const Tensor &other) { class AddFunction : public Function { public: - Tensor forward(std::vector inputs) override { - CHECK_FAIL(inputs.size() == 2, "AddFunction requires exactly two inputs."); - return add(inputs[0], inputs[1]); + std::vector forward() override { + Tensor out; + return {add_out(out, left_, right_)}; + } + void backward(std::vector grad_output) override { } + AddFunction(const Tensor &left, const Tensor &right) : + left_(left), right_(right) { + } + private: - Tensor left_; - Tensor right_; - -} + const Tensor &left_; + const Tensor &right_; +}; Tensor add(const Tensor &left, const Tensor &right) { - Tensor out; - out = add_out(out, left, right); - return out; + auto fn = AddFunction(left, right); + return fn.forward()[0]; } Tensor &sub_out(Tensor &out, const Tensor &left, const Tensor &right) { From 8682f42650fe4f6bd6a00b010717bc9d425ab7fb Mon Sep 17 00:00:00 2001 From: xytpai Date: Fri, 8 Aug 2025 14:39:00 +0800 Subject: [PATCH 3/4] add basic autograd logic --- src/core/binary_ops.cpp | 38 ++++++++++++++++++++++++-------------- src/core/include/tensor.h | 39 ++++++++++++++++++++++++++++++++------- src/core/tensor.cpp | 26 ++++++++++++++++++++++++++ src/register.cpp | 5 ++++- test/test_tensor.py | 15 +++++++++++++++ 5 files changed, 101 insertions(+), 22 deletions(-) diff --git a/src/core/binary_ops.cpp b/src/core/binary_ops.cpp index a197890..366adb5 100644 --- a/src/core/binary_ops.cpp +++ b/src/core/binary_ops.cpp @@ -13,26 +13,36 @@ Tensor &add_(Tensor &self, const Tensor &other) { return add_out(self, self, other); } -class AddFunction : public Function { +class AddGradFunction : public GradFunction { public: - std::vector forward() override { - Tensor out; - return {add_out(out, left_, right_)}; + 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; + 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; } - void backward(std::vector grad_output) override { + AddGradFunction(const Tensor &left, const Tensor &right) { + inputs.reserve(2); + inputs[0] = left; + inputs[1] = right; } - AddFunction(const Tensor &left, const Tensor &right) : - left_(left), right_(right) { - } - -private: - const Tensor &left_; - const Tensor &right_; }; Tensor add(const Tensor &left, const Tensor &right) { - auto fn = AddFunction(left, right); - return fn.forward()[0]; + 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 &sub_out(Tensor &out, const Tensor &left, const Tensor &right) { diff --git a/src/core/include/tensor.h b/src/core/include/tensor.h index 90ce998..29611cd 100644 --- a/src/core/include/tensor.h +++ b/src/core/include/tensor.h @@ -103,10 +103,10 @@ class TensorStorage : public intrusive_ptr_target { } }; -class Function : public intrusive_ptr_target { +class GradFunction : public intrusive_ptr_target { public: - virtual std::vector forward() = 0; - virtual void backward(std::vector grad_output) = 0; + virtual std::vector backward(Tensor grad_output) = 0; + std::vector inputs; }; class Tensor { @@ -117,9 +117,10 @@ class Tensor { int64_t numel_; intrusive_ptr storage_; intrusive_ptr grad_storage_; - intrusive_ptr grad_fn_; + intrusive_ptr grad_fn_; int64_t storage_offset_ = 0; bool is_contiguous_ = true; + bool requires_grad_ = false; void new_storage_(int device); @@ -136,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_; @@ -151,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; @@ -207,9 +215,18 @@ class Tensor { intrusive_ptr storage() const { return storage_; } + intrusive_ptr grad_storage() const { + return grad_storage_; + } + intrusive_ptr 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(); } @@ -224,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 &indices) const; diff --git a/src/core/tensor.cpp b/src/core/tensor.cpp index c9c7cbe..68beae5 100644 --- a/src/core/tensor.cpp +++ b/src/core/tensor.cpp @@ -1,6 +1,8 @@ #include #include #include +#include +#include #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> 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()); } diff --git a/src/register.cpp b/src/register.cpp index 93aa0d4..8505031 100644 --- a/src/register.cpp +++ b/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); } diff --git a/test/test_tensor.py b/test/test_tensor.py index c111d41..e6c3800 100644 --- a/test/test_tensor.py +++ b/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__': From 08d133e351efdcb794f08616409934c904c0c8b0 Mon Sep 17 00:00:00 2001 From: xytpai Date: Wed, 13 Aug 2025 14:32:44 +0800 Subject: [PATCH 4/4] fix bugs --- CMakeLists.txt | 1 + src/core/binary_ops.cpp | 5 ++--- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 73d6c6b..f08d2cf 100644 --- a/CMakeLists.txt +++ b/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}) diff --git a/src/core/binary_ops.cpp b/src/core/binary_ops.cpp index 366adb5..7f5c82d 100644 --- a/src/core/binary_ops.cpp +++ b/src/core/binary_ops.cpp @@ -29,9 +29,8 @@ class AddGradFunction : public GradFunction { return grad_inputs; } AddGradFunction(const Tensor &left, const Tensor &right) { - inputs.reserve(2); - inputs[0] = left; - inputs[1] = right; + inputs.push_back(left); + inputs.push_back(right); } };