saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-24 19:43:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Improve performance of CUDA implementations for GatherElements and Greater, Equal and Less (#4989)

Browse source 
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.
This commit is contained in:
Dmitri Smirnov 2020-09-02 10:17:39 -07:00 committed by GitHub
parent d5d5e37e76
commit e1901a7e10
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 327 additions and 239 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

71
onnxruntime/core/providers/cuda/cu_inc/binary_elementwise_impl.cuh
View file

@ -10,20 +10,21 @@ namespace onnxruntime {
namespace cuda {
// broadcast by computing output coordinate from offset, using fast_divmod
template <typename T, typename T1, typename FuncT, bool lhs_need_compute, bool rhs_need_compute, int NumThreadsPerBlock, int NumElementsPerThread>
template <typename T, typename T1, typename T2, typename FuncT,
bool lhs_need_compute, bool rhs_need_compute, int NumThreadsPerBlock, int NumElementsPerThread>
__global__ void _BinaryElementWise(
int32_t output_rank,
const TArray<int64_t> lhs_padded_strides,
const T* lhs_data,
const T1* lhs_data,
const TArray<int64_t> rhs_padded_strides,
const T1* rhs_data,
const T2* rhs_data,
const TArray<fast_divmod> 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 <bool IncL, bool IncR, typename T, typename T1, typename FuncT, int NumThreadsPerBlock, int NumElementsPerThread>
template <bool IncL, bool IncR, typename T, typename T1, typename T2, typename FuncT, int NumThreadsPerBlock, int NumElementsPerThread>
__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 <typename T, typename T1, typename FuncT, int NumThreadsPerBlock, int NumElementsPerThread>
template <typename T, typename T1, typename T2, typename FuncT, int NumThreadsPerBlock, int NumElementsPerThread>
__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 <typename T, typename T1, typename FuncT, int NumThreadsPerBlock, int NumElementsPerThread>
template <typename T, typename T1, typename T2, typename FuncT, int NumThreadsPerBlock, int NumElementsPerThread>
__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 <typename T, typename T1, typename FuncT>
template <typename T, typename T1, typename T2, typename FuncT>
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<int>(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
CUDA_LONG N = static_cast<CUDA_LONG>(count);
_BinaryElementWiseSimple<true, true, T, T1, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
_BinaryElementWiseSimple<true, true, T, T1, T2, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
lhs_data,
rhs_data,
output_data,
@ -197,13 +198,13 @@ void BinaryElementWiseNoBroadcastImpl(
N);
}
template <typename T, typename T1, typename FuncT>
template <typename T, typename T1, typename T2, typename FuncT>
void BinaryElementWiseImpl(
int32_t output_rank_or_simple_broadcast,
const TArray<int64_t>* lhs_padded_strides,
const T* lhs_data,
const T1* lhs_data,
const TArray<int64_t>* rhs_padded_strides,
const T1* rhs_data,
const T2* rhs_data,
const TArray<fast_divmod>* fdm_output_strides,
const fast_divmod& fdm_H,
const fast_divmod& fdm_C,
@ -216,21 +217,21 @@ void BinaryElementWiseImpl(
int blocksPerGrid = static_cast<int>(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
CUDA_LONG N = static_cast<CUDA_LONG>(count);
if (output_rank_or_simple_broadcast == static_cast<int32_t>(SimpleBroadcast::NoBroadcast)) {
_BinaryElementWiseSimple<true, true, T, T1, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
_BinaryElementWiseSimple<true, true, T, T1, T2, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
lhs_data,
rhs_data,
output_data,
func,
N);
} else if (output_rank_or_simple_broadcast == static_cast<int32_t>(SimpleBroadcast::LeftScalar)) {
_BinaryElementWiseSimple<false, true, T, T1, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
_BinaryElementWiseSimple<false, true, T, T1, T2, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
lhs_data,
rhs_data,
output_data,
func,
N);
} else if (output_rank_or_simple_broadcast == static_cast<int32_t>(SimpleBroadcast::RightScalar)) {
_BinaryElementWiseSimple<true, false, T, T1, FuncT, GridDim::maxThreadsPerBlock,
_BinaryElementWiseSimple<true, false, T, T1, T2, FuncT, GridDim::maxThreadsPerBlock,
GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
lhs_data,
rhs_data,
@ -238,7 +239,7 @@ void BinaryElementWiseImpl(
func,
N);
} else if (output_rank_or_simple_broadcast == static_cast<int32_t>(SimpleBroadcast::RightPerChannelBatch1)) {
_BinaryElementWiseRhsPerChannelBatch1<T, T1, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
_BinaryElementWiseRhsPerChannelBatch1<T, T1, T2, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
lhs_data,
rhs_data,
fdm_H,
@ -246,7 +247,7 @@ void BinaryElementWiseImpl(
func,
N);
} else if (output_rank_or_simple_broadcast == static_cast<int32_t>(SimpleBroadcast::RightPerChannelBatchN)) {
_BinaryElementWiseRhsPerChannelBatchN<T, T1, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
_BinaryElementWiseRhsPerChannelBatchN<T, T1, T2, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
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<T, T1, FuncT, true, true, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
_BinaryElementWise<T, T1, T2, FuncT, true, true, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
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<T, T1, FuncT, true, false, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
_BinaryElementWise<T, T1, T2, FuncT, true, false, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
output_rank_or_simple_broadcast,
*lhs_padded_strides,
lhs_data,
@ -278,7 +279,7 @@ void BinaryElementWiseImpl(
func,
N);
else
_BinaryElementWise<T, T1, FuncT, false, true, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
_BinaryElementWise<T, T1, T2, FuncT, false, true, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
output_rank_or_simple_broadcast,
*lhs_padded_strides,
lhs_data,

16
onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc
View file

@ -369,8 +369,7 @@ template <typename T, typename CudaT>
Status CompareFunction<T, CudaT>::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<T> output_buffer = GetScratchBuffer<T>(output_size);
Impl_Compare(
prepare.output_rank_or_simple_broadcast,
&prepare.lhs_padded_strides,
@ -380,13 +379,8 @@ Status CompareFunction<T, CudaT>::CompareMethod(OpKernelContext* context, ImplCo
&prepare.fdm_output_strides,
prepare.fdm_H,
prepare.fdm_C,
reinterpret_cast<CudaT*>(output_buffer.get()),
prepare.output_tensor->Shape().Size());
Impl_Cast<CudaT, ToCudaType<bool>::MappedType>(
reinterpret_cast<CudaT*>(output_buffer.get()),
reinterpret_cast<ToCudaType<bool>::MappedType*>(prepare.output_tensor->template MutableData<bool>()),
output_size);
prepare.output_tensor->Shape().Size());
return Status::OK();
}
@ -395,14 +389,14 @@ Status CompareFunction<T, CudaT>::CompareMethod(OpKernelContext* context, ImplCo
//for other elementwise ops
template <typename T>
Status Greater<T>::ComputeInternal(OpKernelContext* context) const {
this->CompareMethod(context, &Impl_Greater);
this->CompareMethod(context, &ImplT2_Greater);
return Status::OK();
}
template <typename T>
Status Equal<T>::ComputeInternal(OpKernelContext* context) const {
this->CompareMethod(context, &Impl_Equal);
this->CompareMethod(context, &ImplT2_Equal);
return Status::OK();
}
@ -411,7 +405,7 @@ Status Equal<T>::ComputeInternal(OpKernelContext* context) const {
//for other elementwise ops
template <typename T>
Status Less<T>::ComputeInternal(OpKernelContext* context) const {
this->CompareMethod(context, &Impl_Less);
this->CompareMethod(context, &ImplT2_Less);
return Status::OK();
}

2
onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h
View file

@ -227,7 +227,7 @@ class CompareFunction : public BinaryElementwise<ShouldBroadcast> {
const TArray<fast_divmod>* 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;

52
onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu
View file

@ -21,7 +21,7 @@ namespace cuda {
fdm_H, \
fdm_C, \
output_data, \
OP_##name<T, T>(), \
OP_##name<T, T, T>(), \
count); \
}
@ -36,7 +36,22 @@ namespace cuda {
fdm_H, \
fdm_C, \
output_data, \
OP_##name<T, T1>(), \
OP_##name<T, T, T1>(), \
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<T, T1, T2>(), \
count); \
}
@ -52,6 +67,12 @@ namespace cuda {
const TArray<int64_t>* rhs_padded_strides, const T1* rhs_data, \
const TArray<fast_divmod>* 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<T, T1, T2>(int32_t output_rank, \
const TArray<int64_t>* lhs_padded_strides, const T1* lhs_data, \
const TArray<int64_t>* rhs_padded_strides, const T2* rhs_data, \
const TArray<fast_divmod>* 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

28
onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.h
View file

@ -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 <typename T, typename T1, typename T2> \
void ImplT2_##name( \
int32_t output_rank_or_simple_broadcast, \
const TArray<int64_t>* lhs_padded_strides, \
const T1* lhs_data, \
const TArray<int64_t>* rhs_padded_strides, \
const T2* rhs_data, \
const TArray<fast_divmod>* 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

20
onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl_functors.cuh
View file

@ -10,12 +10,12 @@ namespace cuda {
// define the device functors that perform the computation on scalars
#define OP_FUNCTOR_DEFINITION(name, expr) \
template <class T, class T1> \
struct OP_##name { \
__device__ __inline__ T operator()(T a, T1 b) const { \
return (expr); \
} \
#define OP_FUNCTOR_DEFINITION(name, expr) \
template <class T, class T1, class T2> \
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

2
onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu
View file

@ -17,7 +17,7 @@ struct VariadicElementwiseOpTraits;
#define DEFINE_TRAITS(VariadicElementwiseOpTag, ImplName) \
template <typename T> \
struct VariadicElementwiseOpTraits<T, VariadicElementwiseOpTag> { \
using ScalarComputeFunctor = OP_##ImplName<T, T>; \
using ScalarComputeFunctor = OP_##ImplName<T, T, T>; \
\
static void ComputeFn( \
int32_t output_rank_or_simple_broadcast, \

76
onnxruntime/core/providers/cuda/tensor/gather.cc
View file

@ -33,40 +33,6 @@ ONNX_OPERATOR_KERNEL_EX(
DataTypeImpl::GetTensorType<int64_t>()}),
Gather);
#define TYPED_FUNCTION_CALL(T) \
if (utils::IsPrimitiveDataType<T>(T_type)) { \
T* output_data = p.output_tensor->template MutableData<T>(); \
const T* input_data = p.input_tensor->template Data<T>(); \
if (utils::IsPrimitiveDataType<int32_t>(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<int32_t>(), \
reinterpret_cast<const ToCudaType<T>::MappedType*>(input_data), \
reinterpret_cast<typename ToCudaType<T>::MappedType*>(output_data), \
p.output_tensor->Shape().Size()); \
} \
return Status::OK(); \
} \
if (utils::IsPrimitiveDataType<int64_t>(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<int64_t>(), \
reinterpret_cast<const ToCudaType<T>::MappedType*>(input_data), \
reinterpret_cast<typename ToCudaType<T>::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<int>(output_block_size));
const fast_divmod divmod_block_size(gsl::narrow_cast<int>(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<int32_t>() ||
p.indices_tensor->IsDataType<int64_t>()) {
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.");
}

32
onnxruntime/core/providers/cuda/tensor/gather_elements.cc
View file

@ -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<int64_t>(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<fast_divmod> 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<int>(indices_strides[i]));
fdm_indices_strides[i] = fast_divmod(gsl::narrow_cast<int>(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<int32_t>()) {
const int32_t* indices_data = indices_tensor->template Data<int32_t>();
GatherElementsImpl<int32_t>(
if (indices_tensor->IsDataType<int32_t>() ||
indices_tensor->IsDataType<int64_t>()) {
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<int64_t>()) {
const int64_t* indices_data = indices_tensor->template Data<int64_t>();
GatherElementsImpl<int64_t>(
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");

162
onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu
View file

@ -7,114 +7,138 @@
namespace onnxruntime {
namespace cuda {
template <typename T, typename Tin>
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<const int32_t*>(index_data) + offset);
break;
case sizeof(int64_t):
return *(reinterpret_cast<const int64_t*>(index_data) + offset);
break;
default:
break;
}
// What is a sensible thing to do here?
assert(false);
return std::numeric_limits<int64_t>::max();
}
template <typename T>
__global__ void _GatherElementsKernel(
const int64_t rank,
const T* input_data,
const int64_t input_dim_along_axis,
const TArray<int64_t> input_strides,
const Tin* indices_data,
const void* indices_data,
const int64_t indices_size,
const size_t index_element_size,
const TArray<fast_divmod> 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<int>(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 <typename Tin>
void GatherElementsImpl(
const int64_t rank,
const void* input_data,
const int64_t input_dim_along_axis,
const TArray<int64_t>& input_strides,
const Tin* indices_data,
const void* indices_data,
const int64_t indices_size,
const TArray<fast_divmod>& 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<int>((indices_size + GridDim::maxThreadsPerBlock - 1) / GridDim::maxThreadsPerBlock);
dim3 block(threads_per_block);
dim3 blocksPerGrid((static_cast<int>(indices_size + block.x * thread_worksize - 1) / (block.x * thread_worksize)));
switch (element_size) {
case sizeof(int8_t):
_GatherElementsKernel<int8_t, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
rank, reinterpret_cast<const ToCudaType<int8_t>::MappedType*>(input_data), input_dim_along_axis, input_strides,
indices_data, indices_size, indices_strides,
axis, reinterpret_cast<ToCudaType<int8_t>::MappedType*>(output_data));
break;
case sizeof(int8_t): {
using CudaType = typename ToCudaType<int8_t>::MappedType;
_GatherElementsKernel<<<blocksPerGrid, block, 0>>>(
rank, reinterpret_cast<const CudaType*>(input_data), input_dim_along_axis, input_strides,
indices_data, indices_size, index_element_size, indices_strides,
axis, reinterpret_cast<CudaType*>(output_data));
} break;
case sizeof(int16_t):
_GatherElementsKernel<int16_t, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
rank, reinterpret_cast<const ToCudaType<int16_t>::MappedType*>(input_data), input_dim_along_axis, input_strides,
indices_data, indices_size, indices_strides,
axis, reinterpret_cast<ToCudaType<int16_t>::MappedType*>(output_data));
break;
case sizeof(int16_t): {
using CudaType = typename ToCudaType<int16_t>::MappedType;
_GatherElementsKernel<<<blocksPerGrid, block, 0>>>(
rank, reinterpret_cast<const CudaType*>(input_data), input_dim_along_axis, input_strides,
indices_data, indices_size, index_element_size, indices_strides,
axis, reinterpret_cast<CudaType*>(output_data));
} break;
case sizeof(int32_t):
_GatherElementsKernel<int32_t, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
rank, reinterpret_cast<const ToCudaType<int32_t>::MappedType*>(input_data), input_dim_along_axis, input_strides,
indices_data, indices_size, indices_strides,
axis, reinterpret_cast<ToCudaType<int32_t>::MappedType*>(output_data));
break;
case sizeof(int32_t): {
using CudaType = typename ToCudaType<int32_t>::MappedType;
_GatherElementsKernel<<<blocksPerGrid, block, 0>>>(
rank, reinterpret_cast<const CudaType*>(input_data), input_dim_along_axis, input_strides,
indices_data, indices_size, index_element_size, indices_strides,
axis, reinterpret_cast<CudaType*>(output_data));
} break;
case sizeof(int64_t):
_GatherElementsKernel<int64_t, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
rank, reinterpret_cast<const ToCudaType<int64_t>::MappedType*>(input_data), input_dim_along_axis, input_strides,
indices_data, indices_size, indices_strides,
axis, reinterpret_cast<ToCudaType<int64_t>::MappedType*>(output_data));
break;
case sizeof(int64_t): {
using CudaType = typename ToCudaType<int64_t>::MappedType;
_GatherElementsKernel<<<blocksPerGrid, block, 0>>>(
rank, reinterpret_cast<const CudaType*>(input_data), input_dim_along_axis, input_strides,
indices_data, indices_size, index_element_size, indices_strides,
axis, reinterpret_cast<CudaType*>(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<int32_t>(
const int64_t rank,
const void* input_data,
const int64_t input_dim_along_axis,
const TArray<int64_t>& input_strides,
const int32_t* indices_data,
const int64_t indices_size,
const TArray<fast_divmod>& indices_strides,
const int64_t axis,
void* output_data,
size_t element_size);
template void GatherElementsImpl<int64_t>(
const int64_t rank,
const void* input_data,
const int64_t input_dim_along_axis,
const TArray<int64_t>& input_strides,
const int64_t* indices_data,
const int64_t indices_size,
const TArray<fast_divmod>& indices_strides,
const int64_t axis,
void* output_data,
size_t element_size);
} // namespace cuda
} // namespace cuda
} // namespace onnxruntime

6
onnxruntime/core/providers/cuda/tensor/gather_elements_impl.h
View file

@ -9,18 +9,18 @@
namespace onnxruntime {
namespace cuda {
template <typename Tin>
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<int64_t>& input_strides,
const Tin* indices_data,
const void* indices_data,
const int64_t indices_size,
const TArray<fast_divmod>& 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

90
onnxruntime/core/providers/cuda/tensor/gather_impl.cu
View file

@ -7,13 +7,30 @@
namespace onnxruntime {
namespace cuda {
template <typename T, typename Tin>
__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<const int32_t*>(index_data) + offset);
break;
case sizeof(int64_t):
return *(reinterpret_cast<const int64_t*>(index_data) + offset);
break;
default:
break;
}
// What is a sensible thing to do here?
assert(false);
return std::numeric_limits<int64_t>::max();
}
template <typename T>
__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 <typename T, typename Tin>
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<float>(N) / GridDim::maxThreadsPerBlock));
_GatherKernel<T, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
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<int8_t>::MappedType;
_GatherKernel<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
input_block_size, indices_max, output_block_size, block_size, indices_data, index_element_size,
reinterpret_cast<const CudaType*>(input_data), reinterpret_cast<CudaType*>(output_data), (CUDA_LONG)N);
} break;
case sizeof(int16_t): {
using CudaType = typename ToCudaType<int16_t>::MappedType;
_GatherKernel<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
input_block_size, indices_max, output_block_size, block_size, indices_data, index_element_size,
reinterpret_cast<const CudaType*>(input_data), reinterpret_cast<CudaType*>(output_data), (CUDA_LONG)N);
} break;
case sizeof(int32_t): {
using CudaType = typename ToCudaType<int32_t>::MappedType;
_GatherKernel<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
input_block_size, indices_max, output_block_size, block_size, indices_data, index_element_size,
reinterpret_cast<const CudaType*>(input_data), reinterpret_cast<CudaType*>(output_data), (CUDA_LONG)N);
} break;
case sizeof(int64_t): {
using CudaType = typename ToCudaType<int64_t>::MappedType;
_GatherKernel<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
input_block_size, indices_max, output_block_size, block_size, indices_data, index_element_size,
reinterpret_cast<const CudaType*>(input_data), reinterpret_cast<CudaType*>(output_data), (CUDA_LONG)N);
} break;
default:
ORT_THROW("Unsupported element size by the Gather CUDA kernel");
}
}
#define SPECIALIZED_IMPL(T) \
template void GatherImpl<T, int32_t>(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<T, int64_t>(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

9
onnxruntime/core/providers/cuda/tensor/gather_impl.h
View file

@ -8,15 +8,16 @@
namespace onnxruntime {
namespace cuda {
template <typename T, typename Tin>
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