// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <thrust/device_vector.h>
#include <thrust/sort.h>

#include "paddle/phi/kernels/sparse/unary_kernel.h"

#include "paddle/common/ddim.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
#include "paddle/phi/backends/gpu/gpu_primitives.h"
#include "paddle/phi/common/memory_utils.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/core/visit_type.h"
#include "paddle/phi/kernels/empty_kernel.h"
#include "paddle/phi/kernels/funcs/slice_utils.h"

namespace phi {
namespace sparse {

template <typename IntT>
__global__ void GetCooNonZeroNumberCudaKernel(const IntT* x_indices_data,
                                              const int64_t* axes,
                                              const int64_t* starts,
                                              const int64_t* ends,
                                              const int64_t axes_size,
                                              const int64_t x_nnz,
                                              int* out_nnz,
                                              IntT* out_nnz_indices) {
  CUDA_KERNEL_LOOP_TYPE(j, x_nnz, int64_t) {
    bool hit = true;
    for (size_t ii = 0; ii < axes_size; ++ii) {
      auto item = x_indices_data[axes[ii] * x_nnz + j];
      if (!(starts[ii] <= item && item < ends[ii])) {
        hit = false;
        break;
      }
    }
    if (!hit) continue;
    int old_value = atomicAdd(out_nnz, 1);
    out_nnz_indices[old_value] = j;
  }
}

template <typename T, typename IntT>
__global__ void GetCooOutCudaKernel(const IntT* x_indices_data,
                                    const T* x_values_data,
                                    const int64_t* axes,
                                    const int64_t* starts,
                                    const int64_t axes_size,
                                    const int64_t sparse_dim,
                                    const int64_t x_nnz,
                                    const int64_t out_nnz,
                                    const IntT* out_nnz_indices,
                                    IntT* out_indices_data,
                                    T* out_values_data) {
  CUDA_KERNEL_LOOP_TYPE(index, out_nnz, int64_t) {
    // index is in the order of the non-zero elements in out
    // out_nnz_indices[index] is the valid index in x's non-zero elements, where
    // the `hit` is true.
    IntT j = out_nnz_indices[index];
    // set value
    out_values_data[index] = x_values_data[j];
    // set coordinate
    for (int64_t i = 0; i < sparse_dim; ++i) {
      out_indices_data[i * out_nnz + index] = x_indices_data[i * x_nnz + j];
    }
    for (size_t ii = 0; ii < axes_size; ++ii) {
      auto i = axes[ii];
      out_indices_data[i * out_nnz + index] -= starts[ii];
    }
  }
}

template <typename T, typename IntT, typename Context>
void SliceCooGPUCompute(const Context& dev_ctx,
                        const SparseCooTensor& x,
                        const std::vector<int64_t>& axes,
                        const std::vector<int64_t>& starts,
                        const std::vector<int64_t>& ends,
                        SparseCooTensor* out) {
  const DDim& x_dims = x.dims();

  // Step1: Infer output dims
  auto out_dims = funcs::GetSliceDims<int64_t>(
      x_dims, axes, starts, ends, nullptr, nullptr);

  // Step2: Get the number of non zero elements
  DenseTensor d_out_nnz = Empty<int32_t>(dev_ctx, {1});
  int* d_out_nnz_ptr = d_out_nnz.data<int32_t>();
  backends::gpu::GpuMemsetAsync(
      d_out_nnz_ptr, 0, sizeof(int32_t), dev_ctx.stream());

  // out_nnz_indices is the indices where the data is valid in out
  // the length of the out_nnz_indices must be less than x.nnz()
  DenseTensor d_out_nnz_indices = Empty<IntT>(dev_ctx, {x.nnz()});
  auto* d_out_nnz_indices_ptr = d_out_nnz_indices.data<IntT>();
  backends::gpu::GpuMemsetAsync(
      d_out_nnz_indices_ptr, 0, sizeof(IntT), dev_ctx.stream());

  // copy axes to device
  auto d_axes_tensor = memory_utils::Alloc(
      dev_ctx.GetPlace(),
      sizeof(int64_t) * axes.size(),
      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
  int64_t* d_axes = reinterpret_cast<int64_t*>(d_axes_tensor->ptr());
  memory_utils::Copy(dev_ctx.GetPlace(),
                     d_axes,
                     phi::CPUPlace(),
                     axes.data(),
                     sizeof(int64_t) * axes.size(),
                     dev_ctx.stream());

  // copy starts to device
  auto d_starts_tensor = memory_utils::Alloc(
      dev_ctx.GetPlace(),
      sizeof(int64_t) * starts.size(),
      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
  int64_t* d_starts = reinterpret_cast<int64_t*>(d_starts_tensor->ptr());
  memory_utils::Copy(dev_ctx.GetPlace(),
                     d_starts,
                     phi::CPUPlace(),
                     starts.data(),
                     sizeof(int64_t) * starts.size(),
                     dev_ctx.stream());

  // copy ends to device
  auto d_ends_tensor = memory_utils::Alloc(
      dev_ctx.GetPlace(),
      sizeof(int64_t) * ends.size(),
      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
  int64_t* d_ends = reinterpret_cast<int64_t*>(d_ends_tensor->ptr());
  memory_utils::Copy(dev_ctx.GetPlace(),
                     d_ends,
                     phi::CPUPlace(),
                     ends.data(),
                     sizeof(int64_t) * ends.size(),
                     dev_ctx.stream());

  const auto* x_indices_data = x.indices().data<IntT>();

  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x.nnz() + 1, 1);
  GetCooNonZeroNumberCudaKernel<IntT>
      <<<config.block_per_grid.x,
         config.thread_per_block.x,
         0,
         dev_ctx.stream()>>>(x_indices_data,
                             d_axes,
                             d_starts,
                             d_ends,
                             axes.size(),
                             x.nnz(),
                             d_out_nnz_ptr,
                             d_out_nnz_indices_ptr);

  // copy d_out_nnz from device to host (out_nnz)
  int32_t out_nnz = 0;
  backends::gpu::GpuMemcpyAsync(&out_nnz,
                                d_out_nnz_ptr,
                                sizeof(int32_t),
                                gpuMemcpyDeviceToHost,
                                dev_ctx.stream());
  dev_ctx.Wait();
  // sort `d_out_nnz_indices_ptr`
  d_out_nnz_indices.Resize({out_nnz});
#ifdef PADDLE_WITH_HIP
  thrust::sort(thrust::hip::par.on(dev_ctx.stream()),
#else
  thrust::sort(thrust::cuda::par.on(dev_ctx.stream()),
#endif
               d_out_nnz_indices_ptr,
               d_out_nnz_indices_ptr + out_nnz);

  // Step3: Get the values and indices of output
  auto sparse_dim = static_cast<int64_t>(x.sparse_dim());
  DenseTensor out_indices =
      Empty<IntT, Context>(dev_ctx, {sparse_dim, out_nnz});
  DenseTensor out_values = Empty<T, Context>(dev_ctx, {out_nnz});
  out->SetMember(out_indices, out_values, out_dims, x.coalesced());

  auto* out_indices_data = out_indices.data<IntT>();
  auto* out_values_data = out_values.data<T>();
  const auto* x_values_data = x.values().data<T>();

  config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_nnz + 1, 1);
  GetCooOutCudaKernel<T, IntT>
      <<<config.block_per_grid.x,
         config.thread_per_block.x,
         0,
         dev_ctx.stream()>>>(x_indices_data,
                             x_values_data,
                             d_axes,
                             d_starts,
                             axes.size(),
                             sparse_dim,
                             x.nnz(),
                             static_cast<int64_t>(out_nnz),
                             d_out_nnz_indices_ptr,
                             out_indices_data,
                             out_values_data);
}

template <typename T, typename IntT, typename Context>
void SliceCooGPUKernel(const Context& dev_ctx,
                       const SparseCooTensor& x,
                       const phi::IntArray& axes,
                       const phi::IntArray& starts,
                       const phi::IntArray& ends,
                       SparseCooTensor* out) {
  const DDim& x_dims = x.dims();
  std::vector<int64_t> axes_vec = axes.GetData();
  std::vector<int64_t> starts_vec = starts.GetData();
  std::vector<int64_t> ends_vec = ends.GetData();
  // Check and update attr
  funcs::CheckAndUpdateSparseSliceAttrs<int64_t>(
      x_dims, &axes_vec, &starts_vec, &ends_vec);
  SliceCooGPUCompute<T, IntT, Context>(
      dev_ctx, x, axes_vec, starts_vec, ends_vec, out);
}

template <typename T, typename Context>
void SliceCooKernel(const Context& dev_ctx,
                    const SparseCooTensor& x,
                    const phi::IntArray& axes,
                    const phi::IntArray& starts,
                    const phi::IntArray& ends,
                    SparseCooTensor* out) {
  PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "SliceCooGPUKernel", ([&] {
                                 SliceCooGPUKernel<T, data_t, Context>(
                                     dev_ctx, x, axes, starts, ends, out);
                               }));
}

__global__ void GetCsr2DNonZeroNumberCudaKernel(const int64_t* x_crows_data,
                                                const int64_t* x_cols_data,
                                                const int64_t x_crows_start,
                                                const int64_t x_crows_end,
                                                const int64_t min_col,
                                                const int64_t max_col,
                                                int64_t* out_crows_data) {
  CUDA_KERNEL_LOOP_TYPE(i, x_crows_end - x_crows_start, int64_t) {
    if (i == 0) {
      out_crows_data[0] = 0;
    }
    int64_t st = x_crows_data[x_crows_start + i];
    int64_t ed = x_crows_data[x_crows_start + i + 1];
    out_crows_data[i + 1] = 0;
    for (int64_t jj = st; jj < ed; ++jj) {
      if (x_cols_data[jj] >= min_col && x_cols_data[jj] < max_col) {
        out_crows_data[i + 1] += 1;
      }
    }
  }
}

template <typename T>
__global__ void GetCsr2DCudaKernel(const int64_t* x_crows_data,
                                   const int64_t* x_cols_data,
                                   const T* x_values_data,
                                   const int64_t x_crows_start,
                                   const int64_t x_crows_end,
                                   const int64_t min_col,
                                   const int64_t max_col,
                                   const int64_t* out_crows_data,
                                   int64_t* out_cols_data,
                                   T* out_values_data) {
  CUDA_KERNEL_LOOP_TYPE(i, x_crows_end - x_crows_start, int64_t) {
    int64_t st = x_crows_data[x_crows_start + i];
    int64_t ed = x_crows_data[x_crows_start + i + 1];
    int64_t index = out_crows_data[i];
    for (int64_t jj = st; jj < ed; ++jj) {
      if (x_cols_data[jj] >= min_col && x_cols_data[jj] < max_col) {
        out_cols_data[index] = x_cols_data[jj] - min_col;
        out_values_data[index] = x_values_data[jj];
        index++;
      }
    }
  }
}

template <typename T, typename Context>
void SliceCsrTensor2D(const Context& dev_ctx,
                      const SparseCsrTensor& x,
                      const std::vector<int64_t>& axes,
                      const std::vector<int64_t>& starts,
                      const std::vector<int64_t>& ends,
                      const DDim& out_dims,
                      SparseCsrTensor* out) {
  const auto* x_crows_data = x.crows().data<int64_t>();
  const auto* x_cols_data = x.cols().data<int64_t>();
  const auto* x_values_data = x.values().data<T>();
  // Step1: Get the number of non zero elements for out and out_crows
  int64_t out_n_rows = ends[0] - starts[0];
  DenseTensor out_crows = Empty<int64_t, Context>(dev_ctx, {out_n_rows + 1});
  auto* out_crows_data = out_crows.data<int64_t>();

  auto config =
      backends::gpu::GetGpuLaunchConfig1D(dev_ctx, ends[0] - starts[0] + 1, 1);
  GetCsr2DNonZeroNumberCudaKernel<<<config.block_per_grid.x,
                                    config.thread_per_block.x,
                                    0,
                                    dev_ctx.stream()>>>(x_crows_data,
                                                        x_cols_data,
                                                        starts[0],
                                                        ends[0],
                                                        starts[1],
                                                        ends[1],
                                                        out_crows_data);
#ifdef PADDLE_WITH_HIP
  thrust::inclusive_scan(thrust::hip::par.on(dev_ctx.stream()),
#else
  thrust::inclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
#endif
                         out_crows_data,
                         out_crows_data + out_n_rows + 1,
                         out_crows_data);
  int64_t out_nnz = 0;
  backends::gpu::GpuMemcpyAsync(&out_nnz,
                                &out_crows_data[out_n_rows],
                                sizeof(int64_t),
                                gpuMemcpyDeviceToHost,
                                dev_ctx.stream());
  dev_ctx.Wait();
  // Step2: Set out
  DenseTensor out_cols = Empty<int64_t, Context>(dev_ctx, {out_nnz});
  DenseTensor out_values = Empty<T, Context>(dev_ctx, {out_nnz});
  out->SetMember(out_crows, out_cols, out_values, out_dims);
  config =
      backends::gpu::GetGpuLaunchConfig1D(dev_ctx, ends[0] - starts[0] + 1, 1);
  GetCsr2DCudaKernel<T><<<config.block_per_grid.x,
                          config.thread_per_block.x,
                          0,
                          dev_ctx.stream()>>>(x_crows_data,
                                              x_cols_data,
                                              x_values_data,
                                              starts[0],
                                              ends[0],
                                              starts[1],
                                              ends[1],
                                              out_crows.data<int64_t>(),
                                              out_cols.data<int64_t>(),
                                              out_values.data<T>());
}

__global__ void GetXColsOffsetsCudaKernel(const int64_t* x_crows_data,
                                          const int64_t x_n_rows,
                                          const int64_t x_dim0,
                                          int64_t* x_cols_offsets) {
  CUDA_KERNEL_LOOP_TYPE(i, x_dim0, int64_t) {
    if (i == 0) {
      x_cols_offsets[i] = 0;
    }
    x_cols_offsets[i + 1] = x_crows_data[(i + 1) * (x_n_rows + 1) - 1];
  }
}

__global__ void GetCsr3DNonZeroNumberCudaKernel(const int64_t* x_crows_data,
                                                const int64_t* x_cols_data,
                                                const int64_t x_dim0,
                                                const int64_t x_n_rows,
                                                const int64_t* x_cols_offsets,
                                                const int64_t* starts,
                                                const int64_t* ends,
                                                const int64_t out_n_rows,
                                                int64_t* out_crows_data) {
  CUDA_KERNEL_LOOP_TYPE(i, x_dim0 * (x_n_rows + 1), int64_t) {
    int64_t dim0_i = i / (x_n_rows + 1);
    int64_t dim1_i = i % (x_n_rows + 1);
    if (!(dim0_i >= starts[0] && dim0_i < ends[0])) {
      continue;
    }
    if (!(dim1_i >= starts[1] && dim1_i < ends[1])) {
      continue;
    }
    // the starting index of current 2D Tensor in out_crows
    int64_t out_dim0_start = (dim0_i - starts[0]) * (out_n_rows + 1);
    if (dim1_i == starts[1]) {
      out_crows_data[out_dim0_start] = 0;
    }
    int64_t out_crows_idx = out_dim0_start + (dim1_i - starts[1]);
    int64_t st = x_crows_data[i] + x_cols_offsets[dim0_i];
    int64_t ed = x_crows_data[i + 1] + x_cols_offsets[dim0_i];
    out_crows_data[out_crows_idx + 1] = 0;
    for (int64_t jj = st; jj < ed; ++jj) {
      if (x_cols_data[jj] >= starts[2] && x_cols_data[jj] < ends[2]) {
        out_crows_data[out_crows_idx + 1] += 1;
      }
    }
  }
}

template <typename T>
__global__ void GetCsr3DCudaKernel(const int64_t* x_crows_data,
                                   const int64_t* x_cols_data,
                                   const T* x_values_data,
                                   const int64_t* x_cols_offsets,
                                   const int64_t x_dim0,
                                   const int64_t x_n_rows,
                                   const int64_t* starts,
                                   const int64_t* ends,
                                   const int64_t out_n_rows,
                                   const int64_t* out_cols_offsets,
                                   const int64_t* out_crows_data,
                                   int64_t* out_cols_data,
                                   T* out_values_data) {
  CUDA_KERNEL_LOOP_TYPE(i, x_dim0 * (x_n_rows + 1), int64_t) {
    int dim0_i = i / (x_n_rows + 1);
    int dim1_i = i % (x_n_rows + 1);
    if (!(dim0_i >= starts[0] && dim0_i < ends[0])) {
      continue;
    }
    if (!(dim1_i >= starts[1] && dim1_i < ends[1])) {
      continue;
    }
    // the starting index of current 2D Tensor in out_crows
    int64_t out_dim0_start = (dim0_i - starts[0]) * (out_n_rows + 1);
    int64_t out_crows_idx = out_dim0_start + (dim1_i - starts[1]);
    int64_t st = x_crows_data[i] + x_cols_offsets[dim0_i];
    int64_t ed = x_crows_data[i + 1] + x_cols_offsets[dim0_i];
    int64_t index = out_crows_data[out_crows_idx];
    for (int64_t jj = st; jj < ed; ++jj) {
      if (x_cols_data[jj] >= starts[2] && x_cols_data[jj] < ends[2]) {
        out_cols_data[out_cols_offsets[out_dim0_start] + index] =
            x_cols_data[jj] - starts[2];
        out_values_data[out_cols_offsets[out_dim0_start] + index] =
            x_values_data[jj];
        index++;
      }
    }
  }
}

template <typename T, typename Context>
void SliceCsrTensor3D(const Context& dev_ctx,
                      const SparseCsrTensor& x,
                      const std::vector<int64_t>& axes,
                      const std::vector<int64_t>& starts,
                      const std::vector<int64_t>& ends,
                      const DDim& out_dims,
                      SparseCsrTensor* out) {
  const auto* x_crows_data = x.crows().data<int64_t>();
  const auto* x_cols_data = x.cols().data<int64_t>();
  const auto* x_values_data = x.values().data<T>();
  const int64_t x_dim0 = x.dims()[0], x_n_rows = x.dims()[1];

  // get x_cols_offsets
  DenseTensor x_cols_offsets = Empty<int64_t>(dev_ctx, {x_dim0 + 1});
  auto* x_cols_offsets_data = x_cols_offsets.data<int64_t>();

  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x_dim0 + 1, 1);
  GetXColsOffsetsCudaKernel<<<config.block_per_grid.x,
                              config.thread_per_block.x,
                              0,
                              dev_ctx.stream()>>>(
      x_crows_data, x_n_rows, x_dim0, x_cols_offsets_data);

#ifdef PADDLE_WITH_HIP
  thrust::inclusive_scan(thrust::hip::par.on(dev_ctx.stream()),
#else
  thrust::inclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
#endif
                         x_cols_offsets_data,
                         x_cols_offsets_data + x_dim0 + 1,
                         x_cols_offsets_data);

  // copy starts to device
  auto d_starts_tensor = memory_utils::Alloc(
      dev_ctx.GetPlace(),
      sizeof(int64_t) * starts.size(),
      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
  int64_t* d_starts = reinterpret_cast<int64_t*>(d_starts_tensor->ptr());
  memory_utils::Copy(dev_ctx.GetPlace(),
                     d_starts,
                     phi::CPUPlace(),
                     starts.data(),
                     sizeof(int64_t) * starts.size(),
                     dev_ctx.stream());

  // copy ends to device
  auto d_ends_tensor = memory_utils::Alloc(
      dev_ctx.GetPlace(),
      sizeof(int64_t) * ends.size(),
      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
  int64_t* d_ends = reinterpret_cast<int64_t*>(d_ends_tensor->ptr());
  memory_utils::Copy(dev_ctx.GetPlace(),
                     d_ends,
                     phi::CPUPlace(),
                     ends.data(),
                     sizeof(int64_t) * ends.size(),
                     dev_ctx.stream());

  // get out_nnz
  const int64_t out_dim0 = out_dims[0], out_n_rows = out_dims[1];
  DenseTensor out_crows =
      Empty<int64_t, Context>(dev_ctx, {out_dim0 * (out_n_rows + 1)});
  auto* out_crows_data = out_crows.data<int64_t>();
  config = backends::gpu::GetGpuLaunchConfig1D(
      dev_ctx, x_dim0 * (x_n_rows + 1) + 1, 1);
  GetCsr3DNonZeroNumberCudaKernel<<<config.block_per_grid.x,
                                    config.thread_per_block.x,
                                    0,
                                    dev_ctx.stream()>>>(x_crows_data,
                                                        x_cols_data,
                                                        x_dim0,
                                                        x_n_rows,
                                                        x_cols_offsets_data,
                                                        d_starts,
                                                        d_ends,
                                                        out_n_rows,
                                                        out_crows_data);
  DenseTensor out_cols_offsets =
      Empty<int64_t, Context>(dev_ctx, {out_dim0 * (out_n_rows + 1)});
  auto* out_cols_offsets_data = out_cols_offsets.data<int64_t>();
  backends::gpu::GpuMemcpyAsync(out_cols_offsets_data,
                                out_crows_data,
                                out_dim0 * (out_n_rows + 1) * sizeof(int64_t),
                                gpuMemcpyDeviceToDevice,
                                dev_ctx.stream());
  dev_ctx.Wait();
  int64_t out_nnz =
#ifdef PADDLE_WITH_HIP
      thrust::reduce(thrust::hip::par.on(dev_ctx.stream()),
#else
      thrust::reduce(thrust::cuda::par.on(dev_ctx.stream()),
#endif
                     out_crows_data,
                     out_crows_data + out_dim0 * (out_n_rows + 1));
  for (int64_t i = 0; i < out_dim0; ++i) {
    int64_t st = i * (out_n_rows + 1);
#ifdef PADDLE_WITH_HIP
    thrust::inclusive_scan(thrust::hip::par.on(dev_ctx.stream()),
#else
    thrust::inclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
#endif
                           out_crows_data + st,
                           out_crows_data + st + out_n_rows + 1,
                           out_crows_data + st);
  }
#ifdef PADDLE_WITH_HIP
  thrust::inclusive_scan(thrust::hip::par.on(dev_ctx.stream()),
#else
  thrust::inclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
#endif
                         out_cols_offsets_data,
                         out_cols_offsets_data + out_dim0 * (out_n_rows + 1),
                         out_cols_offsets_data);

  DenseTensor out_cols = Empty<int64_t, Context>(dev_ctx, {out_nnz});
  auto* out_cols_data = out_cols.data<int64_t>();
  DenseTensor out_values = Empty<T, Context>(dev_ctx, {out_nnz});
  auto* out_values_data = out_values.data<T>();
  out->SetMember(out_crows, out_cols, out_values, out_dims);
  config = backends::gpu::GetGpuLaunchConfig1D(
      dev_ctx, x_dim0 * (x_n_rows + 1) + 1, 1);
  GetCsr3DCudaKernel<T><<<config.block_per_grid.x,
                          config.thread_per_block.x,
                          0,
                          dev_ctx.stream()>>>(x_crows_data,
                                              x_cols_data,
                                              x_values_data,
                                              x_cols_offsets_data,
                                              x_dim0,
                                              x_n_rows,
                                              d_starts,
                                              d_ends,
                                              out_n_rows,
                                              out_cols_offsets_data,
                                              out_crows_data,
                                              out_cols_data,
                                              out_values_data);
}

template <typename T, typename Context>
void SliceCsrCompute(const Context& dev_ctx,
                     const SparseCsrTensor& x,
                     const std::vector<int64_t>& axes,
                     const std::vector<int64_t>& starts,
                     const std::vector<int64_t>& ends,
                     SparseCsrTensor* out) {
  const DDim& x_dims = x.dims();

  // Step1: Infer output dims
  auto out_dims = funcs::GetSliceDims<int64_t>(
      x_dims, axes, starts, ends, nullptr, nullptr);

  // Step2: Construct new axes, starts and ends.
  std::vector<int64_t> new_axes(3), new_starts(3), new_ends(3);
  funcs::ConstructNewSliceAttrs(
      x_dims, axes, starts, ends, &new_axes, &new_starts, &new_ends);

  // Step3: Slice csr tensor according to its dimension
  if (x_dims.size() == 2) {
    SliceCsrTensor2D<T, Context>(
        dev_ctx, x, new_axes, new_starts, new_ends, out_dims, out);
  } else if (x_dims.size() == 3) {
    SliceCsrTensor3D<T, Context>(
        dev_ctx, x, new_axes, new_starts, new_ends, out_dims, out);
  } else {
    // throw exception
    common::errors::InvalidArgument(
        "Slice for Sparse CSR Tensor only support 2-D or 3-D, but got %d-D.",
        x_dims.size());
  }
}

template <typename T, typename Context>
void SliceCsrKernel(const Context& dev_ctx,
                    const SparseCsrTensor& x,
                    const phi::IntArray& axes,
                    const phi::IntArray& starts,
                    const phi::IntArray& ends,
                    SparseCsrTensor* out) {
  const DDim& x_dims = x.dims();

  std::vector<int64_t> axes_vec = axes.GetData();
  std::vector<int64_t> starts_vec = starts.GetData();
  std::vector<int64_t> ends_vec = ends.GetData();
  // Check and update attr
  funcs::CheckAndUpdateSparseSliceAttrs<int64_t>(
      x_dims, &axes_vec, &starts_vec, &ends_vec);

  SliceCsrCompute<T, Context>(dev_ctx, x, axes_vec, starts_vec, ends_vec, out);
}

}  // namespace sparse
}  // namespace phi

PD_REGISTER_KERNEL(slice_coo,
                   GPU,
                   ALL_LAYOUT,
                   phi::sparse::SliceCooKernel,
                   float,
                   double,
                   int8_t,
                   uint8_t,
                   int16_t,
                   int,
                   int64_t,
                   bool) {}

PD_REGISTER_KERNEL(slice_csr,
                   GPU,
                   ALL_LAYOUT,
                   phi::sparse::SliceCsrKernel,
                   float,
                   double,
                   int8_t,
                   uint8_t,
                   int16_t,
                   int,
                   int64_t,
                   bool) {}