chore: import upstream snapshot with attribution

paddle/phi/kernels/funcs/sparse/common_shape.h

@@ -0,0 +1,84 @@
+/* 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 <stdint.h>
+
+#include "paddle/common/ddim.h"
+
+namespace phi {
+namespace funcs {
+namespace sparse {
+
+inline const DDim InferDenseDims(const DDim& x_dims,
+                                 const int64_t sparse_dim,
+                                 const int64_t non_zero_num) {
+  auto dense_dim = x_dims.size() - sparse_dim;
+  DDim values_dims;
+  if (dense_dim > 0) {
+    std::vector<int64_t> dense_dim_vec(dense_dim + 1);
+    dense_dim_vec[0] = non_zero_num;
+    memcpy(&dense_dim_vec[1],
+           x_dims.Get() + sparse_dim,
+           dense_dim * sizeof(x_dims[0]));
+    values_dims = make_ddim(dense_dim_vec);
+  } else {
+    values_dims = make_ddim({non_zero_num});
+  }
+  return values_dims;
+}
+
+template <typename IntT>
+inline const IntT HOSTDEVICE IndicesToIndex(const IntT* indices,
+                                            const IntT* sparse_offsets,
+                                            const int64_t non_zero_num,
+                                            const int64_t sparse_dim,
+                                            const int i) {
+  IntT index = 0;
+  for (IntT j = 0; j < sparse_dim; j++) {
+    index += indices[j * non_zero_num + i] * sparse_offsets[j];
+  }
+  return index;
+}
+
+template <typename IntT>
+inline void HOSTDEVICE FlattenIndices(const IntT* indices,
+                                      const IntT* sparse_offsets,
+                                      const int64_t non_zero_num,
+                                      const int64_t sparse_dim,
+                                      const int start,
+                                      const int stride,
+                                      IntT* out) {
+  for (int i = start; i < non_zero_num; i += stride) {
+    out[i] =
+        IndicesToIndex(indices, sparse_offsets, non_zero_num, sparse_dim, i);
+  }
+}
+
+// 1. indices.dims().size() == 2
+template <typename IntT>
+inline void CalcOffsetsPerDim(const DDim& dims,
+                              const int64_t sparse_dim,
+                              std::vector<IntT>* offsets) {
+  IntT offset = 1;
+  for (IntT i = sparse_dim - 1; i >= 0; i--) {
+    (*offsets)[i] = offset;
+    offset *= dims[i];
+  }
+}
+
+}  // namespace sparse
+}  // namespace funcs
+}  // namespace phi

paddle/phi/kernels/funcs/sparse/convolution.h

@@ -0,0 +1,264 @@
+/* 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/common/ddim.h"
+#include "paddle/phi/core/kmap_cache.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+
+#if !defined(PADDLE_WITH_CUDA) || !defined(PADDLE_WITH_CUSTOM_DEVICE)
+#include "paddle/phi/kernels/funcs/sparse/convolution_blas.h"
+#endif
+
+namespace phi {
+namespace funcs {
+namespace sparse {
+
+struct Dims4D {
+  int dims[4];
+  Dims4D(const int batch, const int x, const int y, const int z) {
+    dims[0] = batch;
+    dims[1] = z;
+    dims[2] = y;
+    dims[3] = x;
+  }
+  HOSTDEVICE const int& operator[](int i) const { return dims[i]; }
+};
+
+// Judge whether the current position x is in (lower, upper)
+template <typename IntT = int>
+inline HOSTDEVICE bool Check(const IntT& x,
+                             const int& kx,
+                             const int& pad,
+                             const int& stride,
+                             const int dilation,
+                             const int kdim,
+                             const int xdim) {
+  const IntT lower = x - dilation * kx + pad;
+  const IntT upper = x + (kdim - kx - 1) * dilation - pad;
+  return (lower >= 0 && lower % stride == 0 && upper < xdim);
+}
+
+// Check whether the current position(x, y, z) is legal:
+// Judge the minimum and maximum values at each latitude
+template <typename IntT = int>
+inline HOSTDEVICE bool Check(const Dims4D& dims,
+                             const Dims4D& kernel_dims,
+                             const Dims4D& paddings,
+                             const Dims4D& dilations,
+                             const Dims4D& strides,
+                             const IntT x,
+                             const IntT y,
+                             const IntT z,
+                             const int kx,
+                             const int ky,
+                             const int kz) {
+  bool x_valid = Check(
+      x, kx, paddings[3], strides[3], dilations[3], kernel_dims[3], dims[3]);
+  bool y_valid = Check(
+      y, ky, paddings[2], strides[2], dilations[2], kernel_dims[2], dims[2]);
+  bool z_valid = Check(
+      z, kz, paddings[1], strides[1], dilations[1], kernel_dims[1], dims[1]);
+  return (x_valid && y_valid && z_valid);
+}
+
+template <typename Dim, typename IntT = int>
+inline HOSTDEVICE IntT PointToIndex(const IntT& batch,
+                                    const IntT& x,
+                                    const IntT& y,
+                                    const IntT& z,
+                                    const Dim& dims) {
+  return batch * dims[1] * dims[2] * dims[3] + z * dims[2] * dims[3] +
+         y * dims[3] + x;
+}
+
+// TODO(zhangkaihuo): use division and multiply to optimize
+// modulo operation
+template <typename Dim, typename IntT = int>
+inline HOSTDEVICE void IndexToPoint(
+    const IntT index, const Dim& dims, IntT* batch, IntT* x, IntT* y, IntT* z) {
+  IntT n = index;
+  *x = n % dims[3];
+  n /= dims[3];
+  *y = n % dims[2];
+  n /= dims[2];
+  *z = n % dims[1];
+  n /= dims[1];
+  *batch = n;
+}
+
+inline void GetOutShape(const DDim& x_dims,
+                        const std::vector<int>& kernel_sizes,
+                        const std::vector<int>& paddings,
+                        const std::vector<int>& dilations,
+                        const std::vector<int>& strides,
+                        DDim* out_dims) {
+  const bool is2D = out_dims->size() == 4 ? true : false;
+  if (is2D) {
+    PADDLE_ENFORCE_EQ(x_dims.size(),
+                      4,
+                      common::errors::InvalidArgument(
+                          "the shape of x should be (N, H, W, C)"));
+    PADDLE_ENFORCE_EQ(kernel_sizes.size(),
+                      4,
+                      common::errors::InvalidArgument(
+                          "the shape of kernel should be (H, W, C, OC)"));
+
+    // infer out shape
+    (*out_dims)[0] = x_dims[0];
+    (*out_dims)[3] = kernel_sizes[3];
+    for (int i = 1; i < 3; i++) {
+      (*out_dims)[i] = (x_dims[i] + 2 * paddings[i - 1] -
+                        dilations[i - 1] * (kernel_sizes[i - 1] - 1) - 1) /
+                           strides[i - 1] +
+                       1;
+    }
+  } else {
+    PADDLE_ENFORCE_EQ(x_dims.size(),
+                      5,
+                      common::errors::InvalidArgument(
+                          "the shape of x should be (N, D, H, W, C)"));
+    PADDLE_ENFORCE_EQ(kernel_sizes.size(),
+                      5,
+                      common::errors::InvalidArgument(
+                          "the shape of kernel should be (D, H, W, C, OC)"));
+
+    // infer out shape
+    (*out_dims)[0] = x_dims[0];
+    (*out_dims)[4] = kernel_sizes[4];
+    for (int i = 1; i < 4; i++) {
+      (*out_dims)[i] = (x_dims[i] + 2 * paddings[i - 1] -
+                        dilations[i - 1] * (kernel_sizes[i - 1] - 1) - 1) /
+                           strides[i - 1] +
+                       1;
+    }
+  }
+}
+
+inline void ResetSubmKernelSizeAndStrides(const DDim& kernel_dims,
+                                          std::vector<int>* paddings,
+                                          std::vector<int>* strides) {
+  for (uint64_t i = 0; i < paddings->size(); i++) {
+    (*paddings)[i] = kernel_dims[i] / 2;
+    (*strides)[i] = 1;
+  }
+}
+
+inline const std::vector<int> PoolResetKernel(
+    const std::vector<int>& kernel_sizes,
+    const int in_channels,
+    const int out_channels) {
+  std::vector<int> res(kernel_sizes);
+  res.resize(5);
+  res[3] = in_channels;
+  res[4] = out_channels;
+  return res;
+}
+
+template <typename T>
+inline void PrefixSum(const T* counter, T* offsets, const int n) {
+  T offset = 0;
+  for (int i = 0; i < n; i++) {
+    offsets[i] = offset;
+    offset += counter[i];
+  }
+  offsets[n] = offset;
+}
+
+template <typename IntT>
+inline const IntT* GetRulebookPtr(const SparseCooTensor& coo,
+                                  const DenseTensor& rulebook,
+                                  const std::string& key,
+                                  int* rulebook_len) {
+  if (!key.empty()) {
+    const auto* indices_pairs = coo.IndicesPairs(key);
+    if (indices_pairs != nullptr) {
+      const DenseTensor& tmp_rulebook = indices_pairs->first;
+      *rulebook_len = tmp_rulebook.dims()[1];
+      return tmp_rulebook.data<IntT>();
+    }
+  }
+  *rulebook_len = rulebook.dims()[1];
+  return rulebook.data<IntT>();
+}
+
+inline const int* GetCounterPtr(const SparseCooTensor& coo,
+                                const DenseTensor& counter,
+                                const std::string& key) {
+  if (!key.empty()) {
+    const auto* indices_pairs = coo.IndicesPairs(key);
+    if (indices_pairs != nullptr) {
+      return indices_pairs->second.data<int>();
+    }
+  }
+  return counter.data<int>();
+}
+
+template <typename T, typename IntT, typename Context>
+inline const IntT* PrepareSubm(const Context& dev_ctx,
+                               const SparseCooTensor& x,
+                               const std::string& key,
+                               const DDim& out_dims,
+                               SparseCooTensor* out,
+                               int* counter,
+                               int* offsets,
+                               int* rulebook_len,
+                               bool* need_product_rulebook) {
+  const auto* indices_pairs = x.IndicesPairs(key);
+  if (indices_pairs != nullptr) {
+    *need_product_rulebook = false;
+    const DenseTensor& rulebook = indices_pairs->first;
+    const int64_t counter_size = indices_pairs->second.numel();
+    memcpy(
+        counter, indices_pairs->second.data<int>(), counter_size * sizeof(int));
+    out->SetIndicesDict(x.GetIndicesDict());
+
+    *rulebook_len = rulebook.dims()[1];
+
+    DenseTensor out_indices = EmptyLike<IntT>(dev_ctx, x.non_zero_indices());
+    DenseTensor out_values = EmptyLike<T>(dev_ctx, x.non_zero_elements());
+    phi::Copy(
+        dev_ctx, x.non_zero_indices(), dev_ctx.GetPlace(), false, &out_indices);
+    out->SetMember(out_indices, out_values, out_dims, false);
+    PrefixSum<int>(counter, offsets, counter_size);
+    return rulebook.data<IntT>();
+  }
+  return nullptr;
+}
+
+template <typename Context>
+inline void SaveToTable(const Context& dev_ctx,
+                        const SparseCooTensor& x,
+                        const std::string& key,
+                        const DenseTensor& in_rulebook,
+                        const DenseTensor& h_counter,
+                        SparseCooTensor* out,
+                        DenseTensor* out_rulebook,
+                        DenseTensor* counter) {
+  out->SetIndicesDict(x.GetIndicesDict());
+  if (!key.empty()) {
+    out->SaveIndicesPairs(key, std::make_pair(in_rulebook, h_counter));
+  } else {
+    *out_rulebook = in_rulebook;
+    counter->Resize({h_counter.numel()});
+    int* counter_ptr = dev_ctx.template HostAlloc<int>(counter);
+    memcpy(counter_ptr, h_counter.data<int>(), h_counter.numel() * sizeof(int));
+  }
+}
+
+}  // namespace sparse
+}  // namespace funcs
+}  // namespace phi

paddle/phi/kernels/funcs/sparse/convolution_blas.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/common/ddim.h"
+#include "paddle/phi/core/kmap_cache.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+
+namespace phi {
+namespace funcs {
+namespace sparse {
+
+template <typename T, typename Context>
+inline void SubmPreProcess(const Context& dev_ctx,
+                           const SparseCooTensor& x,
+                           const DenseTensor& kernel,
+                           const DenseTensor& out_grad,
+                           const int in_channels,
+                           const int out_channels,
+                           const int half_kernel_size,
+                           DenseTensor* kernel_grad,
+                           DenseTensor* x_grad) {
+  auto blas = funcs::GetBlas<Context, T>(dev_ctx);
+  const bool is_params_freezing = kernel_grad == nullptr;
+  if (!is_params_freezing) {
+    T* d_kernel_ptr = kernel_grad->data<T>();
+    blas.GEMM(CblasTrans,
+              CblasNoTrans,
+              x.non_zero_elements().dims()[1],
+              out_grad.dims()[1],
+              x.non_zero_elements().dims()[0],
+              static_cast<T>(1),
+              x.non_zero_elements().data<T>(),
+              out_grad.data<T>(),
+              static_cast<T>(0),
+              d_kernel_ptr + half_kernel_size * in_channels * out_channels);
+  }
+
+  // call gemm: d_x = out_grad * transpose(kernel)
+  // (n, out_channels) * (out_channels, in_channels)
+  T* x_grad_ptr = x_grad->data<T>();
+  blas.GEMM(CblasNoTrans,
+            CblasTrans,
+            out_grad.dims()[0],
+            in_channels,
+            out_grad.dims()[1],
+            static_cast<T>(1),
+            out_grad.data<T>(),
+            kernel.data<T>() + half_kernel_size * in_channels * out_channels,
+            static_cast<T>(0),
+            x_grad_ptr);
+}
+
+}  // namespace sparse
+}  // namespace funcs
+}  // namespace phi

paddle/phi/kernels/funcs/sparse/flatten_indices.cu.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/kernels/funcs/sparse/flatten_indices.h"
+
+namespace phi {
+namespace funcs {
+namespace sparse {
+
+template <typename IntT>
+__global__ void FlattenIndicesKernel(const IntT* indices,
+                                     const IntT* sparse_offsets,
+                                     const int64_t non_zero_num,
+                                     const int64_t sparse_dim,
+                                     IntT* out) {
+  int64_t tid =
+      static_cast<int64_t>(threadIdx.x) +
+      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x);
+  funcs::sparse::FlattenIndices<IntT>(indices,
+                                      sparse_offsets,
+                                      non_zero_num,
+                                      sparse_dim,
+                                      tid,
+                                      gridDim.x * blockDim.x,
+                                      out);
+}
+
+template <typename IntT>
+__global__ void IndexToCoordinateKernel(const IntT* index,
+                                        const Dim<DDim::kMaxRank> dims,
+                                        const int64_t non_zero_num,
+                                        const int64_t sparse_dim,
+                                        IntT* indices) {
+  int64_t tid =
+      static_cast<int64_t>(threadIdx.x) +
+      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x);
+  IndexToCoordinate(index,
+                    dims,
+                    non_zero_num,
+                    sparse_dim,
+                    tid,
+                    gridDim.x * blockDim.x,
+                    indices);
+}
+
+}  // namespace sparse
+}  // namespace funcs
+}  // namespace phi

paddle/phi/kernels/funcs/sparse/flatten_indices.h

@@ -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. */
+
+#pragma once
+
+#include <stdint.h>
+
+#include "paddle/common/ddim.h"
+
+namespace phi {
+namespace funcs {
+namespace sparse {
+
+template <typename IntT>
+inline const IntT HOSTDEVICE CoordinateToIndex(const IntT* indices,
+                                               const IntT* sparse_offsets,
+                                               const int64_t non_zero_num,
+                                               const int64_t sparse_dim,
+                                               const int i) {
+  IntT index = 0;
+  for (IntT j = 0; j < sparse_dim; j++) {
+    index += indices[j * non_zero_num + i] * sparse_offsets[j];
+  }
+  return index;
+}
+
+template <typename IntT>
+inline void HOSTDEVICE FlattenIndices(const IntT* indices,
+                                      const IntT* sparse_offsets,
+                                      const int64_t non_zero_num,
+                                      const int64_t sparse_dim,
+                                      const int64_t start,
+                                      const int64_t stride,
+                                      IntT* out) {
+  for (int64_t i = start; i < non_zero_num; i += stride) {
+    out[i] =
+        CoordinateToIndex(indices, sparse_offsets, non_zero_num, sparse_dim, i);
+  }
+}
+
+// 1. indices.dims().size() == 2
+template <typename IntT>
+inline void CalcOffsetsPerDim(const DDim& dims,
+                              const int64_t sparse_dim,
+                              IntT* offsets) {
+  IntT offset = 1;
+  for (IntT i = sparse_dim - 1; i >= 0; i--) {
+    offsets[i] = offset;
+    offset *= dims[i];
+  }
+}
+
+template <typename IntT>
+inline void HOSTDEVICE IndexToCoordinate(const IntT index,
+                                         const Dim<DDim::kMaxRank>& dims,
+                                         const int64_t non_zero_num,
+                                         const int64_t sparse_dim,
+                                         const int indices_offset,
+                                         IntT* indices) {
+  IntT tmp_index = index;
+  for (int j = sparse_dim - 1; j >= 0; j--) {
+    indices[j * non_zero_num + indices_offset] = tmp_index % dims[j];
+    tmp_index /= dims[j];
+  }
+}
+
+template <typename IntT>
+inline void HOSTDEVICE IndexToCoordinate(const IntT* index,
+                                         const Dim<DDim::kMaxRank>& dims,
+                                         const int64_t non_zero_num,
+                                         const int64_t sparse_dim,
+                                         const int64_t start,
+                                         const int64_t stride,
+                                         IntT* indices) {
+  for (int64_t i = start; i < non_zero_num; i += stride) {
+    IntT tmp_index = index[i];
+    IndexToCoordinate(tmp_index, dims, non_zero_num, sparse_dim, i, indices);
+  }
+}
+
+}  // namespace sparse
+}  // namespace funcs
+}  // namespace phi

