[CUDA] Fix MaxPool wrong values/indices for dilation>1 with non-zero begin padding

[Copilot]

Description

The CUDA MaxPool custom index kernel (max_pool_with_index.cu) clamped a negative window start to 0 via _Max<int64_t>(start, 0), discarding the dilation phase. For any axis with begin-pad > 0 and dilation > 1, the loop began at input index 0 (padding) instead of the first valid dilated tap, silently producing wrong Y values and wrong Indices (also corrupting downstream MaxUnpool). CPU is unaffected.

Repro (k=2, s=1, pads=[1,1], dilations=[2], x=[1,9,2,8,3,7,4,6,5]): CUDA returned Y[0]=1/idx[0]=0 vs CPU 9/1.

Changes

• Kernel (onnxruntime/core/providers/cuda/nn/max_pool_with_index.cu): replace each X_start = _Max<int64_t>(X_start, 0) with a phase-preserving advance to the first non-negative dilated tap. *_end is computed from the original start, so it stays correct, and the maxval seed reads a valid in-bounds tap.

  // instead of: h_start = _Max<int64_t>(h_start, 0);
  if (h_start < 0) h_start += ((-h_start + dilation_h - 1) / dilation_h) * dilation_h;

  Applied to d_start, w_start, h_start, so all of 1D/2D/3D are covered.
• Tests (onnxruntime/test/providers/cpu/nn/pool_op_test.cc): un-excluded the CUDA legs from MaxPool_10_DilationPadding_1d/_2d (previously excluded due to this bug), and added MaxPool_DilationPadding_1d_Indices / _2d_Indices covering both value and index outputs with CPU as the oracle.

Motivation and Context

Dilated MaxPool on CUDA always lands on this custom-kernel path (MaxPool<8> selects it whenever I != nullptr || !default_dilations, and dilation > 1 forces !default_dilations), so the wrong-result was unavoidable for dilated + padded MaxPool. Found as an adjacent finding during the AvgPool ceil_mode / asymmetric-pad work; the CUDA AvgPool kernel does not share this bug.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[titaiwangms]

Review summary (multi-model review team)

Reviewed by a 5-reviewer fan-out (readability / code / critical / deep / integration). The core fix is mathematically correct and well-tested — the phase-preserving advance faithfully reproduces the CPU/ONNX oracle, and the new 1D/2D value+index vectors were independently re-derived and cross-checked against the ONNX reference evaluator (exact match). But there is one blocking issue the fix introduces.

🔴 Blocking — empty-window (no valid tap) now reads out of bounds

onnxruntime/core/providers/cuda/nn/max_pool_with_index.cu (the new advance + the maxval seed at p_slice[compute_offset(0,0,h_start,w_start,d_start)] - 1 + the final p_input[... compute_offset(0,0,h_index_max,...)]).

When a window has zero in-range taps, the phase advance pushes *_start past *_end, so the loop never runs, *_index_max stays -1, and:

• the maxval seed reads one element past the input, and
• the final output read uses compute_offset(..., -1, -1, -1) → negative-offset OOB read (and a negative emitted index).

The old _Max(start, 0) clamp kept start < end here, so the loop always ran and the read was in-bounds. So this corner is a memory-safety regression introduced by the PR — it trades a wrong-value bug for an OOB device read.

Reachable with a valid model (passes the only guard pads[dim] < kernel_shape[dim]):

X = [42], shape [1,1,1], kernel_shape=[2], strides=[1], pads=[1,1], dilations=[2]

Output length is 1; the single window's taps are {-1, 1}, both out of range for height 1 → adjusted h_start = 1 (== height). CPU and the ONNX reference define this case as Y = lowest(), Index = -1 (no OOB).

Suggested fix — guard the empty window and make the seed safe, mirroring CPU:

// initialize instead of reading p_slice[start]-1
T maxval = std::numeric_limits<T>::lowest();
...
if (h_index_max < 0 /* && w_/d_ */) {           // no valid tap
  p_output[id] = std::numeric_limits<T>::lowest();
  if (p_indices) p_indices[...] = -1;
} else { ...existing output/index writes... }

This also fixes a fragile pre-existing sentinel (maxval = start - 1 wraps for int8/half type-min).

A compute-sanitizer --tool memcheck run on the repro above will show invalid global read on the current diff; the OOB is otherwise provable directly from the offset arithmetic.

🟡 Minor

• 3D coverage gap: the fix touches d_start/w_start/h_start (all of 1D/2D/3D), but MaxPool_10_DilationPadding_3d still excludes CUDA. Consider un-excluding it (ideally with indices) so the depth branch can't regress independently.
• Missing edge test: add a tiny empty-window regression case (e.g. L=1, k=2, dil=2, pad=1) once the guard above is in.
• Readability: the ceil-division idiom is repeated 3× — a small local lambda (advance_to_first_nonneg_tap(start, dilation)) would name the operation and give the good "why" comment a single home. Also d_end/w_end/h_end are intentionally computed from the raw start before the in-place advance; a one-line note there would protect the subtle ordering from a future refactor. New tests also drop the opset-number prefix used by neighbors (MaxPool_10_...).

✅ Confirmed correct

