diff --git a/paddle/phi/kernels/sparse/elementwise_grad_kernel.h b/paddle/phi/kernels/sparse/elementwise_grad_kernel.h index df3feb597e3c61dee51173b15f469461b48781c9..19b56d3c422720ba243e8f2a23d968e628062f87 100644 --- a/paddle/phi/kernels/sparse/elementwise_grad_kernel.h +++ b/paddle/phi/kernels/sparse/elementwise_grad_kernel.h @@ -15,6 +15,7 @@ limitations under the License. */ #pragma once #include "paddle/phi/core/dense_tensor.h" +#include "paddle/phi/core/sparse_coo_tensor.h" #include "paddle/phi/core/sparse_csr_tensor.h" #include "paddle/phi/kernels/empty_kernel.h" diff --git a/paddle/phi/kernels/sparse/gpu/elementwise_grad_kernel.cu b/paddle/phi/kernels/sparse/gpu/elementwise_grad_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..e434dad588e134f2528294bc88152ce3c828e05f --- /dev/null +++ b/paddle/phi/kernels/sparse/gpu/elementwise_grad_kernel.cu @@ -0,0 +1,56 @@ +/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. */ + +#include "paddle/phi/kernels/sparse/elementwise_grad_kernel.h" +#include "paddle/phi/core/kernel_registry.h" +#include "paddle/phi/core/tensor_utils.h" +#include "paddle/phi/kernels/sparse/empty_kernel.h" + +namespace phi { +namespace sparse { + +template +void ElementWiseAddCooGradKernel(const Context& dev_ctx, + const SparseCooTensor& x, + const SparseCooTensor& y, + const SparseCooTensor& dout, + SparseCooTensor* dx, + SparseCooTensor* dy) { + if (dx) { + EmptyLikeCooKernel(dev_ctx, x, dx); + Copy(dev_ctx, dout, dev_ctx.GetPlace(), false, dx); + } + + if (dy) { + EmptyLikeCooKernel(dev_ctx, y, dy); + Copy(dev_ctx, dout, dev_ctx.GetPlace(), false, dy); + } +} + +} // namespace sparse +} // namespace phi + +PD_REGISTER_KERNEL(add_coo_coo_grad, + GPU, + ALL_LAYOUT, + phi::sparse::ElementWiseAddCooGradKernel, + float, + double, + int16_t, + int, + int64_t, + phi::dtype::float16) { + kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); + kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO); +} diff --git a/paddle/phi/kernels/sparse/gpu/elementwise_kernel.cu b/paddle/phi/kernels/sparse/gpu/elementwise_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..7496f47de894887fc8ee923d00d5cc1966b60f79 --- /dev/null +++ b/paddle/phi/kernels/sparse/gpu/elementwise_kernel.cu @@ -0,0 +1,88 @@ +/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. */ + +#include +#include + +#include "paddle/phi/kernels/elementwise_add_kernel.h" +#include "paddle/phi/kernels/sparse/elementwise_kernel.h" +#include "paddle/phi/kernels/sparse/empty_kernel.h" + +#include "paddle/phi/core/enforce.h" +#include "paddle/phi/core/kernel_registry.h" +#include "paddle/phi/core/visit_type.h" + +namespace phi { +namespace sparse { + +template +void ElementWiseAddCooGPUKernel(const GPUContext& dev_ctx, + const SparseCooTensor& x, + const SparseCooTensor& y, + SparseCooTensor* out) { + const auto& x_indices = x.indices(); + const auto& y_indices = y.indices(); + PADDLE_ENFORCE_EQ( + x_indices.numel(), + y_indices.numel(), + phi::errors::PreconditionNotMet( + "The numel of x.indices() and y.indices() should be equal")); + const IntT* x_indices_ptr = x_indices.data(); + const IntT* y_indices_ptr = y_indices.data(); +#ifdef PADDLE_WITH_HIP + bool is_same = thrust::equal(thrust::hip::par.on(dev_ctx.stream()), +#else + bool is_same = thrust::equal(thrust::cuda::par.on(dev_ctx.stream()), +#endif + x_indices_ptr, + x_indices_ptr + x_indices.numel(), + y_indices_ptr); + PADDLE_ENFORCE_EQ( + is_same, + true, + phi::errors::PreconditionNotMet( + "Currently, ElementWiseAddCooKernel only supports the case " + "where x and y have the same indices")); + EmptyLikeCooKernel(dev_ctx, x, out); + phi::AddKernel( + dev_ctx, x.values(), y.values(), out->mutable_values()); +} + +template +void ElementWiseAddCooKernel(const Context& dev_ctx, + const SparseCooTensor& x, + const SparseCooTensor& y, + SparseCooTensor* out) { + PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "VerifyIndices", ([&] { + ElementWiseAddCooGPUKernel( + dev_ctx, x, y, out); + })); +} + +} // namespace sparse +} // namespace phi + +PD_REGISTER_KERNEL(add_coo_coo, + GPU, + ALL_LAYOUT, + phi::sparse::ElementWiseAddCooKernel, + float, + double, + int16_t, + int, + int64_t, + phi::dtype::float16) { + kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); + kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO); +} diff --git a/python/paddle/fluid/tests/unittests/test_sparse_elementwise_op.py b/python/paddle/fluid/tests/unittests/test_sparse_elementwise_op.py index 149c4cfb22b9931f99e36c41eab856460d69d95a..20f66e5f9a65e060ff2a2bda78fb11ca2ca3e245 100644 --- a/python/paddle/fluid/tests/unittests/test_sparse_elementwise_op.py +++ b/python/paddle/fluid/tests/unittests/test_sparse_elementwise_op.py @@ -18,7 +18,7 @@ from operator import __add__, __sub__, __mul__, __truediv__ import numpy as np import paddle -from paddle.fluid.framework import _test_eager_guard +import paddle.incubate.sparse as sparse op_list = [__add__, __sub__, __mul__, __truediv__] @@ -134,6 +134,35 @@ class TestSparseElementWiseAPI(unittest.TestCase): for op in op_list: self.func_test_coo(op) + def test_add_same_indices(self): + indices_data = [[0, 1], [0, 3]] + values1_data = [[1.0], [2.0]] + values2_data = [[1.0], [2.0]] + shape = [2, 4, 2] + + sp_a = sparse.sparse_coo_tensor(indices_data, + values1_data, + shape, + stop_gradient=False) + sp_b = sparse.sparse_coo_tensor(indices_data, + values2_data, + shape, + stop_gradient=False) + + values1 = paddle.to_tensor(values1_data, stop_gradient=False) + values2 = paddle.to_tensor(values2_data, stop_gradient=False) + + #c.values() = a.values() + b.values() + sp_c = sparse.add(sp_a, sp_b) + sp_c.backward() + ref_c = values1 + values2 + ref_c.backward() + np.testing.assert_allclose(sp_c.values().numpy(), ref_c.numpy()) + np.testing.assert_allclose(sp_a.grad.values().numpy(), + values1.grad.numpy()) + np.testing.assert_allclose(sp_b.grad.values().numpy(), + values2.grad.numpy()) + if __name__ == "__main__": paddle.device.set_device('cpu')