Fix CPU AvgPool ceil_mode + count_include_pad wrong average (opset 7-18 / MLAS path)

[titaiwangms]

Summary

Fixes wrong AveragePool output on the CPU EP when ceil_mode=1 and count_include_pad=1 for opset 7–18 (the float MLAS path). This is the direct fix for the CPU-EP repro in pytorch/pytorch#183528.

Root cause

Pool<float, AveragePool>::Compute routes float AveragePool opset 7–18 to MLAS. With count_include_pad=1, MLAS divides every window by the full kernel size. Under ceil_mode=1 the output grid gains a trailing window whose extent runs past input + pad_tail (phantom cells). MLAS's full-kernel divisor counts those phantom cells, producing an average that is too small on the boundary windows.

The opset-19 reference functor (AveragePool{1,2,3}DTask in pool_functors.h) already handles this correctly: it clamps the window end to input + pad_tail and divides by 1 + (end - start - 1)/dilation, excluding the phantom cells.

Fix

Dispatch-fallback, zero MLAS edits:

1. Extract the opset-19 reference average-pool loop into a shared free function ComputeAveragePoolReference<T>(context, pool_attrs, tp) (reads strides from PoolAttributes, p=0), and make AveragePoolV19<T>::Compute a thin wrapper over it — one canonical loop, no copy-paste drift.
2. In Pool<float, AveragePool>::Compute, route only the buggy combo ceil_mode == 1 && count_include_pad && !global_pooling to the reference loop. Every other case (ceil_mode=0, count_include_pad=0/exclude-pad, global pooling, all MaxPool) keeps the fast MLAS path unchanged.

global_pooling is excluded because it has no ceil/pads (already correct) and leaves strides[] unpopulated; ComputeAveragePoolReference ORT_ENFORCEs that precondition.

Perf

Zero impact on the common path. The reference loop runs only for ceil_mode=1 && count_include_pad=1 (rare); the guard is a cheap early branch.

Relationship to #16752

ORT PR #16752 fixed this behavior for opset ≥ 19 only (via the new v19 reference functor). Opset 7–18 float still went through MLAS and remained wrong. This PR closes that gap by reusing the same already-correct functor for the 7–18 dispatch fallback.

Tests

Added to pool_op_test.cc (CPU-focused; GPU/other EPs excluded so CI stays green):

• AveragePool_18_ceil_count_include_pad_1d — opset-18 clone of the existing v19 test; the direct #183528 repro.
• AveragePool_18_ceil_count_include_pad_2d — the arange(1,17).reshape(1,1,4,4), k3/s2/pad1 case → [1.556, 3.333, 2.0, 6.333, 11.0, 6.0, 4.5, 7.5, 4.0].
• AveragePool_18_ceil_count_include_pad_3d — exercises the 3D functor path.
• AveragePool_18_ceil_count_exclude_pad_2d — count_include_pad=0 no-regression guard (stays on MLAS).

Expected values are derived from the CPU v19 reference (equivalently, hand-computed per the ONNX spec). Full PoolTest suite passes locally.

[Copilot]

Pull request overview

This PR fixes incorrect CPU EP AveragePool results for opset 7–18 when ceil_mode=1 and count_include_pad=1 on the float/MLAS path by selectively falling back to the already-correct reference loop used by opset ≥19, and adds CPU-focused regression tests for the scenario.

Changes:

• Extracts the opset-19 reference average-pool implementation into a shared ComputeAveragePoolReference<T> helper and makes AveragePoolV19<T>::Compute delegate to it.
• Routes only the problematic opset 7–18 float case (ceil_mode=1 && count_include_pad && !global_pooling) away from MLAS to the reference loop.
• Adds opset-18 regression tests for 1D/2D/3D ceil_mode=1 && count_include_pad=1, plus a no-regression case for count_include_pad=0.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.

File	Description
onnxruntime/core/providers/cpu/nn/pool.cc	Adds shared reference AvgPool loop and uses it as a targeted fallback for opset 7–18 float to avoid MLAS’s incorrect divisor under ceil-mode tail windows.
onnxruntime/test/providers/cpu/nn/pool_op_test.cc	Adds opset-18 CPU regression coverage for ceil+include-pad (1D/2D/3D) and a no-regression test for exclude-pad.

