chore: import upstream snapshot with attribution

tests/python/ir/test_datatype_nv_fp4.py

@@ -0,0 +1,56 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+import tvm.tirx as tirx
+from tvm import te
+from tvm.script import tirx as T
+
+try:
+    from ml_dtypes import float4_e2m1fn
+except ImportError:
+    float4_e2m1fn = None
+
+
+nv_fp4_dtypes = [(float4_e2m1fn, "float4_e2m1fn")]
+
+
+@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp4_dtypes)
+def test_create_nv_fp4_nd_array(np_dtype, dtype_str):
+    if np_dtype is None:
+        """Skip test if ml_dtypes is not installed"""
+        return
+    x = np.random.rand(128, 128).astype(np_dtype)
+    x_nd = tvm.runtime.tensor(x)
+    assert x_nd.dtype == dtype_str
+    np.testing.assert_equal(x_nd.numpy(), x)
+
+
+@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp4_dtypes)
+def test_nv_fp4_buffer(np_dtype, dtype_str):
+    m = te.var("m")
+    n = te.var("n")
+    A = tvm.tirx.decl_buffer((m, n), dtype_str)
+    assert A.dtype == dtype_str
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/ir/test_datatype_nv_fp8.py

@@ -0,0 +1,134 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E501, F401
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+import tvm.tirx as tirx
+from tvm import te
+from tvm.script import tirx as T
+
+try:
+    from ml_dtypes import (
+        float8_e3m4,
+        float8_e4m3,
+        float8_e4m3b11fnuz,
+        float8_e4m3fn,
+        float8_e4m3fnuz,
+        float8_e5m2,
+        float8_e5m2fnuz,
+        float8_e8m0fnu,
+    )
+except ImportError:
+    float8_e3m4 = float8_e4m3 = float8_e4m3b11fnuz = float8_e4m3fn = None
+    float8_e4m3fnuz = float8_e5m2 = float8_e5m2fnuz = float8_e8m0fnu = None
+
+
+def fp8_unary(dtype: str):
+    @T.prim_func(s_tir=True)
+    def func(
+        a: T.handle,
+        b: T.handle,
+        a_add_b: T.handle,
+        a_sub_b: T.handle,
+        a_mul_b: T.handle,
+        a_fp32: T.handle,
+        a_roundtrip: T.handle,
+    ) -> None:
+        A = T.match_buffer(a, [128], dtype=dtype)
+        B = T.match_buffer(b, [128], dtype=dtype)
+        A_add_B = T.match_buffer(a_add_b, [128], dtype=dtype)
+        A_sub_B = T.match_buffer(a_sub_b, [128], dtype=dtype)
+        A_mul_B = T.match_buffer(a_mul_b, [128], dtype=dtype)
+        A_fp32 = T.match_buffer(a_fp32, [128], dtype="float32")
+        A_roundtrip = T.match_buffer(a_roundtrip, [128], dtype=dtype)
+        for i in range(128):
+            with T.sblock("fp8_unary"):
+                vi = T.axis.spatial(128, i)
+                A_add_B[vi] = A[vi] + B[vi]
+                A_sub_B[vi] = A[vi] - B[vi]
+                A_mul_B[vi] = A[vi] * B[vi]
+                A_fp32[vi] = A[vi]
+                A_roundtrip[vi] = A_fp32[vi]
+
+    return func
+
+
+nv_fp8_dtypes = [
+    (float8_e3m4, "float8_e3m4"),
+    (float8_e4m3, "float8_e4m3"),
+    (float8_e4m3b11fnuz, "float8_e4m3b11fnuz"),
+    (float8_e4m3fn, "float8_e4m3fn"),
+    (float8_e4m3fnuz, "float8_e4m3fnuz"),
+    (float8_e5m2, "float8_e5m2"),
+    (float8_e5m2fnuz, "float8_e5m2fnuz"),
+    (float8_e8m0fnu, "float8_e8m0fnu"),
+]
+
+
+@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp8_dtypes)
+def test_create_nv_fp8_nd_array(np_dtype, dtype_str):
+    if np_dtype is None:
+        """Skip test if ml_dtypes is not installed"""
+        return
+    x = np.random.rand(128, 128).astype(np_dtype)
+    x_nd = tvm.runtime.tensor(x)
+    assert x_nd.dtype == dtype_str
+    np.testing.assert_equal(x_nd.numpy(), x)
+
+
+@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp8_dtypes)
+def test_fp8_unary_op(np_dtype, dtype_str):
+    func = fp8_unary(dtype_str)
+    if not tvm.testing.device_enabled("llvm"):
+        return
+    if np_dtype is None:
+        """Skip test if ml_dtypes is not installed"""
+        return
+    if dtype_str in ["float8_e8m0fnu", "float8_e4m3b11fnuz", "float8_e4m3fnuz", "float8_e5m2fnuz"]:
+        # float8_e8m0fnu does not support arithmetic operations, and unsigned arithmetic is not tested here
+        return
+
+    f = tvm.compile(func, target="llvm")
+    a = np.random.randn(128).astype(np_dtype)
+    b = np.random.randn(128).astype(np_dtype)
+    a_add_b = np.zeros(128).astype(np_dtype)
+    a_sub_b = np.zeros(128).astype(np_dtype)
+    a_mul_b = np.zeros(128).astype(np_dtype)
+    a_fp32 = np.zeros(128).astype(np.float32)
+    a_roundtrip = np.zeros(128).astype(np_dtype)
+    args = list(
+        map(lambda _: tvm.runtime.tensor(_), [a, b, a_add_b, a_sub_b, a_mul_b, a_fp32, a_roundtrip])
+    )
+    f(*args)
+    expected_a_fp32 = a.astype(np.float32)
+    expected_a_roundtrip = expected_a_fp32.astype(np_dtype)
+    np.testing.assert_equal(args[6].numpy(), expected_a_roundtrip)
+
+
+@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp8_dtypes)
+def test_nv_fp8_buffer(np_dtype, dtype_str):
+    m = te.var("m")
+    n = te.var("n")
+    A = tvm.tirx.decl_buffer((m, n), dtype_str)
+    assert A.dtype == dtype_str
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/ir/test_dtype.py

