Fix CUDA AveragePool wrong results with asymmetric padding and dilation

[titaiwangms]

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): Fix CPU AvgPool ceil_mode + count_include_pad wrong average (opset 7-18 / MLAS path) #29629 — opset 7-18 MLAS ceil_mode + count_include_pad divisor fix.
• JSEP (c): Fix JSEP pooling output shape to honor ceil_mode #29627 — JSEP pooling output shape honoring ceil_mode.
• ONNX export mismatch for avg_pool2d with ceil_mode=True and count_include_pad=True pytorch/pytorch#183528.

[Copilot]

Pull request overview

This PR fixes CUDA EP AveragePool correctness when cuDNN cannot faithfully represent the ONNX pooling request—specifically for asymmetric per-side padding and dilation > 1—by adding a CUDA fallback kernel and routing only the unsupported cases to it, while keeping the cuDNN fast path for the common symmetric/non-dilated cases.

Changes:

• Add a custom CUDA kernel (AveragePoolWithPad) that implements ONNX per-side pads, dilation, and count_include_pad divisor behavior consistent with the CPU reference functors.
• Add a targeted dispatch guard in the CUDA pooling kernel to use the custom implementation only for non-global AveragePool with asymmetric pads or non-default dilations.
• Add CUDA parity tests (kept CUDA un-excluded) to validate asymmetric padding, SAME_* auto_pad asymmetry, dilation>1, and fp16 behavior.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File	Description
onnxruntime/test/providers/cpu/nn/pool_op_test.cc	Adds CUDA parity tests covering asymmetric pads, SAME_* auto_pad, dilation>1, fp16, and a MaxPool unaffected probe.
onnxruntime/core/providers/cuda/nn/pool.cc	Routes AveragePool to a custom CUDA kernel when pads are asymmetric or dilation is non-default (excluding global pooling).
onnxruntime/core/providers/cuda/nn/avg_pool_impl.h	Declares the new AveragePoolWithPad CUDA fallback API.
onnxruntime/core/providers/cuda/nn/avg_pool_impl.cu	Implements the CUDA kernel and instantiations for fp32/fp64/fp16/bf16 (NHWC guarded).

[titaiwangms]

Review summary

Strong PR. The custom kernel is a faithful, verified port of the CPU v19 reference (AveragePool{1,2,3}DTask): window bounds, the min(start + kernel*dilation, in_size + pad_tail) clamp, the per-axis ∏(1 + (end-start-1)/dilation) include-pad divisor, the exclude-pad counted divisor, and the counted==0 → 0 behavior all line up 1:1 with the oracle, including integer-division truncation. All nine new expected_vals arrays were independently re-derived and confirmed exact, and the NCHW/NHWC index math + fp16/bf16-accumulate-in-float were verified. The dispatch guard is evaluated on the post-SetOutputSize pads, so auto_pad=SAME_UPPER/LOWER (naturally asymmetric) routes correctly while symmetric/VALID/global stay on cuDNN. No Critical or Major correctness bugs found.

Verified a potential concern and cleared it: symmetric explicit pad + ceil_mode=1 + count_include_pad=1 is intentionally left on cuDNN and is correct there — the pre-existing, CUDA-enabled AveragePool_19_ceil_count_include_pad_1d (pool_op_test.cc:1076, identical config) asserts the v19-clamped divisor (-0.09990001, divide-by-6 not 7); if cuDNN used a full-kernel divisor it would already be red on CUDA CI. So cuDNN's boundary divisor already equals v19 for symmetric ceil overhang, and the narrow guard is provably sufficient.

🟠 Major — test coverage

New CUDA tests cover only NCHW 1D/2D. The kernel's highest-risk logic is the layout/rank-dependent indexing, but the NHWC decode branch and the 3D traversal branch are never exercised on CUDA. Recommend adding:

• one 3D asymmetric-tail-pad (and/or dilation) AveragePool CUDA case, and
• one NHWC asymmetric/dilated case (internal NHWC domain path).

Also, the pre-existing AveragePool_CountIncludePad_AsymmetricPads and AveragePool3D_CountIncludePad_AsymmetricPads still exclude kCuda*ExecutionProvider — since this PR fixes exactly that path, those exclusions can now be removed to lock in the fix on the canonical cases.

🟡 Minor