[titaiwangms]

Review summary

The core fix is correct and well-tested: the extracted ComputeAveragePoolReference reproduces the opset-19 reference divisor semantics exactly, and all four new expected_vals arrays were hand-derived from the functor algorithm and confirmed exact (including the smoking-gun cells — e.g. 2D (2,2)=4.0 where MLAS's full-kernel divisor would give 16/9≈1.78, and the 1D ph=3 window where the clamped hend correctly drops the divisor from 7 to 6). The refactor is behavior-preserving (the old AveragePoolV19::p_ was uninitialized but unused by the AvgPool functors, so p=0 is safe).

However, the fix is incomplete — two other CPU float paths reintroduce or bypass the same behavior.

🔴 Major

1. The NCHWc optimizer silently defeats this fix.
onnxruntime/core/optimizer/nchwc_transformer.cc TransformPool converts float AveragePool (opsets incl. 7/10/11/19/22) into an NchwcAveragePool node and copies node.GetAttributes() blindly — there is no ceil_mode guard. NchwcAveragePool::Compute (onnxruntime/contrib_ops/cpu/nchwc_ops.cc:261) then calls MlasNchwcPool with MlasAveragePoolingIncludePad, which has the same full-kernel-divisor bug this PR fixes. So for a 4D input whose channel count is a multiple of the NCHWc block size (typical after a Conv), a model this PR "fixes" will still compute the wrong average once Level-3 graph optimization runs.
Suggested fix: in TransformPool, bail out (return;) when ceil_mode == 1 && count_include_pad == 1, forcing execution back onto the now-corrected standard CPU EP path.

2. FP16 CPU path remains broken.
MLFloat16 AveragePool (opsets 11-22) is bound to PoolFp16 (onnxruntime/core/providers/cpu/fp16/fp16_pool.cc), which uses Im2col + a full kernel_size divisor and never routes through the new reference loop. This PR scopes itself to the float MLAS path, so this is arguably out of scope — but the underlying bug is not fully fixed for fp16.
Recommendation: either extend the fix to PoolFp16, or explicitly document that the fp16 CPU AveragePool ceil_mode + count_include_pad case remains a known gap.

3. The diverted path drops PoolBase::Compute's rank validation.
The new guard in Pool<float, AveragePool>::Compute returns into ComputeAveragePoolReference before the pooling_dims > 3 and pooling_dims == kernel_shape.size() checks (pool.cc:200-205). For a malformed rank-mismatched model (e.g. a 4D input with a 1D kernel_shape), SetOutputSize → InferOutputSize indexes strides[dim]/kernel_shape[dim]/dilations[dim] out of bounds, and the 4D-pooling case allocates the output before the late default: rejection. The old MLAS path rejected these cleanly with INVALID_ARGUMENT.
Suggested fix: mirror the two preflight checks at the top of ComputeAveragePoolReference:

const size_t pooling_dims = x_shape.NumDimensions() - 2;
ORT_RETURN_IF_NOT(pooling_dims <= 3, "Unsupported pooling size.");
ORT_RETURN_IF_NOT(pooling_dims == pool_attrs.kernel_shape.size(),
                  "kernel_shape num_dims is not compatible with X num_dims.");

(Lower real-world likelihood since ONNX shape inference may also reject rank mismatch, but it's a defensive regression worth restoring.)

🟡 Minor / Nit

• Overbroad guard (perf): ceil_mode=1 && count_include_pad=1 cases that produce no phantom tail cells (e.g. stride-1 same-padded pools like k3 s1 p1) were already correct on MLAS but are now forced onto the reference loop. Acknowledged as a tradeoff in the PR description; if it matters, gate the fallback on whether any dimension's last window actually extends past input + tail_pad.
• Dead state: AveragePoolV19::p_ (pool.h:39) is now unused after the refactor — remove it.
• Readability: the new stride_h/stride_w/stride_d locals in ComputeAveragePoolReference shadow the identically-named PoolBase::stride_h()/w()/d() accessors but omit their global_pooling ? 1 : ... guard (relying on the ORT_ENFORCE(!global_pooling) instead) — consider distinct names or a clarifying comment. The constexpr int64_t p = 0 reads as arbitrary at the call sites; a name like kUnusedLpNorm would be self-documenting. The test comment's "(a)-gate"/"(b) probe" shorthand has no referent in-file.
• ODR nit: ComputeAveragePoolReference<T> is a file-local template declared directly in namespace onnxruntime; wrapping it in an anonymous namespace avoids any ODR risk.

Cross-EP note

The tests correctly exclude CUDA/TRT/ACL/OpenVINO/DML — those EPs use vendor full-kernel divisors and now diverge from the corrected CPU oracle. Follow-up, not this PR's scope.

Reviewed by a multi-model review team (readability / correctness / adversarial / spec-adherence+numerical / cross-module integration). The four new expected-value arrays were independently re-derived from the functor algorithm and confirmed exact.

[titaiwangms]

Review feedback folded in — head now 97d99b30b

Thanks to the reviewers (Copilot + the team) — the NCHWc-bypass and fp16 catches were the important ones; both are addressed.

Folded in

• M1 — NCHWc optimizer bypass (Major, closed): TransformPool now bails out of the AveragePool → NchwcAveragePool conversion for ceil_mode == 1 && count_include_pad != 0, so optimized x86 graphs stay on the fixed CPU kernel instead of hitting MlasAveragePoolingIncludePad's full-kernel divisor. count_include_pad is gated as != 0 (not == 1) to match how PoolBase and the float/fp16 kernels enable include-pad — an out-of-spec count_include_pad = 2 can no longer escape the guard. Covered by NchwcOptimizerTests.ConvAveragePoolCeilCountIncludePadNotConverted (Conv still converts; the AveragePool stays ONNX).
• M3 — rank validation (OOB fix): mirrored PoolBase::Compute's two rank checks into ComputeAveragePoolReference so a rank-mismatched model can't reach SetOutputSize/InferOutputSize out of bounds.
• N1 — dead field: removed the unused, uninitialized AveragePoolV19::p_.
• Copilot Remove vsts test runner in cmake file #2–5 — test EP exclusions: added kCudaNHWCExecutionProvider to the four new float ceil tests, matching every sibling CPU-reference test.
• M2 — fp16 divisor fix (dtype consistency): the fp16 path (PoolFp16) had the same bug via MlasNhwcAvgPool. Added a self-contained ComputeAveragePoolFp16Reference (clamps the window end, float-accumulate, round to fp16) behind the same ceil_mode && count_include_pad && !global_pooling guard — zero MLAS/float-path edits. This was fixed in-scope per the request that fp16 and float not diverge for the same model.

fp16 merge gate (documented in the PR body): MLAS_F16VEC_INTRINSICS_SUPPORTED is ARM64-only, so PoolFp16 + pool_fp16_op_test.cc compile out on x86 and the fp16 parity tests must run green on the ARM64 / CoreML / XNNPACK CI leg before merge. This is a live-execution confirmation, not a correctness unknown — the oracle values were independently hand-verified and the reference is architecture-independent float math (x86 compile-verified via forced-macro -fsyntax-only).

Declined (with rationale)

• N2 — "overbroad guard" perf: the guard is a per-call early branch on already-loaded attributes; cost is negligible and only the rare buggy combo takes the reference path. No change.
• N4 — anonymous-namespace ODR nit: ComputeAveragePoolReference is a file-local template already reused only within its TU; no ODR exposure to warrant restructuring. No change.

Local validation: PoolTest 46 passed / 2 DML-skipped; NchwcOptimizerTests 36 passed (incl. the new bail-out test) after the fold. clang-format clean.

[titaiwangms]

CI test-scoping fold (WebGPU red leg + defensive EP completeness)

Pushed a test-only follow-up (head 9d42f48bd) that completes the EP scoping of the two
new opset-18 AveragePool ceil_mode + count_include_pad parity tests (and their 3D /
exclude-pad siblings).

• Root cause of the red CI: the WebGPU EP runs these ops but does not implement the
  ceil_mode + count_include_pad clamped-window divisor (PyTorch #183528), so it returned
  the pre-fix full-kernel-size average and the WebGPU legs failed.
• Fix: the shared exclusion set for all four tests is now factored into a single named
  constant kPoolingEpsExcludedFromCeilCountIncludePadTests (prevents per-test drift) and
  covers every EP in the test harness that lacks the clamped divisor — WebGPU (the actual
  red leg) plus DNNL / CoreML / QNN / TensorRT / NvTensorRT-RTX / ACL / OpenVINO / DML,
  added defensively so no later EP CI leg can re-red these CPU-reference gate tests. kCuda /
  kCudaNHWC remain excluded as CPU-reference gates (CUDA has its own parity coverage).
• Scope: test-only, zero kernel logic. clang-format clean; local PoolTest = 46 pass /
  2 DML-skip, no regressions.

Note: a deny-by-default "CPU-only" structuring of these gate tests is a possible future
cleanup, but the explicit shared-constant exclusion is the minimal, reviewable change here.

[tianleiwu]

Review: AvgPool ceil_mode + count_include_pad fix (opset 7-18 / MLAS path)

Verdict: COMMENT — the fix is correct and thoroughly tested. Just a couple of minor maintainability nits below.

Correctness (verified against the code):

• The root cause is accurate: for ceil_mode=1 && count_include_pad, the MLAS include-pad path divides by the full kernel size, counting the ceil-mode phantom tail cells. The reference functors (AveragePool{1,2,3}DTask) clamp hend/wend to input + tail_pad, so the include-pad divisor (1 + (hend - hstart - 1)/dilation) excludes phantom cells. I hand-verified the 1D test (k=7 s=3 pad=3 over {1,2,9}): out0 = 12/7, out1 = 12/6 = 2. ✔
• Scoping is tight and justified: the discrepancy only arises with ceil_mode=1 (without it, hend never exceeds input+tail_pad, so clamped == full divisor), and global_pooling is correctly excluded (no ceil/pads, and its strides[]/dilations[] aren't populated). ✔
• PoolAttributes always resizes strides/dilations to kernel_shape.size() for non-global pooling, so the opset 7-18 Pool<float, AveragePool> fallback safely indexes strides[0..2]/dilations[0..2]. ✔
• The removed dead p_ member of AveragePoolV19 was never initialized — safe removal. ✔

Good decisions:

• Single-source-of-truth extraction into ComputeAveragePoolReferenceCompute (float accumulation) shared by the float and fp16 paths keeps the two dtypes in lockstep and eliminates the duplicated N-D loop, addressing the earlier factorization request.
• The fp16 fallback correctly accumulates in float and rounds once at the float -> MLFloat16 store, and the NHWC<->NCHW transpose indexing is correct.
• Rank checks in ComputeAveragePoolReference mirror PoolBase::Compute, and the fp16 Compute already validates kernel_shape.size() == spatial_dims before reaching the guard, so neither path can read x_shape[3]/[4] out of bounds — the earlier validation concern is fully covered.
• Strong no-regression coverage: count_include_pad=0 control tests confirm the MLAS fast path is untouched, and the NCHWc transformer bail-out (with Level2-vs-Level3 output comparison) prevents the optimizer from silently reintroducing the buggy divisor.

Minor suggestions (non-blocking):

1. The two kPoolingEps... sets carry long cross-referencing comments with hardcoded approximate line numbers (~L21, ~L1150). These drift as the file changes. Prefer referencing the sets by symbol name only and dropping the ~Lxxx anchors (see inline note).
2. ComputeAveragePoolReference (the template wrapper in pool.cc) has external linkage but is only instantiated in this TU. Consider giving it internal linkage (anonymous namespace / static) to avoid unintended ODR exposure; only ComputeAveragePoolReferenceCompute needs to be exported via pool.h.
3. The fp16 fallback runs the cast/transpose loops single-threaded and allocates two full-size float scratch buffers. Acceptable for this rare edge-case combo, but worth a note if profiling ever shows it on a hot model.