paddle/phi/kernels/funcs/sparse/scatter.cu.h

@@ -0,0 +1,162 @@
+/* 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/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+#include "paddle/phi/kernels/funcs/aligned_vector.h"
+
+#define VecBytes 16
+
+namespace phi {
+namespace funcs {
+namespace sparse {
+
+/**
+ * brief: scatter add
+ * input: the inputs
+ * unique_value: refer to UpdateIndexKernel notes
+ * out_index: the output feature index
+ * non_zero_num: the number of output features
+ * rulebook_len: the length of rulebook
+ * channels: the output channel size
+ * out: the outputs
+ **/
+template <typename T, int VecSize>
+__global__ void ScatterKernel(const T* input,
+                              const int* unique_value,
+                              const int* out_index,
+                              const int non_zero_num,
+                              const int rulebook_len,
+                              const int channels,
+                              T* out) {
+  int64_t tid =
+      static_cast<int64_t>(threadIdx.x) +
+      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(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;
+
+    int start = unique_value[indices_i];
+    int end = indices_i == non_zero_num - 1 ? rulebook_len
+                                            : unique_value[indices_i + 1];
+    // max(end-start) = kernel_size
+    StoreT sums = {static_cast<T>(0)};
+    for (int j = start; j < end; j++) {
+      const int out_feature_i = out_index[j];
+      LoadT vec_in;
+      Load<T, VecSize>(input + out_feature_i * channels + channels_i * VecSize,
+                       &vec_in);
+#pragma unroll
+      for (int k = 0; k < VecSize; k++) {
+        sums[k] += vec_in[k];
+      }
+    }
+    Store<T, VecSize>(sums, out + indices_i * channels + channels_i * VecSize);
+  }
+}
+
+// scatter's index has been grouped in advance
+// index_counts record the count of each group
+// index_groups save the index of each group
+template <typename T, int VecSize>
+__global__ void ScatterKernelV2(const T* input,
+                                const int* index_counts,
+                                const int* index_groups,
+                                const int non_zero_num,
+                                const int kernel_size,
+                                const int channels,
+                                const int buffer_counts,
+                                T* out) {
+  int64_t tid =
+      static_cast<int64_t>(threadIdx.x) +
+      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(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;
+
+    StoreT sums = {static_cast<T>(0)};
+    Load<T, VecSize>(out + indices_i * channels + channels_i * VecSize, &sums);
+    for (int it = 0; it < buffer_counts; it++) {
+      int len = index_counts[indices_i + it * non_zero_num];
+      const int group_offset = it * kernel_size * non_zero_num;
+      for (int j = 0; j < len; j++) {
+        const int out_feature_i =
+            index_groups[indices_i * kernel_size + j + group_offset];
+        LoadT vec_in;
+        Load<T, VecSize>(
+            input + out_feature_i * channels + channels_i * VecSize, &vec_in);
+#pragma unroll
+        for (int k = 0; k < VecSize; k++) {
+          sums[k] += vec_in[k];
+        }
+      }
+    }
+    Store<T, VecSize>(sums, out + indices_i * channels + channels_i * VecSize);
+  }
+}
+
+template <typename T>
+void ScatterV2(const GPUContext& dev_ctx,
+               const T* input,
+               const int* index_counts,
+               const int* index_groups,
+               const int non_zero_num,
+               const int kernel_size,
+               const int channels,
+               const int buffer_counts,
+               T* output) {
+  const int VecSize = VecBytes / sizeof(T);
+  if (channels % VecSize == 0) {
+    auto config = phi::backends::gpu::GetGpuLaunchConfig1D(
+        dev_ctx, non_zero_num * channels / VecSize, 1);
+    ScatterKernelV2<T, VecSize><<<config.block_per_grid.x,
+                                  config.thread_per_block.x,
+                                  0,
+                                  dev_ctx.stream()>>>(input,
+                                                      index_counts,
+                                                      index_groups,
+                                                      non_zero_num,
+                                                      kernel_size,
+                                                      channels,
+                                                      buffer_counts,
+                                                      output);
+  } else {
+    auto config = phi::backends::gpu::GetGpuLaunchConfig1D(
+        dev_ctx, non_zero_num * channels, 1);
+    ScatterKernelV2<T, 1><<<config.block_per_grid.x,
+                            config.thread_per_block.x,
+                            0,
+                            dev_ctx.stream()>>>(input,
+                                                index_counts,
+                                                index_groups,
+                                                non_zero_num,
+                                                kernel_size,
+                                                channels,
+                                                buffer_counts,
+                                                output);
+  }
+}
+
+}  // namespace sparse
+}  // namespace funcs
+}  // namespace phi