@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F841
+"""Test data type related API"""
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import DataType
+
+
+@pytest.mark.parametrize(
+    "dtype_str, expected_size",
+    [
+        ("float32", 4),
+        ("float32x4", 16),
+        ("float8_e5m2x4", 4),
+        ("float6_e2m3fnx4", 3),
+        ("float4_e2m1fnx4", 2),
+        ("uint8", 1),
+    ],
+)
+def test_dtype_itemsize(dtype_str, expected_size):
+    dtype = DataType(dtype_str)
+    assert dtype.itemsize == expected_size
+
+
+@pytest.mark.parametrize("dtype_str", ["int32xvscalex4"])
+def test_dtype_itemmize_error(dtype_str):
+    with pytest.raises(ValueError):
+        size = DataType(dtype_str).itemsize
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/ir/test_ir_attrs.py

@@ -0,0 +1,62 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F841
+import pytest
+import tvm_ffi
+
+import tvm
+
+
+def test_dict_attrs():
+    dattr = tvm.ir.make_node("ir.DictAttrs", x=1, y=10, name="xyz", padding=(0, 0))
+    assert dattr.x == 1
+    datrr = tvm.ir.load_json(tvm.ir.save_json(dattr))
+    assert dattr.name == "xyz"
+    assert isinstance(dattr, tvm.ir.DictAttrs)
+    assert "name" in dattr
+    assert dattr["x"] == 1
+    assert len(dattr) == 4
+    assert len([x for x in dattr.keys()]) == 4
+    assert len(dattr.items()) == 4
+
+
+def test_attrs_equal():
+    dattr0 = tvm.ir.make_node("ir.DictAttrs", x=1, y=[10, 20])
+    dattr1 = tvm.ir.make_node("ir.DictAttrs", y=[10, 20], x=1)
+    dattr2 = tvm.ir.make_node("ir.DictAttrs", x=1, y=None)
+    tvm.ir.assert_structural_equal(dattr0, dattr1)
+    assert not tvm_ffi.structural_equal(dattr0, dattr2)
+    assert not tvm_ffi.structural_equal({"x": 1}, tvm.runtime.convert(1))
+    assert not tvm_ffi.structural_equal([1, 2], tvm.runtime.convert(1))
+
+
+def test_assert_structural_equal_reports_mismatch():
+    dattr0 = tvm.ir.make_node("ir.DictAttrs", x=1, y=[10, 20])
+    dattr1 = tvm.ir.make_node("ir.DictAttrs", x=1, y=[10, 30])
+
+    with pytest.raises(ValueError) as err:
+        tvm.ir.assert_structural_equal(dattr0, dattr1)
+
+    message = str(err.value)
+    assert "StructuralEqual check failed" in message
+    assert "caused by lhs at" in message
+    assert "and rhs at" in message
+
+
+if __name__ == "__main__":
+    test_dict_attrs()
+    test_attrs_equal()