• Phase-advance math correct across all residue classes; dilation=1 degenerates to the old max(·,0).
• Fix covers both NCHW and NHWC (CudaNHWC) legs (shared window-bounds code).
• MaxUnpool index contract preserved (flat NCHW offsets unchanged; only the chosen tap is corrected).
• CUDA AvgPool and CPU don't share the bug (they iterate from the raw negative start and skip per-tap). No in-tree ROCm sibling (ROCm hipifies CUDA source, so the fix — and the OOB — propagate there too; port the follow-up guard when it lands).

Note: findings #2 (maxval - 1 sentinel fragility for float -FLT_MAX/-inf and int8/half type-min) is largely pre-existing; the suggested lowest() init resolves it as a side effect.

[Copilot]

Pull request overview

This PR fixes a CUDA-only correctness issue in MaxPool when dilation > 1 and there is non-zero begin padding, where the CUDA custom MaxPool-with-indices kernel could start at the wrong input tap (wrong dilation phase), producing incorrect pooled values and indices (affecting downstream MaxUnpool).

Changes:

• Updates the CUDA MaxPoolWithIndexKernel to advance negative window starts to the first non-negative phase-correct dilated tap (instead of clamping to 0).
• Re-enables CUDA legs for existing dilation+padding MaxPool tests that were previously excluded due to this bug.
• Adds new regression tests that validate both Y and Indices outputs for 1D and 2D cases, using CPU as the oracle.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/core/providers/cuda/nn/max_pool_with_index.cu	Fixes CUDA MaxPool window-start handling for dilated pooling with begin padding to preserve dilation phase.
onnxruntime/test/providers/cpu/nn/pool_op_test.cc	Re-enables CUDA coverage for prior dilation+padding tests and adds new value+index regression tests.

[titaiwangms]

@copilot if reviews and comments make sense to you, can you fix them.

[Copilot]

@copilot if reviews and comments make sense to you, can you fix them.

Addressed in ba67e47.

Blocking (empty-window OOB read): the phase advance can push a window's start past its end when every dilated tap falls in the begin padding. I now seed maxval with NumericLimits<T>::Lowest() (removing the p_slice[start] - 1 read that could go one past the input and that wrapped for integral/half type-min) and guard the no-valid-tap case: it emits Y = lowest() and Index = -1 without reading p_input at the sentinel indices, mirroring the CPU/ONNX reference. This covers both NCHW and NHWC.

Minor:

• Un-excluded CUDA/CudaNHWC from MaxPool_10_DilationPadding_3d for depth-branch coverage.
• Added MaxPool_DilationPadding_1d_EmptyWindow (L=1, k=2, dil=2, pad=1) checking value + index; verified Index = -1 against the ONNX reference evaluator.
• Factored the ceil-division advance into a named local lambda (advance_to_first_nonneg_tap) and added a note that *_end is intentionally computed from the raw start before the in-place advance.

[titaiwangms]

Re-review: update resolves the blocker ✅

The revision addresses every finding from the prior review. Verified against the updated diff:

🔴 Blocking OOB — resolved. The maxval seed no longer reads p_slice[start] (now NumericLimits<T>::Lowest(), which has half/BFloat16 specializations in cuda_utils.h), and the new if (h_index_max < 0) guard returns early with Lowest() value + -1 index before the final compute_offset(...,-1,...) read — so the empty-window path no longer dereferences an out-of-range tap. This matches the CPU/ONNX reference exactly. Checking h_index_max alone is sufficient: the loop is a nested product, so if it never executes, none of the three *_index_max are set.

Bonus fix: dropping the old static_cast<float>(p_slice[pool_offset]) in the max update also removes a latent precision-truncation for double/half element types.

Minors — all addressed: empty-window regression test (MaxPool_DilationPadding_1d_EmptyWindow), 3D CUDA leg un-excluded, ceil-division extracted into the named advance_to_first_nonneg_tap lambda with a clarifying comment, the raw-start ordering note on *_end, and #include <limits>.

No remaining blocking or major concerns from the review team. LGTM pending CI.

[tianleiwu]

The core fix looks correct and is well-motivated. The phase-preserving advance start += ((-start + dilation - 1) / dilation) * dilation correctly snaps a negative window start up to the first non-negative dilated tap while preserving phase, and computing *_end from the raw start keeps the loop bound right. Replacing the p_slice[start] - (T)1 seed read with NumericLimits<T>::Lowest() plus the explicit empty-window early-return also closes the previous out-of-range dereference and the uint8 - 1 wrap, and dropping static_cast<float> fixes latent precision truncation for the double instantiation. Good test coverage too.

One residual correctness issue remains (inline): with the Lowest() seed and the strict > comparison, a valid window whose maximum equals the element type's minimum (uint8_t 0, int8_t -128) is misclassified as empty and emits Indices = -1, which corrupts downstream MaxUnpool. Since the kernel is instantiated for int8_t/uint8_t and all-zero uint8 windows are common in quantized data, this path is reachable. See the inline comment for the suggested one-line fix. A minor follow-up: the new tests are all float; an int8_t/uint8_t Indices case would lock this in.

[titaiwangms]

cc @tianleiwu

[tianleiwu]

The phase-preserving window-start adjustment is correct, and the follow-up changes safely handle empty windows and valid windows whose maximum equals the element type's lowest value. CPU and CUDA behavior are aligned for first-tap selection, the regression coverage exercises values and indices across the relevant edge cases, and all current checks pass.