Support CumSum op in the CUDA EP (#2647)

* Initial commit

* Initial commit

* Updates

* Fix build

* Updates

* PR feedback

* Minor optimization

* Update

* Update

onnxruntime/core/providers/cpu/math/cumsum.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "cumsum.h"
+#include "core/providers/common.h"
 #include "core/providers/cpu/tensor/utils.h"
 #include "core/framework/op_kernel.h"
 #include "core/framework/tensorprotoutils.h"
@@ -50,10 +51,60 @@ void SumSlices(const Tensor& input, Tensor& output,
 
 namespace onnxruntime {
 
-ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum, 11, float, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()), CumSum<float>);
-ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum, 11, double, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<double>()), CumSum<double>);
-ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum, 11, int32_t, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<int32_t>()), CumSum<int32_t>);;
-ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum, 11, int64_t, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<int64_t>()), CumSum<int64_t>);
+namespace cumsum_op {
+Status GetAxis(const Tensor* axis_tensor, int64_t input_rank, int64_t& axis_out) {
+  if (!axis_tensor)
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Axis tensor must be provided to the CumSum op");
+
+  if (axis_tensor->Shape().NumDimensions() > 1)
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Axis tensor should be 0D or 1D");
+
+  if (axis_tensor->IsDataType<int32_t>()) {
+    axis_out = static_cast<int64_t>(axis_tensor->template Data<int32_t>()[0]);
+  } else if (axis_tensor->IsDataType<int64_t>()) {
+    axis_out = axis_tensor->template Data<int64_t>()[0];
+  } else {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Axis tensor should be of type `int32_t` or `int64_t`");
+  }
+
+  axis_out = HandleNegativeAxis(axis_out, input_rank);
+
+  return Status::OK();
+}
+
+}  // namespace cumsum_op
+
+ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum,
+                               11,
+                               float,
+                               KernelDefBuilder()
+                                   .TypeConstraint("T", DataTypeImpl::GetTensorType<float>())
+                                   .TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(), DataTypeImpl::GetTensorType<int64_t>()}),
+                               CumSum<float>);
+
+ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum,
+                               11,
+                               double,
+                               KernelDefBuilder()
+                                   .TypeConstraint("T", DataTypeImpl::GetTensorType<double>())
+                                   .TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(), DataTypeImpl::GetTensorType<int64_t>()}),
+                               CumSum<double>);
+
+ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum,
+                               11,
+                               int32_t,
+                               KernelDefBuilder()
+                                   .TypeConstraint("T", DataTypeImpl::GetTensorType<int32_t>())
+                                   .TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(), DataTypeImpl::GetTensorType<int64_t>()}),
+                               CumSum<int32_t>);
+
+ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum,
+                               11,
+                               int64_t,
+                               KernelDefBuilder()
+                                   .TypeConstraint("T", DataTypeImpl::GetTensorType<int64_t>())
+                                   .TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(), DataTypeImpl::GetTensorType<int64_t>()}),
+                               CumSum<int64_t>);
 
 template <typename T>
 CumSum<T>::CumSum(const OpKernelInfo& info) : OpKernel(info), exclusive_(), reverse_() {
@@ -79,19 +130,13 @@ CumSum<T>::CumSum(const OpKernelInfo& info) : OpKernel(info), exclusive_(), reve
 
 template <typename T>
 Status CumSum<T>::Compute(OpKernelContext* ctx) const {
-  const Tensor* input = ctx->Input<Tensor>(0);                             // input tensor
-  const auto rank = static_cast<int64_t>(input->Shape().NumDimensions());  // the rank of the input/output
-  const Tensor* axis_tensor = ctx->Input<Tensor>(1);                       // axis input tensor
+  const Tensor* input = ctx->Input<Tensor>(0);                       // input tensor
+  auto rank = static_cast<int64_t>(input->Shape().NumDimensions());  // the rank of the input/output
+  if (rank == 0)
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Cannot apply CumSum operator on a scalar");
 
-  if (axis_tensor->Shape().NumDimensions() > 1)
-    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Axis tensor should be 0D or 1D");
+  const Tensor* axis_tensor = ctx->Input<Tensor>(1);  // axis input tensor
 
-  int32_t axis = axis_tensor->template Data<int32_t>()[0];  // the axis on which the accumulation is going to done
-  // validate input
-  if (axis < -rank || axis >= rank)
-    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Axis should be in the range [", -rank, ",", rank, ") but got: ", axis);
-  if (axis < 0)
-    axis = static_cast<int32_t>(rank) + axis;
   TensorShape output_shape(input->Shape());
   auto& output_tensor = *ctx->Output(0, output_shape);  // output tensor
 
@@ -99,6 +144,9 @@ Status CumSum<T>::Compute(OpKernelContext* ctx) const {
   if (output_shape.Size() == 0)
     return Status::OK();
 
+  int64_t axis;
+  ORT_THROW_IF_ERROR(cumsum_op::GetAxis(axis_tensor, rank, axis));
+
   auto dim(output_tensor.Shape()[axis]);    // dimension size for the axis
   TensorShape slice_shape(input->Shape());  // the shape of one slice of input/output for the given value of the axis
   slice_shape[axis] = 1;

onnxruntime/core/providers/cpu/math/cumsum.h

@@ -3,7 +3,6 @@
 
 #include "core/common/common.h"
 #include "core/framework/op_kernel.h"
-#include "core/providers/cpu/tensor/pad.h"
 
 namespace onnxruntime {
 
@@ -19,4 +18,9 @@ class CumSum final : public OpKernel {
   int64_t reverse_;
 };
 
+namespace cumsum_op {
+
+Status GetAxis(const Tensor* axis_tensor, int64_t input_rank, int64_t& axis_out);
+
+}  // namespace cumsum_op
 }  // namespace onnxruntime

onnxruntime/core/providers/cuda/cuda_execution_provider.cc

@@ -667,6 +667,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain,
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, bool, Equal);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, int32_t, Equal);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, int64_t, Equal);
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, CumSum);
 
 static void RegisterCudaKernels(KernelRegistry& kernel_registry) {
   static const BuildKernelCreateInfoFn function_table[] = {
@@ -1115,7 +1116,7 @@ static void RegisterCudaKernels(KernelRegistry& kernel_registry) {
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, MLFloat16, AveragePool)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, float, MaxPool)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, double, MaxPool)>,
-      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, MLFloat16, MaxPool)>,      
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, MLFloat16, MaxPool)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, float, Resize)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, double, Resize)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, MLFloat16, Resize)>,
@@ -1128,6 +1129,7 @@ static void RegisterCudaKernels(KernelRegistry& kernel_registry) {
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, bool, Equal)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, int32_t, Equal)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, int64_t, Equal)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, CumSum)>,
   };
 
   for (auto& function_table_entry : function_table) {

onnxruntime/core/providers/cuda/math/cumsum.cc

@@ -0,0 +1,134 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "cumsum.h"
+#include "cumsum_impl.h"
+#include "core/providers/cpu/math/cumsum.h"
+#include "core/providers/common.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+ONNX_OPERATOR_KERNEL_EX(
+    CumSum,
+    kOnnxDomain,
+    11,
+    kCudaExecutionProvider,
+    KernelDefBuilder()
+        .InputMemoryType<OrtMemTypeCPUInput>(1)  // 'axis' needs to be on CPU
+        .TypeConstraint("T", std::vector<MLDataType>{
+                                 DataTypeImpl::GetTensorType<int32_t>(),
+                                 DataTypeImpl::GetTensorType<int64_t>(),
+                                 DataTypeImpl::GetTensorType<uint32_t>(),
+                                 DataTypeImpl::GetTensorType<uint64_t>(),
+                                 DataTypeImpl::GetTensorType<float>(),
+                                 DataTypeImpl::GetTensorType<double>(),
+                                 DataTypeImpl::GetTensorType<MLFloat16>()})
+        .TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(), DataTypeImpl::GetTensorType<int64_t>()}),
+    CumSum);
+
+Status CumSum::ComputeInternal(OpKernelContext* ctx) const {
+  const Tensor* input = ctx->Input<Tensor>(0);                       // input tensor
+  auto rank = static_cast<int64_t>(input->Shape().NumDimensions());  // the rank of the input/output
+  if (rank == 0)
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Cannot apply CumSum operator on a scalar");
+
+  const Tensor* axis_tensor = ctx->Input<Tensor>(1);  // axis input tensor
+
+  int64_t axis = 0;
+  ORT_THROW_IF_ERROR(cumsum_op::GetAxis(axis_tensor, rank, axis));
+
+  TensorShape output_shape(input->Shape());
+  auto& output = *ctx->Output(0, output_shape);  // output tensor
+
+  // output tensor's size is 0, nothing to fill - return
+  if (output_shape.Size() == 0)
+    return Status::OK();
+
+  const auto& input_dims = input->Shape().GetDims();
+
+  int64_t current_dim = rank - 1;
+  int64_t input_stride_along_axis = 1;
+
+  // axis (and by extension current_dim) can never be negative as this is validated much before
+  // so no need to add the extra check to make sure current_dim is within bounds of the vector size
+  while (current_dim > axis) {
+    input_stride_along_axis *= input_dims[current_dim--];
+  }
+
+  fast_divmod fast_divmod_input_dim_along_axis(static_cast<int>(input_dims[axis]));
+  fast_divmod fast_divmod_input_stride_along_axis(static_cast<int>(input_stride_along_axis));
+
+  if (input->IsDataType<float>()) {
+    CumSumImpl(reinterpret_cast<const typename ToCudaType<float>::MappedType*>(input->Data<float>()),
+               fast_divmod_input_dim_along_axis,
+               fast_divmod_input_stride_along_axis,
+               reinterpret_cast<typename ToCudaType<float>::MappedType*>(output.MutableData<float>()),
+               output_shape.Size(),
+               input->DataType()->Size(),
+               exclusive_,
+               reverse_);
+  } else if (input->IsDataType<double>()) {
+    CumSumImpl(reinterpret_cast<const typename ToCudaType<double>::MappedType*>(input->Data<double>()),
+               fast_divmod_input_dim_along_axis,
+               fast_divmod_input_stride_along_axis,
+               reinterpret_cast<typename ToCudaType<double>::MappedType*>(output.MutableData<double>()),
+               output_shape.Size(),
+               input->DataType()->Size(),
+               exclusive_,
+               reverse_);
+  } else if (input->IsDataType<int32_t>()) {
+    CumSumImpl(reinterpret_cast<const typename ToCudaType<int32_t>::MappedType*>(input->Data<int32_t>()),
+               fast_divmod_input_dim_along_axis,
+               fast_divmod_input_stride_along_axis,
+               reinterpret_cast<typename ToCudaType<int32_t>::MappedType*>(output.MutableData<int32_t>()),
+               output_shape.Size(),
+               input->DataType()->Size(),
+               exclusive_,
+               reverse_);
+  } else if (input->IsDataType<int64_t>()) {
+    CumSumImpl(reinterpret_cast<const typename ToCudaType<int64_t>::MappedType*>(input->Data<int64_t>()),
+               fast_divmod_input_dim_along_axis,
+               fast_divmod_input_stride_along_axis,
+               reinterpret_cast<typename ToCudaType<int64_t>::MappedType*>(output.MutableData<int64_t>()),
+               output_shape.Size(),
+               input->DataType()->Size(),
+               exclusive_,
+               reverse_);
+  } else if (input->IsDataType<uint32_t>()) {
+    CumSumImpl(reinterpret_cast<const typename ToCudaType<uint32_t>::MappedType*>(input->Data<uint32_t>()),
+               fast_divmod_input_dim_along_axis,
+               fast_divmod_input_stride_along_axis,
+               reinterpret_cast<typename ToCudaType<uint32_t>::MappedType*>(output.MutableData<uint32_t>()),
+               output_shape.Size(),
+               input->DataType()->Size(),
+               exclusive_,
+               reverse_);
+  } else if (input->IsDataType<uint64_t>()) {
+    CumSumImpl(reinterpret_cast<const typename ToCudaType<uint64_t>::MappedType*>(input->Data<uint64_t>()),
+               fast_divmod_input_dim_along_axis,
+               fast_divmod_input_stride_along_axis,
+               reinterpret_cast<typename ToCudaType<uint64_t>::MappedType*>(output.MutableData<uint64_t>()),
+               output_shape.Size(),
+               input->DataType()->Size(),
+               exclusive_,
+               reverse_);
+  } else if (input->IsDataType<MLFloat16>()) {
+    CumSumImpl(reinterpret_cast<const typename ToCudaType<MLFloat16>::MappedType*>(input->Data<MLFloat16>()),
+               fast_divmod_input_dim_along_axis,
+               fast_divmod_input_stride_along_axis,
+               reinterpret_cast<typename ToCudaType<MLFloat16>::MappedType*>(output.MutableData<MLFloat16>()),
+               output_shape.Size(),
+               input->DataType()->Size(),
+               exclusive_,
+               reverse_);
+  } else {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Unsupported input data type to the CumSum op: ",
+                           input->DataType());
+  }
+
+  return Status::OK();
+}
+
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/core/providers/cuda/math/cumsum.h