tests/python/ir/test_ir_container.py

@@ -0,0 +1,113 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E712, F401
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+
+
+def test_array():
+    a = tvm.runtime.convert([1, 2, 3])
+    assert len(a) == 3
+    assert a[-1] == 3
+    a_slice = a[-3:-1]
+    assert (a_slice[0], a_slice[1]) == (1, 2)
+
+
+def test_array_save_load_json():
+    a = tvm.runtime.convert([1, 2, 3.5, True])
+    json_str = tvm.ir.save_json(a)
+    a_loaded = tvm.ir.load_json(json_str)
+    assert a_loaded[1] == 2
+    assert a_loaded[2] == 3.5
+    assert a_loaded[3] == True
+    assert isinstance(a_loaded[3], bool)
+
+
+def test_dir_array():
+    a = tvm.runtime.convert([1, 2, 3])
+    assert dir(a)
+
+
+def test_map():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    amap = tvm.runtime.convert({a: 2, b: 3})
+    assert a in amap
+    assert len(amap) == 2
+    dd = dict(amap.items())
+    assert a in dd
+    assert b in dd
+    assert a + 1 not in amap
+    assert {x for x in amap} == {a, b}
+    assert set(amap.keys()) == {a, b}
+    assert set(amap.values()) == {2, 3}
+
+
+def test_str_map():
+    amap = tvm.runtime.convert({"a": 2, "b": 3})
+    assert "a" in amap
+    assert len(amap) == 2
+    dd = dict(amap.items())
+    assert amap["a"] == 2
+    assert "a" in dd
+    assert "b" in dd
+
+
+def test_map_save_load_json():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    amap = tvm.runtime.convert({a: 2, b: 3})
+    json_str = tvm.ir.save_json(amap)
+    amap = tvm.ir.load_json(json_str)
+    assert len(amap) == 2
+    dd = {kv[0].name: kv[1] for kv in amap.items()}
+    assert dd == {"a": 2, "b": 3}
+
+
+def test_dir_map():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    amap = tvm.runtime.convert({a: 2, b: 3})
+    assert dir(amap)
+
+
+def test_getattr_map():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    amap = tvm.runtime.convert({a: 2, b: 3})
+    assert isinstance(amap, tvm_ffi.Map)
+
+
+def test_in_container():
+    arr = tvm.runtime.convert(["a", "b", "c"])
+    assert "a" in arr
+    assert "d" not in arr
+
+
+def test_tensor_container():
+    x = tvm.runtime.tensor([1, 2, 3])
+    arr = tvm.runtime.convert([x, x])
+    assert arr[0].same_as(x)
+    assert arr[1].same_as(x)
+    assert isinstance(arr[0], tvm.runtime.Tensor)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/ir/test_ir_type.py

