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.