chore: import upstream snapshot with attribution

paddle/phi/kernels/sparse/cpu/addmm_grad_kernel.cc

@@ -0,0 +1,70 @@
+/* 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/addmm_grad_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void AddmmCooDenseGradKernel(const Context& dev_ctx UNUSED,
+                             const DenseTensor& input UNUSED,
+                             const SparseCooTensor& x UNUSED,
+                             const DenseTensor& y UNUSED,
+                             const DenseTensor& dout UNUSED,
+                             float alpha UNUSED,
+                             float beta UNUSED,
+                             DenseTensor* dinput UNUSED,
+                             SparseCooTensor* dx UNUSED,
+                             DenseTensor* dy UNUSED) {
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support CPU backward kernel of 'sparse.addmm' now."));
+}
+
+template <typename T, typename Context>
+void AddmmCsrDenseGradKernel(const Context& dev_ctx UNUSED,
+                             const DenseTensor& input UNUSED,
+                             const SparseCsrTensor& x UNUSED,
+                             const DenseTensor& y UNUSED,
+                             const DenseTensor& dout UNUSED,
+                             float alpha UNUSED,
+                             float beta UNUSED,
+                             DenseTensor* dinput UNUSED,
+                             SparseCsrTensor* dx UNUSED,
+                             DenseTensor* dy UNUSED) {
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support CPU backward kernel of 'sparse.addmm' now."));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(addmm_coo_dense_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::AddmmCooDenseGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(addmm_csr_dense_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::AddmmCsrDenseGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}

paddle/phi/kernels/sparse/cpu/addmm_kernel.cc

@@ -0,0 +1,65 @@
+/* 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/addmm_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+namespace phi::sparse {
+
+/* DENSE + COO @ DENSE -> DENSE */
+template <typename T, typename Context>
+void AddmmCooDenseKernel(const Context& dev_ctx UNUSED,
+                         const DenseTensor& input UNUSED,
+                         const SparseCooTensor& x UNUSED,
+                         const DenseTensor& y UNUSED,
+                         float beta UNUSED,
+                         float alpha UNUSED,
+                         DenseTensor* out UNUSED) {
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support CPU kernel of 'sparse.addmm' now."));
+}
+
+/* DENSE + CSR @ DENSE -> DENSE */
+template <typename T, typename Context>
+void AddmmCsrDenseKernel(const Context& dev_ctx UNUSED,
+                         const DenseTensor& input UNUSED,
+                         const SparseCsrTensor& x UNUSED,
+                         const DenseTensor& y UNUSED,
+                         float beta UNUSED,
+                         float alpha UNUSED,
+                         DenseTensor* out UNUSED) {
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support CPU kernel of 'sparse.addmm' now."));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(addmm_coo_dense,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::AddmmCooDenseKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(addmm_csr_dense,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::AddmmCsrDenseKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}

paddle/phi/kernels/sparse/cpu/coalesce_kernel.cc

@@ -0,0 +1,122 @@
+/* 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/coalesce_kernel.h"
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/funcs/sparse/flatten_indices.h"
+
+namespace phi::sparse {
+
+template <typename T, typename IntT>
+void CoalesceCooCPUKernel(const CPUContext& dev_ctx,
+                          const SparseCooTensor& x,
+                          SparseCooTensor* out) {
+  const DenseTensor& x_indices = x.indices();
+  const DenseTensor& x_values = x.values();
+  DenseTensor out_indices = EmptyLike<IntT>(dev_ctx, x_indices);
+  DenseTensor out_values = EmptyLike<T>(dev_ctx, x_values);
+
+  const int64_t sparse_dim = x.indices().dims()[0];
+  std::vector<IntT> sparse_offsets(sparse_dim), x_nnz(x.nnz());
+  funcs::sparse::CalcOffsetsPerDim<IntT>(
+      x.dims(), sparse_dim, sparse_offsets.data());
+
+  funcs::sparse::FlattenIndices(x.indices().data<IntT>(),
+                                sparse_offsets.data(),
+                                x.nnz(),
+                                sparse_dim,
+                                0,
+                                1,
+                                x_nnz.data());
+
+  const T* x_values_ptr = x_values.data<T>();
+  const int64_t stride =
+      x.dims().size() == sparse_dim ? 1 : x.values().dims()[1];
+
+  std::map<IntT, std::vector<int64_t>> indices_to_nnz;
+  for (uint64_t i = 0; i < x_nnz.size(); i++) {
+    IntT index = x_nnz[i];
+    if (indices_to_nnz.find(index) == indices_to_nnz.end()) {
+      std::vector<int64_t> lost_indices;
+      lost_indices.push_back(static_cast<int>(i));
+      indices_to_nnz[index] = lost_indices;
+    } else {
+      indices_to_nnz[index].push_back(i);
+    }
+  }
+
+  const int64_t out_nnz = indices_to_nnz.size();
+
+  out_indices.Resize({x_indices.dims()[0], out_nnz});
+  if (out_values.dims().size() == 1) {
+    out_values.Resize({out_nnz});
+  } else {
+    out_values.Resize({out_nnz, x_values.dims()[1]});
+  }
+
+  IntT* out_indices_ptr = out_indices.data<IntT>();
+  T* out_values_ptr = out_values.data<T>();
+  auto iter = indices_to_nnz.begin();
+
+  Dim<DDim::kMaxRank> const_dims;
+  for (int i = 0; i < x.dims().size(); i++) {
+    const_dims[i] = x.dims()[i];
+  }
+
+  for (int i = 0; iter != indices_to_nnz.end(); iter++, i++) {
+    funcs::sparse::IndexToCoordinate(
+        iter->first, const_dims, out_nnz, sparse_dim, i, out_indices_ptr);
+    memcpy(out_values_ptr + i * stride,
+           x_values_ptr + iter->second[0] * stride,
+           stride * sizeof(T));
+    for (uint64_t j = 1; j < iter->second.size(); j++) {
+      for (int k = 0; k < stride; k++) {
+        out_values_ptr[i * stride + k] +=
+            x_values_ptr[iter->second[j] * stride + k];
+      }
+    }
+  }
+
+  out->SetMember(out_indices, out_values, x.dims(), true);
+}
+
+template <typename T, typename Context>
+void CoalesceCooKernel(const Context& dev_ctx,
+                       const SparseCooTensor& x,
+                       SparseCooTensor* out) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.indices().dtype(), "CoalesceCooCPUKernel", ([&] {
+        CoalesceCooCPUKernel<T, data_t>(dev_ctx, x, out);
+      }));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(coalesce_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CoalesceCooKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/cpu/conv.h

@@ -0,0 +1,267 @@
+/* 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. */
+
+#pragma once
+
+#include <set>
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/tensor_meta.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/sparse/conv_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+using Dims4D = funcs::sparse::Dims4D;
+
+// such as: kernel(3, 3, 3), kernel_size = 27
+// counter_per_weight: (kernel_size)
+// TODO(zhangkaihuo): optimize performance with multithreading
+template <typename T, typename Context, typename IntT = int>
+void ProductRuleBook(const Context& dev_ctx,
+                     const SparseCooTensor& x,
+                     const std::vector<int>& kernel_sizes,
+                     const std::vector<int>& paddings,
+                     const std::vector<int>& dilations,
+                     const std::vector<int>& strides,
+                     const DDim& out_dims,
+                     const bool subm,
+                     DenseTensor* rulebook,
+                     int* counter_per_kernel) {
+  const bool is2D = out_dims.size() == 4 ? true : false;
+  const int64_t non_zero_num = x.nnz();
+  const auto& indices = x.indices();
+  const IntT* indices_ptr = indices.data<IntT>();
+  int kernel_size = is2D ? kernel_sizes[0] * kernel_sizes[1]
+                         : kernel_sizes[0] * kernel_sizes[1] * kernel_sizes[2];
+  memset(counter_per_kernel, 0, kernel_size * sizeof(int));
+
+  int rulebook_len = 0;
+  // calc the rulebook_len
+  const auto& x_dims = x.dims();
+
+  int xdim0, xdim1, xdim2, xdim3;
+  int kdim0, kdim1, kdim2, kdim3;
+  int odim0, odim1, odim2, odim3;
+  int pdim0, pdim1, pdim2, pdim3;
+  int sdim0, sdim1, sdim2, sdim3;
+  int ddim0, ddim1, ddim2, ddim3;
+
+  xdim0 = x_dims[0];
+  xdim1 = is2D ? x_dims[2] : x_dims[3];
+  xdim2 = is2D ? x_dims[1] : x_dims[2];
+  xdim3 = is2D ? 1 : x_dims[1];
+
+  kdim0 = 1;
+  kdim1 = is2D ? kernel_sizes[1] : kernel_sizes[2];
+  kdim2 = is2D ? kernel_sizes[0] : kernel_sizes[1];
+  kdim3 = is2D ? 1 : kernel_sizes[0];
+
+  odim0 = out_dims[0];
+  odim1 = is2D ? out_dims[2] : out_dims[3];
+  odim2 = is2D ? out_dims[1] : out_dims[2];
+  odim3 = is2D ? 1 : out_dims[1];
+
+  pdim0 = 1;
+  pdim1 = is2D ? paddings[1] : paddings[2];
+  pdim2 = is2D ? paddings[0] : paddings[1];
+  pdim3 = is2D ? 1 : paddings[0];
+
+  sdim0 = 1;
+  sdim1 = is2D ? strides[1] : strides[2];
+  sdim2 = is2D ? strides[0] : strides[1];
+  sdim3 = is2D ? 1 : strides[0];
+
+  ddim0 = 1;
+  ddim1 = is2D ? dilations[1] : dilations[2];
+  ddim2 = is2D ? dilations[0] : dilations[1];
+  ddim3 = is2D ? 1 : dilations[0];
+
+  const Dims4D c_x_dims(xdim0, xdim1, xdim2, xdim3);
+  const Dims4D c_kernel_dims(kdim0, kdim1, kdim2, kdim3);
+  const Dims4D c_out_dims(odim0, odim1, odim2, odim3);
+  const Dims4D c_paddings(pdim0, pdim1, pdim2, pdim3);
+  const Dims4D c_strides(sdim0, sdim1, sdim2, sdim3);
+  const Dims4D c_dilations(ddim0, ddim1, ddim2, ddim3);
+
+  std::set<IntT> hash_in;
+  if (subm) {
+    for (int i = 0; i < non_zero_num; i++) {
+      IntT batch = indices_ptr[i];
+      IntT in_z = is2D ? 0 : indices_ptr[i + non_zero_num];
+      IntT in_y = is2D ? indices_ptr[i + non_zero_num]
+                       : indices_ptr[i + 2 * non_zero_num];
+      IntT in_x = is2D ? indices_ptr[i + 2 * non_zero_num]
+                       : indices_ptr[i + 3 * non_zero_num];
+      IntT index = funcs::sparse::PointToIndex<Dims4D>(
+          batch, in_x, in_y, in_z, c_x_dims);
+      hash_in.insert(index);
+    }
+  }
+
+  auto f_calc_rulebook = [&](IntT* rulebook_ptr) {
+    int kernel_index = 0, rulebook_index = 0;
+    int zceil = is2D ? 1 : kernel_sizes[0];
+    int yceil = is2D ? kernel_sizes[0] : kernel_sizes[1];
+    int xceil = is2D ? kernel_sizes[1] : kernel_sizes[2];
+    for (int kz = 0; kz < zceil; kz++) {
+      for (int ky = 0; ky < yceil; ky++) {
+        for (int kx = 0; kx < xceil; kx++) {
+          ++kernel_index;
+          for (int64_t i = 0; i < non_zero_num; i++) {
+            IntT batch = indices_ptr[i];
+            IntT in_z = is2D ? 0 : indices_ptr[i + non_zero_num];
+            IntT in_y = is2D ? indices_ptr[i + non_zero_num]
+                             : indices_ptr[i + 2 * non_zero_num];
+            IntT in_x = is2D ? indices_ptr[i + 2 * non_zero_num]
+                             : indices_ptr[i + 3 * non_zero_num];
+
+            IntT out_z =
+                is2D ? 0
+                     : (in_z + paddings[0] - kz * dilations[0]) / strides[0];
+            IntT out_y =
+                (in_y + c_paddings[2] - ky * c_dilations[2]) / c_strides[2];
+            IntT out_x =
+                (in_x + c_paddings[3] - kx * c_dilations[3]) / c_strides[3];
+            if (funcs::sparse::Check(c_x_dims,
+                                     c_kernel_dims,
+                                     c_paddings,
+                                     c_dilations,
+                                     c_strides,
+                                     in_x,
+                                     in_y,
+                                     in_z,
+                                     kx,
+                                     ky,
+                                     kz)) {
+              if (subm) {
+                IntT out_index = funcs::sparse::PointToIndex<Dims4D>(
+                    batch, out_x, out_y, out_z, c_out_dims);
+                if (hash_in.find(out_index) == hash_in.end()) {
+                  continue;
+                }
+              }
+
+              if (rulebook_ptr == nullptr) {
+                counter_per_kernel[kernel_index - 1] += 1;
+                ++rulebook_len;
+              } else {
+                rulebook_ptr[rulebook_index] = kernel_index - 1;
+                rulebook_ptr[rulebook_index + rulebook_len] = i;  // in_i
+                rulebook_ptr[rulebook_index + rulebook_len * 2] =
+                    funcs::sparse::PointToIndex<Dims4D>(
+                        batch, out_x, out_y, out_z, c_out_dims);  // out_index
+                ++rulebook_index;
+              }
+            }
+          }
+        }
+      }
+    }
+  };
+
+  f_calc_rulebook(nullptr);
+  // alloc the rulebook
+  *rulebook = Empty(dev_ctx,
+                    DenseTensorMeta(phi::CppTypeToDataType<IntT>::Type(),
+                                    {3, rulebook_len},
+                                    DataLayout::NCHW));
+  IntT* rulebook_ptr = rulebook->data<IntT>();
+  f_calc_rulebook(rulebook_ptr);
+}
+
+template <typename T, typename Context, typename IntT = int>
+void UpdateRulebookAndOutIndex(const Context& dev_ctx,
+                               const SparseCooTensor& x,
+                               const int kernel_size UNUSED,
+                               const int out_channels,
+                               const DDim& out_dims,
+                               DenseTensor* rulebook,
+                               SparseCooTensor* out) {
+  const bool is2D = out_dims.size() == 4 ? true : false;
+
+  std::set<IntT> tmp_indices;
+  int64_t n = rulebook->dims()[1];
+  IntT* rulebook_ptr = rulebook->data<IntT>();
+  for (int64_t i = 0; i < n; i++) {
+    tmp_indices.insert(rulebook_ptr[i + n * 2]);
+  }
+
+  int out_non_zero_num = tmp_indices.size();
+  const int64_t sparse_dim = is2D ? 3 : 4;
+  DenseTensorMeta indices_meta(phi::CppTypeToDataType<IntT>::Type(),
+                               {sparse_dim, out_non_zero_num},
+                               DataLayout::NCHW);
+  DenseTensorMeta values_meta(
+      x.dtype(), {out_non_zero_num, out_channels}, x.values().layout());
+  DenseTensor out_indices = Empty(dev_ctx, std::move(indices_meta));
+  DenseTensor out_values = Empty(dev_ctx, std::move(values_meta));
+  IntT* out_indices_ptr = out_indices.data<IntT>();
+  int64_t idx = 0;
+
+  int odim0, odim1, odim2, odim3;
+  odim0 = out_dims[0];
+  odim1 = is2D ? out_dims[2] : out_dims[3];
+  odim2 = is2D ? out_dims[1] : out_dims[2];
+  odim3 = is2D ? 1 : out_dims[1];
+  const Dims4D c_out_dims(odim0, odim1, odim2, odim3);
+
+  for (auto it = tmp_indices.begin(); it != tmp_indices.end(); it++, idx++) {
+    const IntT index = *it;
+    IntT batch, x, y, z;
+    funcs::sparse::IndexToPoint<Dims4D>(index, c_out_dims, &batch, &x, &y, &z);
+    out_indices_ptr[idx] = batch;
+    if (is2D) {
+      out_indices_ptr[idx + out_non_zero_num] = y;
+      out_indices_ptr[idx + out_non_zero_num * 2] = x;
+    } else {
+      out_indices_ptr[idx + out_non_zero_num] = z;
+      out_indices_ptr[idx + out_non_zero_num * 2] = y;
+      out_indices_ptr[idx + out_non_zero_num * 3] = x;
+    }
+  }
+  for (int64_t i = 0; i < n; i++) {
+    IntT out_index = rulebook_ptr[i + n * 2];
+    rulebook_ptr[i + n * 2] =
+        std::distance(tmp_indices.begin(), tmp_indices.find(out_index));
+  }
+
+  out->SetMember(out_indices, out_values, out_dims, true);
+}
+
+template <typename T, typename IntT = int>
+void Gather(
+    const T* x, const IntT* indices, const int n, const int channels, T* out) {
+  for (int i = 0; i < n; i++) {
+    IntT real_i = indices[i];
+    memcpy(out + i * channels, x + real_i * channels, channels * sizeof(T));
+  }
+}
+
+template <typename T, typename IntT = int>
+void Scatter(
+    const T* x, const IntT* indices, const int n, const int channels, T* out) {
+  for (int i = 0; i < n; i++) {
+    IntT real_i = indices[i];
+    for (int j = 0; j < channels; j++) {
+      out[real_i * channels + j] += x[i * channels + j];
+    }
+  }
+}
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/cpu/conv_grad_kernel.cc

@@ -0,0 +1,224 @@
+/* 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/conv_grad_kernel.h"
+
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/sparse/cpu/conv.h"
+
+namespace phi::sparse {
+
+// rulebook:
+//[
+//  [kernel_index],
+//  [in_i],
+//  [out_i],
+//]
+// x_grad = out_grad * transpose(kernel)
+// kernel_grad = transpose(x) * out_grad
+template <typename T, typename IntT = int>
+void Conv3dCooGradCPUKernel(const CPUContext& dev_ctx,
+                            const SparseCooTensor& x,
+                            const DenseTensor& kernel,
+                            const SparseCooTensor& out,
+                            const DenseTensor& rulebook,
+                            const DenseTensor& counter,
+                            const SparseCooTensor& out_grad,
+                            const std::vector<int>& paddings UNUSED,
+                            const std::vector<int>& dilations UNUSED,
+                            const std::vector<int>& strides UNUSED,
+                            const int groups UNUSED,
+                            const bool subm,
+                            const std::string& key,
+                            SparseCooTensor* x_grad,
+                            DenseTensor* kernel_grad) {
+  const auto& kernel_dims = kernel.dims();
+  const bool is2D = kernel_dims.size() == 4 ? true : false;
+  const int kernel_size =
+      static_cast<int>(is2D ? kernel_dims[0] * kernel_dims[1]
+                            : kernel_dims[0] * kernel_dims[1] * kernel_dims[2]);
+  const int in_channels =
+      static_cast<int>(is2D ? kernel_dims[2] : kernel_dims[3]);
+  const int out_channels =
+      static_cast<int>(is2D ? kernel_dims[3] : kernel_dims[4]);
+
+  int rulebook_len = 0;
+  const IntT* rulebook_ptr =
+      funcs::sparse::GetRulebookPtr<IntT>(out, rulebook, key, &rulebook_len);
+  const int* counter_ptr = funcs::sparse::GetCounterPtr(out, counter, key);
+
+  DenseTensorMeta in_features_meta(
+      x.dtype(), {rulebook_len, in_channels}, DataLayout::NCHW);
+  DenseTensorMeta d_x_features_meta(
+      x.dtype(), {rulebook_len, in_channels}, DataLayout::NCHW);
+  DenseTensorMeta out_grad_features_meta(
+      x.dtype(), {rulebook_len, out_channels}, DataLayout::NCHW);
+  DenseTensor in_features = Empty(dev_ctx, std::move(in_features_meta));
+  DenseTensor d_x_features = Empty(dev_ctx, std::move(d_x_features_meta));
+  DenseTensor out_grad_features =
+      Empty(dev_ctx, std::move(out_grad_features_meta));
+
+  T* in_features_ptr = in_features.data<T>();
+  T* d_x_features_ptr = d_x_features.data<T>();
+  T* out_grad_features_ptr = out_grad_features.data<T>();
+  *kernel_grad = EmptyLike<T>(dev_ctx, kernel);
+  T* d_kernel_ptr = kernel_grad->data<T>();
+  memset(d_kernel_ptr, 0, sizeof(T) * kernel_grad->numel());
+
+  int half_kernel_size = kernel_size / 2;
+  auto blas = funcs::GetBlas<CPUContext, T>(dev_ctx);
+  DenseTensor x_grad_indices = EmptyLike<IntT>(dev_ctx, x.indices());
+  DenseTensor x_grad_values = EmptyLike<T>(dev_ctx, x.values());
+  T* x_grad_values_ptr = x_grad_values.data<T>();
+  memset(x_grad_values_ptr, 0, sizeof(T) * x_grad_values.numel());
+  memset(d_x_features_ptr, 0, sizeof(T) * d_x_features.numel());
+  phi::Copy<CPUContext>(
+      dev_ctx, x.indices(), dev_ctx.GetPlace(), false, &x_grad_indices);
+  x_grad->SetMember(x_grad_indices, x_grad_values, x.dims(), true);
+
+  std::vector<IntT> offsets(kernel_size + 1);
+  IntT offset = 0;
+  int max_count = 0;
+  for (int i = 0; i < kernel_size; i++) {
+    offsets[i] = offset;
+    offset += counter_ptr[i];
+    if (i < half_kernel_size) {
+      max_count = std::max(max_count, counter_ptr[i]);
+    }
+  }
+  offsets[kernel_size] = offset;
+
+  if (subm) {
+    funcs::sparse::SubmPreProcess<T, CPUContext>(dev_ctx,
+                                                 x,
+                                                 kernel,
+                                                 out_grad.values(),
+                                                 in_channels,
+                                                 out_channels,
+                                                 half_kernel_size,
+                                                 kernel_grad,
+                                                 &x_grad_values);
+    if (max_count == 0) {
+      return;
+    }
+  }
+
+  Gather<T, IntT>(x.values().data<T>(),
+                  rulebook_ptr + rulebook_len,
+                  rulebook_len,
+                  in_channels,
+                  in_features_ptr);
+  Gather<T, IntT>(out_grad.values().data<T>(),
+                  rulebook_ptr + rulebook_len * 2,
+                  rulebook_len,
+                  out_channels,
+                  out_grad_features_ptr);
+
+  const T* kernel_ptr = kernel.data<T>();
+  for (int i = 0; i < kernel_size; i++) {
+    if (counter_ptr[i] <= 0 || (subm && i == half_kernel_size)) {
+      continue;
+    }
+
+    const int M = counter_ptr[i];
+    const int K = in_channels;
+    const int N = out_channels;
+    T* tmp_in_ptr = in_features_ptr + offsets[i] * in_channels;
+    T* tmp_out_grad_ptr = out_grad_features_ptr + offsets[i] * out_channels;
+    const T* tmp_kernel_ptr = kernel_ptr + i * in_channels * out_channels;
+    T* tmp_d_x_ptr = d_x_features_ptr + offsets[i] * in_channels;
+    T* tmp_d_kernel_ptr = d_kernel_ptr + i * in_channels * out_channels;
+
+    // call gemm: d_kernel = transpose(x) * out_grad
+    // (in_channels, n) * (n, out_channels)
+    blas.GEMM(CblasTrans,
+              CblasNoTrans,
+              K,
+              N,
+              M,
+              static_cast<T>(1),
+              tmp_in_ptr,
+              tmp_out_grad_ptr,
+              static_cast<T>(0),
+              tmp_d_kernel_ptr);
+
+    // call gemm: d_x = out_grad * transpose(kernel)
+    // (n, out_channels) * (out_channels, in_channels)
+    blas.GEMM(CblasNoTrans,
+              CblasTrans,
+              M,
+              K,
+              N,
+              static_cast<T>(1),
+              tmp_out_grad_ptr,
+              tmp_kernel_ptr,
+              static_cast<T>(0),
+              tmp_d_x_ptr);
+  }
+
+  // 4. scatter
+  Scatter<T, IntT>(d_x_features_ptr,
+                   rulebook_ptr + rulebook_len,
+                   rulebook_len,
+                   in_channels,
+                   x_grad_values_ptr);
+}
+
+template <typename T, typename Context>
+void Conv3dCooGradKernel(const Context& dev_ctx,
+                         const SparseCooTensor& x,
+                         const DenseTensor& kernel,
+                         const SparseCooTensor& out,
+                         const DenseTensor& rulebook,
+                         const DenseTensor& counter,
+                         const SparseCooTensor& out_grad,
+                         const std::vector<int>& paddings,
+                         const std::vector<int>& dilations,
+                         const std::vector<int>& strides,
+                         const int groups,
+                         const bool subm,
+                         const std::string& key,
+                         SparseCooTensor* x_grad,
+                         DenseTensor* kernel_grad) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.indices().dtype(), "Conv3dCooGradCPUKernel", ([&] {
+        Conv3dCooGradCPUKernel<T, data_t>(dev_ctx,
+                                          x,
+                                          kernel,
+                                          out,
+                                          rulebook,
+                                          counter,
+                                          out_grad,
+                                          paddings,
+                                          dilations,
+                                          strides,
+                                          groups,
+                                          subm,
+                                          key,
+                                          x_grad,
+                                          kernel_grad);
+      }));
+}
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(conv3d_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::Conv3dCooGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/cpu/conv_kernel.cc

@@ -0,0 +1,214 @@
+/* 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/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_meta.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/sparse/cpu/conv.h"
+
+namespace phi::sparse {
+
+/**
+ * x: (N, D, H, W, C)
+ * kernel: (D, H, W, C, OC)
+ * out: (N, D, H, W, OC)
+ **/
+template <typename T, typename IntT = int>
+void Conv3dCooCPUKernel(const CPUContext& dev_ctx,
+                        const SparseCooTensor& x,
+                        const DenseTensor& kernel,
+                        const std::vector<int>& paddings,
+                        const std::vector<int>& dilations,
+                        const std::vector<int>& strides,
+                        const int groups UNUSED,
+                        const bool subm,
+                        const std::string& key,
+                        SparseCooTensor* out,
+                        DenseTensor* rulebook,
+                        DenseTensor* counter) {
+  // update padding and dilation
+  // Currently, only support x.layout is NDHWC, groups = 1
+  // if x.layout != NDHWC then transpose(x), transpose(weight)
+
+  const auto& x_dims = x.dims();
+  const bool is2D = x_dims.size() == 4 ? true : false;
+  const auto& kernel_dims = kernel.dims();
+  int kernel_size =
+      static_cast<int>(is2D ? kernel_dims[0] * kernel_dims[1]
+                            : kernel_dims[0] * kernel_dims[1] * kernel_dims[2]);
+
+  int count_tmp = is2D ? 4 : 5;
+  std::vector<int> out_dims_vec(count_tmp, 1);
+  DDim out_dims = make_ddim(out_dims_vec);
+
+  std::vector<int> kernel_sizes(kernel_dims.size());
+  for (int i = 0; i < kernel_dims.size(); i++) {
+    kernel_sizes[i] = static_cast<int>(kernel_dims[i]);
+  }
+
+  std::vector<int> subm_paddings(paddings), subm_strides(strides);
+  if (subm) {
+    // the out shape of subm_conv is same as input shape
+    // reset the padding=kernel_size/2 and strides=1
+    funcs::sparse::ResetSubmKernelSizeAndStrides(
+        kernel.dims(), &subm_paddings, &subm_strides);
+  }
+
+  funcs::sparse::GetOutShape(
+      x_dims, kernel_sizes, subm_paddings, dilations, subm_strides, &out_dims);
+  const int in_channels =
+      static_cast<int>(is2D ? kernel_dims[2] : kernel_dims[3]);
+  const int out_channels =
+      static_cast<int>(is2D ? kernel_dims[3] : kernel_dims[4]);
+
+  // Second algorithm:
+  // https://pdfs.semanticscholar.org/5125/a16039cabc6320c908a4764f32596e018ad3.pdf
+  // 1. product rulebook
+  DenseTensor h_counter, h_offsets;
+  h_counter.Resize({kernel_size});
+  h_offsets.Resize({kernel_size + 1});
+  int* h_counter_ptr = dev_ctx.template HostAlloc<int>(&h_counter);
+  int* h_offsets_ptr = dev_ctx.template HostAlloc<int>(&h_offsets);
+
+  // DenseTensor* rulebook = nullptr;
+  const IntT* rulebook_ptr = nullptr;
+  int n = 0;
+  bool need_product_rulebook = true;
+  if (subm && !key.empty()) {
+    rulebook_ptr =
+        funcs::sparse::PrepareSubm<T, IntT, CPUContext>(dev_ctx,
+                                                        x,
+                                                        key,
+                                                        out_dims,
+                                                        out,
+                                                        h_counter_ptr,
+                                                        h_offsets_ptr,
+                                                        &n,
+                                                        &need_product_rulebook);
+  }
+  if (need_product_rulebook) {
+    DenseTensor tmp_rulebook;
+    ProductRuleBook<T, CPUContext, IntT>(dev_ctx,
+                                         x,
+                                         kernel_sizes,
+                                         subm_paddings,
+                                         dilations,
+                                         subm_strides,
+                                         out_dims,
+                                         subm,
+                                         &tmp_rulebook,
+                                         h_counter_ptr);
+
+    UpdateRulebookAndOutIndex<T, CPUContext, IntT>(
+        dev_ctx, x, kernel_size, out_channels, out_dims, &tmp_rulebook, out);
+    n = static_cast<int>(tmp_rulebook.dims()[1]);
+    rulebook_ptr = tmp_rulebook.data<IntT>();
+
+    funcs::sparse::SaveToTable(
+        dev_ctx, x, key, tmp_rulebook, h_counter, out, rulebook, counter);
+  }
+
+  // 2. gather
+  DenseTensorMeta in_features_meta(
+      x.dtype(), {n, in_channels}, DataLayout::NHWC);
+  DenseTensorMeta out_features_meta(
+      x.dtype(), {n, out_channels}, DataLayout::NHWC);
+  DenseTensor in_features = Empty(dev_ctx, std::move(in_features_meta));
+  DenseTensor out_features = Empty(dev_ctx, std::move(out_features_meta));
+  T* in_features_ptr = in_features.data<T>();
+  T* out_features_ptr = out_features.data<T>();
+
+  Gather<T, IntT>(
+      x.values().data<T>(), rulebook_ptr + n, n, in_channels, in_features_ptr);
+
+  // 3. call gemm for every weight
+  auto blas = funcs::GetBlas<CPUContext, T>(dev_ctx);
+  int offset = 0;
+  for (int i = 0; i < kernel_size; i++) {
+    h_offsets_ptr[i] = offset;
+    offset += h_counter_ptr[i];
+  }
+  h_offsets_ptr[kernel_size] = offset;
+
+  const T* kernel_ptr = kernel.data<T>();
+  for (int i = 0; i < kernel_size; i++) {
+    if (h_counter_ptr[i] <= 0) {
+      continue;
+    }
+
+    // call gemm: (n, in_channels) * (in_channels, out_channels)
+    const int M = h_counter_ptr[i];
+    const int K = in_channels;   // in_channels
+    const int N = out_channels;  // out_channels
+    T* tmp_in_ptr = in_features_ptr + h_offsets_ptr[i] * in_channels;
+    const T* tmp_kernel_ptr = kernel_ptr + i * K * N;
+    T* tmp_out_ptr = out_features_ptr + h_offsets_ptr[i] * out_channels;
+    blas.GEMM(CblasNoTrans,
+              CblasNoTrans,
+              M,
+              N,
+              K,
+              static_cast<T>(1),
+              tmp_in_ptr,
+              tmp_kernel_ptr,
+              static_cast<T>(0),
+              tmp_out_ptr);
+  }
+
+  // 4. scatter
+  T* out_values_ptr = out->mutable_values()->data<T>();
+  memset(out_values_ptr, 0, sizeof(T) * out->nnz() * out_channels);
+  Scatter<T, IntT>(
+      out_features_ptr, rulebook_ptr + n * 2, n, out_channels, out_values_ptr);
+}
+
+template <typename T, typename Context>
+void Conv3dCooKernel(const Context& dev_ctx,
+                     const SparseCooTensor& x,
+                     const DenseTensor& kernel,
+                     const std::vector<int>& paddings,
+                     const std::vector<int>& dilations,
+                     const std::vector<int>& strides,
+                     const int groups,
+                     const bool subm,
+                     const std::string& key,
+                     SparseCooTensor* out,
+                     DenseTensor* rulebook,
+                     DenseTensor* counter) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "Conv3dCooCPUKernel", ([&] {
+                                 Conv3dCooCPUKernel<T, data_t>(dev_ctx,
+                                                               x,
+                                                               kernel,
+                                                               paddings,
+                                                               dilations,
+                                                               strides,
+                                                               groups,
+                                                               subm,
+                                                               key,
+                                                               out,
+                                                               rulebook,
+                                                               counter);
+                               }));
+}
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(
+    conv3d_coo, CPU, ALL_LAYOUT, phi::sparse::Conv3dCooKernel, float, double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+  kernel->OutputAt(1).SetDataType(phi::DataType::INT32);
+  kernel->OutputAt(2).SetDataType(phi::DataType::INT32);
+}