1. Guard is overbroad for count_include_pad=0. For exclude-pad, the tail pad affects neither the sum nor the divisor (only in-bounds cells count), and cuDNN represents pad_begin — so asymmetric exclude-pad is already correct on cuDNN (confirmed by this PR's own asymmetric_tail_pad_1d_exclude test, where cuDNN and the kernel agree on [4, 6.5, 8]). Routing it to the one-thread-per-output kernel is an avoidable perf regression. Consider gating asymmetric routing on count_include_pad while still always routing non-default dilations.
2. Imprecise comment. The pool.cc/header comment lists "ceil_mode boundaries" as a cause of asymmetric padding that needs the custom kernel. For explicit (NOTSET) pads this is inaccurate — ComputeSizePadDilations never makes explicit pads asymmetric, and cuDNN is already v19-correct for symmetric ceil overhang (see above). ceil_mode only forces the custom path via auto_pad=SAME_*. Suggest rewording to attribute the trigger to SAME_* and dropping ceil_mode as an independent cause.
3. 32-bit kernel indexing. int id = blockIdx.x*blockDim.x+threadIdx.x with int64_t output_size, and fast_divmod(static_cast<int>(...)), silently overflow for output_size > INT_MAX. This is a latent issue but it matches the existing max_pool_with_index.cu:45 convention exactly, so it's a pre-existing repo-wide limitation rather than a regression in this PR — noting for completeness; a repo-wide CUDA_LONG migration would be a separate change.
4. Regression test doesn't assert the route. symmetric_pad_regression_1d checks only numerics; if a future change wrongly routed it to the custom kernel the test would still pass. Reword the comment or add a route assertion if route coverage matters.
5. NHWC over-instantiation. avg_pool_impl.cu instantiates double/BFloat16 for LAYOUT_NHWC, but AveragePool NHWC (kMSInternalNHWCDomain) is only registered for float/MLFloat16. Harmless binary bloat; trim or leave as-is.

⚪ Nit

• The compute_offset lambda is duplicated verbatim between avg_pool_impl.cu and max_pool_with_index.cu; consider hoisting a shared pooling_compute_offset<Layout>() into a common .cuh.
• Test-comment jargon ("Target (b)", "QA probe", "per design") references an external task list with no in-repo referent — state the purpose directly.

Notes / follow-ups (out of scope)

• ROCm: not applicable — the ROCm provider is no longer present in main, so there is no hipify pipeline to break for the new files. (CUDA build is fine: onnxruntime_providers_cuda.cmake globs core/providers/cuda/*.cu with CONFIGURE_DEPENDS, so avg_pool_impl.cu is picked up automatically.)
• Other EPs: TensorRT/ACL/DML/OpenVINO likely share the same cuDNN-style descriptor deficiency (dropped tail pads / ignored dilation) and are excluded from these tests; a follow-up issue to audit them would be worthwhile.

Reviewed by a multi-model review team (readability / correctness / adversarial / spec-adherence+numerical / cross-module integration). All nine new expected-value arrays were independently re-derived from the CPU v19 reference; the symmetric-ceil coverage question was resolved against a pre-existing CUDA-enabled test.

[titaiwangms]

Review fold-in (post-review fixups pushed: ed114d5a8, 258bc05b3)

Thanks all for the thorough review — the correctness verification of the kernel port and the doc-accuracy catches were both valuable.

Folded in ✅

• Test coverage (Major). Un-excluded the CUDA (NCHW) leg on the two pre-existing asymmetric-pad parity tests AveragePool_CountIncludePad_AsymmetricPads and AveragePool3D_CountIncludePad_AsymmetricPads — this PR fixes exactly that path. The existing expected values are CPU-reference and the custom kernel mirrors CPU, so no re-derivation was needed. This closes the missing 3D-CUDA coverage gap for free. Verified green on an A100 CUDA Debug build (full PoolTest: 53 passed, 2 DML-skipped; the two un-excluded tests now run on the CUDA EP and pass).
• Comment accuracy — dispatch guard (Min2). Reworded the cuda/nn/pool.cc guard comment: dropped ceil_mode boundaries as a cause requiring the custom kernel. Symmetric ceil_mode is v19-correct on cuDNN; the real asymmetric causes are explicit asymmetric pads and SAME_UPPER/SAME_LOWER.
• Comment accuracy — sibling header drift (Min). Applied the same correction to cuda/nn/avg_pool_impl.h, which carried the identical inaccurate wording, so the two files now agree.
• NHWC-exclusion comment reason (Critical Min M1). Corrected the comment justifying the kCudaNHWCExecutionProvider exclusion on the un-excluded tests: NHWC is excluded because the custom kernel's NHWC decode branch is not yet covered by these tests (asymmetric NHWC routes to that branch, not to cuDNN) — not because it falls back to cuDNN. Exclusion itself unchanged.
• De-jargon (Nit2). Reworded internal-jargon test comments (Target (b), QA probe) to describe what the tests verify for external readers.

Respectfully declined (with rationale) 🙏

• Exclude-pad guard-perf micro-opt (Min1). The current form is correct and conservative with no correctness risk; not worth the churn here.
• int32 → CUDA_LONG thread-id widening (Min3). Matches the existing max_pool_with_index convention; belongs in a separate CUDA_LONG-wide PR, not this one.
• compute_offset hoist (Nit1). Cross-file refactor, out of scope for this fix.

Follow-up noted 📌

• NHWC-CUDA asymmetric-pad test coverage (exercise the custom kernel's NHWC decode branch).