@@ -0,0 +1,66 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401, F821
+"""Test type nodes in the IR"""
+
+import pytest
+
+import tvm
+from tvm.script import tirx as T
+
+
+def check_json_roundtrip(node):
+    json_str = tvm.ir.save_json(node)
+    back = tvm.ir.load_json(json_str)
+    tvm.ir.assert_structural_equal(back, node, map_free_vars=True)
+
+
+def test_prim_type():
+    x = tvm.ir.PrimType("int32")
+    assert isinstance(x, tvm.ir.PrimType)
+    assert x.dtype == "int32"
+    with pytest.raises(TypeError, match="PointerType::VoidPointerTy"):
+        tvm.ir.PrimType("handle")
+
+
+def test_func_type():
+    arg_types = tvm.runtime.convert([])
+    ret_type = tvm.ir.PrimType("float32")
+    tf = tvm.ir.FuncType(arg_types, ret_type)
+    assert tf.arg_types == arg_types
+    assert tf.ret_type == ret_type
+    assert tf.span is None
+    # TODO make sure we can set span
+    str(tf)
+    check_json_roundtrip(tf)
+
+
+def test_tuple_type():
+    tf = tvm.ir.FuncType([], tvm.ir.TupleType([]))
+    tt = tvm.ir.PrimType("float32")
+    fields = tvm.runtime.convert([tf, tt])
+
+    tup_ty = tvm.ir.TupleType(fields)
+    assert tup_ty.fields == fields
+    str(tup_ty)
+    check_json_roundtrip(tup_ty)
+
+
+if __name__ == "__main__":
+    test_tensor_type_bad_constructor()
+    test_func_type()
+    test_tuple_type()

tests/python/ir/test_node_reflection.py