@@ -0,0 +1,48 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+#pragma once
+
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+#include "core/providers/cuda/cuda_common.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+class CumSum final : public CudaKernel {
+ public:
+  explicit CumSum(const OpKernelInfo& info) : CudaKernel(info) {
+    // Process exclusive attribute
+    int64_t exclusive = 0;
+    auto status = info.GetAttr("exclusive", &exclusive);
+    if (status.IsOK()) {
+      if (exclusive == 1 || exclusive == 0) {
+        exclusive_ = (exclusive == 1);
+      } else {
+        ORT_ENFORCE("attribute exclusive can only be 0 or 1");
+      }
+    }
+
+    // Process reverse attribute
+    int64_t reverse = 0;
+    status = info.GetAttr("reverse", &reverse);
+    if (status.IsOK()) {
+      if (reverse == 1 || reverse == 0) {
+        reverse_ = (reverse == 1);
+      } else {
+        ORT_ENFORCE("attribute reverse can only be 0 or 1");
+      }
+    }
+  }
+
+  ~CumSum() = default;
+
+  Status ComputeInternal(OpKernelContext* ctx) const override;
+
+ private:
+  bool exclusive_ = false;
+  bool reverse_ = false;
+};
+
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/core/providers/cuda/math/cumsum_impl.cu

