【Hackathon No.60】refactor unary sparse ops and add sparse sqrt, tanh, sin

paddle/phi/api/lib/api_gen_utils.cc

@@ -154,7 +154,7 @@ phi::TensorBase* SetSparseKernelOutput(Tensor* out, TensorType type) {
           std::make_shared<phi::SparseCsrTensor>(phi::DenseTensor(),
                                                  phi::DenseTensor(),
                                                  phi::DenseTensor(),
-                                                 phi::DDim{-1});
+                                                 phi::DDim{-1, -1});
       out->set_impl(sparse_tensor);
       return sparse_tensor.get();
     } else {

paddle/phi/core/sparse_csr_tensor.cc

@@ -27,9 +27,11 @@ SparseCsrTensor::SparseCsrTensor() {
 inline void check_shape(const DDim& dims) {
   bool valid = dims.size() == 2 || dims.size() == 3;
 
-  PADDLE_ENFORCE(valid,
-                 phi::errors::InvalidArgument(
-                     "the SparseCsrTensor only support 2-D Tensor."));
+  PADDLE_ENFORCE(
+      valid,
+      phi::errors::InvalidArgument("the SparseCsrTensor only support 2-D or "
+                                   "3-D Tensor, but get %d-D Tensor",
+                                   dims.size()));
 }
 #define Check(non_zero_crows, non_zero_cols, non_zero_elements, dims)          \
   {                                                                            \

paddle/phi/kernels/activation_grad_kernel.h

@@ -187,6 +187,7 @@ DECLARE_ACTIVATION_GRAD_KERNEL_DEPX(Log1p);
 DECLARE_ACTIVATION_GRAD_KERNEL_DEPOUT(Relu);
 DECLARE_ACTIVATION_GRAD_KERNEL_DEPOUT(Tanh);
 DECLARE_ACTIVATION_GRAD_KERNEL_DEPOUT(Sigmoid);
+DECLARE_ACTIVATION_GRAD_KERNEL_DEPOUT(Sqrt);
 
 DECLARE_ACTIVATION_GRAD_KERNEL_NODEP(Round);
 DECLARE_ACTIVATION_GRAD_KERNEL_NODEP(Floor);

paddle/phi/kernels/sparse/activation_grad_kernel.cc

@@ -13,58 +13,53 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/phi/kernels/sparse/activation_grad_kernel.h"
-#include "paddle/phi/kernels/activation_grad_kernel.h"
-#include "paddle/phi/kernels/copy_kernel.h"
-#include "paddle/phi/kernels/empty_kernel.h"
-
-#include "paddle/phi/backends/cpu/cpu_context.h"
-#include "paddle/phi/backends/gpu/gpu_context.h"
-#include "paddle/phi/core/kernel_registry.h"
-
-namespace phi {
-namespace sparse {
 
-template <typename T, typename Context>
-void SparseReluGradKernel(const Context& dev_ctx,
-                          const SparseCooTensor& x,
-                          const SparseCooTensor& out_grad,
-                          SparseCooTensor* x_grad) {
-  DenseTensor non_zero_indices =
-      phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_indices());
-  DenseTensor non_zero_elements =
-      phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_elements());
-  phi::Copy(dev_ctx,
-            x.non_zero_indices(),
-            dev_ctx.GetPlace(),
-            false,
-            &non_zero_indices);
-  phi::ReluGradKernel<T, Context>(dev_ctx,
-                                  x.non_zero_elements(),
-                                  out_grad.non_zero_elements(),
-                                  &non_zero_elements);
-  x_grad->SetMember(non_zero_indices, non_zero_elements, x.dims(), true);
-}
+#include "paddle/phi/kernels/activation_grad_kernel.h"
+#include "paddle/phi/kernels/sparse/utils.h"
 
-}  // namespace sparse
-}  // namespace phi
+DEFINE_AND_REGISTER_SPARSE_UNARY_GRAD_KERNEL(sqrt_grad, SqrtGradKernel)
 
-PD_REGISTER_KERNEL(sparse_relu_grad,
+// NOTE: the following code is to bypass the restriction of Paddle
+// kernel registration mechanism. Do NOT refactor them unless you
+// know what you are doing.
+// If you want to implement any new kernel, please follow `sqrt_grad` above
+// instead of `relu_grad` following
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(ReluGradKernel)
+PD_REGISTER_KERNEL(sparse_coo_relu_grad,
                    CPU,
                    ALL_LAYOUT,
-                   phi::sparse::SparseReluGradKernel,
+                   phi::sparse::SparseCooReluGradKernel,
                    float,
                    double) {
   kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
 }
-
+PD_REGISTER_KERNEL(sparse_csr_relu_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SparseCsrReluGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-PD_REGISTER_KERNEL(sparse_relu_grad,
+PD_REGISTER_KERNEL(sparse_coo_relu_grad,
                    GPU,
                    ALL_LAYOUT,
-                   phi::sparse::SparseReluGradKernel,
+                   phi::sparse::SparseCooReluGradKernel,
                    float,
                    double,
                    phi::dtype::float16) {
   kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
 }
+
+PD_REGISTER_KERNEL(sparse_csr_relu_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SparseCsrReluGradKernel,
+                   float,
+                   double,
+                   phi::dtype::float16) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
 #endif
+

paddle/phi/kernels/sparse/activation_grad_kernel.h

@@ -15,15 +15,28 @@ limitations under the License. */
 #pragma once
 
 #include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
 
 namespace phi {
 namespace sparse {
 
-template <typename T, typename Context>
-void SparseReluGradKernel(const Context& dev_ctx,
-                          const SparseCooTensor& x,
-                          const SparseCooTensor& out_grad,
-                          SparseCooTensor* x_grad);
+#define DECLARE_SPARSE_ACTIVATION_GRAD_KERNEL(name)                 \
+  template <typename T, typename Context>                           \
+  void SparseCoo##name##GradKernel(const Context& dev_ctx,          \
+                                   const SparseCooTensor& x,        \
+                                   const SparseCooTensor& out_grad, \
+                                   SparseCooTensor* x_grad);        \
+                                                                    \
+  template <typename T, typename Context>                           \
+  void SparseCsr##name##GradKernel(const Context& dev_ctx,          \
+                                   const SparseCsrTensor& x,        \
+                                   const SparseCsrTensor& out_grad, \
+                                   SparseCsrTensor* x_grad);
+
+DECLARE_SPARSE_ACTIVATION_GRAD_KERNEL(Relu)
+DECLARE_SPARSE_ACTIVATION_GRAD_KERNEL(Sqrt)
+
+#undef DECLARE_SPARSE_ACTIVATION_GRAD_KERNEL
 
 }  // namespace sparse
 }  // namespace phi

paddle/phi/kernels/sparse/activation_kernel.cc

@@ -13,54 +13,53 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/phi/kernels/sparse/activation_kernel.h"
-#include "paddle/phi/kernels/copy_kernel.h"
-#include "paddle/phi/kernels/empty_kernel.h"
 
-#include "paddle/phi/backends/cpu/cpu_context.h"
-#include "paddle/phi/backends/gpu/gpu_context.h"
-#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/sparse/utils.h"
 
-namespace phi {
-namespace sparse {
+DEFINE_AND_REGISTER_SPARSE_UNARY_KERNEL(sqrt, SqrtKernel)
 
-template <typename T, typename Context>
-void SparseReluKernel(const Context& dev_ctx,
-                      const SparseCooTensor& x,
-                      SparseCooTensor* out) {
-  DenseTensor non_zero_indices =
-      phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_indices());
-  DenseTensor non_zero_elements =
-      phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_elements());
-  phi::Copy(dev_ctx,
-            x.non_zero_indices(),
-            dev_ctx.GetPlace(),
-            false,
-            &non_zero_indices);
-  phi::ReluKernel<T, Context>(
-      dev_ctx, x.non_zero_elements(), &non_zero_elements);
-  out->SetMember(non_zero_indices, non_zero_elements, x.dims(), true);
-}
-
-}  // namespace sparse
-}  // namespace phi
+// NOTE: the following code is to bypass the restriction of Paddle
+// kernel registration mechanism. Do NOT refactor them unless you
+// know what you are doing.
+// If you want to implement any new kernel, please follow `sqrt` above
+// instead of `relu` following
+DEFINE_SPARSE_UNARY_KERNEL(ReluKernel)
 
