Skip to content

Commit f79df44

Browse files
titaiwangmsCopilot
titaiwangms
and
Copilot
authored
Fix CUDA AveragePool wrong results with asymmetric padding and dilation (#29631)
### Summary Fixes wrong `AveragePool` output on the **CUDA** EP whenever cuDNN's pooling descriptor cannot represent the requested pooling — i.e. **asymmetric padding** or **dilation > 1**. This is the CUDA-EP counterpart of the CPU fix in #29629 and, together with it, the JSEP shape fix in #29627, closes out the CUDA leg of pytorch/pytorch#183528. ### Root cause `CudnnPoolingDescriptor::Set` copies only the **begin** pads (`pads[0..rank)`) into the cuDNN pooling descriptor, which stores a *single symmetric pad value per axis* and applies it to both sides. The ONNX **end** pads (`pads[rank..2*rank)`) are silently dropped. As a result, **all** asymmetric-pad AveragePool on CUDA is wrong: - explicit asymmetric `pads`, - `auto_pad = SAME_UPPER` / `SAME_LOWER` (which produce naturally asymmetric pads), - `ceil_mode = 1` boundary windows. Separately, the cuDNN pooling descriptor has **no dilation parameter at all**, so any `dilations > 1` is also silently ignored — even with symmetric pads. Example divergence from the CPU reference (QA probe): 1D `pad(0,3)`, k7, s3, `ceil_mode=1`, `count_include_pad=1` gave CUDA `[4, 6.5, 8]` vs. correct `[4, 5.571, 4]`; 2D `pad(0,0,3,3)` gave `71.0` vs. correct `17.75`. ### Fix Add a custom CUDA average-pool kernel (`avg_pool_impl.cu` / `.h`, modeled on `max_pool_with_index.cu`) that honors **per-side pads** and **dilation**, and computes the `count_include_pad` divisor exactly like the CPU v19 reference functor (`AveragePool{1,2,3}DTask`): - include-pad divisor clamps the window end to `input + pad_tail` (drops the ceil-mode phantom cells), dividing by `∏ (1 + (end - start - 1)/dilation)`; - exclude-pad divisor counts only in-bounds cells. In `Pool<T, AveragePool, Layout>::ComputeInternal`, a cheap dispatch guard routes to the custom kernel **only** when the pooling is non-global **and** (`asymmetric pads` **or** `!default_dilations`). Every symmetric, non-dilated case — the overwhelmingly common path, including **all** `GlobalAveragePool` — stays on the existing cuDNN fast path, so there is **zero perf regression** on the common path. `GlobalAveragePool` is excluded explicitly because `PoolAttributes` leaves its `kernel_shape`/`strides`/`dilations` unpopulated. `MaxPool` is unaffected (it ignores pad cells; `MaxPool<8>` already routes dilation to its own custom kernel via the same `!default_dilations` check we mirror here). Covers fp32, fp64, fp16, and bf16 (fp16/bf16 accumulate in float). ### Tests Added CUDA-**un-excluded** parity tests in `pool_op_test.cc` — the CUDA leg actually runs and must match the CPU reference oracle: - asymmetric tail-pad 1D/2D, include- and exclude-pad; - `auto_pad=SAME_UPPER` and `SAME_LOWER` (naturally asymmetric); - **symmetric-pad + dilation>1** (wrong on cuDNN, correct via the kernel — locks in the dilation guard); - a symmetric-pad regression case (must stay on cuDNN and remain correct); - fp16; - a `MaxPool` asymmetric-pad case confirming MaxPool is unaffected. The `ceil_mode + count_include_pad` cases use opset 19 so the CPU leg runs the already-correct v19 reference functor and validates the CUDA kernel independently of the separate opset-7..18 CPU MLAS fix (#29629); CUDA routing is opset-independent. Full `PoolTest` suite passes on an A100 (53 passed / 2 DML-skipped / 0 failures). ### Related - CPU (a): #29629 — opset 7-18 MLAS `ceil_mode + count_include_pad` divisor fix. - JSEP (c): #29627 — JSEP pooling output shape honoring `ceil_mode`. - pytorch/pytorch#183528. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent f4aa2b4 commit f79df44

4 files changed

Lines changed: 607 additions & 6 deletions

File tree

‎onnxruntime/core/providers/cuda/nn/avg_pool_impl.cu‎

Lines changed: 247 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,247 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
#include "avg_pool_impl.h"
5+
6+
#include "core/providers/cuda/cu_inc/common.cuh"
7+
#include "core/providers/cuda/shared_inc/fast_divmod.h"
8+
#include "core/providers/cuda/shared_inc/cuda_utils.h"
9+
10+
namespace onnxruntime {
11+
namespace cuda {
12+
13+
// Accumulate half in float for precision; keep native type otherwise.
14+
template <typename T>
15+
struct AveragePoolAccumulator {
16+
using type = T;
17+
};
18+
template <>
19+
struct AveragePoolAccumulator<half> {
20+
using type = float;
21+
};
22+
template <>
23+
struct AveragePoolAccumulator<BFloat16> {
24+
using type = float;
25+
};
26+
27+
template <typename T, bool Layout>
28+
__global__ void AveragePoolWithPadKernel(
29+
int64_t channels,
30+
int64_t height,
31+
int64_t width,
32+
int64_t depth,
33+
int64_t pooled_height,
34+
int64_t pooled_width,
35+
int64_t pooled_depth,
36+
int64_t kernel_h,
37+
int64_t kernel_w,
38+
int64_t kernel_d,
39+
int64_t stride_h,
40+
int64_t stride_w,
41+
int64_t stride_d,
42+
int64_t pad_h_head,
43+
int64_t pad_w_head,
44+
int64_t pad_d_head,
45+
int64_t pad_h_tail,
46+
int64_t pad_w_tail,
47+
int64_t pad_d_tail,
48+
int64_t dilation_h,
49+
int64_t dilation_w,
50+
int64_t dilation_d,
51+
fast_divmod fdm_c,
52+
fast_divmod fdm_h,
53+
fast_divmod fdm_w,
54+
fast_divmod fdm_d,
55+
bool count_include_pad,
56+
const T* p_input,
57+
int64_t output_size,
58+
T* p_output) {
59+
int id = blockIdx.x * blockDim.x + threadIdx.x;
60+
if (id >= output_size) return;
61+
62+
auto compute_offset =
63+
[height, width, depth, channels](int n_index, int c_index, int h_index, int w_index, int d_index) -> int64_t {
64+
if constexpr (Layout == LAYOUT_NCHW) {
65+
return (((n_index * channels + c_index) * height + h_index) * width + w_index) * depth + d_index;
66+
} else if constexpr (Layout == LAYOUT_NHWC) {
67+
return (((n_index * height + h_index) * width + w_index) * depth + d_index) * channels + c_index;
68+
}
69+
};
70+
71+
int d_index, w_index, h_index, c_index, n_index, id_tmp;
72+
if constexpr (Layout == LAYOUT_NCHW) {
73+
fdm_d.divmod(id, id_tmp, d_index);
74+
fdm_w.divmod(id_tmp, id_tmp, w_index);
75+
fdm_h.divmod(id_tmp, id_tmp, h_index);
76+
fdm_c.divmod(id_tmp, n_index, c_index);
77+
} else if constexpr (Layout == LAYOUT_NHWC) {
78+
fdm_c.divmod(id, id_tmp, c_index);
79+
fdm_d.divmod(id_tmp, id_tmp, d_index);
80+
fdm_w.divmod(id_tmp, id_tmp, w_index);
81+
fdm_h.divmod(id_tmp, n_index, h_index);
82+
}
83+
84+
// Window bounds mirror the CPU AveragePool{1,2,3}DTask reference exactly.
85+
int64_t h_start = h_index * stride_h - pad_h_head;
86+
int64_t w_start = w_index * stride_w - pad_w_head;
87+
int64_t d_start = d_index * stride_d - pad_d_head;
88+
89+
int64_t h_end = _Min<int64_t>(h_start + kernel_h * dilation_h, height + pad_h_tail);
90+
int64_t w_end = _Min<int64_t>(w_start + kernel_w * dilation_w, width + pad_w_tail);
91+
int64_t d_end = _Min<int64_t>(d_start + kernel_d * dilation_d, depth + pad_d_tail);
92+
93+
using AccT = typename AveragePoolAccumulator<T>::type;
94+
AccT acc = static_cast<AccT>(0);
95+
int64_t counted = 0;
96+
97+
int64_t offset = compute_offset(n_index, c_index, 0, 0, 0);
98+
const T* p_slice = p_input + offset;
99+
for (int64_t h = h_start; h < h_end; h += dilation_h) {
100+
if (h < 0 || h >= height) continue;
101+
for (int64_t w = w_start; w < w_end; w += dilation_w) {
102+
if (w < 0 || w >= width) continue;
103+
for (int64_t d = d_start; d < d_end; d += dilation_d) {
104+
if (d < 0 || d >= depth) continue;
105+
acc += static_cast<AccT>(p_slice[compute_offset(0, 0, h, w, d)]);
106+
++counted;
107+
}
108+
}
109+
}
110+
111+
AccT result = static_cast<AccT>(0);
112+
if (counted > 0) {
113+
if (count_include_pad) {
114+
int64_t divisor = (1 + (h_end - h_start - 1) / dilation_h) *
115+
(1 + (w_end - w_start - 1) / dilation_w) *
116+
(1 + (d_end - d_start - 1) / dilation_d);
117+
result = acc / static_cast<AccT>(divisor);
118+
} else {
119+
result = acc / static_cast<AccT>(counted);
120+
}
121+
}
122+
p_output[id] = static_cast<T>(result);
123+
}
124+
125+
template <typename T, bool Layout>
126+
void AveragePoolWithPad(
127+
cudaStream_t stream,
128+
const TensorShape& input_shape,
129+
const TensorShape& output_shape,
130+
const gsl::span<const int64_t>& kernel_shape,
131+
const gsl::span<const int64_t>& stride_shape,
132+
const gsl::span<const int64_t>& pads,
133+
const gsl::span<const int64_t>& dilations,
134+
bool count_include_pad,
135+
const T* p_input,
136+
T* p_output) {
137+
int64_t channels, height, width, depth;
138+
int64_t pooled_height, pooled_width, pooled_depth;
139+
if constexpr (Layout == LAYOUT_NCHW) {
140+
channels = input_shape[1];
141+
height = input_shape[2];
142+
width = kernel_shape.size() > 1 ? input_shape[3] : 1;
143+
depth = kernel_shape.size() > 2 ? input_shape[4] : 1;
144+
145+
pooled_height = output_shape[2];
146+
pooled_width = kernel_shape.size() > 1 ? output_shape[3] : 1;
147+
pooled_depth = kernel_shape.size() > 2 ? output_shape[4] : 1;
148+
} else if constexpr (Layout == LAYOUT_NHWC) {
149+
height = input_shape[1];
150+
width = kernel_shape.size() > 1 ? input_shape[2] : 1;
151+
depth = kernel_shape.size() > 2 ? input_shape[3] : 1;
152+
channels = input_shape[input_shape.NumDimensions() - 1];
153+
154+
pooled_height = output_shape[1];
155+
pooled_width = kernel_shape.size() > 1 ? output_shape[2] : 1;
156+
pooled_depth = kernel_shape.size() > 2 ? output_shape[3] : 1;
157+
}
158+
159+
const int64_t rank = static_cast<int64_t>(kernel_shape.size());
160+
int64_t kernel_h = kernel_shape[0];
161+
int64_t kernel_w = rank > 1 ? kernel_shape[1] : 1;
162+
int64_t kernel_d = rank > 2 ? kernel_shape[2] : 1;
163+
int64_t stride_h = stride_shape[0];
164+
int64_t stride_w = rank > 1 ? stride_shape[1] : 1;
165+
int64_t stride_d = rank > 2 ? stride_shape[2] : 1;
166+
167+
// pads: [x1_begin,...,xN_begin, x1_end,...,xN_end]. Begin at [i], end at [rank + i].
168+
int64_t pad_h_head = pads[0];
169+
int64_t pad_w_head = rank > 1 ? pads[1] : 0;
170+
int64_t pad_d_head = rank > 2 ? pads[2] : 0;
171+
int64_t pad_h_tail = pads[rank + 0];
172+
int64_t pad_w_tail = rank > 1 ? pads[rank + 1] : 0;
173+
int64_t pad_d_tail = rank > 2 ? pads[rank + 2] : 0;
174+
175+
int64_t dilation_h = dilations[0];
176+
int64_t dilation_w = rank > 1 ? dilations[1] : 1;
177+
int64_t dilation_d = rank > 2 ? dilations[2] : 1;
178+
179+
int64_t output_size = output_shape.Size();
180+
if (output_size == 0) return;
181+
182+
fast_divmod fdm_c(static_cast<int>(channels));
183+
fast_divmod fdm_h(static_cast<int>(pooled_height));
184+
fast_divmod fdm_w(static_cast<int>(pooled_width));
185+
fast_divmod fdm_d(static_cast<int>(pooled_depth));
186+
187+
int blocksPerGrid = (int)((output_size + GridDim::maxThreadsPerBlock - 1) / GridDim::maxThreadsPerBlock);
188+
AveragePoolWithPadKernel<T, Layout><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
189+
channels,
190+
height,
191+
width,
192+
depth,
193+
pooled_height,
194+
pooled_width,
195+
pooled_depth,
196+
kernel_h,
197+
kernel_w,
198+
kernel_d,
199+
stride_h,
200+
stride_w,
201+
stride_d,
202+
pad_h_head,
203+
pad_w_head,
204+
pad_d_head,
205+
pad_h_tail,
206+
pad_w_tail,
207+
pad_d_tail,
208+
dilation_h,
209+
dilation_w,
210+
dilation_d,
211+
fdm_c,
212+
fdm_h,
213+
fdm_w,
214+
fdm_d,
215+
count_include_pad,
216+
p_input,
217+
output_size,
218+
p_output);
219+
}
220+
221+
#define INSTANTIATE_AVERAGEPOOLWITHPAD(T, Layout) \
222+
template void AveragePoolWithPad<T, Layout>( \
223+
cudaStream_t stream, \
224+
const TensorShape& input_shape, \
225+
const TensorShape& output_shape, \
226+
const gsl::span<const int64_t>& kernel_shape, \
227+
const gsl::span<const int64_t>& stride_shape, \
228+
const gsl::span<const int64_t>& pads, \
229+
const gsl::span<const int64_t>& dilations, \
230+
bool count_include_pad, \
231+
const T* p_input, \
232+
T* p_output);
233+
234+
INSTANTIATE_AVERAGEPOOLWITHPAD(float, LAYOUT_NCHW)
235+
INSTANTIATE_AVERAGEPOOLWITHPAD(double, LAYOUT_NCHW)
236+
INSTANTIATE_AVERAGEPOOLWITHPAD(half, LAYOUT_NCHW)
237+
INSTANTIATE_AVERAGEPOOLWITHPAD(BFloat16, LAYOUT_NCHW)
238+
239+
#ifdef ENABLE_CUDA_NHWC_OPS
240+
INSTANTIATE_AVERAGEPOOLWITHPAD(float, LAYOUT_NHWC)
241+
INSTANTIATE_AVERAGEPOOLWITHPAD(double, LAYOUT_NHWC)
242+
INSTANTIATE_AVERAGEPOOLWITHPAD(half, LAYOUT_NHWC)
243+
INSTANTIATE_AVERAGEPOOLWITHPAD(BFloat16, LAYOUT_NHWC)
244+
#endif
245+
246+
} // namespace cuda
247+
} // namespace onnxruntime

