From fb532bf038e35bee361d86b396769929a8ab8baf Mon Sep 17 00:00:00 2001 From: RunguoLi Date: Mon, 21 Sep 2026 15:43:20 -0500 Subject: [PATCH] [Fix][Relax][ONNX] Support symbolic dims in SAME auto_pad pooling The pool converter materialized every input dimension as a Python int to derive SAME_UPPER/SAME_LOWER padding, so MaxPool/AveragePool/LpPool failed with `int(Var)` on graphs with symbolic spatial extents (e.g. PP-OCR's dynamic-width MaxPool). With unit stride the SAME padding is `dilated_kernel - 1` regardless of the input extent, so compute it without touching the shape and only require a static extent when stride > 1, raising OpAttributeUnImplemented otherwise. Unifying the max/avg paths into one spec-following helper also fixes two silent correctness issues: MaxPool SAME_LOWER used floor(in / stride) instead of ceil for the output extent, producing incorrect pads when the extent is not divisible by the stride, and AveragePool ignored dilations when resolving auto_pad. Fixes #20398 --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 81 +++++++--------- tests/python/relax/test_frontend_onnx.py | 92 +++++++++++++++++++ 2 files changed, 126 insertions(+), 47 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index a3e6140381a0..bc95f051e333 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -4357,17 +4357,39 @@ class Pool(OnnxOpConverter): name = "" @classmethod - def get_pad_pair(cls, input1d, kernel1d, stride1d, mode): - """infer pad size""" - if input1d % stride1d == 0: - pad = max(kernel1d - stride1d, 0) - else: - pad = max(kernel1d - (input1d % stride1d), 0) - pad_before = pad // 2 - pad_after = pad - pad_before - if "LOWER" in mode: - return [pad_after, pad_before] - return [pad_before, pad_after] + def get_same_pads(cls, input_shape, kernel_shape, strides, dilations, auto_pad): + """Compute the explicit padding for SAME_UPPER / SAME_LOWER auto_pad. + + Per the ONNX spec, both SAME modes produce ``ceil(input / stride)`` output + elements along each spatial axis. The total padding of an axis is therefore + ``(ceil(input / stride) - 1) * stride + dilated_kernel - input``, which + simplifies to ``dilated_kernel - 1`` when ``stride == 1``. This lets us + support symbolic spatial extents whenever the stride is 1. + + Returns the padding as ``(begin_0, ..., begin_n, end_0, ..., end_n)``. + """ + pads_begin, pads_end = [], [] + for i, dim in enumerate(list(input_shape)[2:]): + dilated_kernel = (kernel_shape[i] - 1) * dilations[i] + 1 + if strides[i] == 1: + total_pad = dilated_kernel - 1 + elif isinstance(dim, tirx.IntImm | int): + dim = int(dim) + out_dim = (dim + strides[i] - 1) // strides[i] + total_pad = max((out_dim - 1) * strides[i] + dilated_kernel - dim, 0) + else: + raise tvm.error.OpAttributeUnImplemented( + f"{auto_pad} auto_pad with stride {strides[i]} is not supported for " + f"symbolic spatial dimension {dim} in operator {cls.__name__}, " + "since the required padding depends on the runtime extent." + ) + if auto_pad == "SAME_UPPER": + pad_begin = total_pad // 2 + else: + pad_begin = total_pad - total_pad // 2 + pads_begin.append(pad_begin) + pads_end.append(total_pad - pad_begin) + return tuple(pads_begin + pads_end) @classmethod def _impl_v1(cls, bb, inputs, attr, params): @@ -4394,46 +4416,11 @@ def _impl_v1(cls, bb, inputs, attr, params): ], f"Value {auto_pad} in attribute auto_pad is invalid." if auto_pad in ("SAME_UPPER", "SAME_LOWER"): - pads = [] - if cls.name == "avg_pool": - for axis in range(len(input_shape) - 2): - axis_shape = int(input_shape[2 + axis]) - stride = strides[axis] - kernel = kernel_shape[axis] - pad = cls.get_pad_pair(axis_shape, kernel, stride, auto_pad) - pads.append(pad) - else: - input_spatial_shape = cls._get_input_spatial_shape(data) - output_spatial_shape = [0 for _ in input_spatial_shape] - - for i, _ in enumerate(input_spatial_shape): - if auto_pad == "SAME_UPPER": - output_spatial_shape[i] = int(_np.ceil(input_spatial_shape[i] / strides[i])) - else: - output_spatial_shape[i] = int( - _np.floor(input_spatial_shape[i] / strides[i]) - ) - pad_i = ( - (output_spatial_shape[i] - 1) * strides[i] - + ((kernel_shape[i] - 1) * dilations[i] + 1) - - input_spatial_shape[i] - ) - - if auto_pad == "SAME_UPPER": - pads.append([pad_i // 2, pad_i - pad_i // 2]) - else: - pads.append([pad_i - pad_i // 2, pad_i // 2]) - - pads = tuple([val for pair in zip(*pads) for val in pair]) + pads = cls.get_same_pads(input_shape, kernel_shape, strides, dilations, auto_pad) op = getattr(relax.op.nn, cls.name + str(len(kernel_shape)) + "d") return op(data, kernel_shape, strides, pads, dilations, ceil_mode, count_include_pad) - @classmethod - def _get_input_spatial_shape(cls, tensor): - # shape is (N x C x D1 x D2 ... Dn) - return _np.array([int(d) for d in tensor.ty.shape], dtype="int64")[2:] - class MaxPool(Pool): """Converts an onnx MaxPool node into an equivalent Relax expression.""" diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index d7aa987c0527..b5c49ae2e01a 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -9999,6 +9999,98 @@ def main(x: R.Tensor(input_shape, dtype="float32")): ) +def _make_same_pool_model(pool_name, input_shape, auto_pad, kernel_shape, strides, dilations=None): + attrs = {"kernel_shape": kernel_shape, "strides": strides, "auto_pad": auto_pad} + if dilations is not None: + attrs["dilations"] = dilations + node = helper.make_node(pool_name, ["x"], ["y"], **attrs) + graph = helper.make_graph( + [node], + "same_pool_test", + inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape)], + outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, None)], + ) + return helper.make_model(graph, producer_name="same_pool_test") + + +@pytest.mark.parametrize("pool_name", ["MaxPool", "AveragePool"]) +@pytest.mark.parametrize("auto_pad", ["SAME_UPPER", "SAME_LOWER"]) +@pytest.mark.parametrize( + "input_shape, kernel_shape, strides", + [ + # Spatial extents not divisible by the strides. + ([1, 2, 7], [2], [2]), + ([1, 2, 7, 9], [3, 3], [2, 3]), + ([1, 2, 5, 7, 6], [2, 3, 3], [3, 2, 4]), + ], +) +def test_pool_same_padding_numerical(pool_name, auto_pad, input_shape, kernel_shape, strides): + model = _make_same_pool_model(pool_name, input_shape, auto_pad, kernel_shape, strides) + check_correctness(model, opset=18) + + +def _get_pool2d_call(func): + pool_calls = [] + + def visit(expr): + if ( + isinstance(expr, relax.Call) + and isinstance(expr.op, tvm.ir.Op) + and expr.op.name in ("relax.nn.max_pool2d", "relax.nn.avg_pool2d") + ): + pool_calls.append(expr) + + relax.analysis.post_order_visit(func.body, visit) + assert len(pool_calls) == 1 + return pool_calls[0] + + +@pytest.mark.parametrize("pool_name", ["MaxPool", "AveragePool"]) +@pytest.mark.parametrize( + "auto_pad, expected_padding", [("SAME_UPPER", (2, 1, 2, 2)), ("SAME_LOWER", (2, 2, 2, 1))] +) +def test_pool_same_padding_dilation(pool_name, auto_pad, expected_padding): + # The padding must account for the dilated kernel so that the output extent is + # ceil(input / stride), matching ONNX shape inference. onnxruntime ignores the + # dilation when resolving auto_pad, so this is checked structurally. + model = _make_same_pool_model( + pool_name, [1, 2, 9, 10], auto_pad, kernel_shape=[3, 2], strides=[2, 3], dilations=[2, 3] + ) + func = from_onnx(model, opset=19, keep_params_in_input=True)["main"] + + pool_call = _get_pool2d_call(func) + assert tuple(int(value) for value in pool_call.attrs.padding) == expected_padding + assert tuple(int(value) for value in func.ret_ty.shape.values) == (1, 2, 5, 4) + + +@pytest.mark.parametrize("pool_name", ["MaxPool", "AveragePool", "LpPool"]) +@pytest.mark.parametrize( + "auto_pad, expected_padding", [("SAME_UPPER", (0, 0, 1, 1)), ("SAME_LOWER", (1, 1, 0, 0))] +) +def test_pool_same_padding_symbolic_unit_stride(pool_name, auto_pad, expected_padding): + # With unit strides the SAME padding does not depend on the input extent, + # so it can be computed for symbolic spatial dimensions. + model = _make_same_pool_model(pool_name, ["N", 3, 48, "W"], auto_pad, [2, 2], [1, 1]) + func = from_onnx(model, opset=18, keep_params_in_input=True)["main"] + + pool_call = _get_pool2d_call(func) + assert tuple(int(value) for value in pool_call.attrs.padding) == expected_padding + n, _, _, w = func.params[0].ty.shape.values + out_n, out_c, out_h, out_w = func.ret_ty.shape.values + tvm.ir.assert_structural_equal(out_n, n) + tvm.ir.assert_structural_equal(out_w, w) + assert (int(out_c), int(out_h)) == (3, 48) + + x = rg.standard_normal(size=[2, 3, 48, 17]).astype("float32") + check_correctness(model, inputs={"x": x}, opset=18) + + +def test_pool_same_padding_symbolic_non_unit_stride(): + model = _make_same_pool_model("MaxPool", [1, 3, 48, "W"], "SAME_UPPER", [2, 2], [2, 2]) + with pytest.raises(tvm.error.OpAttributeUnImplemented, match="symbolic spatial dimension"): + from_onnx(model, opset=18) + + @pytest.mark.parametrize("p", [1, 3]) def test_lppool_negative_input(p: int): input_data = np.array([[[-1.0, 2.0, -3.0, 4.0]]], dtype="float32")