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 bad796c..7f5c82d 100644 --- a/src/core/binary_ops.cpp +++ b/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 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; + } + 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 6e2a1cd..29611cd 100644 --- a/src/core/include/tensor.h +++ b/src/core/include/tensor.h @@ -103,6 +103,12 @@ class TensorStorage : public intrusive_ptr_target { } }; +class GradFunction : public intrusive_ptr_target { +public: + virtual std::vector backward(Tensor grad_output) = 0; + std::vector inputs; +}; + class Tensor { int dim_; dim_t shape_; @@ -110,8 +116,11 @@ 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; + 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 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(); } @@ -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 &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__':