chore: import upstream snapshot with attribution

paddle/phi/kernels/sparse/addmm_grad_kernel.h

@@ -0,0 +1,77 @@
+/* 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 "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+
+namespace phi {
+namespace sparse {
+
+// TODO(zhouwei25): implement Backward of " COO + COO @ COO -> COO"
+template <typename T, typename Context>
+void AddmmCooCooGradKernel(const Context& dev_ctx,
+                           const SparseCooTensor& input,
+                           const SparseCooTensor& x,
+                           const SparseCooTensor& y,
+                           const SparseCooTensor& dout,
+                           float alpha,
+                           float beta,
+                           SparseCooTensor* dinput,
+                           SparseCooTensor* dx,
+                           SparseCooTensor* dy);
+
+// Backward of "DENSE + COO @ DENSE -> DENSE"
+template <typename T, typename Context>
+void AddmmCooDenseGradKernel(const Context& dev_ctx,
+                             const DenseTensor& input,
+                             const SparseCooTensor& x,
+                             const DenseTensor& y,
+                             const DenseTensor& dout,
+                             float alpha,
+                             float beta,
+                             DenseTensor* dinput,
+                             SparseCooTensor* dx,
+                             DenseTensor* dy);
+
+// TODO(zhouwei25): implement Backward of " CSR + CSR @ CSR -> CSR"
+template <typename T, typename Context>
+void AddmmCsrCsrGradKernel(const Context& dev_ctx,
+                           const SparseCsrTensor& input,
+                           const SparseCsrTensor& x,
+                           const SparseCsrTensor& y,
+                           const SparseCsrTensor& dout,
+                           float alpha,
+                           float beta,
+                           SparseCsrTensor* dinput,
+                           SparseCsrTensor* dx,
+                           SparseCsrTensor* dy);
+
+/* Backward of "DENSE + CSR @ DENSE -> DENSE" */
+template <typename T, typename Context>
+void AddmmCsrDenseGradKernel(const Context& dev_ctx,
+                             const DenseTensor& input,
+                             const SparseCsrTensor& x,
+                             const DenseTensor& y,
+                             const DenseTensor& dout,
+                             float alpha,
+                             float beta,
+                             DenseTensor* dinput,
+                             SparseCsrTensor* dx,
+                             DenseTensor* dy);
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/addmm_kernel.h

@@ -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. */
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+
+namespace phi {
+namespace sparse {
+
+// TODO(zhouwei25): implement " COO + COO @ COO -> COO"
+template <typename T, typename Context>
+void AddmmCooCooKernel(const Context& dev_ctx,
+                       const SparseCooTensor& input,
+                       const SparseCooTensor& x,
+                       const SparseCooTensor& y,
+                       float beta,
+                       float alpha,
+                       SparseCooTensor* out);
+
+/* DENSE + COO @ DENSE -> DENSE */
+template <typename T, typename Context>
+void AddmmCooDenseKernel(const Context& dev_ctx,
+                         const DenseTensor& input,
+                         const SparseCooTensor& x,
+                         const DenseTensor& y,
+                         float beta,
+                         float alpha,
+                         DenseTensor* out);
+
+// TODO(zhouwei25): implement " CSR + CSR @ CSR -> CSR"
+template <typename T, typename Context>
+void AddmmCsrCsrKernel(const Context& dev_ctx,
+                       const SparseCsrTensor& input,
+                       const SparseCsrTensor& x,
+                       const SparseCsrTensor& y,
+                       float beta,
+                       float alpha,
+                       SparseCsrTensor* out);
+
+/* DENSE + CSR @ DENSE -> DENSE */
+template <typename T, typename Context>
+void AddmmCsrDenseKernel(const Context& dev_ctx,
+                         const DenseTensor& input,
+                         const SparseCsrTensor& x,
+                         const DenseTensor& y,
+                         float beta,
+                         float alpha,
+                         DenseTensor* out);
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/batch_norm_grad_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/batch_norm_grad_kernel.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/batch_norm_grad_kernel.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void BatchNormCooGradKernel(const Context& dev_ctx,
+                            const SparseCooTensor& x,
+                            const DenseTensor& scale,
+                            const DenseTensor& bias,
+                            const optional<DenseTensor>& mean,
+                            const optional<DenseTensor>& variance,
+                            const DenseTensor& saved_mean,
+                            const DenseTensor& saved_variance,
+                            const optional<DenseTensor>& reserve_space,
+                            const SparseCooTensor& y_grad,
+                            float momentum,
+                            float epsilon,
+                            const std::string& data_layout,
+                            bool is_test,
+                            bool use_global_stats,
+                            bool trainable_statistics,
+                            SparseCooTensor* x_grad,
+                            DenseTensor* scale_grad,
+                            DenseTensor* bias_grad) {
+  EmptyLikeCooKernel<T, Context>(dev_ctx, x, x_grad);
+
+  // TODO(umiswing): add check for parameter freezing automatically
+  PADDLE_ENFORCE_EQ((scale_grad == nullptr && bias_grad == nullptr) ||
+                        (scale_grad != nullptr && bias_grad != nullptr),
+                    true,
+                    common::errors::InvalidArgument(
+                        "Weight and bias's stop_gradient of BatchNorm must be "
+                        "True or False at the same time."));
+
+  if (scale_grad && bias_grad) {
+    *scale_grad = EmptyLike<T, Context>(dev_ctx, scale);
+    *bias_grad = EmptyLike<T, Context>(dev_ctx, bias);
+  }
+  phi::BatchNormGradKernel<T, Context>(dev_ctx,
+                                       x.values(),
+                                       scale,
+                                       bias,
+                                       mean,
+                                       variance,
+                                       saved_mean,
+                                       saved_variance,
+                                       reserve_space,
+                                       y_grad.values(),
+                                       momentum,
+                                       epsilon,
+                                       data_layout,
+                                       is_test,
+                                       use_global_stats,
+                                       trainable_statistics,
+                                       x_grad->mutable_values(),
+                                       scale_grad,
+                                       bias_grad);
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(batch_norm_coo_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::BatchNormCooGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+#if defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(batch_norm_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::BatchNormCooGradKernel,
+                   float,
+                   phi::float16) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+#endif
+
+#if defined(PADDLE_WITH_CUDA)
+PD_REGISTER_KERNEL(batch_norm_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::BatchNormCooGradKernel,
+                   float,
+                   double,
+                   phi::float16) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  if (kernel_key.dtype() == phi::DataType::FLOAT16) {
+    kernel->OutputAt(0).SetDataType(phi::DataType::FLOAT32);  // x_grad
+    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);  // scale_grad
+    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);  // bias_grad
+  }
+}
+#endif

paddle/phi/kernels/sparse/batch_norm_grad_kernel.h

@@ -0,0 +1,47 @@
+/* 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 <string>
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context>
+void BatchNormCooGradKernel(const Context& dev_ctx,
+                            const SparseCooTensor& x,
+                            const DenseTensor& scale,
+                            const DenseTensor& bias,
+                            const optional<DenseTensor>& mean,
+                            const optional<DenseTensor>& variance,
+                            const DenseTensor& saved_mean,
+                            const DenseTensor& saved_variance,
+                            const optional<DenseTensor>& reserve_space,
+                            const SparseCooTensor& y_grad,
+                            float momentum,
+                            float epsilon,
+                            const std::string& data_layout,
+                            bool is_test,
+                            bool use_global_stats,
+                            bool trainable_statistics,
+                            SparseCooTensor* x_grad,
+                            DenseTensor* scale_grad,
+                            DenseTensor* bias_grad);
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/batch_norm_kernel.cc

@@ -0,0 +1,114 @@
+/* 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/batch_norm_kernel.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/batch_norm_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void BatchNormCooKernel(const Context& dev_ctx,
+                        const SparseCooTensor& x,
+                        const DenseTensor& mean,
+                        const DenseTensor& variance,
+                        const DenseTensor& scale,
+                        const DenseTensor& bias,
+                        bool is_test,
+                        float momentum,
+                        float epsilon,
+                        const std::string& data_layout,
+                        bool use_global_stats,
+                        bool trainable_statistics,
+                        SparseCooTensor* y,
+                        DenseTensor* mean_out,
+                        DenseTensor* variance_out,
+                        DenseTensor* saved_mean,
+                        DenseTensor* saved_variance,
+                        DenseTensor* reserve_space) {
+  EmptyLikeCooKernel<T, Context>(dev_ctx, x, y);
+  phi::BatchNormKernel<T, Context>(dev_ctx,
+                                   x.values(),
+                                   mean,
+                                   variance,
+                                   scale,
+                                   bias,
+                                   is_test,
+                                   momentum,
+                                   epsilon,
+                                   data_layout,
+                                   use_global_stats,
+                                   trainable_statistics,
+                                   y->mutable_values(),
+                                   mean_out,
+                                   variance_out,
+                                   saved_mean,
+                                   saved_variance,
+                                   reserve_space);
+  y->SetIndicesDict(x.GetIndicesDict());
+  y->SetKmaps(x.GetKmaps());
+}
+
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(batch_norm_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::BatchNormCooKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+#if defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(batch_norm_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::BatchNormCooKernel,
+                   float,
+                   phi::float16) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(1).SetDataType(phi::DataType::FLOAT32);
+  kernel->InputAt(2).SetDataType(phi::DataType::FLOAT32);
+  kernel->InputAt(3).SetDataType(phi::DataType::FLOAT32);
+  kernel->InputAt(4).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
+}
+#endif
+
+#if defined(PADDLE_WITH_CUDA)
+PD_REGISTER_KERNEL(batch_norm_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::BatchNormCooKernel,
+                   float,
+                   double,
+                   phi::float16) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  if (kernel_key.dtype() == phi::DataType::FLOAT16) {
+    kernel->InputAt(1).SetDataType(phi::DataType::FLOAT32);
+    kernel->InputAt(2).SetDataType(phi::DataType::FLOAT32);
+    kernel->InputAt(3).SetDataType(phi::DataType::FLOAT32);
+    kernel->InputAt(4).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
+  }
+}
+#endif

paddle/phi/kernels/sparse/batch_norm_kernel.h

@@ -0,0 +1,46 @@
+/* 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 <string>
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context>
+void BatchNormCooKernel(const Context& dev_ctx,
+                        const SparseCooTensor& x,
+                        const DenseTensor& mean,
+                        const DenseTensor& variance,
+                        const DenseTensor& scale,
+                        const DenseTensor& bias,
+                        bool is_test,
+                        float momentum,
+                        float epsilon,
+                        const std::string& data_layout,
+                        bool use_global_stats,
+                        bool trainable_statistics,
+                        SparseCooTensor* y,
+                        DenseTensor* mean_out,
+                        DenseTensor* variance_out,
+                        DenseTensor* saved_mean,
+                        DenseTensor* saved_variance,
+                        DenseTensor* reserve_space);
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/coalesce_kernel.h

@@ -0,0 +1,37 @@
+/* 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 "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context>
+void CoalesceCooKernel(const Context& dev_ctx,
+                       const SparseCooTensor& x,
+                       SparseCooTensor* out);
+
+template <typename T, typename Context>
+SparseCooTensor CoalesceCoo(const Context& dev_ctx, const SparseCooTensor& x) {
+  SparseCooTensor coo;
+  CoalesceCooKernel<T, Context>(dev_ctx, x, &coo);
+  return coo;
+}
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/conv_grad_kernel.h

@@ -0,0 +1,79 @@
+/* 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 "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+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);
+
+template <typename T, typename Context>
+std::tuple<SparseCooTensor, DenseTensor> Conv3dCooGrad(
+    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;
+
+  // TODO(zhangkaihuo): call InferMeta func here
+  Conv3dCooGradKernel<T, Context>(dev_ctx,
+                                  x,
+                                  kernel,
+                                  out,
+                                  rulebook,
+                                  counter,
+                                  out_grad,
+                                  paddings,
+                                  dilations,
+                                  strides,
+                                  groups,
+                                  subm,
+                                  key,
+                                  &x_grad,
+                                  &kernel_grad);
+  return std::make_tuple(x_grad, kernel_grad);
+}
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/conv_kernel.h

@@ -0,0 +1,67 @@
+/* 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 "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/sparse/convolution.h"
+
+namespace phi {
+namespace sparse {
+
+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);
+
+template <typename T, typename Context>
+SparseCooTensor Conv3dCoo(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,
+                          DenseTensor* rulebook,
+                          DenseTensor* counter) {
+  SparseCooTensor coo;
+  Conv3dCooKernel<T, Context>(dev_ctx,
+                              x,
+                              kernel,
+                              paddings,
+                              dilations,
+                              strides,
+                              groups,
+                              subm,
+                              key,
+                              &coo,
+                              rulebook,
+                              counter);
+  return coo;
+}
+}  // namespace sparse
+}  // namespace phi

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);
+}

paddle/phi/kernels/sparse/elementwise_grad_kernel.h

@@ -0,0 +1,148 @@
+/* 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 "paddle/phi/kernels/elementwise_add_grad_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+#include "paddle/phi/infermeta/sparse/unary.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+#define DEFINE_ELEMENTWISE_GRAD_KERNEL_HEAD(name)          \
+  DEFINE_ELEMENTWISE_GRAD_KERNEL_HEAD_WITH_TYPE(name, Csr) \
+                                                           \
+  DEFINE_ELEMENTWISE_GRAD_KERNEL_HEAD_WITH_TYPE(name, Coo)
+
+#define DEFINE_ELEMENTWISE_GRAD_KERNEL_FUNC(name)          \
+  DEFINE_ELEMENTWISE_GRAD_KERNEL_FUNC_WITH_TYPE(name, Csr) \
+                                                           \
+  DEFINE_ELEMENTWISE_GRAD_KERNEL_FUNC_WITH_TYPE(name, Coo)
+
+#define DEFINE_ELEMENTWISE_GRAD_KERNEL_HEAD_WITH_TYPE(name, type)            \
+  template <typename T, typename Context>                                    \
+  void ElementWise##name##type##GradKernel(const Context& dev_ctx,           \
+                                           const Sparse##type##Tensor& x,    \
+                                           const Sparse##type##Tensor& y,    \
+                                           const Sparse##type##Tensor& dout, \
+                                           Sparse##type##Tensor* dx,         \
+                                           Sparse##type##Tensor* dy);
+
+#define DEFINE_ELEMENTWISE_GRAD_KERNEL_FUNC_WITH_TYPE(name, type)  \
+  template <typename T, typename Context>                          \
+  std::vector<Sparse##type##Tensor> ElementWise##name##type##Grad( \
+      const Context& dev_ctx,                                      \
+      const Sparse##type##Tensor& x,                               \
+      const Sparse##type##Tensor& y,                               \
+      const Sparse##type##Tensor& dout) {                          \
+    Sparse##type##Tensor dx;                                       \
+    Sparse##type##Tensor dy;                                       \
+    MetaTensor meta_dx(&dx), meta_dy(&dy);                         \
+    phi::UnchangedInferMeta(x, &meta_dx);                          \
+    phi::UnchangedInferMeta(y, &meta_dy);                          \
+    ElementWise##name##type##GradKernel<T, Context>(               \
+        dev_ctx, x, y, dout, &dx, &dy);                            \
+    return std::vector<Sparse##type##Tensor>{dx, dy};              \
+  }
+
+DEFINE_ELEMENTWISE_GRAD_KERNEL_HEAD(Add)
+DEFINE_ELEMENTWISE_GRAD_KERNEL_HEAD(Subtract)
+DEFINE_ELEMENTWISE_GRAD_KERNEL_HEAD(Multiply)
+
+DEFINE_ELEMENTWISE_GRAD_KERNEL_FUNC(Add)
+DEFINE_ELEMENTWISE_GRAD_KERNEL_FUNC(Subtract)
+DEFINE_ELEMENTWISE_GRAD_KERNEL_FUNC(Multiply)
+
+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);
+
+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);
+
+template <typename T, typename Context>
+std::vector<SparseCsrTensor> ElementWiseDivideCsrGrad(
+    const Context& dev_ctx,
+    const SparseCsrTensor& x,
+    const SparseCsrTensor& y,
+    const SparseCsrTensor& out,
+    const SparseCsrTensor& dout) {
+  SparseCsrTensor dx;
+  SparseCsrTensor dy;
+  MetaTensor meta_dx(&dx), meta_dy(&dy);
+  phi::UnchangedInferMeta(x, &meta_dx);
+  phi::UnchangedInferMeta(y, &meta_dy);
+  ElementWiseDivideCsrGradKernel<T, Context>(
+      dev_ctx, x, y, out, dout, &dx, &dy);
+  return std::vector<SparseCsrTensor>{dx, dy};
+}
+
+template <typename T, typename Context>
+std::vector<SparseCooTensor> ElementWiseDivideCooGrad(
+    const Context& dev_ctx,
+    const SparseCooTensor& x,
+    const SparseCooTensor& y,
+    const SparseCooTensor& out,
+    const SparseCooTensor& dout) {
+  SparseCooTensor dx;
+  SparseCooTensor dy;
+  MetaTensor meta_dx(&dx), meta_dy(&dy);
+  phi::UnchangedInferMeta(x, &meta_dx);
+  phi::UnchangedInferMeta(y, &meta_dy);
+  ElementWiseDivideCooGradKernel<T, Context>(
+      dev_ctx, x, y, out, dout, &dx, &dy);
+  return std::vector<SparseCooTensor>{dx, dy};
+}
+
+template <typename T, typename Context>
+void ElementWiseAddDenseGradKernel(const Context& dev_ctx,
+                                   const SparseCooTensor& x,
+                                   const DenseTensor& y,
+                                   const SparseCooTensor& dout,
+                                   SparseCooTensor* dx,
+                                   DenseTensor* dy) {
+  DenseTensor* x_values_grad = nullptr;
+  DenseTensor* y_grad = nullptr;
+  if (dx) {
+    EmptyLikeCooKernel<T, Context>(dev_ctx, x, dx);
+    x_values_grad = dx->mutable_values();
+  }
+
+  if (dy) {
+    *dy = EmptyLike<T>(dev_ctx, y);
+    y_grad = dy;
+  }
+  phi::AddGradKernel<T, Context>(
+      dev_ctx, x.values(), y, dout.values(), -1, x_values_grad, y_grad);
+}
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/elementwise_kernel.h

@@ -0,0 +1,102 @@
+/* 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 "paddle/phi/kernels/elementwise_add_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+#include "paddle/phi/infermeta/binary.h"
+
+namespace phi {
+namespace sparse {
+
+#define DEFINE_ELEMENTWISE_KERNEL_HEAD(name)          \
+  DEFINE_ELEMENTWISE_KERNEL_HEAD_WITH_TYPE(name, Csr) \
+                                                      \
+  DEFINE_ELEMENTWISE_KERNEL_HEAD_WITH_TYPE(name, Coo)
+
+#define DEFINE_ELEMENTWISE_KERNEL_FUNC(name) \
+  DEFINE_CSR_ELEMENTWISE_KERNEL_FUNC(name)   \
+                                             \
+  DEFINE_COO_ELEMENTWISE_KERNEL_FUNC(name)
+
+#define DEFINE_ELEMENTWISE_KERNEL_HEAD_WITH_TYPE(name, type)          \
+  template <typename T, typename Context>                             \
+  void ElementWise##name##type##Kernel(const Context& dev_ctx,        \
+                                       const Sparse##type##Tensor& x, \
+                                       const Sparse##type##Tensor& y, \
+                                       Sparse##type##Tensor* out);
+
+#define DEFINE_CSR_ELEMENTWISE_KERNEL_FUNC(name)                     \
+  template <typename T, typename Context>                            \
+  SparseCsrTensor ElementWise##name##Csr(const Context& dev_ctx,     \
+                                         const SparseCsrTensor& x,   \
+                                         const SparseCsrTensor& y) { \
+    DenseTensor crows;                                               \
+    DenseTensor cols;                                                \
+    DenseTensor values;                                              \
+    SparseCsrTensor out(crows, cols, values, x.dims());              \
+    MetaTensor meta_out(out);                                        \
+    phi::ElementwiseInferMeta(x, y, &meta_out);                      \
+    ElementWise##name##CsrKernel<T, Context>(dev_ctx, x, y, &out);   \
+    return out;                                                      \
+  }
+
+#define DEFINE_COO_ELEMENTWISE_KERNEL_FUNC(name)                     \
+  template <typename T, typename Context>                            \
+  SparseCooTensor ElementWise##name##Coo(const Context& dev_ctx,     \
+                                         const SparseCooTensor& x,   \
+                                         const SparseCooTensor& y) { \
+    DenseTensor indices;                                             \
+    DenseTensor values;                                              \
+    SparseCooTensor out(indices, values, x.dims());                  \
+    MetaTensor meta_out(out);                                        \
+    phi::ElementwiseInferMeta(x, y, &meta_out);                      \
+    ElementWise##name##CooKernel<T, Context>(dev_ctx, x, y, &out);   \
+    return out;                                                      \
+  }
+
+DEFINE_ELEMENTWISE_KERNEL_HEAD(Add)
+DEFINE_ELEMENTWISE_KERNEL_HEAD(Subtract)
+DEFINE_ELEMENTWISE_KERNEL_HEAD(Multiply)
+DEFINE_ELEMENTWISE_KERNEL_HEAD(Divide)
+
+DEFINE_ELEMENTWISE_KERNEL_FUNC(Add)
+DEFINE_ELEMENTWISE_KERNEL_FUNC(Subtract)
+DEFINE_ELEMENTWISE_KERNEL_FUNC(Multiply)
+DEFINE_ELEMENTWISE_KERNEL_FUNC(Divide)
+
+template <typename T, typename Context>
+void ElementWiseAddDenseKernel(const Context& dev_ctx,
+                               const SparseCooTensor& x,
+                               const DenseTensor& y,
+                               SparseCooTensor* out) {
+  // TODO(zhangkaiuo): to support universal sparse + dense
+  if (y.dims().size() == 1 && y.dims()[0] == x.dims()[x.dims().size() - 1]) {
+    EmptyLikeCooKernel<T, Context>(dev_ctx, x, out);
+    phi::AddKernel<T, Context>(dev_ctx, x.values(), y, out->mutable_values());
+    out->SetIndicesDict(x.GetIndicesDict());
+    out->SetKmaps(x.GetKmaps());
+  } else {
+    PADDLE_THROW(
+        errors::Unimplemented("Not support Sparse + Dense in GPU mode"));
+  }
+}
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/empty_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/empty_kernel.h"
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+
+namespace phi::sparse {
+
+template <typename T, typename Context>
+void EmptyLikeCooKernel(const Context& dev_ctx,
+                        const SparseCooTensor& x,
+                        SparseCooTensor* out) {
+  *(out->mutable_indices()) = x.indices();
+
+  const DenseTensor& x_values = x.values();
+  DenseTensor* out_values = out->mutable_values();
+  out_values->Resize(x_values.dims());
+  out->set_meta(x.meta());
+  dev_ctx.template Alloc<T>(out_values);
+}
+
+template <typename T, typename Context>
+void EmptyLikeCsrKernel(const Context& dev_ctx,
+                        const SparseCsrTensor& x,
+                        SparseCsrTensor* out) {
+  *(out->mutable_crows()) = x.crows();
+  *(out->mutable_cols()) = x.cols();
+
+  const DenseTensor& x_values = x.values();
+  DenseTensor* out_values = out->mutable_values();
+  out_values->Resize(x_values.dims());
+  out->set_meta(x.meta());
+  dev_ctx.template Alloc<T>(out_values);
+}
+}  // namespace phi::sparse
+
+PD_REGISTER_KERNEL(empty_like_coo,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::EmptyLikeCooKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(empty_like_csr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sparse::EmptyLikeCsrKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(empty_like_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::EmptyLikeCooKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(empty_like_csr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::EmptyLikeCsrKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+#endif

paddle/phi/kernels/sparse/empty_kernel.h

@@ -0,0 +1,34 @@
+/* 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 "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context>
+void EmptyLikeCooKernel(const Context& dev_ctx,
+                        const SparseCooTensor& x,
+                        SparseCooTensor* out);
+
+template <typename T, typename Context>
+void EmptyLikeCsrKernel(const Context& dev_ctx,
+                        const SparseCsrTensor& x,
+                        SparseCsrTensor* out);
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/full_kernel.h

@@ -0,0 +1,38 @@
+// 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 "paddle/phi/common/data_type.h"
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void FullLikeCooKernel(const Context& dev_ctx,
+                       const SparseCooTensor& x,
+                       const Scalar& val,
+                       DataType dtype,
+                       SparseCooTensor* out);
+
+template <typename T, typename Context>
+void FullLikeCsrKernel(const Context& dev_ctx,
+                       const SparseCsrTensor& x,
+                       const Scalar& val,
+                       DataType dtype,
+                       SparseCsrTensor* out);
+
+}  // namespace phi

paddle/phi/kernels/sparse/fused_attention_grad_kernel.h

@@ -0,0 +1,35 @@
+/* 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 "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+
+namespace phi {
+namespace 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);
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/fused_attention_kernel.h

@@ -0,0 +1,35 @@
+/* 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 "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+
+namespace phi {
+namespace 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);
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/gpu/addmm_grad_kernel.cu

@@ -0,0 +1,96 @@
+/* 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/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/sparse/matmul_grad_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context>
+void AddmmCooDenseGradKernel(const Context& dev_ctx,
+                             const DenseTensor& input,
+                             const SparseCooTensor& x,
+                             const DenseTensor& y,
+                             const DenseTensor& dout,
+                             float alpha,
+                             float beta,
+                             DenseTensor* dinput,
+                             SparseCooTensor* dx,
+                             DenseTensor* dy) {
+  auto blas = funcs::GetBlas<Context, T>(dev_ctx);
+  if (dinput) {
+    dinput->Resize(input.dims());
+    dev_ctx.template Alloc<T>(dinput);
+
+    blas.VCOPY(input.numel(), dout.data<T>(), dinput->data<T>());
+    blas.SCAL(input.numel(), beta, dinput->data<T>());
+  }
+  DenseTensor dout_scale = EmptyLike<T, Context>(dev_ctx, dout);
+  blas.VCOPY(dout.numel(), dout.data<T>(), dout_scale.data<T>());
+  blas.SCAL(dout.numel(), alpha, dout_scale.data<T>());
+  MatmulCooDenseGradKernel<T, Context>(dev_ctx, x, y, dout_scale, dx, dy);
+}
+
+// Backward of "DENSE + CSR @ DENSE -> DENSE"
+template <typename T, typename Context>
+void AddmmCsrDenseGradKernel(const Context& dev_ctx,
+                             const DenseTensor& input,
+                             const SparseCsrTensor& x,
+                             const DenseTensor& y,
+                             const DenseTensor& dout,
+                             float alpha,
+                             float beta,
+                             DenseTensor* dinput,
+                             SparseCsrTensor* dx,
+                             DenseTensor* dy) {
+  auto blas = funcs::GetBlas<Context, T>(dev_ctx);
+  if (dinput) {
+    dinput->Resize(input.dims());
+    dev_ctx.template Alloc<T>(dinput);
+
+    blas.VCOPY(input.numel(), dout.data<T>(), dinput->data<T>());
+    blas.SCAL(input.numel(), beta, dinput->data<T>());
+  }
+  DenseTensor dout_scale = EmptyLike<T, Context>(dev_ctx, dout);
+  blas.VCOPY(dout.numel(), dout.data<T>(), dout_scale.data<T>());
+  blas.SCAL(dout.numel(), alpha, dout_scale.data<T>());
+  MatmulCsrDenseGradKernel<T, Context>(dev_ctx, x, y, dout_scale, dx, dy);
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(addmm_coo_dense_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::AddmmCooDenseGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(addmm_csr_dense_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::AddmmCsrDenseGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}

paddle/phi/kernels/sparse/gpu/addmm_kernel.cu

@@ -0,0 +1,141 @@
+/* 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/common/ddim.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/sparse/sparse_blas.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context, typename TensorType>
+void AddmmKernelImpl(const Context& dev_ctx,
+                     const DenseTensor& input,
+                     const TensorType& x,
+                     const DenseTensor& y,
+                     float beta,
+                     float alpha,
+                     DenseTensor* out) {
+#if defined(PADDLE_WITH_CUDA)
+  std::vector<int64_t> input_dim = vectorize(input.dims());
+  std::vector<int64_t> x_dim = vectorize(x.dims());
+  std::vector<int64_t> y_dim = vectorize(y.dims());
+  auto rank = input_dim.size();
+
+  PADDLE_ENFORCE_GE(
+      rank,
+      2,
+      common::errors::InvalidArgument(
+          "the dims size of input must be greater than or equal to 2."));
+
+  PADDLE_ENFORCE_EQ(
+      x_dim.size(),
+      rank,
+      common::errors::PreconditionNotMet(
+          "The dims size of Input(input) and Input(x) must be equal."));
+
+  PADDLE_ENFORCE_GE(
+      y_dim.size(),
+      rank,
+      common::errors::InvalidArgument(
+          "the dims size of Input(input) and Input(y) must be equal."));
+
+  for (size_t i = 0; i < rank - 2; ++i) {
+    PADDLE_ENFORCE_EQ(input_dim[i],
+                      x_dim[i],
+                      common::errors::InvalidArgument(
+                          "input.dim[%d] and x.dim[%d] must be equal.", i, i));
+    PADDLE_ENFORCE_EQ(input_dim[i],
+                      y_dim[i],
+                      common::errors::InvalidArgument(
+                          "input.dim[%d] and y.dim[%d] must be equal.", i, i));
+  }
+
+  PADDLE_ENFORCE_GE(
+      input_dim[rank - 2],
+      x_dim[rank - 2],
+      common::errors::PreconditionNotMet(
+          "The shape of Input(input) and Input(x) is not suitable for matmul "
+          "operation, input_dim[-2] must be equal to x_dim[-2]."));
+
+  PADDLE_ENFORCE_GE(
+      input_dim[rank - 1],
+      y_dim[rank - 1],
+      common::errors::PreconditionNotMet(
+          "The shape of Input(input) and Input(y) is not suitable for matmul "
+          "operation, input_dim[-1] must be equal to y_dim[-1]."));
+
+  PADDLE_ENFORCE_GE(
+      x_dim[rank - 1],
+      y_dim[rank - 2],
+      common::errors::PreconditionNotMet(
+          "The shape of Input(x) and Input(y) is not suitable for matmul "
+          "operation, x_dim[-1] must be equal to y_dim[-2]."));
+
+  phi::Copy(dev_ctx, input, dev_ctx.GetPlace(), false, out);
+
+  auto sparse_blas = funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
+  sparse_blas.SPMM(
+      false, false, static_cast<T>(alpha), x, y, static_cast<T>(beta), out);
+#endif
+}
+
+template <typename T, typename Context>
+void AddmmCooDenseKernel(const Context& dev_ctx,
+                         const DenseTensor& input,
+                         const SparseCooTensor& x,
+                         const DenseTensor& y,
+                         float beta,
+                         float alpha,
+                         DenseTensor* out) {
+  AddmmKernelImpl<T>(dev_ctx, input, x, y, beta, alpha, out);
+}
+
+template <typename T, typename Context>
+void AddmmCsrDenseKernel(const Context& dev_ctx,
+                         const DenseTensor& input,
+                         const SparseCsrTensor& x,
+                         const DenseTensor& y,
+                         float beta,
+                         float alpha,
+                         DenseTensor* out) {
+  AddmmKernelImpl<T>(dev_ctx, input, x, y, beta, alpha, out);
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(addmm_coo_dense,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::AddmmCooDenseKernel,
+                   float,
+                   double,
+                   phi::float16) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(addmm_csr_dense,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::AddmmCsrDenseKernel,
+                   float,
+                   double,
+                   phi::float16) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}

paddle/phi/kernels/sparse/gpu/coalesce_kernel.cu

@@ -0,0 +1,203 @@
+/* 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 <thrust/sort.h>
+#include <thrust/unique.h>
+
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/funcs/index_impl.cu.h"
+#include "paddle/phi/kernels/funcs/sparse/flatten_indices.cu.h"
+#include "paddle/phi/kernels/funcs/sparse/scatter.cu.h"
+#include "paddle/phi/kernels/funcs/sparse/utils.cu.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename IntT>
+void CoalesceCooGPUKernel(const GPUContext& 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 nnz = x.nnz();
+  const int64_t sparse_dim = x.indices().dims()[0];
+  std::vector<IntT> sparse_offsets(sparse_dim);
+
+  funcs::sparse::CalcOffsetsPerDim<IntT>(
+      x.dims(), sparse_dim, sparse_offsets.data());
+
+  DenseTensorMeta sparse_offset_meta(
+      phi::CppTypeToDataType<IntT>::Type(), {sparse_dim}, DataLayout::NCHW);
+  DenseTensor d_sparse_offsets =
+      Empty<GPUContext>(dev_ctx, std::move(sparse_offset_meta));
+  DenseTensor indices = Empty(
+      dev_ctx, DenseTensorMeta(x_indices.dtype(), {nnz}, x_indices.layout()));
+  IntT* indices_ptr = indices.data<IntT>();
+
+  backends::gpu::GpuMemcpyAsync(d_sparse_offsets.data<IntT>(),
+                                sparse_offsets.data(),
+                                sizeof(IntT) * sparse_dim,
+                                gpuMemcpyHostToDevice,
+                                dev_ctx.stream());
+
+  // 1. flatten indices
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, nnz, 1);
+  funcs::sparse::FlattenIndicesKernel<<<config.block_per_grid,
+                                        config.thread_per_block,
+                                        0,
+                                        dev_ctx.stream()>>>(
+      x.indices().data<IntT>(),
+      d_sparse_offsets.data<IntT>(),
+      indices.numel(),
+      sparse_dim,
+      indices_ptr);
+
+  // 2. get the address of each non-zero values
+  const T* x_values_ptr = x_values.data<T>();
+  const int64_t stride =
+      x.dims().size() == sparse_dim ? 1 : x.values().dims()[1];
+  DenseTensor values_indices =
+      Empty(dev_ctx, DenseTensorMeta(DataType::INT32, {nnz}, DataLayout::NCHW));
+  int* values_indices_ptr = values_indices.data<int>();
+  DenseTensor public_indices = EmptyLike<int>(dev_ctx, values_indices);
+
+  // values_indices = [0,1,2,,,nnz-1]
+  phi::IndexKernel<int, kps::IdentityFunctor<int>>(
+      dev_ctx, &values_indices, kps::IdentityFunctor<int>());
+  phi::IndexKernel<int, kps::IdentityFunctor<int>>(
+      dev_ctx, &public_indices, kps::IdentityFunctor<int>());
+
+// 3. sort (indices, values index)
+#ifdef PADDLE_WITH_HIP
+  thrust::sort_by_key(thrust::hip::par.on(dev_ctx.stream()),
+#else
+  thrust::sort_by_key(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                      indices_ptr,
+                      indices_ptr + nnz,
+                      values_indices_ptr);
+
+  // 4. unique index
+  thrust::pair<IntT*, int*> new_end =
+#ifdef PADDLE_WITH_HIP
+      thrust::unique_by_key(thrust::hip::par.on(dev_ctx.stream()),
+#else
+      thrust::unique_by_key(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                            indices_ptr,
+                            indices_ptr + nnz,
+                            public_indices.data<int>());
+
+  funcs::sparse::DistanceKernel<<<1, 1, 0, dev_ctx.stream()>>>(
+      indices_ptr, new_end.first, out_indices.data<IntT>());
+
+  IntT out_nnz = 0;
+  backends::gpu::GpuMemcpyAsync(&out_nnz,
+                                out_indices.data<IntT>(),
+                                sizeof(IntT),
+                                gpuMemcpyDeviceToHost,
+                                dev_ctx.stream());
+  dev_ctx.Wait();
+
+  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]});
+  }
+
+  // 5. scatter the values
+  const int VecSize = VecBytes / sizeof(T);
+  if (stride % VecSize == 0) {
+    config =
+        backends::gpu::GetGpuLaunchConfig1D(dev_ctx, nnz * stride / VecSize, 1);
+    funcs::sparse::ScatterKernel<T, VecSize>
+        <<<config.block_per_grid,
+           config.thread_per_block,
+           0,
+           dev_ctx.stream()>>>(x_values_ptr,
+                               public_indices.data<int>(),
+                               values_indices_ptr,
+                               out_nnz,
+                               nnz,
+                               stride,
+                               out_values.data<T>());
+  } else {
+    config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, nnz * stride, 1);
+    funcs::sparse::ScatterKernel<T, 1>
+        <<<config.block_per_grid,
+           config.thread_per_block,
+           0,
+           dev_ctx.stream()>>>(x_values_ptr,
+                               public_indices.data<int>(),
+                               values_indices_ptr,
+                               out_nnz,
+                               nnz,
+                               stride,
+                               out_values.data<T>());
+  }
+
+  // 6. convert index to coordinate
+  Dim<DDim::kMaxRank> const_dims;
+  for (int i = 0; i < x.dims().size(); i++) {
+    const_dims[i] = x.dims()[i];
+  }
+
+  config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_nnz, 1);
+  funcs::sparse::IndexToCoordinateKernel<<<config.block_per_grid,
+                                           config.thread_per_block,
+                                           0,
+                                           dev_ctx.stream()>>>(
+      indices_ptr, const_dims, out_nnz, sparse_dim, out_indices.data<IntT>());
+
+  out->SetMember(out_indices, out_values, x.dims(), true);
+  out->SetIndicesDict(x.GetIndicesDict());
+  out->SetKmaps(x.GetKmaps());
+}
+
+template <typename T, typename Context>
+void CoalesceCooKernel(const Context& dev_ctx,
+                       const SparseCooTensor& x,
+                       SparseCooTensor* out) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.indices().dtype(), "CoalesceCooGPUKernel", ([&] {
+        CoalesceCooGPUKernel<T, data_t>(dev_ctx, x, out);
+      }));
+}
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(coalesce_coo,
+                   GPU,
+                   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/gpu/conv.cu.h

@@ -0,0 +1,746 @@
+/* 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 <thrust/remove.h>
+#include <thrust/unique.h>
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/aligned_vector.h"
+#include "paddle/phi/kernels/funcs/cub.h"
+#include "paddle/phi/kernels/funcs/index_impl.cu.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/sparse/scatter.cu.h"
+#include "paddle/phi/kernels/funcs/sparse/utils.cu.h"
+#include "paddle/phi/kernels/primitive/compute_primitives.h"
+#include "paddle/phi/kernels/sparse/conv_kernel.h"
+#include "paddle/phi/kernels/sparse/gpu/conv_host_buffer.h"
+#include "paddle/phi/kernels/sparse/gpu/conv_with_buffer.cu.h"
+
+namespace phi {
+namespace sparse {
+
+using Dims4D = funcs::sparse::Dims4D;
+
+// Vectorize load and store global memory
+// In the scene of 3D point cloud, the slice_size 4,8,16,32,64 are commonly
+// used.
+template <typename T, typename IndexT = int, int VecSize>
+__global__ void GatherKernel(const T* params,
+                             const IndexT* indices,
+                             T* output,
+                             size_t index_size,
+                             size_t slice_size) {
+  CUDA_KERNEL_LOOP_TYPE(i, index_size * slice_size / VecSize, int64_t) {
+    const int vec_slice_size = slice_size / VecSize;
+    int indices_i = i / vec_slice_size;
+    int slice_i = i - indices_i * vec_slice_size;  // offset inside the slice
+    IndexT gather_i = indices[indices_i];
+    int64_t params_i = gather_i * slice_size + slice_i * VecSize;
+    using LoadT = AlignedVector<T, VecSize>;
+    using StoreT = AlignedVector<T, VecSize>;
+    LoadT params_vec;
+    Load<T, VecSize>(params + params_i, &params_vec);
+    Store<T, VecSize>(params_vec, output + i * VecSize);
+  }
+}
+
+// double sparse, seed GroupIndices
+template <typename T, typename IntT, int VecSize>
+__global__ void GatherKernelV2(const T* inputs,
+                               const int* index_counts,
+                               const int* index_groups,
+                               const int non_zero_num,
+                               const int kernel_size,
+                               const int channels,
+                               const int buffer_count,
+                               T* output) {
+  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  const int vec_channels = channels / VecSize;
+  using LoadT = AlignedVector<T, VecSize>;
+  using StoreT = AlignedVector<T, VecSize>;
+  for (int i = tid; i < non_zero_num * vec_channels;
+       i += gridDim.x * blockDim.x) {
+    int indices_i = i / vec_channels;
+    int channels_i = i - indices_i * vec_channels;
+    LoadT in_vec;
+    Load<T, VecSize>(inputs + indices_i * channels + channels_i * VecSize,
+                     &in_vec);
+#pragma unroll
+    for (int it = 0; it < buffer_count; it++) {
+      int len = index_counts[indices_i + it * non_zero_num];
+      const int group_offset = it * kernel_size * non_zero_num;
+#pragma unroll
+      for (int j = 0; j < len; j++) {
+        int out_i = index_groups[indices_i * kernel_size + j + group_offset];
+        Store<T, VecSize>(in_vec,
+                          output + out_i * channels + channels_i * VecSize);
+      }
+    }
+  }
+}
+
+template <typename T, typename IntT>
+inline void Gather(const GPUContext& dev_ctx,
+                   const T* inputs,
+                   const IntT* indices,
+                   const int indices_size,
+                   const int channels,
+                   T* output) {
+  const int VecSize = VecBytes / sizeof(T);
+  if (channels % VecSize == 0) {
+    auto config = backends::gpu::GetGpuLaunchConfig1D(
+        dev_ctx, indices_size * channels / VecSize, 1);
+    GatherKernel<T, IntT, VecSize>
+        <<<config.block_per_grid.x,
+           config.thread_per_block.x,
+           0,
+           dev_ctx.stream()>>>(inputs, indices, output, indices_size, channels);
+  } else {
+    auto config = backends::gpu::GetGpuLaunchConfig1D(
+        dev_ctx, indices_size * channels, 1);
+    GatherKernel<T, IntT, 1>
+        <<<config.block_per_grid.x,
+           config.thread_per_block.x,
+           0,
+           dev_ctx.stream()>>>(inputs, indices, output, indices_size, channels);
+  }
+}
+
+template <typename T, typename IntT>
+inline void GatherV2(const GPUContext& dev_ctx,
+                     const T* inputs,
+                     const int* index_counts,
+                     const int* index_groups,
+                     const int non_zero_num,
+                     const int kernel_size,
+                     const int channels,
+                     const int buffer_count,
+                     T* output) {
+  const int VecSize = VecBytes / sizeof(T);
+  if (channels % VecSize == 0) {
+    auto config = backends::gpu::GetGpuLaunchConfig1D(
+        dev_ctx, non_zero_num * channels / VecSize, 1);
+    GatherKernelV2<T, IntT, VecSize><<<config.block_per_grid.x,
+                                       config.thread_per_block.x,
+                                       0,
+                                       dev_ctx.stream()>>>(inputs,
+                                                           index_counts,
+                                                           index_groups,
+                                                           non_zero_num,
+                                                           kernel_size,
+                                                           channels,
+                                                           buffer_count,
+                                                           output);
+  } else {
+    auto config = backends::gpu::GetGpuLaunchConfig1D(
+        dev_ctx, non_zero_num * channels, 1);
+    GatherKernelV2<T, IntT, 1><<<config.block_per_grid.x,
+                                 config.thread_per_block.x,
+                                 0,
+                                 dev_ctx.stream()>>>(inputs,
+                                                     index_counts,
+                                                     index_groups,
+                                                     non_zero_num,
+                                                     kernel_size,
+                                                     channels,
+                                                     buffer_count,
+                                                     output);
+  }
+}
+
+// unique the out indices in rulebook
+template <typename IntT>
+__global__ void UniqueKernel(const IntT* in_indices,
+                             const int rulebook_len,
+                             int* index_flags,
+                             int* out_indices,
+                             int* nnz) {
+  extern __shared__ int cache[];
+  __shared__ int count, start;
+  if (threadIdx.x == 0) {
+    count = 0;
+    start = 0;
+  }
+  __syncthreads();
+
+  int i = threadIdx.x + blockDim.x * blockIdx.x;
+  if (i < rulebook_len) {
+    // atomicOr only support int
+    int index = static_cast<int>(in_indices[i]);
+    const bool flag = funcs::sparse::SetBits(index, index_flags);
+    if (!flag) {
+      int j = atomicAdd(&count, 1);
+      cache[j] = index;
+    }
+  }
+  __syncthreads();
+
+  if (threadIdx.x == 0) {
+    start = atomicAdd(nnz, count);
+  }
+  __syncthreads();
+  for (int i = threadIdx.x; i < count; i += blockDim.x) {
+    out_indices[start + i] = cache[i];
+  }
+}
+
+template <typename IntT>
+__global__ void GetOutIndexTable1(const IntT* indices,
+                                  const IntT non_zero_num,
+                                  const Dims4D dims,
+                                  int* index_flags,
+                                  const bool is2D,
+                                  int* out_index_table) {
+  CUDA_KERNEL_LOOP_TYPE(i, non_zero_num, int64_t) {
+    IntT batch = indices[i];
+    IntT in_z = is2D ? 0 : indices[i + non_zero_num];
+    IntT in_y =
+        is2D ? indices[i + non_zero_num] : indices[i + 2 * non_zero_num];
+    IntT in_x =
+        is2D ? indices[i + 2 * non_zero_num] : indices[i + 3 * non_zero_num];
+    IntT index = PointToIndex(batch, in_x, in_y, in_z, dims);
+    funcs::sparse::SetBits(index, index_flags);
+    out_index_table[index] = i;
+  }
+}
+
+template <typename IntT>
+__global__ void CopyRuleBook(const int* counters,
+                             const int* offsets,
+                             const IntT* in_rulebook,
+                             const int len,
+                             const int kernel_size,
+                             const int non_zero_num,
+                             IntT* out_rulebook) {
+  int tid = threadIdx.x + blockDim.x * blockIdx.x;
+  extern __shared__ int cache_counters[];
+  int* cache_offsets = cache_counters + kernel_size;
+  for (int i = threadIdx.x; i < kernel_size; i += blockDim.x) {
+    cache_counters[i] = counters[i];
+    cache_offsets[i] = offsets[i];
+  }
+  __syncthreads();
+  for (int i = tid; i < len; i += gridDim.x * blockDim.x) {
+    // get the kernel index
+    int kernel_index = 0;
+    for (; kernel_index < kernel_size - 1; kernel_index++) {
+      if (i >= offsets[kernel_index] && i < offsets[kernel_index + 1]) {
+        break;
+      }
+    }
+    int inner_index = i - offsets[kernel_index];
+    out_rulebook[i] = in_rulebook[kernel_index * non_zero_num + inner_index];
+    out_rulebook[len + i] =
+        in_rulebook[kernel_size * non_zero_num + kernel_index * non_zero_num +
+                    inner_index];
+  }
+}
+
+template <typename T>
+__global__ void ProductSubmRuleBookKernel(const T* x_indices,
+                                          const Dims4D x_dims,
+                                          const Dims4D kernel_dims,
+                                          const Dims4D out_dims,
+                                          const int64_t non_zero_num,
+                                          const Dims4D paddings,
+                                          const Dims4D dilations,
+                                          const Dims4D strides,
+                                          const bool is2D,
+                                          const int* index_flags,
+                                          const int* out_index_table,
+                                          T* rulebook,
+                                          int* counter) {
+  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  const int kernel_size = kernel_dims[3] * kernel_dims[2] * kernel_dims[1];
+  extern __shared__ int counter_buf[];  // kernel_size
+  int* counter_buf2 = counter_buf + kernel_size;
+  int* rulebook_buf = counter_buf + kernel_size * 2;
+
+  const int offset = kernel_size * non_zero_num;
+  for (int i = threadIdx.x; i < kernel_size; i += blockDim.x) {
+    counter_buf[i] = 0;
+  }
+  __syncthreads();
+
+  for (int i = tid; i < non_zero_num; i += gridDim.x * blockDim.x) {
+    int kernel_index = 0;
+    T batch = x_indices[i];
+    T in_z = is2D ? 0 : x_indices[i + non_zero_num];
+    T in_y =
+        is2D ? x_indices[i + non_zero_num] : x_indices[i + 2 * non_zero_num];
+    T in_x = is2D ? x_indices[i + 2 * non_zero_num]
+                  : x_indices[i + 3 * non_zero_num];
+    for (int kz = 0; kz < kernel_dims[1]; kz++) {
+      for (int ky = 0; ky < kernel_dims[2]; ky++) {
+        for (int kx = 0; kx < kernel_dims[3]; kx++) {
+          int in_i = -1, out_index = -1, kernel_i = -1;
+          if (funcs::sparse::Check(x_dims,
+                                   kernel_dims,
+                                   paddings,
+                                   dilations,
+                                   strides,
+                                   in_x,
+                                   in_y,
+                                   in_z,
+                                   kx,
+                                   ky,
+                                   kz)) {
+            T out_z =
+                is2D ? 0
+                     : (in_z + paddings[1] - kz * dilations[1]) / strides[1];
+            T out_y = (in_y + paddings[2] - ky * dilations[2]) / strides[2];
+            T out_x = (in_x + paddings[3] - kx * dilations[3]) / strides[3];
+            out_index = funcs::sparse::PointToIndex<Dims4D>(
+                batch, out_x, out_y, out_z, out_dims);
+            const bool flag = funcs::sparse::TestBits(out_index, index_flags);
+            if (flag) {
+              int real_out_index = out_index_table[out_index];
+              in_i = i;
+              int buf_i = atomicAdd(&counter_buf[kernel_index], 1);
+              kernel_i = kernel_index;
+              rulebook_buf[kernel_index * blockDim.x + buf_i] = in_i;
+              rulebook_buf[kernel_index * blockDim.x +
+                           kernel_size * blockDim.x + buf_i] = real_out_index;
+            }
+          }
+          ++kernel_index;
+        }
+      }
+    }
+  }
+  __syncthreads();
+  for (int i = threadIdx.x; i < kernel_size; i += blockDim.x) {
+    counter_buf2[i] = atomicAdd(&counter[i], counter_buf[i]);
+  }
+  __syncthreads();
+  for (int i = 0; i < kernel_size; i++) {
+    if (threadIdx.x < counter_buf[i]) {
+      rulebook[i * non_zero_num + counter_buf2[i] + threadIdx.x] =
+          rulebook_buf[i * blockDim.x + threadIdx.x];
+      rulebook[i * non_zero_num + offset + counter_buf2[i] + threadIdx.x] =
+          rulebook_buf[i * blockDim.x + kernel_size * blockDim.x + threadIdx.x];
+    }
+  }
+}
+
+template <typename IntT>
+__global__ void GroupIndices(const int n,
+                             const int kernel_size,
+                             const IntT* indices,
+                             int* index_counts,
+                             int* index_groups) {
+  CUDA_KERNEL_LOOP_TYPE(i, n, int64_t) {
+    IntT index = indices[i];
+    // kernel_size at most
+    int j = atomicAdd(index_counts + index, 1);
+    // nnz * kernel_size
+    index_groups[index * kernel_size + j] = i;
+  }
+}
+
+// double space to reduce atomicAdd conflict
+template <typename IntT>
+__global__ void GroupIndicesV2(const int rulebook_len,
+                               const int non_zero_num,
+                               const int kernel_size,
+                               const int half_kernel_offset,
+                               const IntT* indices,
+                               int* index_counts,
+                               int* index_groups) {
+  CUDA_KERNEL_LOOP_TYPE(i, rulebook_len, int64_t) {
+    IntT index = indices[i];
+    int* counts_ptr =
+        i < half_kernel_offset ? index_counts : index_counts + non_zero_num;
+    int* groups_ptr = i < half_kernel_offset
+                          ? index_groups
+                          : index_groups + non_zero_num * kernel_size;
+    // conflict kernel_size times at most
+    int j = atomicAdd(counts_ptr + index, 1);
+    // nnz * kernel_size
+    groups_ptr[index * kernel_size + j] = i;
+  }
+}
+
+inline void CallThrustScan(const GPUContext& dev_ctx,
+                           const int* counter_ptr,
+                           const int kernel_size,
+                           int* offsets_ptr,
+                           int* h_counter_ptr,
+                           int* h_offsets_ptr) {
+#ifdef PADDLE_WITH_HIP
+  thrust::exclusive_scan(thrust::hip::par.on(dev_ctx.stream()),
+#else
+  thrust::exclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                         counter_ptr,
+                         counter_ptr + kernel_size,
+                         offsets_ptr);
+
+  backends::gpu::GpuMemcpyAsync(h_counter_ptr,
+                                counter_ptr,
+                                kernel_size * sizeof(int),
+                                gpuMemcpyDeviceToHost,
+                                dev_ctx.stream());
+
+  backends::gpu::GpuMemcpyAsync(h_offsets_ptr,
+                                offsets_ptr,
+                                kernel_size * sizeof(int),
+                                gpuMemcpyDeviceToHost,
+                                dev_ctx.stream());
+}
+
+// the basic algorithm can refer to convolution_kernel.cc or
+// the second paper
+// example:
+// 1. the rulebook:
+//  the kernel_index:                       0, 0, 0, 1, 1, 1, 2, 2, ....
+//  the out_index(key):                     20, 30, 33, 30, 33, 20, 25
+// 2. mark the index of out_index(value):   0, 1, 2, 3, 4, 5, 6, ....
+// 3. sorted the (key, value)
+// 4. unique the (key, value):
+//  unique_key:     20, 25, 30, 33
+//  unique_values:  0, 2, 3, 5
+//  the index of unique_values is: 0, 1, 2, 3
+// 5. update the out_index by unique_key, unique_value and the index of
+// unique_value:
+//  the new out_index: 0, 2, 3, 2, 3, 0, 1
+template <typename T, typename Context, typename IntT = int>
+int 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,
+                    DenseTensor* counter_per_kernel,
+                    DenseTensor* offsets_per_kernel,
+                    DenseTensor* out_index,
+                    DenseTensor* unique_value,
+                    SparseCooTensor* out,
+                    int* h_counter,
+                    int* h_offsets) {
+  const bool is2D = out_dims.size() == 4 ? true : false;
+  auto indices_dtype = phi::CppTypeToDataType<IntT>::Type();
+  const int64_t non_zero_num = x.nnz();
+  const auto& indices = x.indices();
+  const IntT* indices_ptr = indices.data<IntT>();
+  int* counter_ptr = counter_per_kernel->data<int>();
+  int* offsets_ptr = offsets_per_kernel->data<int>();
+  int kernel_size = is2D ? kernel_sizes[0] * kernel_sizes[1]
+                         : kernel_sizes[0] * kernel_sizes[1] * kernel_sizes[2];
+
+  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 d_x_dims(xdim0, xdim1, xdim2, xdim3);
+  const Dims4D d_kernel_dims(kdim0, kdim1, kdim2, kdim3);
+  const Dims4D d_out_dims(odim0, odim1, odim2, odim3);
+  const Dims4D d_paddings(pdim0, pdim1, pdim2, pdim3);
+  const Dims4D d_strides(sdim0, sdim1, sdim2, sdim3);
+  const Dims4D d_dilations(ddim0, ddim1, ddim2, ddim3);
+
+  // 1. product rule book
+  backends::gpu::GpuMemsetAsync(counter_ptr,
+                                0,
+                                sizeof(int) * counter_per_kernel->numel(),
+                                dev_ctx.stream());
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num, 1);
+
+  const int rulebook_rows = 2;
+  const int rulebook_cols = kernel_size * non_zero_num;
+  DenseTensorMeta rulebook_meta(
+      indices_dtype, {rulebook_rows, rulebook_cols}, DataLayout::NCHW);
+
+  int table_size = 1;
+  for (int i = 0; i < out_dims.size() - 1; i++) {
+    table_size *= out_dims[i];
+  }
+  DenseTensor out_index_table = Empty<int>(dev_ctx, {table_size});
+  int* out_index_table_ptr = out_index_table.data<int>();
+  // index_flags: flag the indices exist or not
+  int index_flags_size = (table_size + 31) / 32;
+  DenseTensor index_flags = Empty<int>(dev_ctx, {index_flags_size});
+  int* index_flags_ptr = index_flags.data<int>();
+  backends::gpu::GpuMemsetAsync(
+      index_flags_ptr, 0, sizeof(int) * index_flags.numel(), dev_ctx.stream());
+
+  if (subm) {
+    DenseTensor tmp_rulebook = Empty(dev_ctx, std::move(rulebook_meta));
+    IntT* rulebook_ptr = tmp_rulebook.data<IntT>();
+    DenseTensor out_indices = EmptyLike<IntT>(dev_ctx, x.indices());
+    int tmpidx = is2D ? 3 : 4;
+    DenseTensor out_values = Empty<T>(dev_ctx, {x.nnz(), kernel_sizes[tmpidx]});
+
+    phi::Copy(dev_ctx, x.indices(), dev_ctx.GetPlace(), false, &out_indices);
+
+    auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num, 1);
+    GetOutIndexTable1<IntT><<<config.block_per_grid,
+                              config.thread_per_block,
+                              0,
+                              dev_ctx.stream()>>>(out_indices.data<IntT>(),
+                                                  non_zero_num,
+                                                  d_x_dims,
+                                                  index_flags_ptr,
+                                                  is2D,
+                                                  out_index_table_ptr);
+
+    size_t cache_size =
+        kernel_size * 2 * sizeof(int) +
+        kernel_size * config.thread_per_block.x * 2 * sizeof(int);
+    const int MAX_CACHE_SIZE = 48 * 1024;
+    while (cache_size >= MAX_CACHE_SIZE) {
+      config.thread_per_block.x /= 2;
+      config.block_per_grid.x *= 2;
+      PADDLE_ENFORCE_GE(
+          config.thread_per_block.x,
+          32,
+          common::errors::Fatal("the shared memory is not enough"));
+      cache_size = kernel_size * 2 * sizeof(int) +
+                   kernel_size * config.thread_per_block.x * 2 * sizeof(int);
+    }
+    ProductSubmRuleBookKernel<IntT><<<config.block_per_grid.x,
+                                      config.thread_per_block.x,
+                                      cache_size,
+                                      dev_ctx.stream()>>>(indices_ptr,
+                                                          d_x_dims,
+                                                          d_kernel_dims,
+                                                          d_out_dims,
+                                                          non_zero_num,
+                                                          d_paddings,
+                                                          d_dilations,
+                                                          d_strides,
+                                                          is2D,
+                                                          index_flags_ptr,
+                                                          out_index_table_ptr,
+                                                          rulebook_ptr,
+                                                          counter_ptr);
+
+    out->SetMember(out_indices, out_values, out_dims, false);
+
+    CallThrustScan(
+        dev_ctx, counter_ptr, kernel_size, offsets_ptr, h_counter, h_offsets);
+
+    dev_ctx.Wait();
+    int rulebook_len = h_offsets[kernel_size - 1] + h_counter[kernel_size - 1];
+    DenseTensor out_rulebook =
+        Empty<IntT>(dev_ctx, {rulebook_rows, rulebook_len});
+    IntT* out_rulebook_ptr = out_rulebook.data<IntT>();
+    config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rulebook_len, 1);
+    cache_size = kernel_size * 2 * sizeof(int);
+    CopyRuleBook<IntT><<<config.block_per_grid,
+                         config.thread_per_block,
+                         cache_size,
+                         dev_ctx.stream()>>>(counter_ptr,
+                                             offsets_ptr,
+                                             rulebook_ptr,
+                                             rulebook_len,
+                                             kernel_size,
+                                             non_zero_num,
+                                             out_rulebook_ptr);
+    *rulebook = out_rulebook;
+
+    return rulebook_len;
+
+  } else {
+    *rulebook = Empty(dev_ctx, std::move(rulebook_meta));
+    IntT* rulebook_ptr = rulebook->data<IntT>();
+
+    ConvHostBuffer& conv_host_buffer = ConvHostBuffer::getInstance();
+    if (conv_host_buffer.using_buffer()) {
+      return ProductRuleBookWithBuffer<T, GPUContext, IntT>(dev_ctx,
+                                                            indices_ptr,
+                                                            d_x_dims,
+                                                            d_kernel_dims,
+                                                            d_out_dims,
+                                                            d_paddings,
+                                                            d_strides,
+                                                            d_dilations,
+                                                            out_dims,
+                                                            kernel_sizes,
+                                                            non_zero_num,
+                                                            kernel_size,
+                                                            rulebook_rows,
+                                                            rulebook_cols,
+                                                            rulebook_ptr,
+                                                            counter_ptr,
+                                                            offsets_ptr,
+                                                            &index_flags,
+                                                            &out_index_table,
+                                                            rulebook,
+                                                            out_index,
+                                                            unique_value,
+                                                            out,
+                                                            h_counter);
+    }
+
+    ProductRuleBookKernel<IntT><<<config.block_per_grid.x,
+                                  config.thread_per_block.x,
+                                  kernel_size * sizeof(int),
+                                  dev_ctx.stream()>>>(indices_ptr,
+                                                      d_x_dims,
+                                                      d_kernel_dims,
+                                                      d_out_dims,
+                                                      non_zero_num,
+                                                      d_paddings,
+                                                      d_dilations,
+                                                      d_strides,
+                                                      is2D,
+                                                      rulebook_ptr,
+                                                      counter_ptr);
+
+    // 2. remove -1
+#ifdef PADDLE_WITH_HIP
+    IntT* last = thrust::remove(thrust::hip::par.on(dev_ctx.stream()),
+#else
+    IntT* last = thrust::remove(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                                rulebook_ptr,
+                                rulebook_ptr + rulebook_rows * rulebook_cols,
+                                -1);
+
+    IntT rulebook_len = (last - rulebook_ptr) / 2;
+
+    CallThrustScan(
+        dev_ctx, counter_ptr, kernel_size, offsets_ptr, h_counter, h_offsets);
+
+    rulebook->Resize({rulebook_rows, static_cast<int>(rulebook_len)});
+    // 3. sorted or merge the out index
+    out_index->ResizeAndAllocate({static_cast<int>(rulebook_len)});
+    DenseTensor unique_key =
+        Empty<int>(dev_ctx, {static_cast<int>(rulebook_len)});
+    int* out_index_ptr = out_index->data<int>();
+    int* unique_key_ptr = unique_key.data<int>();
+
+    backends::gpu::GpuMemsetAsync(
+        unique_key_ptr, 0, sizeof(int), dev_ctx.stream());
+
+    config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rulebook_len, 1);
+    size_t cache_size = sizeof(int) * config.thread_per_block.x;
+    UniqueKernel<IntT><<<config.block_per_grid,
+                         config.thread_per_block,
+                         cache_size,
+                         dev_ctx.stream()>>>(rulebook_ptr + rulebook_len,
+                                             rulebook_len,
+                                             index_flags_ptr,
+                                             out_index_ptr,
+                                             unique_key_ptr);
+
+    int out_nnz = 0;
+    backends::gpu::GpuMemcpyAsync(&out_nnz,
+                                  unique_key_ptr,
+                                  sizeof(int),
+                                  gpuMemcpyDeviceToHost,
+                                  dev_ctx.stream());
+    dev_ctx.Wait();
+
+    const int threads = 256;
+    const int blocks = (index_flags.numel() + threads - 1) / threads;
+    GetOutIndicesCounter<<<blocks, threads, 0, dev_ctx.stream()>>>(
+        index_flags_ptr, index_flags.numel(), out_index_table_ptr);
+#ifdef PADDLE_WITH_HIP
+    thrust::exclusive_scan(thrust::hip::par.on(dev_ctx.stream()),
+#else
+    thrust::exclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                           out_index_table_ptr,
+                           out_index_table_ptr + blocks,
+                           out_index_table_ptr);
+    GetOutIndices<threads>
+        <<<blocks, threads, 0, dev_ctx.stream()>>>(index_flags_ptr,
+                                                   index_flags.numel(),
+                                                   out_index_table_ptr,
+                                                   out_nnz,
+                                                   out_index_ptr);
+
+    const int64_t sparse_dim = is2D ? 3 : 4;
+    DenseTensor out_indices = Empty<IntT>(dev_ctx, {sparse_dim, out_nnz});
+    DenseTensor out_values =
+        Empty<T>(dev_ctx, {out_nnz, kernel_sizes[sparse_dim]});
+    out->SetMember(out_indices, out_values, out_dims, false);
+
+    IntT* out_indices_ptr = out_indices.data<IntT>();
+
+    config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_nnz, 1);
+    GetOutIndexTable<IntT><<<config.block_per_grid,
+                             config.thread_per_block,
+                             0,
+                             dev_ctx.stream()>>>(out_index_ptr,
+                                                 out_nnz,
+                                                 d_out_dims,
+                                                 is2D,
+                                                 out_index_table_ptr,
+                                                 out_indices_ptr);
+    config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rulebook_len, 1);
+    unique_value->ResizeAndAllocate({static_cast<int>(out_nnz * kernel_size)});
+    int* unique_value_ptr = unique_value->data<int>();
+
+    GroupIndices<<<config.block_per_grid,
+                   config.thread_per_block,
+                   0,
+                   dev_ctx.stream()>>>(out_index_table_ptr,
+                                       rulebook_len,
+                                       kernel_size,
+                                       rulebook_ptr + rulebook_len,
+                                       out_index_ptr,
+                                       unique_value_ptr);
+
+    return rulebook_len;
+  }
+}
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/gpu/conv_grad_kernel.cu

@@ -0,0 +1,271 @@
+/* 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 "glog/logging.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.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/funcs/blas/blas.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/sparse/gpu/conv.cu.h"
+
+namespace phi {
+namespace sparse {
+extern size_t workspace_size;
+
+// rulebook[3, rulebook_len]:
+//[
+//  [kernel_index],
+//  [in_i],
+//  [out_i],
+//]
+// x_grad = out_grad * transpose(kernel)
+// kernel_grad = transpose(x) * out_grad
+template <typename T, typename IntT>
+void Conv3dCooGradGPUKernel(const GPUContext& 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) {
+  const bool is_params_freezing = kernel_grad == nullptr;
+  const auto& kernel_dims = kernel.dims();
+  const bool is2D = kernel_dims.size() == 4 ? true : false;
+  const int kernel_size =
+      is2D ? kernel_dims[0] * kernel_dims[1]
+           : kernel_dims[0] * kernel_dims[1] * kernel_dims[2];
+  const int in_channels = is2D ? kernel_dims[2] : kernel_dims[3];
+  const int out_channels = 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);
+
+  DenseTensor in_features = Empty<T>(dev_ctx, {rulebook_len, in_channels});
+  DenseTensor d_x_features = Empty<T>(dev_ctx, {rulebook_len, in_channels});
+  DenseTensor out_grad_features =
+      Empty<T>(dev_ctx, {rulebook_len, out_channels});
+
+  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>();
+  T* d_kernel_ptr = nullptr;
+  if (!is_params_freezing) {
+    *kernel_grad = EmptyLike<T>(dev_ctx, kernel);
+    d_kernel_ptr = kernel_grad->data<T>();
+    backends::gpu::GpuMemsetAsync(
+        d_kernel_ptr, 0, sizeof(T) * kernel_grad->numel(), dev_ctx.stream());
+  }
+
+  int half_kernel_size = kernel_size / 2;
+  auto blas = funcs::GetBlas<GPUContext, 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>();
+  backends::gpu::GpuMemsetAsync(x_grad_values_ptr,
+                                0,
+                                sizeof(T) * x_grad_values.numel(),
+                                dev_ctx.stream());
+  backends::gpu::GpuMemsetAsync(
+      d_x_features_ptr, 0, sizeof(T) * d_x_features.numel(), dev_ctx.stream());
+  phi::Copy<GPUContext>(
+      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<int> offsets(kernel_size + 1);
+
+  int offset = 0, 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, GPUContext>(dev_ctx,
+                                                 x,
+                                                 kernel,
+                                                 out_grad.values(),
+                                                 in_channels,
+                                                 out_channels,
+                                                 half_kernel_size,
+                                                 kernel_grad,
+                                                 &x_grad_values);
+    if (max_count == 0) {
+      return;
+    }
+  }
+
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rulebook_len, 1);
+  DenseTensor unique_value =
+      Empty<int>(dev_ctx, {static_cast<int>(x_grad->nnz() * kernel_size * 2)});
+  DenseTensor out_index = Empty<int>(dev_ctx, {static_cast<int>(x.nnz() * 2)});
+  int* out_index_ptr = out_index.data<int>();
+  int* unique_value_ptr = unique_value.data<int>();
+  backends::gpu::GpuMemsetAsync(
+      out_index_ptr, 0, sizeof(int) * x.nnz() * 2, dev_ctx.stream());
+
+  GroupIndicesV2<<<config.block_per_grid,
+                   config.thread_per_block,
+                   0,
+                   dev_ctx.stream()>>>(rulebook_len,
+                                       x.nnz(),
+                                       kernel_size,
+                                       offsets[kernel_size / 2],
+                                       rulebook_ptr,
+                                       out_index_ptr,
+                                       unique_value_ptr);
+
+  GatherV2<T, IntT>(dev_ctx,
+                    x.values().data<T>(),
+                    out_index_ptr,
+                    unique_value_ptr,
+                    x.nnz(),
+                    kernel_size,
+                    in_channels,
+                    2,
+                    in_features_ptr);
+
+  Gather<T, IntT>(dev_ctx,
+                  out_grad.values().data<T>(),
+                  rulebook_ptr + rulebook_len,
+                  rulebook_len,
+                  out_channels,
+                  out_grad_features_ptr);
+
+  const T* kernel_ptr = kernel.data<T>();
+  T* tmp_d_x_ptr = nullptr;
+  T* tmp_d_kernel_ptr = nullptr;
+  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;
+    tmp_d_x_ptr = d_x_features_ptr + offsets[i] * in_channels;
+    if (!is_params_freezing) {
+      tmp_d_kernel_ptr = d_kernel_ptr + i * in_channels * out_channels;
+    }
+
+    if (!is_params_freezing) {
+      // 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
+  funcs::sparse::ScatterV2<T>(dev_ctx,
+                              d_x_features_ptr,
+                              out_index.data<int>(),
+                              unique_value.data<int>(),
+                              x_grad->nnz(),
+                              kernel_size,
+                              in_channels,
+                              2,
+                              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(), "Conv3dCooGradGPUKernel", ([&] {
+        Conv3dCooGradGPUKernel<T, data_t>(dev_ctx,
+                                          x,
+                                          kernel,
+                                          out,
+                                          rulebook,
+                                          counter,
+                                          out_grad,
+                                          paddings,
+                                          dilations,
+                                          strides,
+                                          groups,
+                                          subm,
+                                          key,
+                                          x_grad,
+                                          kernel_grad);
+      }));
+}
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(conv3d_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::Conv3dCooGradKernel,
+                   float,
+                   double,
+                   phi::float16) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/gpu/conv_host_buffer.h

@@ -0,0 +1,92 @@
+/* Copyright (c) 2025 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 <numeric>
+#include <vector>
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+class ConvHostBuffer {
+ public:
+  ConvHostBuffer(const ConvHostBuffer&) = delete;
+  ConvHostBuffer& operator=(const ConvHostBuffer&) = delete;
+
+  static ConvHostBuffer& getInstance() {
+    static ConvHostBuffer instance;
+    return instance;
+  }
+
+  void set_host_buffer(int* buffer) { h_buffer_ = buffer; }
+
+  int* get_host_buffer() {
+    PADDLE_ENFORCE_EQ(offset_.empty(),
+                      false,
+                      ::common::errors::InvalidArgument(
+                          "Sparse conv buffer offsets should not be empty."));
+    return h_buffer_ + offset_[current_step_++];
+  }
+
+  void reset() { current_step_ = 0; }
+
+  bool using_buffer() { return use_buffer_; }
+
+  int get_buffer_size() { return buffer_size; }
+
+  int get_max_bound() { return max_bound; }
+
+  void init_from_config(const std::vector<std::vector<int>>& kernels,
+                        const std::vector<std::vector<int>>& strides) {
+    PADDLE_ENFORCE_EQ(kernels.size() == strides.size(),
+                      true,
+                      ::common::errors::InvalidArgument(
+                          "The size of kernels should be equal to the "
+                          "size of strides, but get kernel size:[%d], "
+                          "strides size:[%d].",
+                          kernels.size(),
+                          strides.size()));
+    buffer_size = 0;
+    max_bound = 1;
+    offset_.clear();
+    for (size_t i = 0; i < kernels.size(); ++i) {
+      int kernel_size = std::accumulate(
+          kernels[i].begin(), kernels[i].end(), 1, std::multiplies<int>());
+      buffer_size += 2 * kernel_size + 3;
+      offset_.push_back(buffer_size - (2 * kernel_size + 3));
+      int bound = 1;
+      for (size_t j = 0; j < kernels[i].size(); ++j) {
+        bound *= (kernels[i][j] + strides[i][j] - 1) / strides[i][j];
+      }
+      max_bound = std::max(max_bound, bound);
+    }
+    use_buffer_ = true;
+  }
+
+ private:
+  ConvHostBuffer() {}
+  ~ConvHostBuffer() {}
+
+  int* h_buffer_;
+  int buffer_size;
+  std::vector<int> offset_;
+  int current_step_{0};
+  int max_bound;
+  bool use_buffer_{false};
+};
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/gpu/conv_kernel.cu

@@ -0,0 +1,322 @@
+/* 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_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.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/cast_kernel.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/funcs/scatter.cu.h"
+#include "paddle/phi/kernels/funcs/sparse/scatter.cu.h"
+#include "paddle/phi/kernels/sparse/gpu/conv.cu.h"
+#include "paddle/phi/kernels/sparse/gpu/conv_host_buffer.h"
+
+#include "glog/logging.h"
+
+namespace phi {
+namespace sparse {
+
+#define GATHER_GEMM_SCATTER(arch, input_type, x_nnz, kernel)             \
+  ({                                                                     \
+    const input_type* kernel_ptr = kernel.data<input_type>();            \
+    const input_type* x_nnz_ptr = x_nnz.data<input_type>();              \
+    for (int i = 0; i < kernel_size; i++) {                              \
+      if (h_counter_ptr[i] <= 0) {                                       \
+        continue;                                                        \
+      }                                                                  \
+      const int M = h_counter_ptr[i];                                    \
+      const int K = in_channels;                                         \
+      const int N = out_channels;                                        \
+      const input_type* tmp_kernel_ptr = kernel_ptr + i * K * N;         \
+      const IntT* gather_indices = rulebook_ptr + h_offsets_ptr[i];      \
+      const IntT* scatter_indices =                                      \
+          rulebook_ptr + rulebook_len + h_offsets_ptr[i];                \
+      const size_t key = autotune::GenKey(M / features_num_range, N, K); \
+      GatherGemmScatterDriver<arch, false, false>(                       \
+          dev_ctx,                                                       \
+          key,                                                           \
+          x_nnz_ptr,                                                     \
+          tmp_kernel_ptr,                                                \
+          out_values_ptr,                                                \
+          out_values_ptr,                                                \
+          M,                                                             \
+          N,                                                             \
+          K,                                                             \
+          gather_indices,                                                \
+          static_cast<const IntT*>(nullptr),                             \
+          scatter_indices,                                               \
+          static_cast<T>(1.0),                                           \
+          static_cast<T>(1.0),                                           \
+          nullptr);                                                      \
+    }                                                                    \
+  })
+
+template <typename T, typename IntT>
+void Conv3dCooGPUKernel(const GPUContext& 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) {
+  // 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 auto& kernel_dims = kernel.dims();
+  const bool is2D = x_dims.size() == 4 ? true : false;
+  int kernel_size = is2D ? kernel_dims[0] * kernel_dims[1]
+                         : kernel_dims[0] * kernel_dims[1] * kernel_dims[2];
+
+  int rank = is2D ? 4 : 5;
+  std::vector<int> out_dims_vec(rank, 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] = 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 = is2D ? kernel_dims[2] : kernel_dims[3];
+  const int out_channels = is2D ? kernel_dims[3] : kernel_dims[4];
+
+  int* h_counter_ptr{nullptr};
+  int* h_offsets_ptr{nullptr};
+
+  phi::sparse::ConvHostBuffer& conv_host_buffer =
+      phi::sparse::ConvHostBuffer::getInstance();
+  DenseTensor h_counter, h_offsets;
+  if (conv_host_buffer.using_buffer()) {
+    int* h_buffer_ptr = conv_host_buffer.get_host_buffer();
+    h_counter_ptr = h_buffer_ptr;
+    h_offsets_ptr = h_buffer_ptr + kernel_size;
+  } else {
+    h_counter.Resize({kernel_size});
+    h_offsets.Resize({kernel_size + 1});
+    h_counter_ptr = dev_ctx.template HostAlloc<int>(&h_counter);
+    h_offsets_ptr = dev_ctx.template HostAlloc<int>(&h_offsets);
+  }
+
+  // Second algorithm:
+  // https://pdfs.semanticscholar.org/5125/a16039cabc6320c908a4764f32596e018ad3.pdf
+  // 1. product rulebook
+  DenseTensor counter_per_kernel = Empty<int>(dev_ctx, {kernel_size});
+  DenseTensor offsets_per_kernel = Empty<int>(dev_ctx, {kernel_size});
+  DenseTensor out_index = Empty<int>(dev_ctx, {1});
+  DenseTensor unique_value = Empty<int>(dev_ctx, {1});
+
+  if (is2D) {
+    VLOG(6) << "call SubmConv2D or Conv2D " << subm << " and the key is "
+            << key;
+  } else {
+    VLOG(6) << "call SubmConv3D or Conv3D " << subm << " and the key is "
+            << key;
+  }
+
+  int rulebook_len = 0;
+  const IntT* rulebook_ptr = nullptr;
+  bool need_product_rulebook = true;
+  if (subm && !key.empty()) {
+    rulebook_ptr =
+        funcs::sparse::PrepareSubm<T, IntT, GPUContext>(dev_ctx,
+                                                        x,
+                                                        key,
+                                                        out_dims,
+                                                        out,
+                                                        h_counter_ptr,
+                                                        h_offsets_ptr,
+                                                        &rulebook_len,
+                                                        &need_product_rulebook);
+  }
+
+  if (need_product_rulebook) {
+    DenseTensor tmp_rulebook;
+    rulebook_len = ProductRuleBook<T, GPUContext, IntT>(dev_ctx,
+                                                        x,
+                                                        kernel_sizes,
+                                                        subm_paddings,
+                                                        dilations,
+                                                        subm_strides,
+                                                        out_dims,
+                                                        subm,
+                                                        &tmp_rulebook,
+                                                        &counter_per_kernel,
+                                                        &offsets_per_kernel,
+                                                        &out_index,
+                                                        &unique_value,
+                                                        out,
+                                                        h_counter_ptr,
+                                                        h_offsets_ptr);
+    rulebook_ptr = tmp_rulebook.data<IntT>();
+    DenseTensor h_counter_tensor;
+    h_counter_tensor.Resize({kernel_size});
+    int* h_counter_tensor_ptr =
+        dev_ctx.template HostAlloc<int>(&h_counter_tensor);
+    for (int i = 0; i < kernel_size; ++i) {
+      h_counter_tensor_ptr[i] = h_counter_ptr[i];
+    }
+    funcs::sparse::SaveToTable(dev_ctx,
+                               x,
+                               key,
+                               tmp_rulebook,
+                               h_counter_tensor,
+                               out,
+                               rulebook,
+                               counter);
+  }
+  if (subm) {
+    auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rulebook_len, 1);
+    unique_value.ResizeAndAllocate(
+        {static_cast<int>(out->nnz() * kernel_size)});
+    out_index.ResizeAndAllocate({static_cast<int>(rulebook_len)});
+    int* out_index_ptr = out_index.data<int>();
+    int* unique_value_ptr = unique_value.data<int>();
+    backends::gpu::GpuMemsetAsync(
+        out_index_ptr, 0, sizeof(int) * rulebook_len, dev_ctx.stream());
+    GroupIndices<<<config.block_per_grid,
+                   config.thread_per_block,
+                   0,
+                   dev_ctx.stream()>>>(rulebook_len,
+                                       kernel_size,
+                                       rulebook_ptr + rulebook_len,
+                                       out_index_ptr,
+                                       unique_value_ptr);
+  }
+  // 2. gather
+  DenseTensor in_features = Empty<T>(dev_ctx, {rulebook_len, in_channels});
+  DenseTensor out_features = Empty<T>(dev_ctx, {rulebook_len, out_channels});
+  T* in_features_ptr = in_features.data<T>();
+  T* out_features_ptr = out_features.data<T>();
+  funcs::SetConstant<GPUContext, T> set_zero;
+  set_zero(dev_ctx, &out_features, static_cast<T>(0.0f));
+
+  Gather<T, IntT>(dev_ctx,
+                  x.values().data<T>(),
+                  rulebook_ptr,
+                  rulebook_len,
+                  in_channels,
+                  in_features_ptr);
+
+  // 3. call gemm for every werght
+  auto blas = funcs::GetBlas<GPUContext, T>(dev_ctx);
+  auto* out_values = out->mutable_values();
+  T* out_values_ptr = out_values->data<T>();
+  set_zero(dev_ctx, out_values, static_cast<T>(0.0f));
+
+  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;
+    const int N = 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
+  funcs::sparse::ScatterV2<T>(dev_ctx,
+                              out_features_ptr,
+                              out_index.data<int>(),
+                              unique_value.data<int>(),
+                              out->nnz(),
+                              kernel_size,
+                              out_channels,
+                              1,
+                              out_values_ptr);
+}
+
+/**
+ * x: the input SparseCooTensor, shape is (N, D, H, W, C)
+ * kernel: the weight data, shape is (D, H, W, C, OC)
+ * out: the output SparseCooTensor, shape is (N, D, H, W, OC)
+ * rulebook: return rulebook if key is not valid else return nullptr
+ * counter: return counter if key is not valid else return nullptr
+ **/
+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(), "Conv3dCooGPUKernel", ([&] {
+                                 Conv3dCooGPUKernel<T, data_t>(dev_ctx,
+                                                               x,
+                                                               kernel,
+                                                               paddings,
+                                                               dilations,
+                                                               strides,
+                                                               groups,
+                                                               subm,
+                                                               key,
+                                                               out,
+                                                               rulebook,
+                                                               counter);
+                               }));
+}
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(conv3d_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::Conv3dCooKernel,
+                   float,
+                   double,
+                   phi::float16) {
+  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/gpu/conv_kernel_igemm.cu

@@ -0,0 +1,216 @@
+// Copyright (c) 2024 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/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/funcs/sparse/convolution.h"
+#include "paddle/phi/kernels/funcs/transpose_function.cuh"
+#include "paddle/phi/kernels/sparse/gpu/conv_kernel_impl.cuh"
+#include "paddle/phi/kernels/sparse/gpu/sparse_conv_hashmap.cuh"
+
+#include "glog/logging.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename IntT>
+void Conv3dImplicitGemmGPUKernel(const GPUContext& 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) {
+  // Currently, only support x.layout is NDHWC, subm = true, stride = 1, groups
+  // = 1, dilations = 1
+  PADDLE_ENFORCE_EQ(
+      subm,
+      true,
+      common::errors::InvalidArgument("The subm must be true, but received %s.",
+                                      subm ? "true" : "false"));
+  PADDLE_ENFORCE_EQ(groups,
+                    1,
+                    common::errors::InvalidArgument(
+                        "The group must be 1, but received %d.", groups));
+
+  const auto& x_dims = x.dims();
+  const auto& kernel_dims = kernel.dims();
+  const bool is2D = x_dims.size() == 4 ? true : false;
+
+  if (is2D) {
+    PADDLE_ENFORCE_EQ(
+        (kernel_dims.size() == 4),
+        true,
+        common::errors::InvalidArgument(
+            "For 2D case, the size of kernel_dims must be 4, but received %d.",
+            kernel_dims.size()));
+    PADDLE_ENFORCE_EQ(
+        (strides.size() == 2 && strides[0] == 1 && strides[1] == 1),
+        true,
+        common::errors::InvalidArgument(
+            "The strides must be 1, but received %d, %d.",
+            strides[0],
+            strides[1]));
+    PADDLE_ENFORCE_EQ(
+        (dilations.size() == 2 && dilations[0] == 1 && dilations[1] == 1),
+        true,
+        common::errors::InvalidArgument(
+            "The dilations must be 1, but received %d, %d.",
+            dilations[0],
+            dilations[1]));
+
+  } else {
+    PADDLE_ENFORCE_EQ(
+        (kernel_dims.size() == 5),
+        true,
+        common::errors::InvalidArgument(
+            "For 3D case, the size of kernel_dims must be 5, but received %d.",
+            kernel_dims.size()));
+    PADDLE_ENFORCE_EQ((strides.size() == 3 && strides[0] == 1 &&
+                       strides[1] == 1 && strides[2] == 1),
+                      true,
+                      common::errors::InvalidArgument(
+                          "The strides must be 1, but received %d, %d, %d.",
+                          strides[0],
+                          strides[1],
+                          strides[2]));
+    PADDLE_ENFORCE_EQ((dilations.size() == 3 && dilations[0] == 1 &&
+                       dilations[1] == 1 && dilations[2] == 1),
+                      true,
+                      common::errors::InvalidArgument(
+                          "The dilations must be 1, but received %d, %d, %d.",
+                          dilations[0],
+                          dilations[1],
+                          dilations[2]));
+  }
+
+  int kernel_volume = is2D ? kernel_dims[0] * kernel_dims[1]
+                           : kernel_dims[0] * kernel_dims[1] * kernel_dims[2];
+  int in_channels = is2D ? kernel_dims[2] : kernel_dims[3];
+  int out_channels = is2D ? kernel_dims[3] : kernel_dims[4];
+
+  int rank = is2D ? 4 : 5;
+  std::vector<int> out_dims_vec(rank, 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] = 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);
+
+  // Set the output tensor
+  if (subm) {
+    DenseTensor out_indices = EmptyLike<IntT>(dev_ctx, x.indices());
+    int tmpidx = is2D ? 3 : 4;
+    DenseTensor out_values = Empty<T>(dev_ctx, {x.nnz(), kernel_sizes[tmpidx]});
+    phi::Copy(dev_ctx, x.indices(), dev_ctx.GetPlace(), false, &out_indices);
+    out->SetMember(out_indices, out_values, out_dims, false);
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "The subm must be true, but received %s.", subm ? "true" : "false"));
+  }
+
+  build_sparse_conv_kmap<IntT>(
+      dev_ctx, x, key, kernel_sizes, strides, kernel_volume, is2D, out);
+
+  auto* out_kmap_cache_ptr = out->GetKmapCache(key);
+
+  DenseTensor kernel_transpose = EmptyLike<T, GPUContext>(dev_ctx, kernel);
+  std::vector<int> perm;
+  if (is2D) {
+    perm = {1, 0, 2, 3};
+  } else {
+    perm = {2, 1, 0, 3, 4};
+  }
+  funcs::TransposeGPUKernelDriver<T>(dev_ctx, kernel, perm, &kernel_transpose);
+
+#ifdef PADDLE_WITH_CUDA
+  conv_forward_implicit_gemm_cuda(dev_ctx,
+                                  x.values(),
+                                  kernel_transpose,
+                                  *(out_kmap_cache_ptr->out_in_map),
+                                  out->nnz(),
+                                  out_channels,
+                                  *(out->mutable_values()));
+#else
+  PADDLE_THROW(common::errors::Unimplemented(
+      "conv_forward_implicit_gemm_cuda is only supported on CUDA."));
+#endif
+}
+
+/**
+ * x: the input SparseCooTensor, shape is (N, D, H, W, C)
+ * kernel: the weight data, shape is (D, H, W, C, OC)
+ * out: the output SparseCooTensor, shape is (N, D, H, W, OC)
+ * rulebook: return rulebook if key is not valid else return nullptr
+ * counter: return counter if key is not valid else return nullptr
+ **/
+template <typename T, typename Context>
+void Conv3dImplicitGemmKernel(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) {
+#ifdef PADDLE_WITH_CUDA
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.indices().dtype(), "Conv3dImplicitGemmGPUKernel's indices", ([&] {
+        // Conv3dImplicitGemmGPUKernel<T, data_t>(dev_ctx,
+        Conv3dImplicitGemmGPUKernel<T, int64_t>(dev_ctx,
+                                                x,
+                                                kernel,
+                                                paddings,
+                                                dilations,
+                                                strides,
+                                                groups,
+                                                subm,
+                                                key,
+                                                out);
+      }));
+#else
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Conv3dImplicitGemmKernel is only supported on CUDA."));
+#endif
+}
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(conv3d_implicit_gemm,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::Conv3dImplicitGemmKernel,
+                   float,
+                   phi::float16) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}

paddle/phi/kernels/sparse/gpu/conv_memory_utils.cuh

@@ -0,0 +1,122 @@
+// Copyright (c) 2025 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
+
+#ifdef PADDLE_WITH_CUDA
+
+template <int bytes>
+struct global_load;
+
+template <>
+struct global_load<16> {
+  __device__ __inline__ global_load(uint4 &D,  // NOLINT
+                                    void const *ptr,
+                                    int pred_guard) {
+    uint4 &data = *reinterpret_cast<uint4 *>(&D);
+    asm volatile(
+        "{\n"
+        "  .reg .pred p;\n"
+        "  setp.ne.b32 p, %5, 0;\n"
+        "  mov.b32 %0, %6;\n"
+        "  mov.b32 %1, %7;\n"
+        "  mov.b32 %2, %8;\n"
+        "  mov.b32 %3, %9;\n"
+        "  @p ld.global.v4.u32 {%0, %1, %2, %3}, [%4];\n"
+        "}\n"
+        : "=r"(data.x), "=r"(data.y), "=r"(data.z), "=r"(data.w)
+        : "l"(ptr),
+          "r"((int)(pred_guard & 1)),
+          "r"(data.x),
+          "r"(data.y),
+          "r"(data.z),
+          "r"(data.w));
+  }
+};
+
+template <>
+struct global_load<8> {
+  __device__ __inline__ global_load(uint4 &D,  // NOLINT
+                                    void const *ptr,
+                                    int pred_guard) {
+    uint2 const *ptr_ldg = reinterpret_cast<uint2 const *>(ptr);
+#pragma unroll
+    for (int ldg_idx = 0; ldg_idx < 2; ldg_idx++) {
+      uint2 &data = *(reinterpret_cast<uint2 *>(&D) + ldg_idx);
+      asm volatile(
+          "{\n"
+          "  .reg .pred p;\n"
+          "  setp.ne.b32 p, %3, 0;\n"
+          "  mov.b32 %0, %4;\n"
+          "  mov.b32 %1, %5;\n"
+          "  @p ld.global.v2.u32 {%0, %1}, [%2];\n"
+          "}\n"
+          : "=r"(data.x), "=r"(data.y)
+          : "l"(ptr_ldg + ldg_idx),
+            "r"((int)(pred_guard & (1 << ldg_idx))),
+            "r"(data.x),
+            "r"(data.y));
+    }
+  }
+};
+
+template <>
+struct global_load<4> {
+  __device__ __inline__ global_load(uint4 &D,  // NOLINT
+                                    void const *ptr,
+                                    int pred_guard) {
+    unsigned const *ptr_ldg = reinterpret_cast<unsigned const *>(ptr);
+#pragma unroll
+    for (int ldg_idx = 0; ldg_idx < 4; ldg_idx++) {
+      unsigned &data = *(reinterpret_cast<unsigned *>(&D) + ldg_idx);
+      asm volatile(
+          "{\n"
+          "  .reg .pred p;\n"
+          "  setp.ne.b32 p, %2, 0;\n"
+          "  mov.b32 %0, %3;\n"
+          "  @p ld.global.u32 %0, [%1];\n"
+          "}\n"
+          : "=r"(data)
+          : "l"(ptr_ldg + ldg_idx),
+            "r"((int)(pred_guard & (1 << ldg_idx))),
+            "r"(data));
+    }
+  }
+};
+
+template <>
+struct global_load<2> {
+  __device__ __inline__ global_load(uint4 &D,  // NOLINT
+                                    void const *ptr,
+                                    int pred_guard) {
+    uint16_t const *ptr_ldg = reinterpret_cast<uint16_t const *>(ptr);
+#pragma unroll
+    for (int ldg_idx = 0; ldg_idx < 8; ldg_idx++) {
+      uint16_t &data = *(reinterpret_cast<uint16_t *>(&D) + ldg_idx);
+      asm volatile(
+          "{\n"
+          "  .reg .pred p;\n"
+          "  setp.ne.b32 p, %2, 0;\n"
+          "  mov.b16 %0, %3;\n"
+          "  @p ld.global.u16 %0, [%1];\n"
+          "}\n"
+          : "=h"(data)
+          : "l"(ptr_ldg + ldg_idx),
+            "r"((int)(pred_guard & (1 << ldg_idx))),
+            "h"(data));
+    }
+  }
+};
+
+#endif  // PADDLE_WITH_CUDA

paddle/phi/kernels/sparse/gpu/conv_with_buffer.cu.h

@@ -0,0 +1,541 @@
+/* Copyright (c) 2025 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 "paddle/phi/kernels/sparse/conv_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/aligned_vector.h"
+#include "paddle/phi/kernels/funcs/index_impl.cu.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/sparse/scatter.cu.h"
+#include "paddle/phi/kernels/funcs/sparse/utils.cu.h"
+#include "paddle/phi/kernels/primitive/compute_primitives.h"
+#include "paddle/phi/kernels/sparse/gpu/conv_host_buffer.h"
+
+namespace phi {
+namespace sparse {
+using Dims4D = funcs::sparse::Dims4D;
+
+inline __device__ uint32_t BitCount(const uint32_t data) {
+  uint32_t count = data;
+  count = (count & 0x55555555) + ((count >> 1) & 0x55555555);
+  count = (count & 0x33333333) + ((count >> 2) & 0x33333333);
+  count = (count & 0x0f0f0f0f) + ((count >> 4) & 0x0f0f0f0f);
+  count = (count & 0x00ff00ff) + ((count >> 8) & 0x00ff00ff);
+  count = (count & 0x0000ffff) + ((count >> 16) & 0x0000ffff);
+  return count;
+}
+
+static __global__ void GetOutIndicesCounter(const int* flags,
+                                            const int n,
+                                            int* out) {
+  int tid = threadIdx.x + blockDim.x * blockIdx.x;
+  __shared__ int block_count;
+  if (threadIdx.x == 0) {
+    block_count = 0;
+  }
+  __syncthreads();
+
+  if (tid < n) {
+    // get the count of 1 in flags[tid]
+    uint32_t count = BitCount(static_cast<uint32_t>(flags[tid]));
+    // add to block_count
+    // TODO(zhangkaihuo): replace with block reduce_sum
+    atomicAdd(&block_count, static_cast<int>(count));
+  }
+  __syncthreads();
+  // write to out
+  if (threadIdx.x == 0) {
+    out[blockIdx.x] = block_count;
+  }
+}
+
+// unique the out indices in rulebook
+template <typename IntT>
+__global__ void UniqueKernel(const IntT* in_indices,
+                             int* rulebook_len,
+                             int* index_flags,
+                             int* out_indices,
+                             int* nnz) {
+  extern __shared__ int cache[];
+  __shared__ int count, start;
+  if (threadIdx.x == 0) {
+    count = 0;
+    start = 0;
+  }
+  __syncthreads();
+  int rulebook_len_num = rulebook_len[0] / 2;
+
+  int i = threadIdx.x + blockDim.x * blockIdx.x;
+  if (i < rulebook_len_num) {
+    // atomicOr only support int
+    int index = static_cast<int>((in_indices + rulebook_len_num)[i]);
+    const bool flag = funcs::sparse::SetBits(index, index_flags);
+    if (!flag) {
+      int j = atomicAdd(&count, 1);
+      cache[j] = index;
+    }
+  }
+  __syncthreads();
+
+  if (threadIdx.x == 0) {
+    start = atomicAdd(nnz, count);
+  }
+  __syncthreads();
+  for (int i = threadIdx.x; i < count; i += blockDim.x) {
+    out_indices[start + i] = cache[i];
+  }
+}
+
+struct is_equal {
+  int value;
+  __host__ __device__ bool operator()(int x) const { return x == value; }
+};
+
+template <typename T, typename Pred>
+__global__ void mark_kernel(const T* input, int* flags, Pred pred, int n) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < n) {
+    flags[idx] = !pred(input[idx]);
+  }
+}
+
+template <typename T>
+__global__ void compact_kernel(const T* input,
+                               T* output,
+                               const int* indices,
+                               const int* flags,
+                               int n,
+                               int* out_num) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < n) {
+    if (flags[idx]) output[indices[idx]] = input[idx];
+    if (idx == n - 1) out_num[0] = indices[idx] + flags[idx];
+  }
+}
+
+template <typename T, typename Pred>
+void cuda_remove(const GPUContext& dev_ctx,
+                 DenseTensor& input,  // NOLINT
+                 int n,
+                 Pred pred,
+                 int* out_num_ptr) {
+  const int block_size = 256;
+  const int grid_size = (n + block_size - 1) / block_size;
+  DenseTensor flags = Empty<int>(dev_ctx, {n});
+  DenseTensor indices = Empty<int>(dev_ctx, {n});
+  DenseTensor out = Empty<T>(dev_ctx, {n});
+
+  mark_kernel<<<grid_size, block_size, 0, dev_ctx.stream()>>>(
+      input.data<T>(), flags.data<int>(), pred, n);
+
+  size_t temp_size = 0;
+  cub::DeviceScan::ExclusiveSum(NULL,
+                                temp_size,
+                                flags.data<int>(),
+                                indices.data<int>(),
+                                n,
+                                dev_ctx.stream());
+  phi::Allocator* allocator =
+      const_cast<phi::Allocator*>(&(dev_ctx.GetAllocator()));
+  auto ws = allocator->Allocate(temp_size)->ptr();
+  cub::DeviceScan::ExclusiveSum(ws,
+                                temp_size,
+                                flags.data<int>(),
+                                indices.data<int>(),
+                                n,
+                                dev_ctx.stream());
+
+  compact_kernel<<<grid_size, block_size, 0, dev_ctx.stream()>>>(
+      input.data<T>(),
+      out.data<T>(),
+      indices.data<int>(),
+      flags.data<int>(),
+      n,
+      out_num_ptr);
+
+  backends::gpu::GpuMemcpyAsync(input.data<T>(),
+                                out.data<T>(),
+                                sizeof(T) * n,
+                                gpuMemcpyDeviceToDevice,
+                                dev_ctx.stream());
+}
+
+template <int BS>
+__global__ void GetOutIndices(const int* flags,
+                              const int n,
+                              const int* offsets,
+                              const int out_nnz,
+                              int* out) {
+  int tid = threadIdx.x + blockDim.x * blockIdx.x;
+  __shared__ int block_counts[BS];
+  __shared__ int block_outs[BS * 32];
+
+  int count = 0;
+
+  if (tid < n) {
+    // get the count of 1 in flags[tid]
+    int flag = flags[tid];
+    count = BitCount(static_cast<uint32_t>(flag));
+  }
+
+  // call block prefix_sum
+  // using namespace cub;
+  typedef cub::BlockScan<int, BS> BlockScan;
+  __shared__ typename BlockScan::TempStorage temp_storage;
+  BlockScan(temp_storage).ExclusiveSum(count, count);
+  __syncthreads();
+
+  // write index to out
+  if (tid < n) {
+    // get the count of 1 in flags[tid]
+    int flag = flags[tid];
+    // int j = block_counts[threadIdx.x];
+    int j = count;
+    // TODO(zhangkaihuo): opt the loop
+    for (int i = 0; i < 32; ++i) {
+      if ((1 & (flag >> i)) == 1) {
+        block_outs[j++] = (tid << 5) + i;
+      }
+    }
+  }
+
+  __syncthreads();
+  // write to block_outs
+  int start = offsets[blockIdx.x];
+  int end = blockIdx.x == gridDim.x - 1 ? out_nnz : offsets[blockIdx.x + 1];
+  for (int i = threadIdx.x; i < end - start; i += blockDim.x) {
+    out[start + i] = block_outs[i];
+  }
+}
+
+template <typename IntT>
+__global__ void GroupIndices(const int* out_index_table,
+                             const int n,
+                             const int kernel_size,
+                             IntT* out_indices,
+                             int* out_index_counts,
+                             int* out_index_groups) {
+  CUDA_KERNEL_LOOP_TYPE(i, n, int64_t) {
+    IntT index = out_indices[i];
+    int real_index = out_index_table[index];
+    out_indices[i] = real_index;
+
+    // kernel_size at most
+    int j = atomicAdd(out_index_counts + real_index, 1);
+    // nnz * kernel_size
+    out_index_groups[real_index * kernel_size + j] = i;
+  }
+}
+
+template <typename IntT>
+__global__ void GetOutIndexTable(int* indices,
+                                 const int non_zero_num,
+                                 const Dims4D out_dims,
+                                 const bool is2D,
+                                 int* out_index_table,
+                                 IntT* out_indices) {
+  CUDA_KERNEL_LOOP_TYPE(i, non_zero_num, int64_t) {
+    IntT index = static_cast<IntT>(indices[i]);
+    out_index_table[index] = i;
+    IntT batch, x, y, z;
+    funcs::sparse::IndexToPoint<Dims4D>(index, out_dims, &batch, &x, &y, &z);
+    // get out indices
+    out_indices[i] = batch;
+    if (is2D) {
+      out_indices[i + non_zero_num] = y;
+      out_indices[i + non_zero_num * 2] = x;
+    } else {
+      out_indices[i + non_zero_num] = z;
+      out_indices[i + non_zero_num * 2] = y;
+      out_indices[i + non_zero_num * 3] = x;
+    }
+    indices[i] = 0;
+  }
+}
+
+/**
+ * @brief product rulebook
+ * for input_i in x_indices:
+ *   if input_i participate in the convolution calculation:
+ *       infer the output_i by input_i and kernel_i
+ *       save output_i
+ *
+ * x_indices: the indices of input features
+ * x_dims: the input dims
+ * kernel_dims: the kernel dims
+ * out_dims: the output dims
+ * non_zero_num: the number of input features
+ * rulebook: the rulebook to save the kernel index, input index and output index
+ * counter: save the number of times each location in the kernel participates in
+ * the calculation
+ **/
+template <typename T>
+__global__ void ProductRuleBookKernel(const T* x_indices,
+                                      const Dims4D x_dims,
+                                      const Dims4D kernel_dims,
+                                      const Dims4D out_dims,
+                                      const int64_t non_zero_num,
+                                      const Dims4D paddings,
+                                      const Dims4D dilations,
+                                      const Dims4D strides,
+                                      const bool is2D,
+                                      T* rulebook,
+                                      int* counter) {
+  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  extern __shared__ int counter_buf[];  // kernel_size
+  const int kernel_size = kernel_dims[3] * kernel_dims[2] * kernel_dims[1];
+  const int offset = kernel_size * non_zero_num;
+  for (int i = threadIdx.x; i < kernel_size; i += blockDim.x) {
+    counter_buf[i] = 0;
+  }
+  __syncthreads();
+
+  for (int i = tid; i < non_zero_num; i += gridDim.x * blockDim.x) {
+    int kernel_index = 0;
+    T batch = x_indices[i];
+    T in_z = is2D ? 0 : x_indices[i + non_zero_num];
+    T in_y =
+        is2D ? x_indices[i + non_zero_num] : x_indices[i + 2 * non_zero_num];
+    T in_x = is2D ? x_indices[i + 2 * non_zero_num]
+                  : x_indices[i + 3 * non_zero_num];
+    for (int kz = 0; kz < kernel_dims[1]; kz++) {
+      for (int ky = 0; ky < kernel_dims[2]; ky++) {
+        for (int kx = 0; kx < kernel_dims[3]; kx++) {
+          int in_i = -1, out_index = -1, kernel_i = -1;
+          if (funcs::sparse::Check(x_dims,
+                                   kernel_dims,
+                                   paddings,
+                                   dilations,
+                                   strides,
+                                   in_x,
+                                   in_y,
+                                   in_z,
+                                   kx,
+                                   ky,
+                                   kz)) {
+            T out_z =
+                is2D ? 0
+                     : (in_z + paddings[1] - kz * dilations[1]) / strides[1];
+            T out_y = (in_y + paddings[2] - ky * dilations[2]) / strides[2];
+            T out_x = (in_x + paddings[3] - kx * dilations[3]) / strides[3];
+            in_i = i;
+            out_index = funcs::sparse::PointToIndex<Dims4D>(
+                batch, out_x, out_y, out_z, out_dims);
+            atomicAdd(&counter_buf[kernel_index], 1);
+            kernel_i = kernel_index;
+          }
+          rulebook[kernel_index * non_zero_num + i] = in_i;
+          rulebook[kernel_index * non_zero_num + offset + i] = out_index;
+          ++kernel_index;
+        }
+      }
+    }
+  }
+  __syncthreads();
+  for (int i = threadIdx.x; i < kernel_size; i += blockDim.x) {
+    atomicAdd(&counter[i], counter_buf[i]);
+  }
+}
+
+template <typename T, typename Context, typename IntT = int>
+int ProductRuleBookWithBuffer(const Context& dev_ctx,
+                              const IntT* indices_ptr,
+                              const Dims4D& d_x_dims,
+                              const Dims4D& d_kernel_dims,
+                              const Dims4D& d_out_dims,
+                              const Dims4D& d_paddings,
+                              const Dims4D& d_strides,
+                              const Dims4D& d_dilations,
+                              const DDim& out_dims,
+                              const std::vector<int>& kernel_sizes,
+                              const int64_t& non_zero_num,
+                              const int& kernel_size,
+                              const int& rulebook_rows,
+                              const int& rulebook_cols,
+                              IntT* rulebook_ptr,
+                              int* counter_ptr,
+                              int* offsets_ptr,
+                              DenseTensor* index_flags,
+                              DenseTensor* out_index_table,
+                              DenseTensor* rulebook,
+                              DenseTensor* out_index,
+                              DenseTensor* unique_value,
+                              SparseCooTensor* out,
+                              int* h_buffer) {
+  DenseTensor d_buffer = Empty<int>(dev_ctx, {2 * kernel_size + 3});
+  const bool is2D = out_dims.size() == 4 ? true : false;
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num, 1);
+  ProductRuleBookKernel<IntT><<<config.block_per_grid.x,
+                                config.thread_per_block.x,
+                                kernel_size * sizeof(int),
+                                dev_ctx.stream()>>>(indices_ptr,
+                                                    d_x_dims,
+                                                    d_kernel_dims,
+                                                    d_out_dims,
+                                                    non_zero_num,
+                                                    d_paddings,
+                                                    d_dilations,
+                                                    d_strides,
+                                                    is2D,
+                                                    rulebook_ptr,
+                                                    counter_ptr);
+
+  DenseTensor rulebook_len_tensor = Empty<int>(dev_ctx, {1});
+  cuda_remove<IntT>(dev_ctx,
+                    *rulebook,
+                    rulebook_rows * rulebook_cols,
+                    is_equal{-1},
+                    rulebook_len_tensor.data<int>());
+
+  size_t temp_size = 0;
+  cub::DeviceScan::ExclusiveSum(
+      NULL, temp_size, counter_ptr, offsets_ptr, kernel_size, dev_ctx.stream());
+  phi::Allocator* allocator =
+      const_cast<phi::Allocator*>(&(dev_ctx.GetAllocator()));
+  auto ws = allocator->Allocate(temp_size)->ptr();
+  cub::DeviceScan::ExclusiveSum(
+      ws, temp_size, counter_ptr, offsets_ptr, kernel_size, dev_ctx.stream());
+
+  int64_t max_nnz =
+      phi::sparse::ConvHostBuffer::getInstance().get_max_bound() * non_zero_num;
+  rulebook->Resize({rulebook_rows, static_cast<int>(max_nnz)});
+  // 3. sorted or merge the out index
+
+  out_index->ResizeAndAllocate({static_cast<int>(max_nnz)});
+  DenseTensor unique_key = Empty<int>(dev_ctx, {static_cast<int>(max_nnz)});
+  int* out_index_ptr = out_index->data<int>();
+  int* unique_key_ptr = unique_key.data<int>();
+
+  backends::gpu::GpuMemsetAsync(
+      unique_key_ptr, 0, sizeof(int), dev_ctx.stream());
+  config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, max_nnz, 1);
+  size_t cache_size = sizeof(int) * config.thread_per_block.x;
+  int* index_flags_ptr = index_flags->data<int>();
+  UniqueKernel<IntT><<<config.block_per_grid,
+                       config.thread_per_block,
+                       cache_size,
+                       dev_ctx.stream()>>>(rulebook_ptr,
+                                           rulebook_len_tensor.data<int>(),
+                                           index_flags_ptr,
+                                           out_index_ptr,
+                                           unique_key_ptr);
+
+  backends::gpu::GpuMemcpyAsync(d_buffer.data<int>(),
+                                counter_ptr,
+                                kernel_size * sizeof(int),
+                                gpuMemcpyDeviceToDevice,
+                                dev_ctx.stream());
+  backends::gpu::GpuMemcpyAsync(d_buffer.data<int>() + kernel_size,
+                                offsets_ptr,
+                                kernel_size * sizeof(int),
+                                gpuMemcpyDeviceToDevice,
+                                dev_ctx.stream());
+  backends::gpu::GpuMemcpyAsync(d_buffer.data<int>() + 2 * kernel_size + 1,
+                                rulebook_len_tensor.data<int>(),
+                                sizeof(int),
+                                gpuMemcpyDeviceToDevice,
+                                dev_ctx.stream());
+  backends::gpu::GpuMemcpyAsync(d_buffer.data<int>() + 2 * kernel_size + 2,
+                                unique_key_ptr,
+                                sizeof(int),
+                                gpuMemcpyDeviceToDevice,
+                                dev_ctx.stream());
+  backends::gpu::GpuMemcpyAsync(h_buffer,
+                                d_buffer.data<int>(),
+                                (2 * kernel_size + 3) * sizeof(int),
+                                gpuMemcpyDeviceToHost,
+                                dev_ctx.stream());
+
+  dev_ctx.Wait();
+  int rulebook_len = h_buffer[2 * kernel_size + 1] / 2;
+  int out_nnz = h_buffer[2 * kernel_size + 2];
+
+  rulebook->Resize({rulebook_rows, static_cast<int>(rulebook_len)});
+  out_index->Resize({static_cast<int>(rulebook_len)});
+
+  const int threads = 256;
+  const int blocks = (index_flags->numel() + threads - 1) / threads;
+  int* out_index_table_ptr = out_index_table->data<int>();
+
+  GetOutIndicesCounter<<<blocks, threads, 0, dev_ctx.stream()>>>(
+      index_flags_ptr, index_flags->numel(), out_index_table_ptr);
+
+  size_t temp_size1 = 0;
+  cub::DeviceScan::ExclusiveSum(NULL,
+                                temp_size1,
+                                out_index_table_ptr,
+                                out_index_table_ptr,
+                                blocks,
+                                dev_ctx.stream());
+
+  phi::Allocator* allocator1 =
+      const_cast<phi::Allocator*>(&(dev_ctx.GetAllocator()));
+  auto ws1 = allocator->Allocate(temp_size)->ptr();
+  cub::DeviceScan::ExclusiveSum(ws1,
+                                temp_size1,
+                                out_index_table_ptr,
+                                out_index_table_ptr,
+                                blocks,
+                                dev_ctx.stream());
+
+  GetOutIndices<threads>
+      <<<blocks, threads, 0, dev_ctx.stream()>>>(index_flags_ptr,
+                                                 index_flags->numel(),
+                                                 out_index_table_ptr,
+                                                 out_nnz,
+                                                 out_index_ptr);
+
+  const int64_t sparse_dim = is2D ? 3 : 4;
+  DenseTensor out_indices = Empty<IntT>(dev_ctx, {sparse_dim, out_nnz});
+
+  DenseTensor out_values =
+      Empty<T>(dev_ctx, {out_nnz, kernel_sizes[sparse_dim]});
+  out->SetMember(out_indices, out_values, out_dims, false);
+
+  IntT* out_indices_ptr = out_indices.data<IntT>();
+
+  config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_nnz, 1);
+  GetOutIndexTable<IntT>
+      <<<config.block_per_grid, config.thread_per_block, 0, dev_ctx.stream()>>>(
+          out_index_ptr,
+          out_nnz,
+          d_out_dims,
+          is2D,
+          out_index_table_ptr,
+          out_indices_ptr);
+
+  config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rulebook_len, 1);
+  unique_value->ResizeAndAllocate({static_cast<int>(out_nnz * kernel_size)});
+  int* unique_value_ptr = unique_value->data<int>();
+
+  GroupIndices<<<config.block_per_grid,
+                 config.thread_per_block,
+                 0,
+                 dev_ctx.stream()>>>(out_index_table_ptr,
+                                     rulebook_len,
+                                     kernel_size,
+                                     rulebook_ptr + rulebook_len,
+                                     out_index_ptr,
+                                     unique_value_ptr);
+
+  return rulebook_len;
+}
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/gpu/convolution.cu.h

@@ -0,0 +1,550 @@
+/* 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 <thrust/binary_search.h>
+#include <thrust/execution_policy.h>
+#include <thrust/remove.h>
+#include <thrust/sort.h>
+#include <thrust/unique.h>
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/index_impl.cu.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/sparse/utils.cu.h"
+#include "paddle/phi/kernels/primitive/compute_primitives.h"
+#include "paddle/phi/kernels/sparse/conv_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+using Dims4D = funcs::sparse::Dims4D;
+
+// TODO(zhangkaihuo): After the GatherCUDAKernel is migrated to phi, replace
+// this kernel with phi::GatherCUDAKernel;
+// Vectorization can be used to improve read and write bandwidth
+/**
+ * brief: gather data from params according to indices
+ * params: the inputs
+ * indices: the indices you want to gather
+ * output: the outputs
+ * index_size: the size of indices
+ * slice_size: slice size corresponding to each index, here is the channel size
+ **/
+template <typename T, typename IndexT = int>
+__global__ void GatherKernel(const T* params,
+                             const IndexT* indices,
+                             T* output,
+                             size_t index_size,
+                             size_t slice_size) {
+  CUDA_KERNEL_LOOP_TYPE(i, index_size * slice_size, int64_t) {
+    int64_t indices_i = i / slice_size;
+    int64_t slice_i = i - indices_i * slice_size;  // offset inside the slice
+    IndexT gather_i = indices[indices_i];
+    int64_t params_i = gather_i * slice_size + slice_i;
+    *(output + i) = *(params + params_i);
+  }
+}
+
+template <typename Context, typename IntT = int>
+inline IntT* SortedAndUniqueIndex(const Context& dev_ctx,
+                                  const IntT* rulebook_ptr,
+                                  const int len,
+                                  DenseTensor* out_index,
+                                  DenseTensor* unique_key,
+                                  DenseTensor* unique_value) {
+  phi::IndexKernel<int, kps::IdentityFunctor<int>>(
+      dev_ctx, out_index, kps::IdentityFunctor<int>());
+  phi::IndexKernel<int, kps::IdentityFunctor<int>>(
+      dev_ctx, unique_value, kps::IdentityFunctor<int>());
+
+  backends::gpu::GpuMemcpyAsync(unique_key->data<IntT>(),
+                                rulebook_ptr,
+                                sizeof(IntT) * len,
+#ifdef PADDLE_WITH_HIP
+                                hipMemcpyDeviceToDevice,
+#else
+                                cudaMemcpyDeviceToDevice,
+#endif
+                                dev_ctx.stream());
+// compared with thrust::sort_by_key, thrust::merge_by_key may achieved higher
+// performance, but thrust::merge_by_key limited by data size
+#ifdef PADDLE_WITH_HIP
+  thrust::sort_by_key(thrust::hip::par.on(dev_ctx.stream()),
+#else
+  thrust::sort_by_key(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                      unique_key->data<IntT>(),
+                      unique_key->data<IntT>() + len,
+                      out_index->data<int>());
+
+  // 4. unique
+  thrust::pair<IntT*, int*> new_end =
+#ifdef PADDLE_WITH_HIP
+      thrust::unique_by_key(thrust::hip::par.on(dev_ctx.stream()),
+#else
+      thrust::unique_by_key(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                            unique_key->data<IntT>(),
+                            unique_key->data<IntT>() + len,
+                            unique_value->data<int>());
+  return new_end.first;
+}
+
+/**
+ * @brief: update the out index and indices
+ * unique_keys: save the index of the output feature list
+ * unique_values: indicates the index of key before deduplication
+ * out_indexes: indicates the position of the output index in the rulebook
+ * rulebook_len: indicates the length of rulebook
+ * out_dims: indicates the output dims
+ * out_indices: the indices of output, out_indices = IndexToPoint(unique_keys)
+ * rulebook_out_indices: the output index in rulebook
+ **/
+template <typename T>
+__global__ void UpdateIndexKernel(const T* unique_keys,
+                                  const int* unique_values,
+                                  const int* out_indexes,
+                                  const int64_t non_zero_num,
+                                  const int rulebook_len,
+                                  const Dims4D out_dims,
+                                  T* out_indices,
+                                  T* rulebook_out_indices) {
+  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  for (int i = tid; i < non_zero_num; i += gridDim.x * blockDim.x) {
+    const T index = unique_keys[i];
+    T batch, x, y, z;
+    funcs::sparse::IndexToPoint<Dims4D>(index, out_dims, &batch, &x, &y, &z);
+    // get out indices
+    out_indices[i] = batch;
+    out_indices[i + non_zero_num] = z;
+    out_indices[i + non_zero_num * 2] = y;
+    out_indices[i + non_zero_num * 3] = x;
+
+    // update rulebook
+    int start = unique_values[i];
+    int end = i == non_zero_num - 1 ? rulebook_len : unique_values[i + 1];
+    // max(end-start) = kernel_size
+    for (T j = start; j < end; j++) {
+      rulebook_out_indices[out_indexes[j]] = i;
+    }
+  }
+}
+
+template <typename IntT>
+__global__ void UpdateOutIndexAndCounterAfterLowerBound(
+    const IntT* x_indices,
+    const IntT* bound_out,
+    const int rulebook_len,
+    const int kernel_size,
+    const int64_t non_zero_num,
+    IntT* rulebook_ptr,
+    IntT* out_indices,
+    int* counter_ptr) {
+  extern __shared__ int cache_count[];
+  for (int i = threadIdx.x; i < kernel_size; i += blockDim.x) {
+    cache_count[i] = 0;
+  }
+  __syncthreads();
+
+  CUDA_KERNEL_LOOP_TYPE(i, rulebook_len, int64_t) {
+    int j = bound_out[i];
+    if (j >= 0 && j < non_zero_num && out_indices[i] == x_indices[j]) {
+      out_indices[i] = j;
+    } else {
+      // mask this position will be remove
+      int kernel_index = rulebook_ptr[i];
+      rulebook_ptr[i + rulebook_len] = -1;
+      rulebook_ptr[i + 2 * rulebook_len] = -1;
+      rulebook_ptr[i] = -1;
+      atomicAdd(&cache_count[kernel_index], 1);
+    }
+  }
+  __syncthreads();
+
+  for (int i = threadIdx.x; i < kernel_size; i += blockDim.x) {
+    atomicSub(&counter_ptr[i], cache_count[i]);
+  }
+}
+
+/**
+ * @brief product rulebook
+ * for input_i in x_indices:
+ *   if input_i participate in the convolution calculation:
+ *       infer the output_i by input_i and kernel_i
+ *       save output_i
+ *
+ * x_indices: the indices of input features
+ * x_dims: the input dims
+ * kernel_dims: the kernel dims
+ * out_dims: the output dims
+ * non_zero_num: the number of input features
+ * rulebook: the rulebook to save the kernel index, input index and output index
+ * counter: save the number of times each location in the kernel participates in
+ * the calculation
+ **/
+template <typename T>
+__global__ void ProductRuleBookKernel(const T* x_indices,
+                                      const Dims4D x_dims,
+                                      const Dims4D kernel_dims,
+                                      const Dims4D out_dims,
+                                      const int64_t non_zero_num,
+                                      const Dims4D paddings,
+                                      const Dims4D dilations,
+                                      const Dims4D strides,
+                                      const bool subm,
+                                      T* rulebook,
+                                      int* counter,
+                                      T* in_indices) {
+  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  extern __shared__ int counter_buf[];  // kernel_size
+  const int kernel_size = kernel_dims[3] * kernel_dims[2] * kernel_dims[1];
+  const int offset = kernel_size * non_zero_num;
+  for (int i = threadIdx.x; i < kernel_size; i += blockDim.x) {
+    counter_buf[i] = 0;
+  }
+  __syncthreads();
+
+  for (int i = tid; i < non_zero_num; i += gridDim.x * blockDim.x) {
+    int kernel_index = 0;
+    T batch = x_indices[i];
+    T in_z = x_indices[i + non_zero_num];
+    T in_y = x_indices[i + 2 * non_zero_num];
+    T in_x = x_indices[i + 3 * non_zero_num];
+    if (subm) {
+      in_indices[i] = PointToIndex(batch, in_x, in_y, in_z, x_dims);
+    }
+    for (int kz = 0; kz < kernel_dims[1]; kz++) {
+      for (int ky = 0; ky < kernel_dims[2]; ky++) {
+        for (int kx = 0; kx < kernel_dims[3]; kx++) {
+          int in_i = -1, out_index = -1, kernel_i = -1;
+          if (funcs::sparse::Check(x_dims,
+                                   kernel_dims,
+                                   paddings,
+                                   dilations,
+                                   strides,
+                                   in_x,
+                                   in_y,
+                                   in_z,
+                                   kx,
+                                   ky,
+                                   kz)) {
+            T out_z = (in_z + paddings[1] - kz * dilations[1]) / strides[1];
+            T out_y = (in_y + paddings[2] - ky * dilations[2]) / strides[2];
+            T out_x = (in_x + paddings[3] - kx * dilations[3]) / strides[3];
+            in_i = i;
+            out_index = funcs::sparse::PointToIndex<Dims4D>(
+                batch, out_x, out_y, out_z, out_dims);
+            atomicAdd(&counter_buf[kernel_index], 1);
+            kernel_i = kernel_index;
+          }
+          rulebook[kernel_index * non_zero_num + i] = kernel_i;
+          rulebook[kernel_index * non_zero_num + offset + i] = in_i;
+          rulebook[kernel_index * non_zero_num + offset * 2 + i] = out_index;
+          ++kernel_index;
+        }
+      }
+    }
+  }
+  __syncthreads();
+  for (int i = threadIdx.x; i < kernel_size; i += blockDim.x) {
+    atomicAdd(&counter[i], counter_buf[i]);
+  }
+}
+
+// the basic algorithm can refer to convolution_kernel.cc or
+// the second paper
+// example:
+// 1. the rulebook:
+//  the kernel_index:                       0, 0, 0, 1, 1, 1, 2, 2, ....
+//  the out_index(key):                     20, 30, 33, 30, 33, 20, 25
+// 2. mark the index of out_index(value):   0, 1, 2, 3, 4, 5, 6, ....
+// 3. sorted the (key, value)
+// 4. unique the (key, value):
+//  unique_key:     20, 25, 30, 33
+//  unique_values:  0, 2, 3, 5
+//  the index of unique_values is: 0, 1, 2, 3
+// 5. update the out_index by unique_key, unique_value and the index of
+// unique_value:
+//  the new out_index: 0, 2, 3, 2, 3, 0, 1
+template <typename T, typename Context, typename IntT = int>
+int 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,
+                    DenseTensor* counter_per_kernel,
+                    DenseTensor* offsets_per_kernel,
+                    DenseTensor* out_index,
+                    DenseTensor* unique_value,
+                    SparseCooTensor* out,
+                    std::vector<int>* h_counter,
+                    std::vector<int>* h_offsets) {
+  auto indices_dtype = phi::CppTypeToDataType<IntT>::Type();
+  const int64_t non_zero_num = x.nnz();
+  const auto& indices = x.indices();
+  const IntT* indices_ptr = indices.data<IntT>();
+  DenseTensor in_indices = Empty<Context>(
+      dev_ctx, DenseTensorMeta(indices_dtype, {x.nnz()}, DataLayout::NCHW));
+  int* counter_ptr = counter_per_kernel->data<int>();
+  int* offsets_ptr = offsets_per_kernel->data<int>();
+  int kernel_size = kernel_sizes[0] * kernel_sizes[1] * kernel_sizes[2];
+  const int rulebook_rows = 3;
+  const int rulebook_cols = kernel_size * non_zero_num;
+  DenseTensorMeta rulebook_meta(
+      indices_dtype, {rulebook_rows, rulebook_cols}, DataLayout::NCHW);
+  *rulebook = Empty(dev_ctx, std::move(rulebook_meta));
+  IntT* rulebook_ptr = rulebook->data<IntT>();
+
+  const auto x_dims = x.dims();
+  Dims4D d_x_dims(x_dims[0], x_dims[3], x_dims[2], x_dims[1]);
+  Dims4D d_kernel_dims(1, kernel_sizes[2], kernel_sizes[1], kernel_sizes[0]);
+  Dims4D d_out_dims(out_dims[0], out_dims[3], out_dims[2], out_dims[1]);
+  Dims4D d_paddings(1, paddings[2], paddings[1], paddings[0]);
+  Dims4D d_strides(1, strides[2], strides[1], strides[0]);
+  Dims4D d_dilations(1, dilations[2], dilations[1], dilations[0]);
+  // 1. product rule book
+  funcs::SetConstant<Context, int> set_zero;
+  set_zero(dev_ctx, counter_per_kernel, 0);
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num, 1);
+
+  ProductRuleBookKernel<IntT><<<config.block_per_grid.x,
+                                config.thread_per_block.x,
+                                kernel_size * sizeof(int),
+                                dev_ctx.stream()>>>(indices_ptr,
+                                                    d_x_dims,
+                                                    d_kernel_dims,
+                                                    d_out_dims,
+                                                    non_zero_num,
+                                                    d_paddings,
+                                                    d_dilations,
+                                                    d_strides,
+                                                    subm,
+                                                    rulebook_ptr,
+                                                    counter_ptr,
+                                                    in_indices.data<IntT>());
+
+// 2. remove -1
+#ifdef PADDLE_WITH_HIP
+  IntT* last = thrust::remove(thrust::hip::par.on(dev_ctx.stream()),
+#else
+  IntT* last = thrust::remove(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                              rulebook_ptr,
+                              rulebook_ptr + rulebook_rows * rulebook_cols,
+                              -1);
+
+  funcs::sparse::DistanceKernel<IntT><<<1, 1, 0, dev_ctx.stream()>>>(
+      rulebook_ptr, last, rulebook_ptr + 3 * kernel_size * non_zero_num - 1);
+  IntT rulebook_len = 0;
+  backends::gpu::GpuMemcpyAsync(
+      &rulebook_len,
+      rulebook_ptr + 3 * kernel_size * non_zero_num - 1,
+      sizeof(IntT),
+#ifdef PADDLE_WITH_HIP
+      hipMemcpyDeviceToHost,
+#else
+      cudaMemcpyDeviceToHost,
+#endif
+      dev_ctx.stream());
+  dev_ctx.Wait();
+  rulebook_len /= 3;
+
+  if (subm) {
+    // At present, hashtable is not used to map the input and output indexes.
+    // At present, the intermediate output index is generated by normal
+    // convolution,
+    // and then the intermediate output index is subtracted from the input index
+    // to obtain the rulebook.
+
+    // call lower_bound to get the real index of out_index
+    const IntT* in_indices_ptr = in_indices.data<IntT>();
+    IntT* out_indices_ptr = rulebook_ptr + 2 * rulebook_len;
+    DenseTensor bound = Empty(
+        dev_ctx,
+        DenseTensorMeta(
+            indices_dtype, {static_cast<int>(rulebook_len)}, DataLayout::NCHW));
+    IntT* bound_ptr = bound.data<IntT>();
+#ifdef PADDLE_WITH_HIP
+    thrust::lower_bound(thrust::hip::par.on(dev_ctx.stream()),
+#else
+    thrust::lower_bound(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                        in_indices_ptr,
+                        in_indices_ptr + in_indices.numel(),
+                        out_indices_ptr,
+                        out_indices_ptr + rulebook_len,
+                        bound_ptr);
+
+    config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rulebook_len, 1);
+
+    UpdateOutIndexAndCounterAfterLowerBound<<<config.block_per_grid,
+                                              config.thread_per_block,
+                                              kernel_size * sizeof(int),
+                                              dev_ctx.stream()>>>(
+        in_indices_ptr,
+        bound.data<IntT>(),
+        rulebook_len,
+        kernel_size,
+        x.nnz(),
+        rulebook_ptr,
+        out_indices_ptr,
+        counter_ptr);
+
+// remove -1
+#ifdef PADDLE_WITH_HIP
+    IntT* last = thrust::remove(thrust::hip::par.on(dev_ctx.stream()),
+#else
+    IntT* last = thrust::remove(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                                rulebook_ptr,
+                                rulebook_ptr + 3 * rulebook_len,
+                                -1);
+    funcs::sparse::DistanceKernel<IntT>
+        <<<1, 1, 0, dev_ctx.stream()>>>(rulebook_ptr, last, bound_ptr);
+    backends::gpu::GpuMemcpyAsync(&rulebook_len,
+                                  bound_ptr,
+                                  sizeof(IntT),
+#ifdef PADDLE_WITH_HIP
+                                  hipMemcpyDeviceToHost,
+#else
+                                  cudaMemcpyDeviceToHost,
+#endif
+                                  dev_ctx.stream());
+    dev_ctx.Wait();
+    rulebook_len /= 3;
+  }
+
+#ifdef PADDLE_WITH_HIP
+  thrust::exclusive_scan(thrust::hip::par.on(dev_ctx.stream()),
+#else
+  thrust::exclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                         counter_ptr,
+                         counter_ptr + kernel_size,
+                         offsets_ptr);
+
+  backends::gpu::GpuMemcpyAsync(&(*h_counter)[0],
+                                counter_ptr,
+                                kernel_size * sizeof(int),
+#ifdef PADDLE_WITH_HIP
+                                hipMemcpyDeviceToHost,
+#else
+                                cudaMemcpyDeviceToHost,
+#endif
+                                dev_ctx.stream());
+
+  backends::gpu::GpuMemcpyAsync(&(*h_offsets)[0],
+                                offsets_ptr,
+                                kernel_size * sizeof(int),
+#ifdef PADDLE_WITH_HIP
+                                hipMemcpyDeviceToHost,
+#else
+                                cudaMemcpyDeviceToHost,
+#endif
+                                dev_ctx.stream());
+
+  rulebook->Resize({rulebook_rows, static_cast<int>(rulebook_len)});
+
+  if (!subm) {
+    // 3. sorted or merge the out index
+    out_index->ResizeAndAllocate({static_cast<int>(rulebook_len)});
+    unique_value->ResizeAndAllocate({static_cast<int>(rulebook_len)});
+    DenseTensor unique_key = Empty(
+        dev_ctx,
+        DenseTensorMeta(
+            indices_dtype, {static_cast<int>(rulebook_len)}, DataLayout::NCHW));
+    int* out_index_ptr = out_index->data<int>();
+    int* unique_value_ptr = unique_value->data<int>();
+    IntT* unique_key_ptr = unique_key.data<IntT>();
+
+    IntT* new_end =
+        SortedAndUniqueIndex<Context, IntT>(dev_ctx,
+                                            rulebook_ptr + 2 * rulebook_len,
+                                            rulebook_len,
+                                            out_index,
+                                            &unique_key,
+                                            unique_value);
+    // thrust::distance doesn't support stream parameters
+    // const int out_non_zero_num = thrust::distance(unique_key_ptr,
+    // new_end.first);
+    funcs::sparse::DistanceKernel<IntT><<<1, 1, 0, dev_ctx.stream()>>>(
+        unique_key_ptr,
+        new_end,
+        rulebook_ptr + rulebook_rows * rulebook_cols - 1);
+    IntT out_non_zero_num = 0;
+#ifdef PADDLE_WITH_HIP
+    backends::gpu::GpuMemcpyAsync(
+        &out_non_zero_num,
+        rulebook_ptr + rulebook_rows * rulebook_cols - 1,
+        sizeof(IntT),
+        hipMemcpyDeviceToHost,
+        dev_ctx.stream());
+#else
+    backends::gpu::GpuMemcpyAsync(
+        &out_non_zero_num,
+        rulebook_ptr + rulebook_rows * rulebook_cols - 1,
+        sizeof(IntT),
+        cudaMemcpyDeviceToHost,
+        dev_ctx.stream());
+#endif
+    dev_ctx.Wait();
+
+    // 5. update out_indices and rulebook by unique_value_ptr
+    const int64_t sparse_dim = 4;
+    DenseTensorMeta indices_meta(
+        indices_dtype, {sparse_dim, out_non_zero_num}, DataLayout::NCHW);
+    DenseTensorMeta values_meta(
+        x.dtype(), {out_non_zero_num, kernel_sizes[4]}, 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>();
+
+    config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_non_zero_num, 1);
+    UpdateIndexKernel<IntT>
+        <<<config.block_per_grid.x,
+           config.thread_per_block.x,
+           0,
+           dev_ctx.stream()>>>(unique_key_ptr,
+                               unique_value_ptr,
+                               out_index_ptr,
+                               out_non_zero_num,
+                               rulebook_len,
+                               d_out_dims,
+                               out_indices_ptr,
+                               rulebook_ptr + 2 * rulebook_len);
+    out->SetMember(out_indices, out_values, out_dims, true);
+  } else {
+    DenseTensor out_indices = EmptyLike<IntT>(dev_ctx, x.indices());
+    DenseTensor out_values =
+        Empty(dev_ctx,
+              DenseTensorMeta(
+                  x.dtype(), {x.nnz(), kernel_sizes[4]}, x.values().layout()));
+    phi::Copy(dev_ctx, x.indices(), dev_ctx.GetPlace(), false, &out_indices);
+    out->SetMember(out_indices, out_values, out_dims, true);
+  }
+  return rulebook_len;
+}
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/gpu/cutlass_generator/gather_gemm_scatter_generator.py

@@ -0,0 +1,666 @@
+# 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.
+
+import sys
+
+sys.path.append(sys.argv[1])
+from gather_gemm_scatter_manifest import GatherGemmScatterManifest
+from gather_gemm_scatter_operation import GatherGemmScatterOperation
+from generator import (
+    ComplexTransform,
+    CudaToolkitVersionSatisfies,
+    EpilogueFunctor,
+    GemmKind,
+    SwizzlingFunctor,
+    TensorDescription,
+)
+from library import (
+    DataType,
+    LayoutType,
+    MathInstruction,
+    MathOperation,
+    OpcodeClass,
+    TileDescription,
+)
+from manifest import GeneratorTarget
+
+
+def CreateGatherGemmScatterOperator(
+    manifest,
+    layouts,
+    tile_descriptions,
+    data_type,
+    complex_transforms=None,
+    epilogue_functor=EpilogueFunctor.LinearCombination,
+    swizzling_functor=SwizzlingFunctor.Identity8,
+):
+    # To use StreamK decomposition for basic GEMMs, set `swizzling_functor = SwizzlingFunctor.StreamK`
+
+    if complex_transforms is None:
+        complex_transforms = [
+            (ComplexTransform.none, ComplexTransform.none),
+        ]
+
+    element_a, element_b, element_c, element_epilogue = data_type
+
+    alignment_constraints = [0]
+    if 'f16' == element_a.name or 'bf16' == element_a.name:
+        alignment_constraints = [8]
+    elif 'f32' == element_a.name or 'tf32' == element_a.name:
+        alignment_constraints = [4]
+    elif 'f64' == element_a.name:
+        alignment_constraints = [1]
+
+    operations = []
+
+    for layout in layouts:
+        for tile_description in tile_descriptions:
+            for alignment in alignment_constraints:
+                for complex_transform in complex_transforms:
+                    alignment_c = min(8, alignment)
+
+                    A = TensorDescription(
+                        element_a, layout[0], alignment, complex_transform[0]
+                    )
+                    B = TensorDescription(
+                        element_b, layout[1], alignment, complex_transform[1]
+                    )
+                    C = TensorDescription(element_c, layout[2], alignment_c)
+
+                    new_operation = GatherGemmScatterOperation(
+                        GemmKind.Universal,
+                        tile_description.minimum_compute_capability,
+                        tile_description,
+                        A,
+                        B,
+                        C,
+                        element_epilogue,
+                        epilogue_functor,
+                        swizzling_functor,
+                    )
+
+                    manifest.append(new_operation)
+                    operations.append(new_operation)
+
+    return operations
+
+
+def GenerateSM80_TensorOp_16816(manifest, cuda_version, debug=False):
+    if not CudaToolkitVersionSatisfies(cuda_version, 11, 0):
+        return
+
+    layouts = [
+        (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.RowMajor),
+    ]
+
+    math_instructions = [
+        MathInstruction(
+            [16, 8, 16],
+            DataType.f16,
+            DataType.f16,
+            DataType.f16,
+            OpcodeClass.TensorOp,
+            MathOperation.multiply_add,
+        ),
+    ]
+
+    min_cc = 80
+    max_cc = 1024
+
+    alignment_constraints = [8]
+
+    for math_inst in math_instructions:
+        tile_descriptions = [
+            TileDescription(
+                [256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 64, 32], 3, [4, 1, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 32], 5, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 64, 32], 6, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 128, 32], 6, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 64, 32], 10, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 64, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 256, 64], 4, [1, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 64, 64], 3, [4, 1, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 256, 64], 3, [1, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 64, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 128, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 64, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+        ]
+        if debug:
+            tile_descriptions = [
+                TileDescription(
+                    [256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc
+                ),
+            ]
+
+        data_type = [
+            math_inst.element_a,
+            math_inst.element_b,
+            math_inst.element_accumulator,
+            math_inst.element_accumulator,
+        ]
+
+        CreateGatherGemmScatterOperator(
+            manifest, layouts, tile_descriptions, data_type
+        )
+
+        # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation)
+        if math_inst.element_a != math_inst.element_accumulator:
+            data_type_mixed = [
+                math_inst.element_a,
+                math_inst.element_b,
+                math_inst.element_a,
+                math_inst.element_accumulator,
+            ]
+
+            CreateGatherGemmScatterOperator(
+                manifest, layouts, tile_descriptions, data_type_mixed
+            )
+
+
+def GenerateSM80_TensorOp_1688(manifest, cuda_version, debug=False):
+    if not CudaToolkitVersionSatisfies(cuda_version, 11, 0):
+        return
+
+    layouts = [
+        (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.RowMajor),
+        (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+        (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.RowMajor),
+    ]
+
+    math_instructions = [
+        MathInstruction(
+            [16, 8, 8],
+            DataType.tf32,
+            DataType.tf32,
+            DataType.f32,
+            OpcodeClass.TensorOp,
+            MathOperation.multiply_add,
+        )
+    ]
+
+    min_cc = 80
+    max_cc = 1024
+
+    for math_inst in math_instructions:
+        tile_descriptions = [
+            TileDescription(
+                [256, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 64, 16], 4, [4, 1, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 256, 16], 4, [1, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 64, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 128, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 64, 16], 10, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 64, 32], 5, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+        ]
+
+        if debug:
+            tile_descriptions = [
+                TileDescription(
+                    [256, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc
+                ),
+            ]
+
+        data_type = [
+            math_inst.element_a,
+            math_inst.element_b,
+            math_inst.element_accumulator,
+            math_inst.element_accumulator,
+        ]
+
+        data_type_mixed = [
+            math_inst.element_a,
+            math_inst.element_b,
+            math_inst.element_a,
+            math_inst.element_accumulator,
+        ]
+
+        CreateGatherGemmScatterOperator(
+            manifest, layouts, tile_descriptions, data_type
+        )
+
+        CreateGatherGemmScatterOperator(
+            manifest, layouts, tile_descriptions, data_type_mixed
+        )
+
+
+def GenerateSM80_TensorOp_1688_fast_math(manifest, cuda_version, debug=False):
+    if not CudaToolkitVersionSatisfies(cuda_version, 11, 0):
+        return
+
+    layouts = [
+        (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.RowMajor),
+        (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+        (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.RowMajor),
+    ]
+
+    math_instructions = [
+        MathInstruction(
+            [16, 8, 8],
+            DataType.tf32,
+            DataType.tf32,
+            DataType.f32,
+            OpcodeClass.TensorOp,
+            MathOperation.multiply_add,
+        ),
+    ]
+
+    min_cc = 80
+    max_cc = 1024
+
+    for math_inst in math_instructions:
+        tile_descriptions = [
+            TileDescription(
+                [256, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 64, 16], 4, [4, 1, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 256, 16], 4, [1, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 64, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 128, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 64, 16], 10, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 64, 32], 5, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+        ]
+
+        if debug:
+            tile_descriptions = [
+                TileDescription(
+                    [256, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc
+                ),
+            ]
+
+        data_type = [DataType.f32, DataType.f32, DataType.f32, DataType.f32]
+
+        CreateGatherGemmScatterOperator(
+            manifest, layouts, tile_descriptions, data_type
+        )
+
+
+def GenerateSM80_TensorOp_1688_fast_fp32_math(
+    manifest, cuda_version, debug=False
+):
+    if not CudaToolkitVersionSatisfies(cuda_version, 11, 0):
+        return
+
+    layouts = [
+        (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.RowMajor),
+        (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+        (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.RowMajor),
+    ]
+
+    math_instructions = [
+        MathInstruction(
+            [16, 8, 8],
+            DataType.f32,
+            DataType.f32,
+            DataType.f32,
+            OpcodeClass.TensorOp,
+            MathOperation.multiply_add_fast_f32,
+        ),
+    ]
+
+    min_cc = 80
+    max_cc = 1024
+
+    for math_inst in math_instructions:
+        tile_descriptions = [
+            TileDescription(
+                [128, 128, 16], 4, [4, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 64, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 128, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 64, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+        ]
+
+        if debug:
+            tile_descriptions = [
+                TileDescription(
+                    [128, 128, 16], 4, [4, 2, 1], math_inst, min_cc, max_cc
+                ),
+            ]
+
+        data_type = [DataType.f32, DataType.f32, DataType.f32, DataType.f32]
+
+        CreateGatherGemmScatterOperator(
+            manifest, layouts, tile_descriptions, data_type
+        )
+
+
+def GenerateSM75_TensorOp_1688(manifest, cuda_version, debug=False):
+    if not CudaToolkitVersionSatisfies(cuda_version, 10, 2):
+        return
+
+    layouts = [
+        (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.RowMajor),
+    ]
+
+    math_instructions = [
+        MathInstruction(
+            [16, 8, 8],
+            DataType.f16,
+            DataType.f16,
+            DataType.f32,
+            OpcodeClass.TensorOp,
+            MathOperation.multiply_add,
+        ),
+        MathInstruction(
+            [16, 8, 8],
+            DataType.f16,
+            DataType.f16,
+            DataType.f16,
+            OpcodeClass.TensorOp,
+            MathOperation.multiply_add,
+        ),
+    ]
+
+    min_cc = 75
+    max_cc = 1024
+
+    for math_inst in math_instructions:
+        tile_descriptions = [
+            TileDescription(
+                [256, 128, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 256, 32], 2, [2, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 256, 32], 2, [1, 4, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [256, 64, 32], 2, [4, 1, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [128, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc
+            ),
+            TileDescription(
+                [64, 128, 64], 2, [1, 2, 2], math_inst, min_cc, max_cc
+            ),
+        ]
+
+        if debug:
+            tile_descriptions = [
+                TileDescription(
+                    [256, 128, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc
+                ),
+            ]
+
+        data_type = [
+            math_inst.element_a,
+            math_inst.element_b,
+            math_inst.element_accumulator,
+            math_inst.element_accumulator,
+        ]
+
+        CreateGatherGemmScatterOperator(
+            manifest, layouts, tile_descriptions, data_type
+        )
+
+
+def GenerateSM75(manifest, cuda_version, debug=False):
+    GenerateSM75_TensorOp_1688(manifest, cuda_version, debug)
+
+
+def GenerateSM80(manifest, cuda_version, debug=False):
+    GenerateSM80_TensorOp_16816(manifest, cuda_version, debug)
+    GenerateSM80_TensorOp_1688(manifest, cuda_version, debug)
+    GenerateSM80_TensorOp_1688_fast_math(manifest, cuda_version, debug)
+    GenerateSM80_TensorOp_1688_fast_fp32_math(manifest, cuda_version, debug)
+
+
+class KernelCfg:
+    def __init__(
+        self,
+        architectures,
+        build_dir,
+        cuda_version,
+        curr_build_dir,
+        disable_full_archs_compilation,
+        filter_by_cc,
+        generator_target,
+        ignore_kernels,
+        interface_dir,
+        kernel_filter_file,
+        kernels,
+        operations,
+        selected_kernel_list,
+    ):
+        self.architectures = architectures
+        self.build_dir = build_dir
+        self.cuda_version = cuda_version
+        self.curr_build_dir = curr_build_dir
+        self.disable_full_archs_compilation = disable_full_archs_compilation
+        self.filter_by_cc = filter_by_cc
+        self.generator_target = generator_target
+        self.ignore_kernels = ignore_kernels
+        self.interface_dir = interface_dir
+        self.kernel_filter_file = kernel_filter_file
+        self.kernels = kernels
+        self.operations = operations
+        self.selected_kernel_list = selected_kernel_list
+
+
+if __name__ == "__main__":
+    args = KernelCfg(
+        architectures='80',
+        build_dir=sys.argv[2],
+        cuda_version=sys.argv[3],
+        curr_build_dir=sys.argv[2],
+        disable_full_archs_compilation=False,
+        filter_by_cc='True',
+        generator_target='library',
+        ignore_kernels='',
+        interface_dir=None,
+        kernel_filter_file=None,
+        kernels='',
+        operations='all',
+        selected_kernel_list=None,
+    )
+    manifest = GatherGemmScatterManifest(args)
+
+    debug = False
+    GenerateSM75(manifest, args.cuda_version, debug)
+    GenerateSM80(manifest, args.cuda_version, debug)
+
+    manifest.emit(GeneratorTarget.Library)

paddle/phi/kernels/sparse/gpu/elementwise_grad_kernel.cu

@@ -0,0 +1,75 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/phi/kernels/sparse/elementwise_grad_kernel.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/elementwise_grad_base.h"
+#include "paddle/phi/kernels/funcs/reduce_function.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context>
+void ElementWiseAddCooGradKernel(const Context& dev_ctx,
+                                 const SparseCooTensor& x,
+                                 const SparseCooTensor& y,
+                                 const SparseCooTensor& dout,
+                                 SparseCooTensor* dx,
+                                 SparseCooTensor* dy) {
+  if (dx) {
+    EmptyLikeCooKernel<T, Context>(dev_ctx, x, dx);
+    Copy(dev_ctx, dout, dev_ctx.GetPlace(), false, dx);
+  }
+
+  if (dy) {
+    EmptyLikeCooKernel<T, Context>(dev_ctx, y, dy);
+    Copy(dev_ctx, dout, dev_ctx.GetPlace(), false, dy);
+  }
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(add_coo_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseAddCooGradKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::float16,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(add_coo_dense_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseAddDenseGradKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::float16,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/gpu/elementwise_kernel.cu

@@ -0,0 +1,109 @@
+/* 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 <glog/logging.h>
+#include <thrust/equal.h>
+#include <thrust/execution_policy.h>
+
+#include "paddle/phi/kernels/elementwise_add_kernel.h"
+#include "paddle/phi/kernels/sparse/elementwise_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/visit_type.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename IntT>
+void ElementWiseAddCooGPUKernel(const GPUContext& dev_ctx,
+                                const SparseCooTensor& x,
+                                const SparseCooTensor& y,
+                                SparseCooTensor* out) {
+  // TODO(zhangkaiuo): to support universal sparse + sparse
+  const auto& x_indices = x.indices();
+  const auto& y_indices = y.indices();
+  PADDLE_ENFORCE_EQ(
+      x_indices.numel(),
+      y_indices.numel(),
+      common::errors::PreconditionNotMet(
+          "The numel of x.indices() and y.indices() should be equal"));
+  const IntT* x_indices_ptr = x_indices.data<IntT>();
+  const IntT* y_indices_ptr = y_indices.data<IntT>();
+  if (VLOG_IS_ON(3)) {
+#ifdef PADDLE_WITH_HIP
+    bool is_same = thrust::equal(thrust::hip::par.on(dev_ctx.stream()),
+#else
+    bool is_same = thrust::equal(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                                 x_indices_ptr,
+                                 x_indices_ptr + x_indices.numel(),
+                                 y_indices_ptr);
+    PADDLE_ENFORCE_EQ(
+        is_same,
+        true,
+        common::errors::PreconditionNotMet(
+            "Currently, ElementWiseAddCooKernel only supports the case "
+            "where x and y have the same indices"));
+  }
+  EmptyLikeCooKernel<T, GPUContext>(dev_ctx, x, out);
+  phi::AddKernel<T, GPUContext>(
+      dev_ctx, x.values(), y.values(), out->mutable_values());
+  out->SetIndicesDict(x.GetIndicesDict());
+  out->SetKmaps(x.GetKmaps());
+}
+
+template <typename T, typename Context>
+void ElementWiseAddCooKernel(const Context& dev_ctx,
+                             const SparseCooTensor& x,
+                             const SparseCooTensor& y,
+                             SparseCooTensor* out) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "VerifyIndices", ([&] {
+                                 ElementWiseAddCooGPUKernel<T, data_t>(
+                                     dev_ctx, x, y, out);
+                               }));
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(add_coo_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseAddCooKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::float16,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(add_coo_dense,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ElementWiseAddDenseKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::float16,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/gpu/full_kernel.cu

@@ -0,0 +1,94 @@
+/* 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/gpu/gpu_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/full_kernel.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void FullLikeCooKernel(const Context& dev_ctx,
+                       const SparseCooTensor& x,
+                       const Scalar& val,
+                       DataType dtype,
+                       SparseCooTensor* out) {
+  phi::Copy<Context>(
+      dev_ctx, x.indices(), dev_ctx.GetPlace(), false, out->mutable_indices());
+
+  DenseTensor* values = out->mutable_values();
+  phi::Full<T, Context>(dev_ctx, vectorize(x.values().dims()), val, values);
+  out->set_dims(x.dims());
+}
+
+template <typename T, typename Context>
+void FullLikeCsrKernel(const Context& dev_ctx,
+                       const SparseCsrTensor& x,
+                       const Scalar& val,
+                       DataType dtype,
+                       SparseCsrTensor* out) {
+  phi::Copy<Context>(
+      dev_ctx, x.crows(), dev_ctx.GetPlace(), false, out->mutable_crows());
+
+  phi::Copy<Context>(
+      dev_ctx, x.cols(), dev_ctx.GetPlace(), false, out->mutable_cols());
+
+  DenseTensor* values = out->mutable_values();
+  phi::Full<T, Context>(dev_ctx, vectorize(x.values().dims()), val, values);
+
+  out->set_dims(x.dims());
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(full_like_coo,
+                   GPU,
+                   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,
+                   GPU,
+                   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/gpu/fused_attention_grad_kernel.cu

@@ -0,0 +1,146 @@
+// 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/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
+#include "paddle/phi/kernels/funcs/sparse/sparse_blas.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/matmul_grad_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T>
+__global__ void AttnSoftmaxGpuGradKernel(const int64_t* out_crows,
+                                         const T* out_values,
+                                         const T* dout_values,
+                                         T* dx_values,
+                                         int M,
+                                         int total_row_num,
+                                         float scale,
+                                         int batch_nnz) {
+  // dx = (dout - sum(dout * out)) * out
+  int row = blockIdx.x * blockDim.y + threadIdx.y;
+  if (row >= total_row_num) return;
+
+  int cur_batch = row / M;
+  int crow_idx = cur_batch * (M + 1) + (row % M);
+  int row_first = cur_batch * batch_nnz + static_cast<int>(out_crows[crow_idx]);
+  int row_nnz = static_cast<int>(out_crows[crow_idx + 1] - out_crows[crow_idx]);
+  if (row_nnz == 0) return;
+
+  T mul = 0;
+  for (int idx = threadIdx.x; idx < row_nnz; idx += blockDim.x) {
+    mul += out_values[row_first + idx] * dout_values[row_first + idx];
+  }
+  T mul_sum = funcs::WarpReduceSum<T>(mul, 0xFFFFFFFF);
+
+  for (int idx = threadIdx.x; idx < row_nnz; idx += blockDim.x) {
+    dx_values[row_first + idx] = (dout_values[row_first + idx] - mul_sum) *
+                                 out_values[row_first + idx] / scale;
+  }
+}
+
+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) {
+#if CUDA_VERSION >= 11080
+  /* Step1: Forward: softmax{CSR} * value{Dense} -> out{Dense}, reuse */
+  SparseCsrTensor dsoftmax;
+  MatmulCsrDenseGradKernel<T, Context>(
+      dev_ctx, softmax, value, dout, &dsoftmax, dvalue);
+
+  /* Step2: Calculate grad of sdd_result, manually not reuse */
+  SparseCsrTensor d_sdd_result;
+  EmptyLikeCsrKernel<T, Context>(dev_ctx, dsoftmax, &d_sdd_result);
+  auto q_dim = query.dims();
+  auto q_rank = q_dim.size();
+
+  int total_row_num = 1;
+  int batch_num = 1;
+  for (int i = 0; i < q_rank - 1; ++i) {
+    total_row_num *= q_dim[i];
+    if (i < q_rank - 2) {
+      batch_num *= q_dim[i];
+    }
+  }
+  int M = q_dim[q_rank - 2];
+  int N = q_dim[q_rank - 1];
+  int batch_nnz = softmax.nnz() / batch_num;
+
+  dim3 grid((total_row_num + 7) / 8);
+  dim3 block(WARP_SIZE, 8);
+
+  AttnSoftmaxGpuGradKernel<T><<<grid, block, 0, dev_ctx.stream()>>>(
+      softmax.crows().data<int64_t>(),
+      softmax.values().data<T>(),
+      dsoftmax.mutable_values()->data<T>(),
+      d_sdd_result.mutable_values()->data<T>(),
+      M,
+      total_row_num,
+      std::sqrt(N),
+      batch_nnz);
+
+  /* Step3: Forward: query{Dense} * key'{Dense} -> sdd_result{SparseCsr} */
+  auto sparse_blas = funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
+  // dquery{Dense} = d_sdd_result{SparseCsr} * key{Dense} //
+  dquery->Resize(query.dims());
+  dev_ctx.template Alloc<T>(dquery);
+  sparse_blas.SPMM(false,
+                   false,
+                   static_cast<T>(1.f),
+                   d_sdd_result,
+                   key,
+                   static_cast<T>(0.f),
+                   dquery);
+
+  // dkey{Dense} = d_sdd_result'{SparseCsr} * query{Dense} //
+  dkey->Resize(key.dims());
+  dev_ctx.template Alloc<T>(dkey);
+  sparse_blas.SPMM(true,
+                   false,
+                   static_cast<T>(1.f),
+                   d_sdd_result,
+                   query,
+                   static_cast<T>(0.f),
+                   dkey);
+#else
+  PADDLE_THROW(common::errors::Unimplemented(
+      "backward of 'sparse.nn.functional.attention' "
+      "use 'cusparseCsrSetStridedBatch', which is "
+      "completed supported from CUDA 11.8"));
+#endif
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(fused_attention_csr_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::FusedAttentionCsrGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}

paddle/phi/kernels/sparse/gpu/fused_attention_kernel.cu

@@ -0,0 +1,236 @@
+/* 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/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/activation_functor.h"
+#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
+#include "paddle/phi/kernels/funcs/sparse/sparse_blas.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/matmul_kernel.h"
+#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T>
+__global__ void AttnSoftmaxGpuKernel(const int64_t* x_crows,
+                                     const int64_t* x_cols,
+                                     const T* x_values,
+                                     const T* kp_mask,
+                                     const T* attn_mask,
+                                     T* out_values,
+                                     int M,
+                                     int total_row_num,
+                                     int num_heads,
+                                     int batch_nnz) {
+  // out = exp(x-x_max) / sum(exp(x-x_max))
+  int row = blockIdx.x * blockDim.y + threadIdx.y;
+  if (row >= total_row_num) return;
+
+  int cur_batch = row / M;
+  int cur_row = row % M;
+  int crow_idx = cur_batch * (M + 1) + cur_row;
+  int row_first = cur_batch * batch_nnz + static_cast<int>(x_crows[crow_idx]);
+  int row_nnz = static_cast<int>(x_crows[crow_idx + 1] - x_crows[crow_idx]);
+  if (row_nnz == 0) return;
+
+  T max_val = -std::numeric_limits<T>::infinity();
+  for (int idx = threadIdx.x; idx < row_nnz; idx += blockDim.x) {
+    bool mask = false;
+    int col_idx = static_cast<int>(x_cols[row_first + idx]);
+    if (kp_mask != nullptr &&
+        kp_mask[(cur_batch / num_heads) * M + col_idx] == 0) {
+      mask = true;
+    }
+    if (attn_mask != nullptr && attn_mask[cur_row * M + col_idx] == 0) {
+      mask = true;
+    }
+
+    if (!mask) {
+      T val = x_values[row_first + idx];
+      if (val > max_val) {
+        max_val = val;
+      }
+      out_values[row_first + idx] = val;
+    } else {
+      // Note corner case: when all elements of the row are masked, result
+      // may be wrong because of exp('-inf' - '-inf'), just ignore now.
+      out_values[row_first + idx] = -std::numeric_limits<T>::infinity();
+    }
+  }
+  T row_max_val = funcs::WarpReduceMax<T>(max_val, 0xFFFFFFFF);
+
+  T exp_sum = 0;
+  for (int idx = threadIdx.x; idx < row_nnz; idx += blockDim.x) {
+    auto functor = funcs::CudaExpFunctor<T>();
+    T exp = functor(out_values[row_first + idx] - row_max_val);
+    exp_sum += exp;
+    out_values[row_first + idx] = exp;
+  }
+  T row_exp_sum = funcs::WarpReduceSum<T>(exp_sum, 0xFFFFFFFF);
+
+  for (int idx = threadIdx.x; idx < row_nnz; idx += blockDim.x) {
+    out_values[row_first + idx] = out_values[row_first + idx] / row_exp_sum;
+  }
+}
+
+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) {
+#if CUDA_VERSION >= 11080
+  /* Check Shape */
+  auto q_dim = query.dims();
+  auto q_rank = q_dim.size();
+
+  int total_row_num = 1;
+  int batch_num = 1;
+  for (int i = 0; i < q_rank - 1; ++i) {
+    total_row_num *= q_dim[i];
+    if (i < q_rank - 2) {
+      batch_num *= q_dim[i];
+    }
+  }
+  int M = q_dim[q_rank - 2];
+  int N = q_dim[q_rank - 1];
+
+  PADDLE_ENFORCE_EQ(
+      query.dims().size(),
+      4,
+      common::errors::InvalidArgument(" 'query' must be 4D Tensor"));
+  PADDLE_ENFORCE_EQ(
+      key.dims().size(),
+      4,
+      common::errors::InvalidArgument(" 'key' must be 4D Tensor"));
+  PADDLE_ENFORCE_EQ(
+      value.dims().size(),
+      4,
+      common::errors::InvalidArgument(" 'value' must be 4D Tensor"));
+
+  PADDLE_ENFORCE_EQ(sparse_mask.dims().size(),
+                    3,
+                    common::errors::InvalidArgument(
+                        "dense shape of 'sparse_mask' must be "
+                        "[batch_size*num_heads, seq_len, seq_len]"));
+  PADDLE_ENFORCE_EQ(sparse_mask.dims()[0],
+                    q_dim[0] * q_dim[1],
+                    common::errors::InvalidArgument(
+                        "dense shape of 'sparse_mask' must be "
+                        "[batch_size*num_heads, seq_len, seq_len]"));
+  PADDLE_ENFORCE_EQ(sparse_mask.dims()[1],
+                    M,
+                    common::errors::InvalidArgument(
+                        "dense shape of 'sparse_mask' must be "
+                        "[batch_size*num_heads, seq_len, seq_len]"));
+  PADDLE_ENFORCE_EQ(sparse_mask.dims()[2],
+                    M,
+                    common::errors::InvalidArgument(
+                        "dense shape of 'sparse_mask' must be "
+                        "[batch_size*num_heads, seq_len, seq_len]"));
+
+  const auto kp_mask_ptr = key_padding_mask.get_ptr();
+  if (kp_mask_ptr) {
+    PADDLE_ENFORCE_EQ(
+        kp_mask_ptr->dims().size(),
+        2,
+        common::errors::InvalidArgument(
+            "shape of 'key_padding_mask' must be [batch_size, seq_len]"));
+    PADDLE_ENFORCE_EQ(
+        kp_mask_ptr->dims()[0],
+        q_dim[0],
+        common::errors::InvalidArgument(
+            "shape of 'key_padding_mask' must be [batch_size, seq_len]"));
+    PADDLE_ENFORCE_EQ(
+        kp_mask_ptr->dims()[1],
+        M,
+        common::errors::InvalidArgument(
+            "shape of 'key_padding_mask' must be [batch_size, seq_len]"));
+  }
+
+  const auto attn_mask_ptr = attn_mask.get_ptr();
+  if (attn_mask_ptr) {
+    PADDLE_ENFORCE_EQ(attn_mask_ptr->dims().size(),
+                      2,
+                      common::errors::InvalidArgument(
+                          "shape of 'attn_mask' must be [seq_len, seq_len]"));
+    PADDLE_ENFORCE_EQ(attn_mask_ptr->dims()[0],
+                      M,
+                      common::errors::InvalidArgument(
+                          "shape of 'attn_mask' must be [seq_len, seq_len]"));
+    PADDLE_ENFORCE_EQ(attn_mask_ptr->dims()[1],
+                      M,
+                      common::errors::InvalidArgument(
+                          "shape of 'attn_mask' must be [seq_len, seq_len]"));
+  }
+
+  /* Step1: SDD Matmul, reuse matmul */
+  SparseCsrTensor sdd_result;
+  EmptyLikeCsrKernel<T, Context>(dev_ctx, sparse_mask, &sdd_result);
+  auto sparse_blas = funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
+  sparse_blas.SDDMM(false,
+                    true,
+                    static_cast<T>(1 / std::sqrt(N)),
+                    query,
+                    key,
+                    static_cast<T>(0),
+                    &sdd_result);
+
+  EmptyLikeCsrKernel<T, Context>(dev_ctx, sdd_result, softmax);
+
+  dim3 grid((total_row_num + 7) / 8);
+  dim3 block(WARP_SIZE, 8);
+
+  int batch_nnz = sdd_result.nnz() / batch_num;
+  AttnSoftmaxGpuKernel<T><<<grid, block, 0, dev_ctx.stream()>>>(
+      sdd_result.crows().data<int64_t>(),
+      sdd_result.cols().data<int64_t>(),
+      sdd_result.values().data<T>(),
+      kp_mask_ptr ? kp_mask_ptr->data<T>() : nullptr,
+      attn_mask_ptr ? attn_mask_ptr->data<T>() : nullptr,
+      softmax->mutable_values()->data<T>(),
+      M,
+      total_row_num,
+      q_dim[1],
+      batch_nnz);
+
+  softmax->set_dims(make_ddim({q_dim[0], q_dim[1], q_dim[2], q_dim[2]}));
+  MatmulCsrDenseKernel<T, Context>(dev_ctx, *softmax, value, out);
+#else
+  PADDLE_THROW(common::errors::Unimplemented(
+      "forward of 'sparse.nn.functional.attention' "
+      "use 'cusparseCsrSetStridedBatch', which is "
+      "completed supported from CUDA 11.8"));
+#endif
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(fused_attention_csr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::FusedAttentionCsrKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}

paddle/phi/kernels/sparse/gpu/mask_grad_kernel.cu

@@ -0,0 +1,56 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/phi/kernels/sparse/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,
+                   GPU,
+                   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,
+                   GPU,
+                   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/gpu/mask_kernel.cu

@@ -0,0 +1,579 @@
+/* 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 <thrust/execution_policy.h>
+
+#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/backends/gpu/gpu_info.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.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/aligned_vector.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/sparse/flatten_indices.cu.h"
+#include "paddle/phi/kernels/funcs/sparse/utils.cu.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename IntT>
+__global__ void MaskKernel(const T* x_ptr,
+                           const IntT* indices_ptr,
+                           const int64_t* sparse_offsets,
+                           const int64_t non_zero_num,
+                           const int cols,
+                           const int sparse_dim,
+                           T* out_values_ptr) {
+  CUDA_KERNEL_LOOP_TYPE(i, non_zero_num * cols, int64_t) {
+    int64_t out_i = i / cols;
+    int64_t col_i = i - out_i * cols;
+    int64_t index = 0;
+    for (int j = 0; j < sparse_dim; j++) {
+      index += indices_ptr[j * non_zero_num + out_i] * sparse_offsets[j];
+    }
+    out_values_ptr[out_i * cols + col_i] = x_ptr[index * cols + col_i];
+  }
+}
+
+template <typename T, typename IntT>
+void MaskCooGPUKernel(const GPUContext& 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();
+  DenseTensor out_indices = EmptyLike<IntT>(dev_ctx, indices);
+  DenseTensor out_values = EmptyLike<T>(dev_ctx, values);
+  if (mask.nnz() <= 0) {
+    out->SetMember(out_indices, out_values, dims, true);
+    return;
+  }
+
+  const int sparse_dim = mask.sparse_dim();
+  DenseTensor sparse_offsets = Empty<GPUContext>(
+      dev_ctx,
+      DenseTensorMeta(DataType::INT64, {sparse_dim}, DataLayout::NCHW));
+  std::vector<int64_t> h_sparse_offsets(sparse_dim);
+  funcs::sparse::CalcOffsetsPerDim(dims, sparse_dim, h_sparse_offsets.data());
+
+  backends::gpu::GpuMemcpyAsync(sparse_offsets.data<int64_t>(),
+                                &h_sparse_offsets[0],
+                                sizeof(int64_t) * sparse_dim,
+                                gpuMemcpyHostToDevice,
+                                dev_ctx.stream());
+
+  phi::Copy(dev_ctx, indices, dev_ctx.GetPlace(), false, &out_indices);
+
+  const IntT* indices_ptr = indices.data<IntT>();
+  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 = dims_2d[1];
+
+  auto config =
+      backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num * cols, 1);
+  MaskKernel<T, IntT>
+      <<<config.block_per_grid, config.thread_per_block, 0, dev_ctx.stream()>>>(
+          x_ptr,
+          indices_ptr,
+          sparse_offsets.data<int64_t>(),
+          non_zero_num,
+          cols,
+          sparse_dim,
+          out_values_ptr);
+
+  out->SetMember(out_indices, out_values, dims, true);
+}
+
+template <typename IntT>
+__global__ void ConvertCsrCrowsToCooRows(const IntT* crows_ptr,
+                                         const IntT* crows_offsets,
+                                         IntT* rows_ptr,
+                                         IntT* batch_ptr,
+                                         const int rows) {
+  const int b = blockIdx.y;
+  const int64_t offset = crows_offsets ? crows_offsets[b] : 0;
+  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  for (int i = tid; i < rows; i += gridDim.x * blockDim.x) {
+    for (int j = crows_ptr[b * (rows + 1) + i];
+         j < crows_ptr[b * (rows + 1) + i + 1];
+         j++) {
+      rows_ptr[offset + j] = i;
+      if (batch_ptr) {
+        batch_ptr[offset + j] = b;
+      }
+    }
+  }
+}
+
+template <typename IntT>
+__global__ void GetBatchSizes(const IntT* crows,
+                              const int rows,
+                              const int batches,
+                              IntT* batch_sizes) {
+  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  if (tid < batches) {
+    batch_sizes[tid] = crows[tid * (rows + 1) + rows];
+  }
+}
+
+template <typename T, typename IntT>
+void MaskCsr2DGPUKernel(const GPUContext& 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);
+
+  const DDim& dims = x.dims();
+  const int64_t non_zero_num = mask.nnz();
+  int64_t sparse_dim = 2;
+  DenseTensor sparse_offsets = Empty<IntT>(dev_ctx, {sparse_dim});
+  std::vector<int64_t> h_sparse_offsets(sparse_dim);
+  funcs::sparse::CalcOffsetsPerDim(dims, sparse_dim, h_sparse_offsets.data());
+
+  backends::gpu::GpuMemcpyAsync(sparse_offsets.data<int64_t>(),
+                                &h_sparse_offsets[0],
+                                sizeof(int64_t) * sparse_dim,
+                                gpuMemcpyHostToDevice,
+                                dev_ctx.stream());
+
+  const auto& csr_crows = mask.crows();
+  const auto& csr_cols = mask.cols();
+  const IntT* csr_crows_data = csr_crows.data<IntT>();
+  const IntT* csr_cols_data = csr_cols.data<IntT>();
+
+  const int batches = 1;
+  const int rows = dims[0];
+  auto dims_2d = flatten_to_2d(dims, sparse_dim);
+  const int cols = dims_2d[1];
+
+  DenseTensor indices = Empty<IntT>(dev_ctx, {sparse_dim, non_zero_num});
+  IntT* coo_indices = indices.data<IntT>();
+  IntT* batch_ptr = nullptr;
+  IntT* coo_rows_data = coo_indices;
+  IntT* coo_cols_data = coo_rows_data + non_zero_num;
+  IntT* offsets_ptr = nullptr;
+
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rows, 1);
+  config.block_per_grid.y = batches;
+  ConvertCsrCrowsToCooRows<IntT>
+      <<<config.block_per_grid, config.thread_per_block.x>>>(
+          csr_crows_data, offsets_ptr, coo_rows_data, batch_ptr, rows);
+  backends::gpu::GpuMemcpyAsync(coo_cols_data,
+                                csr_cols_data,
+                                sizeof(IntT) * non_zero_num,
+                                gpuMemcpyDeviceToDevice,
+                                dev_ctx.stream());
+
+  const T* x_ptr = x.data<T>();
+  const IntT* indices_ptr = coo_indices;
+  T* out_values_ptr = out_values.data<T>();
+
+  auto config_mask =
+      backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num * cols, 1);
+  MaskKernel<T, IntT><<<config_mask.block_per_grid,
+                        config_mask.thread_per_block,
+                        0,
+                        dev_ctx.stream()>>>(x_ptr,
+                                            indices_ptr,
+                                            sparse_offsets.data<int64_t>(),
+                                            non_zero_num,
+                                            cols,
+                                            sparse_dim,
+                                            out_values_ptr);
+
+  out->SetMember(out_crows, out_cols, out_values, x.dims());
+}
+
+template <typename T, typename IntT>
+void MaskCsr3DGPUKernel(const GPUContext& 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);
+
+  const DDim& dims = x.dims();
+  const int64_t non_zero_num = mask.nnz();
+  int64_t sparse_dim = 3;
+  DenseTensor sparse_offsets = Empty<IntT>(dev_ctx, {sparse_dim});
+  std::vector<int64_t> h_sparse_offsets(sparse_dim);
+  funcs::sparse::CalcOffsetsPerDim(dims, sparse_dim, h_sparse_offsets.data());
+
+  backends::gpu::GpuMemcpyAsync(sparse_offsets.data<int64_t>(),
+                                &h_sparse_offsets[0],
+                                sizeof(int64_t) * sparse_dim,
+                                gpuMemcpyHostToDevice,
+                                dev_ctx.stream());
+
+  const auto& csr_crows = mask.crows();
+  const auto& csr_cols = mask.cols();
+  const IntT* csr_crows_data = csr_crows.data<IntT>();
+  const IntT* csr_cols_data = csr_cols.data<IntT>();
+
+  const int batches = dims[0];
+  const int rows = dims[1];
+  auto dims_2d = flatten_to_2d(dims, sparse_dim);
+  const int cols = dims_2d[1];
+
+  DenseTensor indices = Empty<IntT>(dev_ctx, {sparse_dim, non_zero_num});
+  DenseTensor offsets = Empty<IntT>(dev_ctx, {batches});
+  IntT* coo_indices = indices.data<IntT>();
+  IntT* batch_ptr = coo_indices;
+  IntT* coo_rows_data = batch_ptr + non_zero_num;
+  IntT* coo_cols_data = coo_rows_data + non_zero_num;
+  IntT* offsets_ptr = offsets.data<IntT>();
+
+  auto config_batch = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, batches, 1);
+  GetBatchSizes<IntT>
+      <<<config_batch.block_per_grid.x, config_batch.thread_per_block.x>>>(
+          csr_crows_data, rows, batches, offsets_ptr);
+
+#ifdef PADDLE_WITH_HIP
+  thrust::exclusive_scan(thrust::hip::par.on(dev_ctx.stream()),
+#else
+  thrust::exclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                         offsets_ptr,
+                         offsets_ptr + batches,
+                         offsets_ptr);
+
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rows, 1);
+  config.block_per_grid.y = batches;
+  ConvertCsrCrowsToCooRows<IntT>
+      <<<config.block_per_grid, config.thread_per_block.x>>>(
+          csr_crows_data, offsets_ptr, coo_rows_data, batch_ptr, rows);
+  backends::gpu::GpuMemcpyAsync(coo_cols_data,
+                                csr_cols_data,
+                                sizeof(IntT) * non_zero_num,
+                                gpuMemcpyDeviceToDevice,
+                                dev_ctx.stream());
+
+  const T* x_ptr = x.data<T>();
+  const IntT* indices_ptr = coo_indices;
+  T* out_values_ptr = out_values.data<T>();
+
+  auto config_mask =
+      backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num * cols, 1);
+  MaskKernel<T, IntT><<<config_mask.block_per_grid,
+                        config_mask.thread_per_block,
+                        0,
+                        dev_ctx.stream()>>>(x_ptr,
+                                            indices_ptr,
+                                            sparse_offsets.data<int64_t>(),
+                                            non_zero_num,
+                                            cols,
+                                            sparse_dim,
+                                            out_values_ptr);
+
+  out->SetMember(out_crows, out_cols, out_values, x.dims());
+}
+
+/**
+ * @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(), "MaskCooGPUKernel", ([&] {
+        MaskCooGPUKernel<T, data_t>(dev_ctx, x, mask, out);
+      }));
+}
+
+/**
+ * @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(), "MaskCsr2DGPUKernel", ([&] {
+          MaskCsr2DGPUKernel<T, data_t>(dev_ctx, x, mask, out);
+        }));
+  } else if (x_dims.size() == 3) {
+    PD_VISIT_BASE_INTEGRAL_TYPES(
+        mask.crows().dtype(), "MaskCsr3DGPUKernel", ([&] {
+          MaskCsr3DGPUKernel<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 IntT>
+__global__ void MaskTable(const IntT* x_indices,
+                          const int n,
+                          int* index_flags,
+                          int* table) {
+  CUDA_KERNEL_LOOP_TYPE(i, n, int64_t) {
+    int index = x_indices[i];
+    funcs::sparse::SetBits(index, index_flags);
+    table[index] = i;
+  }
+}
+
+template <typename T, typename IntT, int VecSize>
+__global__ void MaskCopy(const IntT* mask_indices,
+                         const int* index_flags,
+                         const int* table,
+                         const int n,
+                         const int stride,
+                         const T* x_values,
+                         T* out_values) {
+  using LoadT = AlignedVector<T, VecSize>;
+  using StoreT = AlignedVector<T, VecSize>;
+  CUDA_KERNEL_LOOP_TYPE(i, n, int64_t) {
+    const int mask_index = mask_indices[i];
+    const bool flag = funcs::sparse::TestBits(mask_index, index_flags);
+    if (flag) {
+      int j = table[mask_index];
+      for (int k = 0; k < stride; k += VecSize) {
+        LoadT vec_x;
+        Load<T, VecSize>(x_values + j * stride + k, &vec_x);
+        Store<T, VecSize>(vec_x, out_values + i * stride + k);
+      }
+    }
+  }
+}
+
+template <typename T, typename IntT>
+void MaskHelperCooGPUKernel(const GPUContext& 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();
+  auto indices_dtype = phi::CppTypeToDataType<IntT>::Type();
+
+  std::vector<IntT> sparse_offsets(sparse_dim);
+
+  DenseTensorMeta x_indices_meta(indices_dtype, {x.nnz()}, DataLayout::NCHW);
+  DenseTensorMeta mask_indices_meta(
+      indices_dtype, {mask_indices.dims()[1]}, DataLayout::NCHW);
+  DenseTensorMeta sparse_offset_meta(
+      indices_dtype, {sparse_dim}, DataLayout::NCHW);
+
+  DenseTensor x_indices = Empty<GPUContext>(dev_ctx, std::move(x_indices_meta));
+  DenseTensor mask_meta_indices =
+      Empty<GPUContext>(dev_ctx, std::move(mask_indices_meta));
+  DenseTensor bound_out =
+      Empty<GPUContext>(dev_ctx, std::move(mask_indices_meta));
+  DenseTensor d_sparse_offsets =
+      Empty<GPUContext>(dev_ctx, std::move(sparse_offset_meta));
+  IntT* x_indices_ptr = x_indices.data<IntT>();
+  IntT* mask_indices_ptr = mask_meta_indices.data<IntT>();
+  IntT* bound_out_ptr = bound_out.data<IntT>();
+
+  // 1. calc the offsets of per dim
+  funcs::sparse::CalcOffsetsPerDim(x.dims(), sparse_dim, sparse_offsets.data());
+  // 2. copy sparse_offsets to device
+  backends::gpu::GpuMemcpyAsync(d_sparse_offsets.data<IntT>(),
+                                sparse_offsets.data(),
+                                sizeof(IntT) * sparse_dim,
+                                gpuMemcpyHostToDevice,
+                                dev_ctx.stream());
+
+  // 3. flatten x indices and mask indices
+  auto config =
+      backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x_indices.numel(), 1);
+  funcs::sparse::FlattenIndicesKernel<<<config.block_per_grid,
+                                        config.thread_per_block,
+                                        0,
+                                        dev_ctx.stream()>>>(
+      x.indices().data<IntT>(),
+      d_sparse_offsets.data<IntT>(),
+      x_indices.numel(),
+      sparse_dim,
+      x_indices_ptr);
+
+  config = backends::gpu::GetGpuLaunchConfig1D(
+      dev_ctx, mask_meta_indices.numel(), 1);
+  funcs::sparse::FlattenIndicesKernel<<<config.block_per_grid,
+                                        config.thread_per_block,
+                                        0,
+                                        dev_ctx.stream()>>>(
+      mask_indices.data<IntT>(),
+      d_sparse_offsets.data<IntT>(),
+      mask_meta_indices.numel(),
+      sparse_dim,
+      mask_indices_ptr);
+
+  int table_size = 1;
+  auto x_dims = x.dims();
+  for (int i = 0; i < sparse_dim; i++) {
+    table_size *= x_dims[i];
+  }
+  DenseTensor table = Empty<int>(dev_ctx, {table_size});
+  DenseTensor index_flags = Empty<int>(dev_ctx, {(table_size + 31) / 32});
+  backends::gpu::GpuMemsetAsync(index_flags.data<int>(),
+                                0,
+                                index_flags.numel() * sizeof(int),
+                                dev_ctx.stream());
+  const int64_t stride =
+      x.dims().size() == sparse_dim ? 1 : x.values().dims()[1];
+  *out = EmptyLike<T>(dev_ctx, x.values());
+  funcs::SetConstant<GPUContext, T> set_zero;
+  set_zero(dev_ctx, out, static_cast<T>(0));
+  T* out_ptr = out->data<T>();
+  config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x_indices.numel(), 1);
+  MaskTable<<<config.block_per_grid,
+              config.thread_per_block,
+              0,
+              dev_ctx.stream()>>>(x_indices_ptr,
+                                  x_indices.numel(),
+                                  index_flags.data<int>(),
+                                  table.data<int>());
+  config = backends::gpu::GetGpuLaunchConfig1D(
+      dev_ctx, mask_meta_indices.numel(), 1);
+
+  const int VecBytes = 16;
+  const int VecSize = VecBytes / sizeof(T);
+  if (stride % VecSize == 0) {
+    MaskCopy<T, IntT, VecSize><<<config.block_per_grid,
+                                 config.thread_per_block,
+                                 0,
+                                 dev_ctx.stream()>>>(mask_indices_ptr,
+                                                     index_flags.data<int>(),
+                                                     table.data<int>(),
+                                                     mask_meta_indices.numel(),
+                                                     stride,
+                                                     x.values().data<T>(),
+                                                     out_ptr);
+  } else {
+    MaskCopy<T, IntT, 1><<<config.block_per_grid,
+                           config.thread_per_block,
+                           0,
+                           dev_ctx.stream()>>>(mask_indices_ptr,
+                                               index_flags.data<int>(),
+                                               table.data<int>(),
+                                               mask_meta_indices.numel(),
+                                               stride,
+                                               x.values().data<T>(),
+                                               out_ptr);
+  }
+}
+
+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(), "MaskHelperCooGPUKernel", ([&] {
+        MaskHelperCooGPUKernel<T, data_t>(dev_ctx, x, mask_indices, out);
+      }));
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(mask_helper_coo,
+                   GPU,
+                   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,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaskAsCooKernel,
+                   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,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaskAsCsrKernel,
+                   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/gpu/matmul_grad_kernel.cu

@@ -0,0 +1,277 @@
+/* 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/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/math_function_impl.h"
+#include "paddle/phi/kernels/funcs/sparse/sparse_blas.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h"
+#include "paddle/phi/kernels/sparse/unary_kernel.h"
+#include "paddle/phi/kernels/transpose_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context>
+void MatmulCooDenseGradKernel(const Context& dev_ctx,
+                              const SparseCooTensor& x,
+                              const DenseTensor& y,
+                              const DenseTensor& dout,
+                              SparseCooTensor* dx,
+                              DenseTensor* dy) {
+#if defined(PADDLE_WITH_CUDA) || HIP_VERSION >= 403
+  auto sparse_blas = funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
+
+  // dx{SparseCoo} = dout{Dense} * y'{Dense}
+  if (dx) {
+    // 'cusparseSDDMM' only support CSR now, so use COO->CSR->COO,
+    // which will increase some expenses.
+    EmptyLikeCooKernel<T, Context>(dev_ctx, x, dx);
+    SparseCsrTensor dx_csr = CooToCsr<T, Context>(dev_ctx, *dx);
+#ifdef PADDLE_WITH_HIP
+    funcs::SetConstant<Context, T> set_zero;
+    set_zero(dev_ctx, dx_csr.mutable_non_zero_elements(), static_cast<T>(0.0f));
+#endif
+    sparse_blas.SDDMM(
+        false, true, static_cast<T>(1), dout, y, static_cast<T>(0), &dx_csr);
+
+    CsrToCooKernel<T, Context>(dev_ctx, dx_csr, dx);
+  }
+
+  // dy{Dense} = x'{SparseCoo} * dout{Dense}
+  if (dy) {
+    MetaTensor meta_dy(dy);
+    meta_dy.set_dims(y.dims());
+    meta_dy.set_dtype(y.dtype());
+    dev_ctx.template Alloc<T>(dy);
+
+#ifdef PADDLE_WITH_HIP
+    SparseCsrTensor x_csr = CooToCsr<T, Context>(dev_ctx, x);
+    funcs::SetConstant<Context, T> set_zero;
+    set_zero(dev_ctx, dy, static_cast<T>(0.0f));
+    sparse_blas.SPMM(
+        true, false, static_cast<T>(1), x_csr, dout, static_cast<T>(0), dy);
+#elif defined(PADDLE_WITH_CUDA)
+    sparse_blas.SPMM(
+        true, false, static_cast<T>(1), x, dout, static_cast<T>(0), dy);
+#endif
+  }
+#endif
+}
+
+template <typename T, typename Context>
+void MatmulCsrDenseGradKernel(const Context& dev_ctx,
+                              const SparseCsrTensor& x,
+                              const DenseTensor& y,
+                              const DenseTensor& dout,
+                              SparseCsrTensor* dx,
+                              DenseTensor* dy) {
+#if defined(PADDLE_WITH_CUDA) || HIP_VERSION >= 403
+  auto sparse_blas = funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
+
+  // dx{SparseCsr} = dout{Dense} * y'{Dense}
+  if (dx) {
+    // InferMeta of SparseCsrTensor 'dx', CreateLikeInferMeta
+    EmptyLikeCsrKernel<T, Context>(dev_ctx, x, dx);
+
+    sparse_blas.SDDMM(
+        false, true, static_cast<T>(1), dout, y, static_cast<T>(0), dx);
+  }
+
+  // dy{Dense} = x'{SparseCsr} * dout{Dense}
+  if (dy) {
+    // InferMeta of DenseTensor 'dy'
+    MetaTensor meta_dy(dy);
+    meta_dy.set_dims(y.dims());
+    meta_dy.set_dtype(y.dtype());
+
+    dev_ctx.template Alloc<T>(dy);
+
+#ifdef PADDLE_WITH_HIP
+    funcs::SetConstant<Context, T> set_zero;
+    set_zero(dev_ctx, dy, static_cast<T>(0.0f));
+#endif
+
+    sparse_blas.SPMM(
+        true, false, static_cast<T>(1), x, dout, static_cast<T>(0), dy);
+  }
+#endif
+}
+
+template <typename T, typename Context>
+void MatmulCsrCsrGradKernel(const Context& dev_ctx,
+                            const SparseCsrTensor& x,
+                            const SparseCsrTensor& y,
+                            const SparseCsrTensor& dout,
+                            SparseCsrTensor* dx,
+                            SparseCsrTensor* dy) {
+#if defined(PADDLE_WITH_CUDA)
+  auto sparse_blas = funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
+
+  std::vector<int64_t> xdim_vec = vectorize(x.dims());
+  auto x_ndims = xdim_vec.size();
+  std::vector<int> perm;
+  if (x_ndims == 2) {
+    perm = {1, 0};
+  } else {
+    perm = {0, 2, 1};
+  }
+
+  // dx{SparseCsr} = dout{SparseCsr} * y'{SparseCsr}
+  if (dx) {
+    // cusparseSpGEMM only support CUSPARSE_OPERATION_NON_TRANSPOSE.
+    // transpose y before cusparseSpGEMM computation.
+    SparseCsrTensor trans_y;
+    TransposeCsrKernel<T, Context>(dev_ctx, y, perm, &trans_y);
+
+    sparse_blas.SPGEMM(
+        false, false, static_cast<T>(1), dout, trans_y, static_cast<T>(0), dx);
+  }
+
+  // dy{SparseCsr} = x'{SparseCsr} * dout{SparseCsr}
+  if (dy) {
+    // cusparseSpGEMM only support CUSPARSE_OPERATION_NON_TRANSPOSE.
+    // transpose x before cusparseSpGEMM computation.
+    SparseCsrTensor trans_x;
+    TransposeCsrKernel<T, Context>(dev_ctx, x, perm, &trans_x);
+
+    sparse_blas.SPGEMM(
+        false, false, static_cast<T>(1), trans_x, dout, static_cast<T>(0), dy);
+  }
+#endif
+}
+
+template <typename T, typename Context>
+void MatmulCooCooGradKernel(const Context& dev_ctx,
+                            const SparseCooTensor& x,
+                            const SparseCooTensor& y,
+                            const SparseCooTensor& dout,
+                            SparseCooTensor* dx,
+                            SparseCooTensor* dy) {
+  // cusparseSpGEMM only support CSR now, so use COO->CSR->COO.
+  SparseCsrTensor x_csr, y_csr, dout_csr, dx_csr, dy_csr;
+  CooToCsrKernel<T>(dev_ctx, x, &x_csr);
+  CooToCsrKernel<T>(dev_ctx, y, &y_csr);
+  CooToCsrKernel<T>(dev_ctx, dout, &dout_csr);
+  MetaTensor meta_dx_csr(&dx_csr);
+  phi::UnchangedInferMeta(dx, &meta_dx_csr);
+  MetaTensor meta_dy_csr(&dy_csr);
+  phi::UnchangedInferMeta(dy, &meta_dy_csr);
+  MatmulCsrCsrGradKernel<T>(dev_ctx, x_csr, y_csr, dout_csr, &dx_csr, &dy_csr);
+  CsrToCooKernel<T>(dev_ctx, dx_csr, dx);
+  CsrToCooKernel<T>(dev_ctx, dy_csr, dy);
+}
+
+template <typename T, typename Context>
+void MaskedMatmulCsrGradKernel(const Context& dev_ctx,
+                               const DenseTensor& x,
+                               const DenseTensor& y,
+                               const SparseCsrTensor& dout,
+                               DenseTensor* dx,
+                               DenseTensor* dy) {
+#if defined(PADDLE_WITH_CUDA)
+  auto sparse_blas = funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
+
+  // dx{Dense} = dout{SparseCsr} * y'{Dense}
+  if (dx) {
+    // InferMeta of DenseTensor 'dx'
+    MetaTensor meta_dx(dx);
+    meta_dx.set_dims(x.dims());
+    meta_dx.set_dtype(x.dtype());
+
+    dev_ctx.template Alloc<T>(dx);
+    sparse_blas.SPMM(
+        false, true, static_cast<T>(1), dout, y, static_cast<T>(0), dx);
+  }
+
+  // dy{Dense} = x'{Dense} * dout{SparseCsr}
+  // That is: dy'{Dense} = dout'{SparseCsr} * x{Dense}
+  if (dy) {
+    std::vector<int> trans_dim_vec = vectorize<int>(y.dims());
+    size_t rank = trans_dim_vec.size();
+    std::swap(trans_dim_vec[rank - 1], trans_dim_vec[rank - 2]);
+    DenseTensor trans_dy = Empty<T, Context>(dev_ctx, trans_dim_vec);
+
+    sparse_blas.SPMM(
+        true, false, static_cast<T>(1), dout, x, static_cast<T>(0), &trans_dy);
+
+    // InferMeta of DenseTensor 'dy'
+    MetaTensor meta_dy(dy);
+    meta_dy.set_dims(y.dims());
+    meta_dy.set_dtype(y.dtype());
+
+    dev_ctx.template Alloc<T>(dy);
+
+    size_t y_ndim = y.dims().size();
+    std::vector<int> axis(y_ndim);
+    for (size_t i = 0; i < y_ndim; ++i) {
+      axis[i] = i;
+    }
+    std::swap(axis[y_ndim - 1], axis[y_ndim - 2]);
+    TransposeKernel<T, Context>(dev_ctx, trans_dy, axis, dy);
+  }
+#endif
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(matmul_coo_dense_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MatmulCooDenseGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(matmul_csr_dense_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MatmulCsrDenseGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(matmul_csr_csr_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MatmulCsrCsrGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(matmul_coo_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MatmulCooCooGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(masked_matmul_csr_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaskedMatmulCsrGradKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/sparse/gpu/matmul_kernel.cu

@@ -0,0 +1,296 @@
+/* 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/common/ddim.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/meta_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/math_function_impl.h"
+#include "paddle/phi/kernels/funcs/sparse/sparse_blas.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/impl/unary_kernel_impl.h"
+#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context, typename TensorType>
+void MatmulKernelImpl(const Context& dev_ctx,
+                      const TensorType& x,
+                      const DenseTensor& y,
+                      DenseTensor* out) {
+#if defined(PADDLE_WITH_CUDA) || HIP_VERSION >= 402
+  std::vector<int64_t> xdim_vec = vectorize(x.dims());
+  std::vector<int64_t> ydim_vec = vectorize(y.dims());
+  auto x_ndims = xdim_vec.size();
+  auto y_ndims = ydim_vec.size();
+  PADDLE_ENFORCE_EQ(x_ndims,
+                    y_ndims,
+                    common::errors::PreconditionNotMet(
+                        "The dims size of Input(x) and Input(y) "
+                        "should be equal, But received X's "
+                        "dimensions=%d, Y's dimensions=%d.",
+                        x_ndims,
+                        y_ndims));
+  PADDLE_ENFORCE_GE(
+      x_ndims,
+      2,
+      common::errors::InvalidArgument("the dims size of Input(x) and "
+                                      "Input(y) must be greater than "
+                                      "or equal to 2."));
+
+  for (size_t i = 0; i < x_ndims - 2; ++i) {
+    PADDLE_ENFORCE_EQ(xdim_vec[i],
+                      ydim_vec[i],
+                      common::errors::InvalidArgument(
+                          "x.dim[%d] and x.dim[%d] must be equal.", i, i));
+  }
+
+  PADDLE_ENFORCE_GE(
+      xdim_vec[x_ndims - 1],
+      ydim_vec[y_ndims - 2],
+      common::errors::PreconditionNotMet(
+          "The shape of Input(x) and Input(y) is not suitable for matmul "
+          "operation, x_dim[-1] must be equal to y_dim[-2]."));
+
+  // InferMeta of DenseTensor 'out'
+  std::vector<int64_t> out_dim_vec(ydim_vec);
+  out_dim_vec[y_ndims - 2] = xdim_vec[x_ndims - 2];
+  out_dim_vec[y_ndims - 1] = ydim_vec[y_ndims - 1];
+  MetaTensor meta_out(out);
+  meta_out.set_dims(make_ddim(out_dim_vec));
+  meta_out.set_dtype(y.dtype());
+  // Ensure the output DenseTensor has a proper dense layout, not sparse layout
+  meta_out.set_layout(DataLayout::NCHW);
+
+  dev_ctx.template Alloc<T>(out);
+
+#ifdef PADDLE_WITH_HIP
+  funcs::SetConstant<Context, T> set_zero;
+  set_zero(dev_ctx, out, static_cast<T>(0.0f));
+#endif
+
+  auto sparse_blas = funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
+  sparse_blas.SPMM(
+      false, false, static_cast<T>(1), x, y, static_cast<T>(0), out);
+#endif
+}
+
+template <typename T, typename Context>
+void MatmulCooDenseKernel(const Context& dev_ctx,
+                          const SparseCooTensor& x,
+                          const DenseTensor& y,
+                          DenseTensor* out) {
+  MatmulKernelImpl<T>(dev_ctx, x, y, out);
+}
+
+template <typename T, typename Context>
+void MatmulCsrDenseKernel(const Context& dev_ctx,
+                          const SparseCsrTensor& x,
+                          const DenseTensor& y,
+                          DenseTensor* out) {
+  MatmulKernelImpl<T>(dev_ctx, x, y, out);
+}
+
+template <typename T, typename Context>
+void MatmulCsrCsrKernel(const Context& dev_ctx,
+                        const SparseCsrTensor& x,
+                        const SparseCsrTensor& y,
+                        SparseCsrTensor* out) {
+#if defined(PADDLE_WITH_CUDA)
+  std::vector<int64_t> xdim_vec = vectorize(x.dims());
+  std::vector<int64_t> ydim_vec = vectorize(y.dims());
+  auto x_ndims = xdim_vec.size();
+  auto y_ndims = ydim_vec.size();
+  PADDLE_ENFORCE_EQ(x_ndims,
+                    y_ndims,
+                    common::errors::PreconditionNotMet(
+                        "The dims size of Input(x) and Input(y) "
+                        "should be equal, But received X's "
+                        "dimensions=%d, Y's dimensions=%d.",
+                        x_ndims,
+                        y_ndims));
+  PADDLE_ENFORCE_GE(
+      x_ndims,
+      2,
+      common::errors::InvalidArgument("the dims size of Input(x) and "
+                                      "Input(y) must be greater than "
+                                      "or equal to 2."));
+
+  for (size_t i = 0; i < x_ndims - 2; ++i) {
+    PADDLE_ENFORCE_EQ(xdim_vec[i],
+                      ydim_vec[i],
+                      common::errors::InvalidArgument(
+                          "x.dim[%d] and x.dim[%d] must be equal.", i, i));
+  }
+
+  PADDLE_ENFORCE_GE(
+      xdim_vec[x_ndims - 1],
+      ydim_vec[y_ndims - 2],
+      common::errors::PreconditionNotMet(
+          "The shape of Input(x) and Input(y) is not suitable for matmul "
+          "operation, x_dim[-1] must be equal to y_dim[-2]."));
+
+  auto sparse_blas = funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
+  sparse_blas.SPGEMM(
+      false, false, static_cast<T>(1), x, y, static_cast<T>(0), out);
+#endif
+}
+
+template <typename T, typename Context>
+void MatmulCooCooKernel(const Context& dev_ctx,
+                        const SparseCooTensor& x,
+                        const SparseCooTensor& y,
+                        SparseCooTensor* out) {
+  // 'cusparseSpGEMM' only support CSR now, so use COO->CSR->COO.
+  SparseCsrTensor x_csr = CooToCsr<T, Context>(dev_ctx, x);
+  SparseCsrTensor y_csr = CooToCsr<T, Context>(dev_ctx, y);
+  SparseCsrTensor out_csr;
+  out_csr.set_dims(out->dims());
+  MatmulCsrCsrKernel<T>(dev_ctx, x_csr, y_csr, &out_csr);
+  CsrToCooKernel<T>(dev_ctx, out_csr, out);
+}
+
+template <typename T, typename Context>
+void MaskedMatmulCsrKernel(const Context& dev_ctx,
+                           const DenseTensor& x,
+                           const DenseTensor& y,
+                           const SparseCsrTensor& mask,
+                           SparseCsrTensor* out) {
+#if defined(PADDLE_WITH_CUDA)
+  std::vector<int64_t> xdim_vec = vectorize(x.dims());
+  std::vector<int64_t> ydim_vec = vectorize(y.dims());
+  std::vector<int64_t> maskdim_vec = vectorize(mask.dims());
+
+  auto x_ndims = xdim_vec.size();
+  auto y_ndims = ydim_vec.size();
+  auto mask_ndims = maskdim_vec.size();
+
+  PADDLE_ENFORCE_EQ(x_ndims,
+                    y_ndims,
+                    common::errors::PreconditionNotMet(
+                        "The dims size of Input(x) and Input(y) "
+                        "should be equal, But received X's "
+                        "dimensions=%d, Y's dimensions=%d.",
+                        x_ndims,
+                        y_ndims));
+  PADDLE_ENFORCE_EQ(x_ndims,
+                    mask_ndims,
+                    common::errors::PreconditionNotMet(
+                        "The dims size of Input(x) and Input(mask) "
+                        "should be equal, But received X's "
+                        "dimensions=%d, mask's dimensions=%d.",
+                        x_ndims,
+                        mask_ndims));
+  PADDLE_ENFORCE_GE(
+      x_ndims,
+      2,
+      common::errors::InvalidArgument("the dims size of Input(x) and "
+                                      "Input(y) must be greater than "
+                                      "or equal to 2."));
+
+  for (size_t i = 0; i < x_ndims - 2; ++i) {
+    PADDLE_ENFORCE_EQ(xdim_vec[i],
+                      ydim_vec[i],
+                      common::errors::InvalidArgument(
+                          "x.dim[%d] and x.dim[%d] must match.", i, i));
+    PADDLE_ENFORCE_EQ(xdim_vec[i],
+                      maskdim_vec[i],
+                      common::errors::InvalidArgument(
+                          "x.dim[%d] and mask.dim[%d] must match.", i, i));
+  }
+
+  PADDLE_ENFORCE_GE(
+      xdim_vec[x_ndims - 1],
+      ydim_vec[y_ndims - 2],
+      common::errors::PreconditionNotMet(
+          "The shape of Input(x) and Input(y) is not suitable for matmul "
+          "operation, x_dim[-1] must be equal to y_dim[-2]."));
+
+  PADDLE_ENFORCE_EQ(
+      maskdim_vec[mask_ndims - 2],
+      xdim_vec[x_ndims - 2],
+      common::errors::PreconditionNotMet(
+          "The shape of Input(x) and Input(y) is not suitable for matmul "
+          "operation, mask_dim[-2] must be equal to x_dim[-2]."));
+
+  PADDLE_ENFORCE_EQ(
+      maskdim_vec[mask_ndims - 1],
+      ydim_vec[y_ndims - 1],
+      common::errors::PreconditionNotMet(
+          "The shape of Input(x) and Input(y) is not suitable for matmul "
+          "operation, mask_dim[-1] must be equal to y_dim[-1]."));
+
+  // InferMeta of SparseCsrTensor 'out', CreateLikeInferMeta
+  EmptyLikeCsrKernel<T, Context>(dev_ctx, mask, out);
+
+  auto sparse_blas = funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
+  sparse_blas.SDDMM(
+      false, false, static_cast<T>(1), x, y, static_cast<T>(0), out);
+#endif
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(matmul_csr_dense,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MatmulCsrDenseKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(matmul_coo_dense,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MatmulCooDenseKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(matmul_coo_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MatmulCooCooKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(matmul_csr_csr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MatmulCsrCsrKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+  kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(masked_matmul_csr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaskedMatmulCsrKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/sparse/gpu/mv_grad_kernel.cu

@@ -0,0 +1,155 @@
+/* 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/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/funcs/sparse/sparse_blas.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename IntT>
+__global__ void MvCooGradGpuKernel(const T *dout,
+                                   const T *vec,
+                                   const IntT *dx_indices,
+                                   T *dx_values,
+                                   int nnz) {
+  int idx = blockDim.x * blockIdx.x + threadIdx.x;
+  for (; idx < nnz; idx += blockDim.x * gridDim.x) {
+    int i = dx_indices[idx];
+    int j = dx_indices[idx + nnz];
+    dx_values[idx] = dout[i] * vec[j];
+  }
+}
+
+template <typename T, typename IntT>
+__global__ void MvCsrGradGpuKernel(const T *dout,
+                                   const T *vec,
+                                   const IntT *dx_crows,
+                                   const IntT *dx_cols,
+                                   T *dx_values,
+                                   int row_number) {
+  int i = blockIdx.x * blockDim.x + threadIdx.x;
+  for (; i < row_number; i += gridDim.x * blockDim.x) {
+    int row_first = static_cast<int>(dx_crows[i]);
+    int row_nnz = static_cast<int>(dx_crows[i + 1] - dx_crows[i]);
+
+    int non_zero_idx = blockIdx.y * blockDim.y + threadIdx.y;
+    for (; non_zero_idx < row_nnz; non_zero_idx += gridDim.y * blockDim.y) {
+      int j = dx_cols[row_first + non_zero_idx];
+      dx_values[row_first + non_zero_idx] = dout[i] * vec[j];
+    }
+  }
+}
+
+template <typename T, typename Context>
+void MvCooGradKernel(const Context &dev_ctx,
+                     const SparseCooTensor &x,
+                     const DenseTensor &vec,
+                     const DenseTensor &dout,
+                     SparseCooTensor *dx,
+                     DenseTensor *dvec) {
+  // dx{SparseCoo} = dout{Dense} * vec'{Dense}
+  if (dx) {
+    // InferMeta of SparseCooTensor 'dx', CreateLikeInferMeta
+    EmptyLikeCooKernel<T, Context>(dev_ctx, x, dx);
+    auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, dx->nnz());
+    PD_VISIT_BASE_INTEGRAL_TYPES(
+        dx->indices().dtype(), "MvCooGradKernel", ([&] {
+          MvCooGradGpuKernel<T>
+              <<<config.block_per_grid.x,
+                 config.thread_per_block.x,
+                 0,
+                 dev_ctx.stream()>>>(dout.data<T>(),
+                                     vec.data<T>(),
+                                     dx->indices().data<data_t>(),
+                                     dx->mutable_values()->data<T>(),
+                                     dx->nnz());
+        }));
+  }
+
+  // dvec{Dense} = x'{SparseCoo} * dout{Dense}
+  if (dvec) {
+#if defined(PADDLE_WITH_CUDA)
+    // InferMeta of DenseTensor 'dvec'
+    dvec->Resize(vec.dims());
+    dev_ctx.template Alloc<T>(dvec);
+
+    auto sparse_blas = funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
+    sparse_blas.SPMV(true, static_cast<T>(1), x, dout, static_cast<T>(0), dvec);
+#endif
+  }
+}
+
+template <typename T, typename Context>
+void MvCsrGradKernel(const Context &dev_ctx,
+                     const SparseCsrTensor &x,
+                     const DenseTensor &vec,
+                     const DenseTensor &dout,
+                     SparseCsrTensor *dx,
+                     DenseTensor *dvec) {
+  // dx{SparseCsr} = dout{Dense} * vec'{Dense}
+  if (dx) {
+    // InferMeta of SparseCsrTensor 'dx', CreateLikeInferMeta
+    EmptyLikeCsrKernel<T, Context>(dev_ctx, x, dx);
+
+    int64_t row_number = dx->dims()[0];
+    int64_t col_number = dx->dims()[1];
+    auto config =
+        backends::gpu::GetGpuLaunchConfig2D(dev_ctx, col_number, row_number);
+    PD_VISIT_BASE_INTEGRAL_TYPES(dx->crows().dtype(), "MvCsrGradKernel", ([&] {
+                                   MvCsrGradGpuKernel<T>
+                                       <<<config.block_per_grid.x,
+                                          config.thread_per_block.x,
+                                          0,
+                                          dev_ctx.stream()>>>(
+                                           dout.data<T>(),
+                                           vec.data<T>(),
+                                           dx->crows().data<data_t>(),
+                                           dx->cols().data<data_t>(),
+                                           dx->mutable_values()->data<T>(),
+                                           row_number);
+                                 }));
+  }
+
+  // dvec{Dense} = x'{SparseCsr} * dout{Dense}
+  if (dvec) {
+#if defined(PADDLE_WITH_CUDA)
+    // InferMeta of DenseTensor 'dvec'
+    dvec->Resize(vec.dims());
+    dev_ctx.template Alloc<T>(dvec);
+
+    auto sparse_blas = funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
+    sparse_blas.SPMV(true, static_cast<T>(1), x, dout, static_cast<T>(0), dvec);
+#endif
+  }
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(
+    mv_coo_grad, GPU, ALL_LAYOUT, phi::sparse::MvCooGradKernel, float, double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(
+    mv_csr_grad, GPU, ALL_LAYOUT, phi::sparse::MvCsrGradKernel, float, double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}

paddle/phi/kernels/sparse/gpu/mv_kernel.cu

@@ -0,0 +1,83 @@
+/* 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/common/ddim.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/sparse/sparse_blas.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context, typename TensorType>
+void MvKernelImpl(const Context& dev_ctx,
+                  const TensorType& x,
+                  const DenseTensor& vec,
+                  DenseTensor* out) {
+#if defined(PADDLE_WITH_CUDA)
+  std::vector<int64_t> x_dim = vectorize(x.dims());
+  std::vector<int64_t> vec_dim = vectorize(vec.dims());
+  auto x_ndims = x_dim.size();
+  auto vec_ndims = vec_dim.size();
+  PADDLE_ENFORCE_EQ(x_ndims,
+                    2,
+                    common::errors::InvalidArgument(
+                        "the dims size of Input(x) must be equal to 2."));
+  PADDLE_ENFORCE_EQ(vec_ndims,
+                    1,
+                    common::errors::InvalidArgument(
+                        "the dims size of Input(vec) must be equal to 1."));
+  PADDLE_ENFORCE_EQ(x_dim[x_ndims - 1],
+                    vec_dim[vec_ndims - 1],
+                    common::errors::PreconditionNotMet(
+                        "The shape of Input(x) and Input(vec) is not "
+                        "suitable for mv operation, "
+                        "x_dim[-1] must be equal to vec_dim[-1]."));
+  std::vector<int64_t> out_dim = {x_dim[x_ndims - 2]};
+  out->Resize(out_dim);
+  dev_ctx.template Alloc<T>(out);
+  auto sparse_blas = funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
+  sparse_blas.SPMV(false, static_cast<T>(1), x, vec, static_cast<T>(0), out);
+#endif
+}
+
+template <typename T, typename Context>
+void MvCooKernel(const Context& dev_ctx,
+                 const SparseCooTensor& x,
+                 const DenseTensor& vec,
+                 DenseTensor* out) {
+  MvKernelImpl<T>(dev_ctx, x, vec, out);
+}
+
+template <typename T, typename Context>
+void MvCsrKernel(const Context& dev_ctx,
+                 const SparseCsrTensor& x,
+                 const DenseTensor& vec,
+                 DenseTensor* out) {
+  MvKernelImpl<T>(dev_ctx, x, vec, out);
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(
+    mv_csr, GPU, ALL_LAYOUT, phi::sparse::MvCsrKernel, float, double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(
+    mv_coo, GPU, ALL_LAYOUT, phi::sparse::MvCooKernel, float, double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/gpu/pool_grad_kernel.cu

@@ -0,0 +1,135 @@
+/* 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/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.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/pooling.h"
+#include "paddle/phi/kernels/funcs/sparse/convolution.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename IntT = int>
+__global__ void MaxPoolGradCudaKernel(const T* in_features_ptr,
+                                      const T* out_features_ptr,
+                                      const T* out_grad_ptr,
+                                      const IntT* rulebook_ptr,
+                                      const int n,
+                                      const int rulebook_len,
+                                      const int channels,
+                                      T* x_grad_ptr) {
+  funcs::MaxPoolGrad<T> grad_functor;
+  CUDA_KERNEL_LOOP_TYPE(i, n * channels, int64_t) {
+    int real_i = i / channels;
+    int c = i - real_i * channels;
+    IntT in_i = rulebook_ptr[real_i];
+    IntT out_i = rulebook_ptr[real_i + rulebook_len];
+    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 IntT = int>
+void MaxPoolCooGradGPUKernel(const GPUContext& 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];
+  // TODO(large-tensor): downstream functors may still use int; guard until
+  // upgraded.
+  int64_t in_channels = x.dims()[4];
+
+  int64_t rulebook_len = 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>();
+  funcs::SetConstant<GPUContext, T> set_zero;
+  set_zero(dev_ctx, &x_grad_values, static_cast<T>(0.0f));
+  phi::Copy<GPUContext>(
+      dev_ctx, x.indices(), dev_ctx.GetPlace(), false, &x_grad_indices);
+
+  for (int i = 0; i < kernel_size; i++) {
+    if (counter_ptr[i] <= 0) {
+      continue;
+    }
+
+    auto config = backends::gpu::GetGpuLaunchConfig1D(
+        dev_ctx, counter_ptr[i] * in_channels, 1);
+    MaxPoolGradCudaKernel<T, IntT>
+        <<<config.block_per_grid.x,
+           config.thread_per_block.x,
+           0,
+           dev_ctx.stream()>>>(in_features_ptr,
+                               out_features_ptr,
+                               out_grad_ptr,
+                               rulebook_ptr + offsets[i],
+                               counter_ptr[i],
+                               rulebook_len,
+                               in_channels,
+                               x_grad_ptr);
+  }
+}
+
+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(), "MaxPoolCooGradGPUKernel", ([&] {
+        MaxPoolCooGradGPUKernel<T, data_t>(
+            dev_ctx, x, rulebook, counter, out, out_grad, kernel_sizes, x_grad);
+      }));
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(maxpool_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaxPoolCooGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/gpu/pool_kernel.cu

@@ -0,0 +1,167 @@
+/* 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/gpu/conv.cu.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename IntT = int>
+__global__ void MaxPoolCudaKernel(const T* in_features_ptr,
+                                  const IntT* rulebook_ptr,
+                                  const int n,
+                                  const int rulebook_len,
+                                  const int channels,
+                                  T* out_features_ptr) {
+  funcs::MaxPool<T> max_pool_functor;
+  CUDA_KERNEL_LOOP_TYPE(i, n * channels, int64_t) {
+    int real_i = i / channels;
+    int channel_i = i - real_i * channels;
+    IntT in_i = rulebook_ptr[real_i];
+    IntT out_i = rulebook_ptr[real_i + rulebook_len];
+    max_pool_functor.compute(in_features_ptr[in_i * channels + channel_i],
+                             &out_features_ptr[out_i * channels + channel_i]);
+  }
+}
+
+/**
+ * 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 MaxPoolCooGPUKernel(const GPUContext& 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, x_dims[4], 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> offsets(kernel_size + 1), h_counter(kernel_size);
+  DenseTensorMeta counter_meta(
+      DataType::INT32, {kernel_size}, DataLayout::NCHW);
+  DenseTensor counter_per_kernel = Empty(dev_ctx, std::move(counter_meta));
+  DenseTensor offsets_per_kernel = Empty(dev_ctx, std::move(counter_meta));
+  DenseTensorMeta index_meta(DataType::INT32, {1}, DataLayout::NCHW);
+  DenseTensor out_index = Empty(dev_ctx, std::move(index_meta));
+  DenseTensor unique_value = Empty(dev_ctx, std::move(index_meta));
+
+  // 1. product rulebook
+  int rulebook_len = ProductRuleBook<T, GPUContext, IntT>(dev_ctx,
+                                                          x,
+                                                          real_kernel_sizes,
+                                                          paddings,
+                                                          dilations,
+                                                          strides,
+                                                          out_dims,
+                                                          false,
+                                                          rulebook,
+                                                          &counter_per_kernel,
+                                                          &offsets_per_kernel,
+                                                          &out_index,
+                                                          &unique_value,
+                                                          out,
+                                                          h_counter.data(),
+                                                          offsets.data());
+
+  const IntT* rulebook_ptr = rulebook->data<IntT>();
+
+  T* out_features_ptr = out->mutable_values()->data<T>();
+  const T* in_features_ptr = x.values().data<T>();
+  counter->Resize({kernel_size});
+  int* counter_ptr = dev_ctx.template HostAlloc<int>(counter);
+  memcpy(counter_ptr, h_counter.data(), h_counter.size() * sizeof(int));
+// 2. max pool
+#ifdef PADDLE_WITH_HIP
+  thrust::fill(thrust::hip::par.on(dev_ctx.stream()),
+#else
+  thrust::fill(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+               out_features_ptr,
+               out_features_ptr + out->values().numel(),
+               static_cast<T>(0));
+  // TODO(zhangkaihuo) Replacing multiple calls with one kernel may be faster
+  for (int i = 0; i < kernel_size; i++) {
+    if (h_counter[i] <= 0) {
+      continue;
+    }
+
+    auto config = backends::gpu::GetGpuLaunchConfig1D(
+        dev_ctx, h_counter[i] * in_channels, 1);
+    MaxPoolCudaKernel<T, IntT><<<config.block_per_grid.x,
+                                 config.thread_per_block.x,
+                                 0,
+                                 dev_ctx.stream()>>>(in_features_ptr,
+                                                     rulebook_ptr + offsets[i],
+                                                     h_counter[i],
+                                                     rulebook_len,
+                                                     in_channels,
+                                                     out_features_ptr);
+  }
+}
+
+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(), "MaxPoolCooGPUKernel", ([&] {
+        MaxPoolCooGPUKernel<T, data_t>(dev_ctx,
+                                       x,
+                                       kernel_sizes,
+                                       paddings,
+                                       dilations,
+                                       strides,
+                                       out,
+                                       rulebook,
+                                       counter);
+      }));
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(maxpool_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::MaxPoolCooKernel,
+                   float,
+                   double,
+                   phi::float16) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/gpu/reshape_grad_kernel.cu

@@ -0,0 +1,77 @@
+// 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/gpu/gpu_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 {
+namespace sparse {
+
+// just copy from paddle\phi\kernels\sparse\cpu\reshape_grad_kernel.cc
+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);
+}
+
+// just copy from paddle\phi\kernels\sparse\cpu\reshape_grad_kernel.cc
+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 sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(reshape_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ReshapeCooGradKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(reshape_csr_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ReshapeCsrGradKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/gpu/reshape_kernel.cu

@@ -0,0 +1,178 @@
+// 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/gpu/gpu_context.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/impl/unary_kernel_impl.h"
+
+#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename IntT>
+__global__ void ReshapeCooCudaKernel(const IntT* x_indices_data,
+                                     const int num_x_sparse_part_dims,
+                                     const int num_out_sparse_part_dims,
+                                     const int64_t x_nnz,
+                                     const int64_t* x_sparse_part_strides,
+                                     const int64_t* out_sparse_part_strides,
+                                     IntT* out_indices_data) {
+  CUDA_KERNEL_LOOP_TYPE(j, x_nnz, int64_t) {
+    IntT location = 0;
+    for (int i = 0; i < num_x_sparse_part_dims; ++i) {
+      location += x_indices_data[i * x_nnz + j] *
+                  static_cast<IntT>(x_sparse_part_strides[i]);
+    }
+    for (int i = 0; i < num_out_sparse_part_dims; ++i) {
+      out_indices_data[i * x_nnz + j] =
+          location / static_cast<IntT>(out_sparse_part_strides[i]);
+      location %= static_cast<IntT>(out_sparse_part_strides[i]);
+    }
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void ReshapeCooGPUKernel(const Context& dev_ctx,
+                         const SparseCooTensor& x,
+                         const phi::IntArray& shape,
+                         SparseCooTensor* out) {
+  int64_t x_nnz = x.nnz();
+  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 out 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 *destination_x_sparse_part_strides,
+      *destination_out_sparse_part_strides;
+
+  auto destination_x_sparse_part_strides_tensor = memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int64_t) * x_sparse_part_strides.size(),
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  destination_x_sparse_part_strides = reinterpret_cast<int64_t*>(
+      destination_x_sparse_part_strides_tensor->ptr());
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     reinterpret_cast<void*>(destination_x_sparse_part_strides),
+                     phi::CPUPlace(),
+                     x_sparse_part_strides.Get(),
+                     sizeof(int64_t) * x_sparse_part_strides.size(),
+                     dev_ctx.stream());
+
+  auto destination_out_sparse_part_strides_tensor = memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int64_t) * out_sparse_part_strides.size(),
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  destination_out_sparse_part_strides = reinterpret_cast<int64_t*>(
+      destination_out_sparse_part_strides_tensor->ptr());
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     destination_out_sparse_part_strides,
+                     phi::CPUPlace(),
+                     out_sparse_part_strides.Get(),
+                     sizeof(int64_t) * out_sparse_part_strides.size(),
+                     dev_ctx.stream());
+
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x_nnz, 1);
+  ReshapeCooCudaKernel<<<config.block_per_grid.x,
+                         config.thread_per_block.x,
+                         0,
+                         dev_ctx.stream()>>>(
+      x_indices_data,
+      x_sparse_part_dims.size(),
+      out_sparse_part_dims.size(),
+      x_nnz,
+      destination_x_sparse_part_strides,
+      destination_out_sparse_part_strides,
+      out_indices_data);
+}
+
+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(), "ReshapeCooGPUKernel", ([&] {
+        ReshapeCooGPUKernel<T, data_t, Context>(dev_ctx, x, shape, out);
+      }));
+}
+
+// just copy from paddle\phi\kernels\sparse\cpu\reshape_kernel.cc
+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 sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(reshape_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ReshapeCooKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(reshape_csr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::ReshapeCsrKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/gpu/slice_grad_kernel.cu

@@ -0,0 +1,375 @@
+// 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/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+#include "paddle/phi/backends/gpu/gpu_primitives.h"
+#include "paddle/phi/common/memory_utils.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 {
+namespace sparse {
+
+template <typename T>
+__global__ void GetCooInputGradCudaKernel(const int64_t* out_grad_indices_data,
+                                          const T* out_grad_values_data,
+                                          const int64_t* axes,
+                                          const int64_t* starts,
+                                          const int64_t axes_size,
+                                          const int64_t sparse_dim,
+                                          const int64_t out_grad_nnz,
+                                          int64_t* dx_indices_data,
+                                          T* dx_values_data) {
+  CUDA_KERNEL_LOOP_TYPE(j, out_grad_nnz, int64_t) {
+    // 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];
+  }
+}
+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) {
+  const DDim& x_dims = x.dims();
+
+  // copy axes to device
+  auto d_axes_tensor = memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int64_t) * axes.size(),
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  int64_t* d_axes = reinterpret_cast<int64_t*>(d_axes_tensor->ptr());
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     d_axes,
+                     phi::CPUPlace(),
+                     axes.data(),
+                     sizeof(int64_t) * axes.size(),
+                     dev_ctx.stream());
+
+  // copy starts to device
+  auto d_starts_tensor = memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int64_t) * starts.size(),
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  int64_t* d_starts = reinterpret_cast<int64_t*>(d_starts_tensor->ptr());
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     d_starts,
+                     phi::CPUPlace(),
+                     starts.data(),
+                     sizeof(int64_t) * starts.size(),
+                     dev_ctx.stream());
+
+  // Step2: Set indices and values of 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>();
+
+  x_grad->SetMember(dx_indices, dx_values, x.dims(), x.coalesced());
+
+  auto config =
+      backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_grad_nnz + 1, 1);
+  GetCooInputGradCudaKernel<T><<<config.block_per_grid.x,
+                                 config.thread_per_block.x,
+                                 0,
+                                 dev_ctx.stream()>>>(out_grad_indices_data,
+                                                     out_grad_values_data,
+                                                     d_axes,
+                                                     d_starts,
+                                                     axes.size(),
+                                                     sparse_dim,
+                                                     out_grad_nnz,
+                                                     dx_indices_data,
+                                                     dx_values_data);
+}
+
+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();
+  // Check and update sparse slice attrs
+  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>
+__global__ void GetCsrInputColsValuesCudaKernel(
+    const int64_t* out_grad_cols_data,
+    const T* out_grad_values_data,
+    const int64_t out_grad_nnz,
+    const int64_t cols_start,
+    int64_t* dx_cols_data,
+    T* dx_values_data) {
+  CUDA_KERNEL_LOOP_TYPE(i, out_grad_nnz, int64_t) {
+    dx_cols_data[i] = out_grad_cols_data[i] + cols_start;
+    dx_values_data[i] = out_grad_values_data[i];
+  }
+}
+
+__global__ void GetCsrInputCrowsCudaKernel(
+    const int64_t* out_grad_crows_data,
+    const int64_t out_grad_n_rows,
+    const int64_t out_grad_nnz,
+    const int64_t x_n_rows,
+    const int64_t rows_start,
+    const int64_t rows_end,
+    int64_t* dx_crows_data,
+    const int64_t dx_crows_offset = 0,
+    const int64_t out_grad_crows_offset = 0) {
+  CUDA_KERNEL_LOOP_TYPE(i, x_n_rows + 1, int64_t) {
+    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>();
+  x_grad->SetMember(dx_crows, dx_cols, dx_values, x.dims());
+
+  // set cols and values
+  auto config =
+      backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_grad_nnz + 1, 1);
+  GetCsrInputColsValuesCudaKernel<T><<<config.block_per_grid.x,
+                                       config.thread_per_block.x,
+                                       0,
+                                       dev_ctx.stream()>>>(out_grad_cols_data,
+                                                           out_grad_values_data,
+                                                           out_grad_nnz,
+                                                           starts[1],
+                                                           dx_cols_data,
+                                                           dx_values_data);
+  config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, n_rows + 1, 1);
+  GetCsrInputCrowsCudaKernel<<<config.block_per_grid.x,
+                               config.thread_per_block.x,
+                               0,
+                               dev_ctx.stream()>>>(out_grad_crows_data,
+                                                   out_grad.dims()[0],
+                                                   out_grad_nnz,
+                                                   x.dims()[0],
+                                                   starts[0],
+                                                   ends[0],
+                                                   dx_crows_data,
+                                                   0,
+                                                   0);
+}
+
+__global__ void GetCsrInputCrowsPart1CudaKernel(const int64_t n_rows,
+                                                const int64_t dim0_idx,
+                                                int64_t* dx_crows_data) {
+  CUDA_KERNEL_LOOP_TYPE(j, n_rows + 1, int64_t) {
+    dx_crows_data[dim0_idx * (n_rows + 1) + j] = 0;
+  }
+}
+
+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>();
+  x_grad->SetMember(dx_crows, dx_cols, dx_values, x.dims());
+
+  // set cols and values
+  auto config =
+      backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_grad_nnz + 1, 1);
+  GetCsrInputColsValuesCudaKernel<T><<<config.block_per_grid.x,
+                                       config.thread_per_block.x,
+                                       0,
+                                       dev_ctx.stream()>>>(out_grad_cols_data,
+                                                           out_grad_values_data,
+                                                           out_grad_nnz,
+                                                           starts[2],
+                                                           dx_cols_data,
+                                                           dx_values_data);
+  // 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]) {
+      config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, n_rows + 1, 1);
+      GetCsrInputCrowsPart1CudaKernel<<<config.block_per_grid.x,
+                                        config.thread_per_block.x,
+                                        0,
+                                        dev_ctx.stream()>>>(
+          n_rows, i, dx_crows_data);
+    } else {
+      int64_t dx_crows_offset = i * (n_rows + 1);
+      int64_t out_grad_crows_offset = (i - starts[0]) * (out_grad_n_rows + 1);
+      config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, n_rows + 1, 1);
+      GetCsrInputCrowsCudaKernel<<<config.block_per_grid.x,
+                                   config.thread_per_block.x,
+                                   0,
+                                   dev_ctx.stream()>>>(out_grad_crows_data,
+                                                       out_grad_n_rows,
+                                                       out_grad_nnz,
+                                                       n_rows,
+                                                       starts[1],
+                                                       ends[1],
+                                                       dx_crows_data,
+                                                       dx_crows_offset,
+                                                       out_grad_crows_offset);
+    }
+  }
+}
+
+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 sparse
+}  // namespace phi
+PD_REGISTER_KERNEL(slice_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SliceCooGradKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(slice_csr_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SliceCsrGradKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/gpu/slice_kernel.cu

@@ -0,0 +1,664 @@
+// 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 <thrust/device_vector.h>
+#include <thrust/sort.h>
+
+#include "paddle/phi/kernels/sparse/unary_kernel.h"
+
+#include "paddle/common/ddim.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+#include "paddle/phi/backends/gpu/gpu_primitives.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/funcs/slice_utils.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename IntT>
+__global__ void GetCooNonZeroNumberCudaKernel(const IntT* x_indices_data,
+                                              const int64_t* axes,
+                                              const int64_t* starts,
+                                              const int64_t* ends,
+                                              const int64_t axes_size,
+                                              const int64_t x_nnz,
+                                              int* out_nnz,
+                                              IntT* out_nnz_indices) {
+  CUDA_KERNEL_LOOP_TYPE(j, x_nnz, int64_t) {
+    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;
+    int old_value = atomicAdd(out_nnz, 1);
+    out_nnz_indices[old_value] = j;
+  }
+}
+
+template <typename T, typename IntT>
+__global__ void GetCooOutCudaKernel(const IntT* x_indices_data,
+                                    const T* x_values_data,
+                                    const int64_t* axes,
+                                    const int64_t* starts,
+                                    const int64_t axes_size,
+                                    const int64_t sparse_dim,
+                                    const int64_t x_nnz,
+                                    const int64_t out_nnz,
+                                    const IntT* out_nnz_indices,
+                                    IntT* out_indices_data,
+                                    T* out_values_data) {
+  CUDA_KERNEL_LOOP_TYPE(index, out_nnz, int64_t) {
+    // index is in the order of the non-zero elements in out
+    // out_nnz_indices[index] is the valid index in x's non-zero elements, where
+    // the `hit` is true.
+    IntT j = out_nnz_indices[index];
+    // 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];
+    }
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void SliceCooGPUCompute(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 the number of non zero elements
+  DenseTensor d_out_nnz = Empty<int32_t>(dev_ctx, {1});
+  int* d_out_nnz_ptr = d_out_nnz.data<int32_t>();
+  backends::gpu::GpuMemsetAsync(
+      d_out_nnz_ptr, 0, sizeof(int32_t), dev_ctx.stream());
+
+  // out_nnz_indices is the indices where the data is valid in out
+  // the length of the out_nnz_indices must be less than x.nnz()
+  DenseTensor d_out_nnz_indices = Empty<IntT>(dev_ctx, {x.nnz()});
+  auto* d_out_nnz_indices_ptr = d_out_nnz_indices.data<IntT>();
+  backends::gpu::GpuMemsetAsync(
+      d_out_nnz_indices_ptr, 0, sizeof(IntT), dev_ctx.stream());
+
+  // copy axes to device
+  auto d_axes_tensor = memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int64_t) * axes.size(),
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  int64_t* d_axes = reinterpret_cast<int64_t*>(d_axes_tensor->ptr());
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     d_axes,
+                     phi::CPUPlace(),
+                     axes.data(),
+                     sizeof(int64_t) * axes.size(),
+                     dev_ctx.stream());
+
+  // copy starts to device
+  auto d_starts_tensor = memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int64_t) * starts.size(),
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  int64_t* d_starts = reinterpret_cast<int64_t*>(d_starts_tensor->ptr());
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     d_starts,
+                     phi::CPUPlace(),
+                     starts.data(),
+                     sizeof(int64_t) * starts.size(),
+                     dev_ctx.stream());
+
+  // copy ends to device
+  auto d_ends_tensor = memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int64_t) * ends.size(),
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  int64_t* d_ends = reinterpret_cast<int64_t*>(d_ends_tensor->ptr());
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     d_ends,
+                     phi::CPUPlace(),
+                     ends.data(),
+                     sizeof(int64_t) * ends.size(),
+                     dev_ctx.stream());
+
+  const auto* x_indices_data = x.indices().data<IntT>();
+
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x.nnz() + 1, 1);
+  GetCooNonZeroNumberCudaKernel<IntT>
+      <<<config.block_per_grid.x,
+         config.thread_per_block.x,
+         0,
+         dev_ctx.stream()>>>(x_indices_data,
+                             d_axes,
+                             d_starts,
+                             d_ends,
+                             axes.size(),
+                             x.nnz(),
+                             d_out_nnz_ptr,
+                             d_out_nnz_indices_ptr);
+
+  // copy d_out_nnz from device to host (out_nnz)
+  int32_t out_nnz = 0;
+  backends::gpu::GpuMemcpyAsync(&out_nnz,
+                                d_out_nnz_ptr,
+                                sizeof(int32_t),
+                                gpuMemcpyDeviceToHost,
+                                dev_ctx.stream());
+  dev_ctx.Wait();
+  // sort `d_out_nnz_indices_ptr`
+  d_out_nnz_indices.Resize({out_nnz});
+#ifdef PADDLE_WITH_HIP
+  thrust::sort(thrust::hip::par.on(dev_ctx.stream()),
+#else
+  thrust::sort(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+               d_out_nnz_indices_ptr,
+               d_out_nnz_indices_ptr + out_nnz);
+
+  // Step3: Get the values and indices of output
+  auto sparse_dim = static_cast<int64_t>(x.sparse_dim());
+  DenseTensor out_indices =
+      Empty<IntT, Context>(dev_ctx, {sparse_dim, out_nnz});
+  DenseTensor out_values = Empty<T, Context>(dev_ctx, {out_nnz});
+  out->SetMember(out_indices, out_values, out_dims, x.coalesced());
+
+  auto* out_indices_data = out_indices.data<IntT>();
+  auto* out_values_data = out_values.data<T>();
+  const auto* x_values_data = x.values().data<T>();
+
+  config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_nnz + 1, 1);
+  GetCooOutCudaKernel<T, IntT>
+      <<<config.block_per_grid.x,
+         config.thread_per_block.x,
+         0,
+         dev_ctx.stream()>>>(x_indices_data,
+                             x_values_data,
+                             d_axes,
+                             d_starts,
+                             axes.size(),
+                             sparse_dim,
+                             x.nnz(),
+                             static_cast<int64_t>(out_nnz),
+                             d_out_nnz_indices_ptr,
+                             out_indices_data,
+                             out_values_data);
+}
+
+template <typename T, typename IntT, typename Context>
+void SliceCooGPUKernel(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);
+  SliceCooGPUCompute<T, IntT, Context>(
+      dev_ctx, x, axes_vec, starts_vec, ends_vec, out);
+}
+
+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) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "SliceCooGPUKernel", ([&] {
+                                 SliceCooGPUKernel<T, data_t, Context>(
+                                     dev_ctx, x, axes, starts, ends, out);
+                               }));
+}
+
+__global__ void GetCsr2DNonZeroNumberCudaKernel(const int64_t* x_crows_data,
+                                                const int64_t* x_cols_data,
+                                                const int64_t x_crows_start,
+                                                const int64_t x_crows_end,
+                                                const int64_t min_col,
+                                                const int64_t max_col,
+                                                int64_t* out_crows_data) {
+  CUDA_KERNEL_LOOP_TYPE(i, x_crows_end - x_crows_start, int64_t) {
+    if (i == 0) {
+      out_crows_data[0] = 0;
+    }
+    int64_t st = x_crows_data[x_crows_start + i];
+    int64_t ed = x_crows_data[x_crows_start + i + 1];
+    out_crows_data[i + 1] = 0;
+    for (int64_t jj = st; jj < ed; ++jj) {
+      if (x_cols_data[jj] >= min_col && x_cols_data[jj] < max_col) {
+        out_crows_data[i + 1] += 1;
+      }
+    }
+  }
+}
+
+template <typename T>
+__global__ void GetCsr2DCudaKernel(const int64_t* x_crows_data,
+                                   const int64_t* x_cols_data,
+                                   const T* x_values_data,
+                                   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* out_crows_data,
+                                   int64_t* out_cols_data,
+                                   T* out_values_data) {
+  CUDA_KERNEL_LOOP_TYPE(i, x_crows_end - x_crows_start, int64_t) {
+    int64_t st = x_crows_data[x_crows_start + i];
+    int64_t ed = x_crows_data[x_crows_start + i + 1];
+    int64_t index = out_crows_data[i];
+    for (int64_t jj = st; jj < ed; ++jj) {
+      if (x_cols_data[jj] >= min_col && x_cols_data[jj] < max_col) {
+        out_cols_data[index] = x_cols_data[jj] - min_col;
+        out_values_data[index] = x_values_data[jj];
+        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) {
+  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>();
+  // Step1: Get the number of non zero elements for out and out_crows
+  int64_t out_n_rows = ends[0] - starts[0];
+  DenseTensor out_crows = Empty<int64_t, Context>(dev_ctx, {out_n_rows + 1});
+  auto* out_crows_data = out_crows.data<int64_t>();
+
+  auto config =
+      backends::gpu::GetGpuLaunchConfig1D(dev_ctx, ends[0] - starts[0] + 1, 1);
+  GetCsr2DNonZeroNumberCudaKernel<<<config.block_per_grid.x,
+                                    config.thread_per_block.x,
+                                    0,
+                                    dev_ctx.stream()>>>(x_crows_data,
+                                                        x_cols_data,
+                                                        starts[0],
+                                                        ends[0],
+                                                        starts[1],
+                                                        ends[1],
+                                                        out_crows_data);
+#ifdef PADDLE_WITH_HIP
+  thrust::inclusive_scan(thrust::hip::par.on(dev_ctx.stream()),
+#else
+  thrust::inclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                         out_crows_data,
+                         out_crows_data + out_n_rows + 1,
+                         out_crows_data);
+  int64_t out_nnz = 0;
+  backends::gpu::GpuMemcpyAsync(&out_nnz,
+                                &out_crows_data[out_n_rows],
+                                sizeof(int64_t),
+                                gpuMemcpyDeviceToHost,
+                                dev_ctx.stream());
+  dev_ctx.Wait();
+  // Step2: Set out
+  DenseTensor out_cols = Empty<int64_t, Context>(dev_ctx, {out_nnz});
+  DenseTensor out_values = Empty<T, Context>(dev_ctx, {out_nnz});
+  out->SetMember(out_crows, out_cols, out_values, out_dims);
+  config =
+      backends::gpu::GetGpuLaunchConfig1D(dev_ctx, ends[0] - starts[0] + 1, 1);
+  GetCsr2DCudaKernel<T><<<config.block_per_grid.x,
+                          config.thread_per_block.x,
+                          0,
+                          dev_ctx.stream()>>>(x_crows_data,
+                                              x_cols_data,
+                                              x_values_data,
+                                              starts[0],
+                                              ends[0],
+                                              starts[1],
+                                              ends[1],
+                                              out_crows.data<int64_t>(),
+                                              out_cols.data<int64_t>(),
+                                              out_values.data<T>());
+}
+
+__global__ void GetXColsOffsetsCudaKernel(const int64_t* x_crows_data,
+                                          const int64_t x_n_rows,
+                                          const int64_t x_dim0,
+                                          int64_t* x_cols_offsets) {
+  CUDA_KERNEL_LOOP_TYPE(i, x_dim0, int64_t) {
+    if (i == 0) {
+      x_cols_offsets[i] = 0;
+    }
+    x_cols_offsets[i + 1] = x_crows_data[(i + 1) * (x_n_rows + 1) - 1];
+  }
+}
+
+__global__ void GetCsr3DNonZeroNumberCudaKernel(const int64_t* x_crows_data,
+                                                const int64_t* x_cols_data,
+                                                const int64_t x_dim0,
+                                                const int64_t x_n_rows,
+                                                const int64_t* x_cols_offsets,
+                                                const int64_t* starts,
+                                                const int64_t* ends,
+                                                const int64_t out_n_rows,
+                                                int64_t* out_crows_data) {
+  CUDA_KERNEL_LOOP_TYPE(i, x_dim0 * (x_n_rows + 1), int64_t) {
+    int64_t dim0_i = i / (x_n_rows + 1);
+    int64_t dim1_i = i % (x_n_rows + 1);
+    if (!(dim0_i >= starts[0] && dim0_i < ends[0])) {
+      continue;
+    }
+    if (!(dim1_i >= starts[1] && dim1_i < ends[1])) {
+      continue;
+    }
+    // the starting index of current 2D Tensor in out_crows
+    int64_t out_dim0_start = (dim0_i - starts[0]) * (out_n_rows + 1);
+    if (dim1_i == starts[1]) {
+      out_crows_data[out_dim0_start] = 0;
+    }
+    int64_t out_crows_idx = out_dim0_start + (dim1_i - starts[1]);
+    int64_t st = x_crows_data[i] + x_cols_offsets[dim0_i];
+    int64_t ed = x_crows_data[i + 1] + x_cols_offsets[dim0_i];
+    out_crows_data[out_crows_idx + 1] = 0;
+    for (int64_t jj = st; jj < ed; ++jj) {
+      if (x_cols_data[jj] >= starts[2] && x_cols_data[jj] < ends[2]) {
+        out_crows_data[out_crows_idx + 1] += 1;
+      }
+    }
+  }
+}
+
+template <typename T>
+__global__ void GetCsr3DCudaKernel(const int64_t* x_crows_data,
+                                   const int64_t* x_cols_data,
+                                   const T* x_values_data,
+                                   const int64_t* x_cols_offsets,
+                                   const int64_t x_dim0,
+                                   const int64_t x_n_rows,
+                                   const int64_t* starts,
+                                   const int64_t* ends,
+                                   const int64_t out_n_rows,
+                                   const int64_t* out_cols_offsets,
+                                   const int64_t* out_crows_data,
+                                   int64_t* out_cols_data,
+                                   T* out_values_data) {
+  CUDA_KERNEL_LOOP_TYPE(i, x_dim0 * (x_n_rows + 1), int64_t) {
+    int dim0_i = i / (x_n_rows + 1);
+    int dim1_i = i % (x_n_rows + 1);
+    if (!(dim0_i >= starts[0] && dim0_i < ends[0])) {
+      continue;
+    }
+    if (!(dim1_i >= starts[1] && dim1_i < ends[1])) {
+      continue;
+    }
+    // the starting index of current 2D Tensor in out_crows
+    int64_t out_dim0_start = (dim0_i - starts[0]) * (out_n_rows + 1);
+    int64_t out_crows_idx = out_dim0_start + (dim1_i - starts[1]);
+    int64_t st = x_crows_data[i] + x_cols_offsets[dim0_i];
+    int64_t ed = x_crows_data[i + 1] + x_cols_offsets[dim0_i];
+    int64_t index = out_crows_data[out_crows_idx];
+    for (int64_t jj = st; jj < ed; ++jj) {
+      if (x_cols_data[jj] >= starts[2] && x_cols_data[jj] < ends[2]) {
+        out_cols_data[out_cols_offsets[out_dim0_start] + index] =
+            x_cols_data[jj] - starts[2];
+        out_values_data[out_cols_offsets[out_dim0_start] + index] =
+            x_values_data[jj];
+        index++;
+      }
+    }
+  }
+}
+
+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>();
+  const auto* x_cols_data = x.cols().data<int64_t>();
+  const auto* x_values_data = x.values().data<T>();
+  const int64_t x_dim0 = x.dims()[0], x_n_rows = x.dims()[1];
+
+  // get x_cols_offsets
+  DenseTensor x_cols_offsets = Empty<int64_t>(dev_ctx, {x_dim0 + 1});
+  auto* x_cols_offsets_data = x_cols_offsets.data<int64_t>();
+
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x_dim0 + 1, 1);
+  GetXColsOffsetsCudaKernel<<<config.block_per_grid.x,
+                              config.thread_per_block.x,
+                              0,
+                              dev_ctx.stream()>>>(
+      x_crows_data, x_n_rows, x_dim0, x_cols_offsets_data);
+
+#ifdef PADDLE_WITH_HIP
+  thrust::inclusive_scan(thrust::hip::par.on(dev_ctx.stream()),
+#else
+  thrust::inclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                         x_cols_offsets_data,
+                         x_cols_offsets_data + x_dim0 + 1,
+                         x_cols_offsets_data);
+
+  // copy starts to device
+  auto d_starts_tensor = memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int64_t) * starts.size(),
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  int64_t* d_starts = reinterpret_cast<int64_t*>(d_starts_tensor->ptr());
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     d_starts,
+                     phi::CPUPlace(),
+                     starts.data(),
+                     sizeof(int64_t) * starts.size(),
+                     dev_ctx.stream());
+
+  // copy ends to device
+  auto d_ends_tensor = memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int64_t) * ends.size(),
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  int64_t* d_ends = reinterpret_cast<int64_t*>(d_ends_tensor->ptr());
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     d_ends,
+                     phi::CPUPlace(),
+                     ends.data(),
+                     sizeof(int64_t) * ends.size(),
+                     dev_ctx.stream());
+
+  // get out_nnz
+  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)});
+  auto* out_crows_data = out_crows.data<int64_t>();
+  config = backends::gpu::GetGpuLaunchConfig1D(
+      dev_ctx, x_dim0 * (x_n_rows + 1) + 1, 1);
+  GetCsr3DNonZeroNumberCudaKernel<<<config.block_per_grid.x,
+                                    config.thread_per_block.x,
+                                    0,
+                                    dev_ctx.stream()>>>(x_crows_data,
+                                                        x_cols_data,
+                                                        x_dim0,
+                                                        x_n_rows,
+                                                        x_cols_offsets_data,
+                                                        d_starts,
+                                                        d_ends,
+                                                        out_n_rows,
+                                                        out_crows_data);
+  DenseTensor out_cols_offsets =
+      Empty<int64_t, Context>(dev_ctx, {out_dim0 * (out_n_rows + 1)});
+  auto* out_cols_offsets_data = out_cols_offsets.data<int64_t>();
+  backends::gpu::GpuMemcpyAsync(out_cols_offsets_data,
+                                out_crows_data,
+                                out_dim0 * (out_n_rows + 1) * sizeof(int64_t),
+                                gpuMemcpyDeviceToDevice,
+                                dev_ctx.stream());
+  dev_ctx.Wait();
+  int64_t out_nnz =
+#ifdef PADDLE_WITH_HIP
+      thrust::reduce(thrust::hip::par.on(dev_ctx.stream()),
+#else
+      thrust::reduce(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                     out_crows_data,
+                     out_crows_data + out_dim0 * (out_n_rows + 1));
+  for (int64_t i = 0; i < out_dim0; ++i) {
+    int64_t st = i * (out_n_rows + 1);
+#ifdef PADDLE_WITH_HIP
+    thrust::inclusive_scan(thrust::hip::par.on(dev_ctx.stream()),
+#else
+    thrust::inclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                           out_crows_data + st,
+                           out_crows_data + st + out_n_rows + 1,
+                           out_crows_data + st);
+  }
+#ifdef PADDLE_WITH_HIP
+  thrust::inclusive_scan(thrust::hip::par.on(dev_ctx.stream()),
+#else
+  thrust::inclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                         out_cols_offsets_data,
+                         out_cols_offsets_data + out_dim0 * (out_n_rows + 1),
+                         out_cols_offsets_data);
+
+  DenseTensor out_cols = Empty<int64_t, Context>(dev_ctx, {out_nnz});
+  auto* out_cols_data = out_cols.data<int64_t>();
+  DenseTensor out_values = Empty<T, Context>(dev_ctx, {out_nnz});
+  auto* out_values_data = out_values.data<T>();
+  out->SetMember(out_crows, out_cols, out_values, out_dims);
+  config = backends::gpu::GetGpuLaunchConfig1D(
+      dev_ctx, x_dim0 * (x_n_rows + 1) + 1, 1);
+  GetCsr3DCudaKernel<T><<<config.block_per_grid.x,
+                          config.thread_per_block.x,
+                          0,
+                          dev_ctx.stream()>>>(x_crows_data,
+                                              x_cols_data,
+                                              x_values_data,
+                                              x_cols_offsets_data,
+                                              x_dim0,
+                                              x_n_rows,
+                                              d_starts,
+                                              d_ends,
+                                              out_n_rows,
+                                              out_cols_offsets_data,
+                                              out_crows_data,
+                                              out_cols_data,
+                                              out_values_data);
+}
+
+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 sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(slice_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SliceCooKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(slice_csr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SliceCsrKernel,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/gpu/softmax_grad_kernel.cu

@@ -0,0 +1,318 @@
+/* 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 <thrust/binary_search.h>
+#include <thrust/device_ptr.h>
+#include <thrust/equal.h>
+#include <thrust/iterator/constant_iterator.h>
+#include <thrust/iterator/discard_iterator.h>
+#include <thrust/sequence.h>
+#include <thrust/sort.h>
+#include <thrust/transform.h>
+
+#include "paddle/phi/backends/gpu/gpu_context.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/math_cuda_utils.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/sparse/softmax.cu.h"
+#include "paddle/phi/kernels/softmax_grad_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename IntT = int>
+__global__ void SoftmaxGradGpuKernel(const IntT* out_crows,
+                                     const T* out_values,
+                                     const T* dout_values,
+                                     T* dx_values,
+                                     int row_number,
+                                     int total_row_number) {
+  // dx = (dout - sum(dout * out)) * out
+  int row = blockIdx.x * blockDim.y + threadIdx.y;
+  int non_zero_idx = threadIdx.x;
+  if (row >= total_row_number) return;
+  int cur_batch = row / row_number;
+  int crow_idx = cur_batch * (row_number + 1) + (row % row_number);
+  int cur_batch_offset = 0;
+  for (int i = 1; i < cur_batch + 1; ++i) {
+    cur_batch_offset += out_crows[i * (row_number + 1) - 1];
+  }
+  int row_first = cur_batch_offset + static_cast<int>(out_crows[crow_idx]);
+  int row_nnz = static_cast<int>(out_crows[crow_idx + 1] - out_crows[crow_idx]);
+  if (row_nnz == 0) return;
+
+  int kIteration = (row_nnz + warpSize - 1) / warpSize;
+
+  T mul_result = 0;
+  for (int i = 0; i < kIteration; ++i) {
+    int idx = non_zero_idx + i * warpSize;
+    if (idx >= row_nnz) break;
+
+    mul_result += out_values[row_first + idx] * dout_values[row_first + idx];
+  }
+  T sum = funcs::WarpReduceSum<T>(mul_result, 0xFFFFFFFF);
+
+  for (int i = 0; i < kIteration; ++i) {
+    int idx = non_zero_idx + i * warpSize;
+    if (idx >= row_nnz) break;
+
+    dx_values[row_first + idx] =
+        (dout_values[row_first + idx] - sum) * out_values[row_first + idx];
+  }
+}
+
+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 total_row_number = 1;
+  int row_number = 1;
+  for (int i = 0; i < out_rank - 1; ++i) {
+    total_row_number *= out_dim[i];
+    if (i == out_rank - 2) {
+      row_number = out_dim[i];
+    }
+  }
+
+  dim3 grid((total_row_number + 3) / 4);
+  dim3 block(32, 4);
+
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      out.crows().dtype(), "SoftmaxCsrGradKernel", ([&] {
+        SoftmaxGradGpuKernel<T, data_t><<<grid, block, 0, dev_ctx.stream()>>>(
+            out.crows().data<data_t>(),
+            out.values().data<T>(),
+            dout.values().data<T>(),
+            dx->mutable_values()->data<T>(),
+            row_number,
+            total_row_number);
+      }));
+}
+
+template <typename T, typename IntT>
+__global__ void SoftmaxCooGradGPURawKernel(IntT* sorted_pool_indices,
+                                           IntT size,
+                                           IntT* pool_sizes,
+                                           IntT* pool_offsets,
+                                           IntT nvalues,
+                                           IntT grad_nnz,
+                                           IntT* grad_offsets,
+                                           IntT* out_offsets,
+                                           IntT* lower_bound_values,
+                                           T* values,
+                                           T* out_values,
+                                           T* grad_values,
+                                           int total_rows) {
+  int row = blockIdx.x * blockDim.y + threadIdx.y;
+  if (row >= total_rows) return;
+
+  int tid = threadIdx.x;
+  int index = row / nvalues;
+  int nval = row % nvalues;
+  IntT offset = pool_offsets[index];
+  IntT* pool_indices = sorted_pool_indices + offset;
+  IntT pool_indices_size = pool_sizes[index];
+
+  int kIteration = (pool_indices_size + warpSize - 1) / warpSize;
+  T mul_result = 0;
+  for (int k = 0; k < kIteration; ++k) {
+    int idx = tid + k * warpSize;
+    if (idx >= pool_indices_size) break;
+
+    auto i = pool_indices[idx];
+    auto cur_out_value = out_values + i * nvalues;
+    auto j = lower_bound_values[i];
+    if (j < grad_nnz && (out_offsets[i] == grad_offsets[j])) {
+      auto cur_grad_value = grad_values + j * nvalues;
+      mul_result += (*(cur_out_value + nval)) * (*(cur_grad_value + nval));
+    }
+  }
+  T sum = funcs::WarpReduceSum<T>(mul_result, 0xFFFFFFFF);
+
+  for (int k = 0; k < kIteration; ++k) {
+    int idx = tid + k * warpSize;
+    if (idx >= pool_indices_size) break;
+
+    auto i = pool_indices[idx];
+    auto j = lower_bound_values[i];
+    auto cur_out_value = out_values + i * nvalues;
+    auto cur_value = values + i * nvalues;
+    auto cur_grad_value = grad_values + j * nvalues;
+    if (j < grad_nnz && (out_offsets[i] == grad_offsets[j])) {
+      cur_value[nval] =
+          (*(cur_out_value + nval)) * (*(cur_grad_value + nval) - sum);
+    } else {
+      cur_value[nval] = -(*(cur_out_value + nval)) * sum;
+    }
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void SoftmaxCooGradGPUKernel(const Context& dev_ctx,
+                             const SparseCooTensor& out,
+                             const SparseCooTensor& dout,
+                             int axis,
+                             SparseCooTensor* dx) {
+  using thrust_ptr = thrust::device_ptr<IntT>;
+  auto out_indices = out.indices();
+  auto out_values = out.values();
+  auto out_values_ptr = out_values.data<T>();
+  const auto output_indices_dims = out.indices().dims();
+  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_values_ptr = grad_values.data<T>();
+  auto out_nnz = out.nnz();
+  auto grad_nnz = dout.nnz();
+  auto place = dev_ctx.GetPlace();
+  auto stream = dev_ctx.stream();
+
+  *(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);
+  funcs::SetConstant<GPUContext, T> set_zero;
+  set_zero(dev_ctx, values, static_cast<T>(0.0f));
+
+  DenseTensor out_offsets = funcs::sparse::GetOffsets<IntT, Context>(
+      dev_ctx, out_indices, sizes, static_cast<IntT>(-1));
+  auto out_offsets_ptr = out_offsets.data<IntT>();
+  DenseTensor grad_offsets = funcs::sparse::GetOffsets<IntT, Context>(
+      dev_ctx, grad_indices, sizes, static_cast<IntT>(-1));
+  auto grad_offsets_ptr = grad_offsets.data<IntT>();
+
+#ifdef PADDLE_WITH_HIP
+  const auto& policy = thrust::hip::par.on(dev_ctx.stream());
+  bool is_same_offset = thrust::equal(thrust::hip::par.on(dev_ctx.stream()),
+#else
+  const auto& policy = thrust::cuda::par.on(dev_ctx.stream());
+  bool is_same_offset = thrust::equal(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                                      out_offsets_ptr,
+                                      out_offsets_ptr + out_offsets.numel(),
+                                      grad_offsets_ptr);
+
+  int dim = axis < 0 ? out_dims.size() + axis : axis;
+  if (dim >= sparse_dim) {
+    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 sorted_indices;
+  DenseTensor pool_offsets;
+  DenseTensor pool_sizes;
+
+  std::tie(sorted_indices, pool_offsets, pool_sizes, std::ignore) =
+      funcs::sparse::ComputePoolMax<T, IntT, Context, false>(
+          dev_ctx, out_indices, values_2, sizes, nvalues, dim);
+
+  DenseTensor bound =
+      Empty<IntT>(dev_ctx, {static_cast<IntT>(out_offsets.dims()[0])});
+  IntT* bound_ptr = bound.data<IntT>();
+  thrust::lower_bound(policy,
+                      thrust_ptr(grad_offsets_ptr),
+                      thrust_ptr(grad_offsets_ptr + grad_offsets.dims()[0]),
+                      thrust_ptr(out_offsets_ptr),
+                      thrust_ptr(out_offsets_ptr) + out_offsets.dims()[0],
+                      thrust_ptr(bound.data<IntT>()));
+
+  auto pool_size = pool_offsets.dims()[0];
+  int total_rows = pool_size * nvalues;
+  dim3 grid((total_rows + 15) / 16);
+  dim3 block(32, 16);
+  SoftmaxCooGradGPURawKernel<T, IntT>
+      <<<grid, block, 0, stream>>>(sorted_indices.data<IntT>(),
+                                   pool_size,
+                                   pool_sizes.data<IntT>(),
+                                   pool_offsets.data<IntT>(),
+                                   nvalues,
+                                   grad_nnz,
+                                   grad_offsets.data<IntT>(),
+                                   out_offsets.data<IntT>(),
+                                   bound_ptr,
+                                   values_2.data<T>(),
+                                   out_values.data<T>(),
+                                   grad_values.data<T>(),
+                                   total_rows);
+}
+
+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(), "SoftmaxCooGradGPUKernel", ([&] {
+        SoftmaxCooGradGPUKernel<T, data_t, Context>(
+            dev_ctx, out, dout, axis, dx);
+      }));
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(softmax_csr_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SoftmaxCsrGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(softmax_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SoftmaxCooGradKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/gpu/softmax_kernel.cu

@@ -0,0 +1,276 @@
+/* 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 <thrust/device_ptr.h>
+#include <thrust/iterator/constant_iterator.h>
+#include <thrust/iterator/discard_iterator.h>
+#include <thrust/sequence.h>
+#include <thrust/sort.h>
+#include <thrust/transform.h>
+
+#include "paddle/phi/backends/gpu/gpu_context.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/activation_functor.h"
+#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/reduce_functor.h"
+#include "paddle/phi/kernels/funcs/sparse/softmax.cu.h"
+#include "paddle/phi/kernels/gpu/reduce.h"
+#include "paddle/phi/kernels/softmax_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename IntT = int>
+__global__ void SoftmaxGpuKernel(const IntT* x_crows,
+                                 const T* x_values,
+                                 T* out_values,
+                                 int row_number,
+                                 int total_row_number) {
+  int row = blockIdx.x * blockDim.y + threadIdx.y;
+  int non_zero_idx = threadIdx.x;
+  if (row >= total_row_number) return;
+  int cur_batch = row / row_number;
+  int crow_idx = cur_batch * (row_number + 1) + (row % row_number);
+  int cur_batch_offset = 0;
+  for (int i = 1; i < cur_batch + 1; ++i) {
+    cur_batch_offset += x_crows[i * (row_number + 1) - 1];
+  }
+  int row_first = cur_batch_offset + static_cast<int>(x_crows[crow_idx]);
+  int row_nnz = static_cast<int>(x_crows[crow_idx + 1] - x_crows[crow_idx]);
+  if (row_nnz == 0) return;
+
+  int kIteration = (row_nnz + warpSize - 1) / warpSize;
+
+  T max_val = -std::numeric_limits<T>::infinity();
+  for (int i = 0; i < kIteration; ++i) {
+    int idx = non_zero_idx + i * warpSize;
+    if (idx >= row_nnz) break;
+
+    T val = x_values[row_first + idx];
+    if (val > max_val) {
+      max_val = val;
+    }
+  }
+  T row_max_val = funcs::WarpReduceMax<T>(max_val, 0xFFFFFFFF);
+
+  T exp_sum = 0;
+  for (int i = 0; i < kIteration; ++i) {
+    int idx = non_zero_idx + i * warpSize;
+    if (idx >= row_nnz) break;
+
+    auto functor = funcs::CudaExpFunctor<T>();
+    T exp = functor(x_values[row_first + idx] - row_max_val);
+    exp_sum += exp;
+    out_values[row_first + idx] = exp;
+  }
+  T row_exp_sum = funcs::WarpReduceSum<T>(exp_sum, 0xFFFFFFFF);
+
+  for (int i = 0; i < kIteration; ++i) {
+    int idx = non_zero_idx + i * warpSize;
+    if (idx >= row_nnz) break;
+
+    out_values[row_first + idx] = out_values[row_first + idx] / row_exp_sum;
+  }
+}
+
+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 total_row_number = 1;
+  int row_number = 1;
+  for (int i = 0; i < x_rank - 1; ++i) {
+    total_row_number *= x_dim[i];
+    if (i == x_rank - 2) {
+      row_number = x_dim[i];
+    }
+  }
+
+  dim3 grid((total_row_number + 3) / 4);
+  dim3 block(32, 4);
+
+  PD_VISIT_BASE_INTEGRAL_TYPES(x.crows().dtype(), "CsrSoftmaxKernel", ([&] {
+                                 SoftmaxGpuKernel<T, data_t>
+                                     <<<grid, block, 0, dev_ctx.stream()>>>(
+                                         x.crows().data<data_t>(),
+                                         x.values().data<T>(),
+                                         out->mutable_values()->data<T>(),
+                                         row_number,
+                                         total_row_number);
+                               }));
+}
+
+template <typename T, typename IntT>
+__global__ void SoftmaxCooGPURawKernel(IntT* sorted_pool_indices,
+                                       IntT* pool_sizes,
+                                       IntT* pool_offsets,
+                                       IntT nvalues,
+                                       T* input_values,
+                                       T* output_values,
+                                       int total_rows) {
+  int row = blockIdx.x * blockDim.y + threadIdx.y;
+  if (row >= total_rows) return;
+
+  int tid = threadIdx.x;
+  int index = row / nvalues;
+  int j = row % nvalues;
+  IntT offset = pool_offsets[index];
+  IntT* pool_indices = sorted_pool_indices + offset;
+  IntT pool_indices_size = pool_sizes[index];
+
+  int kIteration = (pool_indices_size + warpSize - 1) / warpSize;
+  T max_val = -std::numeric_limits<T>::infinity();
+  for (int k = 0; k < kIteration; ++k) {
+    int idx = tid + k * warpSize;
+    if (idx >= pool_indices_size) break;
+
+    auto i = pool_indices[idx];
+    auto cur_value = input_values + j + nvalues * i;
+    if (*cur_value > max_val) {
+      max_val = *cur_value;
+    }
+  }
+  T row_max_val = funcs::WarpReduceMax<T>(max_val, 0xFFFFFFFF);
+
+  T exp_sum = 0;
+  for (int k = 0; k < kIteration; ++k) {
+    int idx = tid + k * warpSize;
+    if (idx >= pool_indices_size) break;
+
+    auto i = pool_indices[idx];
+    auto cur_value = input_values + j + nvalues * i;
+    auto cur_out_value = output_values + i * nvalues + j;
+
+    auto functor = funcs::CudaExpFunctor<T>();
+    T exp = functor(*cur_value - row_max_val);
+    exp_sum += exp;
+    *cur_out_value = exp;
+  }
+  T row_exp_sum = funcs::WarpReduceSum<T>(exp_sum, 0xFFFFFFFF);
+  row_exp_sum = 1.0 / row_exp_sum;
+
+  for (int k = 0; k < kIteration; ++k) {
+    int idx = tid + k * warpSize;
+    if (idx >= pool_indices_size) break;
+    auto i = pool_indices[idx];
+    auto cur_out_value = output_values + i * nvalues + j;
+    *cur_out_value *= row_exp_sum;
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void SoftmaxCooGPUKernel(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 std::vector<IntT> sizes = vectorize<IntT>(x_dims);
+  const auto sparse_dim = x.sparse_dim();
+  const IntT x_nnz = x.nnz();
+  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;
+  }
+
+  auto stream = dev_ctx.stream();
+  IntT nvalues = std::accumulate(sizes.begin() + sparse_dim,
+                                 sizes.end(),
+                                 static_cast<IntT>(1),
+                                 std::multiplies<>());
+
+  auto values_2 = values.Resize({x_nnz, nvalues});
+
+  /* Compute independent pools of indices */
+  DenseTensor sorted_indices;
+  DenseTensor pool_offsets;
+  DenseTensor pool_sizes;
+  std::tie(sorted_indices, pool_offsets, pool_sizes, std::ignore) =
+      funcs::sparse::ComputePoolMax<T, IntT, Context, false>(
+          dev_ctx, indices, values_2, sizes, nvalues, static_cast<IntT>(dim));
+
+  auto pool_size = pool_offsets.dims()[0];
+  auto out_values_ptr = out_values.data<T>();
+  auto values_ptr = values.data<T>();
+  int total_rows = pool_size * nvalues;
+  dim3 grid((total_rows + 15) / 16);
+  dim3 block(32, 16);
+  SoftmaxCooGPURawKernel<T, IntT>
+      <<<grid, block, 0, stream>>>(sorted_indices.data<IntT>(),
+                                   pool_sizes.data<IntT>(),
+                                   pool_offsets.data<IntT>(),
+                                   nvalues,
+                                   values_ptr,
+                                   out_values_ptr,
+                                   total_rows);
+}
+
+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(), "SoftmaxCooGPUKernel", ([&] {
+        SoftmaxCooGPUKernel<T, data_t, Context>(dev_ctx, x, axis, out);
+      }));
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(softmax_csr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SoftmaxCsrKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(softmax_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SoftmaxCooKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}

paddle/phi/kernels/sparse/gpu/sparse_attention_kernel.cu

@@ -0,0 +1,932 @@
+// Copyright (c) 2024 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 <math.h>
+
+#include <limits>
+#include <string>
+#include <vector>
+#include "paddle/phi/common/memory_utils.h"
+#if defined(PADDLE_WITH_CUDA)
+#include "paddle/phi/backends/dynload/cusparse.h"
+#endif
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/utils/data_type.h"
+#include "paddle/utils/optional.h"
+
+namespace phi {
+#if defined(PADDLE_WITH_CUDA)
+template <typename T>
+__forceinline__ __device__ T CudaShuffleXorSync(unsigned mask,
+                                                T val,
+                                                int width = warpSize) {
+  return __shfl_xor_sync(mask, val, width);
+}
+
+template <typename T, int batch_size, int warp_size>
+__device__ __forceinline__ void WarpReduceSum(T* sum) {
+#pragma unroll
+  for (int offset = warp_size / 2; offset > 0; offset /= 2) {
+#pragma unroll
+    for (int i = 0; i < batch_size; ++i) {
+      T sum_val = CudaShuffleXorSync(0xFFFFFFFF, sum[i], offset);
+      sum[i] = sum[i] + sum_val;
+    }
+  }
+}
+
+template <typename T, int batch_size, int warp_size>
+__device__ __forceinline__ void WarpReduceMax(T* sum) {
+#pragma unroll
+  for (int offset = warp_size / 2; offset > 0; offset /= 2) {
+#pragma unroll
+    for (int i = 0; i < batch_size; ++i) {
+      T max_val = CudaShuffleXorSync(0xFFFFFFFF, sum[i], offset);
+      sum[i] = max(sum[i], max_val);
+    }
+  }
+}
+
+template <typename T, int BlockSize, int BlockNnzMax>
+__global__ void BlockSparseSoftmaxForward(T* softmax,
+                                          const T* src,
+                                          T scale,
+                                          const T* kp_mask,
+                                          const T* attn_mask,
+                                          const int* layout_rowptr,
+                                          const int* layout_colindex,
+                                          int num_rows) {
+  // current thread related info
+  const int WarpSize = 32;
+  const int cur_row = blockIdx.x * blockDim.y + threadIdx.y;
+  if (cur_row < num_rows) {
+    const int cur_block_row = cur_row / BlockSize;
+    const int cur_block_nnz =
+        layout_rowptr[cur_block_row + 1] - layout_rowptr[cur_block_row];
+
+    T srcdata[(BlockSize * BlockNnzMax + WarpSize - 1) / WarpSize] = {0};
+    T attndata[(BlockSize * BlockNnzMax + WarpSize - 1) / WarpSize] = {0};
+
+    // read tensor data, attn mask
+    const int iter = (cur_block_nnz + WarpSize - 1) / WarpSize;
+    const T* srcptr = src + layout_rowptr[cur_block_row];
+
+    const T* attnptr = (attn_mask == nullptr)
+                           ? nullptr
+                           : (attn_mask + cur_block_row * num_rows);
+    // the column start index in current row
+    const int* colindex = layout_colindex + layout_rowptr[cur_block_row];
+    for (int j = 0; j < iter; j++) {
+      int cur_block_col = j * WarpSize + threadIdx.x;
+      int cur_reg_index = j;
+      if (cur_block_col < cur_block_nnz) {
+        // read kp mask
+        T cur_kp_mask;
+        if ((kp_mask != nullptr) && std::abs(kp_mask[colindex[cur_block_col]]) <
+                                        std::numeric_limits<T>::epsilon()) {
+          cur_kp_mask = -std::numeric_limits<T>::infinity();
+        } else {
+          cur_kp_mask = 0;
+        }
+        // do mask operation
+        if ((attnptr != nullptr) && std::abs(attnptr[colindex[cur_block_col]]) <
+                                        std::numeric_limits<T>::epsilon()) {
+          srcdata[cur_reg_index] =
+              -std::numeric_limits<T>::infinity() * scale + cur_kp_mask;
+        } else {
+          srcdata[cur_reg_index] = scale * srcptr[cur_block_col] + cur_kp_mask;
+        }
+      } else {
+        srcdata[cur_reg_index] = -std::numeric_limits<T>::infinity();
+      }
+    }
+
+    // max value
+    T max_value = srcdata[0];
+    const int kIteration =
+        (cur_block_nnz * BlockSize + WarpSize - 1) / WarpSize;
+#pragma unroll
+    for (int it = 1; it < kIteration; ++it) {
+      max_value = (max_value > srcdata[it]) ? max_value : srcdata[it];
+    }
+    WarpReduceMax<T, 1, WarpSize>(&max_value);
+
+    // exp sum
+    T sum = 0;
+#pragma unroll
+    for (int it = 0; it < kIteration; ++it) {
+      srcdata[it] = std::exp(srcdata[it] - max_value);
+      sum += srcdata[it];
+    }
+    WarpReduceSum<T, 1, WarpSize>(&sum);
+
+    // compute softmax and write out
+    T* softmaxptr = softmax + layout_rowptr[cur_block_row];
+    for (int j = 0; j < iter; j++) {
+      int cur_block_col = j * WarpSize + threadIdx.x;
+      int cur_reg_index = j;
+      if (cur_block_col < cur_block_nnz) {
+        softmaxptr[cur_block_col] = srcdata[cur_reg_index] / sum;
+      }
+    }
+  }
+}
+
+template <typename T, int BlockSize, int BlockNnzMax>
+__global__ void BlockSparseSoftmaxBackward(T* dst,
+                                           const T* grad,
+                                           const T* src,
+                                           T scale,
+                                           const int* layout_rowptr,
+                                           const int* layout_colindex,
+                                           int num_rows) {
+  // current thread related info
+  const int WarpSize = 32;
+  const int cur_row = blockIdx.x * blockDim.y + threadIdx.y;
+  if (cur_row < num_rows) {
+    const int cur_block_row = cur_row / BlockSize;
+    const int cur_block_nnz =
+        layout_rowptr[cur_block_row + 1] - layout_rowptr[cur_block_row];
+
+    T srcdata[(BlockSize * BlockNnzMax + WarpSize - 1) / WarpSize];
+    T graddata[(BlockSize * BlockNnzMax + WarpSize - 1) / WarpSize];
+
+    // read tensor data, attn mask
+    const int iter = (cur_block_nnz + WarpSize - 1) / WarpSize;
+    const T* srcptr = src + layout_rowptr[cur_block_row];
+    const T* gradptr = grad + layout_rowptr[cur_block_row];
+    for (int j = 0; j < iter; j++) {
+      int cur_block_col = j * WarpSize + threadIdx.x;
+      int cur_reg_index = j;
+      if (cur_block_col < cur_block_nnz) {
+        srcdata[cur_reg_index] = srcptr[cur_block_col];
+        graddata[cur_reg_index] = gradptr[cur_block_col];
+      } else {
+        srcdata[cur_reg_index] = 0;
+        graddata[cur_reg_index] = 0;
+      }
+    }
+
+    T sum = 0;
+    const int kIteration =
+        (cur_block_nnz * BlockSize + WarpSize - 1) / WarpSize;
+#pragma unroll
+    for (int it = 0; it < kIteration; ++it) {
+      sum += srcdata[it] * graddata[it];
+    }
+    WarpReduceSum<T, 1, WarpSize>(&sum);
+
+    // compute softmax and write out
+    T* dstptr = dst + layout_rowptr[cur_block_row];
+    for (int j = 0; j < iter; j++) {
+      int cur_block_col = j * WarpSize + threadIdx.x;
+      int cur_reg_index = j;
+      if (cur_block_col < cur_block_nnz) {
+        dstptr[cur_block_col] =
+            scale * srcdata[cur_reg_index] * (graddata[cur_reg_index] - sum);
+      }
+    }
+  }
+}
+
+/*
+input: sparse C in CSR format (num_rows,num_rows)
+output: sparse C after softmax operation
+*/
+template <typename DeviceContext, typename T>
+void SparseSoftmaxForward(const GPUContext& dev_ctx,
+                          const DenseTensor* offset,
+                          const DenseTensor* columns,
+                          DenseTensor* input,
+                          DenseTensor* output,
+                          const int blocksize,
+                          const int num_rows,
+                          const int num_cols,
+                          const DenseTensor* key_padding_mask,
+                          const DenseTensor* attn_mask) {
+  const int* offset_data = offset->data<int>();
+  const int* columns_data = columns->data<int>();
+  T* input_data = input->data<T>();
+  T* output_data = output->data<T>();
+  // Add mask
+  const T* key_padding_mask_data =
+      (key_padding_mask != nullptr) ? key_padding_mask->data<T>() : nullptr;
+  const T* attn_mask_data =
+      (attn_mask != nullptr) ? attn_mask->data<T>() : nullptr;
+
+  const int block_size = 1;
+  dim3 blocks(32, 4, 1);
+  int grid = (num_rows * block_size + 3) / 4;
+  T scaling = static_cast<T>(1.0) / sqrt(static_cast<T>(num_cols));
+
+  if (num_cols <= 4) {
+    BlockSparseSoftmaxForward<T, block_size, 4>
+        <<<grid, blocks>>>(output_data,
+                           input_data,
+                           scaling,
+                           key_padding_mask_data,
+                           attn_mask_data,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else if (num_cols > 4 && num_cols <= 8) {
+    BlockSparseSoftmaxForward<T, block_size, 8>
+        <<<grid, blocks>>>(output_data,
+                           input_data,
+                           scaling,
+                           key_padding_mask_data,
+                           attn_mask_data,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else if (num_cols > 8 && num_cols <= 16) {
+    BlockSparseSoftmaxForward<T, block_size, 16>
+        <<<grid, blocks>>>(output_data,
+                           input_data,
+                           scaling,
+                           key_padding_mask_data,
+                           attn_mask_data,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else if (num_cols > 16 && num_cols <= 32) {
+    BlockSparseSoftmaxForward<T, block_size, 32>
+        <<<grid, blocks>>>(output_data,
+                           input_data,
+                           scaling,
+                           key_padding_mask_data,
+                           attn_mask_data,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else if (num_cols > 32 && num_cols <= 64) {
+    BlockSparseSoftmaxForward<T, block_size, 64>
+        <<<grid, blocks>>>(output_data,
+                           input_data,
+                           scaling,
+                           key_padding_mask_data,
+                           attn_mask_data,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else if (num_cols > 64 && num_cols <= 128) {
+    BlockSparseSoftmaxForward<T, block_size, 128>
+        <<<grid, blocks>>>(output_data,
+                           input_data,
+                           scaling,
+                           key_padding_mask_data,
+                           attn_mask_data,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else if (num_cols > 128 && num_cols <= 256) {
+    BlockSparseSoftmaxForward<T, block_size, 256>
+        <<<grid, blocks>>>(output_data,
+                           input_data,
+                           scaling,
+                           key_padding_mask_data,
+                           attn_mask_data,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else if (num_cols > 256 && num_cols <= 512) {
+    BlockSparseSoftmaxForward<T, block_size, 512>
+        <<<grid, blocks>>>(output_data,
+                           input_data,
+                           scaling,
+                           key_padding_mask_data,
+                           attn_mask_data,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The head_dim of query in sparse_attention op should less or equal "
+        "512"));
+  }
+}
+
+template <typename DeviceContext, typename T>
+void SparseSoftmaxBackward(const GPUContext& dev_ctx,
+                           const DenseTensor* offset,
+                           const DenseTensor* columns,
+                           DenseTensor* dx,
+                           const DenseTensor* dout,
+                           const DenseTensor* out,
+                           const int blocksize,
+                           const int num_rows,
+                           const int num_cols) {
+  const int* offset_data = offset->data<int>();
+  const int* columns_data = columns->data<int>();
+  T* dx_data = dx->data<T>();
+  const T* dout_data = dout->data<T>();
+  const T* out_data = out->data<T>();
+
+  const int block_size = 1;
+  dim3 blocks(32, 4, 1);
+  int grid = (num_rows * block_size + 3) / 4;
+  T scaling = static_cast<T>(1.0) / sqrt(static_cast<T>(num_cols));
+
+  if (num_cols <= 4) {
+    BlockSparseSoftmaxBackward<T, block_size, 4><<<grid, blocks>>>(dx_data,
+                                                                   dout_data,
+                                                                   out_data,
+                                                                   scaling,
+                                                                   offset_data,
+                                                                   columns_data,
+                                                                   num_rows);
+  } else if (num_cols > 4 && num_cols <= 8) {
+    BlockSparseSoftmaxBackward<T, block_size, 8><<<grid, blocks>>>(dx_data,
+                                                                   dout_data,
+                                                                   out_data,
+                                                                   scaling,
+                                                                   offset_data,
+                                                                   columns_data,
+                                                                   num_rows);
+  } else if (num_cols > 8 && num_cols <= 16) {
+    BlockSparseSoftmaxBackward<T, block_size, 16>
+        <<<grid, blocks>>>(dx_data,
+                           dout_data,
+                           out_data,
+                           scaling,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else if (num_cols > 16 && num_cols <= 32) {
+    BlockSparseSoftmaxBackward<T, block_size, 32>
+        <<<grid, blocks>>>(dx_data,
+                           dout_data,
+                           out_data,
+                           scaling,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else if (num_cols > 32 && num_cols <= 64) {
+    BlockSparseSoftmaxBackward<T, block_size, 64>
+        <<<grid, blocks>>>(dx_data,
+                           dout_data,
+                           out_data,
+                           scaling,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else if (num_cols > 64 && num_cols <= 128) {
+    BlockSparseSoftmaxBackward<T, block_size, 128>
+        <<<grid, blocks>>>(dx_data,
+                           dout_data,
+                           out_data,
+                           scaling,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else if (num_cols > 128 && num_cols <= 256) {
+    BlockSparseSoftmaxBackward<T, block_size, 256>
+        <<<grid, blocks>>>(dx_data,
+                           dout_data,
+                           out_data,
+                           scaling,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else if (num_cols > 256 && num_cols <= 512) {
+    BlockSparseSoftmaxBackward<T, block_size, 512>
+        <<<grid, blocks>>>(dx_data,
+                           dout_data,
+                           out_data,
+                           scaling,
+                           offset_data,
+                           columns_data,
+                           num_rows);
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The head_dim of query in sparse_attention op should less or equal "
+        "512"));
+  }
+}
+
+inline cudaDataType_t GetGpuType(const DataType data_type) {
+  if (data_type == DataType::FLOAT32) {
+    return CUDA_R_32F;
+  } else if (data_type == DataType::FLOAT64) {
+    return CUDA_R_64F;
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Not support tensor type in sparse_attention OP: %s",
+        phi::DataTypeToString(data_type)));
+  }
+}
+
+inline cusparseOperation_t GetTransposeOperation(const bool transpose) {
+  if (transpose) {
+    return CUSPARSE_OPERATION_TRANSPOSE;
+  } else {
+    return CUSPARSE_OPERATION_NON_TRANSPOSE;
+  }
+}
+
+void CusparseDestroy(cusparseDnMatDescr_t* dn_mat_first,
+                     cusparseDnMatDescr_t* dn_mat_second,
+                     cusparseSpMatDescr_t* sp_mat) {
+  phi::dynload::cusparseDestroyDnMat(*dn_mat_first);
+  phi::dynload::cusparseDestroyDnMat(*dn_mat_second);
+  phi::dynload::cusparseDestroySpMat(*sp_mat);
+}
+#endif
+
+/*
+input: dense A (num_rows,num_cols), dense B (num_rows,num_cols)
+output: sparse C in CSR format (num_rows,num_rows)
+*/
+template <typename DeviceContext, typename T>
+void DotSdd(const GPUContext& dev_ctx,
+            const DenseTensor* a,
+            const DenseTensor* b,
+            const DenseTensor* c_offset,
+            const DenseTensor* c_columns,
+            DenseTensor* c_value,
+            const int num_rows,
+            const int num_cols,
+            const bool a_transpose,
+            const bool b_transpose) {
+#if defined(PADDLE_WITH_CUDA)
+  const T* a_data = a->data<T>();
+  const T* b_data = b->data<T>();
+  const int* c_offset_data = c_offset->data<int>();
+  const int* c_columns_data = c_columns->data<int>();
+  T* c_value_data = c_value->data<T>();
+
+  cudaDataType_t gpu_type = GetGpuType(c_value->dtype());
+  cusparseHandle_t handle = nullptr;
+  cusparseDnMatDescr_t mat_a, mat_b;
+  cusparseSpMatDescr_t mat_c;
+  phi::dynload::cusparseCreate(&handle);
+
+  // Create dense matrix A
+  phi::dynload::cusparseCreateDnMat(&mat_a,
+                                    num_rows,
+                                    num_cols,
+                                    num_cols,
+                                    const_cast<T*>(a_data),
+                                    gpu_type,
+                                    CUSPARSE_ORDER_ROW);
+  // Create dense matrix B
+  phi::dynload::cusparseCreateDnMat(&mat_b,
+                                    num_rows,
+                                    num_cols,
+                                    num_cols,
+                                    const_cast<T*>(b_data),
+                                    gpu_type,
+                                    CUSPARSE_ORDER_ROW);
+  // Create sparse matrix C in CSR format
+  int64_t c_nnz = c_columns->numel();
+  phi::dynload::cusparseCreateCsr(&mat_c,
+                                  num_rows,
+                                  num_rows,
+                                  c_nnz,
+                                  const_cast<int*>(c_offset_data),
+                                  const_cast<int*>(c_columns_data),
+                                  c_value_data,
+                                  CUSPARSE_INDEX_32I,
+                                  CUSPARSE_INDEX_32I,
+                                  CUSPARSE_INDEX_BASE_ZERO,
+                                  gpu_type);
+
+  T alpha = 1;
+  T beta = 0;
+
+  size_t buffer_size = 0;
+  phi::dynload::cusparseSDDMM_bufferSize(handle,
+                                         GetTransposeOperation(a_transpose),
+                                         GetTransposeOperation(b_transpose),
+                                         &alpha,
+                                         mat_a,
+                                         mat_b,
+                                         &beta,
+                                         mat_c,
+                                         gpu_type,
+                                         CUSPARSE_SDDMM_ALG_DEFAULT,
+                                         &buffer_size);
+  auto d_buffer_ptr = phi::memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      buffer_size,
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  void* d_buffer = static_cast<void*>(d_buffer_ptr->ptr());
+
+  phi::dynload::cusparseSDDMM(handle,
+                              GetTransposeOperation(a_transpose),
+                              GetTransposeOperation(b_transpose),
+                              &alpha,
+                              mat_a,
+                              mat_b,
+                              &beta,
+                              mat_c,
+                              gpu_type,
+                              CUSPARSE_SDDMM_ALG_DEFAULT,
+                              d_buffer);
+
+  CusparseDestroy(&mat_a, &mat_b, &mat_c);
+  phi::dynload::cusparseDestroy(handle);
+#else
+  PADDLE_THROW(common::errors::Unimplemented(
+      "DotSdd use cusparseSDDMM, which is supported "
+      "from CUDA 11.3"));
+#endif
+}
+
+/*
+input: sparse A in CSR format (num_rows,num_rows), dense B (num_rows,num_cols)
+output: dense C (num_rows,num_cols)
+*/
+template <typename DeviceContext, typename T>
+void DotDsd(const GPUContext& dev_ctx,
+            const DenseTensor* a_offset,
+            const DenseTensor* a_columns,
+            const DenseTensor* a_value,
+            const DenseTensor* b,
+            DenseTensor* c,
+            const int num_rows,
+            const int num_cols,
+            const bool a_transpose,
+            const bool b_transpose) {
+#if defined(PADDLE_WITH_CUDA)
+  const int* a_offset_data = a_offset->data<int>();
+  const int* a_columns_data = a_columns->data<int>();
+  const T* a_value_data = a_value->data<T>();
+  const T* b_data = b->data<T>();
+  T* c_data = c->data<T>();
+
+  cudaDataType_t gpu_type = GetGpuType(c->dtype());
+  cusparseHandle_t handle = nullptr;
+  cusparseSpMatDescr_t mat_a;
+  cusparseDnMatDescr_t mat_b, mat_c;
+  phi::dynload::cusparseCreate(&handle);
+
+  // Create sparse matrix A in CSR format
+  int64_t a_nnz = a_columns->numel();
+  phi::dynload::cusparseCreateCsr(&mat_a,
+                                  num_rows,
+                                  num_rows,
+                                  a_nnz,
+                                  const_cast<int*>(a_offset_data),
+                                  const_cast<int*>(a_columns_data),
+                                  const_cast<T*>(a_value_data),
+                                  CUSPARSE_INDEX_32I,
+                                  CUSPARSE_INDEX_32I,
+                                  CUSPARSE_INDEX_BASE_ZERO,
+                                  gpu_type);
+
+  // Create dense matrix B
+  phi::dynload::cusparseCreateDnMat(&mat_b,
+                                    num_rows,
+                                    num_cols,
+                                    num_cols,
+                                    const_cast<T*>(b_data),
+                                    gpu_type,
+                                    CUSPARSE_ORDER_ROW);
+  // Create dense matrix C
+  phi::dynload::cusparseCreateDnMat(&mat_c,
+                                    num_rows,
+                                    num_cols,
+                                    num_cols,
+                                    c_data,
+                                    gpu_type,
+                                    CUSPARSE_ORDER_ROW);
+
+  T alpha = 1;
+  T beta = 0;
+
+  size_t buffer_size = 0;
+  // allocate an external buffer if needed
+  phi::dynload::cusparseSpMM_bufferSize(handle,
+                                        GetTransposeOperation(a_transpose),
+                                        GetTransposeOperation(b_transpose),
+                                        &alpha,
+                                        mat_a,
+                                        mat_b,
+                                        &beta,
+                                        mat_c,
+                                        gpu_type,
+                                        CUSPARSE_SPMM_ALG_DEFAULT,
+                                        &buffer_size);
+  auto d_buffer_ptr = phi::memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      buffer_size,
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  void* d_buffer = static_cast<void*>(d_buffer_ptr->ptr());
+
+  phi::dynload::cusparseSpMM(handle,
+                             GetTransposeOperation(a_transpose),
+                             GetTransposeOperation(b_transpose),
+                             &alpha,
+                             mat_a,
+                             mat_b,
+                             &beta,
+                             mat_c,
+                             gpu_type,
+                             CUSPARSE_SPMM_ALG_DEFAULT,
+                             d_buffer);
+
+  CusparseDestroy(&mat_b, &mat_c, &mat_a);
+  phi::dynload::cusparseDestroy(handle);
+#else
+  PADDLE_THROW(common::errors::Unimplemented(
+      "DotDsd use cusparseSpMM, which is supported "
+      "from CUDA 11.0"));
+#endif
+}
+
+std::vector<DenseTensor> GetSplitTensor(DenseTensor* input) {
+  auto dims = input->dims();
+  int batch_size = dims[0];
+  int num_heads = dims[1];
+  std::vector<int> new_dims(dims.size() - 1);
+  new_dims[0] = batch_size * num_heads;
+  for (int i = 1; i < new_dims.size(); i++) {
+    new_dims[i] = dims[i + 1];
+  }
+  input->Resize(new_dims);
+  return input->Split(1, 0);
+}
+
+template <typename T, typename Context>
+void SparseAttentionCUDAKernel(const Context& dev_ctx,
+                               const DenseTensor& q,
+                               const DenseTensor& k,
+                               const DenseTensor& v,
+                               const DenseTensor& offset,
+                               const DenseTensor& columns,
+                               const optional<DenseTensor>& key_padding_mask,
+                               const optional<DenseTensor>& attn_mask,
+                               DenseTensor* out,
+                               DenseTensor* sparse_dot_sdd,
+                               DenseTensor* softmax) {
+#if defined(PADDLE_WITH_CUDA)
+  auto query = q;
+  auto key = k;
+  auto value = v;
+  auto output_ptr = out;
+  dev_ctx.template Alloc<T>(out);
+  auto sparse_dot_sdd_ptr = sparse_dot_sdd;
+  dev_ctx.template Alloc<T>(sparse_dot_sdd);
+  auto softmax_ptr = softmax;
+  dev_ctx.template Alloc<T>(softmax);
+
+  auto output = *output_ptr;
+  auto result_sdd = *sparse_dot_sdd_ptr;
+  auto result_softmax = *softmax_ptr;
+
+  auto query_dims = query.dims();
+  int batch_size = query_dims[0];
+  int num_heads = query_dims[1];
+  int M = query_dims[2];
+  int N = query_dims[3];
+
+  DenseTensor q2 = q;
+  DenseTensor k2 = k;
+  DenseTensor v2 = v;
+  DenseTensor offset2 = offset;
+  DenseTensor columns2 = columns;
+  std::vector<DenseTensor> query_lists = GetSplitTensor(&q2);
+  std::vector<DenseTensor> key_lists = GetSplitTensor(&k2);
+  std::vector<DenseTensor> value_lists = GetSplitTensor(&v2);
+  std::vector<DenseTensor> offset_lists = GetSplitTensor(&offset2);
+  std::vector<DenseTensor> columns_lists = GetSplitTensor(&columns2);
+  std::vector<DenseTensor> result_sdd_lists = GetSplitTensor(&result_sdd);
+  std::vector<DenseTensor> result_softmax_lists =
+      GetSplitTensor(&result_softmax);
+  std::vector<DenseTensor> output_lists = GetSplitTensor(&output);
+
+  const int iter_num = batch_size * num_heads;
+  for (int i = 0; i < iter_num; i++) {
+    DotSdd<Context, T>(dev_ctx,
+                       &query_lists[i],
+                       &key_lists[i],
+                       &offset_lists[i],
+                       &columns_lists[i],
+                       &result_sdd_lists[i],
+                       M,
+                       N,
+                       false,
+                       true);
+
+    if (key_padding_mask && attn_mask) {
+      SparseSoftmaxForward<Context, T>(
+          dev_ctx,
+          &offset_lists[i],
+          &columns_lists[i],
+          &result_sdd_lists[i],
+          &result_softmax_lists[i],
+          1,
+          M,
+          N,
+          key_padding_mask.get_ptr() + (i / num_heads) * M,
+          attn_mask.get_ptr());
+    } else if (key_padding_mask && !attn_mask.is_initialized()) {
+      SparseSoftmaxForward<Context, T>(
+          dev_ctx,
+          &offset_lists[i],
+          &columns_lists[i],
+          &result_sdd_lists[i],
+          &result_softmax_lists[i],
+          1,
+          M,
+          N,
+          key_padding_mask.get_ptr() + (i / num_heads) * M,
+          nullptr);
+    } else if (!key_padding_mask.is_initialized() && attn_mask) {
+      SparseSoftmaxForward<Context, T>(dev_ctx,
+                                       &offset_lists[i],
+                                       &columns_lists[i],
+                                       &result_sdd_lists[i],
+                                       &result_softmax_lists[i],
+                                       1,
+                                       M,
+                                       N,
+                                       nullptr,
+                                       attn_mask.get_ptr());
+    } else {
+      SparseSoftmaxForward<Context, T>(dev_ctx,
+                                       &offset_lists[i],
+                                       &columns_lists[i],
+                                       &result_sdd_lists[i],
+                                       &result_softmax_lists[i],
+                                       1,
+                                       M,
+                                       N,
+                                       nullptr,
+                                       nullptr);
+    }
+
+    DotDsd<Context, T>(dev_ctx,
+                       &offset_lists[i],
+                       &columns_lists[i],
+                       &result_softmax_lists[i],
+                       &value_lists[i],
+                       &output_lists[i],
+                       M,
+                       N,
+                       false,
+                       false);
+  }
+#endif
+}
+
+template <typename T, typename Context>
+void SparseAttentionGradCUDAKernel(const Context& dev_ctx,
+                                   const DenseTensor& q,
+                                   const DenseTensor& k,
+                                   const DenseTensor& v,
+                                   const DenseTensor& offset,
+                                   const DenseTensor& columns,
+                                   const DenseTensor& sparse_dot_sdd,
+                                   const DenseTensor& softmax,
+                                   const DenseTensor& out_grad,
+                                   DenseTensor* q_grad,
+                                   DenseTensor* k_grad,
+                                   DenseTensor* v_grad) {
+#if defined(PADDLE_WITH_CUDA)
+  auto query = q;
+  auto key = k;
+  auto value = v;
+
+  auto dout = out_grad;
+  auto* dquery_ptr = q_grad;
+  auto* dkey_ptr = k_grad;
+  auto* dvalue_ptr = v_grad;
+  dev_ctx.template Alloc<T>(q_grad);
+  dev_ctx.template Alloc<T>(k_grad);
+  dev_ctx.template Alloc<T>(v_grad);
+
+  auto dquery = *dquery_ptr;
+  auto dkey = *dkey_ptr;
+  auto dvalue = *dvalue_ptr;
+
+  auto query_dims = query.dims();
+  int batch_size = query_dims[0];
+  int num_heads = query_dims[1];
+  int M = query_dims[2];
+  int N = query_dims[3];
+
+  DenseTensor q2 = q;
+  DenseTensor k2 = k;
+  DenseTensor v2 = v;
+  DenseTensor offset2 = offset;
+  DenseTensor columns2 = columns;
+  DenseTensor sparse_dot_sdd2 = sparse_dot_sdd;
+  DenseTensor softmax2 = softmax;
+  DenseTensor dout2 = out_grad;
+  std::vector<DenseTensor> query_lists = GetSplitTensor(&q2);
+  std::vector<DenseTensor> key_lists = GetSplitTensor(&k2);
+  std::vector<DenseTensor> value_lists = GetSplitTensor(&v2);
+  std::vector<DenseTensor> offset_lists = GetSplitTensor(&offset2);
+  std::vector<DenseTensor> columns_lists = GetSplitTensor(&columns2);
+  std::vector<DenseTensor> sparse_dot_sdd_lists =
+      GetSplitTensor(&sparse_dot_sdd2);
+  std::vector<DenseTensor> softmax_lists = GetSplitTensor(&softmax2);
+  std::vector<DenseTensor> dout_lists = GetSplitTensor(&dout2);
+  std::vector<DenseTensor> dquery_lists = GetSplitTensor(&dquery);
+  std::vector<DenseTensor> dkey_lists = GetSplitTensor(&dkey);
+  std::vector<DenseTensor> dvalue_lists = GetSplitTensor(&dvalue);
+
+  const int iter_num = batch_size * num_heads;
+  for (int i = 0; i < iter_num; i++) {
+    // dValue = transpose(result_softmax) * dOut
+    DotDsd<Context, T>(dev_ctx,
+                       &offset_lists[i],
+                       &columns_lists[i],
+                       &softmax_lists[i],
+                       &dout_lists[i],
+                       &dvalue_lists[i],
+                       M,
+                       N,
+                       true,
+                       false);
+
+    // dSoftmax = dOut * transpose(Value)
+    int64_t nnz_num = columns_lists[i].numel();
+    DenseTensor dsoftmax;
+    dsoftmax.Resize({nnz_num});
+    dev_ctx.template Alloc<T>(&dsoftmax);
+    DotSdd<Context, T>(dev_ctx,
+                       &dout_lists[i],
+                       &value_lists[i],
+                       &offset_lists[i],
+                       &columns_lists[i],
+                       &dsoftmax,
+                       M,
+                       N,
+                       false,
+                       true);
+
+    // dSparseDotSdd = dSoftmax * softmax'(SparseDotSdd)
+    DenseTensor dsparse_dot_sdd;
+    dsparse_dot_sdd.Resize({nnz_num});
+    dev_ctx.template Alloc<T>(&dsparse_dot_sdd);
+    SparseSoftmaxBackward<Context, T>(dev_ctx,
+                                      &offset_lists[i],
+                                      &columns_lists[i],
+                                      &dsparse_dot_sdd,
+                                      &dsoftmax,
+                                      &softmax_lists[i],
+                                      1,
+                                      M,
+                                      N);
+
+    // dQuery = dSparseDotSdd * Key
+    DotDsd<Context, T>(dev_ctx,
+                       &offset_lists[i],
+                       &columns_lists[i],
+                       &dsparse_dot_sdd,
+                       &key_lists[i],
+                       &dquery_lists[i],
+                       M,
+                       N,
+                       false,
+                       false);
+
+    // dKey = transpose(dSparseDotSdd) * Query
+    DotDsd<Context, T>(dev_ctx,
+                       &offset_lists[i],
+                       &columns_lists[i],
+                       &dsparse_dot_sdd,
+                       &query_lists[i],
+                       &dkey_lists[i],
+                       M,
+                       N,
+                       true,
+                       false);
+  }
+#endif
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(sparse_attention,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::SparseAttentionCUDAKernel,
+                   float,
+                   double) {
+  kernel->InputAt(3).SetDataType(phi::DataType::INT32);
+  kernel->InputAt(4).SetDataType(phi::DataType::INT32);
+}
+PD_REGISTER_KERNEL(sparse_attention_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::SparseAttentionGradCUDAKernel,
+                   float,
+                   double) {
+  kernel->InputAt(3).SetDataType(phi::DataType::INT32);
+  kernel->InputAt(4).SetDataType(phi::DataType::INT32);
+}

paddle/phi/kernels/sparse/gpu/sparse_conv_hashmap.cuh

@@ -0,0 +1,428 @@
+// Copyright (c) 2025 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 <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <utility>
+#include <vector>
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/math_function_impl.h"
+#include "paddle/phi/kernels/funcs/transpose_function.cuh"
+
+/** Reserved value for indicating "empty". */
+#define EMPTY_CELL (0)
+/** CUDA naive thread block size. */
+#define BLOCK_SIZE (256)
+
+__inline__ __device__ int8_t atomicCAS(int8_t* address,
+                                       int8_t compare,
+                                       int8_t val) {
+  auto address_value = reinterpret_cast<uintptr_t>(address);
+  int32_t* base_address = reinterpret_cast<int32_t*>(
+      reinterpret_cast<char*>(address) - (address_value & 3));
+  int32_t int_val = static_cast<int32_t>(val) << ((address_value & 3) * 8);
+  int32_t int_comp = static_cast<int32_t>(compare) << ((address_value & 3) * 8);
+  return static_cast<int8_t>(atomicCAS(base_address, int_comp, int_val));
+}
+
+// TODO(ShigureNyako): can we do this more efficiently?
+__inline__ __device__ int16_t atomicCAS(int16_t* address,
+                                        int16_t compare,
+                                        int16_t val) {
+  auto address_value = reinterpret_cast<uintptr_t>(address);
+  int32_t* base_address = reinterpret_cast<int32_t*>(
+      reinterpret_cast<char*>(address) - (address_value & 2));
+  int32_t int_val = static_cast<int32_t>(val) << ((address_value & 2) * 8);
+  int32_t int_comp = static_cast<int32_t>(compare) << ((address_value & 2) * 8);
+  return static_cast<int16_t>(atomicCAS(base_address, int_comp, int_val));
+}
+
+__inline__ __device__ int64_t atomicCAS(int64_t* address,
+                                        int64_t compare,
+                                        int64_t val) {
+  using AtomicCAS64Type = unsigned long long;  // NOLINT(runtime/int)
+  return static_cast<int64_t>(
+      atomicCAS(reinterpret_cast<AtomicCAS64Type*>(address),
+                static_cast<AtomicCAS64Type>(compare),
+                static_cast<AtomicCAS64Type>(val)));
+}
+
+namespace phi {
+namespace sparse {
+
+template <typename dtype = int>
+__device__ uint64_t hash_func_64b(dtype* data, int n = 4) {
+  uint64_t hash = UINT64_C(14695981039346656037);
+  for (int j = 0; j < n; j++) {
+    hash ^= static_cast<unsigned int>(data[j]);
+    hash *= UINT64_C(1099511628211);
+  }
+  // hash = (hash >> 60) ^ (hash & 0xFFFFFFFFFFFFFFF);
+  return hash;
+}
+
+template <typename key_type>
+__device__ int hash(key_type key, int _capacity) {
+  return static_cast<uint64_t>(key) % _capacity;
+}
+
+template <typename key_type, typename val_type>
+class GPUHashTable {
+ private:
+  // public:
+  bool free_pointers;
+  const int _capacity;
+  const int _divisor;
+  const int _width;
+  key_type* table_keys;
+  val_type* table_vals;
+  void insert_many_coords(const GPUContext& dev_ctx,
+                          const int* coords,
+                          const int n);
+  void lookup_many_coords(const GPUContext& dev_ctx,
+                          const int* coords,
+                          val_type* results,
+                          const int* kernel_sizes,
+                          const int* tensor_strides,
+                          const int n,
+                          const int kernel_volume);
+
+ public:
+  GPUHashTable(DenseTensor* table_keys,
+               DenseTensor* table_vals,
+               const int divisor,
+               const int width)
+      : _capacity(table_keys->dims()[0]),
+        free_pointers(false),
+        table_keys(table_keys->data<key_type>()),
+        table_vals(table_vals->data<val_type>()),
+        _divisor(divisor),
+        _width(width) {}
+  ~GPUHashTable() {
+    if (free_pointers) {
+      cudaFree(table_keys);
+      cudaFree(table_vals);
+    }
+  }
+  void insert_coords(const GPUContext& dev_ctx, const DenseTensor& coords);
+  void lookup_coords(const GPUContext& dev_ctx,
+                     const DenseTensor& coords,
+                     const int* kernel_sizes,
+                     const int* tensor_strides,
+                     int kernel_volume,
+                     DenseTensor* results);
+  int get_divisor() { return _divisor; }
+  int get_capacity() { return _capacity; }
+};
+
+using hashtable = GPUHashTable<int64_t, int>;
+using hashtable32 = GPUHashTable<int, int>;
+
+template <typename key_type = int64_t, typename val_type = int>
+__global__ void insert_coords_kernel(key_type* table_keys,
+                                     val_type* table_vals,
+                                     const int* coords,
+                                     int n,
+                                     int _capacity,
+                                     int _width) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < n) {
+    key_type key = (key_type)(hash_func_64b(coords + idx * _width, _width));
+    int value = idx + 1;
+    int slot = hash(key, _capacity);
+    while (true) {
+      key_type prev = atomicCAS(&table_keys[slot], EMPTY_CELL, key);
+      if (prev == EMPTY_CELL || prev == key) {
+        table_vals[slot] = value;
+        return;
+      }
+      slot = (slot + 1) % _capacity;
+    }
+  }
+}
+
+template <typename key_type = int64_t, typename val_type = int, bool odd>
+__global__ void lookup_coords_kernel(key_type* table_keys,
+                                     val_type* table_vals,
+                                     const int* coords,
+                                     val_type* vals,
+                                     const int* kernel_sizes,
+                                     const int* strides,
+                                     int n,
+                                     int _capacity,
+                                     int kernel_volume,
+                                     int _width) {
+  int tidx = blockIdx.x * blockDim.x + threadIdx.x;
+  int idx = tidx / kernel_volume;
+  if (idx >= n) return;
+  int _kernel_idx = tidx % kernel_volume;
+  int kernel_idx = _kernel_idx;
+  const int* in_coords = coords + _width * idx;
+  int coords_out[4];
+  // coords_out[2] = in_coords[2];
+  // coords_out[3] = in_coords[3];
+  coords_out[0] = in_coords[0];
+
+  if constexpr (odd) {
+#pragma unroll
+    for (int i = 0; i <= _width - 2; i++) {
+      int cur_offset = _kernel_idx % kernel_sizes[i];
+      cur_offset -= (kernel_sizes[i] - 1) / 2;
+      coords_out[i + 1] = in_coords[i + 1] * strides[i] + cur_offset;
+      _kernel_idx /= kernel_sizes[i];
+    }
+  } else {
+#pragma unroll
+    for (int i = _width - 2; i >= 0; i--) {
+      int cur_offset = _kernel_idx % kernel_sizes[i];
+      cur_offset -= (kernel_sizes[i] - 1) / 2;
+      coords_out[i + 1] = in_coords[i + 1] * strides[i] + cur_offset;
+      _kernel_idx /= kernel_sizes[i];
+    }
+  }
+
+  if (idx < n) {
+    key_type key = (key_type)(hash_func_64b(coords_out, _width));
+    int slot = hash(key, _capacity);
+
+    while (true) {
+      key_type cur_key = table_keys[slot];
+      if (key == cur_key) {
+        vals[idx * kernel_volume + kernel_idx] =
+            table_vals[slot] -
+            1;  // need to subtract 1 to avoid extra operations in python
+      }
+      if (table_keys[slot] == EMPTY_CELL) {
+        return;
+      }
+      slot = (slot + 1) % _capacity;
+    }
+  }
+}
+
+__global__ void set_kernel_sizes_and_strides_tensor(int* kernel_size_tensor,
+                                                    int* strides_tensor,
+                                                    int kernel_size0,
+                                                    int kernel_size1,
+                                                    int kernel_size2,
+                                                    int stride0,
+                                                    int stride1,
+                                                    int stride2) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < 6) {
+    switch (idx) {
+      case 0:
+        kernel_size_tensor[idx] = kernel_size0;
+        break;
+      case 1:
+        kernel_size_tensor[idx] = kernel_size1;
+        break;
+      case 2:
+        kernel_size_tensor[idx] = kernel_size2;
+        break;
+      case 3:
+        strides_tensor[idx - 3] = stride0;
+        break;
+      case 4:
+        strides_tensor[idx - 3] = stride1;
+        break;
+      case 5:
+        strides_tensor[idx - 3] = stride2;
+        break;
+      default:
+        break;
+    }
+  }
+}
+
+template <typename key_type, typename val_type>
+void GPUHashTable<key_type, val_type>::insert_many_coords(
+    const GPUContext& dev_ctx, const int* coords, const int n) {
+  insert_coords_kernel<key_type, val_type>
+      <<<(n + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, dev_ctx.stream()>>>(
+          table_keys, table_vals, coords, n, _capacity, _width);
+}
+
+template <typename key_type, typename val_type>
+void GPUHashTable<key_type, val_type>::insert_coords(
+    const GPUContext& dev_ctx, const DenseTensor& coords) {
+  insert_many_coords(dev_ctx, coords.data<int>(), coords.dims()[0]);
+}
+
+template <typename key_type, typename val_type>
+void GPUHashTable<key_type, val_type>::lookup_many_coords(
+    const GPUContext& dev_ctx,
+    const int* coords,
+    val_type* results,
+    const int* kernel_sizes,
+    const int* strides,
+    const int n,
+    const int kernel_volume) {
+  if (kernel_volume % 2)
+    lookup_coords_kernel<key_type, val_type, true>
+        <<<(n * kernel_volume + BLOCK_SIZE - 1) / BLOCK_SIZE,
+           BLOCK_SIZE,
+           0,
+           dev_ctx.stream()>>>(table_keys,
+                               table_vals,
+                               coords,
+                               results,
+                               kernel_sizes,
+                               strides,
+                               n,
+                               _capacity,
+                               kernel_volume,
+                               _width);
+  else
+    lookup_coords_kernel<key_type, val_type, false>
+        <<<(n * kernel_volume + BLOCK_SIZE - 1) / BLOCK_SIZE,
+           BLOCK_SIZE,
+           0,
+           dev_ctx.stream()>>>(table_keys,
+                               table_vals,
+                               coords,
+                               results,
+                               kernel_sizes,
+                               strides,
+                               n,
+                               _capacity,
+                               kernel_volume,
+                               _width);
+}
+
+template <typename key_type, typename val_type>
+void GPUHashTable<key_type, val_type>::lookup_coords(const GPUContext& dev_ctx,
+                                                     const DenseTensor& coords,
+                                                     const int* kernel_sizes,
+                                                     const int* strides,
+                                                     const int kernel_volume,
+                                                     DenseTensor* results) {
+  int32_t* results_data = results->data<int32_t>();
+  lookup_many_coords(dev_ctx,
+                     coords.data<int>(),
+                     results_data,
+                     kernel_sizes,
+                     strides,
+                     coords.dims()[0],
+                     kernel_volume);
+}
+
+template <typename IntT>
+void build_sparse_conv_kmap(const GPUContext& dev_ctx,
+                            const SparseCooTensor& x,
+                            const std::string& key,
+                            const std::vector<int>& kernel_sizes,
+                            const std::vector<int>& strides,
+                            const int kernel_volume,
+                            const bool is2D,
+                            SparseCooTensor* out) {
+  int nnz = x.nnz();
+  const KmapCache* in_kmap_cache_ptr = x.GetKmapCache(key);
+  out->ClearKmaps();
+  KmapCache* out_kmap_cache_ptr = nullptr;
+  bool to_insert = false;
+  if (in_kmap_cache_ptr == nullptr) {
+    KmapCache kmap_cache;
+    out_kmap_cache_ptr = out->SetKmapCache(key, kmap_cache);
+    if (out_kmap_cache_ptr->hashmap_keys == nullptr) {
+      DenseTensor* tmp_hashmap_keys = new DenseTensor();
+      tmp_hashmap_keys->Resize({2 * x.nnz()});
+      dev_ctx.template Alloc<IntT>(tmp_hashmap_keys);
+      funcs::SetConstant<GPUContext, IntT> set_zero;
+      set_zero(dev_ctx, tmp_hashmap_keys, static_cast<IntT>(0));
+      out_kmap_cache_ptr->hashmap_keys = tmp_hashmap_keys;
+      to_insert = true;
+    }
+    if (out_kmap_cache_ptr->hashmap_values == nullptr) {
+      DenseTensor* tmp_hashmap_values = new DenseTensor();
+      tmp_hashmap_values->Resize({2 * x.nnz()});
+      dev_ctx.template Alloc<int32_t>(tmp_hashmap_values);
+      funcs::SetConstant<GPUContext, int32_t> set_zero;
+      set_zero(dev_ctx, tmp_hashmap_values, static_cast<int32_t>(0));
+      out_kmap_cache_ptr->hashmap_values = tmp_hashmap_values;
+    }
+
+    if (out_kmap_cache_ptr->coords == nullptr) {
+      DenseTensor* tmp_indices = new DenseTensor();
+      tmp_indices->Resize({x.indices().dims()[1], x.indices().dims()[0]});
+      dev_ctx.template Alloc<int32_t>(tmp_indices);
+      // transpose indices
+      std::vector<int> perm = {1, 0};
+      funcs::TransposeGPUKernelDriver<int32_t>(
+          dev_ctx, x.indices(), perm, tmp_indices);
+      out_kmap_cache_ptr->coords = tmp_indices;
+    }
+
+    const int divisor = 128;
+    const int width = is2D ? 3 : 4;
+    auto hashmap =
+        GPUHashTable<IntT, int32_t>(out_kmap_cache_ptr->hashmap_keys,
+                                    out_kmap_cache_ptr->hashmap_values,
+                                    divisor,
+                                    width);
+    if (to_insert) {
+      hashmap.insert_coords(dev_ctx, *(out_kmap_cache_ptr->coords));
+    }
+
+    DenseTensor* tmp_out_in_map = new DenseTensor();
+    tmp_out_in_map->Resize(
+        {(x.nnz() + divisor - 1) / divisor * divisor, kernel_volume});
+    dev_ctx.template Alloc<int32_t>(tmp_out_in_map);
+    out_kmap_cache_ptr->out_in_map = tmp_out_in_map;
+    funcs::SetConstant<GPUContext, int32_t> set_neg_one;
+    set_neg_one(
+        dev_ctx, out_kmap_cache_ptr->out_in_map, static_cast<int32_t>(-1));
+
+    // need to put kernel_sizes and strides to GPU
+    auto kernel_sizes_tensor = phi::Empty<int32_t>(dev_ctx, {3});
+    auto strides_tensor = phi::Empty<int32_t>(dev_ctx, {3});
+    set_kernel_sizes_and_strides_tensor<<<1, 32, 0, dev_ctx.stream()>>>(
+        kernel_sizes_tensor.data<int32_t>(),
+        strides_tensor.data<int32_t>(),
+        kernel_sizes[0],
+        kernel_sizes[1],
+        kernel_sizes[2],
+        strides[0],
+        strides[1],
+        strides[2]);
+    hashmap.lookup_coords(dev_ctx,
+                          *(out_kmap_cache_ptr->coords),
+                          kernel_sizes_tensor.data<int32_t>(),
+                          strides_tensor.data<int32_t>(),
+                          kernel_volume,
+                          out_kmap_cache_ptr->out_in_map);
+
+  } else {
+    // out tensor takes the kmaps from x
+    out->SetKmaps(x.GetKmaps());
+    // force clear the kmaps of x
+    const_cast<SparseCooTensor&>(x).ClearKmaps();
+  }
+  const KmapCache* new_out_kmap_cache_ptr = out->GetKmapCache(key);
+  assert(new_out_kmap_cache_ptr != nullptr);
+  assert(new_out_kmap_cache_ptr->hashmap_keys != nullptr);
+  assert(new_out_kmap_cache_ptr->hashmap_values != nullptr);
+  assert(new_out_kmap_cache_ptr->coords != nullptr);
+  assert(new_out_kmap_cache_ptr->out_in_map != nullptr);
+  return;
+}
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/gpu/sparse_utils_kernel.cu

@@ -0,0 +1,846 @@
+/* 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/sparse_utils_kernel.h"
+
+#include <thrust/execution_policy.h>
+#include <thrust/remove.h>
+
+#ifdef PADDLE_WITH_HIP
+#include "paddle/phi/backends/dynload/rocsparse.h"
+#endif
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.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/visit_type.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/sparse/common_shape.h"
+#include "paddle/phi/kernels/sparse/gpu/conv_host_buffer.h"
+
+#define BUILD_CUDA_TENSOR(T, vector, tensor)                                \
+  if (vector.size() <= 4) {                                                 \
+    switch (vector.size()) {                                                \
+      case 1:                                                               \
+        build_cuda_tensor<<<1, 32, 0, dev_ctx.stream()>>>(tensor.data<T>(), \
+                                                          vector[0]);       \
+        break;                                                              \
+      case 2:                                                               \
+        build_cuda_tensor<<<1, 32, 0, dev_ctx.stream()>>>(                  \
+            tensor.data<T>(), vector[0], vector[1]);                        \
+        break;                                                              \
+      case 3:                                                               \
+        build_cuda_tensor<<<1, 32, 0, dev_ctx.stream()>>>(                  \
+            tensor.data<T>(), vector[0], vector[1], vector[2]);             \
+        break;                                                              \
+      case 4:                                                               \
+        build_cuda_tensor<<<1, 32, 0, dev_ctx.stream()>>>(                  \
+            tensor.data<T>(), vector[0], vector[1], vector[2], vector[3]);  \
+        break;                                                              \
+      default:                                                              \
+        break;                                                              \
+    }                                                                       \
+  } else {                                                                  \
+    backends::gpu::GpuMemcpyAsync(tensor.data<T>(),                         \
+                                  vector.data(),                            \
+                                  vector.size() * sizeof(T),                \
+                                  gpuMemcpyHostToDevice,                    \
+                                  dev_ctx.stream());                        \
+  }
+
+namespace phi {
+namespace sparse {
+
+template <typename T>
+__global__ void build_cuda_tensor(T* data, const T elem0) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < 1) {
+    data[idx] = elem0;
+  }
+}
+
+template <typename T>
+__global__ void build_cuda_tensor(T* data, const T elem0, const T elem1) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < 2) {
+    switch (idx) {
+      case 0:
+        data[idx] = elem0;
+        break;
+      case 1:
+        data[idx] = elem1;
+        break;
+      default:
+        break;
+    }
+  }
+}
+
+template <typename T>
+__global__ void build_cuda_tensor(T* data,
+                                  const T elem0,
+                                  const T elem1,
+                                  const T elem2) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < 3) {
+    switch (idx) {
+      case 0:
+        data[idx] = elem0;
+        break;
+      case 1:
+        data[idx] = elem1;
+        break;
+      case 2:
+        data[idx] = elem2;
+        break;
+      default:
+        break;
+    }
+  }
+}
+
+template <typename T>
+__global__ void build_cuda_tensor(
+    T* data, const T elem0, const T elem1, const T elem2, const T elem3) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < 4) {
+    switch (idx) {
+      case 0:
+        data[idx] = elem0;
+        break;
+      case 1:
+        data[idx] = elem1;
+        break;
+      case 2:
+        data[idx] = elem2;
+        break;
+      case 3:
+        data[idx] = elem3;
+        break;
+      default:
+        break;
+    }
+  }
+}
+
+template <typename T>
+inline __device__ bool DevIsZero(const T* data, const int64_t cols) {
+  const T zero = static_cast<T>(0);
+  // TODO(zhangkaihuo): check the data is zero or not in parallen when cols > 1
+  for (int64_t i = 0; i < cols; i++) {
+    if (data[i] != zero) {
+      return false;
+    }
+  }
+  return true;
+}
+
+template <typename T>
+__global__ void GetNonZeroNums(const T* dense_data,
+                               const int rows,
+                               const int cols,
+                               int* non_zero_num,
+                               int* temp_indices) {
+  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  __shared__ int counter;
+  if (threadIdx.x == 0) counter = 0;
+  __syncthreads();
+
+  for (int i = tid; i < rows; i += gridDim.x * blockDim.x) {
+    int index = -1;
+    // TODO(zhangkaihuo): when cols=1, vectorization can be used
+    if (!DevIsZero(dense_data + i * cols, cols)) {
+      // use reductions?
+      atomicAdd(&counter, 1);
+      index = i;
+    }
+    temp_indices[i] = index;
+  }
+  __syncthreads();
+  if (threadIdx.x == 0) {
+    atomicAdd(non_zero_num, counter);
+  }
+}
+
+template <typename T>
+__global__ void GetNonZeroElementsAndIndices(const T* dense_data,
+                                             const int64_t sparse_dim,
+                                             const int64_t cols,
+                                             const int64_t* x_dims,
+                                             const int non_zero_num,
+                                             const int* sparse_indices,
+                                             int64_t* indices,
+                                             T* sparse_data) {
+  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  for (int i = tid; i < non_zero_num; i += gridDim.x * blockDim.x) {
+    int64_t sparse_index = sparse_indices[i];
+    int64_t x_index = sparse_index;
+    for (int64_t j = sparse_dim - 1; j >= 0; j--) {
+      indices[j * non_zero_num + i] = sparse_index % x_dims[j];
+      sparse_index /= x_dims[j];
+    }
+
+    for (int j = 0; j < cols; j++) {
+      sparse_data[i * cols + j] = dense_data[x_index * cols + j];
+    }
+  }
+}
+
+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"));
+  auto dims_2d = flatten_to_2d(x_dims, sparse_dim);
+  const int rows = dims_2d[0];
+  const int cols = dims_2d[1];
+  DenseTensor nums = Empty<int32_t>(dev_ctx, {1});
+  DenseTensor d_x_dims = Empty<int64_t>(dev_ctx, {x_dims.size()});
+
+  // 1. get numbers of non zero elements, and get the index of non zero elements
+  int* nums_ptr = nums.data<int>();
+  backends::gpu::GpuMemsetAsync(nums_ptr, 0, sizeof(int), dev_ctx.stream());
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rows, 1);
+
+  DenseTensor temp_indices = Empty<int32_t>(dev_ctx, {rows});
+  int* temp_indices_ptr = temp_indices.data<int>();
+
+  GetNonZeroNums<<<config.block_per_grid.x,
+                   config.thread_per_block.x,
+                   0,
+                   dev_ctx.stream()>>>(
+      x_data, rows, cols, nums_ptr, temp_indices_ptr);
+
+#ifdef PADDLE_WITH_HIP
+  thrust::remove(thrust::hip::par.on(dev_ctx.stream()),
+#else
+  thrust::remove(thrust::cuda::par.on(dev_ctx.stream()),
+#endif
+                 temp_indices_ptr,
+                 temp_indices_ptr + rows,
+                 -1);
+
+  // 2. copy non_zero_num to host, copy x_dims to device
+  int non_zero_num = 0;
+  backends::gpu::GpuMemcpyAsync(&non_zero_num,
+                                nums_ptr,
+                                sizeof(int),
+                                gpuMemcpyDeviceToHost,
+                                dev_ctx.stream());
+  backends::gpu::GpuMemcpyAsync(d_x_dims.data<int64_t>(),
+                                x_dims.Get(),
+                                x_dims.size() * sizeof(x_dims[0]),
+                                gpuMemcpyHostToDevice,
+                                dev_ctx.stream());
+
+  dev_ctx.Wait();  // wait the copy
+
+  const auto values_dims =
+      funcs::sparse::InferDenseDims(x_dims, sparse_dim, non_zero_num);
+  DenseTensor indices =
+      Empty<int64_t>(dev_ctx, {sparse_dim, static_cast<int64_t>(non_zero_num)});
+  int64_t* indices_data = indices.data<int64_t>();
+  DenseTensor values;
+  values.Resize(values_dims);
+  T* sparse_data = dev_ctx.template Alloc<T>(&values);
+
+  // 3. calc indices by indices and get values by indices
+  if (non_zero_num > 0) {
+    config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num, 1);
+    GetNonZeroElementsAndIndices<<<config.block_per_grid.x,
+                                   config.thread_per_block.x,
+                                   0,
+                                   dev_ctx.stream()>>>(x_data,
+                                                       sparse_dim,
+                                                       cols,
+                                                       d_x_dims.data<int64_t>(),
+                                                       non_zero_num,
+                                                       temp_indices_ptr,
+                                                       indices_data,
+                                                       sparse_data);
+  }
+
+  out->SetMember(indices, values, x_dims, true);
+}
+
+template <typename IntT>
+__global__ void GetBatchSizes(const IntT* crows,
+                              const int rows,
+                              const int batches,
+                              IntT* batch_sizes) {
+  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  if (tid < batches) {
+    batch_sizes[tid] = crows[tid * (rows + 1) + rows];
+  }
+}
+
+template <typename IntT>
+__global__ void ConvertCsrCrowsToCooRows(const IntT* crows_ptr,
+                                         const IntT* crows_offsets,
+                                         IntT* rows_ptr,
+                                         IntT* batch_ptr,
+                                         const int rows) {
+  const int b = blockIdx.y;
+  const int64_t offset = crows_offsets ? crows_offsets[b] : 0;
+  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  for (int i = tid; i < rows; i += gridDim.x * blockDim.x) {
+    for (int j = crows_ptr[b * (rows + 1) + i];
+         j < crows_ptr[b * (rows + 1) + i + 1];
+         j++) {
+      rows_ptr[offset + j] = i;
+      if (batch_ptr) {
+        batch_ptr[offset + j] = b;
+      }
+    }
+  }
+}
+
+template <typename T, typename IntT>
+void CsrToCooGPUKernel(const GPUContext& 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;
+  }
+
+  if (x.nnz() <= 0) {
+#ifdef PADDLE_WITH_HIP
+    DenseTensor indices = Empty<int>(dev_ctx, {sparse_dim, non_zero_num});
+#else
+    DenseTensor indices = Empty<IntT>(dev_ctx, {sparse_dim, non_zero_num});
+#endif
+    DenseTensor values = EmptyLike<T, GPUContext>(dev_ctx, x.values());
+    out->SetMember(indices, values, x_dims, true);
+    return;
+  }
+
+// rocsparse_csr2coo only support index with type 'rocsparse_int' (aka 'int')
+// now
+#ifdef PADDLE_WITH_HIP
+  const auto& csr_crows = Cast<IntT>(dev_ctx, x.crows(), DataType::INT32);
+  const auto& csr_cols = Cast<IntT>(dev_ctx, x.cols(), DataType::INT32);
+  const int* csr_crows_data = csr_crows.template data<int>();
+  const int* csr_cols_data = csr_cols.template data<int>();
+#else
+  const auto& csr_crows = x.crows();
+  const auto& csr_cols = x.cols();
+  const IntT* csr_crows_data = csr_crows.data<IntT>();
+  const IntT* csr_cols_data = csr_cols.data<IntT>();
+#endif
+  const auto& csr_values = x.values();
+  const T* csr_values_data = csr_values.data<T>();
+
+  int batches = x_dims.size() == 2 ? 1 : x_dims[0];
+  int rows = x_dims.size() == 2 ? x_dims[0] : x_dims[1];
+
+#ifdef PADDLE_WITH_HIP
+  DenseTensor indices = Empty<int>(dev_ctx, {sparse_dim, non_zero_num});
+  int* coo_indices = indices.data<int>();
+  int* coo_rows_data = coo_indices;
+  int* coo_cols_data = coo_rows_data + non_zero_num;
+#else
+  DenseTensor indices = Empty<IntT>(dev_ctx, {sparse_dim, non_zero_num});
+  DenseTensor offsets = Empty<IntT>(dev_ctx, {batches});
+  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;
+  IntT* offsets_ptr = batches == 1 ? nullptr : offsets.data<IntT>();
+#endif
+  DenseTensor values = EmptyLike<T, GPUContext>(dev_ctx, csr_values);
+  T* coo_values_data = values.data<T>();
+
+  if (batches > 1) {
+#ifdef PADDLE_WITH_HIP
+    PADDLE_THROW(common::errors::Unimplemented(
+        "'rocsparse_csr2coo' only supports batches "
+        "with a value of 1 currently."));
+#else
+    auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, batches, 1);
+    GetBatchSizes<IntT><<<config.block_per_grid.x, config.thread_per_block.x>>>(
+        csr_crows_data, rows, batches, offsets_ptr);
+
+    thrust::exclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
+                           offsets_ptr,
+                           offsets_ptr + batches,
+                           offsets_ptr);
+#endif
+  }
+
+#ifdef PADDLE_WITH_HIP
+  dev_ctx.CusparseCall([&](rocsparse_handle handle) {
+    phi::dynload::rocsparse_csr2coo(handle,
+                                    csr_crows_data,
+                                    non_zero_num,
+                                    rows,
+                                    coo_rows_data,
+                                    rocsparse_index_base_zero);
+  });
+#else
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rows, 1);
+  config.block_per_grid.y = batches;
+  ConvertCsrCrowsToCooRows<IntT>
+      <<<config.block_per_grid, config.thread_per_block.x>>>(
+          csr_crows_data, offsets_ptr, coo_rows_data, batch_ptr, rows);
+#endif
+  backends::gpu::GpuMemcpyAsync(coo_cols_data,
+                                csr_cols_data,
+#ifdef PADDLE_WITH_HIP
+                                sizeof(int) * non_zero_num,
+#else
+                                sizeof(IntT) * non_zero_num,
+#endif
+                                gpuMemcpyDeviceToDevice,
+                                dev_ctx.stream());
+  backends::gpu::GpuMemcpyAsync(coo_values_data,
+                                csr_values_data,
+                                sizeof(T) * non_zero_num,
+                                gpuMemcpyDeviceToDevice,
+                                dev_ctx.stream());
+
+#ifdef PADDLE_WITH_HIP
+  if (std::is_same<IntT, int64_t>::value)
+    indices = Cast<int>(dev_ctx, indices, DataType::INT64);
+#endif
+
+  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(), "CsrToCooGPUKernel", ([&] {
+                                 CsrToCooGPUKernel<T, data_t>(dev_ctx, x, out);
+                               }));
+}
+
+template <typename IntT>
+__global__ void GetBatchesOffset(const IntT* batches_ptr,
+                                 const int batches,
+                                 const int non_zero_num,
+                                 int* batches_offset) {
+  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  for (int i = tid; i < non_zero_num; i += gridDim.x * blockDim.x) {
+    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++) {
+        batches_offset[j] = i + 1;
+      }
+    }
+  }
+}
+
+template <typename IntT>
+__global__ void ConvertCooRowsToCsrCrows(
+    const int* batches_offset,  // can be null if batches = 1
+    const IntT* coo_rows_data,
+    IntT* csr_crows_data,
+    const int rows,
+    const int64_t non_zero_num) {
+  const int b = blockIdx.y;
+  int batch_non_zero_num =
+      batches_offset == nullptr ? non_zero_num : batches_offset[b];
+  IntT batch_start = 0;
+  if (b > 0) {
+    batch_start = batches_offset[b - 1];
+    batch_non_zero_num -= batch_start;
+  }
+
+  const IntT* coo_rows_ptr = coo_rows_data + batch_start;
+  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  for (int i = tid; i < batch_non_zero_num; i += gridDim.x * blockDim.x) {
+    if (i == 0) {
+      for (IntT j = 0; j <= coo_rows_ptr[0]; j++) {
+        csr_crows_data[b * (rows + 1) + j] = 0;
+      }
+    } else {
+      for (IntT j = coo_rows_ptr[i - 1]; j < coo_rows_ptr[i]; j++) {
+        csr_crows_data[b * (rows + 1) + j + 1] = i;
+      }
+    }
+    if (i == batch_non_zero_num - 1) {
+      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) {
+    for (int i = tid; i < rows + 1; i += gridDim.x * blockDim.x) {
+      csr_crows_data[b * (rows + 1) + i] = 0;
+    }
+  }
+}
+
+template <typename T, typename IntT>
+void CooToCsrGPUKernel(const GPUContext& 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 = x_dims.size() == 2 ? 1 : x_dims[0];
+  int rows = 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, GPUContext>(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>();
+
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, batches, 1);
+  if (batches > 1) {
+    auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num, 1);
+    DenseTensor batches_offset = Empty<int>(dev_ctx, {batches});
+    int* batches_offset_ptr = batches_offset.data<int>();
+    funcs::SetConstant<GPUContext, int> set_zero;
+    // set zero if the nnz=0 of batches[0]
+    set_zero(dev_ctx, &batches_offset, static_cast<IntT>(0));
+    GetBatchesOffset<IntT><<<config.block_per_grid.x,
+                             config.thread_per_block.x,
+                             0,
+                             dev_ctx.stream()>>>(
+        batches_ptr, batches, non_zero_num, batches_offset_ptr);
+
+    config.block_per_grid.y = batches;
+    ConvertCooRowsToCsrCrows<IntT><<<config.block_per_grid,
+                                     config.thread_per_block.x,
+                                     0,
+                                     dev_ctx.stream()>>>(
+        batches_offset_ptr, coo_rows_data, csr_crows_data, rows, non_zero_num);
+  } else {
+    ConvertCooRowsToCsrCrows<IntT><<<config.block_per_grid.x,
+                                     config.thread_per_block.x,
+                                     0,
+                                     dev_ctx.stream()>>>(
+        nullptr, coo_rows_data, csr_crows_data, rows, non_zero_num);
+  }
+
+  backends::gpu::GpuMemcpyAsync(csr_cols_data,
+                                coo_cols_data,
+                                sizeof(IntT) * non_zero_num,
+                                gpuMemcpyDeviceToDevice,
+                                dev_ctx.stream());
+  backends::gpu::GpuMemcpyAsync(csr_values_data,
+                                coo_values_data,
+                                sizeof(T) * non_zero_num,
+                                gpuMemcpyDeviceToDevice,
+                                dev_ctx.stream());
+  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(), "CooToCsrGPUKernel", ([&] {
+                                 CooToCsrGPUKernel<T, data_t>(dev_ctx, x, out);
+                               }));
+}
+
+template <typename ValueT, typename IndicesT>
+__global__ void KernelCooToDense(const IndicesT* indices,
+                                 const int64_t* sparse_offsets,
+                                 const ValueT* data,
+                                 ValueT* dense_data,
+                                 const IndicesT non_zero_num,
+                                 const int64_t base_offset,
+                                 const int64_t sparse_dim) {
+  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  for (int i = tid; i < non_zero_num; i += gridDim.x * blockDim.x) {
+    int64_t index = 0;
+    for (int j = 0; j < sparse_dim; j++) {
+      index += indices[j * non_zero_num + i] * sparse_offsets[j];
+    }
+
+    for (int j = 0; j < base_offset; j++) {
+      dense_data[index * base_offset + j] = data[i * base_offset + j];
+    }
+  }
+}
+
+template <typename T, typename IntT>
+void CooToDenseGPUKernel(const GPUContext& 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 = indices.dims();
+  int64_t sparse_dim = indices_dims[0];
+  if (indices_dims.size() == 1) {
+    sparse_dim = 1;
+  }
+  const int64_t dense_dim = values.dims().size() - 1;
+
+  const auto place = dev_ctx.GetPlace();
+  dev_ctx.template Alloc<T>(out);
+
+  T* out_data = out->data<T>();
+  backends::gpu::GpuMemsetAsync(
+      out_data, 0, sizeof(T) * out->numel(), dev_ctx.stream());
+
+  if (x.nnz() <= 0) {
+    return;
+  }
+
+  const T* x_data = values.data<T>();
+  int64_t base_offset = 1;
+  for (int64_t i = 0; i < dense_dim; i++) {
+    base_offset *= dense_dims[sparse_dim + i];
+  }
+  std::vector<int64_t> sparse_offsets(sparse_dim);
+  int64_t offset = 1;
+  for (int i = sparse_dim - 1; i >= 0; i--) {
+    sparse_offsets[i] = offset;
+    offset *= dense_dims[i];
+  }
+
+  DenseTensor d_sparse_offsets = Empty<int64_t>(dev_ctx, {sparse_dim});
+
+  BUILD_CUDA_TENSOR(int64_t, sparse_offsets, d_sparse_offsets);
+
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num, 1);
+
+  KernelCooToDense<T, IntT>
+      <<<config.block_per_grid.x,
+         config.thread_per_block.x,
+         0,
+         dev_ctx.stream()>>>(indices.data<IntT>(),
+                             d_sparse_offsets.data<int64_t>(),
+                             x_data,
+                             out_data,
+                             non_zero_num,
+                             base_offset,
+                             sparse_dim);
+  phi::sparse::ConvHostBuffer& conv_host_buffer =
+      phi::sparse::ConvHostBuffer::getInstance();
+  conv_host_buffer.reset();
+}
+
+template <typename T, typename Context>
+void CooToDenseKernel(const Context& dev_ctx,
+                      const SparseCooTensor& x,
+                      DenseTensor* out) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(
+      x.indices().dtype(), "CooToDenseGPUKernel", ([&] {
+        CooToDenseGPUKernel<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 sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(dense_to_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::DenseToCooKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {}
+
+PD_REGISTER_KERNEL(csr_to_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CsrToCooKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::complex64,
+                   phi::complex128) {}
+
+PD_REGISTER_KERNEL(coo_to_csr,
+                   GPU,
+                   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,
+                   GPU,
+                   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,
+                   GPU,
+                   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,
+                   GPU,
+                   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,
+                   GPU,
+                   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(values_csr,
+                   GPU,
+                   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(indices_coo,
+                   GPU,
+                   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(sparse_coo_tensor,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SparseCooTensorKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   phi::complex64,
+                   phi::complex128) {}

paddle/phi/kernels/sparse/gpu/sum_grad_kernel.cu

@@ -0,0 +1,233 @@
+// 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/backends/gpu/gpu_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/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/reduce_sum_grad_kernel.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/unary_grad_kernel.h"
+#include "paddle/phi/kernels/sparse/unary_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T>
+__global__ void SetValueCudaKernel(const T* value,
+                                   const int64_t length,
+                                   T* data) {
+  CUDA_KERNEL_LOOP_TYPE(index, length, int64_t) { data[index] = value[0]; }
+}
+
+template <typename T>
+__global__ void SumCsr2DGradCudaKernel(const int64_t* x_crows_data,
+                                       const T* dout_values_data,
+                                       const int64_t x_dim0,
+                                       T* dx_values_data) {
+  // dout_crows_data[index] should be equal to index;
+  CUDA_KERNEL_LOOP_TYPE(index, x_dim0, int64_t) {
+    T value = dout_values_data[index];
+    for (auto i = x_crows_data[index]; i < x_crows_data[index + 1]; ++i) {
+      dx_values_data[i] = value;
+    }
+  }
+}
+
+template <typename T>
+__global__ void SumCsr3DGradCudaKernel(const int64_t* x_crows_data,
+                                       const T* dout_values_data,
+                                       const int64_t x_dim0,
+                                       const int64_t x_dim1,
+                                       T* dx_values_data) {
+  // dout_crows_data[index] should be equal to number;
+  CUDA_KERNEL_LOOP_TYPE(index, x_dim0 * (x_dim1 + 1) - 1, int64_t) {
+    int64_t batch = index / (x_dim1 + 1);
+    int64_t number = index % (x_dim1 + 1);
+
+    // compute offset of dx_values_data in every batch
+    int64_t batch_offset = 0;
+    for (int64_t b = 1; b <= batch; ++b) {
+      batch_offset += x_crows_data[b * (x_dim1 + 1) - 1];
+    }
+
+    T value = dout_values_data[index - batch];
+    for (auto i = x_crows_data[index]; i < x_crows_data[index + 1]; ++i) {
+      dx_values_data[i + batch_offset] = value;
+    }
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void SumCooGradGPUKernel(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<IntT>();
+  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<IntT>();
+  auto* dx_values_data = dx_values->data<T>();
+
+  if (n_dim == 0) {
+    auto length = dx->nnz();
+    for (auto i = 1; i < x.values().dims().size(); ++i) {
+      length *= x.values().dims()[i];
+    }
+    auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, length, 1);
+
+    SetValueCudaKernel<T>
+        <<<config.block_per_grid.x,
+           config.thread_per_block.x,
+           0,
+           dev_ctx.stream()>>>(dout_values_data, length, dx_values_data);
+
+    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);
+  } else {
+    *dx_values = dout_values;
+  }
+  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,
+                      SparseCsrTensor* dx) {
+  EmptyLikeCsrKernel<T, Context>(dev_ctx, x, dx);
+  size_t n_dim = axis.size();
+
+  const DenseTensor& x_crows = x.crows();
+  const DenseTensor& x_cols = x.cols();
+  const DenseTensor& dout_values = dout.values();
+
+  DenseTensor* dx_crows = dx->mutable_crows();
+  DenseTensor* dx_cols = dx->mutable_cols();
+  DenseTensor* dx_values = dx->mutable_values();
+
+  const auto* x_crows_data = x_crows.data<int64_t>();
+  const auto* dout_values_data = dout_values.data<T>();
+  auto* dx_values_data = dx_values->data<T>();
+
+  *dx_crows = x_crows;
+  *dx_cols = x_cols;
+
+  if (n_dim == 0) {
+    auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, dx->nnz(), 1);
+    SetValueCudaKernel<T>
+        <<<config.block_per_grid.x,
+           config.thread_per_block.x,
+           0,
+           dev_ctx.stream()>>>(dout_values_data, dx->nnz(), dx_values_data);
+
+    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) {
+    auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x.dims()[0], 1);
+    SumCsr2DGradCudaKernel<T><<<config.block_per_grid.x,
+                                config.thread_per_block.x,
+                                0,
+                                dev_ctx.stream()>>>(
+        x_crows_data, dout_values_data, x.dims()[0], dx_values_data);
+  } else {
+    auto config = backends::gpu::GetGpuLaunchConfig1D(
+        dev_ctx, x.dims()[0] * (x.dims()[1] + 1), 1);
+    SumCsr3DGradCudaKernel<T><<<config.block_per_grid.x,
+                                config.thread_per_block.x,
+                                0,
+                                dev_ctx.stream()>>>(x_crows_data,
+                                                    dout_values_data,
+                                                    x.dims()[0],
+                                                    x.dims()[1],
+                                                    dx_values_data);
+  }
+  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(), "SumCooGradGPUKernel", ([&] {
+        SumCooGradGPUKernel<T, data_t, Context>(
+            dev_ctx, x, dout, axis, keep_dim, dx);
+      }));
+}
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(sum_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SumCooGradKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(sum_csr_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SumCsrGradKernel,
+                   float,
+                   double,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/gpu/sum_kernel.cu

@@ -0,0 +1,461 @@
+// 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/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+#include "paddle/phi/backends/gpu/gpu_primitives.h"
+#include "paddle/phi/common/scalar.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/cum_kernel.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/index_select_kernel.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
+#include "paddle/phi/kernels/reshape_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename IntT>
+__global__ void SumCooCudaKernel(const IntT* x_indices_data,
+                                 const T* x_values_data,
+                                 const int64_t x_nnz,
+                                 const int64_t dense_dim,
+                                 const int64_t sparse_dim,
+                                 const int64_t axis,
+                                 const bool keep_dim,
+                                 IntT* out_indices_data,
+                                 T* out_values_data) {
+  CUDA_KERNEL_LOOP_TYPE(index_i, x_nnz, int64_t) {
+    int64_t i = 0;
+    for (int j = 0; j < dense_dim; ++j) {
+      out_values_data[j + index_i * dense_dim] = 0;
+    }
+
+    int64_t _index_j_ =
+        static_cast<int64_t>(blockIdx.y) * blockDim.y + threadIdx.y;
+    for (auto index_j = _index_j_; index_j < x_nnz;
+         index_j += static_cast<int64_t>(blockDim.y) * gridDim.y) {
+      // Determine whether the index_i and index_j elements have the same
+      // indices in all dimensions except for the specified axis dimension.
+      bool same = true;
+      for (int j = 0; j < sparse_dim + !keep_dim; ++j) {
+        if (j != axis && x_indices_data[index_i + j * x_nnz] !=
+                             x_indices_data[index_j + j * x_nnz]) {
+          same = false;
+          break;
+        }
+      }
+      if (same) {
+        for (int j = 0; j < dense_dim; ++j) {
+          CudaAtomicAdd(&out_values_data[j + index_i * dense_dim],
+                        x_values_data[j + index_j * dense_dim]);
+        }
+      }
+    }
+    if (_index_j_ != 0) {
+      return;
+    }
+    if (keep_dim) {
+      for (int j = 0; j < sparse_dim; ++j) {
+        if (j == axis) {
+          out_indices_data[index_i + j * x_nnz] = 0;
+        } else {
+          out_indices_data[index_i + j * x_nnz] =
+              x_indices_data[index_i + j * x_nnz];
+        }
+      }
+      return;
+    }
+    for (int j = 0; j < sparse_dim; ++j) {
+      // out_indices_data [sparse_dim, x.nnz()]
+      int64_t x_indices_data_offset;
+      if (j < axis) {
+        x_indices_data_offset = index_i + j * x_nnz;
+      } else {
+        x_indices_data_offset = index_i + (j + 1) * x_nnz;
+      }
+      out_indices_data[index_i + j * x_nnz] =
+          x_indices_data[x_indices_data_offset];
+    }
+  }
+}
+
+__global__ void SumAllCsrCudaKernel(int64_t* out_crows_data,
+                                    int64_t* out_cols_data) {
+  CUDA_KERNEL_LOOP_TYPE(index, 2, int64_t) {
+    out_crows_data[index] = index;
+    if (index == 0) {
+      out_cols_data[0] = 0;
+    }
+  }
+}
+
+template <typename T>
+__global__ void SumCsr2DCudaKernel(const int64_t* x_crows_data,
+                                   const T* x_values_data,
+                                   const int64_t x_dim0,
+                                   int64_t* out_crows_data,
+                                   int64_t* out_cols_data,
+                                   T* out_values_data) {
+  CUDA_KERNEL_LOOP_TYPE(index, x_dim0 + 1, int64_t) {
+    out_crows_data[index] = index;
+    if (index != x_dim0) {
+      out_cols_data[index] = 0;
+      T sum_value = 0;
+      for (auto j = x_crows_data[index]; j < x_crows_data[index + 1]; ++j) {
+        sum_value += x_values_data[j];
+      }
+      out_values_data[index] = sum_value;
+    }
+  }
+}
+
+template <typename T>
+__global__ void SumCsr3DCudaKernel(const int64_t* x_crows_data,
+                                   const T* x_values_data,
+                                   const int64_t x_dim0,
+                                   const int64_t x_dim1,
+                                   const int64_t* batch_nnz_data,
+                                   int64_t* out_crows_data,
+                                   int64_t* out_cols_data,
+                                   T* out_values_data) {
+  {
+    CUDA_KERNEL_LOOP_TYPE(index, x_dim0 * x_dim1, int64_t) {
+      out_cols_data[index] = 0;
+    }
+  }
+
+  CUDA_KERNEL_LOOP_TYPE(index, x_dim0 * (x_dim1 + 1), int64_t) {
+    int64_t batch = index / (x_dim1 + 1);
+    int64_t number = index % (x_dim1 + 1);
+    out_crows_data[index] = number;
+
+    if (number != x_dim1) {
+      T sum_value = 0;
+      int64_t x_values_data_offset;
+      if (batch == 0) {
+        x_values_data_offset = 0;
+      } else {
+        x_values_data_offset = batch_nnz_data[batch - 1];
+      }
+      for (int64_t j = x_crows_data[index]; j < x_crows_data[index + 1]; ++j) {
+        sum_value += x_values_data[j + x_values_data_offset];
+      }
+
+      // `index - batch` would never exceed x_dim0 * x_dim1.
+      out_values_data[index - batch] = sum_value;
+    }
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void SumCooGPU0Kernel(const Context& dev_ctx,
+                      const SparseCooTensor& x,
+                      const IntArray& axis,
+                      DataType dtype,
+                      bool keep_dim,
+                      SparseCooTensor* out) {
+  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 (keep_dim) {
+    out_dims = make_ddim(std::vector<int64_t>(x_dims.size(), 1));
+    out_indices = Empty<IntT, Context>(dev_ctx, {sparse_dim, 1});
+  } else {
+    out_dims = make_ddim({1});
+    out_indices = Empty<IntT, Context>(dev_ctx, {1, 1});
+  }
+  funcs::SetConstant<Context, IntT> set_out_indices;
+  set_out_indices(dev_ctx, &out_indices, static_cast<IntT>(0));
+  out_values = phi::Sum<T>(dev_ctx, x.values(), {}, dtype, keep_dim);
+  out->SetMember(out_indices, out_values, out_dims, x.coalesced());
+}
+
+template <typename T, typename IntT, typename Context>
+void SumCooGPU1Kernel(const Context& dev_ctx,
+                      const SparseCooTensor& x,
+                      const IntArray& axis,
+                      DataType dtype,
+                      bool keep_dim,
+                      SparseCooTensor* out) {
+  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;
+  auto n_dim = x.dims().size();
+  auto dim = axis[0] < 0 ? x_dims.size() + axis[0] : axis[0];
+
+  std::vector<int64_t> dims;
+  for (int i = 0; i < n_dim; ++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;
+  }
+
+  std::vector<int> out_values_dims;
+  out_values_dims.push_back(x.nnz());
+  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, x.nnz()});
+  out_values = Empty<T, Context>(dev_ctx, out_values_dims);
+
+  const auto* x_indices_data = x_indices.data<IntT>();
+  const auto* x_values_data = x_values.data<T>();
+  auto* out_indices_data = out_indices.data<IntT>();
+  auto* out_values_data = out_values.data<T>();
+
+  auto config = backends::gpu::GetGpuLaunchConfig2D(dev_ctx, x.nnz(), x.nnz());
+  SumCooCudaKernel<T, IntT><<<config.block_per_grid.x,
+                              config.thread_per_block.x,
+                              0,
+                              dev_ctx.stream()>>>(x_indices_data,
+                                                  x_values_data,
+                                                  x.nnz(),
+                                                  dense_dim,
+                                                  sparse_dim,
+                                                  dim,
+                                                  keep_dim,
+                                                  out_indices_data,
+                                                  out_values_data);
+  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 SumCooKernel(const Context& dev_ctx,
+                  const SparseCooTensor& x,
+                  const IntArray& axis,
+                  DataType dtype,
+                  bool keep_dim,
+                  SparseCooTensor* out) {
+  const size_t n_dim = axis.size();
+  if (n_dim == 0) {
+    PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "SumCooGPUKernel", ([&] {
+                                   SumCooGPU0Kernel<T, data_t, Context>(
+                                       dev_ctx, x, axis, dtype, keep_dim, out);
+                                 }));
+  } else {
+    PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "SumCooGPUKernel", ([&] {
+                                   SumCooGPU1Kernel<T, data_t, Context>(
+                                       dev_ctx, x, axis, dtype, keep_dim, out);
+                                 }));
+  }
+}
+
+template <typename T, typename Context>
+void SumCsr0Kernel(const Context& dev_ctx,
+                   const SparseCsrTensor& x,
+                   const IntArray& axis,
+                   DataType dtype,
+                   bool keep_dim,
+                   SparseCsrTensor* out) {
+  auto x_dim0 = x.dims()[0];
+  auto x_dim1 = x.dims()[1];
+  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 (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]
+  out_cols = Empty<int64_t, Context>(dev_ctx, {1});   // crows = [0]
+  auto* out_crows_data = out_crows.data<int64_t>();
+  auto* out_cols_data = out_cols.data<int64_t>();
+
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, 2, 1);
+  SumAllCsrCudaKernel<<<config.block_per_grid.x,
+                        config.thread_per_block.x,
+                        0,
+                        dev_ctx.stream()>>>(out_crows_data, out_cols_data);
+
+  out_values = phi::Sum<T>(dev_ctx, x.values(), {}, dtype, true);
+  out->SetMember(out_crows, out_cols, out_values, out_dims);
+}
+
+template <typename T, typename Context>
+void SumCsr1Kernel(const Context& dev_ctx,
+                   const SparseCsrTensor& x,
+                   const IntArray& axis,
+                   DataType dtype,
+                   bool keep_dim,
+                   SparseCsrTensor* out) {
+  auto x_dim0 = x.dims()[0];
+  auto x_dim1 = x.dims()[1];
+  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;
+  out_crows = EmptyLike<int64_t, Context>(dev_ctx, x.crows());
+  auto* out_crows_data = out_crows.data<int64_t>();
+
+  if (x.dims().size() == 2) {
+    out_cols = Empty<int64_t, Context>(dev_ctx, {x_dim0});
+    out_values = Empty<T, Context>(dev_ctx, {x_dim0});
+    auto* out_cols_data = out_cols.data<int64_t>();
+    auto* out_values_data = out_values.data<T>();
+    out_dims = make_ddim({x_dim0, 1});
+    auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x_dim0 + 1, 1);
+    SumCsr2DCudaKernel<T><<<config.block_per_grid.x,
+                            config.thread_per_block.x,
+                            0,
+                            dev_ctx.stream()>>>(x_crows_data,
+                                                x_values_data,
+                                                x_dim0,
+                                                out_crows_data,
+                                                out_cols_data,
+                                                out_values_data);
+
+  } else {
+    out_cols = Empty<int64_t, Context>(dev_ctx, {x_dim0 * x_dim1});
+    out_values = Empty<T, Context>(dev_ctx, {x_dim0 * x_dim1});
+    auto* out_cols_data = out_cols.data<int64_t>();
+    auto* out_values_data = out_values.data<T>();
+    if (keep_dim) {
+      out_dims = make_ddim({x_dim0, x_dim1, 1});
+    } else {
+      out_dims = make_ddim({x_dim0, x_dim1});
+    }
+
+    DenseTensor x_crows_reshape =
+        Reshape<int64_t, Context>(dev_ctx, x_crows, {x_dim0, x_dim1 + 1});
+    DenseTensor last_indices = Empty<int64_t, Context>(dev_ctx, {1});
+    funcs::SetConstant<Context, int64_t> set_constant;
+    set_constant(dev_ctx, &last_indices, static_cast<int64_t>(x_dim1));
+
+    DenseTensor x_crows_last = Empty<int64_t, Context>(dev_ctx, {x_dim0, 1});
+    IndexSelectKernel<int64_t, Context>(
+        dev_ctx, x_crows_reshape, last_indices, 1, &x_crows_last);
+
+    DenseTensor batch_nnz = Empty<int64_t, Context>(dev_ctx, {x_dim0, 1});
+    CumsumKernel<int64_t, Context>(
+        dev_ctx, x_crows_last, Scalar(0), false, false, false, &batch_nnz);
+    auto* batch_nnz_data = batch_nnz.data<int64_t>();
+
+    auto config = backends::gpu::GetGpuLaunchConfig1D(
+        dev_ctx, x.dims()[0] * (x.dims()[1] + 1), 1);
+    SumCsr3DCudaKernel<T><<<config.block_per_grid.x,
+                            config.thread_per_block.x,
+                            0,
+                            dev_ctx.stream()>>>(x_crows_data,
+                                                x_values_data,
+                                                x_dim0,
+                                                x_dim1,
+                                                batch_nnz_data,
+                                                out_crows_data,
+                                                out_cols_data,
+                                                out_values_data);
+  }
+  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 SumCsrKernel(const Context& dev_ctx,
+                  const SparseCsrTensor& x,
+                  const IntArray& axis,
+                  DataType dtype,
+                  bool keep_dim,
+                  SparseCsrTensor* out) {
+  size_t n_dim = axis.size();
+  if (n_dim == 0) {
+    SumCsr0Kernel<T, Context>(dev_ctx, x, axis, dtype, keep_dim, out);
+  } 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."));
+    SumCsr1Kernel<T, Context>(dev_ctx, x, axis, dtype, keep_dim, out);
+  }
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(sum_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SumCooKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}
+
+PD_REGISTER_KERNEL(sum_csr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SumCsrKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}

paddle/phi/kernels/sparse/gpu/sync_batch_norm_grad_kernel.cu

@@ -0,0 +1,85 @@
+// Copyright (c) 2024 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/sync_batch_norm_grad_kernel.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/sync_batch_norm_utils.h"
+
+// sparse header
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context>
+void SyncBatchNormCooGradKernel(const Context& dev_ctx,
+                                const SparseCooTensor& x,
+                                const DenseTensor& scale,
+                                const DenseTensor& bias,
+                                const DenseTensor& saved_mean,
+                                const DenseTensor& saved_variance,
+                                const optional<DenseTensor>& reserve_space,
+                                const SparseCooTensor& y_grad,
+                                float momentum,
+                                float epsilon,
+                                const std::string& data_layout,
+                                bool is_test,
+                                bool use_global_stats,
+                                bool trainable_statistics,
+                                SparseCooTensor* x_grad,
+                                DenseTensor* scale_grad,
+                                DenseTensor* bias_grad) {
+  EmptyLikeCooKernel<T, Context>(dev_ctx, x, x_grad);
+  *scale_grad = EmptyLike<T, Context>(dev_ctx, scale);
+  *bias_grad = EmptyLike<T, Context>(dev_ctx, bias);
+  phi::SyncBatchNormGradKernel<T, Context>(dev_ctx,
+                                           x.values(),
+                                           scale,
+                                           bias,
+                                           saved_mean,
+                                           saved_variance,
+                                           reserve_space,
+                                           y_grad.values(),
+                                           momentum,
+                                           epsilon,
+                                           data_layout,
+                                           is_test,
+                                           use_global_stats,
+                                           trainable_statistics,
+                                           x_grad->mutable_values(),
+                                           scale_grad,
+                                           bias_grad);
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+#ifdef PADDLE_WITH_HIP
+PD_REGISTER_KERNEL(sync_batch_norm_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SyncBatchNormCooGradKernel,
+                   float,
+                   phi::float16) {}
+#else
+PD_REGISTER_KERNEL(sync_batch_norm_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SyncBatchNormCooGradKernel,
+                   float,
+                   double,
+                   phi::float16) {}
+#endif

paddle/phi/kernels/sparse/gpu/sync_batch_norm_kernel.cu

@@ -0,0 +1,87 @@
+// Copyright (c) 2024 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/sync_batch_norm_kernel.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/sync_batch_norm_utils.h"
+
+// sparse header
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context>
+void SyncBatchNormCooKernel(const Context& dev_ctx,
+                            const SparseCooTensor& x,
+                            const DenseTensor& mean,
+                            const DenseTensor& variance,
+                            const DenseTensor& scale,
+                            const DenseTensor& bias,
+                            bool is_test,
+                            float momentum,
+                            float epsilon,
+                            const std::string& data_layout,
+                            bool use_global_stats,
+                            bool trainable_statistics,
+                            SparseCooTensor* y,
+                            DenseTensor* mean_out,
+                            DenseTensor* variance_out,
+                            DenseTensor* saved_mean,
+                            DenseTensor* saved_variance,
+                            DenseTensor* reserve_space) {
+  EmptyLikeCooKernel<T, Context>(dev_ctx, x, y);
+  phi::SyncBatchNormKernel<T, Context>(dev_ctx,
+                                       x.values(),
+                                       mean,
+                                       variance,
+                                       scale,
+                                       bias,
+                                       is_test,
+                                       momentum,
+                                       epsilon,
+                                       data_layout,
+                                       use_global_stats,
+                                       trainable_statistics,
+                                       y->mutable_values(),
+                                       mean_out,
+                                       variance_out,
+                                       saved_mean,
+                                       saved_variance,
+                                       reserve_space);
+  y->SetIndicesDict(x.GetIndicesDict());
+  y->SetKmaps(x.GetKmaps());
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+#ifdef PADDLE_WITH_HIP
+PD_REGISTER_KERNEL(sync_batch_norm_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SyncBatchNormCooKernel,
+                   float,
+                   phi::float16) {}
+#else
+PD_REGISTER_KERNEL(sync_batch_norm_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SyncBatchNormCooKernel,
+                   float,
+                   double,
+                   phi::float16) {}
+#endif

paddle/phi/kernels/sparse/gpu/transpose_grad_kernel.cu

@@ -0,0 +1,80 @@
+// 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/gpu/gpu_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 {
+namespace sparse {
+
+std::vector<int> get_gpu_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]] = 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_gpu_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_gpu_grad_perm(perm);
+  TransposeCsrKernel<T, Context>(dev_ctx, dout, grad_perm, dx);
+}
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(transpose_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::TransposeCooGradKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(transpose_csr_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::TransposeCsrGradKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/gpu/transpose_kernel.cu

@@ -0,0 +1,355 @@
+// 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/gpu/gpu_context.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/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+__global__ void TransposeCooCudaKernel(const int64_t *x_indices_data,
+                                       const int *perm,
+                                       const std::size_t n_dim,
+                                       const int64_t x_nnz,
+                                       int64_t *out_indices_data) {
+  CUDA_KERNEL_LOOP_TYPE(index, x_nnz * n_dim, int64_t) {
+    int64_t i = index / x_nnz;
+    int64_t j = index % x_nnz;
+    out_indices_data[index] = x_indices_data[j + perm[i] * x_nnz];
+  }
+}
+
+template <typename T, typename IntT>
+__global__ void TransposeCsr2DCudaKernel(const IntT *x_crows_data,
+                                         const IntT *x_cols_data,
+                                         const T *x_values_data,
+                                         const int *perm,
+                                         const int64_t *x_dims,
+                                         const int64_t *out_dims,
+                                         const int64_t x_nnz,
+                                         IntT *out_crows_data,
+                                         IntT *out_cols_data,
+                                         T *out_values_data) {
+  int64_t __index__ =
+      static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  // compute out_crows_data by x_cols_data
+  for (int64_t i = __index__; i <= out_dims[0]; i += blockDim.x * gridDim.x) {
+    out_crows_data[i] = 0;
+  }
+  __syncthreads();
+  if (__index__ == 0) {
+    for (int64_t i = 0; i < x_nnz; ++i) {
+      IntT j = x_cols_data[i];
+      out_crows_data[j + 2]++;
+    }
+    for (int64_t i = 0; i < out_dims[0]; i += 1) {
+      out_crows_data[i + 1] += out_crows_data[i];
+    }
+    // compute out_cols_data and out_values_data by out_crows_data and x
+    for (int i = 0; i < x_dims[0]; ++i) {
+      IntT start = x_crows_data[i];
+      IntT end = x_crows_data[i + 1];
+      for (IntT j = start; j < end; ++j) {
+        IntT x_cols_j = x_cols_data[j] + 1;
+        IntT jjj = out_crows_data[x_cols_j];
+        out_cols_data[jjj] = i;
+        out_values_data[jjj] = x_values_data[j];
+        out_crows_data[x_cols_j]++;
+      }
+    }
+  }
+}
+
+template <typename T, typename IntT>
+__global__ void TransposeCsr3DCudaKernel(const IntT *x_crows_data,
+                                         const IntT *x_cols_data,
+                                         const T *x_values_data,
+                                         const int *perm,
+                                         const int64_t *x_dims,
+                                         const int64_t *out_dims,
+                                         const std::size_t n_dim,
+                                         const int64_t x_nnz,
+                                         IntT *out_crows_data,
+                                         IntT *out_cols_data,
+                                         T *out_values_data) {
+  int64_t __index__ =
+      static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (__index__ == 0) {
+    int out_n_rows = out_dims[1];
+    int x_n_rows = x_dims[1];
+    for (int k = 0; k < out_dims[0]; ++k) {
+      if (perm[0] == 0) {  // dims == {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) {
+          int j = x_cols_data[i];
+          out_crows_data[j + 2]++;
+        }
+        for (int i = 0; i < out_n_rows; ++i) {
+          out_crows_data[i + 1] += out_crows_data[i];
+        }
+        // compute out_cols_data and out_values_data by out_crows_data and x
+        for (int i = 0; i < x_n_rows; ++i) {
+          IntT start = x_crows_data[i];
+          IntT end = x_crows_data[i + 1];
+          for (IntT j = start; j < end; ++j) {
+            IntT x_cols_j = x_cols_data[j] + 1;
+            IntT jjj = out_crows_data[x_cols_j];
+            out_cols_data[jjj] = i;
+            out_values_data[jjj] = x_values_data[j];
+            out_crows_data[x_cols_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;
+        }
+        int x_cols_offset = 0;
+        int out_cols_index = 0;
+        for (int i = 0; i < x_dims[0]; ++i) {
+          int x_crows_index = i * (x_n_rows + 1);
+          int start = x_crows_data[x_crows_index + k];
+          int end = x_crows_data[x_crows_index + 1 + k];
+          out_crows_data[i + 1] = end - start;
+          for (int 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;
+    }
+  }
+}
+
+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();
+  std::size_t n_dim = perm.size();
+  DDim out_dims = x.dims().transpose(perm);
+  DenseTensor out_indices = EmptyLike<int64_t, Context>(dev_ctx, x.indices());
+  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>();
+  int *d_perm;
+
+  auto d_perm_tensor = memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int) * perm.size(),
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  d_perm = reinterpret_cast<int *>(d_perm_tensor->ptr());
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     d_perm,
+                     phi::CPUPlace(),
+                     perm.data(),
+                     sizeof(int) * perm.size(),
+                     dev_ctx.stream());
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x_nnz * n_dim, 1);
+  TransposeCooCudaKernel<<<config.block_per_grid.x,
+                           config.thread_per_block.x,
+                           0,
+                           dev_ctx.stream()>>>(
+      x_indices_data, d_perm, n_dim, x_nnz, out_indices_data);
+}
+
+template <typename T, typename IntT>
+void TransposeCsrGpuKernel(const GPUContext &dev_ctx,
+                           const SparseCsrTensor &x,
+                           const std::vector<int> &perm,
+                           SparseCsrTensor *out) {
+  std::size_t n_dim = perm.size();
+  const DenseTensor &x_crows = x.crows();
+  const DenseTensor &x_cols = x.cols();
+  const DenseTensor &x_values = x.non_zero_elements();
+  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<IntT, GPUContext>(dev_ctx, {out_dims[0] + 1});
+  } else {
+    out_crows =
+        Empty<IntT, GPUContext>(dev_ctx, {out_dims[0] * (out_dims[1] + 1)});
+  }
+  out_cols = EmptyLike<IntT, GPUContext>(dev_ctx, x.cols());
+  out_values = EmptyLike<T, GPUContext>(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, GPUContext>(dev_ctx, x, {1, 0, 2}, &temp);
+    TransposeCsrKernel<T, GPUContext>(dev_ctx, temp, {0, 2, 1}, out);
+    return;
+  } else if (perm[0] == 2 && perm[1] == 0) {  // perm == {2, 0, 1}
+    SparseCsrTensor temp;
+    TransposeCsrKernel<T, GPUContext>(dev_ctx, x, {0, 2, 1}, &temp);
+    TransposeCsrKernel<T, GPUContext>(dev_ctx, temp, {1, 0, 2}, out);
+    return;
+  } else if (perm[0] == 2 && perm[1] == 1) {  // perm == {2, 1, 0}
+    SparseCsrTensor temp;
+    TransposeCsrKernel<T, GPUContext>(dev_ctx, x, {1, 0, 2}, &temp);
+    TransposeCsrKernel<T, GPUContext>(dev_ctx, temp, {2, 0, 1}, out);
+    return;
+  }
+  IntT *out_crows_data = out_crows.data<IntT>();
+  IntT *out_cols_data = out_cols.data<IntT>();
+  T *out_values_data = out_values.data<T>();
+  const IntT *x_crows_data = x_crows.data<IntT>();
+  const IntT *x_cols_data = x_cols.data<IntT>();
+  const T *x_values_data = x_values.data<T>();
+  int *d_perm;
+  int64_t *d_x_dims, *d_out_dims;
+
+  auto d_perm_tensor = memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int) * perm.size(),
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  d_perm = reinterpret_cast<int *>(d_perm_tensor->ptr());
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     d_perm,
+                     phi::CPUPlace(),
+                     perm.data(),
+                     sizeof(int) * perm.size(),
+                     dev_ctx.stream());
+  auto d_x_dims_tensor = memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int64_t) * x.dims().size(),
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  d_x_dims = reinterpret_cast<int64_t *>(d_x_dims_tensor->ptr());
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     d_x_dims,
+                     phi::CPUPlace(),
+                     x.dims().Get(),
+                     sizeof(int64_t) * x.dims().size(),
+                     dev_ctx.stream());
+  auto d_out_dims_tensor = memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int64_t) * out_dims.size(),
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  d_out_dims = reinterpret_cast<int64_t *>(d_out_dims_tensor->ptr());
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     d_out_dims,
+                     phi::CPUPlace(),
+                     out_dims.Get(),
+                     sizeof(int64_t) * out_dims.size(),
+                     dev_ctx.stream());
+
+  int64_t x_nnz = x.nnz();
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_dims[0], 1);
+  if (perm.size() == 2) {
+    TransposeCsr2DCudaKernel<T><<<config.block_per_grid.x,
+                                  config.thread_per_block.x,
+                                  0,
+                                  dev_ctx.stream()>>>(x_crows_data,
+                                                      x_cols_data,
+                                                      x_values_data,
+                                                      d_perm,
+                                                      d_x_dims,
+                                                      d_out_dims,
+                                                      x_nnz,
+                                                      out_crows_data,
+                                                      out_cols_data,
+                                                      out_values_data);
+  } else {
+    TransposeCsr3DCudaKernel<T><<<1, 1, 0, dev_ctx.stream()>>>(x_crows_data,
+                                                               x_cols_data,
+                                                               x_values_data,
+                                                               d_perm,
+                                                               d_x_dims,
+                                                               d_out_dims,
+                                                               perm.size(),
+                                                               x_nnz,
+                                                               out_crows_data,
+                                                               out_cols_data,
+                                                               out_values_data);
+  }
+}
+
+template <typename T, typename Context>
+void TransposeCsrKernel(const Context &dev_ctx,
+                        const SparseCsrTensor &x,
+                        const std::vector<int> &perm,
+                        SparseCsrTensor *out) {
+  PD_VISIT_BASE_INTEGRAL_TYPES(x.crows().dtype(), "TransposeCsrKernel", ([&] {
+                                 TransposeCsrGpuKernel<T, data_t>(
+                                     dev_ctx, x, perm, out);
+                               }));
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(transpose_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::TransposeCooKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(transpose_csr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::TransposeCsrKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/gpu/unary_grad_kernel.cu

@@ -0,0 +1,112 @@
+// 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/gpu/gpu_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_GPU_GRAD_KERNEL(name, prefix)     \
+  PD_REGISTER_KERNEL(name##_coo_grad,                              \
+                     GPU,                                          \
+                     ALL_LAYOUT,                                   \
+                     phi::sparse::prefix##CooGradKernel,           \
+                     phi::float16,                                 \
+                     float,                                        \
+                     double) {                                     \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); \
+  }                                                                \
+                                                                   \
+  PD_REGISTER_KERNEL(name##_csr_grad,                              \
+                     GPU,                                          \
+                     ALL_LAYOUT,                                   \
+                     phi::sparse::prefix##CsrGradKernel,           \
+                     phi::float16,                                 \
+                     float,                                        \
+                     double) {                                     \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \
+  }
+
+#define PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(name, prefix) \
+  PD_REGISTER_KERNEL(name##_coo_grad,                                       \
+                     GPU,                                                   \
+                     ALL_LAYOUT,                                            \
+                     phi::sparse::prefix##CooGradKernel,                    \
+                     phi::float16,                                          \
+                     float,                                                 \
+                     double,                                                \
+                     phi::complex64,                                        \
+                     phi::complex128) {                                     \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);          \
+  }                                                                         \
+                                                                            \
+  PD_REGISTER_KERNEL(name##_csr_grad,                                       \
+                     GPU,                                                   \
+                     ALL_LAYOUT,                                            \
+                     phi::sparse::prefix##CsrGradKernel,                    \
+                     phi::float16,                                          \
+                     float,                                                 \
+                     double,                                                \
+                     phi::complex64,                                        \
+                     phi::complex128) {                                     \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);          \
+  }
+
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(sqrt, Sqrt)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(relu, Relu)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(pow, Pow)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(relu6, Relu6)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(leaky_relu, LeakyRelu)
+
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(asin, Asin)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(asinh, Asinh)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(atanh, Atanh)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(expm1, Expm1)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(log1p, Log1p)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(square, Square)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(sinh, Sinh)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(sin, Sin)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(abs, Abs)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(tan, Tan)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(atan, Atan)
+PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(tanh, Tanh)
+
+PD_REGISTER_KERNEL(cast_coo_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CastCooGradKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(cast_csr_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CastCsrGradKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}

paddle/phi/kernels/sparse/gpu/unary_kernel.cu

@@ -0,0 +1,185 @@
+// 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/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/scale_kernel.h"
+#include "paddle/phi/kernels/sparse/impl/unary_kernel_impl.h"
+
+namespace phi {
+namespace 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);
+
+  phi::ScaleKernel<T, Context>(
+      dev_ctx, x.values(), 1 / scalar, 0.0f, false, out->mutable_values());
+}
+
+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);
+
+  phi::ScaleKernel<T, Context>(
+      dev_ctx, x.values(), 1 / scalar, 0.0f, false, out->mutable_values());
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+#define PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(name, prefix)          \
+  PD_REGISTER_KERNEL(name##_coo,                                   \
+                     GPU,                                          \
+                     ALL_LAYOUT,                                   \
+                     phi::sparse::prefix##CooKernel,               \
+                     phi::float16,                                 \
+                     float,                                        \
+                     double) {                                     \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); \
+  }                                                                \
+                                                                   \
+  PD_REGISTER_KERNEL(name##_csr,                                   \
+                     GPU,                                          \
+                     ALL_LAYOUT,                                   \
+                     phi::sparse::prefix##CsrKernel,               \
+                     phi::float16,                                 \
+                     float,                                        \
+                     double) {                                     \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \
+  }
+
+#define PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(name, prefix) \
+  PD_REGISTER_KERNEL(name##_coo,                                       \
+                     GPU,                                              \
+                     ALL_LAYOUT,                                       \
+                     phi::sparse::prefix##CooKernel,                   \
+                     phi::float16,                                     \
+                     float,                                            \
+                     double,                                           \
+                     phi::complex64,                                   \
+                     phi::complex128) {                                \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);     \
+  }                                                                    \
+                                                                       \
+  PD_REGISTER_KERNEL(name##_csr,                                       \
+                     GPU,                                              \
+                     ALL_LAYOUT,                                       \
+                     phi::sparse::prefix##CsrKernel,                   \
+                     phi::float16,                                     \
+                     float,                                            \
+                     double,                                           \
+                     phi::complex64,                                   \
+                     phi::complex128) {                                \
+    kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);     \
+  }
+
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(sqrt, Sqrt)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(relu, Relu)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(pow, Pow)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(scale, Scale)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(relu6, Relu6)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(leaky_relu, LeakyRelu)
+
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(asin, Asin)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(asinh, Asinh)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(atanh, Atanh)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(expm1, Expm1)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(log1p, Log1p)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(square, Square)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(tanh, Tanh)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(sinh, Sinh)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(tan, Tan)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(sin, Sin)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(abs, Abs)
+PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(atan, Atan)
+
+PD_REGISTER_KERNEL(divide_scalar_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::DivScalarCooKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
+}
+
+PD_REGISTER_KERNEL(divide_scalar_csr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::DivScalarCsrKernel,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}
+
+PD_REGISTER_KERNEL(cast_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CastCooKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(cast_csr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::CastCsrKernel,
+                   phi::float16,
+                   float,
+                   double,
+                   int8_t,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool) {}
+
+PD_REGISTER_KERNEL(isnan_coo,
+                   GPU,
+                   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,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::IsnanCsrKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   int,
+                   int64_t) {
+  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
+}

paddle/phi/kernels/sparse/impl/unary_grad_kernel_impl.h

@@ -0,0 +1,156 @@
+// 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 "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/abs_grad_kernel.h"
+#include "paddle/phi/kernels/activation_grad_kernel.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/impl/unary_kernel_impl.h"
+
+namespace phi {
+namespace sparse {
+
+#define DEFINE_SPARSE_UNARY_GRAD_KERNEL(prefix)                           \
+  template <typename T, typename Context>                                 \
+  void prefix##CooGradKernel(const Context& dev_ctx,                      \
+                             const SparseCooTensor& x_or_out,             \
+                             const SparseCooTensor& dout,                 \
+                             SparseCooTensor* dx) {                       \
+    EmptyLikeCooKernel<T, Context>(dev_ctx, x_or_out, dx);                \
+    if (dx->mutable_non_zero_elements()->numel() == 0) {                  \
+      return;                                                             \
+    }                                                                     \
+    phi::prefix##GradKernel<T, Context>(dev_ctx,                          \
+                                        x_or_out.non_zero_elements(),     \
+                                        dout.non_zero_elements(),         \
+                                        dx->mutable_non_zero_elements()); \
+  }                                                                       \
+                                                                          \
+  template <typename T, typename Context>                                 \
+  void prefix##CsrGradKernel(const Context& dev_ctx,                      \
+                             const SparseCsrTensor& x_or_out,             \
+                             const SparseCsrTensor& dout,                 \
+                             SparseCsrTensor* dx) {                       \
+    EmptyLikeCsrKernel<T, Context>(dev_ctx, x_or_out, dx);                \
+    if (dx->mutable_non_zero_elements()->numel() == 0) {                  \
+      return;                                                             \
+    }                                                                     \
+    phi::prefix##GradKernel<T, Context>(dev_ctx,                          \
+                                        x_or_out.non_zero_elements(),     \
+                                        dout.non_zero_elements(),         \
+                                        dx->mutable_non_zero_elements()); \
+  }
+
+#define DEFINE_SPARSE_UNARY_GRAD_KERNEL_WITH_ONE_ATTR(prefix, attr)       \
+  template <typename T, typename Context>                                 \
+  void prefix##CooGradKernel(const Context& dev_ctx,                      \
+                             const SparseCooTensor& x_or_out,             \
+                             const SparseCooTensor& dout,                 \
+                             float attr,                                  \
+                             SparseCooTensor* dx) {                       \
+    EmptyLikeCooKernel<T, Context>(dev_ctx, x_or_out, dx);                \
+    if (dx->mutable_non_zero_elements()->numel() == 0) {                  \
+      return;                                                             \
+    }                                                                     \
+    phi::prefix##GradKernel<T, Context>(dev_ctx,                          \
+                                        x_or_out.non_zero_elements(),     \
+                                        dout.non_zero_elements(),         \
+                                        attr,                             \
+                                        dx->mutable_non_zero_elements()); \
+  }                                                                       \
+                                                                          \
+  template <typename T, typename Context>                                 \
+  void prefix##CsrGradKernel(const Context& dev_ctx,                      \
+                             const SparseCsrTensor& x_or_out,             \
+                             const SparseCsrTensor& dout,                 \
+                             float attr,                                  \
+                             SparseCsrTensor* dx) {                       \
+    EmptyLikeCsrKernel<T, Context>(dev_ctx, x_or_out, dx);                \
+    if (dx->mutable_non_zero_elements()->numel() == 0) {                  \
+      return;                                                             \
+    }                                                                     \
+    phi::prefix##GradKernel<T, Context>(dev_ctx,                          \
+                                        x_or_out.non_zero_elements(),     \
+                                        dout.non_zero_elements(),         \
+                                        attr,                             \
+                                        dx->mutable_non_zero_elements()); \
+  }
+
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Sin)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Tan)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Asin)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Atan)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Sinh)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Tanh)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Asinh)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Atanh)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Sqrt)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Square)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Log1p)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Abs)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Relu)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Expm1)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL(Relu6)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL_WITH_ONE_ATTR(Pow, factor)
+DEFINE_SPARSE_UNARY_GRAD_KERNEL_WITH_ONE_ATTR(LeakyRelu, alpha)
+
+template <typename T, typename Context>
+void CastCooGradKernel(const Context& dev_ctx,
+                       const SparseCooTensor& x,
+                       const SparseCooTensor& dout,
+                       DataType value_dtype,
+                       SparseCooTensor* dx) {
+  EmptyLikeCooKernel<T, Context>(dev_ctx, x, dx);
+  if (value_dtype == DataType::UNDEFINED) {
+    phi::Copy(dev_ctx,
+              dout.non_zero_elements(),
+              dev_ctx.GetPlace(),
+              false,
+              dx->mutable_non_zero_elements());
+  } else {
+    CastKernel<T, Context>(dev_ctx,
+                           dout.non_zero_elements(),
+                           x.non_zero_elements().dtype(),
+                           dx->mutable_non_zero_elements());
+  }
+}
+
+template <typename T, typename Context>
+void CastCsrGradKernel(const Context& dev_ctx,
+                       const SparseCsrTensor& x,
+                       const SparseCsrTensor& dout,
+                       DataType value_dtype,
+                       SparseCsrTensor* dx) {
+  EmptyLikeCsrKernel<T, Context>(dev_ctx, x, dx);
+  if (value_dtype == DataType::UNDEFINED) {
+    phi::Copy(dev_ctx,
+              dout.non_zero_elements(),
+              dev_ctx.GetPlace(),
+              false,
+              dx->mutable_non_zero_elements());
+  } else {
+    CastKernel<T, Context>(dev_ctx,
+                           dout.non_zero_elements(),
+                           x.non_zero_elements().dtype(),
+                           dx->mutable_non_zero_elements());
+  }
+}
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/impl/unary_kernel_impl.h

@@ -0,0 +1,296 @@
+// 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 "paddle/phi/core/meta_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/abs_kernel.h"
+#include "paddle/phi/kernels/activation_kernel.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+#include "paddle/phi/kernels/isfinite_kernel.h"
+#include "paddle/phi/kernels/scale_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+#define DEFINE_SPARSE_UNARY_KERNEL(prefix)                                   \
+  template <typename T, typename Context>                                    \
+  void prefix##CooKernel(const Context& dev_ctx,                             \
+                         const SparseCooTensor& x,                           \
+                         SparseCooTensor* out) {                             \
+    EmptyLikeCooKernel<T, Context>(dev_ctx, x, out);                         \
+    if (out->mutable_non_zero_elements()->numel() != 0) {                    \
+      phi::prefix##Kernel<T, Context>(                                       \
+          dev_ctx, x.non_zero_elements(), out->mutable_non_zero_elements()); \
+    }                                                                        \
+    out->SetIndicesDict(x.GetIndicesDict());                                 \
+    out->SetKmaps(x.GetKmaps());                                             \
+  }                                                                          \
+                                                                             \
+  template <typename T, typename Context>                                    \
+  void prefix##CsrKernel(const Context& dev_ctx,                             \
+                         const SparseCsrTensor& x,                           \
+                         SparseCsrTensor* out) {                             \
+    EmptyLikeCsrKernel<T, Context>(dev_ctx, x, out);                         \
+    if (out->mutable_non_zero_elements()->numel() == 0) {                    \
+      return;                                                                \
+    }                                                                        \
+    phi::prefix##Kernel<T, Context>(                                         \
+        dev_ctx, x.non_zero_elements(), out->mutable_non_zero_elements());   \
+  }
+
+#define DEFINE_SPARSE_UNARY_KERNEL_WITH_ONE_ATTR(prefix, attr)         \
+  template <typename T, typename Context>                              \
+  void prefix##CooKernel(const Context& dev_ctx,                       \
+                         const SparseCooTensor& x,                     \
+                         float attr,                                   \
+                         SparseCooTensor* out) {                       \
+    EmptyLikeCooKernel<T, Context>(dev_ctx, x, out);                   \
+    if (out->mutable_non_zero_elements()->numel() == 0) {              \
+      return;                                                          \
+    }                                                                  \
+    phi::prefix##Kernel<T, Context>(dev_ctx,                           \
+                                    x.non_zero_elements(),             \
+                                    attr,                              \
+                                    out->mutable_non_zero_elements()); \
+  }                                                                    \
+                                                                       \
+  template <typename T, typename Context>                              \
+  void prefix##CsrKernel(const Context& dev_ctx,                       \
+                         const SparseCsrTensor& x,                     \
+                         float attr,                                   \
+                         SparseCsrTensor* out) {                       \
+    EmptyLikeCsrKernel<T, Context>(dev_ctx, x, out);                   \
+    if (out->mutable_non_zero_elements()->numel() == 0) {              \
+      return;                                                          \
+    }                                                                  \
+    phi::prefix##Kernel<T, Context>(dev_ctx,                           \
+                                    x.non_zero_elements(),             \
+                                    attr,                              \
+                                    out->mutable_non_zero_elements()); \
+  }
+
+#define DEFINE_SPARSE_UNARY_KERNEL_WITH_COMPLEX(prefix)                      \
+  template <typename T, typename Context>                                    \
+  void prefix##CooKernel(const Context& dev_ctx,                             \
+                         const SparseCooTensor& x,                           \
+                         SparseCooTensor* out) {                             \
+    *(out->mutable_indices()) = x.indices();                                 \
+    DenseTensor* out_values = out->mutable_values();                         \
+    const DenseTensor& x_values = x.values();                                \
+    out_values->Resize(x_values.dims());                                     \
+    dev_ctx.template Alloc<T>(out_values);                                   \
+    if (out->mutable_non_zero_elements()->numel() != 0) {                    \
+      phi::prefix##Kernel<T, Context>(                                       \
+          dev_ctx, x.non_zero_elements(), out->mutable_non_zero_elements()); \
+    }                                                                        \
+    out->SetIndicesDict(x.GetIndicesDict());                                 \
+  }                                                                          \
+                                                                             \
+  template <typename T, typename Context>                                    \
+  void prefix##CsrKernel(const Context& dev_ctx,                             \
+                         const SparseCsrTensor& x,                           \
+                         SparseCsrTensor* out) {                             \
+    if (out && out->numel() == 0) {                                          \
+      dev_ctx.template Alloc<T>(out);                                        \
+      return;                                                                \
+    }                                                                        \
+    *(out->mutable_crows()) = x.crows();                                     \
+    *(out->mutable_cols()) = x.cols();                                       \
+    DenseTensor* out_values = out->mutable_values();                         \
+    const DenseTensor& x_values = x.values();                                \
+    out_values->Resize(x_values.dims());                                     \
+    dev_ctx.template Alloc<T>(out_values);                                   \
+    if (out->mutable_non_zero_elements()->numel() == 0) {                    \
+      return;                                                                \
+    }                                                                        \
+    phi::prefix##Kernel<T, Context>(                                         \
+        dev_ctx, x.non_zero_elements(), out->mutable_non_zero_elements());   \
+  }
+
+DEFINE_SPARSE_UNARY_KERNEL(Sqrt)
+DEFINE_SPARSE_UNARY_KERNEL(Relu)
+DEFINE_SPARSE_UNARY_KERNEL(Relu6)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_ONE_ATTR(Pow, factor)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_ONE_ATTR(LeakyRelu, alpha)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_COMPLEX(Abs)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_COMPLEX(Sin)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_COMPLEX(Tan)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_COMPLEX(Sinh)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_COMPLEX(Asin)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_COMPLEX(Asinh)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_COMPLEX(Atan)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_COMPLEX(Atanh)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_COMPLEX(Expm1)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_COMPLEX(Log1p)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_COMPLEX(Square)
+DEFINE_SPARSE_UNARY_KERNEL_WITH_COMPLEX(Tanh)
+
+template <typename T, typename Context>
+void ScaleCooKernel(const Context& dev_ctx,
+                    const SparseCooTensor& x,
+                    float scale,
+                    float bias,
+                    bool bias_after_scale,
+                    SparseCooTensor* out) {
+  EmptyLikeCooKernel<T, Context>(dev_ctx, x, out);
+  phi::ScaleKernel<T, Context>(dev_ctx,
+                               x.non_zero_elements(),
+                               scale,
+                               bias,
+                               bias_after_scale,
+                               out->mutable_non_zero_elements());
+  out->SetIndicesDict(x.GetIndicesDict());
+  out->SetKmaps(x.GetKmaps());
+}
+
+template <typename T, typename Context>
+void ScaleCsrKernel(const Context& dev_ctx,
+                    const SparseCsrTensor& x,
+                    float scale,
+                    float bias,
+                    bool bias_after_scale,
+                    SparseCsrTensor* out) {
+  EmptyLikeCsrKernel<T, Context>(dev_ctx, x, out);
+  phi::ScaleKernel<T, Context>(dev_ctx,
+                               x.non_zero_elements(),
+                               scale,
+                               bias,
+                               bias_after_scale,
+                               out->mutable_non_zero_elements());
+}
+
+template <typename T, typename Context>
+void CastCooKernel(const Context& dev_ctx,
+                   const SparseCooTensor& x,
+                   DataType index_dtype,
+                   DataType value_dtype,
+                   SparseCooTensor* out) {
+  const DenseTensor& x_indices = x.indices();
+  const DenseTensor& x_values = x.non_zero_elements();
+  DenseTensor* out_indices = out->mutable_indices();
+  DenseTensor* out_values = out->mutable_non_zero_elements();
+
+  if (index_dtype == DataType::UNDEFINED) {
+    *out_indices = x_indices;
+  } else {
+    MetaTensor meta(out_indices);
+    meta.set_dims(x_indices.dims());
+    meta.set_dtype(index_dtype);
+
+    PD_VISIT_INTEGRAL_TYPES(x_indices.dtype(), "CastCooKernel", [&] {
+      CastKernel<data_t, Context>(dev_ctx, x_indices, index_dtype, out_indices);
+    });
+  }
+
+  if (value_dtype == DataType::UNDEFINED) {
+    phi::Copy(dev_ctx, x_values, dev_ctx.GetPlace(), false, out_values);
+  } else {
+    MetaTensor meta(out_values);
+    meta.set_dims(x_values.dims());
+    CastKernel<T, Context>(dev_ctx, x_values, value_dtype, out_values);
+  }
+  out->SetIndicesDict(x.GetIndicesDict());
+  out->SetKmaps(x.GetKmaps());
+}
+
+template <typename T, typename Context>
+void CastCsrKernel(const Context& dev_ctx,
+                   const SparseCsrTensor& x,
+                   DataType index_dtype,
+                   DataType value_dtype,
+                   SparseCsrTensor* out) {
+  const DenseTensor& x_crows = x.crows();
+  const DenseTensor& x_cols = x.cols();
+  const DenseTensor& x_values = x.non_zero_elements();
+  DenseTensor* out_crows = out->mutable_crows();
+  DenseTensor* out_cols = out->mutable_cols();
+  DenseTensor* out_values = out->mutable_non_zero_elements();
+
+  if (index_dtype == DataType::UNDEFINED) {
+    *out_crows = x_crows;
+    *out_cols = x_cols;
+  } else {
+    MetaTensor crows_meta(out_crows);
+    crows_meta.set_dims(x_crows.dims());
+    crows_meta.set_dtype(index_dtype);
+
+    PD_VISIT_INTEGRAL_TYPES(x_crows.dtype(), "CastCsrKernel", [&] {
+      CastKernel<data_t, Context>(dev_ctx, x_crows, index_dtype, out_crows);
+    });
+
+    MetaTensor cols_meta(out_cols);
+    cols_meta.set_dims(x_cols.dims());
+    cols_meta.set_dtype(index_dtype);
+
+    PD_VISIT_INTEGRAL_TYPES(x_cols.dtype(), "CastCsrKernel", [&] {
+      CastKernel<data_t, Context>(dev_ctx, x_cols, index_dtype, out_cols);
+    });
+  }
+
+  if (value_dtype == DataType::UNDEFINED) {
+    phi::Copy(dev_ctx, x_values, dev_ctx.GetPlace(), false, out_values);
+  } else {
+    MetaTensor meta(out_values);
+    meta.set_dims(x_values.dims());
+    CastKernel<T, Context>(dev_ctx, x_values, value_dtype, out_values);
+  }
+}
+
+template <typename T, typename Context>
+void IsnanCooKernel(const Context& dev_ctx,
+                    const SparseCooTensor& x,
+                    SparseCooTensor* out) {
+  *(out->mutable_indices()) = x.indices();
+  const DenseTensor& x_values = x.non_zero_elements();
+  DenseTensor* out_values = out->mutable_non_zero_elements();
+
+  MetaTensor meta(out_values);
+  meta.set_dims(x_values.dims());
+  meta.set_dtype(DataType::BOOL);
+
+  phi::IsnanKernel<T, Context>(
+      dev_ctx, x.non_zero_elements(), out->mutable_non_zero_elements());
+  out->SetIndicesDict(x.GetIndicesDict());
+  out->SetKmaps(x.GetKmaps());
+}
+
+template <typename T, typename Context>
+void IsnanCsrKernel(const Context& dev_ctx,
+                    const SparseCsrTensor& x,
+                    SparseCsrTensor* out) {
+  const DenseTensor& x_crows = x.crows();
+  const DenseTensor& x_cols = x.cols();
+  const DenseTensor& x_values = x.non_zero_elements();
+  DenseTensor* out_crows = out->mutable_crows();
+  DenseTensor* out_cols = out->mutable_cols();
+  DenseTensor* out_values = out->mutable_non_zero_elements();
+
+  *out_crows = x_crows;
+  *out_cols = x_cols;
+
+  MetaTensor meta(out_values);
+  meta.set_dims(x_values.dims());
+  meta.set_dtype(DataType::BOOL);
+
+  phi::IsnanKernel<T, Context>(
+      dev_ctx, x.non_zero_elements(), out->mutable_non_zero_elements());
+}
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/mask_grad_kernel.h

@@ -0,0 +1,45 @@
+/* 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 "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+#include "paddle/phi/kernels/sparse/mask_kernel.h"
+#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context>
+void MaskAsCooGradKernel(const Context& dev_ctx,
+                         const DenseTensor& x,
+                         const SparseCooTensor& mask,
+                         const SparseCooTensor& out_grad,
+                         DenseTensor* x_grad) {
+  CooToDenseKernel<T, Context>(dev_ctx, out_grad, x_grad);
+}
+
+template <typename T, typename Context>
+void MaskAsCsrGradKernel(const Context& dev_ctx,
+                         const DenseTensor& x,
+                         const SparseCsrTensor& mask,
+                         const SparseCsrTensor& out_grad,
+                         DenseTensor* x_grad) {
+  CsrToDenseKernel<T, Context>(dev_ctx, out_grad, x_grad);
+}
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/mask_kernel.h

@@ -0,0 +1,43 @@
+/* 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 "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename Context>
+void MaskHelperCooKernel(const Context& dev_ctx,
+                         const SparseCooTensor& x,
+                         const DenseTensor& mask_indices,
+                         DenseTensor* out);
+
+template <typename T, typename Context>
+void MaskAsCooKernel(const Context& dev_ctx,
+                     const DenseTensor& x,
+                     const SparseCooTensor& mask,
+                     SparseCooTensor* out);
+
+template <typename T, typename Context>
+void MaskAsCsrKernel(const Context& dev_ctx,
+                     const DenseTensor& x,
+                     const SparseCsrTensor& mask,
+                     SparseCsrTensor* out);
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/matmul_grad_kernel.h

@@ -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. */
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+
+namespace phi {
+namespace sparse {
+
+// TODO(zhouwei25): implement Backward of " COO @ COO -> COO"
+template <typename T, typename Context>
+void MatmulCooCooGradKernel(const Context& dev_ctx,
+                            const SparseCooTensor& x,
+                            const SparseCooTensor& y,
+                            const SparseCooTensor& dout,
+                            SparseCooTensor* dx,
+                            SparseCooTensor* dy);
+
+// Backward of " COO @ DENSE -> DENSE"
+template <typename T, typename Context>
+void MatmulCooDenseGradKernel(const Context& dev_ctx,
+                              const SparseCooTensor& x,
+                              const DenseTensor& y,
+                              const DenseTensor& dout,
+                              SparseCooTensor* dx,
+                              DenseTensor* dy);
+
+// TODO(zhouwei25): implement Backward of " CSR @ CSR -> CSR"
+template <typename T, typename Context>
+void MatmulCsrCsrGradKernel(const Context& dev_ctx,
+                            const SparseCsrTensor& x,
+                            const SparseCsrTensor& y,
+                            const SparseCsrTensor& dout,
+                            SparseCsrTensor* dx,
+                            SparseCsrTensor* dy);
+
+/* Backward of "CSR @ DENSE -> DENSE" */
+template <typename T, typename Context>
+void MatmulCsrDenseGradKernel(const Context& dev_ctx,
+                              const SparseCsrTensor& x,
+                              const DenseTensor& y,
+                              const DenseTensor& dout,
+                              SparseCsrTensor* dx,
+                              DenseTensor* dy);
+
+/* Backward of "DENSE @ DENSE * CSR_MASK -> CSR" */
+template <typename T, typename Context>
+void MaskedMatmulCsrGradKernel(const Context& dev_ctx,
+                               const DenseTensor& x,
+                               const DenseTensor& y,
+                               const SparseCsrTensor& dout,
+                               DenseTensor* dx,
+                               DenseTensor* dy);
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/matmul_kernel.h

@@ -0,0 +1,61 @@
+/* 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 "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+
+namespace phi {
+namespace sparse {
+
+/* COO @ COO -> COO */
+template <typename T, typename Context>
+void MatmulCooCooKernel(const Context& dev_ctx,
+                        const SparseCooTensor& x,
+                        const SparseCooTensor& y,
+                        SparseCooTensor* out);
+
+/* COO @ DENSE -> DENSE */
+template <typename T, typename Context>
+void MatmulCooDenseKernel(const Context& dev_ctx,
+                          const SparseCooTensor& x,
+                          const DenseTensor& y,
+                          DenseTensor* out);
+
+/* CSR @ CSR -> CSR */
+template <typename T, typename Context>
+void MatmulCsrCsrKernel(const Context& dev_ctx,
+                        const SparseCsrTensor& x,
+                        const SparseCsrTensor& y,
+                        SparseCsrTensor* out);
+
+/* CSR @ DENSE -> DENSE */
+template <typename T, typename Context>
+void MatmulCsrDenseKernel(const Context& dev_ctx,
+                          const SparseCsrTensor& x,
+                          const DenseTensor& y,
+                          DenseTensor* out);
+
+/* DENSE @ DENSE * CSR_MASK -> CSR */
+template <typename T, typename Context>
+void MaskedMatmulCsrKernel(const Context& dev_ctx,
+                           const DenseTensor& x,
+                           const DenseTensor& y,
+                           const SparseCsrTensor& mask,
+                           SparseCsrTensor* out);
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/mv_grad_kernel.h

@@ -0,0 +1,43 @@
+/* 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 "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+
+namespace phi {
+namespace sparse {
+
+/* backward of COO @ DENSE VEC -> DENSE VEC */
+template <typename T, typename Context>
+void MvCooGradKernel(const Context& dev_ctx,
+                     const SparseCooTensor& x,
+                     const DenseTensor& vec,
+                     const DenseTensor& dout,
+                     SparseCooTensor* dx,
+                     DenseTensor* dvec);
+
+/* backward of CSR @ DENSE VEC -> DENSE VEC */
+template <typename T, typename Context>
+void MvCsrGradKernel(const Context& dev_ctx,
+                     const SparseCsrTensor& x,
+                     const DenseTensor& vec,
+                     const DenseTensor& dout,
+                     SparseCsrTensor* dx,
+                     DenseTensor* dvec);
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/mv_kernel.h

@@ -0,0 +1,39 @@
+/* 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 "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+
+namespace phi {
+namespace sparse {
+
+/* COO @ DENSE VEC -> DENSE VEC */
+template <typename T, typename Context>
+void MvCooKernel(const Context& dev_ctx,
+                 const SparseCooTensor& x,
+                 const DenseTensor& vec,
+                 DenseTensor* out);
+
+/* CSR @ DENSE VEC -> DENSE VEC */
+template <typename T, typename Context>
+void MvCsrKernel(const Context& dev_ctx,
+                 const SparseCsrTensor& x,
+                 const DenseTensor& vec,
+                 DenseTensor* out);
+
+}  // namespace sparse
+}  // namespace phi

paddle/phi/kernels/sparse/pool_grad_kernel.h

@@ -0,0 +1,49 @@
+/* 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 "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+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);
+
+template <typename T, typename Context>
+SparseCooTensor MaxPoolCooGrad(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;
+  MaxPoolCooGradKernel<T, Context>(
+      dev_ctx, x, rulebook, counter, out, out_grad, kernel_sizes, &x_grad);
+  return x_grad;
+}
+
+}  // namespace sparse
+}  // namespace phi