From 2f408f757ec54fb1659bca74f358795547030c33 Mon Sep 17 00:00:00 2001 From: Hariharan Seshadri Date: Fri, 16 Jul 2021 09:44:16 -0700 Subject: [PATCH] Improve performance of Pad CUDA kernel (#8408) --- onnxruntime/core/providers/cuda/tensor/pad.cc | 61 +++++++- .../core/providers/cuda/tensor/pad_impl.cu | 136 ++++++++++++++++-- .../core/providers/cuda/tensor/pad_impl.h | 18 ++- 3 files changed, 205 insertions(+), 10 deletions(-) diff --git a/onnxruntime/core/providers/cuda/tensor/pad.cc b/onnxruntime/core/providers/cuda/tensor/pad.cc index 26f5f975ce..1dfccb0ad5 100644 --- a/onnxruntime/core/providers/cuda/tensor/pad.cc +++ b/onnxruntime/core/providers/cuda/tensor/pad.cc @@ -41,6 +41,30 @@ namespace cuda { .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ Pad); +static bool IsNCHWInputWithPaddingAlongHAndW(size_t input_rank, + const TArray& lower_pads, + const TArray& upper_pads) { + if (input_rank == 2) { // N = 1 and C = 1 + return true; + } + + // Is CHW input AND no padding along C dim + if (input_rank == 3 && + lower_pads[0] == 0 && // start padding along C + upper_pads[0] == 0) { // end padding along C + return true; + } + + // Is NCHW input AND no padding along N and C dims + if (input_rank == 4 && + lower_pads[0] == 0 && lower_pads[1] == 0 && // start padding along N and C + upper_pads[0] == 0 && upper_pads[1] == 0) { // end padding along N and C + return true; + } + + return false; +} + template typename ToCudaType::MappedType ToCudaValue(const T& value) { return value; @@ -119,6 +143,7 @@ Status Pad::ComputeInternal(OpKernelContext* ctx) const { upper_pads[i] = (*p_pads)[i + dimension_count] + (*p_slices)[i + dimension_count]; output_dims[i] += lower_pads[i] + upper_pads[i]; } + TensorShape output_shape(output_dims); // special case when there is a dim value of 0 in the shape. behavior depends on mode @@ -127,6 +152,7 @@ Status Pad::ComputeInternal(OpKernelContext* ctx) const { } auto& output_tensor = *ctx->Output(0, output_shape); + if (std::all_of(p_pads->begin(), p_pads->end(), [](const int64_t v) { return v == 0; }) && std::all_of(p_slices->begin(), p_slices->end(), [](const int64_t v) { return v == 0; }) && output_shape.Size() > 0) { @@ -137,6 +163,40 @@ Status Pad::ComputeInternal(OpKernelContext* ctx) const { return Status::OK(); } + if (IsNCHWInputWithPaddingAlongHAndW(static_cast(dimension_count), lower_pads, upper_pads)) { + // If we have entered here, it means the input can only be 4-D (NCHW), 3-D (CHW), or 2-D (HW) + + // NCHW input + int height_dim = 2; + int width_dim = 3; + + if (dimension_count == 3) { // CHW input + height_dim = 1; + width_dim = 2; + } else if (dimension_count == 2) { // HW input + height_dim = 0; + width_dim = 1; + } + + PadNCHWInputWithPaddingAlongHAndWImpl( + Stream(), + dimension_count == 4 ? input_dims[0] : 1, + dimension_count == 4 ? input_dims[1] : (dimension_count == 3 ? input_dims[0] : 1), + input_dims[height_dim], + output_dims[height_dim], + input_dims[width_dim], + output_dims[width_dim], + lower_pads[height_dim], + lower_pads[width_dim], + value, + static_cast(mode_), + reinterpret_cast::MappedType*>(input_tensor.template Data()), + reinterpret_cast::MappedType*>(output_tensor.template MutableData()), + output_tensor.Shape().Size()); + + return Status::OK(); + } + TArray fdm_output_strides(dimension_count); TensorPitches output_strides(output_dims); for (auto i = 0; i < dimension_count; i++) { @@ -149,7 +209,6 @@ Status Pad::ComputeInternal(OpKernelContext* ctx) const { input_dims, input_strides, lower_pads, - upper_pads, value, static_cast(mode_), reinterpret_cast::MappedType*>(input_tensor.template Data()), diff --git a/onnxruntime/core/providers/cuda/tensor/pad_impl.cu b/onnxruntime/core/providers/cuda/tensor/pad_impl.cu index 2e1820f198..62630df7ad 100644 --- a/onnxruntime/core/providers/cuda/tensor/pad_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/pad_impl.cu @@ -20,7 +20,6 @@ __global__ void _PadKernel( const TArray input_dims, const TArray input_strides, const TArray lower_pads, - const TArray upper_pads, const T pad_value, const T* input_data, const TArray fdm_output_strides, @@ -67,6 +66,69 @@ __global__ void _PadKernel( output_data[id] = use_pad_value ? (T)pad_value : input_data[input_index]; } +template +__global__ void _PadNCHWInputWithPaddingAlongHAndWKernel( + const int64_t n, // Batch + const int64_t c, // Channel + const int64_t input_height, + const int64_t output_height, + const int64_t input_width, + const int64_t output_width, + const int64_t pad_height_start, + const int64_t pad_width_start, + const T pad_value, + const T* input_data, + T* output_data, + const size_t N) { + CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N); + + const int current_output_width = id % output_width; + int nc_index = id / output_width; + const int current_output_height = nc_index % output_height; + nc_index /= output_height; + + int current_input_height = current_output_height - pad_height_start; + int current_input_width = current_output_width - pad_width_start; + + switch ((PadMode)pad_mode) { + case PadMode::Constant: + output_data[id] = (current_input_height < 0 || + current_input_width < 0 || + current_input_height >= input_height || + current_input_width >= input_width) + ? pad_value + : input_data[(nc_index * input_height + + current_input_height) * + input_width + + current_input_width]; + break; + + case PadMode::Edge: + current_input_height = std::max(0, std::min(current_input_height, static_cast(input_height - 1))); + current_input_width = std::max(0, std::min(current_input_width, static_cast(input_width - 1))); + output_data[id] = input_data[(nc_index * input_height + + current_input_height) * + input_width + + current_input_width]; + break; + + case PadMode::Reflect: + current_input_height = std::max(current_input_height, -current_input_height); + current_input_height = std::min(static_cast(current_input_height), + 2 * static_cast(input_height) - current_input_height - 2); + + current_input_width = std::max(current_input_width, -current_input_width); + current_input_width = std::min(static_cast(current_input_width), + 2 * static_cast(input_width) - current_input_width - 2); + + output_data[id] = input_data[(nc_index * input_height + + current_input_height) * + input_width + + current_input_width]; + break; + } +} + template void PadImpl( cudaStream_t stream, @@ -74,38 +136,96 @@ void PadImpl( const TArray& input_dims, const TArray& input_strides, const TArray& lower_pads, - const TArray& upper_pads, const T pad_value, const int pad_mode, const T* input_data, const TArray& fdm_output_strides, T* output_data, const size_t N) { - if (N == 0) // special case where there's a dim value of 0 in the output shape + if (N == 0) // special case where there's a dim value of 0 in the output shape return; int blocksPerGrid = (int)(ceil(static_cast(N) / GridDim::maxThreadsPerBlock)); switch (pad_mode) { case 0: _PadKernel<<>>( - shape_rank, input_dims, input_strides, lower_pads, upper_pads, + shape_rank, input_dims, input_strides, lower_pads, pad_value, input_data, fdm_output_strides, output_data, N); break; case 1: _PadKernel<<>>( - shape_rank, input_dims, input_strides, lower_pads, upper_pads, + shape_rank, input_dims, input_strides, lower_pads, pad_value, input_data, fdm_output_strides, output_data, N); break; case 2: _PadKernel<<>>( - shape_rank, input_dims, input_strides, lower_pads, upper_pads, + shape_rank, input_dims, input_strides, lower_pads, pad_value, input_data, fdm_output_strides, output_data, N); break; } } -#define SPECIALIZED_IMPL(T) \ - template void PadImpl(cudaStream_t stream, const size_t shape_rank, const TArray& input_dims, const TArray& input_strides, const TArray& lower_pads, const TArray& upper_pads, const T pad_value, const int pad_mode, const T* input_data, const TArray& fdm_output_strides, T* output_data, const size_t N); +template +void PadNCHWInputWithPaddingAlongHAndWImpl( + cudaStream_t stream, + const int64_t n, // Batch + const int64_t c, // Channel + const int64_t input_height, + const int64_t output_height, + const int64_t input_width, + const int64_t output_width, + const int64_t pad_height_start, + const int64_t pad_width_start, + const T pad_value, + const int pad_mode, + const T* input_data, + T* output_data, + const size_t N) { + if (N == 0) // special case where there's a dim value of 0 in the output shape + return; + + int blocksPerGrid = (int)(ceil(static_cast(N) / GridDim::maxThreadsPerBlock)); + switch (pad_mode) { + case 0: + _PadNCHWInputWithPaddingAlongHAndWKernel<<>>( + n, c, input_height, output_height, input_width, output_width, + pad_height_start, pad_width_start, + pad_value, input_data, output_data, N); + break; + case 1: + _PadNCHWInputWithPaddingAlongHAndWKernel<<>>( + n, c, input_height, output_height, input_width, output_width, + pad_height_start, pad_width_start, + pad_value, input_data, output_data, N); + break; + case 2: + _PadNCHWInputWithPaddingAlongHAndWKernel<<>>( + n, c, input_height, output_height, input_width, output_width, + pad_height_start, pad_width_start, + pad_value, input_data, output_data, N); + break; + } +} + +#define SPECIALIZED_IMPL(T) \ + template void PadImpl(cudaStream_t stream, const size_t shape_rank, \ + const TArray& input_dims, const TArray& input_strides, \ + const TArray& lower_pads, \ + const T pad_value, \ + const int pad_mode, \ + const T* input_data, \ + const TArray& fdm_output_strides, \ + T* output_data, \ + const size_t N); \ + template void PadNCHWInputWithPaddingAlongHAndWImpl(cudaStream_t stream, const int64_t n, const int64_t c, \ + const int64_t input_height, const int64_t output_height, \ + const int64_t input_width, const int64_t output_width, \ + const int64_t pad_height_start, \ + const int64_t pad_width_start, \ + const T pad_value, \ + const int pad_mode, \ + const T* input_data, T* output_data, \ + const size_t N); SPECIALIZED_IMPL(float) SPECIALIZED_IMPL(double) diff --git a/onnxruntime/core/providers/cuda/tensor/pad_impl.h b/onnxruntime/core/providers/cuda/tensor/pad_impl.h index 8be69dcb1f..dc700ea230 100644 --- a/onnxruntime/core/providers/cuda/tensor/pad_impl.h +++ b/onnxruntime/core/providers/cuda/tensor/pad_impl.h @@ -8,6 +8,23 @@ namespace onnxruntime { namespace cuda { +template +void PadNCHWInputWithPaddingAlongHAndWImpl( + cudaStream_t stream, + const int64_t n, // Batch + const int64_t c, // Channel + const int64_t input_height, + const int64_t output_height, + const int64_t input_width, + const int64_t output_width, + const int64_t pad_height_start, + const int64_t pad_width_start, + const T pad_value, + const int pad_mode, + const T* input_data, + T* output_data, + const size_t N); + template void PadImpl( cudaStream_t stream, @@ -15,7 +32,6 @@ void PadImpl( const TArray& input_dims, const TArray& input_strides, const TArray& lower_pads, - const TArray& upper_pads, const T pad_value, const int pad_mode, const T* input_data,