@@ -0,0 +1,167 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/cuda/cu_inc/common.cuh"
+#include "core/providers/cuda/shared_inc/fast_divmod.h"
+
+#include "cumsum_impl.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+template <typename T>
+__global__ void _CumSumKernel(
+    const T* input_data,
+    const fast_divmod fast_divmod_input_dim_along_axis,
+    const fast_divmod fast_divmod_input_stride_along_axis,
+    T* output_data,
+    const int64_t output_size,
+    const bool exclusive,
+    const bool reverse) {
+  CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(indices_index, output_size);
+
+  int input_dim_along_axis = fast_divmod_input_dim_along_axis.d_;
+  int input_stride_along_axis = fast_divmod_input_stride_along_axis.d_;
+
+  int axis_dim = 0;
+  int div = fast_divmod_input_stride_along_axis.div(static_cast<int>(indices_index));
+  fast_divmod_input_dim_along_axis.divmod(div, div, axis_dim);
+
+  int start = 0;
+  int end = 0;
+
+  if (!reverse && !exclusive) {
+    start = 0;
+    end = axis_dim;
+  
+  } else if (reverse && !exclusive) {
+    start = axis_dim;
+    end = input_dim_along_axis - 1;
+
+  } else if (!reverse && exclusive) {
+    start = 0;
+    end = axis_dim - 1;
+
+  } else { // reverse && exclusive
+    start = axis_dim + 1;
+    end = input_dim_along_axis - 1;
+
+  }
+
+  // count the number of elements to accumulate the sum
+  int count = end - start + 1;
+  if (count <= 0) {
+    output_data[indices_index] = 0;
+    return;  
+  }
+
+  // adjust start index based on the above identified start dim value along the axis of interest
+  int data_index = static_cast<int>(indices_index) + (start - axis_dim) * input_stride_along_axis;
+  T sum = 0;
+
+  // keep accumulating values from the start index for 'count' times and skip appropriately 
+  while (count != 0) {
+    sum += input_data[data_index];
+    data_index += input_stride_along_axis;
+    --count;
+  }
+
+  output_data[indices_index] = sum;
+}
+
+template<typename T>
+void CumSumImpl(
+    const T* input_data,
+    const fast_divmod& input_dim_along_axis,
+    const fast_divmod& input_stride_along_axis,
+    T* output_data,
+    const int64_t output_size,
+    const size_t element_size,
+    const bool exclusive,
+    const bool reverse) {
+  if (output_size > 0) {
+    int blocksPerGrid = static_cast<int>((output_size + GridDim::maxThreadsPerBlock - 1) / GridDim::maxThreadsPerBlock);
+
+    _CumSumKernel<T><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(input_data,
+                                                                        input_dim_along_axis,
+                                                                        input_stride_along_axis,
+                                                                        output_data,
+                                                                        output_size,
+                                                                        exclusive,
+                                                                        reverse);
+  }
+}
+
+template void CumSumImpl<int32_t>(
+    const int32_t* input_data,
+    const fast_divmod& input_dim_along_axis,
+    const fast_divmod& input_stride_along_axis,
+    int32_t* output_data,
+    const int64_t output_size,
+    const size_t element_size,
+    const bool exclusive,
+    const bool reverse);
+
+template void CumSumImpl<int64_t>(
+    const int64_t* input_data,
+    const fast_divmod& input_dim_along_axis,
+    const fast_divmod& input_stride_along_axis,
+    int64_t* output_data,
+    const int64_t output_size,
+    const size_t element_size,
+    const bool exclusive,
+    const bool reverse);
+
+template void CumSumImpl<uint32_t>(
+    const uint32_t* input_data,
+    const fast_divmod& input_dim_along_axis,
+    const fast_divmod& input_stride_along_axis,
+    uint32_t* output_data,
+    const int64_t output_size,
+    const size_t element_size,
+    const bool exclusive,
+    const bool reverse);
+
+template void CumSumImpl<uint64_t>(
+    const uint64_t* input_data,
+    const fast_divmod& input_dim_along_axis,
+    const fast_divmod& input_stride_along_axis,
+    uint64_t* output_data,
+    const int64_t output_size,
+    const size_t element_size,
+    const bool exclusive,
+    const bool reverse);
+
+template void CumSumImpl<float>(
+    const float* input_data,
+    const fast_divmod& input_dim_along_axis,
+    const fast_divmod& input_stride_along_axis,
+    float* output_data,
+    const int64_t output_size,
+    const size_t element_size,
+    const bool exclusive,
+    const bool reverse);
+
+template void CumSumImpl<double>(
+    const double* input_data,
+    const fast_divmod& input_dim_along_axis,
+    const fast_divmod& input_stride_along_axis,
+    double* output_data,
+    const int64_t output_size,
+    const size_t element_size,
+    const bool exclusive,
+    const bool reverse);
+
+template void CumSumImpl<half>(
+    const half* input_data,
+    const fast_divmod& input_dim_along_axis,
+    const fast_divmod& input_stride_along_axis,
+    half* output_data,
+    const int64_t output_size,
+    const size_t element_size,
+    const bool exclusive,
+    const bool reverse);
+
+}  // namespace cuda
+}  // namespace onnxruntime
+

