[Fix][Relax][ONNX] Support symbolic dims in SAME auto_pad pooling

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."""

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")