#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
#include <ATen/core/Tensor.h>
#include <ATen/Dispatch.h>
#include <ATen/ExpandUtils.h>
#include <ATen/InitialTensorOptions.h>
#include <ATen/SparseCsrTensorImpl.h>
#include <ATen/SparseCsrTensorUtils.h>
#include <ATen/SparseTensorUtils.h>
#include <ATen/WrapDimUtilsMulti.h>
#include <ATen/native/BinaryOps.h>
#include <ATen/native/Resize.h>
#include <algorithm>

#ifndef AT_PER_OPERATOR_HEADERS
#include <ATen/NativeFunctions.h>
#else
#include <ATen/ops/_convert_indices_from_coo_to_csr_native.h>
#include <ATen/ops/_convert_indices_from_csr_to_coo_native.h>
#include <ATen/ops/_sparse_csr_tensor_unsafe_native.h>
#include <ATen/ops/_unique.h>
#include <ATen/ops/add_native.h>
#include <ATen/ops/resize_as_sparse_native.h>
#include <ATen/ops/tensor.h>
#include <ATen/ops/zeros.h>
#endif

#include <cuda_runtime.h>
#include <type_traits>


#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/CUDAUtils.h>
#include <ATen/cuda/ThrustAllocator.h>
#include <c10/cuda/CUDACachingAllocator.h>

#include <ATen/native/cuda/Reduce.cuh>
#include <ATen/native/sparse/cuda/SparseBlasImpl.h>
#include <ATen/native/sparse/cuda/SparseCUDABlas.h>
#include <ATen/native/sparse/cuda/SparseCUDATensorMath.cuh>

#include <thrust/device_ptr.h>
#include <thrust/execution_policy.h>
#include <thrust/fill.h>
#include <thrust/for_each.h>
#include <thrust/sequence.h>

