onnxruntime
f79df447 - Fix CUDA AveragePool wrong results with asymmetric padding and dilation (#29631)

Commit
8 days ago
Fix CUDA AveragePool wrong results with asymmetric padding and dilation (#29631) ### Summary Fixes wrong `AveragePool` output on the **CUDA** EP whenever cuDNN's pooling descriptor cannot represent the requested pooling — i.e. **asymmetric padding** or **dilation > 1**. This is the CUDA-EP counterpart of the CPU fix in #29629 and, together with it, the JSEP shape fix in #29627, closes out the CUDA leg of pytorch/pytorch#183528. ### Root cause `CudnnPoolingDescriptor::Set` copies only the **begin** pads (`pads[0..rank)`) into the cuDNN pooling descriptor, which stores a *single symmetric pad value per axis* and applies it to both sides. The ONNX **end** pads (`pads[rank..2*rank)`) are silently dropped. As a result, **all** asymmetric-pad AveragePool on CUDA is wrong: - explicit asymmetric `pads`, - `auto_pad = SAME_UPPER` / `SAME_LOWER` (which produce naturally asymmetric pads), - `ceil_mode = 1` boundary windows. Separately, the cuDNN pooling descriptor has **no dilation parameter at all**, so any `dilations > 1` is also silently ignored — even with symmetric pads. Example divergence from the CPU reference (QA probe): 1D `pad(0,3)`, k7, s3, `ceil_mode=1`, `count_include_pad=1` gave CUDA `[4, 6.5, 8]` vs. correct `[4, 5.571, 4]`; 2D `pad(0,0,3,3)` gave `71.0` vs. correct `17.75`. ### Fix Add a custom CUDA average-pool kernel (`avg_pool_impl.cu` / `.h`, modeled on `max_pool_with_index.cu`) that honors **per-side pads** and **dilation**, and computes the `count_include_pad` divisor exactly like the CPU v19 reference functor (`AveragePool{1,2,3}DTask`): - include-pad divisor clamps the window end to `input + pad_tail` (drops the ceil-mode phantom cells), dividing by `∏ (1 + (end - start - 1)/dilation)`; - exclude-pad divisor counts only in-bounds cells. In `Pool<T, AveragePool, Layout>::ComputeInternal`, a cheap dispatch guard routes to the custom kernel **only** when the pooling is non-global **and** (`asymmetric pads` **or** `!default_dilations`). Every symmetric, non-dilated case — the overwhelmingly common path, including **all** `GlobalAveragePool` — stays on the existing cuDNN fast path, so there is **zero perf regression** on the common path. `GlobalAveragePool` is excluded explicitly because `PoolAttributes` leaves its `kernel_shape`/`strides`/`dilations` unpopulated. `MaxPool` is unaffected (it ignores pad cells; `MaxPool<8>` already routes dilation to its own custom kernel via the same `!default_dilations` check we mirror here). Covers fp32, fp64, fp16, and bf16 (fp16/bf16 accumulate in float). ### Tests Added CUDA-**un-excluded** parity tests in `pool_op_test.cc` — the CUDA leg actually runs and must match the CPU reference oracle: - asymmetric tail-pad 1D/2D, include- and exclude-pad; - `auto_pad=SAME_UPPER` and `SAME_LOWER` (naturally asymmetric); - **symmetric-pad + dilation>1** (wrong on cuDNN, correct via the kernel — locks in the dilation guard); - a symmetric-pad regression case (must stay on cuDNN and remain correct); - fp16; - a `MaxPool` asymmetric-pad case confirming MaxPool is unaffected. The `ceil_mode + count_include_pad` cases use opset 19 so the CPU leg runs the already-correct v19 reference functor and validates the CUDA kernel independently of the separate opset-7..18 CPU MLAS fix (#29629); CUDA routing is opset-independent. Full `PoolTest` suite passes on an A100 (53 passed / 2 DML-skipped / 0 failures). ### Related - CPU (a): #29629 — opset 7-18 MLAS `ceil_mode + count_include_pad` divisor fix. - JSEP (c): #29627 — JSEP pooling output shape honoring `ceil_mode`. - pytorch/pytorch#183528. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Author
Parents
Loading