diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 87675a3f9c12..855972b9489d 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -1170,6 +1170,21 @@ def _impl_v13(cls, bb, inputs, attr, params): return relax.op.astype(inputs[0], to_type) +class CastLike(OnnxOpConverter): + """Convert an onnx CastLike node into an equivalent Relax expression.""" + + @classmethod + def _impl_v15(cls, bb, inputs, attr, params): + data = inputs[0] + target = inputs[1] + target_dtype = getattr(getattr(getattr(target, "ty", None), "dtype", None), "dtype", None) + if target_dtype is None: + target_dtype = getattr(target.struct_info, "dtype", None) + if target_dtype is None: + raise ValueError(f"CastLike: unable to determine dtype from target {target}") + return relax.op.astype(data, target_dtype) + + class Gather(OnnxOpConverter): """Convert an onnx Gather node into an equivalent Relax expression.""" @@ -1506,19 +1521,26 @@ def _impl_v14(cls, bb, inputs, attr, params): x = inputs[0] k = inputs[1] if len(inputs) > 1 else 0 - if len(inputs) > 1: - k = get_constant(inputs[1], params) - if isinstance(k, relax.Constant): - k = int(k.data.numpy().item()) - else: - raise ValueError("Currently only support constant k for Trilu op.") - else: - k = 0 + if isinstance(k, relax.Constant): + k = int(k.data.numpy().item()) + if isinstance(k, int): + if upper: + return relax.op.triu(x, k) + return relax.op.tril(x, k) + # Dynamic k: build the mask explicitly so it works with any scalar k. + shape = x.ty.shape + m, n = shape[-2], shape[-1] + row_idx = relax.op.reshape(relax.op.arange(0, m, dtype="int64"), (m, 1)) + col_idx = relax.op.reshape(relax.op.arange(0, n, dtype="int64"), (1, n)) + diff = relax.op.subtract(col_idx, row_idx) + k_int64 = relax.op.astype(k, "int64") if upper: - return relax.op.triu(x, k) + mask = relax.op.greater_equal(diff, k_int64) else: - return relax.op.tril(x, k) + mask = relax.op.less_equal(diff, k_int64) + mask = relax.op.broadcast_to(mask, shape) + return relax.op.where(mask, x, relax.const(0, x.ty.dtype.dtype)) class Relu(OnnxOpConverter): @@ -5241,6 +5263,7 @@ def _get_convert_map(): "Max": Max, "Mean": Mean, "Cast": Cast, + "CastLike": CastLike, "Gemm": Gemm, "MatMul": MatMul, "MatMulInteger": MatMulInteger, diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 972bc4830729..0ad958c62d6c 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -1358,6 +1358,36 @@ def test_cast_nan_inf_to_int8(): np.testing.assert_array_equal(out_np, expected) +def test_castlike_ir(): + castlike_node = helper.make_node("CastLike", ["a", "b"], ["c"]) + graph = helper.make_graph( + [castlike_node], + "castlike_test", + inputs=[ + helper.make_tensor_value_info("a", TensorProto.FLOAT, [1, 32]), + helper.make_tensor_value_info("b", TensorProto.INT32, [1]), + ], + outputs=[helper.make_tensor_value_info("c", TensorProto.INT32, [1, 32])], + ) + model = helper.make_model(graph, producer_name="castlike_test") + tvm_model = from_onnx(model, opset=15, keep_params_in_input=True) + + @I.ir_module + class Expected: + @R.function + def main( + a: R.Tensor((1, 32), dtype="float32"), + b: R.Tensor((1,), dtype="int32"), + ) -> R.Tensor((1, 32), dtype="int32"): + R.func_attr({"num_input": 2}) + with R.dataflow(): + gv: R.Tensor((1, 32), dtype="int32") = R.astype(a, "int32") + R.output(gv) + return gv + + tvm.ir.assert_structural_equal(tvm_model, Expected) + + def test_gather(): def _verify_gather(data_shape, indices, out_shape, expected, axis=0): gather_node = helper.make_node("Gather", ["data", "indices"], ["y"], axis=axis) @@ -3088,6 +3118,74 @@ def test_trilu_with_const_k(k_value: int): check_correctness(model) +@pytest.mark.parametrize("upper", [True, False]) +def test_trilu_dynamic_k_ir(upper: bool): + graph = helper.make_graph( + [ + helper.make_node("Trilu", inputs=["x", "k"], outputs=["y"], upper=upper), + ], + "trilu_dynamic_k_graph", + inputs=[ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("k", TensorProto.INT64, []), + ], + outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + ) + model = helper.make_model(graph, producer_name="trilu_dynamic_k_graph") + tvm_model = from_onnx(model, opset=14, keep_params_in_input=True) + + if upper: + + @I.ir_module + class ExpectedTriu: + @R.function + def main( + x: R.Tensor((2, 3), dtype="float32"), + k: R.Tensor((), dtype="int64"), + ) -> R.Tensor((2, 3), dtype="float32"): + R.func_attr({"num_input": 2}) + with R.dataflow(): + lv: R.Tensor((3,), dtype="int64") = R.arange(0, 3, 1, dtype="int64") + lv1: R.Tensor((1, 3), dtype="int64") = R.reshape(lv, R.shape([1, 3])) + lv2: R.Tensor((2,), dtype="int64") = R.arange(0, 2, 1, dtype="int64") + lv3: R.Tensor((2, 1), dtype="int64") = R.reshape(lv2, R.shape([2, 1])) + lv4: R.Tensor((2, 3), dtype="int64") = R.subtract(lv1, lv3) + lv5: R.Tensor((), dtype="int64") = R.astype(k, dtype="int64") + lv6: R.Tensor((2, 3), dtype="bool") = R.greater_equal(lv4, lv5) + lv7: R.Tensor((2, 3), dtype="bool") = R.broadcast_to(lv6, R.shape([2, 3])) + gv: R.Tensor((2, 3), dtype="float32") = R.where(lv7, x, R.const(0.0, "float32")) + R.output(gv) + return gv + + expected = ExpectedTriu + else: + + @I.ir_module + class ExpectedTril: + @R.function + def main( + x: R.Tensor((2, 3), dtype="float32"), + k: R.Tensor((), dtype="int64"), + ) -> R.Tensor((2, 3), dtype="float32"): + R.func_attr({"num_input": 2}) + with R.dataflow(): + lv: R.Tensor((3,), dtype="int64") = R.arange(0, 3, 1, dtype="int64") + lv1: R.Tensor((1, 3), dtype="int64") = R.reshape(lv, R.shape([1, 3])) + lv2: R.Tensor((2,), dtype="int64") = R.arange(0, 2, 1, dtype="int64") + lv3: R.Tensor((2, 1), dtype="int64") = R.reshape(lv2, R.shape([2, 1])) + lv4: R.Tensor((2, 3), dtype="int64") = R.subtract(lv1, lv3) + lv5: R.Tensor((), dtype="int64") = R.astype(k, dtype="int64") + lv6: R.Tensor((2, 3), dtype="bool") = R.less_equal(lv4, lv5) + lv7: R.Tensor((2, 3), dtype="bool") = R.broadcast_to(lv6, R.shape([2, 3])) + gv: R.Tensor((2, 3), dtype="float32") = R.where(lv7, x, R.const(0.0, "float32")) + R.output(gv) + return gv + + expected = ExpectedTril + + tvm.ir.assert_structural_equal(tvm_model, expected) + + def test_selu(): model = make_unary_model("Selu", [2, 3]) tvm_model = from_onnx(model, keep_params_in_input=True) diff --git a/tests/python/relax/test_frontend_onnx_backend.py b/tests/python/relax/test_frontend_onnx_backend.py index 9f9ec8779cdd..2a6c7bdbcdcf 100644 --- a/tests/python/relax/test_frontend_onnx_backend.py +++ b/tests/python/relax/test_frontend_onnx_backend.py @@ -19,13 +19,13 @@ ONNX Backend Tests =================== Systematically verify the Relax ONNX importer using the official ONNX -Backend Test Suite (node-level tests only). Each test loads a small -ONNX model with protobuf reference inputs/outputs and checks that the -Relax-imported model produces numerically correct results. +Backend Test Suite. Each test loads a small ONNX model with protobuf +reference inputs/outputs and checks that the Relax-imported model +produces numerically correct results. -Only ``onnx.backend.test.data.node`` tests are registered here; real, -simple, and PyTorch model tests are out of scope for importer-level -semantic verification. +Currently ``_INCLUDE_OPS`` selects node-level operator tests. Other +test classes (real/simple/PyTorch models) remain available in +``backend_test.test_cases`` and can be enabled explicitly in the future. """ @@ -81,9 +81,9 @@ def run(self, inputs, **kwargs): self._vm.invoke_stateful("main") output = self._vm.get_outputs("main") - if isinstance(output, (tvm.runtime.Tensor, np.ndarray)): # noqa: UP038 + if isinstance(output, tvm.runtime.Tensor | np.ndarray): return (output.numpy() if hasattr(output, "numpy") else output,) - if isinstance(output, (tuple, list)): # noqa: UP038 + if isinstance(output, tuple | list): return tuple(o.numpy() if hasattr(o, "numpy") else np.array(o) for o in output) return (np.array(output),) @@ -172,6 +172,7 @@ def supports_device(cls, device: str) -> bool: "less", "less_equal", "lrn", + "logsoftmax", "matmul", "matmulinteger", "mean", @@ -183,6 +184,7 @@ def supports_device(cls, device: str) -> bool: "not", "or", "reciprocal", + "relu", "round", "scatternd", "sigmoid", @@ -191,6 +193,7 @@ def supports_device(cls, device: str) -> bool: "sinh", "size", "slice", + "softmax", "spacetodepth", "sqrt", "squeeze", @@ -200,6 +203,8 @@ def supports_device(cls, device: str) -> bool: "tanh", "tile", "transpose", + "tril", + "triu", "unique", "unsqueeze", "where", @@ -209,4 +214,15 @@ def supports_device(cls, device: str) -> bool: for _op in _INCLUDE_OPS: backend_test.include(rf"^test_{_op}(?:_.*)?(?:_cpu|_cuda)$") +# A small number of model-level tests (e.g. from PyTorch converted models) +# have names that collide with the node-level include patterns above. The +# current adapter is focused on node-level protobuf test cases, so exclude +# those known collisions explicitly rather than limiting the test classes. +_EXCLUDE_PATTERNS = [ + r"^test_softmax_functional_dim3_cpu$", + r"^test_softmax_lastdim_cpu$", +] +for _pattern in _EXCLUDE_PATTERNS: + backend_test.exclude(_pattern) + globals().update(backend_test.test_cases)