[Fix][Relax][ONNX] Fix Upsample and Resize-10 import

python/tvm/relax/frontend/onnx/onnx_frontend.py

@@ -3821,6 +3821,13 @@ def resize3d_dyn(d, r, s0, s1, s2):
 class Resize(OnnxOpConverter):
     """Converts an onnx Resize node into an equivalent Relax expression."""
 
+    @classmethod
+    def _impl_v10(cls, bb, inputs, attr, params):
+        # Resize-10 takes (X, scales) and has the same semantics as Upsample-9.
+        scales = _get_constant_scales(inputs[1], params, "Resize")
+        mode = attr.get("mode", b"nearest")
+        return _legacy_upsample(bb, inputs[0], scales, mode, "Resize")
+
     @classmethod
     def _impl_v18(cls, bb, inputs, attr, params):
         # Extract the many attributes of resize.
@@ -3950,6 +3957,9 @@ def _impl_v18(cls, bb, inputs, attr, params):
                 extrapolation_value=extrapolation_value,
             )
 
+    # Resize-11 through Resize-18 share the (X, roi, scales, sizes) signature.
+    _impl_v11 = _impl_v18
+
 
 class AffineGrid(OnnxOpConverter):
     """Converts an onnx AffineGrid node into an equivalent Relax expression."""
@@ -5456,33 +5466,98 @@ def _impl_v9(cls, bb, inputs, attr, params):
         )
 
 
-class Upsample(OnnxOpConverter):
-    """Operator converter for Upsample (nearest mode)."""
+def _legacy_upsample(bb, data, scales, mode, op_name):
+    """Lower Upsample (opset 7-9) and Resize-10, which share the same semantics.
 
-    @classmethod
-    def _impl_v9(cls, bb, inputs, attr, params):
-        scales = attr.get("scales")
-        assert len(scales) == 4
-        assert scales[0] == scales[1] == 1
+    Each output extent is ``floor(input * scale)`` and coordinates are mapped with
+    ``x_in = x_out / scale`` (asymmetric). In nearest mode the source index is
+    rounded down when upsampling and up when downsampling.
+    """
+    mode = mode.decode("ascii") if isinstance(mode, bytes) else mode
+    if mode not in ("nearest", "linear"):
+        raise tvm.error.OpAttributeInvalid(
+            f'Value {mode} in attribute "mode" of operator {op_name} is not valid.'
+        )
+    scales = [float(s) for s in scales]
+    shape = list(data.ty.shape)
+    ndims = len(shape)
+    if len(scales) != ndims:
+        raise tvm.error.OpAttributeInvalid(
+            f"{op_name} expects {ndims} scales for a rank-{ndims} input, got {len(scales)}."
+        )
+    if ndims not in (3, 4, 5) or scales[0] != 1.0 or scales[1] != 1.0:
+        raise tvm.error.OpAttributeUnImplemented(
+            f"{op_name} is only supported for 3-D/4-D/5-D inputs with unit batch and "
+            f"channel scales, got scales {scales} for a rank-{ndims} input."
+        )
 
-        inp_shape = [int(x) for x in inputs[0].ty.shape]
-        assert len(inp_shape) == 4
-        out_shape2d = [int(dim * scale) for dim, scale in zip(inp_shape[2:], scales[2:])]
+    sizes = []
+    for dim, scale in zip(shape[2:], scales[2:]):
+        if isinstance(dim, tirx.IntImm):
+            sizes.append(math.floor(int(dim) * scale))
+        elif scale.is_integer():
+            sizes.append(dim * int(scale))
+        else:
+            sizes.append((dim.astype("float32") * scale).astype("int64"))
+
+    spatial_scales = scales[2:]
+    if mode == "linear":
+        method, rounding_method = "linear", ""
+    elif all(s >= 1.0 for s in spatial_scales):
+        method, rounding_method = "nearest_neighbor", "floor"
+    elif all(s < 1.0 for s in spatial_scales):
+        method, rounding_method = "nearest_neighbor", "ceil"
+    else:
+        raise tvm.error.OpAttributeUnImplemented(
+            f"{op_name} nearest mode cannot mix upsampling and downsampling axes, "
+            f"got scales {scales}."
+        )
 
