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).

[titaiwangms]

CI follow-up (test-only, head 725fb3c77)

Thanks again for the reviews. Pushed two test-only commits (pool_op_test.cc) that fix the red non-CUDA CI legs — no kernel logic changed; every CUDA leg was already green.

Root cause: the new AveragePool_CUDA_* parity tests were running on execution providers that legitimately don't implement ONNX's asymmetric-pad / dilated / ceil_mode + count_include_pad clamped-divisor pooling, so they produced wrong values (or, for CoreML, rejected asymmetric + ceil_mode).

Fixes:

• EP-scoping corrected. The fp16 case now runs CUDA-only (CPU has no fp16 AveragePool kernel on x64, and the Arm64 NEON fp16 kernel mishandles the divisor, so CPU can't be its oracle). The float parity tests exclude the non-implementing EPs via a single shared constant kPoolingEpsWithoutCeilCip (eliminates per-test drift) — which also lists WebGPU defensively (it auto-skips in a CUDA build, but naming it prevents a future WebGPU build leg from re-triggering the same failure).
• NHWC-CUDA asymmetric is now confirmed covered + green. The parity tests keep kCudaNHWCExecutionProvider un-excluded (USE_CUDA_NHWC_OPS defaults ON), so they run on the NHWC EP and pass. Updated the earlier "NHWC untested / follow-up" notes on the mirror asymmetric tests accordingly (3D/NDHWC NHWC remains the only follow-up).

Not our regression: the Windows GPU DML red is infra-flaky — a DXGI_ERROR_DEVICE_REMOVED (device-suspended) cascade that took out unrelated Einsum/Conv/Softmax tests; DML runs none of these pool tests (it's excluded). It just needs a CI re-run.

Verified locally on A100 (CUDA Debug, ENABLE_CUDA_NHWC_OPS on): PoolTest = 53 pass / 2 DML-skip; the 9 parity tests run on CPU + CUDA + CUDA-NHWC and pass; clang-format clean.