From 971bc439b5a57b38c82c3e855f341676cbe4ede4 Mon Sep 17 00:00:00 2001 From: Hariharan Seshadri Date: Wed, 18 Dec 2019 16:49:59 -0800 Subject: [PATCH] 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 | 78 ++++++-- onnxruntime/core/providers/cpu/math/cumsum.h | 6 +- .../providers/cuda/cuda_execution_provider.cc | 4 +- .../core/providers/cuda/math/cumsum.cc | 134 ++++++++++++++ onnxruntime/core/providers/cuda/math/cumsum.h | 48 +++++ .../core/providers/cuda/math/cumsum_impl.cu | 167 ++++++++++++++++++ .../core/providers/cuda/math/cumsum_impl.h | 24 +++ .../test/providers/cpu/math/cumsum_test.cc | 16 +- 8 files changed, 459 insertions(+), 18 deletions(-) create mode 100644 onnxruntime/core/providers/cuda/math/cumsum.cc create mode 100644 onnxruntime/core/providers/cuda/math/cumsum.h create mode 100644 onnxruntime/core/providers/cuda/math/cumsum_impl.cu create mode 100644 onnxruntime/core/providers/cuda/math/cumsum_impl.h diff --git a/onnxruntime/core/providers/cpu/math/cumsum.cc b/onnxruntime/core/providers/cpu/math/cumsum.cc index 65a65857c9..9370b6942b 100644 --- a/onnxruntime/core/providers/cpu/math/cumsum.cc +++ b/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()), CumSum); -ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum, 11, double, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), CumSum); -ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum, 11, int32_t, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), CumSum);; -ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum, 11, int64_t, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), CumSum); +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()) { + axis_out = static_cast(axis_tensor->template Data()[0]); + } else if (axis_tensor->IsDataType()) { + axis_out = axis_tensor->template Data()[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()) + .TypeConstraint("T2", std::vector{DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}), + CumSum); + +ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum, + 11, + double, + KernelDefBuilder() + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", std::vector{DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}), + CumSum); + +ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum, + 11, + int32_t, + KernelDefBuilder() + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", std::vector{DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}), + CumSum); + +ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum, + 11, + int64_t, + KernelDefBuilder() + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", std::vector{DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}), + CumSum); template CumSum::CumSum(const OpKernelInfo& info) : OpKernel(info), exclusive_(), reverse_() { @@ -79,19 +130,13 @@ CumSum::CumSum(const OpKernelInfo& info) : OpKernel(info), exclusive_(), reve template Status CumSum::Compute(OpKernelContext* ctx) const { - const Tensor* input = ctx->Input(0); // input tensor - const auto rank = static_cast(input->Shape().NumDimensions()); // the rank of the input/output - const Tensor* axis_tensor = ctx->Input(1); // axis input tensor + const Tensor* input = ctx->Input(0); // input tensor + auto rank = static_cast(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(1); // axis input tensor - int32_t axis = axis_tensor->template Data()[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(rank) + axis; TensorShape output_shape(input->Shape()); auto& output_tensor = *ctx->Output(0, output_shape); // output tensor @@ -99,6 +144,9 @@ Status CumSum::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; diff --git a/onnxruntime/core/providers/cpu/math/cumsum.h b/onnxruntime/core/providers/cpu/math/cumsum.h index 60865a6468..fa1c1ceb0d 100644 --- a/onnxruntime/core/providers/cpu/math/cumsum.h +++ b/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 diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index 1f900d3a7c..6c15a7ffe2 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/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, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -1128,6 +1129,7 @@ static void RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, }; for (auto& function_table_entry : function_table) { diff --git a/onnxruntime/core/providers/cuda/math/cumsum.cc b/onnxruntime/core/providers/cuda/math/cumsum.cc new file mode 100644 index 0000000000..febd749827 --- /dev/null +++ b/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(1) // 'axis' needs to be on CPU + .TypeConstraint("T", std::vector{ + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}) + .TypeConstraint("T2", std::vector{DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}), + CumSum); + +Status CumSum::ComputeInternal(OpKernelContext* ctx) const { + const Tensor* input = ctx->Input(0); // input tensor + auto rank = static_cast(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(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(input_dims[axis])); + fast_divmod fast_divmod_input_stride_along_axis(static_cast(input_stride_along_axis)); + + if (input->IsDataType()) { + CumSumImpl(reinterpret_cast::MappedType*>(input->Data()), + fast_divmod_input_dim_along_axis, + fast_divmod_input_stride_along_axis, + reinterpret_cast::MappedType*>(output.MutableData()), + output_shape.Size(), + input->DataType()->Size(), + exclusive_, + reverse_); + } else if (input->IsDataType()) { + CumSumImpl(reinterpret_cast::MappedType*>(input->Data()), + fast_divmod_input_dim_along_axis, + fast_divmod_input_stride_along_axis, + reinterpret_cast::MappedType*>(output.MutableData()), + output_shape.Size(), + input->DataType()->Size(), + exclusive_, + reverse_); + } else if (input->IsDataType()) { + CumSumImpl(reinterpret_cast::MappedType*>(input->Data()), + fast_divmod_input_dim_along_axis, + fast_divmod_input_stride_along_axis, + reinterpret_cast::MappedType*>(output.MutableData()), + output_shape.Size(), + input->DataType()->Size(), + exclusive_, + reverse_); + } else if (input->IsDataType()) { + CumSumImpl(reinterpret_cast::MappedType*>(input->Data()), + fast_divmod_input_dim_along_axis, + fast_divmod_input_stride_along_axis, + reinterpret_cast::MappedType*>(output.MutableData()), + output_shape.Size(), + input->DataType()->Size(), + exclusive_, + reverse_); + } else if (input->IsDataType()) { + CumSumImpl(reinterpret_cast::MappedType*>(input->Data()), + fast_divmod_input_dim_along_axis, + fast_divmod_input_stride_along_axis, + reinterpret_cast::MappedType*>(output.MutableData()), + output_shape.Size(), + input->DataType()->Size(), + exclusive_, + reverse_); + } else if (input->IsDataType()) { + CumSumImpl(reinterpret_cast::MappedType*>(input->Data()), + fast_divmod_input_dim_along_axis, + fast_divmod_input_stride_along_axis, + reinterpret_cast::MappedType*>(output.MutableData()), + output_shape.Size(), + input->DataType()->Size(), + exclusive_, + reverse_); + } else if (input->IsDataType()) { + CumSumImpl(reinterpret_cast::MappedType*>(input->Data()), + fast_divmod_input_dim_along_axis, + fast_divmod_input_stride_along_axis, + reinterpret_cast::MappedType*>(output.MutableData()), + 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 diff --git a/onnxruntime/core/providers/cuda/math/cumsum.h b/onnxruntime/core/providers/cuda/math/cumsum.h new file mode 100644 index 0000000000..0345e91ad6 --- /dev/null +++ b/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 diff --git a/onnxruntime/core/providers/cuda/math/cumsum_impl.cu b/onnxruntime/core/providers/cuda/math/cumsum_impl.cu new file mode 100644 index 0000000000..c698330fcb --- /dev/null +++ b/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 +__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(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(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 +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((output_size + GridDim::maxThreadsPerBlock - 1) / GridDim::maxThreadsPerBlock); + + _CumSumKernel<<>>(input_data, + input_dim_along_axis, + input_stride_along_axis, + output_data, + output_size, + exclusive, + reverse); + } +} + +template void CumSumImpl( + 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( + 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( + 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( + 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( + 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( + 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( + 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 + diff --git a/onnxruntime/core/providers/cuda/math/cumsum_impl.h b/onnxruntime/core/providers/cuda/math/cumsum_impl.h new file mode 100644 index 0000000000..890652c3ac --- /dev/null +++ b/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 +#include "core/providers/cuda/shared_inc/cuda_utils.h" + +namespace onnxruntime { +namespace cuda { + +template +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 diff --git a/onnxruntime/test/providers/cpu/math/cumsum_test.cc b/onnxruntime/test/providers/cpu/math/cumsum_test.cc index aa3166d0b2..bc624025e1 100644 --- a/onnxruntime/test/providers/cpu/math/cumsum_test.cc +++ b/onnxruntime/test/providers/cpu/math/cumsum_test.cc @@ -193,5 +193,19 @@ TEST(CumSumTest, _1DTestInt64) { test.AddOutput("y", {5}, {1, 3, 6, 10, 15}); test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); } +TEST(CumSumTest, _1DTestdouble) { + OpTester test("CumSum", 11, onnxruntime::kOnnxDomain); + test.AddInput("x", {5}, {1., 2., 3., 4., 5.}); + test.AddInput("axis", {1}, {0}); + test.AddOutput("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("x", {5}, {1., 2., 3., 4., 5.}); + test.AddInput("axis", {1}, {0}); + test.AddOutput("y", {5}, {1., 3., 6., 10., 15.}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} } // namespace test -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime