ITADN

[Bug] Pool ceil_mode=1 causes tensor.cast shape mismatch compilation failure (AveragePool, MaxPool, 1D/2D/3D)

#24579Openwuyii8941 创建于 2026-06-08
bug 🐞
W
wuyii8941commented
### 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 条评论