Fix JSEP pooling output shape to honor ceil_mode

[titaiwangms]

Summary

The ORT-Web JSEP pooling output-shape helper (PoolConvUtil.computePoolOutputShape → computeShapeHelper → adjustPadAndReturnShape) was floor-only: it had no ceilMode parameter and silently ignored ceil_mode. Under ceil_mode=1 this allocated the output tensor one element too small along each affected spatial axis, producing a wrong output shape (and downstream shape mismatches) — independent of any pooling divisor concern.

The same floor-only helper is duplicated in the legacy js/web/lib/onnxjs/util.ts; both copies are fixed.

Fix

• Thread a ceilMode parameter (default 0, i.e. floor) through computePoolOutputShape → computeShapeHelper → adjustPadAndReturnShape in both js/web/lib/wasm/jsep/util.ts and js/web/lib/onnxjs/util.ts.
• Route the NOTSET / VALID / SAME_UPPER / SAME_LOWER branches through a new computeOutputSize() helper that produces results identical to the C++ reference PoolAttributes::ComputeOutputSize (onnxruntime/core/providers/cpu/nn/pool_attributes.h), including the ceil_mode "shrink the last window if it starts entirely in the trailing padding" rule (ref: Fix corner case where output size need to reduce by one in MaxPool onnx/onnx#5741).
• The floor path (ceilMode default) is algebraically unchanged, so Conv and auto_pad pad-adjustment are unaffected.
• js/web/lib/wasm/jsep/webgpu/ops/pool.ts now passes attributes.ceilMode into the shape computation.

Scope: SHAPE-only

This PR fixes the output-shape computation only. End-to-end ceil_mode execution remains gated by the existing throw guards in parseAveragePoolAttributes / parseMaxPoolAttributes, because the WebGPU pooling kernel does not yet implement ceil_mode trailing-padding handling (and, for AveragePool, the count_include_pad divisor). Removing those throws + adding kernel support is a tracked follow-up; the now-correct shape path is exercised directly by the added unit tests until then. Comments at the throw sites and in the test header document this.

Tests

Adds js/web/test/unittests/pool-output-shape.ts (registered in unittests/index.ts), asserting output shapes for both implementations (jsep + onnxjs) against the C++ CPU reference ground truth:

• test_maxpool_2d_ceil → [1,1,2,2], AveragePool_10_ceil1_2d → [1,1,2,3]
• AvgPool ceil 1D [1,2,4] / 2D [1,1,3,3] / 3D [1,1,2,2,2] (matching the CPU pool_op_test.cc cases)
• ceil_mode with dilation, VALID / SAME_UPPER / SAME_LOWER auto_pad
• a shrink-rule discriminator (naive ceil() would give 3; correct = 2)
• floor-mode no-regression case

Related

• PyTorch context: ONNX export mismatch for avg_pool2d with ceil_mode=True and count_include_pad=True pytorch/pytorch#183528
• Related ORT PR: Align AvgPool ceil_mode on last value to torch #16752

Post-review fixups

• SAME_UPPER/SAME_LOWER integer division (external-review Major). The legacyTargetSize = (inSize + stride - 1) / stride computation now uses Math.floor(...) to match C++ pool_attributes.h ComputeSizePadDilations, which uses integer division. This fixes a latent float-division divergence that mis-rounded the auto_pad pad distribution. The correction applies to all ceil modes (not only ceil_mode=1), aligning JSEP/onnxjs with the CPU/CUDA EPs. The floor-path (ceil_mode=0) output shape is unchanged — only the SAME_* pad split is corrected — so nothing previously-correct regresses. Added ceil_mode=0 SAME_UPPER/SAME_LOWER non-divisible regression tests that assert the corrected pad out-param (SAME_UPPER → [0,1], SAME_LOWER → [1,0]).
• Throw-message / TODO clarity. The retained ceil_mode throw statements and the pool op TODO banner were reworded to make explicit that the output shape is now computed correctly, while ceil_mode kernel execution (padding/divisor) remains the pending WebGPU follow-up.

[Copilot]

Pull request overview

This PR fixes ORT-Web pooling output-shape computation to correctly honor ceil_mode (including the ONNX “shrink the last window if it starts entirely in trailing padding” rule) in both the JSEP and legacy onnxjs copies of PoolConvUtil.computePoolOutputShape. It also adds unit tests that exercise the pure shape helper directly (since end-to-end WebGPU pooling still throws on ceil_mode != 0 pending kernel support).

Changes:

• Thread ceilMode through computePoolOutputShape → computeShapeHelper → adjustPadAndReturnShape and centralize per-dimension sizing in a computeOutputSize() helper.
• Pass attributes.ceilMode into the WebGPU pooling output-shape computation.
• Add unit tests covering ceil-mode shapes (including the shrink rule) for both implementations (jsep + onnxjs).

Reviewed changes

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

Show a summary per file

File	Description
js/web/lib/wasm/jsep/util.ts	Adds ceilMode plumbing and computeOutputSize() for pooling/conv output shape calculation.
js/web/lib/onnxjs/util.ts	Mirrors the same ceilMode and computeOutputSize() logic for the legacy onnxjs copy.
js/web/lib/wasm/jsep/webgpu/ops/pool.ts	Passes ceilMode into shape computation and expands comments around current ceil_mode execution guards.
js/web/test/unittests/pool-output-shape.ts	New unit tests validating ceil-mode output shapes (and shrink rule) against CPU-reference expectations.
js/web/test/unittests/index.ts	Registers the new unit test module.

[titaiwangms]

Review summary

Solid, well-tested, and correctly scoped (shape-only, execution still gated). The ceil-mode math matches the C++ reference (PoolAttributes::ComputeOutputSize), all test expectations check out (including the shrink-rule boundary tested exactly at equality), the two computeOutputSize copies are byte-identical, and all other callers (Conv, WebGL, GlobalPool) safely keep the floor default. One Major to address before the ceil path can be fully trusted; the rest are Minor/Nit.

🔴 Major (1)

SAME_UPPER/SAME_LOWER uses float division for legacyTargetSize, diverging from the C++ reference in ceil mode.
js/web/lib/wasm/jsep/util.ts:523 and js/web/lib/onnxjs/util.ts:1420

const legacyTargetSize = (inSize + stride - 1) / stride;   // float; C++ uses int64 division

Reproduced with inSize=2, stride=2, kernel=3, SAME_UPPER, ceilMode=1:

• JS: legacyTargetSize = 1.5 → padNeeded = 2 → output […,2]
• C++ (legacy_target_size = 1) → pad_needed = 1 → output […,1]

Floor mode masks this (the extra pad is absorbed by floor), so it's a pre-existing latent bug that the new ceil path newly exposes — and it breaks the "identical to the C++ reference" claim for the SAME_* branch. The added SAME_* tests only use in=5, stride=2, where the division is integral and hides the issue.

Fix: use integer division to match C++, in both files:

const legacyTargetSize = Math.floor((inSize + stride - 1) / stride);

and add a non-divisible regression case, e.g. inputDims=[1,1,2], strides=[2], kernelShape=[3], autoPad='SAME_UPPER', ceilMode=1 → expected [1,1,1].

🟡 Minor (2)

1. float32 vs float64 ceil. C++ does std::ceil(static_cast<float>(numerator) / stride) (float32); JS uses double Math.ceil (jsep/util.ts:488, onnxjs/util.ts:1385). These diverge only for numerator > 2^24, which is unreachable for real pooling spatial dims — so no practical impact. Either emulate float32 with Math.fround(...) for strict parity, or soften the "identical to the C++ reference" comment to note JS intentionally uses exact double math. Non-blocking.
2. Stale banner TODO at the top of pool.ts still lists ceil_mode as an unqualified TODO, which now contradicts the precise per-function comments in parseAveragePoolAttributes/parseMaxPoolAttributes. Suggest updating it to note the output-shape math is supported and only kernel execution is gated.

⚪ Nits

• The "Keep in sync" comment on the duplicated computeOutputSize isn't enforced by anything but hand-picked case parity. Since the test already imports both PoolConvUtils, consider an explicit "both implementations agree" test over a matrix of random inputs — turns the promise into CI coverage.
• computeOutputSize's numerator param is semantically indirect; and the pool-output-shape.ts "Reviewer-verified" comment on the SAME_UPPER case would be more traceable as the actual hand-derivation.

Follow-up (out of scope, already tracked)

End-to-end ceil_mode execution remains correctly gated by the throws in parseAveragePool/MaxPoolAttributes pending the kernel trailing-padding (and AveragePool count_include_pad) support — the tracked follow-up.

Reviewed by a multi-model review team (readability / correctness / adversarial / spec-adherence / cross-module integration).

[titaiwangms]

Review fold-in (post-review fixups pushed: 134ab29f8, 1ea30f903)

Thanks all for the thorough passes — and a special thanks to the external review team for catching the SAME_* integer-division issue.

Folded in ✅

• SAME_UPPER/SAME_LOWER integer-division Major (external team). legacyTargetSize now uses Math.floor((inSize + stride - 1) / stride) to match C++ pool_attributes.h ComputeSizePadDilations (integer division). This was a latent float-division divergence that mis-rounded the auto_pad pad split. It's live on all ceil modes — on the ceil_mode=0 floor path the output shape is unchanged but the pad distribution was wrong (e.g. in=2, stride=2, kernel=3, SAME_UPPER gave pads [1,1] instead of C++ [0,1]). Added ceil_mode=0 SAME_UPPER/SAME_LOWER non-divisible pad-distribution regression tests (assert the corrected out-param, not just the shape); they fail against the pre-fix float code and pass with the fix.
• Throw-message / TODO clarity (Copilot C3/C4 + team Minor 2). Reworded the retained ceil_mode throw messages and the pool op TODO banner to make clear the output shape is computed, while ceil_mode kernel execution (padding/divisor) is the pending WebGPU follow-up.

Respectfully declined (with rationale) 🙏

• Copilot C1/C2 — Math.floor → Math.trunc on the floor path. These differ only for a negative numerator with a remainder, i.e. kernel > inSize + pads producing an invalid <= 0 output dim — unreachable for valid models. It's also the ceil_mode=0 path this PR intentionally scopes out. Happy to file separately if desired.
• Team Minor 1 — float32 vs double for the ceil computation. Diverges only for numerator > 2^24, impossible for real spatial dims; non-blocking, as the team itself noted.
• Team Nits A/B — random-matrix cross-impl test + numerator rename. Optional polish; left out to keep the fixup tight and reviewable.