Improve performance of Pad CUDA kernel (#8408)

onnxruntime/core/providers/cuda/tensor/pad.cc

@@ -41,6 +41,30 @@ namespace cuda {
           .TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
       Pad<T>);
 
+static bool IsNCHWInputWithPaddingAlongHAndW(size_t input_rank,
+                                             const TArray<int64_t>& lower_pads,
+                                             const TArray<int64_t>& 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 T>
 typename ToCudaType<T>::MappedType ToCudaValue(const T& value) {
   return value;
@@ -119,6 +143,7 @@ Status Pad<T>::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<T>::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<T>::ComputeInternal(OpKernelContext* ctx) const {
     return Status::OK();
   }
 
+  if (IsNCHWInputWithPaddingAlongHAndW(static_cast<size_t>(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<int>(mode_),
+        reinterpret_cast<const typename ToCudaType<T>::MappedType*>(input_tensor.template Data<T>()),
+        reinterpret_cast<typename ToCudaType<T>::MappedType*>(output_tensor.template MutableData<T>()),
+        output_tensor.Shape().Size());
+
+    return Status::OK();
+  }
+
   TArray<fast_divmod> fdm_output_strides(dimension_count);
   TensorPitches output_strides(output_dims);
   for (auto i = 0; i < dimension_count; i++) {
@@ -149,7 +209,6 @@ Status Pad<T>::ComputeInternal(OpKernelContext* ctx) const {
       input_dims,
       input_strides,
       lower_pads,
-      upper_pads,
       value,
       static_cast<int>(mode_),
       reinterpret_cast<const typename ToCudaType<T>::MappedType*>(input_tensor.template Data<T>()),

onnxruntime/core/providers/cuda/tensor/pad_impl.cu

@@ -20,7 +20,6 @@ __global__ void _PadKernel(
     const TArray<int64_t> input_dims,
     const TArray<int64_t> input_strides,
     const TArray<int64_t> lower_pads,
-    const TArray<int64_t> upper_pads,
     const T pad_value,
     const T* input_data,
     const TArray<fast_divmod> fdm_output_strides,
@@ -67,6 +66,69 @@ __global__ void _PadKernel(
   output_data[id] = use_pad_value ? (T)pad_value : input_data[input_index];
 }
 
+template <typename T, int pad_mode>
+__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<int>(input_height - 1)));
+      current_input_width = std::max(0, std::min(current_input_width, static_cast<int>(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<int>(current_input_height),
+                                      2 * static_cast<int>(input_height) - current_input_height - 2);
+
+      current_input_width = std::max(current_input_width, -current_input_width);
+      current_input_width = std::min(static_cast<int>(current_input_width),
+                                     2 * static_cast<int>(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 <typename T>
 void PadImpl(
     cudaStream_t stream,
@@ -74,38 +136,96 @@ void PadImpl(
     const TArray<int64_t>& input_dims,
     const TArray<int64_t>& input_strides,
     const TArray<int64_t>& lower_pads,
-    const TArray<int64_t>& upper_pads,
     const T pad_value,
     const int pad_mode,
     const T* input_data,
     const TArray<fast_divmod>& 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<float>(N) / GridDim::maxThreadsPerBlock));
   switch (pad_mode) {
     case 0:
       _PadKernel<T, 0><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
-          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<T, 1><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
-          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<T, 2><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
-          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<T>(cudaStream_t stream, const size_t shape_rank, const TArray<int64_t>& input_dims, const TArray<int64_t>& input_strides, const TArray<int64_t>& lower_pads, const TArray<int64_t>& upper_pads, const T pad_value, const int pad_mode, const T* input_data, const TArray<fast_divmod>& fdm_output_strides, T* output_data, const size_t N);
+template <typename T>
+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<float>(N) / GridDim::maxThreadsPerBlock));
+  switch (pad_mode) {
+    case 0:
+      _PadNCHWInputWithPaddingAlongHAndWKernel<T, 0><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
+          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<T, 1><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
+          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<T, 2><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
+          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<T>(cudaStream_t stream, const size_t shape_rank,                                          \
+                           const TArray<int64_t>& input_dims, const TArray<int64_t>& input_strides,               \
+                           const TArray<int64_t>& lower_pads,                                                     \
+                           const T pad_value,                                                                     \
+                           const int pad_mode,                                                                    \
+                           const T* input_data,                                                                   \
+                           const TArray<fast_divmod>& fdm_output_strides,                                         \
+                           T* output_data,                                                                        \
+                           const size_t N);                                                                       \
+  template void PadNCHWInputWithPaddingAlongHAndWImpl<T>(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)

onnxruntime/core/providers/cuda/tensor/pad_impl.h

@@ -8,6 +8,23 @@
 namespace onnxruntime {
 namespace cuda {
 
+template <typename T>
+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 <typename T>
 void PadImpl(
     cudaStream_t stream,
@@ -15,7 +32,6 @@ void PadImpl(
     const TArray<int64_t>& input_dims,
     const TArray<int64_t>& input_strides,
     const TArray<int64_t>& lower_pads,
-    const TArray<int64_t>& upper_pads,
     const T pad_value,
     const int pad_mode,
     const T* input_data,