[Relax][ONNX] Add CastLike support and dynamic-k Trilu to expand backend coverage

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

@@ -1170,6 +1170,17 @@ 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 = target.ty.dtype.dtype
+        return relax.op.astype(data, target_dtype)
+
+
 class Gather(OnnxOpConverter):
     """Convert an onnx Gather node into an equivalent Relax expression."""
 
@@ -1506,19 +1517,29 @@ 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(
+            relax.op.broadcast_to(col_idx, (m, n)),
+            relax.op.broadcast_to(row_idx, (m, n)),
+        )
+        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))
 
 
 class Relu(OnnxOpConverter):
@@ -5241,6 +5262,7 @@ def _get_convert_map():
         "Max": Max,
         "Mean": Mean,
         "Cast": Cast,
+        "CastLike": CastLike,
         "Gemm": Gemm,
         "MatMul": MatMul,
         "MatMulInteger": MatMulInteger,

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,78 @@ 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, 3), dtype="int64") = R.broadcast_to(lv1, R.shape([2, 3]))
+                    lv3: R.Tensor((2,), dtype="int64") = R.arange(0, 2, 1, dtype="int64")
+                    lv4: R.Tensor((2, 1), dtype="int64") = R.reshape(lv3, R.shape([2, 1]))
+                    lv5: R.Tensor((2, 3), dtype="int64") = R.broadcast_to(lv4, R.shape([2, 3]))
+                    lv6: R.Tensor((2, 3), dtype="int64") = R.subtract(lv2, lv5)
+                    lv7: R.Tensor((), dtype="int64") = R.astype(k, dtype="int64")
+                    lv8: R.Tensor((2, 3), dtype="bool") = R.greater_equal(lv6, lv7)
+                    lv9: R.Tensor((2, 3), dtype="bool") = R.broadcast_to(lv8, R.shape([2, 3]))
+                    gv: R.Tensor((2, 3), dtype="float32") = R.where(lv9, 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, 3), dtype="int64") = R.broadcast_to(lv1, R.shape([2, 3]))
+                    lv3: R.Tensor((2,), dtype="int64") = R.arange(0, 2, 1, dtype="int64")
+                    lv4: R.Tensor((2, 1), dtype="int64") = R.reshape(lv3, R.shape([2, 1]))
+                    lv5: R.Tensor((2, 3), dtype="int64") = R.broadcast_to(lv4, R.shape([2, 3]))
+                    lv6: R.Tensor((2, 3), dtype="int64") = R.subtract(lv2, lv5)
+                    lv7: R.Tensor((), dtype="int64") = R.astype(k, dtype="int64")
+                    lv8: R.Tensor((2, 3), dtype="bool") = R.less_equal(lv6, lv7)
+                    lv9: R.Tensor((2, 3), dtype="bool") = R.broadcast_to(lv8, R.shape([2, 3]))
+                    gv: R.Tensor((2, 3), dtype="float32") = R.where(lv9, 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)

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)