[Bug] Pool ceil_mode=1 causes tensor.cast shape mismatch compilation failure (AveragePool, MaxPool, 1D/2D/3D)
bug 🐞
### What happened?
### What happened?
When `AveragePool` or `MaxPool` has `ceil_mode=1` and `(input_size + 2*pad - kernel_size) % stride != 0`, the ONNX-to-torch lowering computes the correct ceil-mode output shape in the type system, but the actual pooling computation only produces floor-mode number of elements, causing a `tensor.cast` incompatibility error during compilation.
```
error: 'tensor.cast' op operand type 'tensor<1x3x4x4xf32>' and result type 'tensor<1x3x5x5xf32>' are cast incompatible
```
This affects AveragePool and MaxPool across all dimensionalities (1D, 2D, 3D). ONNX Runtime handles all these cases correctly.
### Affected configurations
All configurations where `(H + 2*pad - kernel) % stride != 0` trigger this bug:
| Op | Dims | H | kernel | stride | pad | Floor shape | Ceil shape (expected) |
|---|---|---|---|---|---|---|---|
| AveragePool | 2D | 7 | 2 | 2 | 1 | 4x4 | **5x5** |
| AveragePool | 2D | 9 | 2 | 2 | 1 | 5x5 | **6x6** |
| MaxPool | 2D | 7 | 2 | 3 | 1 | 3x3 | **4x4** |
| MaxPool | 1D | 7 | 2 | 2 | 1 | 4 | **5** |
| MaxPool | 3D | 7 | 2 | 2 | 1 | 4x4x4 | **5x5x5** |
Tested 40+ configurations — all follow the same pattern.
### Root cause
The pooling lowering computes the output tensor type using the correct ceil formula `ceil((H + 2*pad - kernel) / stride) + 1`, but the underlying linalg pooling operation only produces `floor((H + 2*pad - kernel) / stride) + 1` elements. The mismatch triggers the `tensor.cast` verification failure.
### Relationship to IREE #20432
IREE #20432 reports numerical errors in AveragePool with `ceil_mode=1, count_include_pad=0`. That issue likely hits the same root cause: for certain shape/kernel/stride combinations where `(H+2*pad-kernel) % stride != 0`, the ceil_mode lowering is incorrect. The difference is that #20432's shapes happen to not trigger the `tensor.cast` crash (the floor and ceil output shapes coincide for those specific parameters), while this bug demonstrates cases where they diverge and cause a hard compilation failure.
### Expected behavior
Compilation should succeed and produce correct ceil-mode pooling output matching ONNX Runtime.
### Environment
- IREE compiler: `3.12.0rc20260515`
- Target: `llvm-cpu` (generic)
- ONNX opset: 17
### Steps to reproduce your issue
### Steps to reproduce
```python
import numpy as np
import onnx
from onnx import helper, TensorProto
# AveragePool 2D, ceil_mode=1
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 7, 7])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, None)
nodes = [
helper.make_node("AveragePool", ["X"], ["Y"],
kernel_shape=[2, 2], strides=[2, 2],
pads=[1, 1, 1, 1], ceil_mode=1),
]
graph = helper.make_graph(nodes, "test", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
onnx.save(model, "/tmp/pool_ceil.onnx")
```
```bash
iree-import-onnx /tmp/pool_ceil.onnx -o /tmp/pool_ceil.mlir
iree-compile --iree-hal-target-backends=llvm-cpu --iree-llvmcpu-target-cpu=generic \
/tmp/pool_ceil.mlir -o /tmp/pool_ceil.vmfb
# Error: 'tensor.cast' op operand type 'tensor<1x3x4x4xf32>' and result type 'tensor<1x3x5x5xf32>' are cast incompatible
```
### What component(s) does this issue relate to?
_No response_
### Version information
_No response_
### Additional context
_No response_
0 条评论