@@ -0,0 +1,186 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E712, F401, F841
+import json
+import sys
+
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm import te
+
+
+def test_const_saveload_json():
+    # save load json
+    x = tvm.tirx.const(1, "int32")
+    y = tvm.tirx.const(10, "int32")
+    z = x + y
+    z = z + z
+    json_str = tvm.ir.save_json(z)
+    zz = tvm.ir.load_json(json_str)
+    tvm.ir.assert_structural_equal(zz, z, map_free_vars=True)
+
+
+def test_save_json_metadata_version():
+    obj = tvm.runtime.convert([1, 2])
+    json_str = tvm.ir.save_json(obj)
+    assert json.loads(json_str)["metadata"]["tvm_version"] == tvm.__version__
+    assert list(tvm.ir.load_json(json_str)) == [1, 2]
+
+
+def _test_infinity_value(value, dtype):
+    x = tvm.tirx.const(value, dtype)
+    json_str = tvm.ir.save_json(x)
+    tvm.ir.assert_structural_equal(x, tvm.ir.load_json(json_str))
+
+
+def test_infinity_value():
+    _test_infinity_value(float("inf"), "float64")
+    _test_infinity_value(float("-inf"), "float64")
+    _test_infinity_value(float("inf"), "float32")
+    _test_infinity_value(float("-inf"), "float32")
+
+
+def _test_minmax_value(value):
+    json_str = tvm.ir.save_json(value)
+    tvm.ir.assert_structural_equal(value, tvm.ir.load_json(json_str))
+
+
+def test_minmax_value():
+    _test_minmax_value(tvm.tirx.min_value("float32"))
+    _test_minmax_value(tvm.tirx.max_value("float32"))
+
+
+def test_make_smap():
+    # save load json
+    x = tvm.tirx.const(1, "int32")
+    y = tvm.tirx.const(10, "int32")
+    z = tvm.tirx.Add(x, y)
+    smap = tvm.runtime.convert({"z": z, "x": x})
+    json_str = tvm.ir.save_json(tvm.runtime.convert([smap]))
+    arr = tvm.ir.load_json(json_str)
+    assert len(arr) == 1
+    assert arr[0]["z"].a == arr[0]["x"]
+    tvm.ir.assert_structural_equal(arr, [smap], map_free_vars=True)
+
+
+def test_make_node():
+    x = tvm.ir.make_node("ir.IntImm", ty=tvm.ir.PrimType("int32"), value=10, span=None)
+    assert isinstance(x, tvm.tirx.IntImm)
+    assert x.value == 10
+    A = te.placeholder((10,), name="A")
+    AA = tvm.ir.make_node(
+        "te.Tensor", shape=A.shape, dtype=A.dtype, op=A.op, value_index=A.value_index
+    )
+    assert AA.op == A.op
+    assert AA.value_index == A.value_index
+
+    y = tvm.ir.make_node(
+        "ir.IntImm", ty=tvm.ir.PrimType(tvm_ffi.core.String("int32")), value=10, span=None
+    )
+    assert isinstance(y, tvm.tirx.IntImm)
+    assert y.value == 10
+
+
+def test_make_sum():
+    A = te.placeholder((2, 10), name="A")
+    k = te.reduce_axis((0, 10), "k")
+    B = te.compute((2,), lambda i: te.sum(A[i, k], axis=k), name="B")
+    json_str = tvm.ir.save_json(B)
+    BB = tvm.ir.load_json(json_str)
+    assert B.op.body[0].combiner is not None
+    assert BB.op.body[0].combiner is not None
+
+
+def test_string():
+    # non printable str, need to store by b64
+    s1 = tvm_ffi.core.String("xy\x01z")
+    s2 = tvm.ir.load_json(tvm.ir.save_json(s1))
+    tvm.ir.assert_structural_equal(s1, s2)
+
+    # printable str, need to store by repr_str
+    s1 = tvm_ffi.core.String("xyz")
+    s2 = tvm.ir.load_json(tvm.ir.save_json(s1))
+    tvm.ir.assert_structural_equal(s1, s2)
+
+
+def test_pass_config():
+    cfg = tvm.transform.PassContext(
+        opt_level=1,
+        config={
+            "tirx.UnrollLoop": {
+                "auto_max_step": 10,
+            }
+        },
+    )
+    cfg.opt_level == 1
+
+    assert cfg.config["tirx.UnrollLoop"].auto_max_step == 10
+    # default option
+    assert cfg.config["tirx.UnrollLoop"].explicit_unroll == True
+
+    # schema checking for specific config key
+    with pytest.raises(TypeError):
+        cfg = tvm.transform.PassContext(config={"tirx.UnrollLoop": {"invalid": 1}})
+
+    # schema check for un-registered config
+    with pytest.raises(AttributeError):
+        cfg = tvm.transform.PassContext(config={"inavlid-opt": True})
+
+    # schema check for wrong type
+    with pytest.raises(AttributeError):
+        cfg = tvm.transform.PassContext(config={"tirx.UnrollLoop": 1})
+
+
+def test_dict():
+    x = tvm.tirx.const(1)  # a class that has Python-defined methods
+    # instances should see the full class dict
+    assert set(dir(x.__class__)) <= set(dir(x))
+
+
+def test_tensor():
+    dev = tvm.cpu(0)
+    tvm_arr = tvm.runtime.tensor(np.random.rand(4), device=dev)
+    tvm_arr2 = tvm.ir.load_json(tvm.ir.save_json(tvm_arr))
+    tvm.ir.assert_structural_equal(tvm_arr, tvm_arr2)
+    np.testing.assert_array_equal(tvm_arr.numpy(), tvm_arr2.numpy())
+
+
+def test_tensor_dict():
+    dev = tvm.cpu(0)
+    m1 = {
+        "key1": tvm.runtime.tensor(np.random.rand(4), device=dev),
+        "key2": tvm.runtime.tensor(np.random.rand(4), device=dev),
+    }
+    m2 = tvm.ir.load_json(tvm.ir.save_json(m1))
+    tvm.ir.assert_structural_equal(m1, m2)
+
+
+def test_free_var_equal():
+    x = tvm.tirx.Var("x", dtype="int32")
+    y = tvm.tirx.Var("y", dtype="int32")
+    z = tvm.tirx.Var("z", dtype="int32")
+    v1 = x + y
+    v1 = y + z
+    tvm.ir.assert_structural_equal(x, z, map_free_vars=True)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/ir/test_pass_instrument.py