paddle/phi/kernels/funcs/sparse/softmax.cu.h

@@ -0,0 +1,203 @@
+/* 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. */
+
+#pragma once
+
+#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
+
+namespace phi {
+namespace funcs {
+namespace sparse {
+
+/* Given the indices of a sparse tensor, return a vector of offsets
+for the entries in the equivalent dense tensor. */
+template <typename IntT, typename Context>
+inline DenseTensor GetOffsets(const Context& dev_ctx,
+                              const DenseTensor& indices,
+                              const std::vector<IntT>& sizes,
+                              const IntT dim) {
+#ifdef __HIPCC__
+  const auto& policy = thrust::hip::par.on(dev_ctx.stream());
+#else
+  const auto& policy = thrust::cuda::par.on(dev_ctx.stream());
+#endif
+
+  auto ndim = indices.dims()[0];
+  auto nnz = indices.dims()[1];
+  std::vector<IntT> host_strides(ndim, 1);
+  if (ndim > 1) {
+    for (IntT i = ndim - 2; i >= 0; i--) {
+      host_strides[i] = host_strides[i + 1] * (i + 1 == dim ? 1 : sizes[i + 1]);
+    }
+  }
+
+  const IntArray strides_shape(vectorize<IntT>(indices.dims()));
+  DenseTensor strides = Empty<IntT>(dev_ctx, strides_shape);
+  auto strides_ptr = strides.data<IntT>();
+#if defined(__NVCC__) || defined(__HIPCC__)
+  const IntT* stable_st =
+      phi::backends::gpu::RestoreHostMemIfCapturingCUDAGraph(
+          host_strides.data(), host_strides.size());
+#else
+  const IntT* stable_st = host_strides.data();
+#endif
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     strides_ptr,
+                     CPUPlace(),
+                     stable_st,
+                     sizeof(IntT) * host_strides.size(),
+                     dev_ctx.stream());
+
+  DenseTensor offsets = Empty<IntT>(dev_ctx, {nnz});
+  auto indices_ptr = indices.data<IntT>();
+
+  thrust::transform(
+      policy,
+      thrust::make_counting_iterator(IntT(0)),
+      thrust::make_counting_iterator(IntT(nnz)),
+      thrust::device_ptr<IntT>(offsets.data<IntT>()),
+      [strides_ptr, indices_ptr, nnz, dim, ndim] __device__(IntT x) {
+        IntT pool_index = 0;
+        for (IntT j = 0; j < ndim; j++) {
+          if (j != dim) {
+            auto indice_cur_ptr = indices_ptr + j * nnz + x;
+            auto stride = strides_ptr[j];
+            pool_index += stride * (*indice_cur_ptr);
+          }
+        }
+        return pool_index;
+      });
+  return offsets;
+}
+
+/* Return pools of indices that align with the given dimension and the
+corresponding max values for each pool. */
+template <typename T,
+          typename IntT,
+          typename Context,
+          bool requireMxRows = true>
+std::tuple<DenseTensor, DenseTensor, DenseTensor, DenseTensor> ComputePoolMax(
+    const Context& dev_ctx,
+    const DenseTensor& indices,
+    const DenseTensor& values,
+    const std::vector<IntT>& sizes,
+    IntT nvalues,
+    const IntT dim) {
+#ifdef __HIPCC__
+  const auto& policy = thrust::hip::par.on(dev_ctx.stream());
+#else
+  const auto& policy = thrust::cuda::par.on(dev_ctx.stream());
+#endif
+  using thrust_ptr = thrust::device_ptr<IntT>;
+  auto nnz = indices.dims()[1];
+  DenseTensor offsets =
+      funcs::sparse::GetOffsets<IntT, Context>(dev_ctx, indices, sizes, dim);
+  auto offsets_ptr = offsets.data<IntT>();
+
+  DenseTensor sorted_indices = Empty<IntT>(dev_ctx, {nnz});
+  thrust_ptr sorted_indices_thrust_ptr(sorted_indices.data<IntT>());
+  thrust::sequence(
+      policy, sorted_indices_thrust_ptr, sorted_indices_thrust_ptr + nnz, 0);
+
+  /* sort indices corresponding to offsets */
+  thrust::sort(policy,
+               sorted_indices_thrust_ptr,
+               sorted_indices_thrust_ptr + nnz,
+               [offsets_ptr] __device__(IntT x, IntT y) {
+                 return offsets_ptr[x] < offsets_ptr[y];
+               });
+
+  DenseTensor pool_sizes = Empty<IntT>(dev_ctx, {nnz});
+
+  /* reduce the elements which are grouped by pool index,
+  returns all the pool indexes with unique offset value for each. */
+  auto new_end =
+      thrust::reduce_by_key(policy,
+                            sorted_indices_thrust_ptr,
+                            sorted_indices_thrust_ptr + nnz,
+                            thrust::make_constant_iterator(IntT(1)),
+                            thrust::make_discard_iterator(),
+                            thrust_ptr(pool_sizes.data<IntT>()),
+                            [offsets_ptr] __device__(IntT x, IntT y) {
+                              return offsets_ptr[x] == offsets_ptr[y];
+                            });
+  auto new_sz =
+      thrust::distance(thrust_ptr(pool_sizes.data<IntT>()), new_end.second);
+  pool_sizes.Resize({new_sz});
+
+  DenseTensor pool_offsets;
+  pool_offsets.Resize({new_sz});
+  dev_ctx.template Alloc<T>(&pool_offsets);
+  phi::Copy(dev_ctx, pool_sizes, dev_ctx.GetPlace(), false, &pool_offsets);
+
+  /* accumulate value for each pool index */
+  thrust_ptr pool_offsets_thrust_ptr(pool_offsets.data<IntT>());
+  thrust::exclusive_scan(policy,
+                         pool_offsets_thrust_ptr,
+                         pool_offsets_thrust_ptr + new_sz,
+                         pool_offsets_thrust_ptr);
+
+  DenseTensor mx_buffer;
+  if (requireMxRows) {
+    mx_buffer = phi::Full<T>(
+        dev_ctx, {new_sz * nvalues}, -std::numeric_limits<T>::infinity());
+    auto mx_buffer_ptr = mx_buffer.data<T>();
+
+    auto pool_sizes_ptr = pool_sizes.data<IntT>();
+    auto sorted_indices_ptr = sorted_indices.data<IntT>();
+    auto pool_offsets_ptr = pool_offsets.data<IntT>();
+    auto values_ptr = values.data<T>();
+
+    /* calculate max value in each pool. */
+    thrust::for_each(policy,
+                     thrust::make_counting_iterator(IntT(0)),
+                     thrust::make_counting_iterator(IntT(new_sz)),
+                     [sorted_indices_ptr,
+                      pool_sizes_ptr,
+                      pool_offsets_ptr,
+                      mx_buffer_ptr,
+                      values_ptr,
+                      nvalues] __device__(IntT index) {
+                       IntT curr_pool_size = pool_sizes_ptr[index];
+                       auto mx_row = mx_buffer_ptr + index * nvalues;
+                       IntT offset = pool_offsets_ptr[index];
+                       for (IntT p = 0; p < curr_pool_size; p++) {
+                         IntT i = *(sorted_indices_ptr + offset + p);
+                         for (IntT j = 0; j < nvalues; j++) {
+                           auto value_tmp = *(values_ptr);
+                           mx_row[j] = std::max(mx_row[j], value_tmp);
+                         }
+                       }
+                     });
+  }
+  return std::make_tuple(sorted_indices, pool_offsets, pool_sizes, mx_buffer);
+}
+
+inline int GetNumThreads(int nElem) {
+#if defined(PADLDE_WITH_ROCM)
+  int threadSizes[5] = {16, 32, 64, 128, 256};
+#else
+  int threadSizes[5] = {32, 64, 128, 256, 512};
+#endif
+  for (int i = 0; i != 5; ++i) {
+    if (nElem <= threadSizes[i]) {
+      return threadSizes[i];
+    }
+  }
+  return threadSizes[4];
+}
+
+}  // namespace sparse
+}  // namespace funcs
+}  // namespace phi

paddle/phi/kernels/funcs/sparse/softmax.h

@@ -0,0 +1,88 @@
+/* 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. */
+
+#pragma once
+
+#include "paddle/common/ddim.h"
+#include "paddle/phi/core/tensor_utils.h"
+
+namespace phi {
+namespace funcs {
+namespace sparse {
+
+template <typename IntT>
+inline void GetPoolsSoftmax(const DenseTensor& indices,
+                            const std::vector<IntT>& sizes,
+                            const int dim,
+                            std::map<IntT, std::vector<IntT>>* pools) {
+  auto ndim = indices.dims()[0];
+  auto nnz = indices.dims()[1];
+  std::vector<IntT> strides(ndim, 1);
+
+  if (ndim > 1) {
+    for (IntT i = ndim - 2; i >= 0; i--) {
+      strides[i] = strides[i + 1] * (i + 1 == dim ? 1 : sizes[i + 1]);
+    }
+  }
+
+  auto* indices_data = indices.data<IntT>();
+  for (IntT i = 0; i < nnz; i++) {
+    IntT pool_index = 0;
+    for (IntT j = 0; j < ndim; j++) {
+      if (j == dim) continue;
+      pool_index += strides[j] * indices_data[j * nnz + i];
+    }
+
+    if (pools->find(pool_index) == pools->end()) {
+      std::vector<IntT> vec;
+      (*pools)[pool_index] = vec;
+    }
+    (*pools)[pool_index].push_back(i);
+  }
+}
+
+template <typename IntT>
+inline std::vector<IntT> GetOffsets(const DenseTensor& indices,
+                                    const std::vector<IntT>& sizes,
+                                    const int dim) {
+  auto ndim = indices.dims()[0];
+  auto nnz = indices.dims()[1];
+  std::vector<IntT> offsets(nnz);
+  std::vector<IntT> strides(ndim, 1);
+  auto indices_ptr = indices.data<IntT>();
+
+  if (ndim > 1) {
+    for (IntT i = ndim - 2; i >= 0; i--) {
+      strides[i] = strides[i + 1] * sizes[i + 1];
+    }
+  }
+
+  for (int i = 0; i < nnz; i++) {
+    IntT acc = 0;
+    for (int j = 0; j < ndim; j++) {
+      auto indices_cur = indices_ptr + j * nnz + i;
+      auto stride = strides[j];
+      if (j != dim) {
+        acc += stride * (*indices_cur);
+      }
+    }
+    offsets[i] = acc;
+  }
+
+  return offsets;
+}
+
+}  // namespace sparse
+}  // namespace funcs
+}  // namespace phi