‎onnxruntime/core/providers/cuda/nn/avg_pool_impl.h‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
#pragma once
5+
6+
#include "core/framework/tensor_shape.h"
7+
8+
namespace onnxruntime {
9+
namespace cuda {
10+
11+
// Custom average-pooling CUDA kernel that honors per-side (asymmetric) padding.
12+
//
13+
// cuDNN's pooling descriptor stores a single symmetric pad value per axis, so it cannot
14+
// represent ONNX asymmetric padding (pad_begin != pad_end), which arises from explicit
15+
// asymmetric `pads` or `auto_pad = SAME_UPPER/SAME_LOWER` resolving to asymmetric pads. This
16+
// kernel is the CUDA fallback for that case and mirrors the CPU reference functor
17+
// (AveragePool{1,2,3}DTask) exactly:
18+
// start = out_idx * stride - pad_begin
19+
// end = min(start + kernel * dilation, in_size + pad_end)
20+
// sum over cells [start, end) with dilation step that are in [0, in_size)
21+
// count_include_pad == 1: divisor = product of (1 + (end - start - 1) / dilation)
22+
// count_include_pad == 0: divisor = number of summed in-bounds cells
23+
//
24+
// `pads` is the full ONNX layout [x1_begin,...,xN_begin, x1_end,...,xN_end] (2 * rank).
25+
template <typename T, bool Layout>
26+
void AveragePoolWithPad(
27+
cudaStream_t stream,
28+
const TensorShape& input_shape,
29+
const TensorShape& output_shape,
30+
const gsl::span<const int64_t>& kernel_shape,
31+
const gsl::span<const int64_t>& stride_shape,
32+
const gsl::span<const int64_t>& pads,
33+
const gsl::span<const int64_t>& dilations,
34+
bool count_include_pad,
35+
const T* p_input,
36+
T* p_output);
37+
38+
} // namespace cuda
39+
} // namespace onnxruntime