-PD_REGISTER_KERNEL(sparse_relu,
+PD_REGISTER_KERNEL(sparse_coo_relu,
                    CPU,
                    ALL_LAYOUT,
-                   phi::sparse::SparseReluKernel,
+                   phi::sparse::SparseCooReluKernel,
                    float,
                    double) {
   kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
 }
+PD_REGISTER_KERNEL(sparse_csr_relu,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SparseCsrReluKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
 
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-PD_REGISTER_KERNEL(sparse_relu,
+PD_REGISTER_KERNEL(sparse_coo_relu,
                    GPU,
                    ALL_LAYOUT,
-                   phi::sparse::SparseReluKernel,
+                   phi::sparse::SparseCooReluKernel,
                    float,
                    double,
                    phi::dtype::float16) {
   kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
 }
+
+PD_REGISTER_KERNEL(sparse_csr_relu,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SparseCsrReluKernel,
+                   float,
+                   double,
+                   phi::dtype::float16) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
 #endif

paddle/phi/kernels/sparse/activation_kernel.h

@@ -16,22 +16,40 @@ limitations under the License. */
 
 #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/activation_kernel.h"
 #include "paddle/phi/kernels/empty_kernel.h"
 
 namespace phi {
 namespace sparse {
 
-template <typename T, typename Context>
-void SparseReluKernel(const Context& dev_ctx,
-                      const SparseCooTensor& x,
-                      SparseCooTensor* out);
+#define DECLARE_SPARSE_ACTIVATION_KERNEL(name)                                 \
+  template <typename T, typename Context>                                      \
+  void SparseCoo##name##Kernel(                                                \
+      const Context& dev_ctx, const SparseCooTensor& x, SparseCooTensor* out); \
+                                                                               \
+  template <typename T, typename Context>                                      \
+  void SparseCsr##name##Kernel(                                                \
+      const Context& dev_ctx, const SparseCsrTensor& x, SparseCsrTensor* out);
+
+DECLARE_SPARSE_ACTIVATION_KERNEL(Relu)
+DECLARE_SPARSE_ACTIVATION_KERNEL(Sqrt)
+
+#undef DECLARE_SPARSE_ACTIVATION_KERNEL
 
 template <typename T, typename Context>
 SparseCooTensor SparseRelu(const Context& dev_ctx, const SparseCooTensor& x) {
   DenseTensor indices, values;
   SparseCooTensor coo(indices, values, x.dims());
-  SparseReluKernel<T, Context>(dev_ctx, x, &coo);
+  SparseCooReluKernel<T, Context>(dev_ctx, x, &coo);
+  return coo;
+}
+
+template <typename T, typename Context>
+SparseCooTensor SparseSqrt(const Context& dev_ctx, const SparseCooTensor& x) {
+  DenseTensor indices, values;
+  SparseCooTensor coo(indices, values, x.dims());
+  SparseCooSqrtKernel<T, Context>(dev_ctx, x, &coo);
   return coo;
 }