saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Optimize Slice Kernel by Removing If-statement (#4753)

Browse source 
* Slice kernel optimization.

* remove space

Co-authored-by: Vincent Wang <weicwang@AiFramework2080ti2.corp.microsoft.com>
This commit is contained in:
Vincent Wang 2020-08-12 16:36:03 +08:00 committed by GitHub
parent b7254551f0
commit 7e955960f1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 72 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

124
onnxruntime/core/providers/cuda/tensor/slice_impl.cu
View file

@ -8,36 +8,57 @@
namespace onnxruntime {
namespace cuda {
template <bool is_grad, typename T>
__global__ void _SliceKernel(const int32_t dimension_count,
const TArray<int64_t> starts,
template <bool is_grad, int DIMS, int NumThreadsPerBlock, int NumElementsPerThread, typename T>
__global__ void _SliceKernel(const TArray<int64_t> starts,
const TArray<int64_t> steps,
const TArray<int64_t> input_strides,
const TArray<fast_divmod> output_strides,
const T* input_data,
T* output_data,
const CUDA_LONG N) {
CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
CUDA_LONG input_index = 0;
int div;
int mod = id;
int value = id;
int dim = 0;
#pragma unroll
for (; dim < starts.Capacity(); ++dim) {
if (dim >= dimension_count - 1) {
break;
}
CUDA_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
CUDA_LONG input_indices[NumElementsPerThread];
output_strides[dim].divmod(value, div, mod);
input_index += (starts[dim] + div * steps[dim]) * input_strides[dim];
value = mod;
CUDA_LONG id = start;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
CUDA_LONG input_index = 0;
int div;
int mod = id;
int value = id;
int dim = 0;
#pragma unroll
for (; dim < DIMS - 1; ++dim) {
output_strides[dim].divmod(value, div, mod);
input_index += (starts[dim] + div * steps[dim]) * input_strides[dim];
value = mod;
}
input_index += starts[dim] + mod * steps[dim];
input_indices[i] = input_index;
id += NumThreadsPerBlock;
}
}
if (is_grad) {
id = start;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
output_data[input_indices[i]] = input_data[id];
id += NumThreadsPerBlock;
}
}
} else {
id = start;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
output_data[id] = input_data[input_indices[i]];
id += NumThreadsPerBlock;
}
}
}
input_index += starts[dim] + mod * steps[dim];
if (is_grad)
output_data[input_index] = input_data[id];
else
output_data[id] = input_data[input_index];
}
Status SliceImpl(const size_t element_size,
@ -66,6 +87,30 @@ Status SliceImplGrad(const size_t element_size,
output_data, N);
}
#define HANDLE_DIMS(ELEMENT_TYPE, DIMS) \
case DIMS: { \
_SliceKernel<is_grad, DIMS, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread> \
<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>( \
starts, steps, input_strides, output_strides, \
reinterpret_cast<const ToCudaType<ELEMENT_TYPE>::MappedType*>(input_data), \
reinterpret_cast<ToCudaType<ELEMENT_TYPE>::MappedType*>(output_data), \
(CUDA_LONG)N); \
} break
#define HANDLE_ELEMENT_TYPE(ELEMENT_TYPE) \
case sizeof(ELEMENT_TYPE): { \
switch (dimension_count) { \
HANDLE_DIMS(ELEMENT_TYPE, 1); \
HANDLE_DIMS(ELEMENT_TYPE, 2); \
HANDLE_DIMS(ELEMENT_TYPE, 3); \
HANDLE_DIMS(ELEMENT_TYPE, 4); \
HANDLE_DIMS(ELEMENT_TYPE, 5); \
HANDLE_DIMS(ELEMENT_TYPE, 6); \
HANDLE_DIMS(ELEMENT_TYPE, 7); \
HANDLE_DIMS(ELEMENT_TYPE, 8); \
} \
} break
template <bool is_grad>
Status SliceImplEx(const size_t element_size,
const int32_t dimension_count,
@ -76,37 +121,12 @@ Status SliceImplEx(const size_t element_size,
const void* input_data,
void* output_data,
const size_t N) {
int blocksPerGrid = (int)(ceil(static_cast<float>(N) / GridDim::maxThreadsPerBlock));
int blocksPerGrid = static_cast<int>(CeilDiv(N, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
switch (element_size) {
case sizeof(int8_t):
_SliceKernel<is_grad><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
dimension_count, starts, steps, input_strides, output_strides,
reinterpret_cast<const ToCudaType<int8_t>::MappedType*>(input_data),
reinterpret_cast<ToCudaType<int8_t>::MappedType*>(output_data),
(CUDA_LONG)N);
break;
case sizeof(int16_t):
_SliceKernel<is_grad><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
dimension_count, starts, steps, input_strides, output_strides,
reinterpret_cast<const ToCudaType<int16_t>::MappedType*>(input_data),
reinterpret_cast<ToCudaType<int16_t>::MappedType*>(output_data),
(CUDA_LONG)N);
break;
case sizeof(int32_t):
_SliceKernel<is_grad><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
dimension_count, starts, steps, input_strides, output_strides,
reinterpret_cast<const ToCudaType<int32_t>::MappedType*>(input_data),
reinterpret_cast<ToCudaType<int32_t>::MappedType*>(output_data),
(CUDA_LONG)N);
break;
case sizeof(int64_t):
_SliceKernel<is_grad><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
dimension_count, starts, steps, input_strides, output_strides,
reinterpret_cast<const ToCudaType<int64_t>::MappedType*>(input_data),
reinterpret_cast<ToCudaType<int64_t>::MappedType*>(output_data),
(CUDA_LONG)N);
break;
HANDLE_ELEMENT_TYPE(int8_t);
HANDLE_ELEMENT_TYPE(int16_t);
HANDLE_ELEMENT_TYPE(int32_t);
HANDLE_ELEMENT_TYPE(int64_t);
default:
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Type not supported for Slice operator");
}