paddle/phi/kernels/funcs/sparse/sparse_blas.h

@@ -0,0 +1,115 @@
+//   Copyright (c) 2018 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/device_context.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+
+namespace phi {
+namespace funcs {
+namespace sparse {
+template <typename DeviceContext>
+class SparseBlas {
+ public:
+  explicit SparseBlas(const DeviceContext& dev_ctx) : dev_ctx_(dev_ctx) {}
+
+  template <typename T, typename TensorType>
+  void SPMM(bool transa,
+            bool transb,
+            T alpha,
+            const TensorType& mat_a,
+            const DenseTensor& mat_b,
+            T beta,
+            DenseTensor* mat_out) const;
+
+  template <typename T, typename TensorType>
+  void SPMV(bool transa,
+            T alpha,
+            const TensorType& mat_a,
+            const DenseTensor& vec_x,
+            T beta,
+            DenseTensor* vec_out) const;
+
+  template <typename T, typename TensorType>
+  void SDDMM(bool transa,
+             bool transb,
+             T alpha,
+             const DenseTensor& mat_a,
+             const DenseTensor& mat_b,
+             T beta,
+             TensorType* mat_out) const;
+
+  template <typename T>
+  void SPGEMM(bool transa,
+              bool transb,
+              T alpha,
+              const SparseCsrTensor& mat_a,
+              const SparseCsrTensor& mat_b,
+              T beta,
+              SparseCsrTensor* mat_out) const;
+
+ private:
+  const DeviceContext& dev_ctx_;
+};
+
+template <typename DeviceContext, typename T>
+class SparseBlasT : private SparseBlas<DeviceContext> {
+ public:
+  using SparseBlas<DeviceContext>::SparseBlas;
+
+  template <typename... ARGS>
+  void SPMM(ARGS... args) const {
+    Base()->template SPMM<T>(args...);
+  }
+
+  template <typename... ARGS>
+  void SPMV(ARGS... args) const {
+    Base()->template SPMV<T>(args...);
+  }
+
+  template <typename... ARGS>
+  void SDDMM(ARGS... args) const {
+    Base()->template SDDMM<T>(args...);
+  }
+
+  template <typename... ARGS>
+  void SPGEMM(ARGS... args) const {
+    Base()->template SPGEMM<T>(args...);
+  }
+
+ private:
+  const SparseBlas<DeviceContext>* Base() const {
+    return static_cast<const SparseBlas<DeviceContext>*>(this);
+  }
+};
+
+template <typename DeviceContext, typename T>
+inline SparseBlasT<DeviceContext, T> GetSparseBlas(
+    const DeviceContext& dev_ctx) {
+  return SparseBlasT<DeviceContext, T>(dev_ctx);
+}
+
+}  // namespace sparse
+}  // namespace funcs
+}  // namespace phi
+
+#if defined(PADDLE_WITH_CUDA)
+#include "paddle/phi/kernels/funcs/sparse/sparse_blas_impl.cu.h"
+#endif
+#if defined(PADDLE_WITH_HIP) && HIP_VERSION >= 402
+#include "paddle/phi/kernels/funcs/sparse/sparse_blas_impl.hip.h"
+#endif

paddle/phi/kernels/funcs/sparse/sparse_blas_impl.cu.h