namespace at {
namespace native {

namespace {

template <typename input_t, typename output_t>
__global__ void convert_indices_from_coo_to_csr_cuda_kernel(output_t* data_out, const input_t* data_in, const int64_t size, const int64_t numel) {
  int64_t tid = blockDim.x * blockIdx.x + threadIdx.x;
  if (tid == 0) {
    for (int64_t i = 0; i <= data_in[0]; i++)
      data_out[i] = static_cast<output_t>(0);
  } else if (tid < numel) {
    for (int64_t i = data_in[tid - 1]; i < data_in[tid]; i++)
      data_out[i + 1] = static_cast<output_t>(tid);
  } else if (tid == numel) {
    for (int64_t i = data_in[numel - 1] + 1; i < size + 1; i++)
      data_out[i] = static_cast<output_t>(numel);
  }
}

template <typename input_t, typename output_t>
void convert_indices_from_coo_to_csr_cuda(const Tensor& result, const Tensor& input, const int64_t size) {
  int64_t numel = input.numel();
  const input_t* data_in = input.data_ptr<input_t>();
  output_t* data_out = result.data_ptr<output_t>();

  if (numel == 0) {
    result.zero_();
    return;
  }

  // Run (numel + 1) threads...
  int64_t THREADS = at::cuda::getCurrentDeviceProperties()->maxThreadsPerBlock;
  int64_t BLOCKS = (numel + THREADS) / THREADS;
  at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream();
  convert_indices_from_coo_to_csr_cuda_kernel<<<BLOCKS, THREADS, 0, stream>>>(data_out, data_in, size, numel);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

template <typename input_t, typename output_t>
__global__ void convert_indices_from_csr_to_coo_cuda_kernel(output_t* data_out, const input_t* data_in, const int64_t nrows) {
  int64_t tid = blockDim.x * blockIdx.x + threadIdx.x;

  if (tid < nrows) {
    for (int64_t i = data_in[tid]; i < data_in[tid + 1]; i++)
      data_out[i] = static_cast<output_t>(tid);
  }
}

template <typename input_t, typename output_t>
void convert_indices_from_csr_to_coo_cuda(const Tensor& indices, const Tensor& crow_indices, const Tensor& col_indices, const bool transpose=false) {
  int64_t nrows = crow_indices.numel() - 1;
  if (nrows == 0) {
    indices.zero_();
    return;
  }

  auto crow_indices_ = crow_indices.expect_contiguous();
  const input_t* crow_indices_data_in = crow_indices_->data_ptr<input_t>();
  TORCH_INTERNAL_ASSERT(indices.is_contiguous());
  auto row0 = indices.select(0, transpose?1:0);
  auto row1 = indices.select(0, transpose?0:1);
  output_t* data_out = row0.data_ptr<output_t>();

  // Run nrows threads...
  int64_t THREADS = at::cuda::getCurrentDeviceProperties()->maxThreadsPerBlock;
  int64_t BLOCKS = (nrows + THREADS) / THREADS;
  at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream();
  row1.copy_(*col_indices.expect_contiguous());
  convert_indices_from_csr_to_coo_cuda_kernel<<<BLOCKS, THREADS, 0, stream>>>(data_out, crow_indices_data_in, nrows);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

} // namespace

using namespace at::sparse_csr;
// certain utiliy functions are usable from sparse COO.
using namespace at::sparse;

Tensor& add_out_dense_sparse_csr_cuda(
    Tensor& output,
    const Tensor& dense,
    const SparseCsrTensor& src,
    const Scalar& alpha) {
  TORCH_INTERNAL_ASSERT(dense.layout() == kStrided);
  TORCH_INTERNAL_ASSERT(src.is_sparse_csr());
  TORCH_INTERNAL_ASSERT(dense.is_cuda());

  TORCH_CHECK(
      output.is_contiguous(),
      "out argument must be contiguous, but got: ",
      output.suggest_memory_format());
  TORCH_CHECK(
      output.is_cuda(),
      "add: expected 'out' to be CUDA tensor, but got tensor on device: ",
      output.device());

  TORCH_CHECK(
      src.is_cuda(),
      "add: expected 'other' to be a CUDA tensor, but got tensor on device: ",
      src.device());

  TORCH_CHECK(
      dense.sizes().equals(src.sizes()),
      "add: expected 'self' and 'other' to have same size, but self has size ",
      dense.sizes(),
      " while other has size ",
      src.sizes(),
      " (FYI: dense-sparse addition does not currently support broadcasting)");

  auto commonDtype = promoteTypes(dense.scalar_type(), src.scalar_type());
  TORCH_CHECK(
      canCast(commonDtype, output.scalar_type()),
      "Can't convert result type ",
      commonDtype,
      " to output ",
      output.scalar_type(),
      " in add operation");

  Tensor src_values = src.values();

  resize_output(output, dense.sizes());

  Tensor resultBuffer = output;

  if (output.scalar_type() != commonDtype) {
    resultBuffer = dense.to(commonDtype);
  } else if (!is_same_tensor(output, dense)) {
    resultBuffer.copy_(dense);
  }

  if (src._nnz() == 0) {
    return output;
  }

  auto valuesBuffer = src_values.to(commonDtype).view({-1, src_values.size(-1)});
  resultBuffer = resultBuffer.view({-1, output.size(-2), output.size(-1)});
  auto src_crow_indices = src.crow_indices().view({-1, src.crow_indices().size(-1)});
  auto src_col_indices = src.col_indices().view({-1, src.col_indices().size(-1)});

  AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND4(
      kComplexHalf, kHalf, kBool, kBFloat16,
      commonDtype,
      "add_out_op2_sparse_csr",
      [&valuesBuffer, &resultBuffer, &alpha, &src_crow_indices, &src_col_indices]() {
        AT_DISPATCH_INDEX_TYPES(
            src_crow_indices.scalar_type(),
            "csr_add_out_crow_indices",
              [&valuesBuffer, &resultBuffer, &alpha, &src_crow_indices, &src_col_indices]() {
                auto batch_count = resultBuffer.dim() > 2 ? resultBuffer.size(-3) : 1;
                scalar_t* values_accessor = valuesBuffer.data_ptr<scalar_t>();
                scalar_t* out_ptr = resultBuffer.data_ptr<scalar_t>();
                scalar_t cast_value = alpha.to<scalar_t>();

                index_t* crow_indices_accessor = src_crow_indices.data_ptr<index_t>();
                index_t* col_indices_accessor = src_col_indices.data_ptr<index_t>();
                int64_t out_storage_offset = resultBuffer.storage_offset();

                auto out_strides = resultBuffer.strides();
                auto out_strides0 = out_strides[0];
                auto out_strides1 = out_strides[1];
                auto crow_stride0 = src_crow_indices.stride(0);
                auto col_stride0 = src_col_indices.stride(0);
                auto val_stride0 = valuesBuffer.stride(0);

                cudaStream_t stream = at::cuda::getCurrentCUDAStream();
                at::cuda::ThrustAllocator allocator;
                auto policy = thrust::cuda::par(allocator).on(stream);

               // Note that this could be wildly imbalanced if the sparsity pattern varies a lot between rows.
               thrust::for_each(
                    policy,
                    thrust::make_counting_iterator(int64_t(0)),
                    thrust::make_counting_iterator(int64_t(src_crow_indices.size(-1) - 1)),
                    [values_accessor,
                    crow_indices_accessor,
                    col_indices_accessor,
                    out_ptr,
                    cast_value,
                    out_strides0,
                    out_strides1,
                    crow_stride0,
                    col_stride0,
                    val_stride0,
                    batch_count
                    ]__device__(int64_t irow) {
                      for (index_t batch_idx = 0; batch_idx < batch_count; batch_idx++) {
                        index_t start_index = crow_indices_accessor[batch_idx*crow_stride0 + irow];
                        index_t end_index = crow_indices_accessor[batch_idx*crow_stride0 + irow + 1];

                        for (index_t i = start_index; i < end_index; ++i) {
                            auto icol = col_indices_accessor[batch_idx*col_stride0 + i];
                            auto index = batch_idx * out_strides0 + irow * out_strides1 + icol;
                            out_ptr[index] += cast_value * values_accessor[batch_idx*val_stride0 + i];
                        }
                      }
                    });
              });
      });
  if (output.scalar_type() != commonDtype) {
    output.copy_(resultBuffer);
  }
  return output;
}

Tensor& add_out_sparse_csr_cuda(
    const Tensor& self,
    const SparseCsrTensor& other,
    const Scalar& alpha,
    SparseCsrTensor& out) {
  if (self.layout() == kStrided) {
    add_out_dense_sparse_csr_cuda(out, self, other, alpha);
  } else {
    TORCH_CHECK(
        self.sizes().equals(other.sizes()),
        "torch.add: Expected input tensors to have the same shape, but got tensor `self` with shape ",
        self.sizes(),
        " and tensor `other` with shape ",
        other.sizes());
    at::native::resize_as_sparse_compressed_(out, self);
    sparse::impl::cuda::add_out_sparse_csr(self, other, Scalar(1), alpha, out);
  }
  return out;
}

TORCH_IMPL_FUNC(_convert_indices_from_coo_to_csr_structured_cuda) (
  const Tensor& input, const int64_t size, const bool out_int32, const Tensor& result
) {
  if (out_int32) {
    AT_DISPATCH_INTEGRAL_TYPES(input.scalar_type(), "convert_indices_from_coo_to_csr_cuda", [&] {
      convert_indices_from_coo_to_csr_cuda<scalar_t, int>(result, input, size);
    });
  } else {
    AT_DISPATCH_INTEGRAL_TYPES(input.scalar_type(), "convert_indices_from_coo_to_csr_cuda", [&] {
      convert_indices_from_coo_to_csr_cuda<scalar_t, int64_t>(result, input, size);
    });
  }
}

TORCH_IMPL_FUNC(_convert_indices_from_csr_to_coo_structured_cuda) (
  const Tensor& crow_indices, const Tensor& col_indices, const bool out_int32, const bool transpose, const Tensor& result
) {
  if (out_int32) {
    AT_DISPATCH_INTEGRAL_TYPES(crow_indices.scalar_type(), "convert_indices_from_csr_to_coo_cuda", [&] {
      convert_indices_from_csr_to_coo_cuda<scalar_t, int32_t>(result, crow_indices, col_indices, transpose);
    });
  } else {
    AT_DISPATCH_INTEGRAL_TYPES(crow_indices.scalar_type(), "convert_indices_from_csr_to_coo_cuda", [&] {
      convert_indices_from_csr_to_coo_cuda<scalar_t, int64_t>(result, crow_indices, col_indices, transpose);
    });
  }
}

  /*
    Reductions on sparse CSR tensors using masked semantics.

    - To support a reduction operator on a CSR tensor with CUDA storage, define

template <typename scalar_t>
struct Reduction...Op {
  __device__ __forceinline__ scalar_t operator()(const scalar_t a, const scalar_t b) const {
    return a ... b;
  }
  __device__ __forceinline__ scalar_t identity() const { return ...; }
  __forceinline__ scalar_t identity_cpu() const { return ...; }
};


Tensor _sparse_csr_..._cuda(const Tensor& input, IntArrayRef dims_to_sum, bool keepdim, c10::optional<ScalarType> dtype) {
  ...
      result = reduce_sparse_csr_cuda_template<scalar_t>(input_, dims_to_sum, keepdim, Reduction...Op<scalar_t>());
  ...
  return result;
}

      and add the following

        - func: _sparse_csr_op.dim_dtype(Tensor self, int[1] dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor
          dispatch:
            SparseCsrCUDA: _sparse_csr_..._cuda

      to native_functions.yaml
  */

namespace {

template <typename scalar_t, typename index_t, typename ReductionOp>
__global__ void reduce_sparse_csr_dim0_cuda_kernel(scalar_t* new_values,
                                                   const index_t* new_col_indices,
                                                   const int64_t new_nnz,
                                                   const scalar_t* values,
                                                   const index_t* col_indices,
                                                   const int64_t nnz,
                                                   ReductionOp rop
                                                   ) {
  int64_t tid = blockDim.x * blockIdx.x + threadIdx.x;
  if (tid < new_nnz) {
    index_t col = new_col_indices[tid];
    scalar_t v = rop.identity();
    for (int64_t j=0; j < nnz; j++) {
      if (col == col_indices[j]) {
        v = rop(v, values[j]);
      }
    }
    new_values[tid] = v;
  }
}

template <typename scalar_t, typename ReductionOp>
Tensor reduce_sparse_csr_dim0_cuda_template(const Tensor& sparse, ReductionOp rop) {
  /*
    Consider the following sparse tensor:

      1 * * * *
      * * * 2 *
      * * 3 * *
      * * * * *
      4 * 5 * *

    that has CSR representation

      crow_indices = [0, 1, 2, 3, 3, 5]
      col_indices = [0, 3, 2, 0, 2]
      values = [1, 2, 3, 4, 5]

    Reduction with dim=0 results:

      rop(1,4) * rop(3,5) 2 *

    that has CSR representation

      new_crow_indices = [0, 3]
      new_col_indices = [0, 2, 3]
      new_values = [rop(1, 4], rop(3, 5), 2]

    In general, the CSR representation data can be computed as follows:

      nnz = col_indices.numel()
      new_col_indices = col_indices.unique(sorted=True, return_inverse=False)
      new_nnz = new_col_indices.numel()
      new_crow_indices = [0, new_nnz]
      new_values.resize(new_nnz)

      for i in range(new_nnz):
          v = identity
          col = new_col_indices[i]
          for j in range(nnz):
              if col == col_indices[j]:
                  v = rop(v, values[j])
          new_values[i] = v

    Notice this algorithm is different from the one used on CPU data.
  */

  Tensor col_indices = sparse.col_indices();
  Tensor values = sparse.values();
  auto ncols = sparse.size(1);
  auto nnz = col_indices.numel();
  Tensor new_col_indices;

  std::tie(new_col_indices, std::ignore) = at::_unique(col_indices, true, false);
  auto new_nnz = new_col_indices.numel();
  Tensor new_crow_indices = at::tensor(ArrayRef<int64_t>{0, new_nnz}, col_indices.options());
  Tensor new_values = at::empty({new_nnz}, values.options());

  scalar_t* values_ptr = values.data_ptr<scalar_t>();
  scalar_t* new_values_ptr = new_values.data_ptr<scalar_t>();
  int64_t THREADS = at::cuda::getCurrentDeviceProperties()->maxThreadsPerBlock;
  int64_t BLOCKS = (new_nnz + THREADS) / THREADS;
  at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream();
  AT_DISPATCH_INDEX_TYPES(col_indices.scalar_type(), "reduce_sparse_csr_dim0_cuda_indices",
                          [&]() {
                            index_t* col_indices_ptr = col_indices.data_ptr<index_t>();
                            index_t* new_col_indices_ptr = new_col_indices.data_ptr<index_t>();
                            reduce_sparse_csr_dim0_cuda_kernel<<<BLOCKS, THREADS, 0, stream>>>(new_values_ptr,
                                                                                               new_col_indices_ptr,
                                                                                               new_nnz,
                                                                                               values_ptr,
                                                                                               col_indices_ptr,
                                                                                               nnz,
                                                                                               rop
                                                                                               );
                          });
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return at::native::_sparse_csr_tensor_unsafe(new_crow_indices, new_col_indices, new_values,
                                               {1, ncols},
                                               new_values.scalar_type(),
                                               sparse.layout(),
                                               new_values.device());
}

template <typename index_t>
__global__ void reduce_crow_indices_dim1_cuda_kernel(index_t* new_crow_indices,
                                                     index_t* row_map,
                                                     const index_t* crow_indices,
                                                     const int64_t nrows
                                                     ) {
  int64_t nnz = 0;
  new_crow_indices[0] = 0;
  for(int64_t i=0; i<nrows; i++) {
    if (crow_indices[i] != crow_indices[i + 1]) {
      row_map[i] = nnz;
      nnz++;
    }
    new_crow_indices[i + 1] = nnz;
  }
}

template <typename scalar_t, typename index_t, typename ReductionOp>
__global__ void reduce_sparse_csr_dim1_cuda_kernel(scalar_t* new_values,
                                                   const scalar_t* values,
                                                   const index_t* crow_indices,
                                                   const index_t* row_map,
                                                   const int64_t nrows,
                                                   ReductionOp rop
                                                   ) {
  int64_t tid = blockDim.x * blockIdx.x + threadIdx.x;
  if (tid < nrows) {
    index_t i_start = crow_indices[tid];
    index_t i_end = crow_indices[tid+1];
    if (i_start != i_end) {
      scalar_t acc = rop.identity();
      for (index_t i = i_start; i < i_end; i++) {
        acc = rop(acc, values[i]);
      }
      new_values[row_map[tid]] = acc;
    }
  }
}

template <typename scalar_t, typename ReductionOp>
Tensor reduce_sparse_csr_dim1_cuda_template(const Tensor& sparse, ReductionOp rop) {
  /*
    The algorithm of computing reduce of a CSR tensor along the last
    dimension is explained in the comment of the
    reduce_sparse_csr_dim1_cpu_template function.
  */
  Tensor crow_indices = sparse.crow_indices();
  auto ioptions = crow_indices.options();
  Tensor values = sparse.values();
  auto nrows = sparse.size(0);
  auto numel = values.numel();

  Tensor new_crow_indices = at::empty({crow_indices.numel()}, ioptions);
  Tensor new_col_indices = at::empty({}, ioptions);
  Tensor new_values = at::empty({}, values.options());
  Tensor row_map = at::empty({nrows}, ioptions);

  at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream();
  int64_t THREADS = at::cuda::getCurrentDeviceProperties()->maxThreadsPerBlock;
  int64_t BLOCKS = (nrows + THREADS) / THREADS;

  AT_DISPATCH_INDEX_TYPES(crow_indices.scalar_type(), "reduce_sparse_csr_dim1_cuda_indices",
                          [&]() {
                            index_t* crow_indices_ptr = crow_indices.data_ptr<index_t>();
                            index_t* new_crow_indices_ptr = new_crow_indices.data_ptr<index_t>();
                            index_t* row_map_ptr = row_map.data_ptr<index_t>();
                            reduce_crow_indices_dim1_cuda_kernel<<<1, 1, 0, stream>>>(new_crow_indices_ptr,
                                                                                      row_map_ptr,
                                                                                      crow_indices_ptr,
                                                                                      nrows);
                            C10_CUDA_KERNEL_LAUNCH_CHECK();
                            index_t new_nnz = new_crow_indices[-1].item<index_t>();
                            new_col_indices.resize_(new_nnz);
                            new_col_indices.fill_(index_t(0));
                            new_values.resize_(new_nnz);

                            scalar_t* values_ptr = values.data_ptr<scalar_t>();
                            scalar_t* new_values_ptr = new_values.data_ptr<scalar_t>();
                            reduce_sparse_csr_dim1_cuda_kernel<<<BLOCKS, THREADS, 0, stream>>>(new_values_ptr,
                                                                                               values_ptr,
                                                                                               crow_indices_ptr,
                                                                                               row_map_ptr,
                                                                                               nrows,
                                                                                               rop);
                            C10_CUDA_KERNEL_LAUNCH_CHECK();
                          });

  return at::native::_sparse_csr_tensor_unsafe(new_crow_indices, new_col_indices, new_values,
                                               {sparse.size(0), 1},
                                               new_values.scalar_type(),
                                               sparse.layout(),
                                               new_values.device());
}

template <typename scalar_t, typename ReductionOp>
Tensor reduce_sparse_csr_dim01_cuda_template(const Tensor& sparse, ReductionOp rop) {

  auto ioptions = sparse.col_indices().options();
  Tensor values = sparse.values();
  auto numel = values.numel();
  auto nnz = std::min<int64_t>(1, numel);

  Tensor new_values;
  if (numel > 0) {
    new_values = at::empty({1}, values.options());
    auto iter = TensorIterator::reduce_op(new_values, values);
    gpu_reduce_kernel<scalar_t, scalar_t>(iter, func_wrapper<scalar_t>(rop), rop.identity_cpu());
  } else {
    new_values = at::empty({}, values.options());
  }
  Tensor new_col_indices = at::zeros({nnz}, ioptions);
  Tensor new_crow_indices = at::tensor(ArrayRef<int64_t>{0, nnz}, ioptions);
  return at::native::_sparse_csr_tensor_unsafe(new_crow_indices, new_col_indices, new_values,
                                               {1, std::min<int64_t>(1, sparse.size(1))},
                                               new_values.scalar_type(),
                                               sparse.layout(),
                                               new_values.device());
}

template <typename scalar_t, typename ReductionOp>
Tensor reduce_sparse_csr_cuda_template(const Tensor& sparse, std::vector<int64_t> dims, ReductionOp rop) {
  if (dims.size() == 1) {
    if (dims[0] == 0) {
      return reduce_sparse_csr_dim0_cuda_template<scalar_t>(sparse, rop);
    } else {
      TORCH_INTERNAL_ASSERT(dims[0] == 1);
      return reduce_sparse_csr_dim1_cuda_template<scalar_t>(sparse, rop);
    }
  } else if (dims.size() == 2) {
    TORCH_INTERNAL_ASSERT(((dims[0] == 0 && dims[1] == 1) || (dims[0] == 1 && dims[1] == 0)));
    return reduce_sparse_csr_dim01_cuda_template<scalar_t>(sparse, rop);
  }
  TORCH_INTERNAL_ASSERT(dims.size() == 0);
  // effective after gh-29137 has been resolved
  return sparse.clone();
}

template <typename scalar_t, typename ReductionOp>
Tensor reduce_sparse_csr_cuda_template(const Tensor& sparse, IntArrayRef dims_to_sum, bool keepdim, ReductionOp rop) {
  TORCH_INTERNAL_ASSERT(sparse.is_sparse_csr());
  TORCH_CHECK(keepdim, "reduction operations on CSR tensors with keepdim=False is unsupported");
  TORCH_INTERNAL_ASSERT(sparse.is_cuda());

  const int64_t input_dim = sparse.dim();
  TORCH_INTERNAL_ASSERT(input_dim == 2);
  auto dims = dims_to_sum.vec();
  maybe_wrap_dims(dims, input_dim);
  if (dims.size() == 0) {
    // after gh-29137 is resolved, delete this if-block
    dims.emplace_back(0);
    dims.emplace_back(1);
  }
  return reduce_sparse_csr_cuda_template<scalar_t>(sparse, dims, rop);
}

template <typename scalar_t>
struct ReductionAddOp {
  __device__ __forceinline__ scalar_t operator()(const scalar_t a, const scalar_t b) const {
    return a + b;
  }
  __device__ __forceinline__ scalar_t identity() const { return 0; }
  __forceinline__ scalar_t identity_cpu() const { return 0; }
};

template <typename scalar_t>
struct ReductionMulOp {
  __device__ __forceinline__ scalar_t operator()(const scalar_t a, const scalar_t b) const {
    return a * b;
  }
  __device__ __forceinline__ scalar_t identity() const { return 1; }
  __forceinline__ scalar_t identity_cpu() const { return 1; }
};

} // namespace

Tensor _sparse_csr_sum_cuda(const Tensor& input, IntArrayRef dims_to_sum, bool keepdim, c10::optional<ScalarType> dtype) {
  ScalarType dtype_ = dtype.value_or(input.scalar_type());
  Tensor input_ = input.to(dtype_);
  Tensor result;
  AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(
    kHalf, kBFloat16, input_.scalar_type(), "_sparse_csr_sum_cuda",
    [&] {
      result = reduce_sparse_csr_cuda_template<scalar_t>(input_, dims_to_sum, keepdim, ReductionAddOp<scalar_t>());
    });
  return result;
}

Tensor _sparse_csr_prod_cuda(const Tensor& input, IntArrayRef dims_to_reduce, bool keepdim, c10::optional<ScalarType> dtype) {
  ScalarType dtype_ = dtype.value_or(input.scalar_type());
  Tensor input_ = input.to(dtype_);
  Tensor result;
  AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(
    kHalf, kBFloat16, input_.scalar_type(), "_sparse_csr_prod_cuda",
    [&] {
      result = reduce_sparse_csr_cuda_template<scalar_t>(input_, dims_to_reduce, keepdim, ReductionMulOp<scalar_t>());
    });
  return result;
}

} // namespace native
} // namespace at
