//   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