From ce9d13495284ae77435738fe65f3db349268bec2 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Thu, 1 Jul 2021 14:30:00 +0800 Subject: [PATCH] gather elements optimization (#8154) --- .../providers/cuda/tensor/gather_elements.cc | 14 +++-- .../cuda/tensor/gather_elements_impl.cu | 53 +++++++++---------- .../cuda/tensor/gather_elements_impl.h | 4 +- 3 files changed, 35 insertions(+), 36 deletions(-) diff --git a/onnxruntime/core/providers/cuda/tensor/gather_elements.cc b/onnxruntime/core/providers/cuda/tensor/gather_elements.cc index 08312899de..8e653e651e 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather_elements.cc +++ b/onnxruntime/core/providers/cuda/tensor/gather_elements.cc @@ -61,8 +61,11 @@ Status GatherElements::ComputeInternal(OpKernelContext* context) const { if (indices_shape.Size() == 0) return Status::OK(); - TensorPitches input_strides(input_dims); - TArray gpu_input_strides(input_strides); + // Set stride along axis to 0 so we don't need IF statment to check in kernel. + TensorPitches masked_input_strides(input_dims); + int64_t input_stride_along_axis = masked_input_strides[axis]; + masked_input_strides[axis] = 0; + TArray gpu_masked_input_strides(masked_input_strides); TArray fdm_indices_strides(indices_rank); TensorPitches indices_strides(indices_dims); @@ -77,14 +80,15 @@ Status GatherElements::ComputeInternal(OpKernelContext* context) const { indices_tensor->IsDataType()) { GatherElementsImpl( Stream(), - input_rank, + // Save one divmod in kernel if axis is the last dim. + input_rank == axis + 1 ? input_rank - 1 : input_rank, input_tensor->DataRaw(), input_dims[axis], - gpu_input_strides, + input_stride_along_axis, + gpu_masked_input_strides, indices_tensor->DataRaw(), indices_size, fdm_indices_strides, - axis, output_tensor->MutableDataRaw(), element_size, index_element_size); diff --git a/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu b/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu index 87920a7fcb..4886193f56 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu @@ -33,12 +33,12 @@ __global__ void _GatherElementsKernel( const int64_t rank, const T* input_data, const int64_t input_dim_along_axis, - const TArray input_strides, + const int64_t input_stride_along_axis, + const TArray masked_input_strides, // the stride along axis is 0 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) { CUDA_LONG indices_index = threads_per_block * thread_worksize * blockIdx.x + threadIdx.x; @@ -46,19 +46,20 @@ __global__ void _GatherElementsKernel( #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) { + #pragma unroll + for (auto i = 0; i < indices_strides.Capacity(); i++) { + if (i >= rank) { + break; + } + indices_strides[i].divmod(remain, dim, remain); - data_idx += input_strides[i] * dim; + data_idx += masked_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 @@ -68,13 +69,7 @@ __global__ void _GatherElementsKernel( 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; - } + data_idx += input_stride_along_axis * dim; output_data[indices_index] = input_data[data_idx]; indices_index += threads_per_block; @@ -87,11 +82,11 @@ void GatherElementsImpl( const int64_t rank, const void* input_data, const int64_t input_dim_along_axis, - const TArray& input_strides, + const int64_t input_stride_along_axis, + const TArray& masked_input_strides, 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 index_element_size) { @@ -105,33 +100,33 @@ void GatherElementsImpl( 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)); + rank, reinterpret_cast(input_data), input_dim_along_axis, input_stride_along_axis, + masked_input_strides, indices_data, indices_size, index_element_size, indices_strides, + reinterpret_cast(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)); + rank, reinterpret_cast(input_data), input_dim_along_axis, input_stride_along_axis, + masked_input_strides, indices_data, indices_size, index_element_size, indices_strides, + reinterpret_cast(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)); + rank, reinterpret_cast(input_data), input_dim_along_axis, input_stride_along_axis, + masked_input_strides, indices_data, indices_size, index_element_size, indices_strides, + reinterpret_cast(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)); + rank, reinterpret_cast(input_data), input_dim_along_axis, input_stride_along_axis, + masked_input_strides, indices_data, indices_size, index_element_size, indices_strides, + reinterpret_cast(output_data)); } break; // should not reach here as we validate if the all relevant types are supported in the Compute method diff --git a/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.h b/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.h index 920415678b..f973ecc684 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.h +++ b/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.h @@ -14,11 +14,11 @@ 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 int64_t input_stride_along_axis, + const TArray& masked_input_strides, 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 index_element_size);