paddle/phi/kernels/sparse/cpu/elementwise_grad_kernel.cc

@@ -0,0 +1,690 @@
+/* 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/kernels/sparse/elementwise_kernel.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_meta.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/activation_kernel.h"
+#include "paddle/phi/kernels/complex_kernel.h"
+#include "paddle/phi/kernels/elementwise_kernel.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/sparse/flatten_indices.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi::sparse {
+
+template <typename T, typename IntT, typename Context>
+void AllocCsrPtr(const Context& dev_ctx,
+                 const SparseCsrTensor& x,
+                 SparseCsrTensor* dx) {
+  DenseTensor dx_crows = EmptyLike<IntT>(dev_ctx, x.crows());
+  DenseTensor dx_cols = EmptyLike<IntT>(dev_ctx, x.cols());
+  DenseTensor dx_values = EmptyLike<T>(dev_ctx, x.values());
+  dx->set_meta(x.meta());  // NOLINT
+  dx->SetMember(dx_crows, dx_cols, dx_values, x.dims());
+}
+
+template <typename T, typename IntT, typename Context>
+void AllocCooPtr(const Context& dev_ctx,
+                 const SparseCooTensor& x,
+                 SparseCooTensor* dx) {
+  DenseTensor dx_indices = EmptyLike<IntT>(dev_ctx, x.indices());
+  DenseTensor dx_values = EmptyLike<T>(dev_ctx, x.values());
+  dx->set_meta(x.meta());  // NOLINT
+  dx->SetMember(dx_indices, dx_values, x.dims(), x.coalesced());
+}
+
+template <typename T, typename IntT, typename Context>
+void CopyCooValues(const Context& dev_ctx,
+                   const SparseCooTensor& dout,
+                   const SparseCooTensor& x,
+                   SparseCooTensor* dx) {
+  Copy(dev_ctx, x.indices(), dev_ctx.GetPlace(), false, dx->mutable_indices());
+
+  const int sparse_dim = x.sparse_dim();
+  std::vector<IntT> sparse_offsets(sparse_dim), dout_indices(dout.nnz()),
+      x_indices(x.nnz());
+
+  funcs::sparse::CalcOffsetsPerDim<IntT>(
+      dout.dims(), sparse_dim, sparse_offsets.data());
+
+  funcs::sparse::FlattenIndices(dout.indices().data<IntT>(),
+                                sparse_offsets.data(),
+                                dout.nnz(),
+                                sparse_dim,
+                                0,
+                                1,
+                                dout_indices.data());
+
+  funcs::sparse::FlattenIndices(x.indices().data<IntT>(),
+                                sparse_offsets.data(),
+                                x.nnz(),
+                                sparse_dim,
+                                0,
+                                1,
+                                x_indices.data());
+
+  size_t i = 0, j = 0;
+  T* dx_values_ptr = dx->mutable_values()->data<T>();
+  const T* dout_values_ptr = dout.values().data<T>();
+
+  int64_t element_size = 1;
+  for (auto j = 1; j < x.values().dims().size(); ++j) {
+    element_size *= x.values().dims()[j];
+  }
+
+  while (i < dout_indices.size() && j < x_indices.size()) {
+    if (dout_indices[i] == x_indices[j]) {
+      memcpy(dx_values_ptr + j * element_size,
+             dout_values_ptr + i * element_size,
+             element_size * sizeof(T));
+      ++i;
+      ++j;
+    } else if (dout_indices[i] > x_indices[j]) {
+      memset(dx_values_ptr + j * element_size, 0, element_size * sizeof(T));
+      ++j;
+    } else {
+      ++i;
+    }
+  }
+  while (j < x_indices.size()) {
+    memset(dx_values_ptr + j * element_size, 0, element_size * sizeof(T));
+    ++j;
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void CopyCsrValues(const Context& dev_ctx,
+                   const SparseCsrTensor& dout,
+                   const SparseCsrTensor& x,
+                   SparseCsrTensor* dx) {
+  Copy(dev_ctx, x.crows(), dev_ctx.GetPlace(), false, dx->mutable_crows());
+  Copy(dev_ctx, x.cols(), dev_ctx.GetPlace(), false, dx->mutable_cols());
+
+  const auto& x_dims = x.dims();
+  int batch = static_cast<int>(x_dims.size() == 2 ? 1 : x_dims[0]);
+  int rows = static_cast<int>(x_dims.size() == 2 ? x_dims[0] : x_dims[1]);
+
+  const IntT* x_crows_ptr = x.crows().data<IntT>();
+  const IntT* x_cols_ptr = x.cols().data<IntT>();
+
+  const IntT* dout_crows_ptr = dout.crows().data<IntT>();
+  const IntT* dout_cols_ptr = dout.cols().data<IntT>();
+  const T* dout_values_ptr = dout.values().data<T>();
+
+  T* dx_values_ptr = dx->mutable_values()->data<T>();
+
+  for (int b = 0; b < batch; b++) {
+    for (int r = 0; r < rows; r++) {
+      int x_start = x_crows_ptr[b * (rows + 1) + r];
+      int dout_start = dout_crows_ptr[b * (rows + 1) + r];
+      int x_row_nnz = x_crows_ptr[b * (rows + 1) + r + 1] - x_start;
+      int dout_row_nnz = dout_crows_ptr[b * (rows + 1) + r + 1] - dout_start;
+      int i = 0, j = 0;
+      while (i < x_row_nnz && j < dout_row_nnz) {
+        if (x_cols_ptr[x_start + i] == dout_cols_ptr[dout_start + j]) {
+          dx_values_ptr[x_start + i] = dout_values_ptr[dout_start + j];
+          ++i;
+          ++j;
+        } else if (x_cols_ptr[x_start + i] < dout_cols_ptr[dout_start + j]) {
+          dx_values_ptr[x_start + i] = static_cast<T>(0);
+          ++i;
+        } else {
+          ++j;
+        }
+      }
+      while (i < x_row_nnz) {
+        dx_values_ptr[x_start + i] = static_cast<T>(0);
+        ++i;
+      }
+    }
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void ConjugateCsrValues(const Context& dev_ctx,
+                        const SparseCsrTensor& x,
+                        SparseCsrTensor* x_conj) {
+  AllocCsrPtr<T, IntT>(dev_ctx, x, x_conj);
+  CopyCsrValues<T, IntT, Context>(dev_ctx, x, x, x_conj);
+  DenseTensor x_conj_values = x_conj->values();
+  x_conj_values = phi::Conj<T, Context>(dev_ctx, x_conj_values);
+  DenseTensor x_conj_crows = x_conj->crows();
+  DenseTensor x_conj_cols = x_conj->cols();
+  x_conj->SetMember(x_conj_crows, x_conj_cols, x_conj_values, x_conj->dims());
+}
+
+template <typename T, typename IntT, typename Context>
+void ConjugateCooValues(const Context& dev_ctx,
+                        const SparseCooTensor& x,
+                        SparseCooTensor* x_conj) {
+  AllocCooPtr<T, IntT>(dev_ctx, x, x_conj);
+  CopyCooValues<T, IntT, Context>(dev_ctx, x, x, x_conj);
+  DenseTensor x_conj_values = x_conj->values();
+  x_conj_values = phi::Conj<T, Context>(dev_ctx, x_conj_values);
+  DenseTensor x_conj_indices = x_conj->indices();
+  x_conj->SetMember(
+      x_conj_indices, x_conj_values, x_conj->dims(), x_conj->coalesced());
+}
+
+template <typename T, typename IntT, typename Context>
+void ElementWiseAddCsrGradCPUKernel(const Context& dev_ctx,
+                                    const SparseCsrTensor& x,
+                                    const SparseCsrTensor& y,
+                                    const SparseCsrTensor& dout,
+                                    SparseCsrTensor* dx,
+                                    SparseCsrTensor* dy) {
+  // Special case when y_grad is not needed
+  if (dx != nullptr && dy == nullptr) {
+    VLOG(4) << "Special case when dy is not needed";
+    AllocCsrPtr<T, IntT>(dev_ctx, x, dx);
+    CopyCsrValues<T, IntT, Context>(dev_ctx, dout, x, dx);
+  } else if (dx == nullptr && dy != nullptr) {
+    VLOG(4) << "Special case when dx is not needed";
+    AllocCsrPtr<T, IntT>(dev_ctx, y, dy);
+    CopyCsrValues<T, IntT, Context>(dev_ctx, dout, y, dy);
+  } else {
+    AllocCsrPtr<T, IntT>(dev_ctx, x, dx);
+    AllocCsrPtr<T, IntT>(dev_ctx, y, dy);
+    CopyCsrValues<T, IntT, Context>(dev_ctx, dout, x, dx);
+    CopyCsrValues<T, IntT, Context>(dev_ctx, dout, y, dy);
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void ElementWiseSubtractCsrGradCPUKernel(const Context& dev_ctx,
+                                         const SparseCsrTensor& x,
+                                         const SparseCsrTensor& y,
+                                         const SparseCsrTensor& dout,
+                                         SparseCsrTensor* dx,
+                                         SparseCsrTensor* dy) {
+  if (dx) {
+    AllocCsrPtr<T, IntT>(dev_ctx, x, dx);
+    CopyCsrValues<T, IntT, Context>(dev_ctx, dout, x, dx);
+  }
+
+  if (dy) {
+    AllocCsrPtr<T, IntT>(dev_ctx, y, dy);
+    CopyCsrValues<T, IntT, Context>(dev_ctx, dout, y, dy);
+    phi::NegativeKernel<T, Context>(
+        dev_ctx, dout.values(), dy->mutable_values());
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void ElementWiseMultiplyCsrGradCPUKernel(const Context& dev_ctx,
+                                         const SparseCsrTensor& x,
+                                         const SparseCsrTensor& y,
+                                         const SparseCsrTensor& dout,
+                                         SparseCsrTensor* dx,
+                                         SparseCsrTensor* dy) {
+  if (dx) {
+    AllocCsrPtr<T, IntT>(dev_ctx, x, dx);
+    SparseCsrTensor tmp_dx;
+    AllocCsrPtr<T, IntT>(dev_ctx, x, &tmp_dx);
+    if (std::is_same<T, phi::complex64>::value ||
+        std::is_same<T, phi::complex128>::value) {
+      //    dout*y_conj
+      SparseCsrTensor y_conj;
+      ConjugateCsrValues<T, IntT, Context>(dev_ctx, y, &y_conj);
+      sparse::ElementWiseMultiplyCsrKernel<T, Context>(
+          dev_ctx, dout, y_conj, &tmp_dx);
+    } else {
+      //    dout*y
+      sparse::ElementWiseMultiplyCsrKernel<T, Context>(
+          dev_ctx, dout, y, &tmp_dx);
+    }
+    CopyCsrValues<T, IntT, Context>(dev_ctx, tmp_dx, x, dx);
+  }
+
+  if (dy) {
+    AllocCsrPtr<T, IntT>(dev_ctx, y, dy);
+    SparseCsrTensor tmp_dy;
+    AllocCsrPtr<T, IntT>(dev_ctx, y, &tmp_dy);
+    if (std::is_same<T, phi::complex64>::value ||
+        std::is_same<T, phi::complex128>::value) {
+      //    dout*x_conj
+      SparseCsrTensor x_conj;
+      ConjugateCsrValues<T, IntT, Context>(dev_ctx, x, &x_conj);
+      sparse::ElementWiseMultiplyCsrKernel<T, Context>(
+          dev_ctx, dout, x_conj, &tmp_dy);
+    } else {
+      //    dout*x
+      sparse::ElementWiseMultiplyCsrKernel<T, Context>(
+          dev_ctx, dout, x, &tmp_dy);
+    }
+    CopyCsrValues<T, IntT, Context>(dev_ctx, tmp_dy, y, dy);
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void ElementWiseDivideCsrGradCPUKernel(const Context& dev_ctx,
+                                       const SparseCsrTensor& x,
+                                       const SparseCsrTensor& y,
+                                       const SparseCsrTensor& out,
+                                       const SparseCsrTensor& dout,
+                                       SparseCsrTensor* dx,
+                                       SparseCsrTensor* dy) {
+  if (dx) {
+    AllocCsrPtr<T, IntT>(dev_ctx, x, dx);
+    SparseCsrTensor tmp_dx;
+    AllocCsrPtr<T, IntT>(dev_ctx, x, &tmp_dx);
+    if (std::is_same<T, phi::complex64>::value ||
+        std::is_same<T, phi::complex128>::value) {
+      //    dout/y_conj
+      SparseCsrTensor y_conj;
+      ConjugateCsrValues<T, IntT, Context>(dev_ctx, y, &y_conj);
+      sparse::ElementWiseDivideCsrKernel<T, Context>(
+          dev_ctx, dout, y_conj, &tmp_dx);
+    } else {
+      //    dout/y
+      sparse::ElementWiseDivideCsrKernel<T, Context>(dev_ctx, dout, y, &tmp_dx);
+    }
+    CopyCsrValues<T, IntT, Context>(dev_ctx, tmp_dx, x, dx);
+  }
+
+  if (dy) {
+    //    -dout * out / y
+    AllocCsrPtr<T, IntT>(dev_ctx, y, dy);
+    SparseCsrTensor tmp_dy;
+    AllocCsrPtr<T, IntT>(dev_ctx, y, &tmp_dy);
+
+    Copy(dev_ctx, dout, dev_ctx.GetPlace(), false, &tmp_dy);
+    phi::NegativeKernel<T, Context>(
+        dev_ctx, dout.values(), tmp_dy.mutable_values());
+    if (std::is_same<T, phi::complex64>::value ||
+        std::is_same<T, phi::complex128>::value) {
+      //    -dout * (out / y)_conj = -dout * out_conj / y_conj
+      SparseCsrTensor out_conj;
+      ConjugateCsrValues<T, IntT, Context>(dev_ctx, out, &out_conj);
+      SparseCsrTensor y_conj;
+      ConjugateCsrValues<T, IntT, Context>(dev_ctx, y, &y_conj);
+      auto tmp =
+          sparse::ElementWiseMultiplyCsr<T, Context>(dev_ctx, tmp_dy, out_conj);
+      sparse::ElementWiseDivideCsrKernel<T, Context>(
+          dev_ctx, tmp, y_conj, &tmp_dy);
+    } else {
+      auto tmp =
+          sparse::ElementWiseMultiplyCsr<T, Context>(dev_ctx, tmp_dy, out);
+      sparse::ElementWiseDivideCsrKernel<T, Context>(dev_ctx, tmp, y, &tmp_dy);
+    }
+    CopyCsrValues<T, IntT, Context>(dev_ctx, tmp_dy, y, dy);
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void ElementWiseAddCooGradCPUKernel(const Context& dev_ctx,
+                                    const SparseCooTensor& x,
+                                    const SparseCooTensor& y,
+                                    const SparseCooTensor& dout,
+                                    SparseCooTensor* dx,
+                                    SparseCooTensor* dy) {
+  //     Special case when y_grad is not needed*/
+  if (dx != nullptr && dy == nullptr) {
+    VLOG(4) << "Special case when dy is not needed";
+    AllocCooPtr<T, IntT>(dev_ctx, x, dx);
+    CopyCooValues<T, IntT, Context>(dev_ctx, dout, x, dx);
+  } else if (dx == nullptr && dy != nullptr) {
+    VLOG(4) << "Special case when dx is not needed";
+    AllocCooPtr<T, IntT>(dev_ctx, y, dy);
+    CopyCooValues<T, IntT, Context>(dev_ctx, dout, y, dy);
+  } else {
+    AllocCooPtr<T, IntT>(dev_ctx, x, dx);
+    AllocCooPtr<T, IntT>(dev_ctx, y, dy);
+    CopyCooValues<T, IntT, Context>(dev_ctx, dout, x, dx);
+    CopyCooValues<T, IntT, Context>(dev_ctx, dout, y, dy);
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void ElementWiseSubtractCooGradCPUKernel(const Context& dev_ctx,
+                                         const SparseCooTensor& x,
+                                         const SparseCooTensor& y,
+                                         const SparseCooTensor& dout,
+                                         SparseCooTensor* dx,
+                                         SparseCooTensor* dy) {
+  if (dx) {
+    AllocCooPtr<T, IntT>(dev_ctx, x, dx);
+    CopyCooValues<T, IntT, Context>(dev_ctx, dout, x, dx);
+  }
+
+  if (dy) {
+    AllocCooPtr<T, IntT>(dev_ctx, y, dy);
+    CopyCooValues<T, IntT, Context>(dev_ctx, dout, y, dy);
+    phi::NegativeKernel<T, Context>(
+        dev_ctx, dout.values(), dy->mutable_values());
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void ElementWiseMultiplyCooGradCPUKernel(const Context& dev_ctx,
+                                         const SparseCooTensor& x,
+                                         const SparseCooTensor& y,
+                                         const SparseCooTensor& dout,
+                                         SparseCooTensor* dx,
+                                         SparseCooTensor* dy) {
+  if (dx) {
+    AllocCooPtr<T, IntT>(dev_ctx, x, dx);
+    SparseCooTensor tmp_dx;
+    AllocCooPtr<T, IntT>(dev_ctx, x, &tmp_dx);
+    if (std::is_same<T, phi::complex64>::value ||
+        std::is_same<T, phi::complex128>::value) {
+      //    dout*y_conj
+      SparseCooTensor y_conj;
+      ConjugateCooValues<T, IntT, Context>(dev_ctx, y, &y_conj);
+      sparse::ElementWiseMultiplyCooKernel<T, Context>(
+          dev_ctx, dout, y_conj, &tmp_dx);
+    } else {
+      //    dout*y
+      sparse::ElementWiseMultiplyCooKernel<T, Context>(
+          dev_ctx, dout, y, &tmp_dx);
+    }
+    CopyCooValues<T, IntT, Context>(dev_ctx, tmp_dx, x, dx);
+  }
+
+  if (dy) {
+    AllocCooPtr<T, IntT>(dev_ctx, y, dy);
+    SparseCooTensor tmp_dy;
+    AllocCooPtr<T, IntT>(dev_ctx, y, &tmp_dy);
+    if (std::is_same<T, phi::complex64>::value ||
+        std::is_same<T, phi::complex128>::value) {
+      //    dout*x_conj
+      SparseCooTensor x_conj;
+      ConjugateCooValues<T, IntT, Context>(dev_ctx, x, &x_conj);
+      sparse::ElementWiseMultiplyCooKernel<T, Context>(
+          dev_ctx, dout, x_conj, &tmp_dy);
+    } else {
+      //    dout*x
+      sparse::ElementWiseMultiplyCooKernel<T, Context>(
+          dev_ctx, dout, x, &tmp_dy);
+    }
+    CopyCooValues<T, IntT, Context>(dev_ctx, tmp_dy, y, dy);
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void ElementWiseDivideCooGradCPUKernel(const Context& dev_ctx,
+                                       const SparseCooTensor& x,
+                                       const SparseCooTensor& y,
+                                       const SparseCooTensor& out,
+                                       const SparseCooTensor& dout,
+                                       SparseCooTensor* dx,
+                                       SparseCooTensor* dy) {
+  if (dx) {
+    AllocCooPtr<T, IntT>(dev_ctx, x, dx);
+    SparseCooTensor tmp_dx;
+    AllocCooPtr<T, IntT>(dev_ctx, x, &tmp_dx);
+    if (std::is_same<T, phi::complex64>::value ||
+        std::is_same<T, phi::complex128>::value) {
+      //    dout/y_conj
+      SparseCooTensor y_conj;
+      ConjugateCooValues<T, IntT, Context>(dev_ctx, y, &y_conj);
+      sparse::ElementWiseDivideCooKernel<T, Context>(
+          dev_ctx, dout, y_conj, &tmp_dx);
+    } else {
+      //    dout/y
+      sparse::ElementWiseDivideCooKernel<T, Context>(dev_ctx, dout, y, &tmp_dx);
+    }
+    CopyCooValues<T, IntT, Context>(dev_ctx, tmp_dx, x, dx);
+  }
+
+  if (dy) {
+    //    -dout * out / y
+    AllocCooPtr<T, IntT>(dev_ctx, y, dy);
+    SparseCooTensor tmp_dy;
+    AllocCooPtr<T, IntT>(dev_ctx, y, &tmp_dy);
+    Copy(dev_ctx, dout, dev_ctx.GetPlace(), false, &tmp_dy);
+    phi::NegativeKernel<T, Context>(
+        dev_ctx, dout.values(), tmp_dy.mutable_values());
+    if (std::is_same<T, phi::complex64>::value ||
+        std::is_same<T, phi::complex128>::value) {
+      //    -dout * (out / y)_conj = -dout * out_conj / y_conj
+      SparseCooTensor out_conj;
+      ConjugateCooValues<T, IntT, Context>(dev_ctx, out, &out_conj);
+      SparseCooTensor y_conj;
+      ConjugateCooValues<T, IntT, Context>(dev_ctx, y, &y_conj);
+      auto tmp =
+          sparse::ElementWiseMultiplyCoo<T, Context>(dev_ctx, tmp_dy, out_conj);
+      sparse::ElementWiseDivideCooKernel<T, Context>(
+          dev_ctx, tmp, y_conj, &tmp_dy);
+    } else {
+      auto tmp =
+          sparse::ElementWiseMultiplyCoo<T, Context>(dev_ctx, tmp_dy, out);
+      sparse::ElementWiseDivideCooKernel<T, Context>(dev_ctx, tmp, y, &tmp_dy);
+    }
+    CopyCooValues<T, IntT, Context>(dev_ctx, tmp_dy, y, dy);
+  }
+}
+
+template <typename T, typename Context>
+void ElementWiseDivideCsrGradKernel(const Context& dev_ctx,
+                                    const SparseCsrTensor& x,
+                                    const SparseCsrTensor& y,
+                                    const SparseCsrTensor& out,
+                                    const SparseCsrTensor& dout,
+                                    SparseCsrTensor* dx,
+                                    SparseCsrTensor* dy) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.crows().dtype(), "ElementWiseDivideCsrGradCPUKernel", ([&] {
+        ElementWiseDivideCsrGradCPUKernel<T, data_t>(
+            dev_ctx, x, y, out, dout, dx, dy);
+      }));
+}
+template <typename T, typename Context>
+void ElementWiseDivideCooGradKernel(const Context& dev_ctx,
+                                    const SparseCooTensor& x,
+                                    const SparseCooTensor& y,
+                                    const SparseCooTensor& out,
+                                    const SparseCooTensor& dout,
+                                    SparseCooTensor* dx,
+                                    SparseCooTensor* dy) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.indices().dtype(), "ElementWiseDivideCooGradCPUKernel", ([&] {
+        ElementWiseDivideCooGradCPUKernel<T, data_t>(
+            dev_ctx, x, y, out, dout, dx, dy);
+      }));
+}
+
+#define DEFINE_ELEMENTWISE_GRAD_KERNEL(name) \
+  DEFINE_ELEMENTWISE_GRAD_KERNEL_CSR(name)   \
+                                             \
+  DEFINE_ELEMENTWISE_GRAD_KERNEL_COO(name)
+
+#define DEFINE_ELEMENTWISE_GRAD_KERNEL_CSR(name)                         \
+  template <typename T, typename Context>                                \
+  void ElementWise##name##CsrGradKernel(const Context& dev_ctx,          \
+                                        const SparseCsrTensor& x,        \
+                                        const SparseCsrTensor& y,        \
+                                        const SparseCsrTensor& dout,     \
+                                        SparseCsrTensor* dx,             \
+                                        SparseCsrTensor* dy) {           \
+    PD_VISIT_BASE_INTEGRAL_TYPES(                                        \
+        x.crows().dtype(), "ElementWise##name##CsrGradCPUKernel", ([&] { \
+          ElementWise##name##CsrGradCPUKernel<T, data_t>(                \
+              dev_ctx, x, y, dout, dx, dy);                              \
+        }));                                                             \
+  }
+
+#define DEFINE_ELEMENTWISE_GRAD_KERNEL_COO(name)                           \
+  template <typename T, typename Context>                                  \
+  void ElementWise##name##CooGradKernel(const Context& dev_ctx,            \
+                                        const SparseCooTensor& x,          \
+                                        const SparseCooTensor& y,          \
+                                        const SparseCooTensor& dout,       \
+                                        SparseCooTensor* dx,               \
+                                        SparseCooTensor* dy) {             \
+    PD_VISIT_BASE_INTEGRAL_TYPES(                                          \
+        x.indices().dtype(), "ElementWise##name##CooGradCPUKernel", ([&] { \
+          ElementWise##name##CooGradCPUKernel<T, data_t>(                  \
+              dev_ctx, x, y, dout, dx, dy);                                \
+        }));                                                               \
+  }
+
+DEFINE_ELEMENTWISE_GRAD_KERNEL(Add)
+DEFINE_ELEMENTWISE_GRAD_KERNEL(Subtract)
+DEFINE_ELEMENTWISE_GRAD_KERNEL(Multiply)
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(add_csr_csr_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseAddCsrGradKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(2).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(subtract_csr_csr_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseSubtractCsrGradKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(2).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(multiply_csr_csr_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseMultiplyCsrGradKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(2).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(divide_csr_csr_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseDivideCsrGradKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(2).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(3).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(add_coo_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseAddCooGradKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(2).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(subtract_coo_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseSubtractCooGradKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(2).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(multiply_coo_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseMultiplyCooGradKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(2).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(divide_coo_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseDivideCooGradKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(2).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(3).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(add_coo_dense_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseAddDenseGradKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/cpu/elementwise_kernel.cc

@@ -0,0 +1,463 @@
+/* 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_kernel.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_meta.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/elementwise_add_kernel.h"
+#include "paddle/phi/kernels/elementwise_kernel.h"
+#include "paddle/phi/kernels/funcs/elementwise_functor.h"
+#include "paddle/phi/kernels/funcs/sparse/flatten_indices.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Functor>
+struct BinaryOPWithZeroCompareFunctor {
+  explicit BinaryOPWithZeroCompareFunctor(Functor functor)
+      : functor_(functor) {}
+  inline HOSTDEVICE void operator()(const T* a,
+                                    const T* b,
+                                    T* result,
+                                    const int64_t len) const {
+    for (int64_t i = 0; i < len; ++i) {
+      result[i] = functor_(a[i], b[i]);
+    }
+  }
+  Functor functor_;
+};
+
+template <typename T, typename IntT, typename Functor>
+void Merge(const IntT el_len,
+           const IntT* a_index,
+           const T* a_values,
+           const IntT len_a,
+           const IntT* b_index_org,
+           const T* b_values_org,
+           const IntT len_b,
+           const IntT len_b_max,
+           IntT* c_index,
+           T* c_values,
+           IntT* out_nnz,
+           const Functor& functor_org,
+           const bool is_divide) {
+  IntT a = 0;
+  IntT b = 0;
+  IntT& nnz = (*out_nnz);
+  nnz = 0;
+  const IntT* b_index = nullptr;
+  std::vector<IntT> b_full_index;
+  const std::vector<T> zero(el_len, 0);
+  auto functor = BinaryOPWithZeroCompareFunctor<T, Functor>(functor_org);
+
+  std::vector<const T*> b_values(len_b_max, zero.data());
+  for (auto i = 0; i < len_b; ++i) {
+    b_values[b_index_org[i]] = b_values_org + i * el_len;
+  }
+  //  if is divide expend b_index_org to b_full_index
+  if (is_divide) {
+    b_full_index = std::vector<IntT>(len_b_max);
+    for (int64_t j = 0; j < static_cast<int64_t>(b_full_index.size()); ++j) {
+      b_full_index[j] = j;
+    }
+    b_index = b_full_index.data();
+  } else {
+    b_index = b_index_org;
+  }
+  // merge
+  while (a < len_a && b < (is_divide ? len_b_max : len_b)) {
+    if (a_index[a] == b_index[b]) {
+      functor(a_values + a * el_len,
+              b_values[b_index[b]],
+              c_values + nnz * el_len,
+              el_len);
+      c_index[nnz] = a_index[a];
+      ++nnz;
+      ++a;
+      ++b;
+    } else if (a_index[a] < b_index[b]) {  // coordinate x[a] < coordinate y[b]
+      functor(
+          a_values + a * el_len, zero.data(), c_values + nnz * el_len, el_len);
+      c_index[nnz] = a_index[a];
+      ++nnz;
+      ++a;
+    } else if (a_index[a] > b_index[b]) {  // coordinate x[a] > coordinate y[b]
+      functor(
+          zero.data(), b_values[b_index[b]], c_values + nnz * el_len, el_len);
+      c_index[nnz] = b_index[b];
+      ++nnz;
+      ++b;
+    }
+  }
+  // a tail
+  while (a < len_a) {
+    functor(
+        a_values + a * el_len, zero.data(), c_values + nnz * el_len, el_len);
+    c_index[nnz] = a_index[a];
+    ++nnz;
+    ++a;
+  }
+  //  b tail
+  while (b < (is_divide ? len_b_max : len_b)) {
+    functor(zero.data(), b_values[b_index[b]], c_values + nnz * el_len, el_len);
+    c_index[nnz] = b_index[b];
+    ++nnz;
+    ++b;
+  }
+}
+
+// SparseCooTensor elementwise op, only support same shape tensor now
+template <typename T, typename IntT, typename Context, typename Functor>
+void ElementWiseCooKernelImpl(const Context& dev_ctx,
+                              const SparseCooTensor& x,
+                              const SparseCooTensor& y,
+                              SparseCooTensor* out,
+                              const Functor& functor) {
+  PADDLE_ENFORCE_EQ(x.dims(),
+                    y.dims(),
+                    common::errors::InvalidArgument(
+                        "Currently only support same shape elementwise "
+                        "compute. The input tensor X's shape "
+                        "should be identical with Y's shape. But received X's "
+                        "shape = [%s], Y's shape = [%s].",
+                        x.dims(),
+                        y.dims()));
+
+  // temporary policy: for broadcast add
+  // TODO(zhangkaihuo): implement a correct function
+  const bool is_add = std::is_same<Functor, funcs::AddFunctor<T>>::value;
+  if (is_add && x.indices().numel() == y.indices().numel()) {
+    int compare_indices = memcmp(x.indices().data<IntT>(),
+                                 y.indices().data<IntT>(),
+                                 sizeof(IntT) * x.indices().numel());
+    if (compare_indices == 0) {
+      EmptyLikeCooKernel<T, Context>(dev_ctx, x, out);
+      phi::AddKernel<T, Context>(
+          dev_ctx, x.values(), y.values(), out->mutable_values());
+      return;
+    }
+  }
+  int64_t element_size = 1;
+  for (auto j = 1; j < x.values().dims().size(); ++j) {
+    element_size *= x.values().dims()[j];
+  }
+  IntT nnz = 0;
+  const auto x_values = x.values().data<T>();
+  const auto y_values = y.values().data<T>();
+  const auto sparse_dim = x.indices().dims()[0];
+  const bool is_divide = std::is_same<Functor, funcs::DivideFunctor<T>>::value;
+
+  int64_t max_len = 1;
+  for (auto j = 0; j < sparse_dim; ++j) {
+    max_len *= x.dims()[j];
+  }
+
+  std::vector<IntT> sparse_offsets(sparse_dim), x_indices(x.nnz()),
+      y_indices(y.nnz());
+
+  funcs::sparse::CalcOffsetsPerDim<IntT>(
+      x.dims(), sparse_dim, sparse_offsets.data());
+
+  funcs::sparse::FlattenIndices(x.indices().data<IntT>(),
+                                sparse_offsets.data(),
+                                x.nnz(),
+                                sparse_dim,
+                                0,
+                                1,
+                                x_indices.data());
+
+  funcs::sparse::FlattenIndices(y.indices().data<IntT>(),
+                                sparse_offsets.data(),
+                                y.nnz(),
+                                sparse_dim,
+                                0,
+                                1,
+                                y_indices.data());
+
+  std::vector<IntT> out_indices;
+  std::vector<T> out_values_vec;
+  if (is_divide) {
+    out_indices.reserve(max_len);
+  } else {
+    out_indices.reserve(x.nnz() + y.nnz());
+  }
+  out_values_vec.reserve(max_len * element_size);
+
+  //  merge x and y
+  Merge<T, IntT, Functor>(element_size,
+                          x_indices.data(),
+                          x_values,
+                          x_indices.size(),
+                          y_indices.data(),
+                          y_values,
+                          y_indices.size(),
+                          max_len,
+                          out_indices.data(),
+                          out_values_vec.data(),
+                          &nnz,
+                          functor,
+                          is_divide);
+
+  std::vector<IntT> out_indices_vec;
+  out_indices_vec.resize(nnz * sparse_dim);
+
+  Dim<DDim::kMaxRank> const_dims;
+  for (auto i = 0; i < x.dims().size(); i++) {
+    const_dims[i] = x.dims()[i];
+  }
+
+  funcs::sparse::IndexToCoordinate<IntT>(out_indices.data(),
+                                         const_dims,
+                                         nnz,
+                                         sparse_dim,
+                                         0,
+                                         1,
+                                         out_indices_vec.data());
+
+  if (nnz == 0) {
+    DenseTensor out_indices = EmptyLike<IntT>(dev_ctx, x.indices());
+    DenseTensor out_values = EmptyLike<T>(dev_ctx, x.values());
+    out->SetMember(out_indices, out_values, x.dims());
+  } else {
+    DenseTensorMeta indices_meta(phi::CppTypeToDataType<IntT>::Type(),
+                                 make_ddim({static_cast<int64_t>(sparse_dim),
+                                            static_cast<int64_t>(nnz)}),
+                                 DataLayout::NCHW);
+    auto indices_dim =
+        vectorize(slice_ddim(x.values().dims(), 1, x.values().dims().size()));
+    indices_dim.insert(indices_dim.begin(), nnz);
+    DenseTensorMeta values_meta(
+        x.dtype(), make_ddim(indices_dim), DataLayout::NCHW);
+    DenseTensor out_indices = Empty(dev_ctx, std::move(indices_meta));
+    DenseTensor out_values = Empty(dev_ctx, std::move(values_meta));
+
+    std::memcpy(out_indices.data<IntT>(),
+                out_indices_vec.data(),
+                sizeof(IntT) * sparse_dim * nnz);
+    std::memcpy(out_values.data<T>(),
+                out_values_vec.data(),
+                sizeof(T) * nnz * element_size);
+
+    out->SetMember(out_indices, out_values, x.dims());
+  }
+}
+
+#define DEFINE_CSR_ELEMENTWISE_CPU_KERNEL(name)                               \
+  template <typename T, typename IntT, typename Context>                      \
+  void ElementWise##name##CsrCPUKernel(const Context& dev_ctx,                \
+                                       const SparseCsrTensor& x,              \
+                                       const SparseCsrTensor& y,              \
+                                       SparseCsrTensor* out) {                \
+    auto coo_x = CsrToCoo<T>(dev_ctx, x);                                     \
+    auto coo_y = CsrToCoo<T>(dev_ctx, y);                                     \
+    auto coo_out = ElementWise##name##Coo<T, Context>(dev_ctx, coo_x, coo_y); \
+    CooToCsrKernel<T>(dev_ctx, coo_out, out);                                 \
+  }
+
+#define DEFINE_CSR_ELEMENTWISE_KERNEL(name)                               \
+  template <typename T, typename Context>                                 \
+  void ElementWise##name##CsrKernel(const Context& dev_ctx,               \
+                                    const SparseCsrTensor& x,             \
+                                    const SparseCsrTensor& y,             \
+                                    SparseCsrTensor* out) {               \
+    PD_VISIT_BASE_INTEGRAL_TYPES(                                         \
+        x.crows().dtype(), "ElementWise##name##CsrCPUKernel", ([&] {      \
+          ElementWise##name##CsrCPUKernel<T, data_t>(dev_ctx, x, y, out); \
+        }));                                                              \
+  }
+
+#define DEFINE_COO_ELEMENTWISE_CPU_KERNEL(name)                          \
+  template <typename T, typename IntT, typename Context>                 \
+  void ElementWise##name##CooCPUKernel(const Context& dev_ctx,           \
+                                       const SparseCooTensor& x,         \
+                                       const SparseCooTensor& y,         \
+                                       SparseCooTensor* out) {           \
+    funcs::name##Functor<T> functor;                                     \
+    ElementWiseCooKernelImpl<T, IntT, Context, funcs::name##Functor<T>>( \
+        dev_ctx, x, y, out, functor);                                    \
+  }
+
+#define DEFINE_COO_ELEMENTWISE_KERNEL(name)                               \
+  template <typename T, typename Context>                                 \
+  void ElementWise##name##CooKernel(const Context& dev_ctx,               \
+                                    const SparseCooTensor& x,             \
+                                    const SparseCooTensor& y,             \
+                                    SparseCooTensor* out) {               \
+    PD_VISIT_BASE_INTEGRAL_TYPES(                                         \
+        x.indices().dtype(), "ElementWise##name##CooCPUKernel", ([&] {    \
+          ElementWise##name##CooCPUKernel<T, data_t>(dev_ctx, x, y, out); \
+        }));                                                              \
+  }
+
+DEFINE_CSR_ELEMENTWISE_CPU_KERNEL(Add)
+DEFINE_CSR_ELEMENTWISE_CPU_KERNEL(Subtract)
+DEFINE_CSR_ELEMENTWISE_CPU_KERNEL(Multiply)
+DEFINE_CSR_ELEMENTWISE_CPU_KERNEL(Divide)
+
+DEFINE_CSR_ELEMENTWISE_KERNEL(Add)
+DEFINE_CSR_ELEMENTWISE_KERNEL(Subtract)
+DEFINE_CSR_ELEMENTWISE_KERNEL(Multiply)
+DEFINE_CSR_ELEMENTWISE_KERNEL(Divide)
+
+DEFINE_COO_ELEMENTWISE_CPU_KERNEL(Add)
+DEFINE_COO_ELEMENTWISE_CPU_KERNEL(Subtract)
+DEFINE_COO_ELEMENTWISE_CPU_KERNEL(Multiply)
+DEFINE_COO_ELEMENTWISE_CPU_KERNEL(Divide)
+
+DEFINE_COO_ELEMENTWISE_KERNEL(Add)
+DEFINE_COO_ELEMENTWISE_KERNEL(Subtract)
+DEFINE_COO_ELEMENTWISE_KERNEL(Multiply)
+DEFINE_COO_ELEMENTWISE_KERNEL(Divide)
+
+}  // namespace phi::sparse
+
+using complex64 = phi::complex64;
+using complex128 = phi::complex128;
+
+PD_REGISTER_KERNEL(add_csr_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseAddCsrKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   complex64,
+                   complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(add_coo_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseAddCooKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   complex64,
+                   complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(subtract_csr_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseSubtractCsrKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   complex64,
+                   complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(subtract_coo_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseSubtractCooKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   complex64,
+                   complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(multiply_csr_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseMultiplyCsrKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   complex64,
+                   complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(multiply_coo_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseMultiplyCooKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   complex64,
+                   complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(divide_csr_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseDivideCsrKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   complex64,
+                   complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(divide_coo_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseDivideCooKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   complex64,
+                   complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(add_coo_dense,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseAddDenseKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   complex64,
+                   complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/cpu/full_kernel.cc

@@ -0,0 +1,115 @@
+/* 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/full_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void FullValue(const Context& dev_ctx, DenseTensor* tensor, T val) {
+  dev_ctx.template Alloc<T>(tensor);
+  auto t = EigenVector<T>::Flatten(*tensor);
+  t.device(*dev_ctx.eigen_device()) = t.constant(val);
+}
+
+template <typename T, typename Context>
+void FullLikeCooKernel(const Context& dev_ctx,
+                       const SparseCooTensor& x,
+                       const Scalar& val,
+                       DataType dtype UNUSED,
+                       SparseCooTensor* out) {
+  phi::Copy<Context>(dev_ctx,
+                     x.non_zero_indices(),
+                     dev_ctx.GetPlace(),
+                     false,
+                     out->mutable_non_zero_indices());
+
+  DenseTensor* values = out->mutable_non_zero_elements();
+  values->Resize(x.non_zero_elements().dims());
+  dev_ctx.template Alloc<T>(values);
+  FullValue<T, Context>(dev_ctx, values, val.to<T>());
+
+  out->set_dims(x.dims());
+}
+
+template <typename T, typename Context>
+void FullLikeCsrKernel(const Context& dev_ctx,
+                       const SparseCsrTensor& x,
+                       const Scalar& val,
+                       DataType dtype UNUSED,
+                       SparseCsrTensor* out) {
+  phi::Copy<Context>(dev_ctx,
+                     x.non_zero_crows(),
+                     dev_ctx.GetPlace(),
+                     false,
+                     out->mutable_non_zero_crows());
+
+  phi::Copy<Context>(dev_ctx,
+                     x.non_zero_cols(),
+                     dev_ctx.GetPlace(),
+                     false,
+                     out->mutable_non_zero_cols());
+
+  DenseTensor* values = out->mutable_non_zero_elements();
+  values->Resize(x.non_zero_elements().dims());
+  dev_ctx.template Alloc<T>(values);
+  FullValue<T, Context>(dev_ctx, values, val.to<T>());
+
+  out->set_dims(x.dims());
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(full_like_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::FullLikeCooKernel,
+                   float,
+                   double,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::bfloat16,
+                   phi::float16,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(full_like_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::FullLikeCsrKernel,
+                   float,
+                   double,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::bfloat16,
+                   phi::float16,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}

paddle/phi/kernels/sparse/cpu/fused_attention_grad_kernel.cc

@@ -0,0 +1,36 @@
+/* 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/fused_attention_grad_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void FusedAttentionCsrGradKernel(const Context& dev_ctx,
+                                 const DenseTensor& query,
+                                 const DenseTensor& key,
+                                 const DenseTensor& value,
+                                 const SparseCsrTensor& softmax,
+                                 const DenseTensor& dout,
+                                 DenseTensor* dquery,
+                                 DenseTensor* dkey,
+                                 DenseTensor* dvalue) {
+  PD_THROW(
+      "Not support CPU kernel of 'sparse.nn.functional.fused_attention' now");
+}
+
+}  // namespace phi::sparse

paddle/phi/kernels/sparse/cpu/fused_attention_kernel.cc

@@ -0,0 +1,36 @@
+/* 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/fused_attention_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void FusedAttentionCsrKernel(const Context& dev_ctx,
+                             const DenseTensor& query,
+                             const DenseTensor& key,
+                             const DenseTensor& value,
+                             const SparseCsrTensor& sparse_mask,
+                             const optional<DenseTensor>& key_padding_mask,
+                             const optional<DenseTensor>& attn_mask,
+                             DenseTensor* out,
+                             SparseCsrTensor* softmax) {
+  PD_THROW(
+      "Not support CPU kernel of 'sparse.nn.functional.fused_attention' now");
+}
+
+}  // namespace phi::sparse

paddle/phi/kernels/sparse/cpu/mask_grad_kernel.cc

@@ -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/mask_grad_kernel.h"
+#include "paddle/phi/kernels/sparse/mask_kernel.h"
+#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(mask_as_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaskAsCooGradKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(mask_as_csr_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaskAsCsrGradKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}

paddle/phi/kernels/sparse/cpu/mask_kernel.cc

@@ -0,0 +1,307 @@
+/* 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/mask_kernel.h"
+#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h"
+
+#include "paddle/common/ddim.h"
+#include "paddle/phi/api/ext/dispatch.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/sparse/flatten_indices.h"
+
+namespace phi::sparse {
+
+template <typename T, typename IntT>
+void MaskCooCPUKernel(const CPUContext& dev_ctx,
+                      const DenseTensor& x,
+                      const SparseCooTensor& mask,
+                      SparseCooTensor* out) {
+  const DDim& dims = x.dims();
+  PADDLE_ENFORCE_EQ(x.dims(),
+                    mask.dims(),
+                    common::errors::InvalidArgument(
+                        "the input x and mask must have the shape"));
+  const DenseTensor& indices = mask.indices();
+  const DenseTensor& values = mask.values();
+  const int sparse_dim = mask.sparse_dim();
+
+  DenseTensor out_indices = EmptyLike<T>(dev_ctx, indices);
+  DenseTensor out_values = EmptyLike<T>(dev_ctx, values);
+
+  // the out_indices is same as indices of mask
+  phi::Copy(dev_ctx, indices, dev_ctx.GetPlace(), false, &out_indices);
+
+  T* out_values_ptr = out_values.data<T>();
+  const T* x_ptr = x.data<T>();
+
+  const int64_t non_zero_num = mask.nnz();
+  auto dims_2d = flatten_to_2d(dims, sparse_dim);
+  const int cols = static_cast<int>(dims_2d[1]);
+  const IntT* indices_ptr = indices.data<IntT>();
+
+  std::vector<IntT> sparse_offsets(sparse_dim);
+
+  funcs::sparse::CalcOffsetsPerDim<IntT>(
+      dims, sparse_dim, sparse_offsets.data());
+
+  for (int64_t i = 0; i < non_zero_num; i++) {
+    int64_t index = funcs::sparse::CoordinateToIndex<IntT>(
+        indices_ptr, sparse_offsets.data(), non_zero_num, sparse_dim, i);
+    memcpy(out_values_ptr + i * cols, x_ptr + index * cols, cols * sizeof(T));
+  }
+
+  out->SetMember(out_indices, out_values, dims, true);
+}
+
+/**
+ * @brief Filter the DenseTensor x by the
+ * mask.indices() and output a SparseCooTensor
+ * x and mask must have the same shape.
+ **/
+template <typename T, typename Context>
+void MaskAsCooKernel(const Context& dev_ctx,
+                     const DenseTensor& x,
+                     const SparseCooTensor& mask,
+                     SparseCooTensor* out) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      mask.indices().dtype(), "MaskCooCPUKernel", ([&] {
+        MaskCooCPUKernel<T, data_t>(dev_ctx, x, mask, out);
+      }));
+}
+
+template <typename T, typename IntT>
+void MaskCsr2DCPUKernel(const CPUContext& dev_ctx,
+                        const DenseTensor& x,
+                        const SparseCsrTensor& mask,
+                        SparseCsrTensor* out) {
+  const DenseTensor& mask_cols = mask.cols();
+  const DenseTensor& mask_crows = mask.crows();
+  int64_t num_non_zeros = mask.nnz();
+
+  DenseTensor out_cols = EmptyLike<IntT>(dev_ctx, mask_cols);
+  DenseTensor out_crows = EmptyLike<IntT>(dev_ctx, mask_crows);
+  DenseTensor out_values = Empty<T>(dev_ctx, {num_non_zeros});
+
+  phi::Copy(dev_ctx, mask_cols, dev_ctx.GetPlace(), false, &out_cols);
+  phi::Copy(dev_ctx, mask_crows, dev_ctx.GetPlace(), false, &out_crows);
+
+  int64_t numel = 0;
+  for (int64_t i = 0; i < mask_crows.numel() - 1; ++i) {
+    for (int64_t j = mask_crows.data<IntT>()[i];
+         j < mask_crows.data<IntT>()[i + 1];
+         ++j) {
+      IntT col_idx = mask_cols.data<IntT>()[numel];
+
+      out_values.data<T>()[numel] =
+          x.data<T>()[(i / x.dims()[0]) * x.dims()[1] +
+                      (i % x.dims()[0]) * x.dims()[1] + col_idx];
+
+      ++numel;
+    }
+  }
+
+  out->SetMember(out_crows, out_cols, out_values, x.dims());
+}
+
+template <typename T, typename IntT>
+void MaskCsr3DCPUKernel(const CPUContext& dev_ctx,
+                        const DenseTensor& x,
+                        const SparseCsrTensor& mask,
+                        SparseCsrTensor* out) {
+  const DenseTensor& mask_cols = mask.cols();
+  const DenseTensor& mask_crows = mask.crows();
+  int64_t num_non_zeros = mask.nnz();
+
+  DenseTensor out_cols = EmptyLike<IntT>(dev_ctx, mask_cols);
+  DenseTensor out_crows = EmptyLike<IntT>(dev_ctx, mask_crows);
+  DenseTensor out_values = Empty<T>(dev_ctx, {num_non_zeros});
+
+  phi::Copy(dev_ctx, mask_cols, dev_ctx.GetPlace(), false, &out_cols);
+  phi::Copy(dev_ctx, mask_crows, dev_ctx.GetPlace(), false, &out_crows);
+
+  int64_t numel = 0;
+  for (int64_t i = 0; i < mask_crows.numel() - 1; ++i) {
+    for (int64_t j = mask_crows.data<IntT>()[i];
+         j < mask_crows.data<IntT>()[i + 1];
+         ++j) {
+      IntT col_idx = mask_cols.data<IntT>()[numel];
+
+      out_values.data<T>()[numel] =
+          x.data<T>()[(i / (mask_crows.numel() / x.dims()[0])) *
+                          (x.dims()[1] * x.dims()[2]) +
+                      (i % (mask_crows.numel() / x.dims()[0])) * x.dims()[2] +
+                      col_idx];
+
+      ++numel;
+    }
+  }
+
+  out->SetMember(out_crows, out_cols, out_values, x.dims());
+}
+
+/**
+ * @brief Filter the DenseTensor x by the
+ * mask.crows(), mask.cols() and output a SparseCsrTensor
+ * x and mask must have the same shape.
+ **/
+template <typename T, typename Context>
+void MaskAsCsrKernel(const Context& dev_ctx,
+                     const DenseTensor& x,
+                     const SparseCsrTensor& mask,
+                     SparseCsrTensor* out) {
+  const DDim& x_dims = x.dims();
+  if (x_dims.size() == 2) {
+    PD_VISIT_BASE_INTEGRAL_TYPES(
+        mask.crows().dtype(), "MaskCsr2DCPUKernel", ([&] {
+          MaskCsr2DCPUKernel<T, data_t>(dev_ctx, x, mask, out);
+        }));
+  } else if (x_dims.size() == 3) {
+    PD_VISIT_BASE_INTEGRAL_TYPES(
+        mask.crows().dtype(), "MaskCsr3DCPUKernel", ([&] {
+          MaskCsr3DCPUKernel<T, data_t>(dev_ctx, x, mask, out);
+        }));
+  } else {
+    // throw exception
+    common::errors::InvalidArgument(
+        "mask_as for Sparse CSR Tensor only support 2-D or 3-D, but got "
+        "%d-D.",
+        x_dims.size());
+  }
+}
+
+template <typename T, typename IntT>
+void MaskHelperCooCPUKernel(const CPUContext& dev_ctx,
+                            const SparseCooTensor& x,
+                            const DenseTensor& mask_indices,
+                            DenseTensor* out) {
+  PADDLE_ENFORCE_EQ(
+      mask_indices.dims().size(),
+      2,
+      common::errors::InvalidArgument("the mask_indices must be 2-D tensor"));
+
+  const int32_t sparse_dim = x.sparse_dim();
+
+  std::vector<IntT> sparse_offsets(sparse_dim), x_indices(x.nnz()),
+      mask_out_indices(mask_indices.dims()[1]);
+  funcs::sparse::CalcOffsetsPerDim<IntT>(
+      x.dims(), sparse_dim, sparse_offsets.data());
+
+  funcs::sparse::FlattenIndices(x.indices().data<IntT>(),
+                                sparse_offsets.data(),
+                                x.nnz(),
+                                sparse_dim,
+                                0,
+                                1,
+                                x_indices.data());
+  funcs::sparse::FlattenIndices(mask_indices.data<IntT>(),
+                                sparse_offsets.data(),
+                                x.nnz(),
+                                sparse_dim,
+                                0,
+                                1,
+                                mask_out_indices.data());
+
+  std::unordered_map<IntT, uint64_t> x_indices_map;
+  for (uint64_t i = 0; i < x_indices.size(); i++) {
+    x_indices_map[x_indices[i]] = i;
+  }
+
+  *out = EmptyLike<T>(dev_ctx, x.values());
+  funcs::SetConstant<CPUContext, T> set_zero;
+  set_zero(dev_ctx, out, static_cast<T>(0));
+  T* out_ptr = out->data<T>();
+  const int64_t stride =
+      x.dims().size() == sparse_dim ? 1 : x.values().dims()[1];
+  const T* in_ptr = x.values().data<T>();
+  // TODO(zhangkaihuo): multithreading can be used for acceleration
+  for (uint64_t i = 0; i < mask_out_indices.size(); i++) {
+    auto iter = x_indices_map.find(mask_out_indices[i]);
+    if (iter != x_indices_map.end()) {
+      memcpy(out_ptr + i * stride,
+             in_ptr + iter->second * stride,
+             stride * sizeof(T));
+    }
+  }
+}
+
+/**
+ * @brief filter values from x.values() using mask_indices
+ */
+template <typename T, typename Context>
+void MaskHelperCooKernel(const Context& dev_ctx,
+                         const SparseCooTensor& x,
+                         const DenseTensor& mask_indices,
+                         DenseTensor* out) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.indices().dtype(), "MaskHelperCooCPUKernel", ([&] {
+        MaskHelperCooCPUKernel<T, data_t>(dev_ctx, x, mask_indices, out);
+      }));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(mask_helper_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaskHelperCooKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(mask_as_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaskAsCooKernel,
+                   float,
+                   double,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(mask_as_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaskAsCsrKernel,
+                   float,
+                   double,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}

paddle/phi/kernels/sparse/cpu/matmul_grad_kernel.cc

@@ -0,0 +1,62 @@
+/* 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/matmul_grad_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi::sparse {
+
+// TODO(zhouwei25): implement CPU backward kernel of " CSR @ DENSE -> DENSE"
+template <typename T, typename Context>
+void MatmulCsrDenseGradKernel(const Context& dev_ctx UNUSED,
+                              const SparseCsrTensor& x UNUSED,
+                              const DenseTensor& y UNUSED,
+                              const DenseTensor& dout UNUSED,
+                              SparseCsrTensor* dx UNUSED,
+                              DenseTensor* dy UNUSED) {
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support CPU backward kernel of 'sparse.matmul' now."));
+}
+
+// TODO(zhouwei25): implement CPU kernel of " DENSE @ DENSE * CSR_MASK -> CSR"
+template <typename T, typename Context>
+void MaskedMatmulCsrGradKernel(const Context& dev_ctx UNUSED,
+                               const DenseTensor& x UNUSED,
+                               const DenseTensor& y UNUSED,
+                               const SparseCsrTensor& dout UNUSED,
+                               DenseTensor* dx UNUSED,
+                               DenseTensor* dy UNUSED) {
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support CPU backward kernel of 'sparse.masked_matmul' now."));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(matmul_csr_dense_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MatmulCsrDenseGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(masked_matmul_csr_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaskedMatmulCsrGradKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/sparse/cpu/matmul_kernel.cc

@@ -0,0 +1,59 @@
+/* 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/matmul_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi::sparse {
+
+// TODO(zhouwei25): implement CPU kernel of " CSR @ DENSE -> DENSE"
+template <typename T, typename Context>
+void MatmulCsrDenseKernel(const Context& dev_ctx UNUSED,
+                          const SparseCsrTensor& x UNUSED,
+                          const DenseTensor& y UNUSED,
+                          DenseTensor* out UNUSED) {
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support CPU kernel of 'sparse.matmul' now."));
+}
+
+// TODO(zhouwei25): implement CPU kernel of " DENSE @ DENSE * CSR_MASK -> CSR"
+template <typename T, typename Context>
+void MaskedMatmulCsrKernel(const Context& dev_ctx UNUSED,
+                           const DenseTensor& x UNUSED,
+                           const DenseTensor& y UNUSED,
+                           const SparseCsrTensor& mask UNUSED,
+                           SparseCsrTensor* out UNUSED) {
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support CPU kernel of 'sparse.masked_matmul' now."));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(matmul_csr_dense,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MatmulCsrDenseKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(masked_matmul_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaskedMatmulCsrKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/sparse/cpu/mv_grad_kernel.cc

@@ -0,0 +1,54 @@
+/* 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/mv_grad_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void MvCooGradKernel(const Context& dev_ctx UNUSED,
+                     const SparseCooTensor& x UNUSED,
+                     const DenseTensor& vec UNUSED,
+                     const DenseTensor& dout UNUSED,
+                     SparseCooTensor* dx UNUSED,
+                     DenseTensor* dvec UNUSED) {
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support CPU backward kernel of 'sparse.mv' now."));
+}
+
+template <typename T, typename Context>
+void MvCsrGradKernel(const Context& dev_ctx UNUSED,
+                     const SparseCsrTensor& x UNUSED,
+                     const DenseTensor& vec UNUSED,
+                     const DenseTensor& dout UNUSED,
+                     SparseCsrTensor* dx UNUSED,
+                     DenseTensor* dvec UNUSED) {
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support CPU backward kernel of 'sparse.mv' now."));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(
+    mv_coo_grad, CPU, ALL_LAYOUT, phi::sparse::MvCooGradKernel, float, double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(
+    mv_csr_grad, CPU, ALL_LAYOUT, phi::sparse::MvCsrGradKernel, float, double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}

paddle/phi/kernels/sparse/cpu/mv_kernel.cc

@@ -0,0 +1,50 @@
+/* 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/mv_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void MvCsrKernel(const Context& dev_ctx UNUSED,
+                 const SparseCsrTensor& x UNUSED,
+                 const DenseTensor& vec UNUSED,
+                 DenseTensor* out UNUSED) {
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support CPU kernel of 'sparse.mv' now."));
+}
+
+template <typename T, typename Context>
+void MvCooKernel(const Context& dev_ctx UNUSED,
+                 const SparseCooTensor& x UNUSED,
+                 const DenseTensor& vec UNUSED,
+                 DenseTensor* out UNUSED) {
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support CPU kernel of 'sparse.mv' now."));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(
+    mv_csr, CPU, ALL_LAYOUT, phi::sparse::MvCsrKernel, float, double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(
+    mv_coo, CPU, ALL_LAYOUT, phi::sparse::MvCooKernel, float, double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/cpu/pool_grad_kernel.cc

@@ -0,0 +1,97 @@
+/* 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/pool_grad_kernel.h"
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/pooling.h"
+#include "paddle/phi/kernels/funcs/sparse/convolution.h"
+
+namespace phi::sparse {
+
+template <typename T, typename IntT = int>
+void MaxPoolCooGradCPUKernel(const CPUContext& dev_ctx,
+                             const SparseCooTensor& x,
+                             const DenseTensor& rulebook,
+                             const DenseTensor& counter,
+                             const SparseCooTensor& out,
+                             const SparseCooTensor& out_grad,
+                             const std::vector<int>& kernel_sizes,
+                             SparseCooTensor* x_grad) {
+  int kernel_size = kernel_sizes[0] * kernel_sizes[1] * kernel_sizes[2];
+  const int channels = static_cast<int>(x.dims()[4]);
+  int rulebook_len = static_cast<int>(rulebook.dims()[1]);
+  const IntT* rulebook_ptr = rulebook.data<IntT>();
+  std::vector<int> offsets(kernel_size + 1);
+  const int* counter_ptr = counter.data<int>();
+
+  funcs::sparse::PrefixSum(counter_ptr, &offsets[0], kernel_size);
+
+  const T* in_features_ptr = x.values().data<T>();
+  const T* out_features_ptr = out.values().data<T>();
+  const T* out_grad_ptr = out_grad.values().data<T>();
+  // TODO(zhangkaihuo): call phi::sparse::EmptyLike
+  DenseTensor x_grad_indices = EmptyLike<IntT>(dev_ctx, x.indices());
+  DenseTensor x_grad_values = EmptyLike<T>(dev_ctx, x.values());
+  x_grad->SetMember(x_grad_indices, x_grad_values, x.dims(), true);
+  T* x_grad_ptr = x_grad_values.data<T>();
+  memset(x_grad_ptr, 0, sizeof(T) * x_grad_values.numel());
+  phi::Copy<CPUContext>(
+      dev_ctx, x.indices(), dev_ctx.GetPlace(), false, &x_grad_indices);
+
+  funcs::MaxPoolGrad<T> grad_functor;
+  for (int i = 0; i < kernel_size; i++) {
+    for (int j = 0; j < counter_ptr[i]; j++) {
+      IntT in_i = rulebook_ptr[rulebook_len + offsets[i] + j];
+      IntT out_i = rulebook_ptr[rulebook_len * 2 + offsets[i] + j];
+      for (int c = 0; c < channels; c++) {
+        grad_functor.compute(in_features_ptr[in_i * channels + c],
+                             out_features_ptr[out_i * channels + c],
+                             out_grad_ptr[out_i * channels + c],
+                             1,
+                             &x_grad_ptr[in_i * channels + c]);
+      }
+    }
+  }
+}
+
+template <typename T, typename Context>
+void MaxPoolCooGradKernel(const Context& dev_ctx,
+                          const SparseCooTensor& x,
+                          const DenseTensor& rulebook,
+                          const DenseTensor& counter,
+                          const SparseCooTensor& out,
+                          const SparseCooTensor& out_grad,
+                          const std::vector<int>& kernel_sizes,
+                          SparseCooTensor* x_grad) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.indices().dtype(), "MaxPoolCooGradCPUKernel", ([&] {
+        MaxPoolCooGradCPUKernel<T, data_t>(
+            dev_ctx, x, rulebook, counter, out, out_grad, kernel_sizes, x_grad);
+      }));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(maxpool_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaxPoolCooGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/cpu/pool_kernel.cc

@@ -0,0 +1,134 @@
+/* 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/pool_kernel.h"
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_meta.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/funcs/pooling.h"
+#include "paddle/phi/kernels/funcs/sparse/convolution.h"
+#include "paddle/phi/kernels/sparse/cpu/conv.h"
+
+namespace phi::sparse {
+
+/**
+ * x: (N, D, H, W, C)
+ * kernel: (D, H, W, C, OC)
+ * out: (N, D, H, W, OC)
+ **/
+template <typename T, typename IntT = int>
+void MaxPoolCooCPUKernel(const CPUContext& dev_ctx,
+                         const SparseCooTensor& x,
+                         const std::vector<int>& kernel_sizes,
+                         const std::vector<int>& paddings,
+                         const std::vector<int>& dilations,
+                         const std::vector<int>& strides,
+                         SparseCooTensor* out,
+                         DenseTensor* rulebook,
+                         DenseTensor* counter) {
+  const auto& x_dims = x.dims();
+  int kernel_size = kernel_sizes[0] * kernel_sizes[1] * kernel_sizes[2];
+  const std::vector<int>& real_kernel_sizes = funcs::sparse::PoolResetKernel(
+      kernel_sizes, static_cast<int>(x_dims[4]), static_cast<int>(x_dims[4]));
+  DDim out_dims = {1, 1, 1, 1, 1};
+  funcs::sparse::GetOutShape(
+      x_dims, real_kernel_sizes, paddings, dilations, strides, &out_dims);
+  const int in_channels = real_kernel_sizes[3];
+
+  std::vector<int> counter_per_kernel(kernel_size, 0);
+
+  const T* in_features_ptr = x.values().data<T>();
+  // 1. product rule book
+  ProductRuleBook<T, CPUContext, IntT>(dev_ctx,
+                                       x,
+                                       real_kernel_sizes,
+                                       paddings,
+                                       dilations,
+                                       strides,
+                                       out_dims,
+                                       false,
+                                       rulebook,
+                                       counter_per_kernel.data());
+
+  UpdateRulebookAndOutIndex<T, CPUContext, IntT>(
+      dev_ctx, x, kernel_size, in_channels, out_dims, rulebook, out);
+
+  int rulebook_len = static_cast<int>(rulebook->dims()[1]);
+  const IntT* rulebook_ptr = rulebook->data<IntT>();
+
+  counter->Resize({kernel_size});
+  int* counter_ptr = dev_ctx.template HostAlloc<int>(counter);
+  memcpy(counter_ptr, counter_per_kernel.data(), kernel_size * sizeof(int));
+
+  std::vector<int> offsets(kernel_size + 1);
+  funcs::sparse::PrefixSum(counter_ptr, &offsets[0], kernel_size);
+  std::vector<bool> out_flags(out->nnz(), false);
+
+  // 2. max pool
+  T* out_features_ptr = out->mutable_values()->data<T>();
+  funcs::MaxPool<T> max_pool_functor;
+  for (int i = 0; i < kernel_size; i++) {
+    for (int j = 0; j < counter_ptr[i]; j++) {
+      IntT in_i = rulebook_ptr[rulebook_len + offsets[i] + j];
+      IntT out_i = rulebook_ptr[rulebook_len * 2 + offsets[i] + j];
+      if (!out_flags[out_i]) {
+        out_flags[out_i] = true;
+        memcpy(&out_features_ptr[out_i * in_channels],
+               &in_features_ptr[in_i * in_channels],
+               in_channels * sizeof(T));
+      } else {
+        for (int c = 0; c < in_channels; c++) {
+          max_pool_functor.compute(in_features_ptr[in_i * in_channels + c],
+                                   &out_features_ptr[out_i * in_channels + c]);
+        }
+      }
+    }
+  }
+}
+
+template <typename T, typename Context>
+void MaxPoolCooKernel(const Context& dev_ctx,
+                      const SparseCooTensor& x,
+                      const std::vector<int>& kernel_sizes,
+                      const std::vector<int>& paddings,
+                      const std::vector<int>& dilations,
+                      const std::vector<int>& strides,
+                      SparseCooTensor* out,
+                      DenseTensor* rulebook,
+                      DenseTensor* counter) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.indices().dtype(), "MaxPoolCooCPUKernel", ([&] {
+        MaxPoolCooCPUKernel<T, data_t>(dev_ctx,
+                                       x,
+                                       kernel_sizes,
+                                       paddings,
+                                       dilations,
+                                       strides,
+                                       out,
+                                       rulebook,
+                                       counter);
+      }));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(maxpool_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaxPoolCooKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/cpu/reshape_grad_kernel.cc

@@ -0,0 +1,71 @@
+// 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/unary_grad_kernel.h"
+#include "paddle/phi/kernels/sparse/unary_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/impl/unary_grad_kernel_impl.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void ReshapeCooGradKernel(const Context& dev_ctx,
+                          const SparseCooTensor& x,
+                          const SparseCooTensor& dout,
+                          SparseCooTensor* dx) {
+  EmptyLikeCooKernel<T, Context>(dev_ctx, x, dx);
+  phi::IntArray x_shape(vectorize(x.dims()));
+  ReshapeCooKernel<T, Context>(dev_ctx, dout, x_shape, dx);
+}
+
+template <typename T, typename Context>
+void ReshapeCsrGradKernel(const Context& dev_ctx,
+                          const SparseCsrTensor& x,
+                          const SparseCsrTensor& dout,
+                          SparseCsrTensor* dx) {
+  EmptyLikeCsrKernel<T, Context>(dev_ctx, x, dx);
+  phi::IntArray x_shape(vectorize(x.dims()));
+  ReshapeCsrKernel<T, Context>(dev_ctx, dout, x_shape, dx);
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(reshape_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ReshapeCooGradKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(reshape_csr_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ReshapeCsrGradKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/cpu/reshape_kernel.cc

@@ -0,0 +1,126 @@
+// 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/unary_kernel.h"
+
+#include "paddle/common/ddim.h"
+#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/impl/unary_grad_kernel_impl.h"
+#include "paddle/phi/kernels/sparse/impl/unary_kernel_impl.h"
+
+namespace phi::sparse {
+
+template <typename T, typename IntT, typename Context>
+void ReshapeCooCPUKernel(const Context& dev_ctx,
+                         const SparseCooTensor& x,
+                         const phi::IntArray& shape,
+                         SparseCooTensor* out) {
+  // TODO(OccupyMars2025): Currently, reshape is only applicable to sparse dims
+  int64_t x_nnz = x.nnz();
+
+  // Use DDim::reshape to handle -1 and 0 in the argument "shape"
+  std::vector<int> new_shape(shape.GetData().begin(), shape.GetData().end());
+  DDim out_dims = x.dims().reshape(new_shape);
+  // get sparse part dimensions of x and out
+  std::vector<int64_t> x_sparse_part_dims;
+  std::vector<int64_t> out_sparse_part_dims;
+  for (int i = 0; i < x.sparse_dim(); ++i) {
+    x_sparse_part_dims.push_back(x.dims()[i]);
+  }
+  for (int i = 0; i < out_dims.size() - x.dense_dim(); ++i) {
+    out_sparse_part_dims.push_back(out_dims[i]);
+  }
+  DenseTensor out_indices = Empty<IntT, Context>(
+      dev_ctx, {static_cast<int64_t>(out_sparse_part_dims.size()), x_nnz});
+  DenseTensor out_values(x.values());
+  out->SetMember(out_indices, out_values, out_dims, x.coalesced());
+
+  // compute values of indices
+  const DenseTensor& x_indices = x.indices();
+  const auto* x_indices_data = x_indices.data<IntT>();
+  auto* out_indices_data = out_indices.data<IntT>();
+
+  const DDim& x_sparse_part_strides =
+      common::stride(make_ddim(x_sparse_part_dims));
+  const DDim& out_sparse_part_strides =
+      common::stride(make_ddim(out_sparse_part_dims));
+  int64_t location = 0;
+  for (int64_t j = 0; j < x_nnz; ++j) {
+    location = 0;
+    for (int i = 0; i < x.sparse_dim(); ++i) {
+      location += x_indices_data[i * x_nnz + j] * x_sparse_part_strides[i];
+    }
+    for (int i = 0; i < static_cast<int>(out_sparse_part_dims.size()); ++i) {
+      out_indices_data[i * x_nnz + j] = location / out_sparse_part_strides[i];
+      location %= out_sparse_part_strides[i];
+    }
+  }
+}
+
+template <typename T, typename Context>
+void ReshapeCooKernel(const Context& dev_ctx,
+                      const SparseCooTensor& x,
+                      const phi::IntArray& shape,
+                      SparseCooTensor* out) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.indices().dtype(), "ReshapeCooCPUKernel", ([&] {
+        ReshapeCooCPUKernel<T, data_t, Context>(dev_ctx, x, shape, out);
+      }));
+}
+
+template <typename T, typename Context>
+void ReshapeCsrKernel(const Context& dev_ctx,
+                      const SparseCsrTensor& x,
+                      const phi::IntArray& shape,
+                      SparseCsrTensor* out) {
+  // transform csr format to coo format, and then use coo kernel
+  const SparseCooTensor x_coo = CsrToCoo<T, Context>(dev_ctx, x);
+  SparseCooTensor out_coo;
+  ReshapeCooKernel<T, Context>(dev_ctx, x_coo, shape, &out_coo);
+  CooToCsrKernel<T, Context>(dev_ctx, out_coo, out);
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(reshape_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ReshapeCooKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(reshape_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ReshapeCsrKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/cpu/slice_grad_kernel.cc

@@ -0,0 +1,274 @@
+// Copyright (c) 2023 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/unary_grad_kernel.h"
+#include "paddle/phi/kernels/sparse/unary_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/slice_utils.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void SliceCooGradCompute(const Context& dev_ctx,
+                         const SparseCooTensor& x,
+                         const SparseCooTensor& out_grad,
+                         const std::vector<int64_t>& axes,
+                         const std::vector<int64_t>& starts,
+                         const std::vector<int64_t>& ends,
+                         SparseCooTensor* x_grad) {
+  // set x_grad
+  const int64_t out_grad_nnz = out_grad.nnz();
+  auto sparse_dim = static_cast<int64_t>(out_grad.sparse_dim());
+  DenseTensor dx_indices =
+      Empty<int64_t, Context>(dev_ctx, {sparse_dim, out_grad_nnz});
+  DenseTensor dx_values = Empty<T, Context>(dev_ctx, {out_grad_nnz});
+  auto* dx_indices_data = dx_indices.data<int64_t>();
+  auto* dx_values_data = dx_values.data<T>();
+
+  const auto* out_grad_indices_data = out_grad.indices().data<int64_t>();
+  const auto* out_grad_values_data = out_grad.values().data<T>();
+
+  for (int64_t j = 0; j < out_grad_nnz; ++j) {
+    // set indices
+    for (int64_t i = 0; i < sparse_dim; ++i) {
+      dx_indices_data[i * out_grad_nnz + j] =
+          out_grad_indices_data[i * out_grad_nnz + j];
+    }
+    for (size_t ii = 0; ii < axes.size(); ++ii) {
+      int64_t i = axes[ii];
+      dx_indices_data[i * out_grad_nnz + j] += starts[ii];
+    }
+    // set value
+    dx_values_data[j] = out_grad_values_data[j];
+  }
+
+  x_grad->SetMember(dx_indices, dx_values, x.dims(), x.coalesced());
+}
+
+template <typename T, typename Context>
+void SliceCooGradKernel(const Context& dev_ctx,
+                        const SparseCooTensor& x,
+                        const SparseCooTensor& out_grad,
+                        const phi::IntArray& axes,
+                        const phi::IntArray& starts,
+                        const phi::IntArray& ends,
+                        SparseCooTensor* x_grad) {
+  const DDim& x_dims = x.dims();
+  std::vector<int64_t> axes_vec = axes.GetData();
+  std::vector<int64_t> starts_vec = starts.GetData();
+  std::vector<int64_t> ends_vec = ends.GetData();
+
+  // update starts and ends
+  funcs::CheckAndUpdateSparseSliceAttrs<int64_t>(
+      x_dims, &axes_vec, &starts_vec, &ends_vec);
+
+  SliceCooGradCompute<T, Context>(
+      dev_ctx, x, out_grad, axes_vec, starts_vec, ends_vec, x_grad);
+}
+
+template <typename T>
+void GetCsrInputGradCrows(const int64_t* out_grad_crows_data,
+                          const int64_t out_grad_n_rows,
+                          const int64_t x_n_rows,
+                          const int64_t rows_start,
+                          int64_t* dx_crows_data,
+                          const int64_t out_grad_crows_offset = 0,
+                          const int64_t dx_crows_offset = 0) {
+  for (int64_t i = 0; i < x_n_rows + 1; ++i) {
+    int64_t idx = i + dx_crows_offset;
+    if (i < rows_start) {
+      dx_crows_data[idx] = 0;
+    } else if (i < rows_start + out_grad_n_rows + 1) {
+      int64_t out_grad_idx = out_grad_crows_offset + (i - rows_start);
+      dx_crows_data[idx] = out_grad_crows_data[out_grad_idx];
+    } else {
+      int64_t out_grad_idx = out_grad_crows_offset + out_grad_n_rows;
+      dx_crows_data[idx] = out_grad_crows_data[out_grad_idx];
+    }
+  }
+}
+
+template <typename T, typename Context>
+void SliceCsrGrad2D(const Context& dev_ctx,
+                    const SparseCsrTensor& x,
+                    const SparseCsrTensor& out_grad,
+                    const std::vector<int64_t>& axes,
+                    const std::vector<int64_t>& starts,
+                    const std::vector<int64_t>& ends,
+                    SparseCsrTensor* x_grad) {
+  const int64_t out_grad_nnz = out_grad.nnz();
+  const int64_t n_rows = x.dims()[0];
+  const auto* out_grad_crows_data = out_grad.crows().data<int64_t>();
+  const auto* out_grad_cols_data = out_grad.cols().data<int64_t>();
+  const auto* out_grad_values_data = out_grad.values().data<T>();
+
+  DenseTensor dx_crows = Empty<int64_t>(dev_ctx, {n_rows + 1});
+  DenseTensor dx_cols = Empty<int64_t>(dev_ctx, {out_grad_nnz});
+  DenseTensor dx_values = Empty<T, Context>(dev_ctx, {out_grad_nnz});
+  auto* dx_crows_data = dx_crows.data<int64_t>();
+  auto* dx_cols_data = dx_cols.data<int64_t>();
+  auto* dx_values_data = dx_values.data<T>();
+
+  // set cols
+  for (int64_t i = 0; i < out_grad_nnz; ++i) {
+    dx_cols_data[i] = out_grad_cols_data[i] + starts[1];
+  }
+  // set values
+  for (int64_t i = 0; i < out_grad_nnz; ++i) {
+    dx_values_data[i] = out_grad_values_data[i];
+  }
+  // set crows
+  const int64_t out_grad_n_rows = out_grad.dims()[0];
+  GetCsrInputGradCrows<T>(out_grad_crows_data,
+                          out_grad_n_rows,
+                          n_rows,
+                          starts[0],
+                          dx_crows_data,
+                          0,
+                          0);
+  x_grad->SetMember(dx_crows, dx_cols, dx_values, x.dims());
+}
+
+template <typename T, typename Context>
+void SliceCsrGrad3D(const Context& dev_ctx,
+                    const SparseCsrTensor& x,
+                    const SparseCsrTensor& out_grad,
+                    const std::vector<int64_t>& axes,
+                    const std::vector<int64_t>& starts,
+                    const std::vector<int64_t>& ends,
+                    SparseCsrTensor* x_grad) {
+  const int64_t dim0 = x.dims()[0], n_rows = x.dims()[1];
+  const int64_t out_grad_nnz = out_grad.nnz();
+  const auto* out_grad_crows_data = out_grad.crows().data<int64_t>();
+  const auto* out_grad_cols_data = out_grad.cols().data<int64_t>();
+  const auto* out_grad_values_data = out_grad.values().data<T>();
+
+  DenseTensor dx_crows = Empty<int64_t>(dev_ctx, {dim0 * (n_rows + 1)});
+  DenseTensor dx_cols = Empty<int64_t>(dev_ctx, {out_grad_nnz});
+  DenseTensor dx_values = Empty<T, Context>(dev_ctx, {out_grad_nnz});
+  auto* dx_crows_data = dx_crows.data<int64_t>();
+  auto* dx_cols_data = dx_cols.data<int64_t>();
+  auto* dx_values_data = dx_values.data<T>();
+
+  // set cols
+  for (int64_t i = 0; i < out_grad_nnz; ++i) {
+    dx_cols_data[i] = out_grad_cols_data[i] + starts[2];
+  }
+  // set values
+  for (int64_t i = 0; i < out_grad_nnz; ++i) {
+    dx_values_data[i] = out_grad_values_data[i];
+  }
+  // set crows
+  int64_t out_grad_n_rows = out_grad.dims()[1];
+  for (int64_t i = 0; i < dim0; ++i) {
+    if (i < starts[0] || i >= ends[0]) {
+      for (int64_t j = 0; j < n_rows + 1; ++j) {
+        dx_crows_data[i * (n_rows + 1) + j] = 0;
+      }
+    } else {
+      int64_t out_grad_crows_offset = (i - starts[0]) * (out_grad_n_rows + 1);
+      int64_t dx_crows_offset = i * (n_rows + 1);
+      GetCsrInputGradCrows<T>(out_grad_crows_data,
+                              out_grad_n_rows,
+                              n_rows,
+                              starts[1],
+                              dx_crows_data,
+                              out_grad_crows_offset,
+                              dx_crows_offset);
+    }
+  }
+  x_grad->SetMember(dx_crows, dx_cols, dx_values, x.dims());
+}
+
+template <typename T, typename Context>
+void SliceCsrGradCompute(const Context& dev_ctx,
+                         const SparseCsrTensor& x,
+                         const SparseCsrTensor& out_grad,
+                         const std::vector<int64_t>& axes,
+                         const std::vector<int64_t>& starts,
+                         const std::vector<int64_t>& ends,
+                         SparseCsrTensor* x_grad) {
+  const DDim& x_dims = x.dims();
+
+  // Construct new axes, starts, and ends
+  std::vector<int64_t> new_axes(3), new_starts(3), new_ends(3);
+  funcs::ConstructNewSliceAttrs(
+      x_dims, axes, starts, ends, &new_axes, &new_starts, &new_ends);
+
+  const int64_t sparse_dim = x_dims.size();
+  if (sparse_dim == 2) {
+    SliceCsrGrad2D<T, Context>(
+        dev_ctx, x, out_grad, new_axes, new_starts, new_ends, x_grad);
+  } else if (sparse_dim == 3) {
+    SliceCsrGrad3D<T, Context>(
+        dev_ctx, x, out_grad, new_axes, new_starts, new_ends, x_grad);
+  } else {
+    // throw exception
+    common::errors::InvalidArgument(
+        "Slice grad for Sparse CSR Tensor only support 2-D or 3-D, but got "
+        "%d-D.",
+        x_dims.size());
+  }
+}
+
+template <typename T, typename Context>
+void SliceCsrGradKernel(const Context& dev_ctx,
+                        const SparseCsrTensor& x,
+                        const SparseCsrTensor& out_grad,
+                        const phi::IntArray& axes,
+                        const phi::IntArray& starts,
+                        const phi::IntArray& ends,
+                        SparseCsrTensor* x_grad) {
+  const DDim& x_dims = x.dims();
+  std::vector<int64_t> axes_vec = axes.GetData();
+  std::vector<int64_t> starts_vec = starts.GetData();
+  std::vector<int64_t> ends_vec = ends.GetData();
+
+  // Update starts and ends
+  funcs::CheckAndUpdateSparseSliceAttrs<int64_t>(
+      x_dims, &axes_vec, &starts_vec, &ends_vec);
+
+  SliceCsrGradCompute<T, Context>(
+      dev_ctx, x, out_grad, axes_vec, starts_vec, ends_vec, x_grad);
+}
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(slice_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SliceCooGradKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(slice_csr_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SliceCsrGradKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/cpu/slice_kernel.cc

@@ -0,0 +1,330 @@
+// Copyright (c) 2023 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/unary_kernel.h"
+
+#include "paddle/common/ddim.h"
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/slice_utils.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void SliceCooCompute(const Context& dev_ctx,
+                     const SparseCooTensor& x,
+                     const std::vector<int64_t>& axes,
+                     const std::vector<int64_t>& starts,
+                     const std::vector<int64_t>& ends,
+                     SparseCooTensor* out) {
+  const DDim& x_dims = x.dims();
+
+  // Step1: Infer output dims
+  auto out_dims = funcs::GetSliceDims<int64_t>(
+      x_dims, axes, starts, ends, nullptr, nullptr);
+
+  // Step2: Get out_nnz (the number of non-zero elements in output)
+  const int64_t x_nnz = x.nnz();
+  int64_t out_nnz = 0;
+  const auto* x_indices_data = x.indices().data<int64_t>();
+  for (int64_t j = 0; j < x_nnz; ++j) {
+    bool hit = true;
+    for (size_t ii = 0; ii < axes.size(); ++ii) {
+      auto item = x_indices_data[axes[ii] * x_nnz + j];
+      if (!(starts[ii] <= item && item < ends[ii])) {
+        hit = false;
+        break;
+      }
+    }
+    if (!hit) continue;
+    out_nnz++;
+  }
+
+  // Step3: Get the values and indices of output
+  auto sparse_dim = static_cast<int64_t>(x.sparse_dim());
+  DenseTensor out_indices =
+      Empty<int64_t, Context>(dev_ctx, {sparse_dim, out_nnz});
+  DenseTensor out_values = Empty<T, Context>(dev_ctx, {out_nnz});
+
+  auto* out_indices_data = out_indices.data<int64_t>();
+  auto* out_values_data = out_values.data<T>();
+  const auto* x_values_data = x.values().data<T>();
+  int64_t index = 0;
+  for (int64_t j = 0; j < x_nnz && index < out_nnz; ++j) {
+    bool hit = true;
+    for (size_t ii = 0; ii < axes.size(); ++ii) {
+      auto item = x_indices_data[axes[ii] * x_nnz + j];
+      if (!(starts[ii] <= item && item < ends[ii])) {
+        hit = false;
+        break;
+      }
+    }
+    if (!hit) continue;
+    // set value
+    out_values_data[index] = x_values_data[j];
+    // set coordinate
+    for (int64_t i = 0; i < sparse_dim; ++i) {
+      out_indices_data[i * out_nnz + index] = x_indices_data[i * x_nnz + j];
+    }
+    for (size_t ii = 0; ii < axes.size(); ++ii) {
+      auto i = axes[ii];
+      out_indices_data[i * out_nnz + index] -= starts[ii];
+    }
+    index++;
+  }
+  out->SetMember(out_indices, out_values, out_dims, x.coalesced());
+}
+
+template <typename T, typename Context>
+void SliceCooKernel(const Context& dev_ctx,
+                    const SparseCooTensor& x,
+                    const phi::IntArray& axes,
+                    const phi::IntArray& starts,
+                    const phi::IntArray& ends,
+                    SparseCooTensor* out) {
+  const DDim& x_dims = x.dims();
+  std::vector<int64_t> axes_vec = axes.GetData();
+  std::vector<int64_t> starts_vec = starts.GetData();
+  std::vector<int64_t> ends_vec = ends.GetData();
+
+  // Check and update attr
+  funcs::CheckAndUpdateSparseSliceAttrs<int64_t>(
+      x_dims, &axes_vec, &starts_vec, &ends_vec);
+
+  SliceCooCompute<T, Context>(dev_ctx, x, axes_vec, starts_vec, ends_vec, out);
+}
+
+int64_t GetCsrNonZeroNumber(const SparseCsrTensor& x,
+                            const int64_t x_crows_start,
+                            const int64_t x_crows_end,
+                            const int64_t min_col,
+                            const int64_t max_col,
+                            const int64_t x_cols_offset = 0) {
+  const auto* x_crows_data = x.crows().data<int64_t>();
+  const auto* x_cols_data = x.cols().data<int64_t>();
+  int64_t out_nnz = 0;
+  for (int64_t i = x_crows_start; i < x_crows_end; ++i) {
+    int64_t st = x_crows_data[i] + x_cols_offset;
+    int64_t ed = x_crows_data[i + 1] + x_cols_offset;
+    for (int64_t jj = st; jj < ed; ++jj) {
+      if (x_cols_data[jj] >= min_col && x_cols_data[jj] < max_col) {
+        out_nnz++;
+      }
+    }
+  }
+  return out_nnz;
+}
+
+template <typename T>
+void GetCsrSubMatrix(const SparseCsrTensor& x,
+                     const int64_t x_crows_start,
+                     const int64_t x_crows_end,
+                     const int64_t min_col,
+                     const int64_t max_col,
+                     DenseTensor* out_crows,
+                     DenseTensor* out_cols,
+                     DenseTensor* out_values,
+                     const int64_t x_cols_offset = 0,
+                     const int64_t out_crows_offset = 0,
+                     const int64_t out_cols_offset = 0) {
+  const auto* x_crows_data = x.crows().data<int64_t>();
+  const auto* x_cols_data = x.cols().data<int64_t>();
+  const auto* x_values_data = x.values().data<T>();
+
+  auto* out_crows_data = out_crows->data<int64_t>();
+  auto* out_cols_data = out_cols->data<int64_t>();
+  auto* out_values_data = out_values->data<T>();
+  out_crows_data[out_crows_offset] = 0;
+  int64_t index = 0, out_n_rows = x_crows_end - x_crows_start;
+  for (int i = 0; i < out_n_rows; ++i) {
+    int64_t st = x_crows_data[x_crows_start + i] + x_cols_offset;
+    int64_t ed = x_crows_data[x_crows_start + i + 1] + x_cols_offset;
+    for (int64_t jj = st; jj < ed; ++jj) {
+      if (x_cols_data[jj] >= min_col && x_cols_data[jj] < max_col) {
+        out_cols_data[out_cols_offset + index] = x_cols_data[jj] - min_col;
+        out_values_data[out_cols_offset + index] = x_values_data[jj];
+        index++;
+      }
+    }
+    out_crows_data[out_crows_offset + i + 1] = index;
+  }
+}
+
+template <typename T, typename Context>
+void SliceCsrTensor2D(const Context& dev_ctx,
+                      const SparseCsrTensor& x,
+                      const std::vector<int64_t>& axes,
+                      const std::vector<int64_t>& starts,
+                      const std::vector<int64_t>& ends,
+                      const DDim& out_dims,
+                      SparseCsrTensor* out) {
+  // Step1: Get nnz of out
+  int64_t out_nnz =
+      GetCsrNonZeroNumber(x, starts[0], ends[0], starts[1], ends[1], 0);
+  // Step2: Set out
+  int64_t out_n_rows = ends[0] - starts[0];
+  DenseTensor out_crows = Empty<int64_t, Context>(dev_ctx, {out_n_rows + 1});
+  DenseTensor out_cols = Empty<int64_t, Context>(dev_ctx, {out_nnz});
+  DenseTensor out_values = Empty<T, Context>(dev_ctx, {out_nnz});
+  GetCsrSubMatrix<T>(x,
+                     starts[0],
+                     ends[0],
+                     starts[1],
+                     ends[1],
+                     &out_crows,
+                     &out_cols,
+                     &out_values,
+                     0,
+                     0,
+                     0);
+  out->SetMember(out_crows, out_cols, out_values, out_dims);
+}
+
+template <typename T, typename Context>
+void SliceCsrTensor3D(const Context& dev_ctx,
+                      const SparseCsrTensor& x,
+                      const std::vector<int64_t>& axes,
+                      const std::vector<int64_t>& starts,
+                      const std::vector<int64_t>& ends,
+                      const DDim& out_dims,
+                      SparseCsrTensor* out) {
+  const auto* x_crows_data = x.crows().data<int64_t>();
+  // Step1: Get nnz of out
+  const int64_t x_dim0 = x.dims()[0], x_n_rows = x.dims()[1];
+  int64_t x_cols_offset = 0, out_nnz = 0;
+  // all_nnzs stores the nnz along with out_dim0, which will be used to set out.
+  std::vector<int64_t> all_nnzs(ends[0] - starts[0]);
+  for (int64_t i = 0; i < x_dim0; ++i) {
+    if (i >= starts[0] && i < ends[0]) {  // slice dim 0
+      int64_t x_crows_st = i * (x_n_rows + 1) + starts[1];
+      int64_t x_crows_ed = i * (x_n_rows + 1) + ends[1];
+      int64_t nnz = GetCsrNonZeroNumber(
+          x, x_crows_st, x_crows_ed, starts[2], ends[2], x_cols_offset);
+      out_nnz += nnz;
+      all_nnzs[i - starts[0]] = nnz;
+    }
+    // get the start index in non_zero_cols_
+    x_cols_offset += x_crows_data[(i + 1) * (x_n_rows + 1) - 1];
+  }
+
+  // Step2: Set out
+  const int64_t out_dim0 = out_dims[0], out_n_rows = out_dims[1];
+  DenseTensor out_crows =
+      Empty<int64_t, Context>(dev_ctx, {out_dim0 * (out_n_rows + 1)});
+  DenseTensor out_cols = Empty<int64_t, Context>(dev_ctx, {out_nnz});
+  DenseTensor out_values = Empty<T, Context>(dev_ctx, {out_nnz});
+
+  x_cols_offset = 0;
+  int64_t out_crows_offset = 0, out_cols_offset = 0;
+  for (int64_t i = 0; i < x_dim0; ++i) {
+    if (i >= starts[0] && i < ends[0]) {  // slice dim 0
+      int64_t x_crows_start = i * (x_n_rows + 1) + starts[1];
+      int64_t x_crows_end = i * (x_n_rows + 1) + ends[1];
+      GetCsrSubMatrix<T>(x,
+                         x_crows_start,
+                         x_crows_end,
+                         starts[2],
+                         ends[2],
+                         &out_crows,
+                         &out_cols,
+                         &out_values,
+                         x_cols_offset,
+                         out_crows_offset,
+                         out_cols_offset);
+      out_crows_offset += (out_n_rows + 1);
+      out_cols_offset += all_nnzs[i - starts[0]];
+    }
+    x_cols_offset += x_crows_data[(i + 1) * (x_n_rows + 1) - 1];
+  }
+  out->SetMember(out_crows, out_cols, out_values, out_dims);
+}
+
+template <typename T, typename Context>
+void SliceCsrCompute(const Context& dev_ctx,
+                     const SparseCsrTensor& x,
+                     const std::vector<int64_t>& axes,
+                     const std::vector<int64_t>& starts,
+                     const std::vector<int64_t>& ends,
+                     SparseCsrTensor* out) {
+  const DDim& x_dims = x.dims();
+
+  // Step1: Infer output dims
+  auto out_dims = funcs::GetSliceDims<int64_t>(
+      x_dims, axes, starts, ends, nullptr, nullptr);
+
+  // Step2: Construct new axes, starts and ends.
+  std::vector<int64_t> new_axes(3), new_starts(3), new_ends(3);
+  funcs::ConstructNewSliceAttrs(
+      x_dims, axes, starts, ends, &new_axes, &new_starts, &new_ends);
+
+  // Step3: Slice csr tensor according to its dimension
+  if (x_dims.size() == 2) {
+    SliceCsrTensor2D<T, Context>(
+        dev_ctx, x, new_axes, new_starts, new_ends, out_dims, out);
+  } else if (x_dims.size() == 3) {
+    SliceCsrTensor3D<T, Context>(
+        dev_ctx, x, new_axes, new_starts, new_ends, out_dims, out);
+  } else {
+    // throw exception
+    common::errors::InvalidArgument(
+        "Slice for Sparse CSR Tensor only support 2-D or 3-D, but got %d-D.",
+        x_dims.size());
+  }
+}
+
+template <typename T, typename Context>
+void SliceCsrKernel(const Context& dev_ctx,
+                    const SparseCsrTensor& x,
+                    const phi::IntArray& axes,
+                    const phi::IntArray& starts,
+                    const phi::IntArray& ends,
+                    SparseCsrTensor* out) {
+  const DDim& x_dims = x.dims();
+  std::vector<int64_t> axes_vec = axes.GetData();
+  std::vector<int64_t> starts_vec = starts.GetData();
+  std::vector<int64_t> ends_vec = ends.GetData();
+
+  // Check and update attr
+  funcs::CheckAndUpdateSparseSliceAttrs<int64_t>(
+      x_dims, &axes_vec, &starts_vec, &ends_vec);
+  SliceCsrCompute<T, Context>(dev_ctx, x, axes_vec, starts_vec, ends_vec, out);
+}
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(slice_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SliceCooKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(slice_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SliceCsrKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/cpu/softmax_grad_kernel.cc

@@ -0,0 +1,223 @@
+/* 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/softmax_grad_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/backends/cpu/cpu_info.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
+#include "paddle/phi/kernels/funcs/cpu_vec.h"
+#include "paddle/phi/kernels/funcs/sparse/softmax.h"
+#include "paddle/phi/kernels/softmax_grad_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void SoftmaxCsrGradKernel(const Context& dev_ctx,
+                          const SparseCsrTensor& out,
+                          const SparseCsrTensor& dout,
+                          int axis,
+                          SparseCsrTensor* dx) {
+  PADDLE_ENFORCE_EQ(axis,
+                    -1,
+                    common::errors::Unimplemented(
+                        "SparseCsrTensor only support axis=-1 for softmax, "
+                        "which is faster when reading data by row (axis=-1)"));
+  EmptyLikeCsrKernel<T, Context>(dev_ctx, dout, dx);
+  auto out_dim = out.dims();
+  auto out_rank = out_dim.size();
+
+  int batch_size = 1;
+  int row_number = 1;
+  for (int i = 0; i < out_rank - 1; ++i) {
+    if (i < out_rank - 2) {
+      batch_size *= static_cast<int>(out_dim[i]);
+    } else if (i == out_rank - 2) {
+      row_number = static_cast<int>(out_dim[i]);
+    }
+  }
+
+  const DenseTensor& out_crows = out.non_zero_crows();
+  const DenseTensor& out_values = out.non_zero_elements();
+  const DenseTensor& dout_values = dout.non_zero_elements();
+  DenseTensor* dx_values = dx->mutable_non_zero_elements();
+
+  int row_nnz = 0;
+  const T* out_data = out_values.data<T>();
+  const T* dout_data = dout_values.data<T>();
+  T* dx_data = dx_values->data<T>();
+
+  // dx = (dout - sum(dout * out)) * out
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      out.non_zero_crows().dtype(), "SoftmaxCsrGradKernel", ([&] {
+        const data_t* out_crows_data = out_crows.data<data_t>();
+        for (int i = 0; i < batch_size; ++i) {
+          for (int j = 0; j < row_number; ++j) {
+            int crow_idx = i * (row_number + 1) + j;
+            row_nnz = static_cast<int>(out_crows_data[crow_idx + 1] -
+                                       out_crows_data[crow_idx]);
+
+            T sum = 0;
+            funcs::vec_mul_reduce<T, backends::cpu::avx>(
+                row_nnz, dout_data, out_data, &sum);
+            funcs::vec_add_bias<T, backends::cpu::avx>(
+                row_nnz, static_cast<T>(-1) * sum, dout_data, dx_data);
+            funcs::vec_mul<T, backends::cpu::avx>(
+                row_nnz, dx_data, out_data, dx_data);
+
+            out_data = out_data + row_nnz;
+            dout_data = dout_data + row_nnz;
+            dx_data = dx_data + row_nnz;
+          }
+        }
+      }));
+}
+
+template <typename T, typename IntT, typename Context>
+void SoftmaxCooGradCPUKernel(const Context& dev_ctx,
+                             const SparseCooTensor& out,
+                             const SparseCooTensor& dout,
+                             int axis,
+                             SparseCooTensor* dx) {
+  auto out_indices = out.indices();
+  auto out_values = out.values();
+  const auto out_dims = out.dims();
+  auto sparse_dim = out.sparse_dim();
+  auto sizes = vectorize<IntT>(out_dims);
+  auto grad_indices = dout.indices();
+  auto grad_values = dout.values();
+  auto grad_nnz = dout.nnz();
+
+  *(dx->mutable_indices()) = out_indices;
+  DenseTensor* values = dx->mutable_values();
+  values->Resize(out_dims);
+  values->set_meta(out_values.meta());
+  dev_ctx.template Alloc<T>(values);
+
+  auto out_offsets = funcs::sparse::GetOffsets(out_indices, sizes, -1);
+  auto grad_offsets = funcs::sparse::GetOffsets(grad_indices, sizes, -1);
+
+  int dim = axis < 0 ? out_dims.size() + axis : axis;
+
+  if (dim >= sparse_dim) {
+    bool is_same_offset = out_offsets == grad_offsets;
+    PADDLE_ENFORCE_EQ(
+        is_same_offset,
+        true,
+        common::errors::Unimplemented(
+            "SparseCooTensor only support same offsets for softmax."));
+
+    SoftmaxGradKernel<T, Context>(
+        dev_ctx, out_values, grad_values, dim - sparse_dim + 1, values);
+    return;
+  }
+
+  auto nnz = out.nnz();
+  IntT nvalues = std::accumulate(sizes.begin() + sparse_dim,
+                                 sizes.end(),
+                                 static_cast<IntT>(1),
+                                 std::multiplies<>());
+
+  DenseTensor values_2(*values);
+  values_2.Resize({nnz, nvalues});
+
+  DenseTensor out_values_2(out_values);
+  out_values_2.Resize({nnz, nvalues});
+
+  DenseTensor grad_values_2(grad_values);
+  grad_values_2.Resize({nnz, nvalues});
+  std::map<IntT, std::vector<IntT>> pools;
+  funcs::sparse::GetPoolsSoftmax(out_indices, sizes, dim, &pools);
+
+  for (size_t p = 0; p < pools.size(); p++) {
+    auto pool_indices = pools[p];
+    if (pool_indices.empty()) continue;
+
+    std::vector<T> tmp_row(nvalues, 0);
+
+    /* Compute tmp = - sum_j output_j * grad_j */
+    for (IntT i : pool_indices) {
+      auto out_values_row = out_values_2.data<T>() + i * nvalues;
+      auto low = std::lower_bound(
+          grad_offsets.begin(), grad_offsets.end(), out_offsets[i]);
+      auto j = low - grad_offsets.begin();
+
+      if (j < grad_nnz && (out_offsets[i] == grad_offsets[j])) {
+        auto grad_values_row = grad_values_2.data<T>() + j * nvalues;
+        for (IntT k = 0; k < nvalues; k++) {
+          tmp_row[k] -= (*(out_values_row + k)) * (*(grad_values_row + k));
+        }
+      }
+    }
+
+    /* Compute grad_input = output * (grad + tmp)*/
+    for (IntT i : pool_indices) {
+      auto out_values_row = out_values_2.data<T>() + i * nvalues;
+      auto values_row = values_2.data<T>() + i * nvalues;
+      auto low = std::lower_bound(
+          grad_offsets.begin(), grad_offsets.end(), out_offsets[i]);
+      auto j = low - grad_offsets.begin();
+
+      if (j < grad_nnz && (out_offsets[i] == grad_offsets[j])) {
+        auto grad_values_row = grad_values_2.data<T>() + j * nvalues;
+        for (IntT k = 0; k < nvalues; k++) {
+          *(values_row + k) =
+              (*(out_values_row + k)) * ((*(grad_values_row + k)) + tmp_row[k]);
+        }
+      } else {
+        for (IntT k = 0; k < nvalues; k++) {
+          *(values_row + k) = (*out_values_row + k) * (tmp_row[k]);
+        }
+      }
+    }
+  }
+}
+
+template <typename T, typename Context>
+void SoftmaxCooGradKernel(const Context& dev_ctx,
+                          const SparseCooTensor& out,
+                          const SparseCooTensor& dout,
+                          int axis,
+                          SparseCooTensor* dx) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      out.indices().dtype(), "SoftmaxCooGradCPUKernel", ([&] {
+        SoftmaxCooGradCPUKernel<T, data_t, Context>(
+            dev_ctx, out, dout, axis, dx);
+      }));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(softmax_csr_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SoftmaxCsrGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(softmax_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SoftmaxCooGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/cpu/softmax_kernel.cc

@@ -0,0 +1,192 @@
+/* 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/softmax_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/backends/cpu/cpu_info.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/cpu_vec.h"
+#include "paddle/phi/kernels/funcs/sparse/softmax.h"
+#include "paddle/phi/kernels/softmax_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void SoftmaxCsrKernel(const Context& dev_ctx,
+                      const SparseCsrTensor& x,
+                      int axis,
+                      SparseCsrTensor* out) {
+  PADDLE_ENFORCE_EQ(axis,
+                    -1,
+                    common::errors::Unimplemented(
+                        "SparseCsrTensor only support axis=-1 for softmax, "
+                        "which is faster when reading data by row (axis=-1)"));
+  EmptyLikeCsrKernel<T, Context>(dev_ctx, x, out);
+  auto x_dim = x.dims();
+  auto x_rank = x_dim.size();
+
+  int batch_size = 1;
+  int row_number = 1;
+  for (int i = 0; i < x_rank - 1; ++i) {
+    if (i < x_rank - 2) {
+      batch_size *= static_cast<int>(x_dim[i]);
+    } else if (i == x_rank - 2) {
+      row_number = static_cast<int>(x_dim[i]);
+    }
+  }
+
+  const DenseTensor& x_crows = x.non_zero_crows();
+  const DenseTensor& x_values = x.non_zero_elements();
+  DenseTensor* out_values = out->mutable_non_zero_elements();
+
+  int row_nnz = 0;
+  T row_max_val = 0;
+  const T* x_data = x_values.data<T>();
+  T* out_data = out_values->data<T>();
+
+  // out = exp(x-x_max) / sum( exp(x-x_max ))
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.non_zero_crows().dtype(), "CsrSoftmaxKernel", ([&] {
+        const data_t* x_crows_data = x_crows.data<data_t>();
+        for (int i = 0; i < batch_size; ++i) {
+          for (int j = 0; j < row_number; ++j) {
+            int crow_idx = i * (row_number + 1) + j;
+            row_nnz = static_cast<int>(x_crows_data[crow_idx + 1] -
+                                       x_crows_data[crow_idx]);
+
+            row_max_val = *std::max_element(x_data, x_data + row_nnz);
+            funcs::vec_add_bias<T, backends::cpu::avx>(
+                row_nnz, static_cast<T>(-1) * row_max_val, x_data, out_data);
+
+            funcs::vec_exp<T>(row_nnz, out_data, out_data);
+
+            T sum = 0;
+            funcs::vec_sum<T, backends::cpu::avx>(row_nnz, out_data, &sum);
+            funcs::vec_scal<T, backends::cpu::avx>(
+                row_nnz, static_cast<T>(1) / sum, out_data, out_data);
+
+            x_data = x_data + row_nnz;
+            out_data = out_data + row_nnz;
+          }
+        }
+      }));
+}
+
+template <typename T, typename IntT, typename Context>
+void SoftmaxCooCPUKernel(const Context& dev_ctx,
+                         const SparseCooTensor& x,
+                         int axis,
+                         SparseCooTensor* out) {
+  auto indices = x.indices();
+  auto values = x.values();
+  const auto x_dims = x.dims();
+  const auto sparse_dim = x.sparse_dim();
+  DenseTensor out_indices(indices);
+  DenseTensor out_values = EmptyLike<T, Context>(dev_ctx, values);
+  out->SetMember(out_indices, out_values, x.dims(), x.coalesced());
+
+  int dim = axis < 0 ? x_dims.size() + axis : axis;
+
+  /* If dim is greater than or equal to sparse_dim, the dense softmax is used.
+   */
+  if (dim >= sparse_dim) {
+    SoftmaxKernel<T, Context>(
+        dev_ctx, values, dim - sparse_dim + 1, &out_values);
+    return;
+  }
+
+  const std::vector<IntT> sizes = vectorize<IntT>(x_dims);
+  std::map<IntT, std::vector<IntT>> pools;
+  IntT nvalues = std::accumulate(sizes.begin() + sparse_dim,
+                                 sizes.end(),
+                                 static_cast<IntT>(1),
+                                 std::multiplies<>());
+  funcs::sparse::GetPoolsSoftmax(out_indices, sizes, dim, &pools);
+
+  auto values_ptr = values.data<T>();
+  auto out_values_ptr = out_values.data<T>();
+  for (size_t p = 0; p < pools.size(); p++) {
+    auto pool_indices = pools[p];
+    if (pool_indices.empty()) {
+      continue;
+    }
+
+    std::vector<T> mx_row(nvalues, -std::numeric_limits<T>::infinity());
+    std::vector<T> exp_sums_row(nvalues, 0);
+    IntT pool_size = static_cast<IntT>(pool_indices.size());
+
+    // Compute max for each pool
+    for (IntT i = 0; i < pool_size; i++) {
+      auto values_row = values_ptr + pool_indices[i] * nvalues;
+      for (IntT j = 0; j < nvalues; j++) {
+        mx_row[j] = std::max(mx_row[j], *(values_row + j));
+      }
+    }
+
+    // exp to (v - mx) and sum the results
+    for (IntT i = 0; i < pool_size; i++) {
+      auto values_row = values_ptr + pool_indices[i] * nvalues;
+      auto out_values_row = out_values_ptr + pool_indices[i] * nvalues;
+      for (IntT j = 0; j < nvalues; j++) {
+        auto v = std::exp(*(values_row + j) - mx_row[j]);
+        out_values_row[j] = v;
+        exp_sums_row[j] += v;
+      }
+    }
+
+    /* Normalize with the sum of exponents */
+    for (IntT i = 0; i < pool_size; i++) {
+      auto out_values_row = out_values_ptr + pool_indices[i] * nvalues;
+      for (IntT j = 0; j < nvalues; j++) {
+        out_values_row[j] *= 1.0 / exp_sums_row[j];
+      }
+    }
+  }
+}
+
+// cpu kernel
+template <typename T, typename Context>
+void SoftmaxCooKernel(const Context& dev_ctx,
+                      const SparseCooTensor& x,
+                      int axis,
+                      SparseCooTensor* out) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.indices().dtype(), "SoftmaxCooCPUKernel", ([&] {
+        SoftmaxCooCPUKernel<T, data_t, Context>(dev_ctx, x, axis, out);
+      }));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(softmax_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SoftmaxCsrKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(softmax_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SoftmaxCooKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/cpu/sparse_utils_kernel.cc

@@ -0,0 +1,482 @@
+/* Copyright (c) 2023 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/sparse_utils_kernel.h"
+
+#include "paddle/phi/api/lib/utils/allocator.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_meta.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/funcs/sparse/common_shape.h"
+
+namespace phi::sparse {
+
+template <typename T>
+inline bool IsZero(const T* data, const size_t n) {
+  const T zero = static_cast<T>(0);
+  for (size_t i = 0; i < n; i++) {
+    if (data[i] != zero) {
+      return false;
+    }
+  }
+  return true;
+}
+
+// TODO(zhangkaihuo): implement a kernel to count the number of non-zero
+// elements in tensor
+template <typename T>
+inline int64_t GetNonZeroNum(const DenseTensor& dense,
+                             const int64_t sparse_dim) {
+  const auto& dims = dense.dims();
+  PADDLE_ENFORCE_GE(
+      dims.size(),
+      sparse_dim,
+      common::errors::InvalidArgument(
+          "sparse_dim(%d) should be less than or equal to dense.dim(%d)",
+          sparse_dim,
+          dims.size()));
+
+  auto dims_2d = flatten_to_2d(dims, static_cast<int>(sparse_dim));
+  const int rows = static_cast<int>(dims_2d[0]);
+  const int cols = static_cast<int>(dims_2d[1]);
+
+  const T* data = dense.data<T>();
+  int64_t non_zero_num = 0;
+  for (int64_t i = 0; i < rows; i++) {
+    if (!IsZero(data + i * cols, cols)) {
+      non_zero_num = non_zero_num + 1;
+    }
+  }
+  return non_zero_num;
+}
+
+template <typename T, typename Context>
+void DenseToCooKernel(const Context& dev_ctx,
+                      const DenseTensor& x,
+                      const int64_t sparse_dim,
+                      SparseCooTensor* out) {
+  const T* x_data = x.data<T>();
+  const auto& x_dims = x.dims();
+  PADDLE_ENFORCE_LE(sparse_dim,
+                    x_dims.size(),
+                    common::errors::InvalidArgument(
+                        "sparse_dim must be less than the size of x.dims()"));
+  PADDLE_ENFORCE_GT(
+      sparse_dim, 0, common::errors::InvalidArgument("sparse_dim must be >0"));
+
+  int64_t non_zero_num = GetNonZeroNum<T>(x, sparse_dim);
+
+  const auto values_dims =
+      funcs::sparse::InferDenseDims(x_dims, sparse_dim, non_zero_num);
+  DenseTensorMeta values_meta(x.meta().dtype, values_dims, x.meta().layout);
+  DenseTensor indices = Empty<int64_t>(dev_ctx, {sparse_dim, non_zero_num});
+  DenseTensor values = Empty(dev_ctx, std::move(values_meta));
+  int64_t* indices_data = indices.data<int64_t>();
+  T* values_data = values.data<T>();
+
+  auto dims_2d = flatten_to_2d(x_dims, static_cast<int>(sparse_dim));
+  const int rows = static_cast<int>(dims_2d[0]);
+  const int cols = static_cast<int>(dims_2d[1]);
+
+  int index = 0;
+  for (int i = 0; i < rows; i++) {
+    if (!IsZero(x_data + i * cols, cols)) {
+      int64_t sparse_index = i;
+      for (int j = static_cast<int>(sparse_dim - 1); j >= 0; j--) {
+        indices_data[j * non_zero_num + index] = sparse_index % x_dims[j];
+        sparse_index /= x_dims[j];
+      }
+      memcpy(values_data + index * cols, x_data + i * cols, cols * sizeof(T));
+      ++index;
+    }
+  }
+
+  out->SetMember(indices, values, x_dims, true);
+}
+
+template <typename T, typename IntT>
+void CsrToCooCPUKernel(const CPUContext& dev_ctx,
+                       const SparseCsrTensor& x,
+                       SparseCooTensor* out) {
+  const DDim& x_dims = x.dims();
+  const int64_t non_zero_num = x.cols().numel();
+  int64_t sparse_dim = 2;
+  if (x_dims.size() == 3) {
+    sparse_dim = 3;
+  }
+  DenseTensor indices = Empty<IntT>(dev_ctx, {sparse_dim, non_zero_num});
+  DenseTensor values = Empty<T>(dev_ctx, {non_zero_num});
+  if (x.nnz() <= 0) {
+    out->SetMember(indices, values, x_dims, true);
+    return;
+  }
+  const auto& csr_crows = x.crows();
+  const auto& csr_cols = x.cols();
+  const auto& csr_values = x.values();
+  const IntT* csr_crows_data = csr_crows.data<IntT>();
+  const IntT* csr_cols_data = csr_cols.data<IntT>();
+  const T* csr_values_data = csr_values.data<T>();
+
+  IntT* coo_indices = indices.data<IntT>();
+  IntT* batch_ptr = x_dims.size() == 2 ? nullptr : coo_indices;
+  IntT* coo_rows_data =
+      x_dims.size() == 2 ? coo_indices : batch_ptr + non_zero_num;
+  IntT* coo_cols_data = coo_rows_data + non_zero_num;
+  T* coo_values_data = values.data<T>();
+
+  int batch = static_cast<int>(x_dims.size() == 2 ? 1 : x_dims[0]);
+  int rows = static_cast<int>(x_dims.size() == 2 ? x_dims[0] : x_dims[1]);
+
+  int index = 0;
+  for (int b = 0; b < batch; b++) {
+    for (int i = 0; i < rows; i++) {
+      for (IntT j = csr_crows_data[b * (rows + 1) + i];
+           j < csr_crows_data[b * (rows + 1) + i + 1];
+           j++) {
+        coo_rows_data[index] = i;
+        if (batch_ptr) {
+          batch_ptr[index] = b;
+        }
+        ++index;
+      }
+    }
+  }
+
+  memcpy(coo_cols_data, csr_cols_data, sizeof(IntT) * non_zero_num);
+  memcpy(coo_values_data, csr_values_data, sizeof(T) * non_zero_num);
+  out->SetMember(indices, values, x_dims, true);
+}
+
+template <typename T, typename Context>
+void CsrToCooKernel(const Context& dev_ctx,
+                    const SparseCsrTensor& x,
+                    SparseCooTensor* out) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(x.crows().dtype(), "CsrToCooCPUKernel", ([&] {
+                                 CsrToCooCPUKernel<T, data_t>(dev_ctx, x, out);
+                               }));
+}
+
+template <typename T, typename IntT>
+void CooToCsrCPUKernel(const CPUContext& dev_ctx,
+                       const SparseCooTensor& x,
+                       SparseCsrTensor* out) {
+  const auto& x_dims = x.dims();
+  bool valid = x_dims.size() == 2 || x_dims.size() == 3;
+  PADDLE_ENFORCE_EQ(valid,
+                    true,
+                    common::errors::InvalidArgument(
+                        "SparseCsrTensor only support 2-D or 3-D matrix"));
+  const int64_t non_zero_num = x.nnz();
+
+  int batches = static_cast<int>(x_dims.size() == 2 ? 1 : x_dims[0]);
+  int rows = static_cast<int>(x_dims.size() == 2 ? x_dims[0] : x_dims[1]);
+
+  DenseTensor crows = Empty<IntT>(dev_ctx, {batches * (rows + 1)});
+  DenseTensor cols = Empty<IntT>(dev_ctx, {non_zero_num});
+  DenseTensor values = EmptyLike<T, CPUContext>(dev_ctx, x.values());
+  if (non_zero_num <= 0) {
+    out->SetMember(crows, cols, values, x_dims);
+    return;
+  }
+  IntT* csr_crows_data = crows.data<IntT>();
+  IntT* csr_cols_data = cols.data<IntT>();
+  T* csr_values_data = values.data<T>();
+
+  const auto& coo_indices = x.indices();
+  const auto& coo_values = x.values();
+  const IntT* batches_ptr = coo_indices.data<IntT>();
+  const IntT* coo_rows_data =
+      x_dims.size() == 2 ? batches_ptr : batches_ptr + non_zero_num;
+  const IntT* coo_cols_data = coo_rows_data + non_zero_num;
+  const T* coo_values_data = coo_values.data<T>();
+
+  std::vector<int64_t> offsets(batches, 0);
+  if (batches > 1) {
+    for (int i = 0; i < non_zero_num; i++) {
+      if (i == non_zero_num - 1 || batches_ptr[i] != batches_ptr[i + 1]) {
+        const int start = batches_ptr[i];
+        const int end = i == non_zero_num - 1 ? batches : batches_ptr[i + 1];
+        for (int j = start; j < end; j++) {
+          offsets[j] = i + 1;
+        }
+      }
+    }
+  } else {
+    offsets[0] = non_zero_num;
+  }
+
+  for (int b = 0; b < batches; b++) {
+    int batch_start = 0;
+    int batch_non_zero_num = static_cast<int>(offsets[b]);
+    if (b > 0) {
+      batch_start = static_cast<int>(offsets[b - 1]);
+      batch_non_zero_num -= batch_start;
+    }
+    auto* coo_rows_ptr = coo_rows_data + batch_start;
+    for (int i = 0; i <= coo_rows_ptr[0]; i++) {
+      csr_crows_data[b * (rows + 1) + i] = 0;
+    }
+    for (int64_t i = 1; i < batch_non_zero_num; i++) {
+      for (IntT j = coo_rows_ptr[i - 1]; j < coo_rows_ptr[i]; j++) {
+        csr_crows_data[b * (rows + 1) + j + 1] = i;
+      }
+    }
+    for (IntT i = coo_rows_ptr[batch_non_zero_num - 1] + 1; i < rows + 1; i++) {
+      csr_crows_data[b * (rows + 1) + i] = batch_non_zero_num;
+    }
+    if (batch_non_zero_num == 0) {
+      memset(csr_crows_data + b * (rows + 1), 0, sizeof(IntT) * (rows + 1));
+    }
+  }
+
+  memcpy(csr_cols_data, coo_cols_data, sizeof(IntT) * non_zero_num);
+  memcpy(csr_values_data, coo_values_data, sizeof(T) * non_zero_num);
+  out->SetMember(crows, cols, values, x_dims);
+}
+
+template <typename T, typename Context>
+void CooToCsrKernel(const Context& dev_ctx,
+                    const SparseCooTensor& x,
+                    SparseCsrTensor* out) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "CooToCsrCPUKernel", ([&] {
+                                 CooToCsrCPUKernel<T, data_t>(dev_ctx, x, out);
+                               }));
+}
+
+template <typename T, typename IntT>
+void CooToDenseCPUKernel(const CPUContext& dev_ctx,
+                         const SparseCooTensor& x,
+                         DenseTensor* out) {
+  const auto non_zero_num = x.nnz();
+  const auto& dense_dims = x.dims();
+  const auto& indices = x.indices();
+  const auto& values = x.values();
+  const auto indices_dims = vectorize<int>(indices.dims());
+  int64_t sparse_dim = indices_dims[0];
+  if (indices_dims.size() == 1) {
+    sparse_dim = 1;
+  }
+  const int64_t dense_dim = x.dense_dim();
+
+  const T* x_data = values.data<T>();
+  dev_ctx.template Alloc<T>(out);
+  T* out_data = out->data<T>();
+  memset(out_data, 0, sizeof(T) * out->numel());
+
+  if (x.nnz() <= 0) {
+    return;
+  }
+
+  int64_t base_offset = 1;
+  for (int64_t i = 0; i < dense_dim; i++) {
+    base_offset *= dense_dims[static_cast<int>(sparse_dim + i)];
+  }
+  std::vector<int64_t> sparse_offsets(sparse_dim);
+  int64_t offset = 1;
+  for (int i = static_cast<int>(sparse_dim - 1); i >= 0; i--) {
+    sparse_offsets[i] = offset;
+    offset *= dense_dims[i];
+  }
+
+  for (auto i = 0; i < non_zero_num; i++) {
+    int64_t index = 0;
+    for (int j = 0; j < sparse_dim; j++) {
+      index += indices.data<IntT>()[j * non_zero_num + i] * sparse_offsets[j];
+    }
+
+    for (int j = 0; j < base_offset; j++) {
+      out_data[index * base_offset + j] = x_data[i * base_offset + j];
+    }
+  }
+}
+
+template <typename T, typename Context>
+void CooToDenseKernel(const Context& dev_ctx,
+                      const SparseCooTensor& x,
+                      DenseTensor* out) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.indices().dtype(), "CooToDenseCPUKernel", ([&] {
+        CooToDenseCPUKernel<T, data_t>(dev_ctx, x, out);
+      }));
+
+  // Set proper dense layout after conversion from sparse
+  // SparseCooTensor uses SPARSE_COO layout, but DenseTensor should use
+  // a standard dense layout (NCHW, NHWC, etc.)
+  if (out->meta().layout == DataLayout::SPARSE_COO ||
+      out->meta().layout == DataLayout::SPARSE_CSR) {
+    // Default to NCHW for dense tensors
+    out->set_meta(DenseTensorMeta(out->dtype(), out->dims(), DataLayout::NCHW));
+  }
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(dense_to_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::DenseToCooKernel,
+                   float,
+                   double,
+                   paddle::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {}
+
+PD_REGISTER_KERNEL(csr_to_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CsrToCooKernel,
+                   float,
+                   double,
+                   paddle::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {}
+
+PD_REGISTER_KERNEL(coo_to_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CooToCsrKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {}
+
+PD_REGISTER_KERNEL(dense_to_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::DenseToCsrKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {}
+
+PD_REGISTER_KERNEL(coo_to_dense,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CooToDenseKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {}
+
+PD_REGISTER_KERNEL(csr_to_dense,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CsrToDenseKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {}
+
+PD_REGISTER_KERNEL(values_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ValuesCooKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(indices_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::IndicesCooKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(values_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ValuesCsrKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(sparse_coo_tensor,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SparseCooTensorKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {}

paddle/phi/kernels/sparse/cpu/sum_grad_kernel.cc

@@ -0,0 +1,218 @@
+// Copyright (c) 2023 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/unary_grad_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/reduce_sum_grad_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/impl/unary_grad_kernel_impl.h"
+
+namespace phi::sparse {
+
+template <typename T, typename IntT, typename Context>
+void SumCooGradCPUKernel(const Context& dev_ctx,
+                         const SparseCooTensor& x,
+                         const SparseCooTensor& dout,
+                         const IntArray& axis,
+                         bool keep_dim,
+                         SparseCooTensor* dx) {
+  EmptyLikeCooKernel<T, Context>(dev_ctx, x, dx);
+  unsigned int n_dim = axis.size();
+
+  const DenseTensor& x_indices = x.indices();
+  const DenseTensor& dout_indices = dout.indices();
+  const DenseTensor& dout_values = dout.values();
+  const auto* dout_indices_data = dout_indices.data<int64_t>();
+  const auto* dout_values_data = dout_values.data<T>();
+
+  DenseTensor* dx_indices = dx->mutable_indices();
+  DenseTensor* dx_values = dx->mutable_values();
+  *dx_indices = x_indices;
+
+  const auto* dx_indices_data = dx_indices->data<int64_t>();
+  auto* dx_values_data = dx_values->data<T>();
+
+  funcs::SetConstant<Context, T> set_constant;
+  if (n_dim == 0) {
+    T value = dout_values.data<T>()[0];
+    set_constant(dev_ctx, dx_values, value);
+    if (dx_values->dtype() != dx->dtype()) {
+      *dx_values = Cast<T, Context>(dev_ctx, *dx_values, dx->dtype());
+    }
+    return;
+  }
+
+  auto dim = axis[0] < 0 ? x.dims().size() + axis[0] : axis[0];
+  auto sparse_dim = x.sparse_dim();
+  if (dim >= sparse_dim) {
+    dim = dim - sparse_dim + 1;
+    phi::ReduceSumGradKernel<T, Context>(
+        dev_ctx, x.values(), dout.values(), {dim}, keep_dim, false, dx_values);
+    if (dx_values->dtype() != dx->dtype()) {
+      *dx_values = Cast<T, Context>(dev_ctx, *dx_values, dx->dtype());
+    }
+    return;
+  }
+  // Ensure the sparse_dim is not less than 1.
+  if (sparse_dim == 1) {
+    keep_dim = true;
+  }
+
+  int64_t dense_dim = 1;
+  for (auto i = 1; i < x.values().dims().size(); ++i) {
+    dense_dim *= x.values().dims()[i];
+  }
+
+  std::map<std::vector<IntT>, int64_t> indices_map;
+  for (auto j = 0; j < dout_indices.dims()[1]; ++j) {
+    std::vector<IntT> pos;
+    pos.reserve(dout_indices.dims()[0]);
+    for (int i = 0; i < dout_indices.dims()[0]; ++i) {
+      pos.push_back(dout_indices_data[j + i * dout_indices.dims()[1]]);
+    }
+    indices_map[pos] = j;
+  }
+
+  for (auto j = 0; j < dx_indices->dims()[1]; ++j) {
+    std::vector<IntT> pos;
+    for (int i = 0; i < dx_indices->dims()[0]; ++i) {
+      if (i != dim) {
+        pos.push_back(dx_indices_data[j + i * dx_indices->dims()[1]]);
+      } else if (keep_dim) {
+        pos.push_back(0);
+      }
+    }
+    for (int i = 0; i < dense_dim; ++i) {
+      dx_values_data[i + j * dense_dim] =
+          dout_values_data[i + indices_map[pos] * dense_dim];
+    }
+  }
+  if (dx_values->dtype() != dx->dtype()) {
+    *dx_values = Cast<T, Context>(dev_ctx, *dx_values, dx->dtype());
+  }
+}
+
+template <typename T, typename Context>
+void SumCsrGradKernel(const Context& dev_ctx,
+                      const SparseCsrTensor& x,
+                      const SparseCsrTensor& dout,
+                      const IntArray& axis,
+                      bool keep_dim UNUSED,
+                      SparseCsrTensor* dx) {
+  EmptyLikeCsrKernel<T, Context>(dev_ctx, x, dx);
+  unsigned int n_dim = axis.size();
+
+  const DenseTensor& x_crows = x.crows();
+  const DenseTensor& x_cols = x.cols();
+  const DenseTensor& dout_values = dout.values();
+  const auto* x_crows_data = x_crows.data<int64_t>();
+
+  DenseTensor* dx_crows = dx->mutable_crows();
+  DenseTensor* dx_cols = dx->mutable_cols();
+  DenseTensor* dx_values = dx->mutable_values();
+
+  *dx_crows = x_crows;
+  *dx_cols = x_cols;
+
+  funcs::SetConstant<Context, T> set_constant;
+  if (n_dim == 0) {
+    T value = dout_values.data<T>()[0];
+    set_constant(dev_ctx, dx_values, value);
+    if (dx_values->dtype() != dx->dtype()) {
+      *dx_values = Cast<T, Context>(dev_ctx, *dx_values, dx->dtype());
+    }
+    return;
+  }
+  PADDLE_ENFORCE_EQ(axis[0],
+                    -1,
+                    common::errors::Unimplemented(
+                        "`axis` of SumCsrKernel only support None or -1 now."
+                        "More number will be supported in the future."));
+
+  if (x.dims().size() == 2) {
+    int value_index = 0;
+    for (int k = 0; k < x.dims()[0]; ++k) {
+      if (x_crows_data[k] == x_crows_data[k + 1]) {
+        continue;
+      }
+      T value = dout_values.data<T>()[value_index];
+      set_constant(dev_ctx, dx_values, value);
+      value_index += 1;
+    }
+  } else {
+    int dout_value_index = 0;
+    int dx_value_index = 0;
+    for (auto batch = 0; batch < x.dims()[0]; ++batch) {
+      for (auto k = batch * (x.dims()[1] + 1);
+           k < batch * (x.dims()[1] + 1) + x.dims()[1];
+           ++k) {
+        if (x_crows_data[k] == x_crows_data[k + 1]) {
+          continue;
+        }
+        T value = dout_values.data<T>()[dout_value_index];
+        for (auto i = x_crows_data[k]; i < x_crows_data[k + 1]; ++i) {
+          dx_values->data<T>()[dx_value_index] = value;
+          dx_value_index++;
+        }
+        dout_value_index++;
+      }
+    }
+  }
+
+  if (dx_values->dtype() != dx->dtype()) {
+    *dx_values = Cast<T, Context>(dev_ctx, *dx_values, dx->dtype());
+  }
+}
+
+template <typename T, typename Context>
+void SumCooGradKernel(const Context& dev_ctx,
+                      const SparseCooTensor& x,
+                      const SparseCooTensor& dout,
+                      const IntArray& axis,
+                      bool keep_dim,
+                      SparseCooTensor* dx) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.indices().dtype(), "SumCooGradCPUKernel", ([&] {
+        SumCooGradCPUKernel<T, data_t, Context>(
+            dev_ctx, x, dout, axis, keep_dim, dx);
+      }));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(sum_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SumCooGradKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(sum_csr_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SumCsrGradKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/cpu/sum_kernel.cc

@@ -0,0 +1,281 @@
+// Copyright (c) 2023 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/unary_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi::sparse {
+
+template <typename T, typename IntT, typename Context>
+void SumCooCPUKernel(const Context& dev_ctx,
+                     const SparseCooTensor& x,
+                     const IntArray& axis,
+                     DataType dtype,
+                     bool keep_dim,
+                     SparseCooTensor* out) {
+  size_t n_dim = axis.size();
+  auto sparse_dim = x.sparse_dim();
+  // create out sparse tensor
+  const auto& x_dims = x.dims();
+  const auto& x_indices = x.indices();
+  const auto& x_values = x.values();
+  DDim out_dims;
+  DenseTensor out_indices;
+  DenseTensor out_values;
+  if (n_dim == 0) {
+    std::vector<int64_t> out_indices_shape;
+    if (keep_dim) {
+      out_dims = make_ddim(std::vector<int64_t>(x_dims.size(), 1));
+      out_indices_shape = {sparse_dim, 1};
+    } else {
+      out_dims = make_ddim({1});
+      out_indices_shape = {1};
+    }
+    out_indices = Empty<IntT, Context>(dev_ctx, out_indices_shape);
+    auto* out_indices_data = out_indices.data<IntT>();
+    std::fill(out_indices_data, out_indices_data + out_indices.numel(), 0);
+    out_values = phi::Sum<T>(dev_ctx, x.values(), {}, dtype, keep_dim);
+    out->SetMember(out_indices, out_values, out_dims, x.coalesced());
+    return;
+  }
+
+  auto dim = axis[0] < 0 ? x_dims.size() + axis[0] : axis[0];
+  const auto* x_indices_data = x_indices.data<IntT>();
+  const auto* x_values_data = x_values.data<T>();
+
+  std::vector<int64_t> dims;
+  for (int i = 0; i < x.dims().size(); ++i) {
+    if (i != dim) {
+      dims.emplace_back(x.dims()[i]);
+    } else if (keep_dim || (dim < sparse_dim && sparse_dim == 1)) {
+      dims.emplace_back(1);
+    }
+  }
+  out_dims = make_ddim(dims);
+
+  if (dim >= sparse_dim) {
+    out_indices = x_indices;
+    dim = dim - sparse_dim + 1;
+    out_values = phi::Sum<T>(dev_ctx, x.values(), {dim}, dtype, keep_dim);
+    out->SetMember(out_indices, out_values, out_dims, x.coalesced());
+    return;
+  }
+
+  // Ensure the sparse_dim is not less than 1.
+  if (sparse_dim == 1) {
+    keep_dim = true;
+  }
+  // if axis in sparse_dim and keep_dim, sparse_dim will be reduced.
+  if (!keep_dim) {
+    sparse_dim -= 1;
+  }
+
+  // indices_map is a mapping from output's position to values to be summed.
+  std::map<std::vector<IntT>, std::vector<int64_t>> indices_map;
+  for (int64_t j = 0; j < x_indices.dims()[1]; ++j) {
+    std::vector<IntT> pos;
+    for (int64_t i = 0; i < x_indices.dims()[0]; ++i) {
+      if (dim != i) {
+        pos.emplace_back(x_indices_data[j + i * x_indices.dims()[1]]);
+      } else if (keep_dim) {
+        pos.emplace_back(0);
+      }
+    }
+    indices_map[pos].emplace_back(j);
+  }
+
+  std::vector<int> out_values_dims;
+  out_values_dims.push_back(static_cast<int>(indices_map.size()));
+  for (auto i = 1; i < x.values().dims().size(); ++i) {
+    out_values_dims.push_back(static_cast<int>(x.values().dims()[i]));
+  }
+  int64_t dense_dim = std::accumulate(out_values_dims.begin() + 1,
+                                      out_values_dims.end(),
+                                      1,
+                                      std::multiplies<int64_t>());
+
+  out_indices = Empty<IntT, Context>(
+      dev_ctx, {sparse_dim, static_cast<int>(indices_map.size())});
+  out_values = Empty<T, Context>(dev_ctx, out_values_dims);
+
+  auto* out_indices_data = out_indices.data<IntT>();
+  auto* out_values_data = out_values.data<T>();
+
+  auto iter_indices_map = indices_map.begin();
+  for (size_t j = 0; j < indices_map.size(); ++j) {
+    std::vector<IntT> pos = iter_indices_map->first;
+    std::vector<int64_t> values_index = iter_indices_map->second;
+    iter_indices_map++;
+    for (auto i = 0; i < sparse_dim; ++i) {
+      out_indices_data[j + i * indices_map.size()] = pos[i];
+    }
+    for (auto i = 0; i < dense_dim; ++i) {
+      T out_value = 0;
+      for (auto index : values_index) {
+        out_value += x_values_data[i + index * dense_dim];
+      }
+      out_values_data[i + j * dense_dim] = out_value;
+    }
+  }
+
+  if (dtype != phi::DataType::UNDEFINED && dtype != x.dtype()) {
+    out_values = Cast<T, Context>(dev_ctx, out_values, dtype);
+  }
+  out->SetMember(out_indices, out_values, out_dims, x.coalesced());
+}
+
+template <typename T, typename Context>
+void SumCsrKernel(const Context& dev_ctx,
+                  const SparseCsrTensor& x,
+                  const IntArray& axis,
+                  DataType dtype,
+                  bool keep_dim,
+                  SparseCsrTensor* out) {
+  size_t n_dim = axis.size();
+  const auto& x_crows = x.crows();
+  const auto& x_values = x.values();
+  const auto* x_crows_data = x_crows.data<int64_t>();
+  const auto* x_values_data = x_values.data<T>();
+
+  DenseTensor out_crows, out_cols, out_values;
+  DDim out_dims;
+  if (n_dim == 0) {
+    if (keep_dim && x.dims().size() == 3) {
+      out_dims = make_ddim({1, 1, 1});
+    } else {
+      out_dims = make_ddim({1, 1});
+    }
+    out_crows = Empty<int64_t, Context>(dev_ctx, {2});  // crows = [0, 1]
+    auto* out_crows_data = out_crows.data<int64_t>();
+    out_crows_data[0] = 0;
+    out_crows_data[1] = 1;
+
+    out_cols = Empty<int64_t, Context>(dev_ctx, {1});  // crows = [0]
+    auto* out_cols_data = out_cols.data<int64_t>();
+    out_cols_data[0] = 0;
+    out_values = phi::Sum<T>(dev_ctx, x.values(), {}, dtype, true);
+  } else {
+    PADDLE_ENFORCE_EQ(axis[0],
+                      -1,
+                      common::errors::Unimplemented(
+                          "`axis` of SumCsrKernel only support None or -1 now."
+                          "More number will be supported in the future."));
+    out_crows = EmptyLike<int64_t, Context>(dev_ctx, x.crows());
+    auto* out_crows_data = out_crows.data<int64_t>();
+    std::vector<T> out_data;
+    if (x.dims().size() == 2) {
+      out_crows_data[0] = 0;
+      out_dims = make_ddim({x.dims()[0], 1});
+      for (int i = 0; i < x.dims()[0]; ++i) {
+        if (x_crows_data[i] != x_crows_data[i + 1]) {
+          T sum_value = 0;
+          for (auto j = x_crows_data[i]; j < x_crows_data[i + 1]; ++j) {
+            sum_value += x_values_data[j];
+          }
+          out_crows_data[i + 1] = out_crows_data[i] + 1;
+          out_data.emplace_back(sum_value);
+        } else {
+          out_crows_data[i + 1] = out_crows_data[i];
+        }
+      }
+    } else {
+      if (keep_dim) {
+        out_dims = make_ddim({x.dims()[0], x.dims()[1], 1});
+      } else {
+        out_dims = make_ddim({x.dims()[0], x.dims()[1]});
+      }
+      int j = 0;
+      for (int batch = 0; batch < x.dims()[0]; ++batch) {
+        auto* cur_x_crows_data = x_crows_data + batch * x.dims()[2];
+        auto* cur_out_crows_data = out_crows_data + batch * x.dims()[2];
+        for (int i = 0; i < x.dims()[1]; ++i) {
+          cur_out_crows_data[0] = 0;
+          if (cur_x_crows_data[i] != cur_x_crows_data[i + 1]) {
+            T sum_value = 0;
+            for (auto k = cur_x_crows_data[i]; k < cur_x_crows_data[i + 1];
+                 ++k) {
+              sum_value += x_values_data[j++];
+            }
+            out_data.emplace_back(sum_value);
+            cur_out_crows_data[i + 1] = cur_out_crows_data[i] + 1;
+          } else {
+            cur_out_crows_data[i + 1] = cur_out_crows_data[i];
+          }
+        }
+      }
+    }
+    out_cols =
+        Empty<int64_t, Context>(dev_ctx, {static_cast<int>(out_data.size())});
+    out_values =
+        Empty<T, Context>(dev_ctx, {static_cast<int>(out_data.size())});
+    auto* out_cols_data = out_cols.data<int64_t>();
+    T* out_values_data = out_values.data<T>();
+    for (size_t i = 0; i < out_data.size(); ++i) {
+      out_cols_data[i] = 0;
+      out_values_data[i] = out_data[i];
+    }
+    if (dtype != phi::DataType::UNDEFINED && dtype != x.dtype()) {
+      out_values = Cast<T, Context>(dev_ctx, out_values, dtype);
+    }
+  }
+  out->SetMember(out_crows, out_cols, out_values, out_dims);
+}
+
+template <typename T, typename Context>
+void SumCooKernel(const Context& dev_ctx,
+                  const SparseCooTensor& x,
+                  const IntArray& axis,
+                  DataType dtype,
+                  bool keep_dim,
+                  SparseCooTensor* out) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "SumCooCPUKernel", ([&] {
+                                 SumCooCPUKernel<T, data_t, Context>(
+                                     dev_ctx, x, axis, dtype, keep_dim, out);
+                               }));
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(sum_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SumCooKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}
+
+PD_REGISTER_KERNEL(sum_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SumCsrKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}

paddle/phi/kernels/sparse/cpu/transpose_grad_kernel.cc

@@ -0,0 +1,76 @@
+// 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/unary_grad_kernel.h"
+#include "paddle/phi/kernels/sparse/unary_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/impl/unary_grad_kernel_impl.h"
+
+namespace phi::sparse {
+
+std::vector<int> get_cpu_grad_perm(std::vector<int> perm) {
+  std::vector<int> grad_perm(perm.size());
+  for (unsigned int i = 0; i < perm.size(); ++i) {
+    grad_perm[perm[i]] = static_cast<int>(i);
+  }
+  return grad_perm;
+}
+
+template <typename T, typename Context>
+void TransposeCooGradKernel(const Context& dev_ctx,
+                            const SparseCooTensor& dout,
+                            const std::vector<int>& perm,
+                            SparseCooTensor* dx) {
+  std::vector<int> grad_perm = get_cpu_grad_perm(perm);
+  TransposeCooKernel<T, Context>(dev_ctx, dout, grad_perm, dx);
+}
+
+template <typename T, typename Context>
+void TransposeCsrGradKernel(const Context& dev_ctx,
+                            const SparseCsrTensor& dout,
+                            const std::vector<int>& perm,
+                            SparseCsrTensor* dx) {
+  std::vector<int> grad_perm = get_cpu_grad_perm(perm);
+  TransposeCsrKernel<T, Context>(dev_ctx, dout, grad_perm, dx);
+}
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(transpose_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::TransposeCooGradKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(transpose_csr_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::TransposeCsrGradKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/cpu/transpose_kernel.cc

@@ -0,0 +1,229 @@
+// 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/unary_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void TransposeCooKernel(const Context& dev_ctx,
+                        const SparseCooTensor& x,
+                        const std::vector<int>& perm,
+                        SparseCooTensor* out) {
+  // create out sparse tensor
+  int64_t x_nnz = x.nnz();
+  DDim out_dims = x.dims().transpose(perm);
+  DenseTensor out_indices = EmptyLike<int64_t, Context>(dev_ctx, x.indices());
+  const DenseTensor& out_values(x.values());
+  out->SetMember(out_indices, out_values, out_dims, x.coalesced());
+
+  // compute values of indices
+  const DenseTensor& x_indices = x.indices();
+  const auto* x_indices_data = x_indices.data<int64_t>();
+  auto* out_indices_data = out_indices.data<int64_t>();
+  for (unsigned int i = 0; i < perm.size(); ++i) {
+    for (int64_t j = 0; j < x_nnz; ++j) {
+      out_indices_data[j + i * x_nnz] = x_indices_data[j + perm[i] * x_nnz];
+    }
+  }
+}
+
+template <typename T, typename Context>
+void TransposeCsrKernel(const Context& dev_ctx,
+                        const SparseCsrTensor& x,
+                        const std::vector<int>& perm,
+                        SparseCsrTensor* out) {
+  unsigned int n_dim = perm.size();
+  const DenseTensor& x_crows = x.crows();
+  const DenseTensor& x_cols = x.cols();
+  const DenseTensor& x_values = x.values();
+  DenseTensor out_crows, out_cols, out_values;
+  // return a copy of x
+  if (perm[0] == 0 && perm[1] == 1 && (n_dim == 2 || perm[2] == 2)) {
+    out_crows = x_crows;
+    out_cols = x_cols;
+    out_values = x_values;
+    out->SetMember(out_crows, out_cols, out_values, x.dims());
+    return;
+  }
+  // create out sparse tensor
+  DDim out_dims = x.dims().transpose(perm);
+  if (n_dim == 2) {
+    out_crows = Empty<int64_t, Context>(dev_ctx, {out_dims[0] + 1});
+  } else {
+    out_crows =
+        Empty<int64_t, Context>(dev_ctx, {out_dims[0] * (out_dims[1] + 1)});
+  }
+  out_cols = EmptyLike<int64_t, Context>(dev_ctx, x.cols());
+  out_values = EmptyLike<T, Context>(dev_ctx, x.values());
+  out->SetMember(out_crows, out_cols, out_values, out_dims);
+  // transpose by two stages
+  if (perm[0] == 1 && perm[1] == 2) {  // perm == {1, 2, 0}
+    SparseCsrTensor temp;
+    TransposeCsrKernel<T, Context>(dev_ctx, x, {1, 0, 2}, &temp);
+    TransposeCsrKernel<T, Context>(dev_ctx, temp, {0, 2, 1}, out);
+    return;
+  } else if (perm[0] == 2 && perm[1] == 0) {  // perm == {2, 0, 1}
+    SparseCsrTensor temp;
+    TransposeCsrKernel<T, Context>(dev_ctx, x, {0, 2, 1}, &temp);
+    TransposeCsrKernel<T, Context>(dev_ctx, temp, {1, 0, 2}, out);
+    return;
+  } else if (perm[0] == 2 && perm[1] == 1) {  // perm == {2, 1, 0}
+    SparseCsrTensor temp;
+    TransposeCsrKernel<T, Context>(dev_ctx, x, {1, 0, 2}, &temp);
+    TransposeCsrKernel<T, Context>(dev_ctx, temp, {2, 0, 1}, out);
+    return;
+  }
+
+  int64_t* out_crows_data = out_crows.data<int64_t>();
+  int64_t* out_cols_data = out_cols.data<int64_t>();
+  T* out_values_data = out_values.data<T>();
+  const int64_t* x_crows_data = x_crows.data<int64_t>();
+  const int64_t* x_cols_data = x_cols.data<int64_t>();
+  const T* x_values_data = x_values.data<T>();
+
+  int64_t x_nnz = x.nnz();
+  if (n_dim == 2) {  // perm == {1, 0}
+    // compute out_crows_data by x_cols_data
+    for (int i = 0; i < out_dims[0]; ++i) {
+      out_crows_data[i] = 0;
+    }
+    for (int i = 0; i < x_nnz; ++i) {
+      int64_t j = x_cols_data[i];
+      out_crows_data[j + 1]++;
+    }
+    out_crows_data[out_dims[0]] = x_nnz;
+    for (int i = 1; i < out_dims[0]; ++i) {
+      out_crows_data[i] += out_crows_data[i - 1];
+    }
+    // compute out_cols_data and out_values_data by out_crows_data and x
+    std::unordered_map<int64_t, int> cols_offset;
+    for (int i = 0; i < x.dims()[0]; ++i) {
+      int64_t start = x_crows_data[i];
+      int64_t end = x_crows_data[i + 1];
+      for (int64_t j = start; j < end; ++j) {
+        int64_t x_cols_j = x_cols_data[j];
+        int64_t jjj = out_crows_data[x_cols_j];
+        if (cols_offset.count(jjj)) {
+          cols_offset[jjj]++;
+        } else {
+          cols_offset[jjj] = 0;
+        }
+        int64_t jjj_offset = jjj + cols_offset[jjj];
+        out_cols_data[jjj_offset] = i;
+        out_values_data[jjj_offset] = x_values_data[j];
+      }
+    }
+  } else {  // n_dim == 3
+    int64_t out_n_rows = out_dims[1];
+    int64_t x_n_rows = x.dims()[1];
+    for (int k = 0; k < out_dims[0]; ++k) {
+      if (perm[0] == 0) {  // perm == {0, 2, 1}
+        // compute out_crows_data by x_cols_data
+        for (int i = 0; i < out_n_rows; ++i) {
+          out_crows_data[i] = 0;
+        }
+        for (int i = 0; i < x_crows_data[x_n_rows]; ++i) {
+          int64_t j = x_cols_data[i];
+          out_crows_data[j + 1]++;
+        }
+        out_crows_data[out_n_rows] = x_crows_data[x_n_rows];
+        for (int i = 1; i < out_n_rows; ++i) {
+          out_crows_data[i] += out_crows_data[i - 1];
+        }
+        // compute out_cols_data and out_values_data by out_crows_data and x
+        std::unordered_map<int64_t, int> cols_offset;
+        for (int i = 0; i < x_n_rows; ++i) {
+          int64_t start = x_crows_data[i];
+          int64_t end = x_crows_data[i + 1];
+          for (int64_t j = start; j < end; ++j) {
+            int64_t x_cols_j = x_cols_data[j];
+            int64_t jjj = out_crows_data[x_cols_j];
+            if (cols_offset.count(jjj)) {
+              cols_offset[jjj]++;
+            } else {
+              cols_offset[jjj] = 0;
+            }
+            int64_t jjj_offset = jjj + cols_offset[jjj];
+            out_cols_data[jjj_offset] = i;
+            out_values_data[jjj_offset] = x_values_data[j];
+          }
+        }
+        // x offset
+        x_cols_data += x_crows_data[x_n_rows];
+        x_values_data += x_crows_data[x_n_rows];
+        x_crows_data += x_n_rows + 1;
+      } else if (perm[0] == 1 && perm[1] == 0) {  // perm == {1, 0, 2}
+        for (int i = 0; i < out_n_rows; ++i) {
+          out_crows_data[i] = 0;
+        }
+        int64_t x_cols_offset = 0;
+        int out_cols_index = 0;
+        for (int i = 0; i < x.dims()[0]; ++i) {
+          int x_crows_index = static_cast<int>(i * (x_n_rows + 1));
+          int64_t start = x_crows_data[x_crows_index + k];
+          int64_t end = x_crows_data[x_crows_index + 1 + k];
+          out_crows_data[i + 1] = end - start;
+          for (int64_t j = start; j < end; ++j) {
+            out_cols_data[out_cols_index] = x_cols_data[x_cols_offset + j];
+            out_values_data[out_cols_index] = x_values_data[x_cols_offset + j];
+            out_cols_index++;
+          }
+          x_cols_offset += x_crows_data[x_crows_index + x_n_rows];
+        }
+        for (int i = 1; i <= out_n_rows; ++i) {
+          out_crows_data[i] += out_crows_data[i - 1];
+        }
+      }
+      // out offset
+      out_cols_data += out_crows_data[out_n_rows];
+      out_values_data += out_crows_data[out_n_rows];
+      out_crows_data += out_n_rows + 1;
+    }
+  }
+}
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(transpose_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::TransposeCooKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(transpose_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::TransposeCsrKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/cpu/unary_grad_kernel.cc

@@ -0,0 +1,106 @@
+// 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/unary_grad_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/sparse/impl/unary_grad_kernel_impl.h"
+
+#define PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(name, prefix)     \
+  PD_REGISTER_KERNEL(name##_coo_grad,                              \
+                     CPU,                                          \
+                     ALL_LAYOUT,                                   \
+                     phi::sparse::prefix##CooGradKernel,           \
+                     float,                                        \
+                     double) {                                     \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); \
+  }                                                                \
+                                                                   \
+  PD_REGISTER_KERNEL(name##_csr_grad,                              \
+                     CPU,                                          \
+                     ALL_LAYOUT,                                   \
+                     phi::sparse::prefix##CsrGradKernel,           \
+                     float,                                        \
+                     double) {                                     \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \
+  }
+
+#define PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(name, prefix) \
+  PD_REGISTER_KERNEL(name##_coo_grad,                                       \
+                     CPU,                                                   \
+                     ALL_LAYOUT,                                            \
+                     phi::sparse::prefix##CooGradKernel,                    \
+                     float,                                                 \
+                     double,                                                \
+                     phi::complex64,                                        \
+                     phi::complex128) {                                     \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);          \
+  }                                                                         \
+                                                                            \
+  PD_REGISTER_KERNEL(name##_csr_grad,                                       \
+                     CPU,                                                   \
+                     ALL_LAYOUT,                                            \
+                     phi::sparse::prefix##CsrGradKernel,                    \
+                     float,                                                 \
+                     double,                                                \
+                     phi::complex64,                                        \
+                     phi::complex128) {                                     \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);          \
+  }
+
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(sqrt, Sqrt)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(relu, Relu)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(pow, Pow)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(relu6, Relu6)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(leaky_relu, LeakyRelu)
+
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(atan, Atan)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(asin, Asin)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(asinh, Asinh)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(atanh, Atanh)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(expm1, Expm1)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(log1p, Log1p)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(square, Square)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(sinh, Sinh)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(tan, Tan)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(sin, Sin)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(abs, Abs)
+PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(tanh, Tanh)
+
+PD_REGISTER_KERNEL(cast_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CastCooGradKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(cast_csr_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CastCsrGradKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/cpu/unary_kernel.cc

@@ -0,0 +1,186 @@
+// 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/unary_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
+#include "paddle/phi/kernels/sparse/impl/unary_grad_kernel_impl.h"
+#include "paddle/phi/kernels/sparse/impl/unary_kernel_impl.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void DivScalarCooKernel(const Context& dev_ctx,
+                        const SparseCooTensor& x,
+                        float scalar,
+                        SparseCooTensor* out) {
+  EmptyLikeCooKernel<T, Context>(dev_ctx, x, out);
+
+  auto eigen_out = EigenVector<T>::Flatten(*(out->mutable_non_zero_elements()));
+  auto eigen_x = EigenVector<T>::Flatten(x.non_zero_elements());
+  auto& dev = *dev_ctx.eigen_device();
+
+  funcs::EigenDiv<std::decay_t<decltype(dev)>, T>::Eval(
+      dev, eigen_out, eigen_x, static_cast<T>(scalar));
+}
+
+template <typename T, typename Context>
+void DivScalarCsrKernel(const Context& dev_ctx,
+                        const SparseCsrTensor& x,
+                        float scalar,
+                        SparseCsrTensor* out) {
+  EmptyLikeCsrKernel<T, Context>(dev_ctx, x, out);
+
+  auto eigen_out = EigenVector<T>::Flatten(*(out->mutable_non_zero_elements()));
+  auto eigen_x = EigenVector<T>::Flatten(x.non_zero_elements());
+  auto& dev = *dev_ctx.eigen_device();
+
+  funcs::EigenDiv<std::decay_t<decltype(dev)>, T>::Eval(
+      dev, eigen_out, eigen_x, static_cast<T>(scalar));
+}
+
+}  // namespace phi::sparse
+
+#define PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(name, prefix)          \
+  PD_REGISTER_KERNEL(name##_coo,                                   \
+                     CPU,                                          \
+                     ALL_LAYOUT,                                   \
+                     phi::sparse::prefix##CooKernel,               \
+                     float,                                        \
+                     double) {                                     \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); \
+  }                                                                \
+                                                                   \
+  PD_REGISTER_KERNEL(name##_csr,                                   \
+                     CPU,                                          \
+                     ALL_LAYOUT,                                   \
+                     phi::sparse::prefix##CsrKernel,               \
+                     float,                                        \
+                     double) {                                     \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \
+  }
+
+#define PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(name, prefix) \
+  PD_REGISTER_KERNEL(name##_coo,                                       \
+                     CPU,                                              \
+                     ALL_LAYOUT,                                       \
+                     phi::sparse::prefix##CooKernel,                   \
+                     float,                                            \
+                     double,                                           \
+                     phi::complex64,                                   \
+                     phi::complex128) {                                \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);     \
+  }                                                                    \
+                                                                       \
+  PD_REGISTER_KERNEL(name##_csr,                                       \
+                     CPU,                                              \
+                     ALL_LAYOUT,                                       \
+                     phi::sparse::prefix##CsrKernel,                   \
+                     float,                                            \
+                     double,                                           \
+                     phi::complex64,                                   \
+                     phi::complex128) {                                \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);     \
+  }
+
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(sqrt, Sqrt)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(relu, Relu)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(pow, Pow)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(scale, Scale)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(relu6, Relu6)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(leaky_relu, LeakyRelu)
+
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(asin, Asin)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(asinh, Asinh)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(atanh, Atanh)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(expm1, Expm1)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(log1p, Log1p)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(square, Square)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(tanh, Tanh)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(sinh, Sinh)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(tan, Tan)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(sin, Sin)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(abs, Abs)
+PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(atan, Atan)
+
+PD_REGISTER_KERNEL(divide_scalar_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::DivScalarCooKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(divide_scalar_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::DivScalarCsrKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(cast_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CastCooKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(cast_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CastCsrKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(isnan_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::IsnanCooKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   int,
+                   int64_t) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(isnan_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::IsnanCsrKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   int,
+                   int64_t) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}