‎onnxruntime/core/providers/cuda/nn/pool.cc‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#include "core/providers/cuda/nn/pool.h"
77
#include "core/providers/cuda/cudnn_common.h"
88
#include "core/providers/cuda/nn/max_pool_with_index.h"
9+
#include "core/providers/cuda/nn/avg_pool_impl.h"
910
#include "core/providers/cuda/math/unary_elementwise_ops_impl.h"
1011

1112
using namespace onnxruntime::common;
@@ -205,6 +206,37 @@ Status Pool<T, PoolType, Layout>::ComputeInternal(OpKernelContext* context) cons
205206
auto x_data = reinterpret_cast<const CudaT*>(X->Data<T>());
206207
auto y_data = reinterpret_cast<CudaT*>(Y->MutableData<T>());
207208

209+
// cuDNN's pooling descriptor cannot represent two ONNX features:
210+
// (1) Asymmetric padding: it stores a single symmetric pad value per axis and applies it to
211+
// both sides, so it silently drops ONNX end pads when pad_begin != pad_end (explicit
212+
// asymmetric pads, or auto_pad=SAME_UPPER/SAME_LOWER resolving to asymmetric pads).
213+
// (2) Dilation: the pooling descriptor has no dilation parameter at all, so any dilation > 1
214+
// is silently ignored.
215+
// Either case produces wrong sums and divisors on cuDNN, so route it to the custom kernel,
216+
// which honors per-side pads AND dilation and matches the CPU reference divisor exactly.
217+
// Symmetric, non-dilated pooling (the common case, including all global pooling) keeps the
218+
// fast cuDNN path unchanged, so there is zero perf regression. (MaxPool<8> guards dilation the
219+
// same way via !default_dilations.) Global pooling is always symmetric and never dilated, and
220+
// its kernel_shape/strides/dilations are left unpopulated (PoolAttributes returns early), so it
221+
// is excluded here and stays on cuDNN.
222+
if constexpr (PoolType::type == onnxruntime::PoolType::kAveragePool) {
223+
if (!pool_attrs_.global_pooling) {
224+
const size_t spatial_rank = kernel_shape.size();
225+
bool asymmetric_pads = false;
226+
for (size_t i = 0; i < spatial_rank; ++i) {
227+
if (pads[i] != pads[spatial_rank + i]) {
228+
asymmetric_pads = true;
229+
break;
230+
}
231+
}
232+
if (asymmetric_pads || !pool_attrs_.default_dilations) {
233+
AveragePoolWithPad<CudaT, Layout>(Stream(context), x_shape, y_shape, kernel_shape, strides, pads,
234+
pool_attrs_.dilations, pool_attrs_.count_include_pad, x_data, y_data);
235+
return Status::OK();
236+
}
237+
}
238+
}
239+
208240
TensorShapeVector x_dims_cudnn(x_dims.begin(), x_dims.end());
209241
TensorShapeVector y_dims_cudnn(y_dims);
210242
if (kernel_shape.size() < 2) {

0 commit comments

Comments
 (0)