@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E741
+"""Instrument test cases."""
+
+import tvm
+from tvm import relax
+from tvm.ir.instrument import PrintAfterAll, PrintBeforeAll
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tirx as T
+
+# pylint: disable=invalid-name,missing-function-docstring,no-value-for-parameter
+
+
+def test_tir_print_all_passes(capsys):
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle) -> None:
+        A = T.match_buffer(a, (128, 128, 128, 128))
+        B = T.match_buffer(b, (128, 128, 128, 128))
+        for i, j, k, l in T.grid(128, 128, 128, 128):
+            with T.sblock("B"):
+                vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+                B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
+
+    with tvm.transform.PassContext(opt_level=3, instruments=[PrintBeforeAll(), PrintAfterAll()]):
+        tvm.compile(func)
+    all_passes_output = capsys.readouterr().out
+    assert "Before Running Pass:" in all_passes_output
+    assert "After Running Pass:" in all_passes_output
+    assert 'name="tirx.' in all_passes_output
+
+
+def test_relax_print_all_passes(capsys):
+    @I.ir_module(s_tir=True)
+    class Module:
+        @R.function
+        def func(x: R.Tensor((16,), "float32"), y: R.Tensor((16,), "float32")):
+            z = R.add(x, y)
+            y = z
+            return y
+
+    pipeline = relax.get_pipeline("default_build")
+    with tvm.transform.PassContext(opt_level=3, instruments=[PrintBeforeAll(), PrintAfterAll()]):
+        pipeline(Module)
+    all_passes_output = capsys.readouterr().out
+    assert "Before Running Pass:" in all_passes_output
+    assert "After Running Pass:" in all_passes_output
+    assert 'name="_pipeline"' in all_passes_output

tests/python/ir/test_roundtrip_runtime_module.py

@@ -0,0 +1,37 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401
+"""Test roundtrip of runtime modules"""
+# pylint: disable=missing-docstring
+
+import pytest
+
+import tvm
+import tvm.testing
+
+
+def test_csource_module():
+    mod = tvm.runtime._ffi_api.CSourceModuleCreate("", "cc", [], [])
+    assert mod.kind == "c"
+    assert mod.is_binary_serializable()
+    new_mod = tvm.ir.load_json(tvm.ir.save_json(mod))
+    assert new_mod.kind == "c"
+    assert new_mod.is_binary_serializable()
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/ir/test_unique_name_supply.py

@@ -0,0 +1,62 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import tvm
+import tvm.testing
+from tvm import relax as rx
+from tvm.ir.supply import UniqueNameSupply
+
+
+def _empty_relax_func():
+    return rx.Function([], rx.Tuple([]))
+
+
+def test_fresh_name_empty_string():
+    """Empty name should produce a valid variable name, not an empty string."""
+    ns = UniqueNameSupply("")
+    name = ns.fresh_name("", add_prefix=False)
+    assert name == "v"
+    name2 = ns.fresh_name("", add_prefix=False)
+    assert name2 == "v_1"
+
+
+def test_fresh_name_empty_string_with_prefix():
+    """Empty name with prefix should produce a valid variable name."""
+    ns = UniqueNameSupply("prefix")
+    name = ns.fresh_name("", add_prefix=True)
+    assert name == "prefix_v"
+    name2 = ns.fresh_name("", add_prefix=True)
+    assert name2 == "prefix_v_1"
+
+
+def test_ir_module_from_expr_freshens_main_collision():
+    main_gv = tvm.ir.GlobalVar("main")
+    mod = tvm.IRModule.from_expr(_empty_relax_func(), {main_gv: _empty_relax_func()})
+
+    assert sorted(gvar.name_hint for gvar in mod.get_global_vars()) == ["main", "main_1"]
+
+
+def test_ir_module_from_expr_reuses_existing_global_symbol():
+    foo_gv = tvm.ir.GlobalVar("foo")
+    func = _empty_relax_func().with_attr("global_symbol", "foo")
+    mod = tvm.IRModule.from_expr(func, {foo_gv: _empty_relax_func()})
+
+    assert mod.get_global_var("foo").same_as(foo_gv)
+    assert [gvar.name_hint for gvar in mod.get_global_vars()] == ["foo"]
+
+
+if __name__ == "__main__":
+    tvm.testing.main()