@@ -0,0 +1,839 @@
+//   Copyright (c) 2018 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 "glog/logging.h"
+
+#include "paddle/common/ddim.h"
+#include "paddle/phi/backends/dynload/cusparse.h"
+#ifdef PADDLE_WITH_CUDA
+#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
+#endif
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+#include "paddle/phi/kernels/concat_kernel.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+
+namespace phi {
+namespace funcs {
+namespace sparse {
+
+template <typename T>
+cudaDataType_t GetGpuDataType() {
+  if (std::is_same<T, float>::value) {
+    return CUDA_R_32F;
+  } else if (std::is_same<T, double>::value) {
+    return CUDA_R_64F;
+  } else if (std::is_same<T, phi::float16>::value) {
+    return CUDA_R_16F;
+  }
+}
+
+template <typename T>
+cusparseIndexType_t GetCusparseIndexType() {
+  if (std::is_same<T, int32_t>::value) {
+    return CUSPARSE_INDEX_32I;
+  } else if (std::is_same<T, int64_t>::value) {
+    return CUSPARSE_INDEX_64I;
+  }
+}
+
+inline cusparseOperation_t GetTransposeOperation(const bool trans) {
+  if (trans) {
+    return CUSPARSE_OPERATION_TRANSPOSE;
+  } else {
+    return CUSPARSE_OPERATION_NON_TRANSPOSE;
+  }
+}
+
+inline cusparseSpMMAlg_t GetSpMMAlgorithm(const SparseCsrTensor& x) {
+  // TODO(zhouwei): will change to 'CUSPARSE_SPMM_CSR_ALG2' when support batch
+  return CUSPARSE_SPMM_CSR_ALG2;
+}
+
+inline cusparseSpMMAlg_t GetSpMMAlgorithm(const SparseCooTensor& x) {
+  return CUSPARSE_SPMM_ALG_DEFAULT;
+}
+
+/************* SPARSE MATRIX DESCRIPTOR (COO/CSR) ************/
+
+template <typename T, typename IntT>
+inline void CreateCsrDescriptor(const SparseCsrTensor& x,
+                                const GPUContext& dev_ctx,
+                                cusparseSpMatDescr_t* descriptor) {
+  std::vector<int64_t> xdim_vec = vectorize(x.dims());
+  auto x_ndims = xdim_vec.size();
+  PADDLE_ENFORCE_GE(
+      x_ndims,
+      2,
+      common::errors::InvalidArgument("the dim size of SparseCsrTensor must be "
+                                      "greater than or equal to 2."));
+  int64_t M = xdim_vec[x_ndims - 2];
+  int64_t N = xdim_vec[x_ndims - 1];
+  int batch_size = 1;
+  for (int i = 0; i < x_ndims - 2; i++) {
+    batch_size *= xdim_vec[i];
+  }
+  PADDLE_ENFORCE_EQ(x.non_zero_crows().numel(),
+                    batch_size * (M + 1),
+                    common::errors::PreconditionNotMet(
+                        "the length of SparseCsrTensor crows is not right."));
+
+  const IntT* crows_data = x.non_zero_crows().data<IntT>();
+  const IntT* cols_data = x.non_zero_cols().data<IntT>();
+  const T* values_data = x.non_zero_elements().data<T>();
+  int64_t batch_nnz = x.nnz() / batch_size;
+  cudaDataType_t gpu_type = GetGpuDataType<T>();
+  cusparseIndexType_t index_type = GetCusparseIndexType<IntT>();
+  dev_ctx.CusparseCall([&](cusparseHandle_t handle) {
+    phi::dynload::cusparseCreateCsr(descriptor,
+                                    M,
+                                    N,
+                                    batch_nnz,
+                                    const_cast<IntT*>(crows_data),
+                                    const_cast<IntT*>(cols_data),
+                                    const_cast<T*>(values_data),
+                                    index_type,
+                                    index_type,
+                                    CUSPARSE_INDEX_BASE_ZERO,
+                                    gpu_type);
+  });
+  if (batch_size > 1) {
+#if CUDA_VERSION >= 11080
+    dev_ctx.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseCsrSetStridedBatch(
+          *descriptor, batch_size, M + 1, batch_nnz);
+    });
+#else
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Batch Sparse matmul use 'cusparseCsrSetStridedBatch', which is "
+        "supported from CUDA 11.8"));
+#endif
+  }
+}
+
+template <typename T, typename IntT>
+inline void CreateCooDescriptor(const SparseCooTensor& x,
+                                const GPUContext& dev_ctx,
+                                cusparseSpMatDescr_t* descriptor) {
+  std::vector<int64_t> xdim_vec = vectorize(x.dims());
+  auto x_ndims = xdim_vec.size();
+  PADDLE_ENFORCE_GE(
+      x_ndims,
+      2,
+      common::errors::InvalidArgument("the dim size of SparseCsrTensor must be "
+                                      "greater than or equal to 2."));
+
+  int64_t M = xdim_vec[x_ndims - 2];
+  int64_t N = xdim_vec[x_ndims - 1];
+  int batch_size = 1;
+  for (int i = 0; i < x_ndims - 2; i++) {
+    batch_size *= xdim_vec[i];
+  }
+  int64_t nnz = x.nnz();
+
+  const IntT* indices_data = x.non_zero_indices().data<IntT>();
+  const T* values_data = x.non_zero_elements().data<T>();
+  auto rows_data = indices_data + (x_ndims - 2) * nnz;
+  auto cols_data = indices_data + (x_ndims - 1) * nnz;
+
+  int64_t batch_nnz = nnz / batch_size;
+  cudaDataType_t gpu_type = GetGpuDataType<T>();
+  cusparseIndexType_t index_type = GetCusparseIndexType<IntT>();
+  dev_ctx.CusparseCall([&](cusparseHandle_t handle) {
+    phi::dynload::cusparseCreateCoo(descriptor,
+                                    M,
+                                    N,
+                                    batch_nnz,
+                                    const_cast<IntT*>(rows_data),
+                                    const_cast<IntT*>(cols_data),
+                                    const_cast<T*>(values_data),
+                                    index_type,
+                                    CUSPARSE_INDEX_BASE_ZERO,
+                                    gpu_type);
+  });
+
+  if (batch_size > 1) {
+#if CUDA_VERSION >= 11080
+    dev_ctx.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseCooSetStridedBatch(
+          *descriptor, batch_size, batch_nnz);
+    });
+#else
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Batch Sparse matmul use 'cusparseCooSetStridedBatch', which is "
+        "supported from CUDA 11.8"));
+#endif
+  }
+}
+
+template <typename T>
+class CuSparseSpMatDescriptor {
+ public:
+  explicit CuSparseSpMatDescriptor(const SparseCsrTensor& x,
+                                   const GPUContext& dev_ctx)
+      : dev_ctx_(dev_ctx) {
+    PD_VISIT_BASE_INTEGRAL_TYPES(
+        x.non_zero_crows().dtype(), "Csr CuSparseSpMatDescriptor", ([&] {
+          CreateCsrDescriptor<T, data_t>(x, dev_ctx_, &descriptor_);
+        }));
+    VLOG(6) << "Create csr cusparseSpMatDescr_t " << &descriptor_;
+  }
+
+  explicit CuSparseSpMatDescriptor(const SparseCooTensor& x,
+                                   const GPUContext& dev_ctx)
+      : dev_ctx_(dev_ctx) {
+    PD_VISIT_BASE_INTEGRAL_TYPES(
+        x.non_zero_indices().dtype(), "Coo CuSparseSpMatDescriptor", ([&] {
+          CreateCooDescriptor<T, data_t>(x, dev_ctx_, &descriptor_);
+        }));
+    VLOG(6) << "Create coo cusparseSpMatDescr_t " << &descriptor_;
+  }
+
+  ~CuSparseSpMatDescriptor() {
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseDestroySpMat(descriptor_);
+    });
+    VLOG(6) << "Destroy cusparseSpMatDescr_t " << &descriptor_;
+  }
+
+  const cusparseSpMatDescr_t& descriptor() const { return descriptor_; }
+
+ private:
+  const GPUContext& dev_ctx_;
+  cusparseSpMatDescr_t descriptor_;
+};
+
+/************* DENSE MATRIX DESCRIPTOR ************/
+template <typename T>
+class CuSparseDnMatDescriptor {
+ public:
+  explicit CuSparseDnMatDescriptor(const DenseTensor& x,
+                                   const GPUContext& dev_ctx)
+      : dev_ctx_(dev_ctx) {
+    std::vector<int64_t> xdim_vec = vectorize(x.dims());
+    auto x_ndims = xdim_vec.size();
+    PADDLE_ENFORCE_GE(
+        x_ndims,
+        2,
+        common::errors::InvalidArgument("the dim size of DenseTensor must be "
+                                        "greater than or equal to 2."));
+
+    int64_t M = xdim_vec[x_ndims - 2];
+    int64_t N = xdim_vec[x_ndims - 1];
+    int batch_size = 1;
+    for (int i = 0; i < x_ndims - 2; i++) {
+      batch_size *= xdim_vec[i];
+    }
+
+    const T* x_data = x.data<T>();
+    cudaDataType_t gpu_type = GetGpuDataType<T>();
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseCreateDnMat(&descriptor_,
+                                        M,
+                                        N,
+                                        N,
+                                        const_cast<T*>(x_data),
+                                        gpu_type,
+                                        CUSPARSE_ORDER_ROW);
+    });
+
+    PADDLE_ENFORCE_EQ(x.numel(), batch_size * M * N);
+    if (batch_size > 1) {
+#if CUDA_VERSION >= 11080
+      dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+        phi::dynload::cusparseDnMatSetStridedBatch(
+            descriptor_, batch_size, M * N);
+      });
+#else
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Batch Sparse matmul use 'cusparseDnMatSetStridedBatch', which is "
+          "supported from CUDA 11.8"));
+#endif
+    }
+    VLOG(6) << "Create cusparseDnMatDescr_t " << &descriptor_;
+  }
+
+  ~CuSparseDnMatDescriptor() {
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseDestroyDnMat(descriptor_);
+    });
+    VLOG(6) << "Destroy cusparseDnMatDescr_t " << &descriptor_;
+  }
+
+  const cusparseDnMatDescr_t& descriptor() const { return descriptor_; }
+
+ private:
+  const GPUContext& dev_ctx_;
+  cusparseDnMatDescr_t descriptor_;
+};
+
+/************* DENSE VECTOR DESCRIPTOR ************/
+template <typename T>
+class CuSparseDnVecDescriptor {
+ public:
+  explicit CuSparseDnVecDescriptor(const DenseTensor& x,
+                                   const GPUContext& dev_ctx)
+      : dev_ctx_(dev_ctx) {
+    std::vector<int64_t> xdim_vec = vectorize(x.dims());
+    auto x_ndims = xdim_vec.size();
+    PADDLE_ENFORCE_GE(x_ndims,
+                      1,
+                      common::errors::InvalidArgument(
+                          "the dim size of Vec must be equal to 1."));
+
+    const T* x_data = x.data<T>();
+    cudaDataType_t gpu_type = GetGpuDataType<T>();
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseCreateDnVec(
+          &descriptor_, x.numel(), const_cast<T*>(x_data), gpu_type);
+    });
+
+    VLOG(6) << "Create cusparseDnVecDescr_t " << &descriptor_;
+  }
+
+  ~CuSparseDnVecDescriptor() {
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseDestroyDnVec(descriptor_);
+    });
+    VLOG(6) << "Destroy cusparseDnVecDescr_t " << &descriptor_;
+  }
+
+  const cusparseDnVecDescr_t& descriptor() const { return descriptor_; }
+
+ private:
+  const GPUContext& dev_ctx_;
+  cusparseDnVecDescr_t descriptor_;
+};
+
+/************* SPARSE*DENSE->DENSE MATMUL ************/
+template <>
+template <typename T, typename TensorType>
+void SparseBlas<GPUContext>::SPMM(bool transa,
+                                  bool transb,
+                                  T alpha,
+                                  const TensorType& mat_a,
+                                  const DenseTensor& mat_b,
+                                  T beta,
+                                  DenseTensor* mat_out) const {
+  auto a_descriptor = CuSparseSpMatDescriptor<T>(mat_a, dev_ctx_);
+  auto b_descriptor = CuSparseDnMatDescriptor<T>(mat_b, dev_ctx_);
+  auto out_descriptor = CuSparseDnMatDescriptor<T>(*mat_out, dev_ctx_);
+
+  cudaDataType_t gpu_type = GetGpuDataType<T>();
+  size_t buffer_size = 0;
+  dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+    phi::dynload::cusparseSpMM_bufferSize(handle,
+                                          GetTransposeOperation(transa),
+                                          GetTransposeOperation(transb),
+                                          &alpha,
+                                          a_descriptor.descriptor(),
+                                          b_descriptor.descriptor(),
+                                          &beta,
+                                          out_descriptor.descriptor(),
+                                          gpu_type,
+                                          GetSpMMAlgorithm(mat_a),
+                                          &buffer_size);
+  });
+
+  phi::Allocator::AllocationPtr tmp_buffer = phi::memory_utils::Alloc(
+      dev_ctx_.GetPlace(),
+      buffer_size,
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx_.stream())));
+  void* tmp_buffer_ptr = tmp_buffer->ptr();
+  dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+    phi::dynload::cusparseSpMM(handle,
+                               GetTransposeOperation(transa),
+                               GetTransposeOperation(transb),
+                               &alpha,
+                               a_descriptor.descriptor(),
+                               b_descriptor.descriptor(),
+                               &beta,
+                               out_descriptor.descriptor(),
+                               gpu_type,
+                               GetSpMMAlgorithm(mat_a),
+                               tmp_buffer_ptr);
+  });
+}
+
+/************* SPARSE*DENSE->DENSE MV ************/
+template <>
+template <typename T, typename TensorType>
+void SparseBlas<GPUContext>::SPMV(bool transa,
+                                  T alpha,
+                                  const TensorType& mat_a,
+                                  const DenseTensor& vec_x,
+                                  T beta,
+                                  DenseTensor* vec_out) const {
+  auto a_descriptor = CuSparseSpMatDescriptor<T>(mat_a, dev_ctx_);
+  auto x_descriptor = CuSparseDnVecDescriptor<T>(vec_x, dev_ctx_);
+  auto out_descriptor = CuSparseDnVecDescriptor<T>(*vec_out, dev_ctx_);
+
+  cudaDataType_t gpu_type = GetGpuDataType<T>();
+  size_t buffer_size = 0;
+  dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+    phi::dynload::cusparseSpMV_bufferSize(handle,
+                                          GetTransposeOperation(transa),
+                                          &alpha,
+                                          a_descriptor.descriptor(),
+                                          x_descriptor.descriptor(),
+                                          &beta,
+                                          out_descriptor.descriptor(),
+                                          gpu_type,
+                                          CUSPARSE_SPMV_ALG_DEFAULT,
+                                          &buffer_size);
+  });
+
+  phi::Allocator::AllocationPtr tmp_buffer = phi::memory_utils::Alloc(
+      dev_ctx_.GetPlace(),
+      buffer_size,
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx_.stream())));
+  void* tmp_buffer_ptr = tmp_buffer->ptr();
+  dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+    phi::dynload::cusparseSpMV(handle,
+                               GetTransposeOperation(transa),
+                               &alpha,
+                               a_descriptor.descriptor(),
+                               x_descriptor.descriptor(),
+                               &beta,
+                               out_descriptor.descriptor(),
+                               gpu_type,
+                               CUSPARSE_SPMV_ALG_DEFAULT,
+                               tmp_buffer_ptr);
+  });
+}
+
+/************* DENSE*DENSE->SPARSE MATMUL ************/
+template <>
+template <typename T, typename TensorType>
+void SparseBlas<GPUContext>::SDDMM(bool transa,
+                                   bool transb,
+                                   T alpha,
+                                   const DenseTensor& mat_a,
+                                   const DenseTensor& mat_b,
+                                   T beta,
+                                   TensorType* mat_out) const {
+  auto a_descriptor = CuSparseDnMatDescriptor<T>(mat_a, dev_ctx_);
+  auto b_descriptor = CuSparseDnMatDescriptor<T>(mat_b, dev_ctx_);
+  auto out_descriptor = CuSparseSpMatDescriptor<T>(*mat_out, dev_ctx_);
+
+  cudaDataType_t gpu_type = GetGpuDataType<T>();
+  size_t buffer_size = 0;
+  dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+    phi::dynload::cusparseSDDMM_bufferSize(handle,
+                                           GetTransposeOperation(transa),
+                                           GetTransposeOperation(transb),
+                                           &alpha,
+                                           a_descriptor.descriptor(),
+                                           b_descriptor.descriptor(),
+                                           &beta,
+                                           out_descriptor.descriptor(),
+                                           gpu_type,
+                                           CUSPARSE_SDDMM_ALG_DEFAULT,
+                                           &buffer_size);
+  });
+
+  phi::Allocator::AllocationPtr tmp_buffer = phi::memory_utils::Alloc(
+      dev_ctx_.GetPlace(),
+      buffer_size,
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx_.stream())));
+  void* tmp_buffer_ptr = tmp_buffer->ptr();
+
+  dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+    phi::dynload::cusparseSDDMM_preprocess(handle,
+                                           GetTransposeOperation(transa),
+                                           GetTransposeOperation(transb),
+                                           &alpha,
+                                           a_descriptor.descriptor(),
+                                           b_descriptor.descriptor(),
+                                           &beta,
+                                           out_descriptor.descriptor(),
+                                           gpu_type,
+                                           CUSPARSE_SDDMM_ALG_DEFAULT,
+                                           tmp_buffer_ptr);
+  });
+
+  dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+    phi::dynload::cusparseSDDMM(handle,
+                                GetTransposeOperation(transa),
+                                GetTransposeOperation(transb),
+                                &alpha,
+                                a_descriptor.descriptor(),
+                                b_descriptor.descriptor(),
+                                &beta,
+                                out_descriptor.descriptor(),
+                                gpu_type,
+                                CUSPARSE_SDDMM_ALG_DEFAULT,
+                                tmp_buffer_ptr);
+  });
+}
+
+/************* SPARSE*SPARSE->SPARSE MATMUL ************/
+template <typename T>
+__global__ void GetCsrBatchNnz(const int32_t* crow_data,
+                               int64_t rows,
+                               int32_t* batch_nnz) {
+  int64_t i = static_cast<int64_t>(threadIdx.x);
+  batch_nnz[i] = crow_data[(i + 1) * (rows + 1) - 1];
+}
+
+template <>
+template <typename T>
+void SparseBlas<GPUContext>::SPGEMM(bool transa,
+                                    bool transb,
+                                    T alpha,
+                                    const SparseCsrTensor& mat_a,
+                                    const SparseCsrTensor& mat_b,
+                                    T beta,
+                                    SparseCsrTensor* mat_out) const {
+  DenseTensor* mat_out_crows = mat_out->mutable_crows();
+  DenseTensor* mat_out_cols = mat_out->mutable_cols();
+  DenseTensor* mat_out_values = mat_out->mutable_values();
+
+  MetaTensor out_crows_meta(mat_out_crows);
+  out_crows_meta.set_dtype(DataType::INT32);
+  out_crows_meta.set_dims(mat_a.crows().dims());
+  dev_ctx_.template Alloc<int32_t>(mat_out_crows);
+
+  std::vector<int64_t> a_dim_vec = vectorize(mat_a.dims());
+  auto a_ndims = a_dim_vec.size();
+  const int64_t a_rows = a_dim_vec[a_ndims - 2];
+  const int64_t a_cols = a_dim_vec[a_ndims - 1];
+  int a_batch_size = 1;
+  for (int i = 0; i < a_ndims - 2; i++) {
+    a_batch_size *= a_dim_vec[i];
+  }
+
+  std::vector<int64_t> b_dim_vec = vectorize(mat_b.dims());
+  auto b_ndims = b_dim_vec.size();
+  const int64_t b_rows = b_dim_vec[b_ndims - 2];
+  const int64_t b_cols = b_dim_vec[b_ndims - 1];
+
+  // cusparseSpGEMM only support 32-bit indices.
+  const int32_t *a_crows_data = nullptr, *a_cols_data = nullptr,
+                *b_crows_data = nullptr, *b_cols_data = nullptr;
+  std::shared_ptr<DenseTensor> a_crows_int = nullptr, a_cols_int = nullptr,
+                               b_crows_int = nullptr, b_cols_int = nullptr;
+
+  if (mat_a.crows().dtype() == DataType::INT32) {
+    a_crows_data = mat_a.crows().data<int32_t>();
+    a_cols_data = mat_a.cols().data<int32_t>();
+  } else {
+    a_crows_int = std::make_shared<DenseTensor>();
+    a_cols_int = std::make_shared<DenseTensor>();
+    MetaTensor crows_meta(a_crows_int.get());
+    crows_meta.set_dims(mat_a.crows().dims());
+    MetaTensor cols_meta(a_cols_int.get());
+    cols_meta.set_dims(mat_a.cols().dims());
+
+    CastKernel<int64_t>(
+        dev_ctx_, mat_a.crows(), DataType::INT32, a_crows_int.get());
+    CastKernel<int64_t>(
+        dev_ctx_, mat_a.cols(), DataType::INT32, a_cols_int.get());
+
+    a_crows_data = a_crows_int->data<int32_t>();
+    a_cols_data = a_cols_int->data<int32_t>();
+  }
+
+  if (mat_b.crows().dtype() == DataType::INT32) {
+    b_crows_data = mat_b.crows().data<int32_t>();
+    b_cols_data = mat_b.cols().data<int32_t>();
+  } else {
+    b_crows_int = std::make_shared<DenseTensor>();
+    b_cols_int = std::make_shared<DenseTensor>();
+    MetaTensor crows_meta(b_crows_int.get());
+    crows_meta.set_dims(mat_b.crows().dims());
+    MetaTensor cols_meta(b_cols_int.get());
+    cols_meta.set_dims(mat_b.cols().dims());
+
+    CastKernel<int64_t>(
+        dev_ctx_, mat_b.crows(), DataType::INT32, b_crows_int.get());
+    CastKernel<int64_t>(
+        dev_ctx_, mat_b.cols(), DataType::INT32, b_cols_int.get());
+
+    b_crows_data = b_crows_int->data<int32_t>();
+    b_cols_data = b_cols_int->data<int32_t>();
+  }
+
+  const T* a_values_data = mat_a.values().data<T>();
+  const T* b_values_data = mat_b.values().data<T>();
+  const int32_t* out_crows_data = mat_out->crows().data<int32_t>();
+
+  const int batch_size = a_batch_size;
+  std::vector<int32_t> a_batch_nnz_vec(batch_size);
+  std::vector<int32_t> b_batch_nnz_vec(batch_size);
+
+  if (batch_size == 1) {
+    a_batch_nnz_vec[0] = mat_a.nnz();
+    b_batch_nnz_vec[0] = mat_b.nnz();
+  } else {
+    phi::Allocator::AllocationPtr tmp_buffer = phi::memory_utils::Alloc(
+        dev_ctx_.GetPlace(),
+        batch_size * sizeof(int32_t),
+        phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx_.stream())));
+    void* tmp_buffer_ptr = tmp_buffer->ptr();
+
+    GetCsrBatchNnz<T><<<1, batch_size, 0, dev_ctx_.stream()>>>(
+        a_crows_data, a_rows, static_cast<int32_t*>(tmp_buffer_ptr));
+#ifdef PADDLE_WITH_CUDA
+    PADDLE_ENFORCE_EQ(
+        phi::backends::gpu::IsCUDAGraphCapturing(),
+        false,
+        common::errors::InvalidArgument(
+            "SparseBlas CsrMM does not support CUDA Graph capture: async D2H "
+            "copy to local vectors 'a_batch_nnz_vec' / 'b_batch_nnz_vec' will "
+            "bake the destination addresses into the graph; on replay the "
+            "vectors are re-created at different addresses, causing "
+            "dangling-pointer writes."));
+#endif
+    phi::backends::gpu::GpuMemcpyAsync(a_batch_nnz_vec.data(),
+                                       tmp_buffer_ptr,
+                                       batch_size * sizeof(int32_t),
+                                       gpuMemcpyDeviceToHost,
+                                       dev_ctx_.stream());
+
+    GetCsrBatchNnz<T><<<1, batch_size, 0, dev_ctx_.stream()>>>(
+        b_crows_data, b_rows, static_cast<int32_t*>(tmp_buffer_ptr));
+    phi::backends::gpu::GpuMemcpyAsync(b_batch_nnz_vec.data(),
+                                       tmp_buffer_ptr,
+                                       batch_size * sizeof(int32_t),
+                                       gpuMemcpyDeviceToHost,
+                                       dev_ctx_.stream());
+  }
+
+  std::vector<DenseTensor> out_batch_cols_vec(batch_size);
+  std::vector<DenseTensor> out_batch_values_vec(batch_size);
+  cudaDataType_t gpu_type = GetGpuDataType<T>();
+
+  const int32_t* a_batch_crows_data = a_crows_data;
+  const int32_t* a_batch_cols_data = a_cols_data;
+  const T* a_batch_values_data = a_values_data;
+
+  const int32_t* b_batch_crows_data = b_crows_data;
+  const int32_t* b_batch_cols_data = b_cols_data;
+  const T* b_batch_values_data = b_values_data;
+
+  const int32_t* out_batch_crows_data = out_crows_data;
+
+  for (int i = 0; i < batch_size; ++i) {
+    int32_t a_batch_nnz = a_batch_nnz_vec[i];
+    int32_t b_batch_nnz = b_batch_nnz_vec[i];
+
+    cusparseSpMatDescr_t a_batch_desc, b_batch_desc, out_batch_desc;
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseCreateCsr(&a_batch_desc,
+                                      a_rows,
+                                      a_cols,
+                                      a_batch_nnz,
+                                      const_cast<int32_t*>(a_batch_crows_data),
+                                      const_cast<int32_t*>(a_batch_cols_data),
+                                      const_cast<T*>(a_batch_values_data),
+                                      CUSPARSE_INDEX_32I,
+                                      CUSPARSE_INDEX_32I,
+                                      CUSPARSE_INDEX_BASE_ZERO,
+                                      gpu_type);
+    });
+
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseCreateCsr(&b_batch_desc,
+                                      b_rows,
+                                      b_cols,
+                                      b_batch_nnz,
+                                      const_cast<int32_t*>(b_batch_crows_data),
+                                      const_cast<int32_t*>(b_batch_cols_data),
+                                      const_cast<T*>(b_batch_values_data),
+                                      CUSPARSE_INDEX_32I,
+                                      CUSPARSE_INDEX_32I,
+                                      CUSPARSE_INDEX_BASE_ZERO,
+                                      gpu_type);
+    });
+
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseCreateCsr(&out_batch_desc,
+                                      a_rows,
+                                      b_cols,
+                                      0,
+                                      nullptr,
+                                      nullptr,
+                                      nullptr,
+                                      CUSPARSE_INDEX_32I,
+                                      CUSPARSE_INDEX_32I,
+                                      CUSPARSE_INDEX_BASE_ZERO,
+                                      gpu_type);
+    });
+
+    size_t buffer_a_size = 0, buffer_b_size = 0;
+    cusparseSpGEMMDescr_t spgemm_desc;
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseSpGEMM_createDescr(&spgemm_desc);
+    });
+
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseSpGEMM_workEstimation(handle,
+                                                  GetTransposeOperation(transa),
+                                                  GetTransposeOperation(transb),
+                                                  &alpha,
+                                                  a_batch_desc,
+                                                  b_batch_desc,
+                                                  &beta,
+                                                  out_batch_desc,
+                                                  gpu_type,
+                                                  CUSPARSE_SPGEMM_DEFAULT,
+                                                  spgemm_desc,
+                                                  &buffer_a_size,
+                                                  nullptr);
+    });
+
+    phi::Allocator::AllocationPtr tmp_buffer_a = phi::memory_utils::Alloc(
+        dev_ctx_.GetPlace(),
+        buffer_a_size,
+        phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx_.stream())));
+    void* tmp_buffer_a_ptr = tmp_buffer_a->ptr();
+
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseSpGEMM_workEstimation(handle,
+                                                  GetTransposeOperation(transa),
+                                                  GetTransposeOperation(transb),
+                                                  &alpha,
+                                                  a_batch_desc,
+                                                  b_batch_desc,
+                                                  &beta,
+                                                  out_batch_desc,
+                                                  gpu_type,
+                                                  CUSPARSE_SPGEMM_DEFAULT,
+                                                  spgemm_desc,
+                                                  &buffer_a_size,
+                                                  tmp_buffer_a_ptr);
+    });
+
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseSpGEMM_compute(handle,
+                                           GetTransposeOperation(transa),
+                                           GetTransposeOperation(transb),
+                                           &alpha,
+                                           a_batch_desc,
+                                           b_batch_desc,
+                                           &beta,
+                                           out_batch_desc,
+                                           gpu_type,
+                                           CUSPARSE_SPGEMM_DEFAULT,
+                                           spgemm_desc,
+                                           &buffer_b_size,
+                                           nullptr);
+    });
+
+    phi::Allocator::AllocationPtr tmp_buffer_b = phi::memory_utils::Alloc(
+        dev_ctx_.GetPlace(),
+        buffer_b_size,
+        phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx_.stream())));
+    void* tmp_buffer_b_ptr = tmp_buffer_b->ptr();
+
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseSpGEMM_compute(handle,
+                                           GetTransposeOperation(transa),
+                                           GetTransposeOperation(transb),
+                                           &alpha,
+                                           a_batch_desc,
+                                           b_batch_desc,
+                                           &beta,
+                                           out_batch_desc,
+                                           gpu_type,
+                                           CUSPARSE_SPGEMM_DEFAULT,
+                                           spgemm_desc,
+                                           &buffer_b_size,
+                                           tmp_buffer_b_ptr);
+    });
+
+    int64_t out_num_crows, out_num_cols, out_num_values;
+
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseSpMatGetSize(
+          out_batch_desc, &out_num_crows, &out_num_cols, &out_num_values);
+    });
+
+    out_batch_cols_vec[i].Resize(common::make_dim(out_num_values));
+    dev_ctx_.template Alloc<int32_t>(&out_batch_cols_vec[i]);
+    out_batch_values_vec[i].Resize(common::make_dim(out_num_values));
+    dev_ctx_.template Alloc<T>(&out_batch_values_vec[i]);
+
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseCsrSetPointers(
+          out_batch_desc,
+          const_cast<int32_t*>(out_batch_crows_data),
+          const_cast<int32_t*>(out_batch_cols_vec[i].data<int32_t>()),
+          const_cast<T*>(out_batch_values_vec[i].data<T>()));
+    });
+
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseSpGEMM_copy(handle,
+                                        GetTransposeOperation(transa),
+                                        GetTransposeOperation(transb),
+                                        &alpha,
+                                        a_batch_desc,
+                                        b_batch_desc,
+                                        &beta,
+                                        out_batch_desc,
+                                        gpu_type,
+                                        CUSPARSE_SPGEMM_DEFAULT,
+                                        spgemm_desc);
+    });
+
+    dev_ctx_.CusparseCall([&](cusparseHandle_t handle) {
+      phi::dynload::cusparseSpGEMM_destroyDescr(spgemm_desc);
+    });
+
+    a_batch_crows_data += a_rows + 1;
+    a_batch_cols_data += a_batch_nnz;
+    a_batch_values_data += a_batch_nnz;
+
+    b_batch_crows_data += b_rows + 1;
+    b_batch_cols_data += b_batch_nnz;
+    b_batch_values_data += b_batch_nnz;
+
+    out_batch_crows_data += a_rows + 1;
+  }
+
+  if (batch_size == 1) {
+    *(mat_out->mutable_cols()) = std::move(out_batch_cols_vec[0]);
+    *(mat_out->mutable_values()) = std::move(out_batch_values_vec[0]);
+
+  } else {
+    std::vector<const DenseTensor*> cols_vec, values_vec;
+
+    for (int i = 0; i < batch_size; ++i) {
+      cols_vec.push_back(&out_batch_cols_vec[i]);
+      values_vec.push_back(&out_batch_values_vec[i]);
+    }
+
+    phi::ConcatKernel<int32_t>(dev_ctx_, cols_vec, 0, mat_out->mutable_cols());
+    phi::ConcatKernel<T>(dev_ctx_, values_vec, 0, mat_out->mutable_values());
+  }
+
+  if (mat_a.crows().dtype() == DataType::INT64 ||
+      mat_b.crows().dtype() == DataType::INT64) {
+    CastKernel<int32_t>(
+        dev_ctx_, *mat_out_crows, DataType::INT64, mat_out_crows);
+    CastKernel<int32_t>(dev_ctx_, *mat_out_cols, DataType::INT64, mat_out_cols);
+  }
+}
+}  // namespace sparse
+}  // namespace funcs
+}  // namespace phi