-        mode = attr.get("mode", b"nearest").decode("ascii")
-        if mode == "nearest":
-            mode = "nearest_neighbor"
-        msg = f'Value {mode} in attribute "mode" of operator Upsample is not valid.'
-        assert mode in ("linear", "nearest_neighbor", "cubic"), msg
+    if ndims == 3:
+        return bb.emit_te(
+            topi.image.resize1d,
+            data,
+            [0.0, 0.0],
+            sizes,
+            "NCW",
+            method,
+            "asymmetric",
+            rounding_method,
+        )
+    resize_op = relax.op.image.resize2d if ndims == 4 else relax.op.image.resize3d
+    return resize_op(
+        data,
+        size=relax.ShapeExpr(sizes),
+        layout="NCHW" if ndims == 4 else "NCDHW",
+        method=method,
+        coordinate_transformation_mode="asymmetric",
+        rounding_method=rounding_method,
+    )
 
-        return relax.op.image.resize2d(
-            data=inputs[0],
-            roi=None,
-            size=relax.ShapeExpr(out_shape2d),  # (H, W)
-            layout="NCHW",
-            method=mode,
-            coordinate_transformation_mode="asymmetric",  # Align with Upsample
+
+def _get_constant_scales(scales, params, op_name):
+    scales = get_constant(scales, params)
+    if not isinstance(scales, tvm.ir.GenericConst):
+        raise tvm.error.OpAttributeUnImplemented(
+            f"{op_name} with non-constant scales is not supported."
         )
+    return scales.value.numpy().tolist()
+
+
+class Upsample(OnnxOpConverter):
+    """Converts an onnx Upsample node into an equivalent Relax expression."""
+
+    @classmethod
+    def _impl_v7(cls, bb, inputs, attr, params):
+        mode = attr.get("mode", b"nearest")
+        return _legacy_upsample(bb, inputs[0], attr["scales"], mode, "Upsample")
+
+    @classmethod
+    def _impl_v9(cls, bb, inputs, attr, params):
+        # Since opset 9 the scales are an input rather than an attribute.
+        scales = _get_constant_scales(inputs[1], params, "Upsample")
+        mode = attr.get("mode", b"nearest")
+        return _legacy_upsample(bb, inputs[0], scales, mode, "Upsample")
 
 
 class HardSigmoid(OnnxOpConverter):

tests/python/relax/test_frontend_onnx.py

@@ -9432,6 +9432,97 @@ def _visit(expr):
     assert seen_resize3d
 
 
+def _make_legacy_upsample_model(op_name, opset, input_shape, scales, mode):
+    """Upsample-7 takes scales as an attribute; Upsample-9 and Resize-10 take an input."""
+    inputs, initializer, attrs = ["X"], [], {"mode": mode}
+    if op_name == "Upsample" and opset < 9:
+        attrs["scales"] = [float(s) for s in scales]
+    else:
+        inputs.append("scales")
+        initializer.append(helper.make_tensor("scales", TensorProto.FLOAT, [len(scales)], scales))
+    node = helper.make_node(op_name, inputs, ["Y"], **attrs)
+    graph = helper.make_graph(
+        [node],
+        "legacy_upsample",
+        inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)],
+        initializer=initializer,
+        outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, None)],
+    )
+    return helper.make_model(graph, producer_name="legacy_upsample")
+
+
+@pytest.mark.parametrize("op_name, opset", [("Upsample", 7), ("Upsample", 9), ("Resize", 10)])
+@pytest.mark.parametrize(
+    "input_shape, scales, mode",
+    [
+        ([1, 2, 4, 5], [1.0, 1.0, 2.0, 2.0], "nearest"),
+        ([1, 2, 4, 5], [1.0, 1.0, 3.0, 2.0], "nearest"),
+        # floor(5 * 2.5) = 12: the nearest source index uses the given scale.
+        ([1, 2, 4, 5], [1.0, 1.0, 1.5, 2.5], "nearest"),
+        ([1, 2, 4, 5], [1.0, 1.0, 2.0, 3.0], "linear"),
+        ([1, 2, 5], [1.0, 1.0, 2.0], "nearest"),
+        ([1, 2, 3, 4, 5], [1.0, 1.0, 2.0, 1.0, 3.0], "nearest"),
+        ([1, 2, 3, 4, 5], [1.0, 1.0, 2.0, 2.0, 2.0], "linear"),
+    ],
+)
+def test_legacy_upsample(op_name, opset, input_shape, scales, mode):
+    model = _make_legacy_upsample_model(op_name, opset, input_shape, scales, mode)
+    check_correctness(model, opset=opset)
+
+
+@pytest.mark.parametrize("mode", ["nearest", "linear"])
+def test_resize_v10_downsample(mode):
+    # Resize-10 allows scales below one; nearest then rounds the source index up.
+    model = _make_legacy_upsample_model("Resize", 10, [1, 2, 6, 5], [1.0, 1.0, 0.5, 0.6], mode)
+    check_correctness(model, opset=10)
+
+
+@pytest.mark.parametrize("op_name, opset", [("Upsample", 7), ("Upsample", 9), ("Resize", 10)])
+def test_legacy_upsample_symbolic_shape(op_name, opset):
+    model = _make_legacy_upsample_model(
+        op_name, opset, ["N", 2, "H", "W"], [1.0, 1.0, 2.0, 1.5], "nearest"
+    )
+    func = from_onnx(model, opset=opset, keep_params_in_input=True)["main"]
+    n, _, h, _ = func.params[0].ty.shape.values
+    out_n, out_c, out_h, _ = func.ret_ty.shape.values
+    tvm.ir.assert_structural_equal(out_n, n)
+    assert int(out_c) == 2
+    # Integer scales keep the output extent an exact integer expression.
+    assert tvm.sym.Analyzer().can_prove_equal(out_h, h * 2)
+
+    x = generate_random_value([2, 2, 4, 6], TensorProto.FLOAT)
+    check_correctness(model, inputs={"X": x}, opset=opset)
+
+
+def test_resize_v11_still_uses_roi_scales_sizes_signature():
+    # Adding Resize-10 must not change which converter handles opsets 11-17.
+    node = helper.make_node("Resize", ["X", "roi", "scales"], ["Y"], mode="nearest")
+    graph = helper.make_graph(
+        [node],
+        "resize_v11",
+        inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 2, 4, 5])],
+        initializer=[
+            helper.make_tensor("roi", TensorProto.FLOAT, [0], []),
+            helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.0, 2.0]),
+        ],
+        outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, None)],
+    )
+    model = helper.make_model(graph, producer_name="resize_v11")
+    check_correctness(model, opset=11)
+
+
+def test_legacy_upsample_unsupported_scales():
+    model = _make_legacy_upsample_model(
+        "Upsample", 9, [1, 2, 4, 5], [1.0, 2.0, 2.0, 2.0], "nearest"
+    )
+    with pytest.raises(tvm.error.OpAttributeUnImplemented, match="unit batch and channel"):
+        from_onnx(model, opset=9)
+
+    model = _make_legacy_upsample_model("Resize", 10, [1, 2, 4, 5], [1.0, 1.0, 0.5, 2.0], "nearest")
+    with pytest.raises(tvm.error.OpAttributeUnImplemented, match="cannot mix"):
+        from_onnx(model, opset=10)
+
+
 def test_einsum():
     eqn = "ij->i"
     einsum_node = helper.make_node("Einsum", ["x"], ["y"], equation=eqn)