onnxruntime/core/providers/cuda/math/cumsum_impl.h

@@ -0,0 +1,24 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <stdint.h>
+#include "core/providers/cuda/shared_inc/cuda_utils.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+template <typename T>
+void CumSumImpl(
+    const T* input_data,
+    const fast_divmod& input_dim_along_axis,
+    const fast_divmod& input_stride_along_axis,
+    T* output_data,
+    const int64_t output_size,
+    const size_t element_size,
+    const bool exclusive,
+    const bool reverse);
+
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/test/providers/cpu/math/cumsum_test.cc

@@ -193,5 +193,19 @@ TEST(CumSumTest, _1DTestInt64) {
   test.AddOutput<int64_t>("y", {5}, {1, 3, 6, 10, 15});
   test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
 }
+TEST(CumSumTest, _1DTestdouble) {
+  OpTester test("CumSum", 11, onnxruntime::kOnnxDomain);
+  test.AddInput<double>("x", {5}, {1., 2., 3., 4., 5.});
+  test.AddInput<int32_t>("axis", {1}, {0});
+  test.AddOutput<double>("y", {5}, {1., 3., 6., 10., 15.});
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
+}
+TEST(CumSumTest, _1DTestdouble_WithInt64Axis) {
+  OpTester test("CumSum", 11, onnxruntime::kOnnxDomain);
+  test.AddInput<double>("x", {5}, {1., 2., 3., 4., 5.});
+  test.AddInput<int64_t>("axis", {1}, {0});
+  test.AddOutput<double>("y", {5}, {1., 3., 6., 10., 15.});
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
+}
 }  // namespace test
-}  // namespace onnxruntime
+}  // namespace onnxruntime