paddle/phi/kernels/funcs/sparse/sparse_blas_impl.hip.h

@@ -0,0 +1,405 @@
+// 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.
+
+#pragma once
+
+#include "paddle/common/ddim.h"
+#include "paddle/phi/backends/dynload/rocsparse.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
+#include "paddle/phi/core/visit_type.h"
+namespace phi {
+namespace funcs {
+namespace sparse {
+
+template <typename IntT>
+rocsparse_indextype GetGpuIndexType() {
+  if (std::is_same<IntT, int32_t>::value) {
+    return rocsparse_indextype_i32;
+  } else if (std::is_same<IntT, int64_t>::value) {
+    return rocsparse_indextype_i64;
+  }
+}
+
+template <typename T>
+rocsparse_datatype GetGpuDataType() {
+  if (std::is_same<T, float>::value) {
+    return rocsparse_datatype_f32_r;
+  } else if (std::is_same<T, double>::value) {
+    return rocsparse_datatype_f64_r;
+  }
+}
+
+inline rocsparse_operation GetTransposeOperation(const bool trans) {
+  if (trans) {
+    return rocsparse_operation_transpose;
+  } else {
+    return rocsparse_operation_none;
+  }
+}
+
+template <typename TensorType>
+inline rocsparse_spmm_alg GetSpMMAlgorithm(const TensorType& x) {
+  return rocsparse_spmm_alg_default;
+}
+
+/************* SPARSE MATRIX DESCRIPTOR (COO/CSR) ************/
+template <typename T, typename IntT>
+inline void CreateCsrDescriptor(const SparseCsrTensor& x,
+                                const GPUContext& dev_ctx,
+                                rocsparse_spmat_descr* descriptor) {
+  std::vector<int64_t> xdim_vec = vectorize(x.dims());
+  auto x_ndims = xdim_vec.size();
+  PADDLE_ENFORCE_GE(
+      x_ndims,
+      2,
+      common::errors::InvalidArgument("the dim size of SparseCsrTensor must be "
+                                      "greater than or equal to 2."));
+  int64_t M = xdim_vec[x_ndims - 2];
+  int64_t N = xdim_vec[x_ndims - 1];
+  int batch_size = 1;
+  for (int i = 0; i < x_ndims - 2; i++) {
+    batch_size *= xdim_vec[i];
+  }
+  PADDLE_ENFORCE_EQ(x.non_zero_crows().numel(),
+                    batch_size * (M + 1),
+                    common::errors::PreconditionNotMet(
+                        "the length of SparseCsrTensor crows is not right."));
+
+  const IntT* crows_data = x.non_zero_crows().data<IntT>();
+  const IntT* cols_data = x.non_zero_cols().data<IntT>();
+  const T* values_data = x.non_zero_elements().data<T>();
+
+  int64_t batch_nnz = x.nnz() / batch_size;
+  rocsparse_indextype itype = GetGpuIndexType<int64_t>();
+  rocsparse_indextype jtype = GetGpuIndexType<int64_t>();
+  rocsparse_datatype ttype = GetGpuDataType<T>();
+  dev_ctx.CusparseCall([&](rocsparse_handle handle) {
+    phi::dynload::rocsparse_create_csr_descr(descriptor,
+                                             M,
+                                             N,
+                                             batch_nnz,
+                                             const_cast<IntT*>(crows_data),
+                                             const_cast<IntT*>(cols_data),
+                                             const_cast<T*>(values_data),
+                                             itype,
+                                             jtype,
+                                             rocsparse_index_base_zero,
+                                             ttype);
+  });
+  if (batch_size > 1) {
+    // TODO(umiswing): Add batch sparse matmul support for ROCM after 5.2.0
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Batch Sparse matmul use 'rocsparse_coo_set_strided_batch', which is "
+        "supported from ROCM 5.2.0"));
+  }
+}
+
+template <typename T, typename IntT>
+inline void CreateCooDescriptor(const SparseCooTensor& x,
+                                const GPUContext& dev_ctx,
+                                rocsparse_spmat_descr* descriptor) {
+  std::vector<int64_t> xdim_vec = vectorize(x.dims());
+  auto x_ndims = xdim_vec.size();
+  PADDLE_ENFORCE_GE(
+      x_ndims,
+      2,
+      common::errors::InvalidArgument("the dim size of SparseCooTensor must be "
+                                      "greater than or equal to 2."));
+
+  int64_t M = xdim_vec[x_ndims - 2];
+  int64_t N = xdim_vec[x_ndims - 1];
+  int batch_size = 1;
+  for (int i = 0; i < x_ndims - 2; i++) {
+    batch_size *= xdim_vec[i];
+  }
+  int64_t nnz = x.nnz();
+
+  const IntT* indices_data = x.non_zero_indices().data<IntT>();
+  const T* values_data = x.non_zero_elements().data<T>();
+  auto rows_data = indices_data + (x_ndims - 2) * nnz;
+  auto cols_data = indices_data + (x_ndims - 1) * nnz;
+
+  int64_t batch_nnz = nnz / batch_size;
+  rocsparse_indextype itype = GetGpuIndexType<int64_t>();
+  rocsparse_datatype ttype = GetGpuDataType<T>();
+  dev_ctx.CusparseCall([&](rocsparse_handle handle) {
+    phi::dynload::rocsparse_create_coo_descr(descriptor,
+                                             M,
+                                             N,
+                                             batch_nnz,
+                                             const_cast<IntT*>(rows_data),
+                                             const_cast<IntT*>(cols_data),
+                                             const_cast<T*>(values_data),
+                                             itype,
+                                             rocsparse_index_base_zero,
+                                             ttype);
+  });
+
+  if (batch_size > 1) {
+    // TODO(umiswing): Add batch sparse matmul support for ROCM after 5.2.0
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Batch Sparse matmul use 'rocsparse_coo_set_strided_batch', which is "
+        "supported from ROCM 5.2.0"));
+  }
+}
+
+template <typename T>
+class RocSparseSpMatDescriptor {
+ public:
+  explicit RocSparseSpMatDescriptor(const SparseCsrTensor& x,
+                                    const GPUContext& dev_ctx)
+      : dev_ctx_(dev_ctx) {
+    PD_VISIT_BASE_INTEGRAL_TYPES(
+        x.non_zero_crows().dtype(), "Csr RocSparseSpMatDescriptor", ([&] {
+          CreateCsrDescriptor<T, data_t>(x, dev_ctx_, &descriptor_);
+        }));
+    VLOG(6) << "Create csr rocsparse_spmat_descr " << &descriptor_;
+  }
+  explicit RocSparseSpMatDescriptor(const SparseCooTensor& x,
+                                    const GPUContext& dev_ctx)
+      : dev_ctx_(dev_ctx) {
+    PD_VISIT_BASE_INTEGRAL_TYPES(
+        x.non_zero_indices().dtype(), "Coo RocSparseSpMatDescriptor", ([&] {
+          CreateCooDescriptor<T, data_t>(x, dev_ctx_, &descriptor_);
+        }));
+    VLOG(6) << "Create coo rocsparse_spmat_descr " << &descriptor_;
+  }
+
+  ~RocSparseSpMatDescriptor() {
+    dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
+      phi::dynload::rocsparse_destroy_spmat_descr(descriptor_);
+    });
+    VLOG(6) << "Destroy roscparse_spmat_descr " << &descriptor_;
+  }
+
+  const rocsparse_spmat_descr& descriptor() const { return descriptor_; }
+
+ private:
+  const GPUContext& dev_ctx_;
+  rocsparse_spmat_descr descriptor_;
+};
+
+/************* DENSE MATRIX DESCRIPTOR ************/
+template <typename T>
+class RocSparseDnMatDescriptor {
+ public:
+  explicit RocSparseDnMatDescriptor(const DenseTensor& x,
+                                    const GPUContext& dev_ctx)
+      : dev_ctx_(dev_ctx) {
+    std::vector<int64_t> xdim_vec = vectorize(x.dims());
+    auto x_ndims = xdim_vec.size();
+    PADDLE_ENFORCE_GE(
+        x_ndims,
+        2,
+        common::errors::InvalidArgument("the dim size of DenseTensor must be "
+                                        "greater than or equal to 2."));
+
+    int64_t M = xdim_vec[x_ndims - 2];
+    int64_t N = xdim_vec[x_ndims - 1];
+    int batch_size = 1;
+    for (int i = 0; i < x_ndims - 2; i++) {
+      batch_size *= xdim_vec[i];
+    }
+
+    const T* x_data = x.data<T>();
+    rocsparse_datatype ttype = GetGpuDataType<T>();
+    dev_ctx.CusparseCall([&](rocsparse_handle handle) {
+      phi::dynload::rocsparse_create_dnmat_descr(&descriptor_,
+                                                 M,
+                                                 N,
+                                                 N,
+                                                 const_cast<T*>(x_data),
+                                                 ttype,
+                                                 rocsparse_order_row);
+    });
+
+    PADDLE_ENFORCE_EQ(
+        x.numel(),
+        batch_size * M * N,
+        common::errors::InvalidArgument("The number of elements in DenseTensor "
+                                        "must equals to batch_size * M * N."));
+    if (batch_size > 1) {
+      // TODO(umiswing): Add batch sparse matmul support for ROCM after 5.2.0
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Batch Sparse matmul use 'rocsparse_dnmat_set_strided_batch', which "
+          "is supported from ROCM 5.2.0"));
+    }
+    VLOG(6) << "Create cusparseDnMatDescr_t " << &descriptor_;
+  }
+
+  ~RocSparseDnMatDescriptor() {
+    dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
+      phi::dynload::rocsparse_destroy_dnmat_descr(descriptor_);
+    });
+    VLOG(6) << "Destroy rocsparse_dnmat_descr " << &descriptor_;
+  }
+
+  const rocsparse_dnmat_descr& descriptor() const { return descriptor_; }
+
+ private:
+  const GPUContext& dev_ctx_;
+  rocsparse_dnmat_descr descriptor_;
+};
+
+/************* SPARSE*DENSE->DENSE MATMUL ************/
+template <>
+template <typename T, typename TensorType>
+void SparseBlas<GPUContext>::SPMM(bool transa,
+                                  bool transb,
+                                  T alpha,
+                                  const TensorType& mat_a,
+                                  const DenseTensor& mat_b,
+                                  T beta,
+                                  DenseTensor* mat_out) const {
+  auto a_descriptor = RocSparseSpMatDescriptor<T>(mat_a, dev_ctx_);
+  auto b_descriptor = RocSparseDnMatDescriptor<T>(mat_b, dev_ctx_);
+  auto out_descriptor = RocSparseDnMatDescriptor<T>(*mat_out, dev_ctx_);
+
+  rocsparse_datatype ttype = GetGpuDataType<T>();
+  size_t buffer_size = 0;
+
+  // Query SpMM buffer
+  dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
+    phi::dynload::rocsparse_spmm(handle,
+                                 GetTransposeOperation(transa),
+                                 GetTransposeOperation(transb),
+                                 &alpha,
+                                 a_descriptor.descriptor(),
+                                 b_descriptor.descriptor(),
+                                 &beta,
+                                 out_descriptor.descriptor(),
+                                 ttype,
+                                 GetSpMMAlgorithm(mat_a),
+                                 rocsparse_spmm_stage_buffer_size,
+                                 &buffer_size,
+                                 nullptr);
+  });
+
+  // Allocate buffer
+  phi::Allocator::AllocationPtr tmp_buffer = phi::memory_utils::Alloc(
+      dev_ctx_.GetPlace(),
+      buffer_size,
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx_.stream())));
+  void* tmp_buffer_ptr = tmp_buffer->ptr();
+
+  // Preprocess data
+  dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
+    phi::dynload::rocsparse_spmm(handle,
+                                 GetTransposeOperation(transa),
+                                 GetTransposeOperation(transb),
+                                 &alpha,
+                                 a_descriptor.descriptor(),
+                                 b_descriptor.descriptor(),
+                                 &beta,
+                                 out_descriptor.descriptor(),
+                                 ttype,
+                                 GetSpMMAlgorithm(mat_a),
+                                 rocsparse_spmm_stage_preprocess,
+                                 &buffer_size,
+                                 tmp_buffer_ptr);
+  });
+
+  // Performs the actual SpMM computation
+  dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
+    phi::dynload::rocsparse_spmm(handle,
+                                 GetTransposeOperation(transa),
+                                 GetTransposeOperation(transb),
+                                 &alpha,
+                                 a_descriptor.descriptor(),
+                                 b_descriptor.descriptor(),
+                                 &beta,
+                                 out_descriptor.descriptor(),
+                                 ttype,
+                                 GetSpMMAlgorithm(mat_a),
+                                 rocsparse_spmm_stage_compute,
+                                 &buffer_size,
+                                 tmp_buffer_ptr);
+  });
+}
+
+/************* DENSE*DENSE->SPARSE MATMUL ************/
+#if HIP_VERSION >= 403
+template <>
+template <typename T, typename TensorType>
+void SparseBlas<GPUContext>::SDDMM(bool transa,
+                                   bool transb,
+                                   T alpha,
+                                   const DenseTensor& mat_a,
+                                   const DenseTensor& mat_b,
+                                   T beta,
+                                   TensorType* mat_out) const {
+  auto a_descriptor = RocSparseDnMatDescriptor<T>(mat_a, dev_ctx_);
+  auto b_descriptor = RocSparseDnMatDescriptor<T>(mat_b, dev_ctx_);
+  auto out_descriptor = RocSparseSpMatDescriptor<T>(*mat_out, dev_ctx_);
+
+  rocsparse_datatype gpu_type = GetGpuDataType<T>();
+  size_t buffer_size = 0;
+  dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
+    phi::dynload::rocsparse_sddmm_buffer_size(handle,
+                                              GetTransposeOperation(transa),
+                                              GetTransposeOperation(transb),
+                                              &alpha,
+                                              a_descriptor.descriptor(),
+                                              b_descriptor.descriptor(),
+                                              &beta,
+                                              out_descriptor.descriptor(),
+                                              gpu_type,
+                                              rocsparse_sddmm_alg_default,
+                                              &buffer_size);
+  });
+
+  phi::Allocator::AllocationPtr tmp_buffer = phi::memory_utils::Alloc(
+      dev_ctx_.GetPlace(),
+      buffer_size,
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx_.stream())));
+  void* tmp_buffer_ptr = tmp_buffer->ptr();
+
+  dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
+    phi::dynload::rocsparse_sddmm_preprocess(handle,
+                                             GetTransposeOperation(transa),
+                                             GetTransposeOperation(transb),
+                                             &alpha,
+                                             a_descriptor.descriptor(),
+                                             b_descriptor.descriptor(),
+                                             &beta,
+                                             out_descriptor.descriptor(),
+                                             gpu_type,
+                                             rocsparse_sddmm_alg_default,
+                                             tmp_buffer_ptr);
+  });
+
+  dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
+    phi::dynload::rocsparse_sddmm(handle,
+                                  GetTransposeOperation(transa),
+                                  GetTransposeOperation(transb),
+                                  &alpha,
+                                  a_descriptor.descriptor(),
+                                  b_descriptor.descriptor(),
+                                  &beta,
+                                  out_descriptor.descriptor(),
+                                  gpu_type,
+                                  rocsparse_sddmm_alg_default,
+                                  tmp_buffer_ptr);
+  });
+}
+#endif
+}  // namespace sparse
+}  // namespace funcs
+}  // namespace phi

paddle/phi/kernels/funcs/sparse/utils.cu.h

@@ -0,0 +1,44 @@
+/* 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
+
+namespace phi {
+namespace funcs {
+namespace sparse {
+
+// brief: calculation the distance between start and end
+template <typename T>
+__global__ void DistanceKernel(const T* start, const T* end, T* distance) {
+  if (threadIdx.x == 0) {
+    *distance = end - start;
+  }
+}
+
+inline __device__ bool SetBits(const int value, int* ptr) {
+  const int index = value >> 5;
+  const int mask = 1 << (value & 31);
+  const int old = atomicOr(ptr + index, mask);
+  return (mask & old) != 0;
+}
+
+inline __device__ bool TestBits(const int value, const int* ptr) {
+  const int index = value >> 5;
+  const int mask = 1 << (value & 31);
+  return (mask & ptr[index]) != 0;
+}
+
+}  // namespace sparse
+}  // namespace funcs
+}  // namespace phi