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

[titaiwangms]

Summary

The CUDA MaxPool custom index kernel (onnxruntime/core/providers/cuda/nn/max_pool_with_index.cu) produces incorrect max values AND incorrect indices when dilation > 1 is combined with a non-zero begin (head) padding on that axis. The failure is silent — no error or crash, just wrong results.

• EP: CUDA only (CPU is correct).
• Outputs affected: both Y (pooled values) and Indices (which also corrupts any downstream MaxUnpool).
• Severity: correctness / silent-wrong-result — P1.

This was found during the AvgPool ceil_mode / asymmetric-pad work (PRs #29629 / #29631 / #29627) as an out-of-scope adjacent finding. It is a distinct bug from the AvgPool divisor issue — the AvgPool CUDA kernel does not share it (cross-checked, see below).

Reproduction (definitive)

MaxPool, opset 12, kernel_shape=[2], strides=[1], pads=[1,1], dilations=[2], input x = [1, 9, 2, 8, 3, 7, 4, 6, 5] (shape [1,1,9]):

CPU  Y  =[[[9. 2. 9. 3. 8. 4. 7. 5. 6.]]]   idx=[[[1 2 1 4 3 6 5 8 7]]]
CUDA Y  =[[[1. 2. 9. 3. 8. 4. 7. 5. 6.]]]   idx=[[[0 2 1 4 3 6 5 8 7]]]
        ^ output[0] wrong: 1 vs 9           ^ index 0 vs 1     max_abs_diff = 8

Reproduces with the single-output (common inference) path AND with the indices output; 1D and 2D.

Control case: same dilations=[2] but pads=[0,0] → matches CPU. This isolates the trigger to a negative window start produced by the begin padding.

Hand-check (CPU is correct per ONNX): for output[0], window start = 0*stride - pad_begin = -1. Dilated taps (start + m*dilation, m = 0,1) = {-1, 1}. Tap -1 is padding (skipped) → the only real tap is index 1 = x[1] = 9. CPU returns 9 / idx 1. CUDA returns 1 / idx 0 (it wrongly read x[0]).

Root cause (read-verified)

In max_pool_with_index.cu (lines ~70–89), the window start is computed, then clamped up to 0 without preserving the dilation phase:

int64_t h_start = h_index * stride_h - pad_h;                        // -1
int64_t h_end   = _Min(h_start + (kernel_h - 1) * dilation_h + 1, height);
...
h_start = _Max<int64_t>(h_start, 0);                                 // <-- BUG: -1 -> 0
...
T maxval = p_slice[compute_offset(0, 0, h_start, w_start, d_start)] - (T)1;  // seeds from x[0]
for (int64_t h = h_start; h < h_end; h += dilation_h) { ... }        // visits tap 0 (wrong phase)

_Max(start, 0) collapses a negative start to 0, so the loop starts at input index 0 instead of the first valid dilated tap (1 = -1 + dilation). The kernel then reads x[0] and reports index 0. The same clamp applies to w_start and d_start, so 2D/3D are affected on every axis that has begin-pad > 0 and dilation > 1.

Scope / impact

• Trigger: any axis with begin-pad > 0 AND dilation > 1 (symmetric or asymmetric padding — the repro above uses symmetric [1,1]).
• Path: the custom-kernel path, which MaxPool<8> selects when I != nullptr || !default_dilations (cuda/nn/pool.cc:346). Since dilation > 1 forces !default_dilations, dilated MaxPool on CUDA always lands here → the bug is unavoidable for dilated + padded MaxPool (both single-output and indices).
• Corrupts both pooled values and returned indices (breaks MaxUnpool downstream).
• CPU EP is correct. Cross-checked: the CUDA AvgPool custom kernel (AveragePoolWithPad, added by the recent AvgPool fix) does not share this bug — dilation + pad matches CPU.

Proposed fix (kernel fix, not a dispatch guard)

Replace each X_start = _Max<int64_t>(X_start, 0); (lines ~78–80) with an advance to the first non-negative dilated tap, preserving phase:

if (h_start < 0) {
  int64_t steps = (-h_start + dilation_h - 1) / dilation_h;  // ceil(-start / dilation)
  h_start += steps * dilation_h;                             // smallest start + m*dilation >= 0
}

(Analogously for w_start and d_start.) h_end is computed from the original start before the clamp, so it stays correct; the maxval seed then reads the first valid in-bounds tap. This matches the CPU reference's kernel-offset iteration.

Add a CUDA MaxPool dilation + pad parity test (MaxPool k=2 s=1 pads=[1,1] dilations=[2], 1D + 2D), CUDA leg un-excluded, with CPU as the oracle.

References

• Found during the AvgPool ceil_mode / asymmetric-pad fixes: PRs Fix CPU AvgPool ceil_mode + count_include_pad wrong average (opset 7-18 / MLAS path) #29629 (CPU), Fix CUDA AveragePool wrong results with asymmetric padding and dilation #29631 (CUDA), Fix JSEP pooling output shape to honor ceil_mode #29627 (JSEP).
• Related upstream context: ONNX export mismatch for avg_pool2d with ceil_mode=True and count_include_pad=True pytorch/pytorch#183528 (AvgPool ceil_mode divisor), which motivated the pooling audit that surfaced this adjacent MaxPool bug.
• Full QA evidence + repro scripts (qa_maxpool_cuda_probe.py, qa_maxpool_dilation_focus.py) captured during the maxpool-cudnn-asym probe.