From e1901a7e10b4a156093288d0c0a200700f3caa5e Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Wed, 2 Sep 2020 10:17:39 -0700 Subject: [PATCH] Improve performance of CUDA implementations for GatherElements and Greater, Equal and Less (#4989) Make GatherElements kernel process 16 items each. unroll the constant loop. Quit loops early for zero dividend. Optimize Binary CompareFunction and remove Impl_Cast invocation. --- .../cuda/cu_inc/binary_elementwise_impl.cuh | 71 ++++---- .../cuda/math/binary_elementwise_ops.cc | 16 +- .../cuda/math/binary_elementwise_ops.h | 2 +- .../cuda/math/binary_elementwise_ops_impl.cu | 52 +++++- .../cuda/math/binary_elementwise_ops_impl.h | 28 ++- .../binary_elementwise_ops_impl_functors.cuh | 20 ++- .../math/variadic_elementwise_ops_impl.cu | 2 +- .../core/providers/cuda/tensor/gather.cc | 76 +++----- .../providers/cuda/tensor/gather_elements.cc | 32 ++-- .../cuda/tensor/gather_elements_impl.cu | 162 ++++++++++-------- .../cuda/tensor/gather_elements_impl.h | 6 +- .../core/providers/cuda/tensor/gather_impl.cu | 90 ++++++---- .../core/providers/cuda/tensor/gather_impl.h | 9 +- 13 files changed, 327 insertions(+), 239 deletions(-) diff --git a/onnxruntime/core/providers/cuda/cu_inc/binary_elementwise_impl.cuh b/onnxruntime/core/providers/cuda/cu_inc/binary_elementwise_impl.cuh index 8fa38813e7..c7f6c7a83f 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/binary_elementwise_impl.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/binary_elementwise_impl.cuh @@ -10,20 +10,21 @@ namespace onnxruntime { namespace cuda { // broadcast by computing output coordinate from offset, using fast_divmod -template +template __global__ void _BinaryElementWise( int32_t output_rank, const TArray lhs_padded_strides, - const T* lhs_data, + const T1* lhs_data, const TArray rhs_padded_strides, - const T1* rhs_data, + const T2* rhs_data, const TArray fdm_output_strides, T* output_data, const FuncT& functor, CUDA_LONG N) { CUDA_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x; - T lvalue[NumElementsPerThread]; - T1 rvalue[NumElementsPerThread]; + T1 lvalue[NumElementsPerThread]; + T2 rvalue[NumElementsPerThread]; CUDA_LONG id = start; #pragma unroll @@ -68,16 +69,16 @@ __global__ void _BinaryElementWise( } // for scalar broadcast or non-broadcast case -template +template __global__ void _BinaryElementWiseSimple( - const T* lhs_data, - const T1* rhs_data, + const T1* lhs_data, + const T2* rhs_data, T* output_data, const FuncT& func, CUDA_LONG N) { CUDA_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x; - T lvalue[NumElementsPerThread]; - T1 rvalue[NumElementsPerThread]; + T1 lvalue[NumElementsPerThread]; + T2 rvalue[NumElementsPerThread]; CUDA_LONG id = start; #pragma unroll @@ -102,17 +103,17 @@ __global__ void _BinaryElementWiseSimple( } // for rhs per-channel broadcast case -template +template __global__ void _BinaryElementWiseRhsPerChannelBatch1( - const T* lhs_data, - const T1* rhs_data, + const T1* lhs_data, + const T2* rhs_data, const fast_divmod fdm_H, T* output_data, FuncT func, CUDA_LONG N) { CUDA_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x; - T lvalue[NumElementsPerThread]; - T1 rvalue[NumElementsPerThread]; + T1 lvalue[NumElementsPerThread]; + T2 rvalue[NumElementsPerThread]; CUDA_LONG id = start; #pragma unroll @@ -137,18 +138,18 @@ __global__ void _BinaryElementWiseRhsPerChannelBatch1( } } -template +template __global__ void _BinaryElementWiseRhsPerChannelBatchN( - const T* lhs_data, - const T1* rhs_data, + const T1* lhs_data, + const T2* rhs_data, const fast_divmod fdm_H, const fast_divmod fdm_C, T* output_data, FuncT func, CUDA_LONG N) { CUDA_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x; - T lvalue[NumElementsPerThread]; - T1 rvalue[NumElementsPerThread]; + T1 lvalue[NumElementsPerThread]; + T2 rvalue[NumElementsPerThread]; CUDA_LONG id = start; #pragma unroll @@ -177,10 +178,10 @@ __global__ void _BinaryElementWiseRhsPerChannelBatchN( } } -template +template void BinaryElementWiseNoBroadcastImpl( - const T* lhs_data, - const T1* rhs_data, + const T1* lhs_data, + const T2* rhs_data, T* output_data, const FuncT& func, size_t count) { @@ -189,7 +190,7 @@ void BinaryElementWiseNoBroadcastImpl( int blocksPerGrid = static_cast(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread)); CUDA_LONG N = static_cast(count); - _BinaryElementWiseSimple<<>>( + _BinaryElementWiseSimple<<>>( lhs_data, rhs_data, output_data, @@ -197,13 +198,13 @@ void BinaryElementWiseNoBroadcastImpl( N); } -template +template void BinaryElementWiseImpl( int32_t output_rank_or_simple_broadcast, const TArray* lhs_padded_strides, - const T* lhs_data, + const T1* lhs_data, const TArray* rhs_padded_strides, - const T1* rhs_data, + const T2* rhs_data, const TArray* fdm_output_strides, const fast_divmod& fdm_H, const fast_divmod& fdm_C, @@ -216,21 +217,21 @@ void BinaryElementWiseImpl( int blocksPerGrid = static_cast(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread)); CUDA_LONG N = static_cast(count); if (output_rank_or_simple_broadcast == static_cast(SimpleBroadcast::NoBroadcast)) { - _BinaryElementWiseSimple<<>>( + _BinaryElementWiseSimple<<>>( lhs_data, rhs_data, output_data, func, N); } else if (output_rank_or_simple_broadcast == static_cast(SimpleBroadcast::LeftScalar)) { - _BinaryElementWiseSimple<<>>( + _BinaryElementWiseSimple<<>>( lhs_data, rhs_data, output_data, func, N); } else if (output_rank_or_simple_broadcast == static_cast(SimpleBroadcast::RightScalar)) { - _BinaryElementWiseSimple<<>>( lhs_data, rhs_data, @@ -238,7 +239,7 @@ void BinaryElementWiseImpl( func, N); } else if (output_rank_or_simple_broadcast == static_cast(SimpleBroadcast::RightPerChannelBatch1)) { - _BinaryElementWiseRhsPerChannelBatch1<<>>( + _BinaryElementWiseRhsPerChannelBatch1<<>>( lhs_data, rhs_data, fdm_H, @@ -246,7 +247,7 @@ void BinaryElementWiseImpl( func, N); } else if (output_rank_or_simple_broadcast == static_cast(SimpleBroadcast::RightPerChannelBatchN)) { - _BinaryElementWiseRhsPerChannelBatchN<<>>( + _BinaryElementWiseRhsPerChannelBatchN<<>>( lhs_data, rhs_data, fdm_H, @@ -256,7 +257,7 @@ void BinaryElementWiseImpl( N); } else { if (lhs_padded_strides && rhs_padded_strides && lhs_padded_strides->Size() && rhs_padded_strides->Size()) - _BinaryElementWise<<>>( + _BinaryElementWise<<>>( output_rank_or_simple_broadcast, *lhs_padded_strides, lhs_data, @@ -267,7 +268,7 @@ void BinaryElementWiseImpl( func, N); else if (lhs_padded_strides && lhs_padded_strides->Size()) - _BinaryElementWise<<>>( + _BinaryElementWise<<>>( output_rank_or_simple_broadcast, *lhs_padded_strides, lhs_data, @@ -278,7 +279,7 @@ void BinaryElementWiseImpl( func, N); else - _BinaryElementWise<<>>( + _BinaryElementWise<<>>( output_rank_or_simple_broadcast, *lhs_padded_strides, lhs_data, diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc index d92486e071..30a4c8caa5 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc @@ -369,8 +369,7 @@ template Status CompareFunction::CompareMethod(OpKernelContext* context, ImplCompare Impl_Compare) const { BinaryElementwisePreparation prepare; ORT_RETURN_IF_ERROR(Prepare(context, &prepare)); - size_t output_size = prepare.output_tensor->Shape().Size(); - IAllocatorUniquePtr output_buffer = GetScratchBuffer(output_size); + Impl_Compare( prepare.output_rank_or_simple_broadcast, &prepare.lhs_padded_strides, @@ -380,13 +379,8 @@ Status CompareFunction::CompareMethod(OpKernelContext* context, ImplCo &prepare.fdm_output_strides, prepare.fdm_H, prepare.fdm_C, - reinterpret_cast(output_buffer.get()), - prepare.output_tensor->Shape().Size()); - - Impl_Cast::MappedType>( - reinterpret_cast(output_buffer.get()), reinterpret_cast::MappedType*>(prepare.output_tensor->template MutableData()), - output_size); + prepare.output_tensor->Shape().Size()); return Status::OK(); } @@ -395,14 +389,14 @@ Status CompareFunction::CompareMethod(OpKernelContext* context, ImplCo //for other elementwise ops template Status Greater::ComputeInternal(OpKernelContext* context) const { - this->CompareMethod(context, &Impl_Greater); + this->CompareMethod(context, &ImplT2_Greater); return Status::OK(); } template Status Equal::ComputeInternal(OpKernelContext* context) const { - this->CompareMethod(context, &Impl_Equal); + this->CompareMethod(context, &ImplT2_Equal); return Status::OK(); } @@ -411,7 +405,7 @@ Status Equal::ComputeInternal(OpKernelContext* context) const { //for other elementwise ops template Status Less::ComputeInternal(OpKernelContext* context) const { - this->CompareMethod(context, &Impl_Less); + this->CompareMethod(context, &ImplT2_Less); return Status::OK(); } diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h index afa4b165b1..488b3d6d24 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h @@ -227,7 +227,7 @@ class CompareFunction : public BinaryElementwise { const TArray* fdm_output_strides, const fast_divmod& fdm_H, const fast_divmod& fdm_C, - CudaT* output_data, + bool* output_data, size_t count); Status CompareMethod(OpKernelContext* context, ImplCompare Impl_Compare) const; diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu index 9d10d9bb76..d1f5b84c74 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu @@ -21,7 +21,7 @@ namespace cuda { fdm_H, \ fdm_C, \ output_data, \ - OP_##name(), \ + OP_##name(), \ count); \ } @@ -36,7 +36,22 @@ namespace cuda { fdm_H, \ fdm_C, \ output_data, \ - OP_##name(), \ + OP_##name(), \ + count); \ + } + +#define BINARY_ELEMENTWISE_IMPL_T2(name) \ + BINARY_ELEMENTWISE_IMPL_DECLARATION_T2(name) { \ + BinaryElementWiseImpl(output_rank_or_simple_broadcast, \ + lhs_padded_strides, \ + lhs_data, \ + rhs_padded_strides, \ + rhs_data, \ + fdm_output_strides, \ + fdm_H, \ + fdm_C, \ + output_data, \ + OP_##name(), \ count); \ } @@ -52,6 +67,12 @@ namespace cuda { const TArray* rhs_padded_strides, const T1* rhs_data, \ const TArray* fdm_output_strides, const fast_divmod& fdm_H, const fast_divmod& fdm_C, T* output_data, size_t count); +#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(x, T, T1, T2) \ + template void ImplT2_##x(int32_t output_rank, \ + const TArray* lhs_padded_strides, const T1* lhs_data, \ + const TArray* rhs_padded_strides, const T2* rhs_data, \ + const TArray* fdm_output_strides, const fast_divmod& fdm_H, const fast_divmod& fdm_C, T* output_data, size_t count); + #define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(x) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, uint32_t) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, uint64_t) \ @@ -103,11 +124,8 @@ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(And, bool) SPECIALIZED_BINARY_ELEMENTWISE_IMPL(Or, bool) SPECIALIZED_BINARY_ELEMENTWISE_IMPL(Xor, bool) SPECIALIZED_BINARY_ELEMENTWISE_IMPL_HFD(PRelu) -SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(Greater) -SPECIALIZED_BINARY_ELEMENTWISE_IMPL_OIL(Equal) SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(Max) SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(Min) -SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(Less) // create declarations for impl for Pow BINARY_ELEMENTWISE_IMPL_T1(Pow) @@ -132,5 +150,29 @@ SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1(Pow, double, int64_t) SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1(Pow, double, float) SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1(Pow, double, double) +// create declarations for impl2 +#define BINARY_OP_NAME_EXPR2(name, expr) \ + BINARY_ELEMENTWISE_IMPL_T2(name) + +BINARY_OPS2() +#undef BINARY_OP_NAME_EXPR2 + +#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD2(name) \ + SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, uint32_t, uint32_t) \ + SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, uint64_t, uint64_t) \ + SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, int32_t, int32_t) \ + SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, int64_t, int64_t) \ + SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, half, half) \ + SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, float, float) \ + SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, double, double) + +SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD2(Greater) + +SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(Equal, bool, bool, bool) +SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(Equal, bool, int32_t, int32_t) +SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(Equal, bool, int64_t, int64_t) + +SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD2(Less) + } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.h b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.h index ec47fadb86..dbc7e89a03 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.h +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.h @@ -25,11 +25,8 @@ namespace cuda { BINARY_OP_NAME_EXPR(Or, (a | b)) \ BINARY_OP_NAME_EXPR(Xor, (a ^ b)) \ BINARY_OP_NAME_EXPR(PRelu, (a > (T)0 ? a : a * b)) \ - BINARY_OP_NAME_EXPR(Greater, (a > b) ? 1 : 0) \ - BINARY_OP_NAME_EXPR(Equal, ((a == b) ? 1 : 0)) \ BINARY_OP_NAME_EXPR(Max, _Max(a, b)) \ - BINARY_OP_NAME_EXPR(Min, _Min(a, b)) \ - BINARY_OP_NAME_EXPR(Less, (a < b) ? 1 : 0) + BINARY_OP_NAME_EXPR(Min, _Min(a, b)) // NOTE that cu files are compiled with nvcc and should not refer to any onnxruntime headers // so struct BinaryElementwisePreparation cannot be used here @@ -68,5 +65,28 @@ BINARY_OPS() BINARY_ELEMENTWISE_IMPL_DECLARATION_T1(Pow); +#define BINARY_ELEMENTWISE_IMPL_DECLARATION_T2(name) \ + template \ + void ImplT2_##name( \ + int32_t output_rank_or_simple_broadcast, \ + const TArray* lhs_padded_strides, \ + const T1* lhs_data, \ + const TArray* rhs_padded_strides, \ + const T2* rhs_data, \ + const TArray* fdm_output_strides, \ + const fast_divmod& fdm_H, \ + const fast_divmod& fdm_C, \ + T* output_data, \ + size_t count) + +#define BINARY_OPS2() \ + BINARY_OP_NAME_EXPR2(Greater, (a > b)) \ + BINARY_OP_NAME_EXPR2(Equal, (a == b)) \ + BINARY_OP_NAME_EXPR2(Less, (a < b)) + +#define BINARY_OP_NAME_EXPR2(name, expr) BINARY_ELEMENTWISE_IMPL_DECLARATION_T2(name); +BINARY_OPS2() +#undef BINARY_OP_NAME_EXPR2 + } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl_functors.cuh b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl_functors.cuh index 471a42d342..d49cd43c90 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl_functors.cuh +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl_functors.cuh @@ -10,12 +10,12 @@ namespace cuda { // define the device functors that perform the computation on scalars -#define OP_FUNCTOR_DEFINITION(name, expr) \ - template \ - struct OP_##name { \ - __device__ __inline__ T operator()(T a, T1 b) const { \ - return (expr); \ - } \ +#define OP_FUNCTOR_DEFINITION(name, expr) \ + template \ + struct OP_##name { \ + __device__ __inline__ T operator()(T1 a, T2 b) const { \ + return (expr); \ + } \ }; #define BINARY_OP_NAME_EXPR(name, expr) \ @@ -26,6 +26,14 @@ BINARY_OPS() OP_FUNCTOR_DEFINITION(Pow, _Pow(a, b)) #undef BINARY_OP_NAME_EXPR + +#define BINARY_OP_NAME_EXPR2(name, expr) \ + OP_FUNCTOR_DEFINITION(name, expr) + +BINARY_OPS2() + +#undef BINARY_OP_NAME_EXPR2 + #undef OP_FUNCTOR_DEFINITION } // namespace cuda diff --git a/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu b/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu index 86af4b82ba..a960226524 100644 --- a/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu +++ b/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu @@ -17,7 +17,7 @@ struct VariadicElementwiseOpTraits; #define DEFINE_TRAITS(VariadicElementwiseOpTag, ImplName) \ template \ struct VariadicElementwiseOpTraits { \ - using ScalarComputeFunctor = OP_##ImplName; \ + using ScalarComputeFunctor = OP_##ImplName; \ \ static void ComputeFn( \ int32_t output_rank_or_simple_broadcast, \ diff --git a/onnxruntime/core/providers/cuda/tensor/gather.cc b/onnxruntime/core/providers/cuda/tensor/gather.cc index b7811dff42..1f1a86b2a0 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather.cc +++ b/onnxruntime/core/providers/cuda/tensor/gather.cc @@ -33,40 +33,6 @@ ONNX_OPERATOR_KERNEL_EX( DataTypeImpl::GetTensorType()}), Gather); -#define TYPED_FUNCTION_CALL(T) \ - if (utils::IsPrimitiveDataType(T_type)) { \ - T* output_data = p.output_tensor->template MutableData(); \ - const T* input_data = p.input_tensor->template Data(); \ - if (utils::IsPrimitiveDataType(Tin_type)) { \ - if (p.output_tensor->Shape().Size() > 0) { \ - GatherImpl( \ - input_block_size, \ - indices_max, \ - divmod_output_block_size, \ - divmod_block_size, \ - p.indices_tensor->template Data(), \ - reinterpret_cast::MappedType*>(input_data), \ - reinterpret_cast::MappedType*>(output_data), \ - p.output_tensor->Shape().Size()); \ - } \ - return Status::OK(); \ - } \ - if (utils::IsPrimitiveDataType(Tin_type)) { \ - if (p.output_tensor->Shape().Size() > 0) { \ - GatherImpl( \ - input_block_size, \ - indices_max, \ - divmod_output_block_size, \ - divmod_block_size, \ - p.indices_tensor->template Data(), \ - reinterpret_cast::MappedType*>(input_data), \ - reinterpret_cast::MappedType*>(output_data), \ - p.output_tensor->Shape().Size()); \ - } \ - return Status::OK(); \ - } \ - } - Status Gather::ComputeInternal(OpKernelContext* context) const { Prepare p; ORT_RETURN_IF_ERROR(PrepareForCompute(context, p)); @@ -79,24 +45,38 @@ Status Gather::ComputeInternal(OpKernelContext* context) const { const int64_t output_block_size = N * block_size; const int64_t indices_max = input_shape[p.axis]; + const void* input_data = p.input_tensor->DataRaw(); + const void* indices_data = p.indices_tensor->DataRaw(); + void* output_data = p.output_tensor->MutableDataRaw(); + + if (p.output_tensor->Shape().Size() == 0) { + return Status::OK(); + } + const fast_divmod divmod_output_block_size(gsl::narrow_cast(output_block_size)); const fast_divmod divmod_block_size(gsl::narrow_cast(block_size)); - MLDataType T_type = p.input_tensor->DataType(); - MLDataType Tin_type = p.indices_tensor->DataType(); + const size_t element_size = p.input_tensor->DataType()->Size(); + const size_t index_element_size = p.indices_tensor->DataType()->Size(); - TYPED_FUNCTION_CALL(int8_t) - TYPED_FUNCTION_CALL(int16_t) - TYPED_FUNCTION_CALL(int32_t) - TYPED_FUNCTION_CALL(int64_t) - TYPED_FUNCTION_CALL(uint8_t) - TYPED_FUNCTION_CALL(uint16_t) - TYPED_FUNCTION_CALL(uint32_t) - TYPED_FUNCTION_CALL(uint64_t) - TYPED_FUNCTION_CALL(MLFloat16) - TYPED_FUNCTION_CALL(float) - TYPED_FUNCTION_CALL(double) - TYPED_FUNCTION_CALL(bool) + // CUDA Kernel implementation supports element sizes of: + // int8_t, int16_t, int32_t and int64_t which covers all supported + // types since there is no computations necessary just data movement + if (p.indices_tensor->IsDataType() || + p.indices_tensor->IsDataType()) { + GatherImpl( + input_block_size, + indices_max, + divmod_output_block_size, + divmod_block_size, + indices_data, + index_element_size, + input_data, + element_size, + output_data, + p.output_tensor->Shape().Size()); + return Status::OK(); + } return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "Type for Tind not supported yet in Gather."); } diff --git a/onnxruntime/core/providers/cuda/tensor/gather_elements.cc b/onnxruntime/core/providers/cuda/tensor/gather_elements.cc index 3609a4d96e..6c333e27b9 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather_elements.cc +++ b/onnxruntime/core/providers/cuda/tensor/gather_elements.cc @@ -36,7 +36,7 @@ Status GatherElements::ComputeInternal(OpKernelContext* context) const { const int64_t indices_size = indices_shape.Size(); // Handle negative axis if any - const int64_t axis = static_cast(HandleNegativeAxis(axis_, input_rank)); + const int64_t axis = HandleNegativeAxis(axis_, input_rank); // Validate input shapes and ranks (invoke the static method in the CPU GatherElements kernel that hosts the shared checks) auto status = onnxruntime::GatherElements::ValidateInputShapes(input_shape, indices_shape, axis); @@ -56,38 +56,26 @@ Status GatherElements::ComputeInternal(OpKernelContext* context) const { TArray fdm_indices_strides(indices_rank); TensorPitches indices_strides(indices_dims); for (auto i = 0; i < indices_rank; i++) { - fdm_indices_strides[i] = fast_divmod(static_cast(indices_strides[i])); + fdm_indices_strides[i] = fast_divmod(gsl::narrow_cast(indices_strides[i])); } - size_t element_size = input_tensor->DataType()->Size(); + const size_t element_size = input_tensor->DataType()->Size(); + const size_t index_element_size = indices_tensor->DataType()->Size(); - if (indices_tensor->IsDataType()) { - const int32_t* indices_data = indices_tensor->template Data(); - GatherElementsImpl( + if (indices_tensor->IsDataType() || + indices_tensor->IsDataType()) { + GatherElementsImpl( input_rank, input_tensor->DataRaw(), input_dims[axis], gpu_input_strides, - indices_data, + indices_tensor->DataRaw(), indices_size, fdm_indices_strides, axis, output_tensor->MutableDataRaw(), - element_size); - return Status::OK(); - } else if (indices_tensor->IsDataType()) { - const int64_t* indices_data = indices_tensor->template Data(); - GatherElementsImpl( - input_rank, - input_tensor->DataRaw(), - input_dims[axis], - gpu_input_strides, - indices_data, - indices_size, - fdm_indices_strides, - axis, - output_tensor->MutableDataRaw(), - element_size); + element_size, + index_element_size); return Status::OK(); } else { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "GatherElements op: Type of 'indices' must be int32 or int64"); diff --git a/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu b/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu index 1af643dd46..fc4cc644c1 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu @@ -7,114 +7,138 @@ namespace onnxruntime { namespace cuda { -template +namespace { +constexpr int threads_per_block = GridDim::maxThreadsPerBlock; +constexpr int thread_worksize = 16; +} // namespace + +__host__ __device__ inline int64_t GetIndexValue(const void* index_data, size_t index_element_size, size_t offset) { + switch (index_element_size) { + case sizeof(int32_t): + return *(reinterpret_cast(index_data) + offset); + break; + case sizeof(int64_t): + return *(reinterpret_cast(index_data) + offset); + break; + default: + break; + } + // What is a sensible thing to do here? + assert(false); + return std::numeric_limits::max(); +} + +template __global__ void _GatherElementsKernel( const int64_t rank, const T* input_data, const int64_t input_dim_along_axis, const TArray input_strides, - const Tin* indices_data, + const void* indices_data, const int64_t indices_size, + const size_t index_element_size, const TArray indices_strides, const int64_t axis, T* output_data) { - CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(indices_index, indices_size); - int dim = 0; - int remain = indices_index; - size_t data_idx = 0; - for (int i = 0; i < rank; ++i) { - indices_strides[i].divmod(remain, dim, remain); - if (i == axis) { - dim = static_cast(indices_data[indices_index]); - if (dim < -input_dim_along_axis || dim >= input_dim_along_axis) { - return; // Invalid index + + CUDA_LONG indices_index = threads_per_block * thread_worksize * blockIdx.x + threadIdx.x; + + #pragma unroll + for (int work = 0; work < thread_worksize; ++work) { + if (indices_index < indices_size) { + + int dim = 0; + int remain = indices_index; + int64_t data_idx = 0; + + int i = 0; + for (; i < axis && remain > 0; ++i) { + indices_strides[i].divmod(remain, dim, remain); + data_idx += input_strides[i] * dim; } + + i = axis; + indices_strides[i].divmod(remain, dim, remain); + dim = GetIndexValue(indices_data, index_element_size, indices_index); + if (dim < -input_dim_along_axis || dim >= input_dim_along_axis) { + return; // Invalid index + } + if (dim < 0) { dim += input_dim_along_axis; } + + data_idx += input_strides[i] * dim; + + ++i; // past axis + for (; i < rank && remain > 0; ++i) { + indices_strides[i].divmod(remain, dim, remain); + data_idx += input_strides[i] * dim; + } + output_data[indices_index] = input_data[data_idx]; + + indices_index += threads_per_block; } - data_idx += input_strides[i] * dim; } - output_data[indices_index] = input_data[data_idx]; } -template void GatherElementsImpl( const int64_t rank, const void* input_data, const int64_t input_dim_along_axis, const TArray& input_strides, - const Tin* indices_data, + const void* indices_data, const int64_t indices_size, const TArray& indices_strides, const int64_t axis, void* output_data, - size_t element_size) { + size_t element_size, + size_t index_element_size) { if (indices_size > 0) { - int blocksPerGrid = static_cast((indices_size + GridDim::maxThreadsPerBlock - 1) / GridDim::maxThreadsPerBlock); + dim3 block(threads_per_block); + dim3 blocksPerGrid((static_cast(indices_size + block.x * thread_worksize - 1) / (block.x * thread_worksize))); switch (element_size) { - case sizeof(int8_t): - _GatherElementsKernel<<>>( - rank, reinterpret_cast::MappedType*>(input_data), input_dim_along_axis, input_strides, - indices_data, indices_size, indices_strides, - axis, reinterpret_cast::MappedType*>(output_data)); - break; + case sizeof(int8_t): { + using CudaType = typename ToCudaType::MappedType; + _GatherElementsKernel<<>>( + rank, reinterpret_cast(input_data), input_dim_along_axis, input_strides, + indices_data, indices_size, index_element_size, indices_strides, + axis, reinterpret_cast(output_data)); + } break; - case sizeof(int16_t): - _GatherElementsKernel<<>>( - rank, reinterpret_cast::MappedType*>(input_data), input_dim_along_axis, input_strides, - indices_data, indices_size, indices_strides, - axis, reinterpret_cast::MappedType*>(output_data)); - break; + case sizeof(int16_t): { + using CudaType = typename ToCudaType::MappedType; + _GatherElementsKernel<<>>( + rank, reinterpret_cast(input_data), input_dim_along_axis, input_strides, + indices_data, indices_size, index_element_size, indices_strides, + axis, reinterpret_cast(output_data)); + } break; - case sizeof(int32_t): - _GatherElementsKernel<<>>( - rank, reinterpret_cast::MappedType*>(input_data), input_dim_along_axis, input_strides, - indices_data, indices_size, indices_strides, - axis, reinterpret_cast::MappedType*>(output_data)); - break; + case sizeof(int32_t): { + using CudaType = typename ToCudaType::MappedType; + _GatherElementsKernel<<>>( + rank, reinterpret_cast(input_data), input_dim_along_axis, input_strides, + indices_data, indices_size, index_element_size, indices_strides, + axis, reinterpret_cast(output_data)); + } break; - case sizeof(int64_t): - _GatherElementsKernel<<>>( - rank, reinterpret_cast::MappedType*>(input_data), input_dim_along_axis, input_strides, - indices_data, indices_size, indices_strides, - axis, reinterpret_cast::MappedType*>(output_data)); - break; + case sizeof(int64_t): { + using CudaType = typename ToCudaType::MappedType; + _GatherElementsKernel<<>>( + rank, reinterpret_cast(input_data), input_dim_along_axis, input_strides, + indices_data, indices_size, index_element_size, indices_strides, + axis, reinterpret_cast(output_data)); + } break; - // should not reach here as we validate if the all relevant types are supported in the Compute method + // should not reach here as we validate if the all relevant types are supported in the Compute method default: ORT_THROW("Unsupported element size by the GatherElements CUDA kernel"); } } -} - -template void GatherElementsImpl( - const int64_t rank, - const void* input_data, - const int64_t input_dim_along_axis, - const TArray& input_strides, - const int32_t* indices_data, - const int64_t indices_size, - const TArray& indices_strides, - const int64_t axis, - void* output_data, - size_t element_size); - -template void GatherElementsImpl( - const int64_t rank, - const void* input_data, - const int64_t input_dim_along_axis, - const TArray& input_strides, - const int64_t* indices_data, - const int64_t indices_size, - const TArray& indices_strides, - const int64_t axis, - void* output_data, - size_t element_size); +} // namespace cuda } // namespace cuda } // namespace onnxruntime - diff --git a/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.h b/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.h index 2b086ab0d8..1caaea647c 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.h +++ b/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.h @@ -9,18 +9,18 @@ namespace onnxruntime { namespace cuda { -template void GatherElementsImpl( const int64_t rank, // both inputs have same rank and this is validated in the main Compute const void* input_data, const int64_t input_dim_along_axis, const TArray& input_strides, - const Tin* indices_data, + const void* indices_data, const int64_t indices_size, const TArray& indices_strides, const int64_t axis, void* output_data, - size_t element_size); + size_t element_size, + size_t index_element_size); } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/tensor/gather_impl.cu b/onnxruntime/core/providers/cuda/tensor/gather_impl.cu index ab66da982b..aa42e6d6fb 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/gather_impl.cu @@ -7,13 +7,30 @@ namespace onnxruntime { namespace cuda { -template +__host__ __device__ inline int64_t GetIndexValue(const void* index_data, size_t index_element_size, size_t offset) { + switch (index_element_size) { + case sizeof(int32_t): + return *(reinterpret_cast(index_data) + offset); + break; + case sizeof(int64_t): + return *(reinterpret_cast(index_data) + offset); + break; + default: + break; + } + // What is a sensible thing to do here? + assert(false); + return std::numeric_limits::max(); +} + +template __global__ void _GatherKernel( const int64_t input_block_size, const int64_t indices_max, const fast_divmod output_block_size, const fast_divmod block_size, - const Tin* indices_data, + const void* indices_data, + const size_t index_element_size, const T* input_data, T* output_data, const CUDA_LONG N) { @@ -23,7 +40,7 @@ __global__ void _GatherKernel( output_block_size.divmod(id, input_block_index, block_offset); int indices_index, offset; block_size.divmod(block_offset, indices_index, offset); - int64_t idx = indices_data[indices_index]; + int64_t idx = GetIndexValue(indices_data, index_element_size, indices_index); idx = idx < 0 ? idx + indices_max : idx; if (idx < 0 || idx >= indices_max) { output_data[id] = 0; @@ -34,41 +51,54 @@ __global__ void _GatherKernel( output_data[id] = input_data[input_index]; } -template void GatherImpl( const int64_t input_block_size, const int64_t indices_max, const fast_divmod& output_block_size, const fast_divmod& block_size, - const Tin* indices_data, - const T* input_data, - T* output_data, + const void* indices_data, + size_t index_element_size, + const void* input_data, + size_t element_size, + void* output_data, const size_t N) { + int blocksPerGrid = (int)(ceil(static_cast(N) / GridDim::maxThreadsPerBlock)); - _GatherKernel<<>>( - input_block_size, indices_max, output_block_size, block_size, indices_data, input_data, output_data, (CUDA_LONG)N); + + switch (element_size) { + case sizeof(int8_t): { + using CudaType = typename ToCudaType::MappedType; + _GatherKernel<<>>( + input_block_size, indices_max, output_block_size, block_size, indices_data, index_element_size, + reinterpret_cast(input_data), reinterpret_cast(output_data), (CUDA_LONG)N); + + } break; + case sizeof(int16_t): { + using CudaType = typename ToCudaType::MappedType; + _GatherKernel<<>>( + input_block_size, indices_max, output_block_size, block_size, indices_data, index_element_size, + reinterpret_cast(input_data), reinterpret_cast(output_data), (CUDA_LONG)N); + + } break; + case sizeof(int32_t): { + using CudaType = typename ToCudaType::MappedType; + _GatherKernel<<>>( + input_block_size, indices_max, output_block_size, block_size, indices_data, index_element_size, + reinterpret_cast(input_data), reinterpret_cast(output_data), (CUDA_LONG)N); + + } break; + case sizeof(int64_t): { + using CudaType = typename ToCudaType::MappedType; + _GatherKernel<<>>( + input_block_size, indices_max, output_block_size, block_size, indices_data, index_element_size, + reinterpret_cast(input_data), reinterpret_cast(output_data), (CUDA_LONG)N); + + } break; + + default: + ORT_THROW("Unsupported element size by the Gather CUDA kernel"); + } } -#define SPECIALIZED_IMPL(T) \ - template void GatherImpl(const int64_t input_block_size, const int64_t indices_max, \ - const fast_divmod& output_block_size, const fast_divmod& block_size, \ - const int32_t* indices_data, const T* input_data, T* output_data, const size_t N); \ - template void GatherImpl(const int64_t input_block_size, const int64_t indices_max, \ - const fast_divmod& output_block_size, const fast_divmod& block_size, \ - const int64_t* indices_data, const T* input_data, T* output_data, const size_t N); - -SPECIALIZED_IMPL(int8_t) -SPECIALIZED_IMPL(int16_t) -SPECIALIZED_IMPL(int32_t) -SPECIALIZED_IMPL(int64_t) -SPECIALIZED_IMPL(uint8_t) -SPECIALIZED_IMPL(uint16_t) -SPECIALIZED_IMPL(uint32_t) -SPECIALIZED_IMPL(uint64_t) -SPECIALIZED_IMPL(half) -SPECIALIZED_IMPL(float) -SPECIALIZED_IMPL(double) -SPECIALIZED_IMPL(bool) - } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/tensor/gather_impl.h b/onnxruntime/core/providers/cuda/tensor/gather_impl.h index 3e74f81695..11af5c3888 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather_impl.h +++ b/onnxruntime/core/providers/cuda/tensor/gather_impl.h @@ -8,15 +8,16 @@ namespace onnxruntime { namespace cuda { -template void GatherImpl( const int64_t input_block_size, const int64_t indices_max, const fast_divmod& output_block_size, const fast_divmod& block_size, - const Tin* indices_data, - const T* input_data, - T* output_data, + const void* indices_data, + size_t index_element_size, + const void* input_data, + size_t element_size, + void* output_data, const size_t N); } // namespace cuda