chore: import upstream snapshot with attribution

tests/python/tirx-base/test_tir_base.py

@@ -0,0 +1,197 @@
+# 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: E711, F841
+import itertools
+
+import numpy as np
+import pytest
+
+import tvm
+from tvm import tirx
+from tvm.ir.transform import PassContext
+from tvm.script import tirx as T
+
+
+def build_tir_func(func):
+    func = func.with_attr("global_symbol", "main")
+    pass_ctx = PassContext.current()
+    if pass_ctx.config.get("tirx.noalias", True):
+        func = func.with_attr("tirx.noalias", True)
+    mod = tvm.IRModule({"main": func})
+    func = tvm.compile(mod)
+    return func
+
+
+def test_scalar_add():
+    # All these types should be interchangeable with each other
+    # E.g. float16 + float32 upconverts the float16 --> float32
+    # Meanwhile if an int or float or together the int will be
+    # cast to the float type.
+    lhs_types = ["float32", "float16", "int32", "int64"]
+    rhs_types = ["float32", "float16"]
+    for lhs_type, rhs_type in itertools.product(lhs_types, rhs_types):
+        # Input vars should be float32, we will cast to test for upcasting between them
+        lhs_input = tirx.Var("lhs", "float32")
+        rhs_input = tirx.Var("rhs", "float32")
+        lhs = tirx.Cast(lhs_type, lhs_input)
+        rhs = tirx.Cast(rhs_type, rhs_input)
+        output = lhs + rhs
+        output = tirx.ret(output)
+        output = tirx.Evaluate(output)
+        func = tirx.PrimFunc([lhs_input, rhs_input], output)
+        func = build_tir_func(func)
+        out = func(1.0, 2.0)
+        assert out == 3.0
+
+
+def assignment_helper(store_dtype, value_dtype):
+    store = tirx.Var("store", dtype=store_dtype)
+    value = tirx.Var("value", dtype=value_dtype)
+    tirx.Let(store, value, body=store)
+
+
+def test_fail_implicit_downcasts_same_type():
+    # These lists should be sorted
+    bits = [8, 16, 32, 64]
+    for type in ["float", "int", "uint"]:
+        for i in range(len(bits) - 1):
+            with pytest.raises(RuntimeError):
+                assignment_helper(
+                    store_dtype=f"{type}{bits[i]}", value_dtype=f"{type}{bits[i + 1]}"
+                )
+
+
+def test_cast_between_types():
+    # We should only be able to assign values with the same types
+    bits = [16, 32]
+    types = ["float", "int", "uint"]
+    for store_type, store_bits, value_type, value_bits in itertools.product(
+        types, bits, types, bits
+    ):
+        store_dtype = f"{store_type}{store_bits}"
+        value_dtype = f"{value_type}{value_bits}"
+        if store_dtype == value_dtype:
+            assignment_helper(store_dtype, value_dtype)
+        else:
+            # TODO: we might want to allow casts between uint and int types
+            with pytest.raises(RuntimeError):
+                assignment_helper(store_dtype, value_dtype)
+
+
+def test_ret_const():
+    a = tirx.const(0)
+    b = tirx.ret(a)
+    b = tirx.Evaluate(b)
+    func = tirx.PrimFunc([], b)
+    func = build_tir_func(func)
+    out = func()
+    assert out == 0
+
+
+def test_control_flow_jump():
+    @T.prim_func(s_tir=True)
+    def func(a: T.float32, b: T.float32):
+        if True:
+            T.evaluate(T.ret(a))
+        T.evaluate(T.ret(b))
+
+    func = build_tir_func(func)
+    out = func(1.0, 2.0)
+    assert out == 1.0
+
+
+def test_break_loop():
+    @T.prim_func(s_tir=True)
+    def func(In: T.Buffer((2,), "int32"), Out: T.Buffer((2,), "int32")):
+        Out[0] = 0
+        Out[1] = 1
+        for i in range(10):
+            for j in range(10):
+                if i * 10 + j == In[0]:
+                    Out[0] = i + j
+                    break
+            if Out[0] > 0:
+                break
+        while Out[1] > 0:
+            Out[1] = Out[1] + 1
+            if Out[1] > In[1]:
+                break
+
+    func = build_tir_func(func)
+    a = np.asarray([49, 8], "int32")
+    b = np.zeros([2], "int32")
+    if not hasattr(b, "__dlpack__"):
+        return
+    func(a, b)
+    assert b[0] == 13
+    assert b[1] == 9
+
+
+def test_continue_loop():
+    @T.prim_func(s_tir=True)
+    def func(Out: T.Buffer((2,), "int32")):
+        T.func_attr({"global_symbol": "main"})
+        Out[0] = 0
+        Out[1] = 0
+        for i in range(10):
+            for j in range(10):
+                if (i * 10 + j) % 3 != 0:
+                    continue
+                Out[0] = Out[0] + 1
+        k = T.decl_buffer([], "int32")
+        k[()] = 0
+        while k[()] < Out[0]:
+            k[()] = k[()] + 1
+            if k[()] % 6 == 0:
+                Out[1] = Out[1] + 1
+                continue
+
+    func = build_tir_func(func)
+    b = np.zeros([2], "int32")
+    if not hasattr(b, "__dlpack__"):
+        return
+    func(b)
+    assert b[0] == 34
+    assert b[1] == 5
+
+
+def test_exception():
+    with pytest.raises(TypeError):
+        x = tirx.Var(name=1, dtype="int")
+
+
+def test_eq_ops():
+    # NOTE: the `== None` / `!= None` below are intentional and must NOT be
+    # rewritten as `is None` / `is not None`. This test exercises the overloaded
+    # `__eq__` / `__ne__` operators on `IntImm` / `StringImm`; the `is` operators
+    # bypass those overloads and would defeat the test.
+    a = tirx.IntImm("int8", 1)
+    with pytest.raises(ValueError):
+        assert a != None
+    with pytest.raises(ValueError):
+        assert not a == None
+    b = tirx.StringImm("abc")
+    assert b != None
+    assert not b == None
+
+
+if __name__ == "__main__":
+    test_scalar_add()
+    test_ret_const()
+    test_control_flow_jump()
+    test_exception()
+    test_eq_ops()

tests/python/tirx-base/test_tir_buffer.py

@@ -0,0 +1,219 @@
+# 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, F401, F841
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.tirx import Buffer
+
+
+def test_buffer():
+    m = tvm.tirx.Var("m", "int32")
+    n = tvm.tirx.Var("n", "int32")
+    l = tvm.tirx.Var("l", "int32")
+    Ab = tvm.tirx.decl_buffer((m, n), "float32")
+    Bb = tvm.tirx.decl_buffer((n, l), "float32")
+
+    assert isinstance(Ab, tvm.tirx.Buffer)
+    assert Ab.dtype == tvm.ir.PrimType("float32")
+    assert tuple(Ab.shape) == (m, n)
+
+
+def test_buffer_access_ptr():
+    m = tvm.tirx.Var("m", "int32")
+    n = tvm.tirx.Var("n", "int32")
+    Ab = tvm.tirx.decl_buffer((m, n), "float32", strides=[n + 1, 1])
+    aptr = Ab.access_ptr("rw")
+    assert isinstance(aptr.ty, tvm.ir.PointerType)
+    assert aptr.ty.element_type == tvm.ir.PrimType("void")
+    tvm.ir.assert_structural_equal(aptr.args[3], Ab.strides[0] * m)
+    assert aptr.args[0].ty == Ab.dtype
+    assert aptr.args[4].value == Buffer.READ | Buffer.WRITE
+    typed_ptr = Ab.access_ptr("r", ptr_type="uint8")
+    assert typed_ptr.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint8"))
+    shared = tvm.tirx.decl_buffer((m, n), "float32", scope="shared")
+    assert shared.access_ptr("r").ty == tvm.ir.PointerType(tvm.ir.PrimType("void"), "shared")
+    assert shared.access_ptr("r", ptr_type="uint8").ty == tvm.ir.PointerType(
+        tvm.ir.PrimType("uint8"), "shared"
+    )
+    aptr = Ab.access_ptr("w")
+    assert aptr.args[4].value == Buffer.WRITE
+
+
+def test_buffer_access_ptr_offset():
+    m = tvm.tirx.Var("m", "int32")
+    n = tvm.tirx.Var("n", "int32")
+    Ab = tvm.tirx.decl_buffer((m, n), "float32")
+    aptr = Ab.access_ptr("rw", offset=100)
+    tvm.testing.assert_prim_expr_equal(aptr.args[2], 100)
+    assert aptr.args[4].value == Buffer.READ | Buffer.WRITE
+    v = tvm.tirx.Var("int32", "int32")
+    aptr = Ab.access_ptr("rw", offset=100 + 100 + v)
+    tvm.testing.assert_prim_expr_equal(aptr.args[2], 200 + v)
+    assert aptr.args[4].value == Buffer.READ | Buffer.WRITE
+    aptr = Ab.access_ptr("rw", offset=tvm.tirx.call_extern("int32", "test_call", 100 + 100 + v))
+    tvm.testing.assert_prim_expr_equal(
+        aptr.args[2], tvm.tirx.call_extern("int32", "test_call", 200 + v)
+    )
+    assert aptr.args[4].value == Buffer.READ | Buffer.WRITE
+
+
+def test_buffer_access_ptr_extent():
+    m = tvm.tirx.Var("m", "int32")
+    n = tvm.tirx.Var("n", "int32")
+    Ab = tvm.tirx.decl_buffer((m, n), "float32")
+    aptr = Ab.access_ptr("rw")
+    tvm.ir.assert_structural_equal(aptr.args[3], m * n)
+    aptr = Ab.access_ptr("rw", offset=100)
+    tvm.ir.assert_structural_equal(aptr.args[3], m * n - 100)
+    Ab = tvm.tirx.decl_buffer((m, n), "float32", strides=[n + 1, 1])
+    aptr = Ab.access_ptr("rw", offset=100)
+    tvm.ir.assert_structural_equal(aptr.args[3], Ab.strides[0] * m - 100)
+
+    # Test extent from input params
+    aptr = Ab.access_ptr("rw", extent=200)
+    tvm.ir.assert_structural_equal(aptr.args[3], T.int32(200))
+    aptr = Ab.access_ptr("rw", offset=100, extent=100)
+    tvm.ir.assert_structural_equal(aptr.args[3], T.int32(100))
+
+
+def test_buffer_vload():
+    m = tvm.tirx.Var("m", "int32")
+    n = tvm.tirx.Var("n", "int32")
+    Ab = tvm.tirx.decl_buffer((m, n), "float32", elem_offset=100)
+    load = Ab.vload([2, 3])
+    tvm.ir.assert_structural_equal(load.indices, [T.int32(2), T.int32(3)])
+
+
+def test_buffer_offset_of():
+    m = tvm.tirx.Var("m", "int32")
+    n = tvm.tirx.Var("n", "int32")
+    Ab = tvm.tirx.decl_buffer((m, n), "float32", elem_offset=100)
+    offset = Ab.offset_of([2, 3])
+    tvm.ir.assert_structural_equal(offset, [n * 2 + 103])
+
+
+def test_buffer_index_merge_mult_mod():
+    m = tvm.tirx.Var("m", "int32")
+    n = tvm.tirx.Var("n", "int32")
+    s = tvm.tirx.Var("s", "int32")
+    k0 = tvm.tirx.Var("k0", "int32")
+    k1 = tvm.tirx.Var("k1", "int32")
+    A = tvm.tirx.decl_buffer((m, n), "float32")
+    A_stride = tvm.tirx.decl_buffer((m, n), "float32", strides=(s, 1))
+
+    def assert_simplified_equal(index_simplified, index_direct):
+        (
+            tvm.ir.assert_structural_equal(index_simplified, index_direct),
+            f"index_simplified={index_simplified}, index_direct={index_direct}",
+        )
+
+    idxd = tvm.tirx.indexdiv
+    idxm = tvm.tirx.indexmod
+
+    # Test Case1
+    index_simplified = A_stride.offset_of(
+        (idxd(idxm(k0, k1), s), idxm(idxm(k0, k1), s) + idxd(k0, k1) * k1)
+    )
+    index_direct = A_stride.offset_of((0, k0))
+    assert_simplified_equal(index_simplified, index_direct)
+
+    # Test Case2
+    index_simplified = A.offset_of(
+        (idxd(idxm(k0, idxd(k1, s)), n), idxm(idxm(k0, idxd(k1, s)), n) + idxm(k0, k1))
+    )
+    index_direct = A.offset_of((0, idxm(k0, idxd(k1, s)) + idxm(k0, k1)))
+    assert_simplified_equal(index_simplified, index_direct)
+    # Test Case3
+    index_simplified = A.offset_of(
+        (
+            idxd((idxd(k0, idxd(k1, s)) * idxd(k1, s)), n) + idxd(idxm(k0, idxd(k1, s)), n),
+            idxm((idxd(k0, idxd(k1, s)) * idxd(k1, s)), n) + idxm(idxm(k0, idxd(k1, s)), n),
+        )
+    )
+    index_direct = A.offset_of((0, k0))
+    assert_simplified_equal(index_simplified, index_direct)
+    # Test Case4 (not able to simplify)
+    index_simplified = A.offset_of(
+        (idxd(idxm(k0, idxd(k1, s)), n), idxm(idxm(k0, idxd(k1, n)), n) + idxm(k0, k1))
+    )
+    index_direct = A.offset_of(
+        (0, idxd(idxm(k0, idxd(k1, s)), n) * n + (idxm(idxm(k0, idxd(k1, n)), n) + idxm(k0, k1)))
+    )
+    assert_simplified_equal(index_simplified, index_direct)
+
+    # Test Case5
+    B = tvm.tirx.decl_buffer((1, 14, 14, 1024))
+    i = tvm.tirx.Var("i", "int32")
+    j = tvm.tirx.Var("j", "int32")
+    k = tvm.tirx.Var("k", "int32")
+
+    index_simplified1 = B.offset_of(
+        (
+            idxd(idxd(idxd((i * 50176 + j * 28672 + k), 1024), 14), 14),
+            idxm(idxd(idxd((i * 50176 + j * 28672 + k), 1024), 14), 14),
+            idxm(idxd((i * 50176 + j * 28672 + k), 1024), 14),
+            idxm((i * 50176 + j * 28672 + k), 1024),
+        )
+    )
+    index_simplified2 = B.offset_of(
+        (
+            idxd(idxd(i * 49 + j * 28 + idxd(k, 1024), 14), 14),
+            idxm(idxd(i * 49 + j * 28 + idxd(k, 1024), 14), 14),
+            idxm(i * 7 + idxd(k, 1024), 14),
+            idxm(k, 1024),
+        )
+    )
+    index_direct = B.offset_of((0, 0, 0, (i * 50176 + j * 28672 + k)))
+    assert_simplified_equal(index_simplified1, index_direct)
+    assert_simplified_equal(index_simplified2, index_direct)
+
+
+def test_buffer_flatten():
+    """A buffer should flatten to a 1-d shape"""
+    buf = tvm.tirx.decl_buffer([16, 32])
+    flat = buf.get_flattened_buffer()
+    assert buf.data.same_as(flat.data)
+    tvm.ir.assert_structural_equal(flat.shape, [T.int32(16 * 32)])
+
+
+def test_buffer_flatten_preserves_identity():
+    """Flattening a 1-d buffer should return the original"""
+    buf = tvm.tirx.decl_buffer([16])
+    flat = buf.get_flattened_buffer()
+    assert buf.same_as(flat)
+
+
+def test_buffer_flatten_uses_axis_separators():
+    """Flattening to N-d physical buffers uses the axis separators"""
+    buf = tvm.tirx.decl_buffer([4, 16, 32], axis_separators=[2])
+    flat = buf.get_flattened_buffer()
+    tvm.ir.assert_structural_equal(flat.axis_separators, [T.int32(1)])
+    tvm.ir.assert_structural_equal(flat.shape, [T.int32(4 * 16), T.int32(32)])
+
+
+def test_invalid_axis_separators_raises_exception():
+    with pytest.raises(ValueError):
+        tvm.tirx.decl_buffer([1], axis_separators=[1, 2])
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-base/test_tir_constructor.py

@@ -0,0 +1,285 @@
+# 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 pytest
+import tvm_ffi
+
+import tvm
+from tvm import te, topi
+from tvm.tirx.analysis import expr_deep_equal
+from tvm.tirx.expr_functor import ExprMutator
+
+
+class ReplaceVar(ExprMutator):
+    def __init__(self, old_var, new_var):
+        super().__init__()
+        self.old_var = old_var
+        self.new_var = new_var
+
+    def visit_var_(self, op):
+        if op.same_as(self.old_var):
+            return self.new_var
+        return op
+
+
+def test_expr_constructor():
+    x = tvm.tirx.Var("xx", "float32")
+    assert isinstance(x, tvm.tirx.Var)
+    assert x.name == "xx"
+
+    x = tvm.tirx.Reduce(None, [1], [tvm.tirx.IterVar((0, 1), "x", 2)], None, 0)
+    assert isinstance(x, tvm.tirx.Reduce)
+    assert x.combiner is None
+    assert x.value_index == 0
+
+    x = tvm.tirx.FloatImm("float32", 1.0)
+    assert isinstance(x, tvm.tirx.FloatImm)
+    assert x.value == 1.0
+    assert x.ty == tvm.ir.PrimType("float32")
+
+    x = tvm.tirx.IntImm("int64", 2)
+    assert isinstance(x, tvm.tirx.IntImm)
+    assert x.value == 2
+    assert x.ty == tvm.ir.PrimType("int64")
+
+    x = tvm.tirx.StringImm("xyza")
+    assert isinstance(x, tvm.tirx.StringImm)
+    assert x.value == "xyza"
+
+    x = tvm.tirx.Cast("float32", tvm.tirx.IntImm("uint32", 1))
+    assert isinstance(x, tvm.tirx.Cast)
+    assert x.ty == tvm.ir.PrimType("float32")
+    assert x.value.value == 1
+
+    a = tvm.tirx.const(1.0, dtype="float32")
+    b = tvm.tirx.Var("x", "float32")
+
+    for cls in [
+        tvm.tirx.Add,
+        tvm.tirx.Sub,
+        tvm.tirx.Mul,
+        tvm.tirx.Div,
+        tvm.tirx.Mod,
+        tvm.tirx.Min,
+        tvm.tirx.Max,
+        tvm.tirx.LT,
+        tvm.tirx.LE,
+        tvm.tirx.GT,
+        tvm.tirx.GE,
+    ]:
+        x = cls(a, b)
+        assert isinstance(x, cls)
+        assert x.a == a
+        assert x.b.same_as(b)
+
+    a = tvm.runtime.convert(tvm.tirx.Var("x", "int32") > 1)
+    b = tvm.runtime.convert(tvm.tirx.Var("x", "int32") == 1)
+
+    for cls in [tvm.tirx.And, tvm.tirx.Or]:
+        x = cls(a, b)
+        assert isinstance(x, cls)
+        assert x.a == a
+        assert x.b.same_as(b)
+
+    x = tvm.tirx.Not(a)
+    assert isinstance(x, tvm.tirx.Not)
+    assert x.a == a
+
+    x = tvm.tirx.Select(a, a, b)
+    assert isinstance(x, tvm.tirx.Select)
+    assert x.true_value == a
+    assert x.false_value == b
+    assert x.condition == a
+
+    buffer_var = tvm.tirx.Var("buf", tvm.ir.PointerType(tvm.ir.PrimType("float32")))
+    buffer = tvm.tirx.decl_buffer([16], "float32", data=buffer_var)
+    x = tvm.tirx.BufferLoad(buffer, [1])
+    assert isinstance(x, tvm.tirx.BufferLoad)
+    assert x.ty == tvm.ir.PrimType("float32")
+    assert x.buffer == buffer
+    assert x.buffer.data == buffer_var
+    assert list(x.indices) == [1]
+
+    x = tvm.tirx.Ramp(1, 2, 10)
+    assert isinstance(x, tvm.tirx.Ramp)
+    assert x.base.value == 1
+    assert x.stride.value == 2
+    assert x.lanes == 10
+
+    x = tvm.tirx.Broadcast(a, 10)
+    assert isinstance(x, tvm.tirx.Broadcast)
+    assert x.value == a
+    assert x.lanes == 10
+
+    x = tvm.tirx.Shuffle([a], [0])
+    assert isinstance(x, tvm.tirx.Shuffle)
+    assert x.vectors[0] == a
+    assert x.indices[0].value == 0
+
+    x = tvm.ir.Call("tirx.call_extern", [tvm.tirx.StringImm("xyz"), a], ret_ty="float32")
+    assert isinstance(x, tvm.ir.Call)
+    assert tvm.ir.is_prim_expr(x)
+    assert x.ty == tvm.ir.PrimType("float32")
+    assert x.op.name == "tirx.call_extern"
+    assert x.args[1] == a
+    assert x.attrs is None
+
+    attr_arg = tvm.tirx.Var("attr_arg", "float32")
+    x_with_attrs = tvm.ir.Call(
+        "tirx.call_extern",
+        [tvm.tirx.StringImm("xyz"), attr_arg],
+        attrs={"disable_tma": True},
+        ret_ty="float32",
+    )
+    assert x_with_attrs.attrs["disable_tma"] is True
+    assert not tvm_ffi.structural_equal(x, x_with_attrs)
+    script = tvm.tirx.Evaluate(x_with_attrs).script()
+    assert "attrs" in script
+    assert "disable_tma" in script
+    func = tvm.tirx.PrimFunc([], tvm.tirx.Evaluate(x_with_attrs))
+    assert tvm.script.from_source(func.script()).script() == func.script()
+
+    y = tvm.tirx.Var("y", "float32")
+    mutated = ReplaceVar(attr_arg, y)(x_with_attrs)
+    assert mutated.attrs["disable_tma"] is True
+    assert mutated.args[1].same_as(y)
+
+    x_from_intrin = tvm.tirx.call_intrin(
+        "float32", "tirx.call_extern", tvm.tirx.StringImm("xyz"), attrs={"disable_tma": True}
+    )
+    assert x_from_intrin.attrs["disable_tma"] is True
+    x_with_other_attrs = tvm.ir.Call(
+        "tirx.call_extern",
+        [tvm.tirx.StringImm("xyz"), attr_arg],
+        attrs={"disable_tma": False},
+        ret_ty="float32",
+    )
+    assert not expr_deep_equal(x_with_attrs, x_with_other_attrs)
+
+    cond0 = tvm.tirx.Var("cond0", "bool")
+    cond1 = tvm.tirx.Var("cond1", "bool")
+    inner_if = tvm.ir.Call(
+        "tirx.if_then_else",
+        [cond1, tvm.tirx.IntImm("int32", 1), tvm.tirx.IntImm("int32", 0)],
+        ret_ty="int32",
+    )
+    outer_if = tvm.ir.Call(
+        "tirx.if_then_else",
+        [cond0, inner_if, tvm.tirx.IntImm("int32", 0)],
+        attrs={"keep": True},
+        ret_ty="int32",
+    )
+    simplified = tvm.tirx.transform.StmtSimplify()(
+        tvm.IRModule({"main": tvm.tirx.PrimFunc([], tvm.tirx.Evaluate(outer_if))})
+    )["main"].body.value
+    assert simplified.attrs["keep"] is True
+
+    v = tvm.tirx.Var("aa", "int32")
+    x = tvm.tirx.Let(v, 1, v)
+    assert x.var == v
+    assert x.value.value == 1
+    assert x.body == v
+
+
+def test_stmt_constructor():
+    v = tvm.tirx.Var("aa", "int32")
+    nop = tvm.tirx.Evaluate(1)
+    x = tvm.tirx.Bind(v, 1)
+    assert isinstance(x, tvm.tirx.Bind)
+    assert x.var == v
+    assert x.value.value == 1
+
+    x = tvm.tirx.AttrStmt(v == 1, "xx", 1, tvm.tirx.Evaluate(1))
+    assert isinstance(x, tvm.tirx.AttrStmt)
+    assert x.value.value == 1
+
+    x = tvm.tirx.AssertStmt(
+        tvm.tirx.const(1, "bool"),
+        tvm.tirx.StringImm("RuntimeError"),
+        [tvm.tirx.StringImm("hellow")],
+    )
+    assert isinstance(x, tvm.tirx.AssertStmt)
+    assert x.error_kind.value == "RuntimeError"
+    assert len(x.message_parts) == 1
+    assert x.message_parts[0].value == "hellow"
+
+    x = tvm.tirx.For(tvm.tirx.Var("x", "int32"), 0, 10, tvm.tirx.ForKind.SERIAL, nop)
+    assert isinstance(x, tvm.tirx.For)
+    assert x.min.value == 0
+    assert x.extent.value == 10
+    assert x.body == nop
+
+    buffer_var = tvm.tirx.Var("buf", tvm.ir.PointerType(tvm.ir.PrimType("bool")))
+    buffer = tvm.tirx.decl_buffer([16], "bool", data=buffer_var)
+    x = tvm.tirx.BufferStore(buffer, tvm.tirx.IntImm("bool", 1), [10])
+    assert isinstance(x, tvm.tirx.BufferStore)
+    assert x.buffer == buffer
+    assert x.buffer.data == buffer_var
+    assert list(x.indices) == [10]
+    assert x.value.value == 1
+
+    buf = tvm.tirx.decl_buffer([10], "float32")
+    x = tvm.tirx.AllocBuffer(buf)
+    assert isinstance(x, tvm.tirx.AllocBuffer)
+    assert x.buffer == buf
+
+    x = tvm.tirx.AttrStmt(buffer_var, "xyz", 1, nop)
+    assert isinstance(x, tvm.tirx.AttrStmt)
+    assert x.node == buffer_var
+    assert x.attr_key == "xyz"
+    assert x.body == nop
+
+    x = tvm.tirx.IfThenElse(tvm.tirx.const(1, "bool"), tvm.tirx.Evaluate(11), nop)
+    assert isinstance(x, tvm.tirx.IfThenElse)
+    assert x.then_case.value.value == 11
+    assert x.else_case == nop
+
+
+def test_float_constructor_requires_float_dtype():
+    # FloatImm dtype validation raises a builtin ValueError.
+    with pytest.raises(ValueError):
+        tvm.tirx.FloatImm("int32", 1.0)
+
+
+def test_math_unary_constructor_requires_float_dtype():
+    x = tvm.tirx.Var("x", "int32")
+
+    with pytest.raises(TypeError, match=r"tirx\.tan only supports floating-point inputs"):
+        tvm.tirx.tan(x)
+
+    with pytest.raises(TypeError, match=r"tirx\.sin only supports floating-point inputs"):
+        tvm.tirx.sin(x)
+
+    y = tvm.tirx.Var("y", "float32")
+    assert tvm.tirx.tan(y).ty == tvm.ir.PrimType("float32")
+
+
+def test_topi_tan_requires_float_dtype():
+    x = te.placeholder((2, 2), dtype="int32", name="x")
+
+    with pytest.raises(TypeError, match=r"tirx\.tan only supports floating-point inputs"):
+        topi.tan(x)
+
+
+def test_math_unary_constructor_preserves_bfloat16():
+    x = tvm.tirx.Var("x", "bfloat16")
+    y = tvm.tirx.exp(x)
+    assert y.ty == tvm.ir.PrimType("bfloat16")
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-base/test_tir_expr_functor.py

@@ -0,0 +1,829 @@
+# 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 tirx as tir
+from tvm.ir import Call, Op
+from tvm.ir.base import assert_structural_equal
+from tvm.tirx.expr import (
+    EQ,
+    GE,
+    GT,
+    LE,
+    LT,
+    NE,
+    Add,
+    And,
+    Broadcast,
+    BufferLoad,
+    Cast,
+    Div,
+    FloatImm,
+    FloorDiv,
+    FloorMod,
+    IntImm,
+    Let,
+    Max,
+    Min,
+    Mod,
+    Mul,
+    Not,
+    Or,
+    ProducerLoad,
+    Ramp,
+    Reduce,
+    Select,
+    Shuffle,
+    StringImm,
+    Sub,
+    Var,
+)
+from tvm.tirx.expr_functor import ExprMutator, ExprVisitor
+
+# Basic example variables for testing
+n = tir.Var("n", "int32")
+m = tir.Var("m", "int32")
+x = tir.Var("x", "float32")
+y = tir.Var("y", "float32")
+
+
+class BasicVisitor(ExprVisitor):
+    """Default ExprVisitor"""
+
+
+class ASTLog:
+    """Helper class to log AST"""
+
+    def __init__(self) -> None:
+        self.log = []
+        self.indent = "\t"
+        self.level = 0
+
+    def push_scope(self):
+        self.level += 1
+
+    def pop_scope(self):
+        self.level -= 1
+
+    def add(self, s: str):
+        self.log.append(self.indent * self.level + s)
+
+    def __str__(self) -> str:
+        return "\n".join(self.log)
+
+
+class ASTPrinter(ExprVisitor):
+    """Print TIR AST in structured format."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.log = ASTLog()
+
+    def visit_var_(self, op: Var) -> None:
+        self.log.add("Var")
+
+    def visit_buffer_load_(self, op: BufferLoad) -> None:
+        self.log.add("BufferLoad")
+        self.log.push_scope()
+        for idx in op.indices:
+            self.visit_expr(idx)
+        self.log.pop_scope()
+
+    def visit_producer_load_(self, op: ProducerLoad) -> None:
+        self.log.add("ProducerLoad")
+        self.log.push_scope()
+        for idx in op.indices:
+            self.visit_expr(idx)
+        self.log.pop_scope()
+
+    def visit_let_(self, op: Let) -> None:
+        self.log.add("Let")
+        self.log.push_scope()
+        self.visit_expr(op.var)
+        self.visit_expr(op.value)
+        self.visit_expr(op.body)
+        self.log.pop_scope()
+
+    def visit_call_(self, op: Call) -> None:
+        self.log.add("Call")
+        self.log.push_scope()
+        if isinstance(op.op, Op):
+            self.log.add("Op")
+        else:
+            self.visit_expr(op.op)
+        for arg in op.args:
+            self.visit_expr(arg)
+        self.log.pop_scope()
+
+    def visit_add_(self, op: Add) -> None:
+        self.log.add("Add")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_sub_(self, op: Sub) -> None:
+        self.log.add("Sub")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_mul_(self, op: Mul) -> None:
+        self.log.add("Mul")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_div_(self, op: Div) -> None:
+        self.log.add("Div")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_mod_(self, op: Mod) -> None:
+        self.log.add("Mod")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_floordiv_(self, op: FloorDiv) -> None:
+        self.log.add("FloorDiv")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_floormod_(self, op: FloorMod) -> None:
+        self.log.add("FloorMod")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_min_(self, op: Min) -> None:
+        self.log.add("Min")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_max_(self, op: Max) -> None:
+        self.log.add("Max")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_eq_(self, op: EQ) -> None:
+        self.log.add("EQ")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_ne_(self, op: NE) -> None:
+        self.log.add("NE")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_lt_(self, op: LT) -> None:
+        self.log.add("LT")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_le_(self, op: LE) -> None:
+        self.log.add("LE")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_gt_(self, op: GT) -> None:
+        self.log.add("GT")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_ge_(self, op: GE) -> None:
+        self.log.add("GE")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_and_(self, op: And) -> None:
+        self.log.add("And")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_or_(self, op: Or) -> None:
+        self.log.add("Or")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_reduce_(self, op: Reduce) -> None:
+        self.log.add("Reduce")
+        self.log.push_scope()
+        for source in op.source:
+            self.visit_expr(source)
+        for axis in op.axis:
+            self.visit_expr(axis.var)
+        self.visit_expr(op.condition)
+        self.log.pop_scope()
+
+    def visit_cast_(self, op: Cast) -> None:
+        self.log.add("Cast")
+        self.log.push_scope()
+        self.visit_expr(op.value)
+        self.log.pop_scope()
+
+    def visit_not_(self, op: Not) -> None:
+        self.log.add("Not")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.log.pop_scope()
+
+    def visit_select_(self, op: Select) -> None:
+        self.log.add("Select")
+        self.log.push_scope()
+        self.visit_expr(op.condition)
+        self.visit_expr(op.true_value)
+        self.visit_expr(op.false_value)
+        self.log.pop_scope()
+
+    def visit_ramp_(self, op: Ramp) -> None:
+        self.log.add("Ramp")
+        self.log.push_scope()
+        self.visit_expr(op.base)
+        self.visit_expr(op.stride)
+        self.visit_expr(op.lanes)
+        self.log.pop_scope()
+
+    def visit_broadcast_(self, op: Broadcast) -> None:
+        self.log.add("Broadcast")
+        self.log.push_scope()
+        self.visit_expr(op.value)
+        self.visit_expr(op.lanes)
+        self.log.pop_scope()
+
+    def visit_shuffle_(self, op: Shuffle) -> None:
+        self.log.add("Shuffle")
+        self.log.push_scope()
+        for vec in op.vectors:
+            self.visit_expr(vec)
+        for idx in op.indices:
+            self.visit_expr(idx)
+        self.log.pop_scope()
+
+    def visit_int_imm_(self, op: IntImm) -> None:
+        self.log.add("IntImm")
+
+    def visit_float_imm_(self, op: FloatImm) -> None:
+        self.log.add("FloatImm")
+
+    def visit_string_imm_(self, op: StringImm) -> None:
+        self.log.add("StringImm")
+
+
+class BasicMutator(ExprMutator):
+    """Default ExprMutator"""
+
+
+class ASTPostPrinterMutator(ExprMutator):
+    """Print TIR AST in the post order format."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.log = ASTLog()
+
+    def visit_var_(self, op: Var) -> tir.Expr:
+        result = super().visit_var_(op)
+        self.log.add("Var")
+        return result
+
+    def visit_buffer_load_(self, op: BufferLoad) -> tir.Expr:
+        result = super().visit_buffer_load_(op)
+        self.log.add("BufferLoad")
+        return result
+
+    def visit_producer_load_(self, op: ProducerLoad) -> tir.Expr:
+        result = super().visit_producer_load_(op)
+        self.log.add("ProducerLoad")
+        return result
+
+    def visit_let_(self, op: Let) -> tir.Expr:
+        result = super().visit_let_(op)
+        self.log.add("Let")
+        return result
+
+    def visit_call_(self, op: Call) -> tir.Expr:
+        result = super().visit_call_(op)
+        self.log.add("Call")
+        return result
+
+    def visit_add_(self, op: Add) -> tir.Expr:
+        result = super().visit_add_(op)
+        self.log.add("Add")
+        return result
+
+    def visit_sub_(self, op: Sub) -> tir.Expr:
+        result = super().visit_sub_(op)
+        self.log.add("Sub")
+        return result
+
+    def visit_mul_(self, op: Mul) -> tir.Expr:
+        result = super().visit_mul_(op)
+        self.log.add("Mul")
+        return result
+
+    def visit_div_(self, op: Div) -> tir.Expr:
+        result = super().visit_div_(op)
+        self.log.add("Div")
+        return result
+
+    def visit_mod_(self, op: Mod) -> tir.Expr:
+        result = super().visit_mod_(op)
+        self.log.add("Mod")
+        return result
+
+    def visit_floordiv_(self, op: FloorDiv) -> tir.Expr:
+        result = super().visit_floordiv_(op)
+        self.log.add("FloorDiv")
+        return result
+
+    def visit_floormod_(self, op: FloorMod) -> tir.Expr:
+        result = super().visit_floormod_(op)
+        self.log.add("FloorMod")
+        return result
+
+    def visit_min_(self, op: Min) -> tir.Expr:
+        result = super().visit_min_(op)
+        self.log.add("Min")
+        return result
+
+    def visit_max_(self, op: Max) -> tir.Expr:
+        result = super().visit_max_(op)
+        self.log.add("Max")
+        return result
+
+    def visit_eq_(self, op: EQ) -> tir.Expr:
+        result = super().visit_eq_(op)
+        self.log.add("EQ")
+        return result
+
+    def visit_ne_(self, op: NE) -> tir.Expr:
+        result = super().visit_ne_(op)
+        self.log.add("NE")
+        return result
+
+    def visit_lt_(self, op: LT) -> tir.Expr:
+        result = super().visit_lt_(op)
+        self.log.add("LT")
+        return result
+
+    def visit_le_(self, op: LE) -> tir.Expr:
+        result = super().visit_le_(op)
+        self.log.add("LE")
+        return result
+
+    def visit_gt_(self, op: GT) -> tir.Expr:
+        result = super().visit_gt_(op)
+        self.log.add("GT")
+        return result
+
+    def visit_ge_(self, op: GE) -> tir.Expr:
+        result = super().visit_ge_(op)
+        self.log.add("GE")
+        return result
+
+    def visit_and_(self, op: And) -> tir.Expr:
+        result = super().visit_and_(op)
+        self.log.add("And")
+        return result
+
+    def visit_or_(self, op: Or) -> tir.Expr:
+        result = super().visit_or_(op)
+        self.log.add("Or")
+        return result
+
+    def visit_reduce_(self, op: Reduce) -> tir.Expr:
+        result = super().visit_reduce_(op)
+        self.log.add("Reduce")
+        return result
+
+    def visit_cast_(self, op: Cast) -> tir.Expr:
+        result = super().visit_cast_(op)
+        self.log.add("Cast")
+        return result
+
+    def visit_not_(self, op: Not) -> tir.Expr:
+        result = super().visit_not_(op)
+        self.log.add("Not")
+        return result
+
+    def visit_select_(self, op: Select) -> tir.Expr:
+        result = super().visit_select_(op)
+        self.log.add("Select")
+        return result
+
+    def visit_ramp_(self, op: Ramp) -> tir.Expr:
+        result = super().visit_ramp_(op)
+        self.log.add("Ramp")
+        return result
+
+    def visit_broadcast_(self, op: Broadcast) -> tir.Expr:
+        result = super().visit_broadcast_(op)
+        self.log.add("Broadcast")
+        return result
+
+    def visit_shuffle_(self, op: Shuffle) -> tir.Expr:
+        result = super().visit_shuffle_(op)
+        self.log.add("Shuffle")
+        return result
+
+    def visit_int_imm_(self, op: IntImm) -> tir.Expr:
+        result = super().visit_int_imm_(op)
+        self.log.add("IntImm")
+        return result
+
+    def visit_float_imm_(self, op: FloatImm) -> tir.Expr:
+        result = super().visit_float_imm_(op)
+        self.log.add("FloatImm")
+        return result
+
+    def visit_string_imm_(self, op: StringImm) -> tir.Expr:
+        result = super().visit_string_imm_(op)
+        self.log.add("StringImm")
+        return result
+
+
+def basic_check(expr, visitor_str, mutator_str):
+    """Helper function to check visitor and mutator on an expression"""
+
+    # Check visitor
+    basic_visitor = BasicVisitor()
+    basic_visitor.visit_expr(expr)
+    # Check AST printer visitor
+    log_visitor = ASTPrinter()
+    log_visitor.visit_expr(expr)
+    assert str(log_visitor.log) == visitor_str
+
+    # Check basic mutator
+    basic_mutator = BasicMutator()
+    mutated_expr = basic_mutator.visit_expr(expr)
+    assert_structural_equal(mutated_expr, expr)
+
+    # Check post-order printer mutator
+    post_log_mutator = ASTPostPrinterMutator()
+    mutated_expr = post_log_mutator.visit_expr(expr)
+    assert_structural_equal(mutated_expr, expr)
+    assert str(post_log_mutator.log) == mutator_str
+
+
+def test_var():
+    basic_check(n, "Var", "Var")
+
+
+def test_int_imm():
+    basic_check(tir.IntImm("int32", 10), "IntImm", "IntImm")
+
+
+def test_float_imm():
+    basic_check(tir.FloatImm("float32", 1.5), "FloatImm", "FloatImm")
+
+
+def test_string_imm():
+    basic_check(tir.StringImm("hello"), "StringImm", "StringImm")
+
+
+def test_add():
+    add_node = tir.Add(n, m)
+    basic_check(add_node, "\n".join(["Add", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Add"]))
+
+
+def test_sub():
+    sub_node = tir.Sub(n, m)
+    basic_check(sub_node, "\n".join(["Sub", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Sub"]))
+
+
+def test_mul():
+    mul_node = tir.Mul(n, m)
+    basic_check(mul_node, "\n".join(["Mul", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Mul"]))
+
+
+def test_div():
+    div_node = tir.Div(n, m)
+    basic_check(div_node, "\n".join(["Div", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Div"]))
+
+
+def test_floor_div():
+    floor_div_node = tir.FloorDiv(n, m)
+    basic_check(
+        floor_div_node,
+        "\n".join(["FloorDiv", "\tVar", "\tVar"]),
+        "\n".join(["Var", "Var", "FloorDiv"]),
+    )
+
+
+def test_floor_mod():
+    floor_mod_node = tir.FloorMod(n, m)
+    basic_check(
+        floor_mod_node,
+        "\n".join(["FloorMod", "\tVar", "\tVar"]),
+        "\n".join(["Var", "Var", "FloorMod"]),
+    )
+
+
+def test_min():
+    min_node = tir.Min(n, m)
+    basic_check(min_node, "\n".join(["Min", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Min"]))
+
+
+def test_max():
+    max_node = tir.Max(n, m)
+    basic_check(max_node, "\n".join(["Max", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Max"]))
+
+
+def test_eq():
+    eq_node = tir.EQ(n, m)
+    basic_check(eq_node, "\n".join(["EQ", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "EQ"]))
+
+
+def test_ne():
+    ne_node = tir.NE(n, m)
+    basic_check(ne_node, "\n".join(["NE", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "NE"]))
+
+
+def test_lt():
+    lt_node = tir.LT(n, m)
+    basic_check(lt_node, "\n".join(["LT", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "LT"]))
+
+
+def test_le():
+    le_node = tir.LE(n, m)
+    basic_check(le_node, "\n".join(["LE", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "LE"]))
+
+
+def test_gt():
+    gt_node = tir.GT(n, m)
+    basic_check(gt_node, "\n".join(["GT", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "GT"]))
+
+
+def test_ge():
+    ge_node = tir.GE(n, m)
+    basic_check(ge_node, "\n".join(["GE", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "GE"]))
+
+
+def test_and():
+    and_node = tir.And(tir.EQ(n, m), tir.LT(n, 10))
+    basic_check(
+        and_node,
+        "\n".join(["And", "\tEQ", "\t\tVar", "\t\tVar", "\tLT", "\t\tVar", "\t\tIntImm"]),
+        "\n".join(["Var", "Var", "EQ", "Var", "IntImm", "LT", "And"]),
+    )
+
+
+def test_or():
+    or_node = tir.Or(tir.EQ(n, m), tir.LT(n, 10))
+    basic_check(
+        or_node,
+        "\n".join(["Or", "\tEQ", "\t\tVar", "\t\tVar", "\tLT", "\t\tVar", "\t\tIntImm"]),
+        "\n".join(["Var", "Var", "EQ", "Var", "IntImm", "LT", "Or"]),
+    )
+
+
+def test_not():
+    not_node = tir.Not(tir.EQ(n, m))
+    basic_check(
+        not_node,
+        "\n".join(["Not", "\tEQ", "\t\tVar", "\t\tVar"]),
+        "\n".join(["Var", "Var", "EQ", "Not"]),
+    )
+
+
+def test_select():
+    select_node = tir.Select(tir.EQ(n, m), n, m)
+    basic_check(
+        select_node,
+        "\n".join(["Select", "\tEQ", "\t\tVar", "\t\tVar", "\tVar", "\tVar"]),
+        "\n".join(["Var", "Var", "EQ", "Var", "Var", "Select"]),
+    )
+
+
+def test_cast():
+    cast_node = tir.Cast("float32", n)
+    basic_check(cast_node, "\n".join(["Cast", "\tVar"]), "\n".join(["Var", "Cast"]))
+
+
+def test_let():
+    let_node = tir.Let(n, tir.IntImm("int32", 10), n + 1)
+    basic_check(
+        let_node,
+        "\n".join(["Let", "\tVar", "\tIntImm", "\tAdd", "\t\tVar", "\t\tIntImm"]),
+        "\n".join(["Var", "IntImm", "Var", "IntImm", "Add", "Let"]),
+    )
+
+
+def test_ramp():
+    ramp_node = tir.Ramp(n, 1, 4)
+    basic_check(
+        ramp_node,
+        "\n".join(["Ramp", "\tVar", "\tIntImm", "\tIntImm"]),
+        "\n".join(["Var", "IntImm", "IntImm", "Ramp"]),
+    )
+
+
+def test_broadcast():
+    broadcast_node = tir.Broadcast(n, 4)
+    basic_check(
+        broadcast_node,
+        "\n".join(["Broadcast", "\tVar", "\tIntImm"]),
+        "\n".join(["Var", "IntImm", "Broadcast"]),
+    )
+
+
+def test_inherit():
+    # The internal class is not instantiated.
+    class InternalVisitor(ExprVisitor):
+        def __init__(self) -> None:
+            super().__init__()
+            self.log = ASTLog()
+
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("InternalAdd")
+            self.log.push_scope()
+            self.visit_expr(op.a)
+            self.visit_expr(op.b)
+            self.log.pop_scope()
+
+        def visit_var_(self, op: Var) -> None:
+            self.log.add("InternalVar")
+
+    class LeafVisitor(InternalVisitor):
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("LeafAdd")
+            self.log.push_scope()
+            self.visit_expr(op.a)
+            self.visit_expr(op.b)
+            self.log.pop_scope()
+
+    add_node = tir.Add(n, m)
+    lv = LeafVisitor()
+    lv.visit_expr(add_node)
+    assert str(lv.log) == "\n".join(["LeafAdd", "\tInternalVar", "\tInternalVar"])
+
+
+def test_inherit_with_cls():
+    class InternalVisitor(ExprVisitor):
+        def __init__(self) -> None:
+            super().__init__()
+            self.log = ASTLog()
+
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("InternalAdd")
+            self.log.push_scope()
+            self.visit_expr(op.a)
+            self.visit_expr(op.b)
+            self.log.pop_scope()
+
+        def visit_var_(self, op: Var) -> None:
+            self.log.add("InternalVar")
+
+    class LeafVisitor(InternalVisitor):
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("LeafAdd")
+            self.log.push_scope()
+            self.visit_expr(op.a)
+            self.visit_expr(op.b)
+            self.log.pop_scope()
+
+    add_node = tir.Add(n, m)
+    iv = InternalVisitor()
+    iv.visit_expr(add_node)
+    assert str(iv.log) == "\n".join(["InternalAdd", "\tInternalVar", "\tInternalVar"])
+
+    lv = LeafVisitor()
+    lv.visit_expr(add_node)
+    assert str(lv.log) == "\n".join(["LeafAdd", "\tInternalVar", "\tInternalVar"])
+
+
+def test_call_visitor_super():
+    class InternalVisitor(ExprVisitor):
+        def __init__(self) -> None:
+            super().__init__()
+            self.log = ASTLog()
+
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("InternalAdd")
+            super().visit_add_(op)  # call ExprVisitor.visit_add_
+
+        def visit_var_(self, op: Var) -> None:
+            self.log.add("InternalVar")
+
+        def visit_int_imm_(self, op: IntImm) -> None:
+            self.log.add("InternalIntImm")
+
+    class LeafVisitor(InternalVisitor):
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("LeafAdd")
+            super().visit_add_(op)  # call InternalVisitor.visit_add_
+
+    add_node = tir.Add(n, tir.IntImm("int32", 10))
+    iv = InternalVisitor()
+    iv.visit_expr(add_node)
+    assert str(iv.log) == "\n".join(["InternalAdd", "InternalVar", "InternalIntImm"])
+
+    lv = LeafVisitor()
+    lv.visit_expr(add_node)
+    assert str(lv.log) == "\n".join(["LeafAdd", "InternalAdd", "InternalVar", "InternalIntImm"])
+
+
+def test_call_mutator_super():
+    class InternalMutator(ExprMutator):
+        def __init__(self) -> None:
+            super().__init__()
+            self.log = ASTLog()
+
+        def visit_add_(self, op: Add) -> tir.Expr:
+            self.log.add("InternalAdd")
+            return super().visit_add_(op)  # call ExprMutator.visit_add_
+
+        def visit_var_(self, op: Var) -> tir.Expr:
+            self.log.add("InternalVar")
+            return super().visit_var_(op)  # call ExprMutator.visit_var_
+
+        def visit_int_imm_(self, op: IntImm) -> tir.Expr:
+            self.log.add("InternalIntImm")
+            return super().visit_int_imm_(op)  # call ExprMutator.visit_int_imm_
+
+    class LeafMutator(InternalMutator):
+        def visit_add_(self, op: Add) -> tir.Expr:
+            self.log.add("LeafAdd")
+            return super().visit_add_(op)  # call InternalMutator.visit_add_
+
+    add_node = tir.Add(n, tir.IntImm("int32", 10))
+    im = InternalMutator()
+    im.visit_expr(add_node)
+    assert str(im.log) == "\n".join(["InternalAdd", "InternalVar", "InternalIntImm"])
+
+    lm = LeafMutator()
+    lm.visit_expr(add_node)
+    assert str(lm.log) == "\n".join(["LeafAdd", "InternalAdd", "InternalVar", "InternalIntImm"])
+
+
+def test_var_mutation():
+    """Test mutating variables in a TIR expression"""
+
+    class VarMutator(ExprMutator):
+        def __init__(self, var_map):
+            super().__init__()
+            self.var_map = var_map
+
+        def visit_var_(self, op: Var) -> tir.Expr:
+            if op.name in self.var_map:
+                return self.var_map[op.name]
+            return op
+
+    # Create a simple expression
+    expr = n + m
+
+    # Create a mutator that replaces 'n' with a constant
+    var_map = {"n": tir.IntImm("int32", 42)}
+    mutator = VarMutator(var_map)
+    result = mutator.visit_expr(expr)
+
+    # The result should be 42 + m
+    expected = tir.Add(tir.IntImm("int32", 42), m)
+    assert_structural_equal(result, expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-base/test_tir_host_func.py

@@ -0,0 +1,80 @@
+# 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
+from tvm.s_tir.meta_schedule.testing import te_workload
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,missing-class-docstring,missing-function-docstring
+# fmt: off
+
+
+@I.ir_module(s_tir=True)
+class Module:
+    @T.prim_func(s_tir=True)
+    def main(
+        A: T.Buffer((729, 729), "float32"),
+        B: T.Buffer((729, 729), "float32"),
+        C: T.Buffer((729, 729), "float32"),
+    ):
+        T.func_attr(
+            {
+                "global_symbol": "test",
+                "target": tvm.target.Target("llvm", host="llvm"),
+                "tirx.noalias": True,
+            }
+        )
+        # with T.sblock("root"):
+        for i, j, k in T.grid(729, 729, 729):
+            with T.sblock("C"):
+                v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
+                T.reads(A[v_i, v_k], B[v_k, v_j])
+                T.writes(C[v_i, v_j])
+                with T.init():
+                    C[v_i, v_j] = T.float32(0)
+                C[v_i, v_j] = C[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]
+
+# fmt: on
+# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,missing-class-docstring,missing-function-docstring
+
+
+def test_host_func():
+    """Test that host functions are not split."""
+    # te schedule copied from test_tir_transform_split_host_device.py
+
+    func = tvm.te.create_prim_func(
+        te_workload.matmul(729, 729, 729, in_dtype="float32", out_dtype="float32")
+    )
+    mod = tvm.ir.IRModule({"main": func})
+    target = tvm.target.Target("cuda")
+    mod = tvm.tirx.transform.Apply(
+        lambda f: f.with_attr(
+            {
+                "global_symbol": "test",
+                "tirx.is_host_func": True,
+            }
+        )
+    )(mod)
+    mod = tvm.tirx.transform.BindTarget(target)(mod)
+    tvm.ir.assert_structural_equal(mod, Module)
+    assert "tirx.is_host_func" not in mod["main"].attrs, (
+        """Target and is_host_func attributes should be mutually exclusive"""
+    )
+
+
+if __name__ == "__main__":
+    test_host_func()

tests/python/tirx-base/test_tir_imm_values.py

@@ -0,0 +1,597 @@
+# 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
+import math
+import random
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@pytest.mark.parametrize(
+    "dtype, literals",
+    [
+        ["int8", [-128, 0, 127]],
+        ["uint8", [0, 255]],
+        ["int32", [-2147483648, 2147483647]],
+        ["uint32", [0, 4294967295]],
+        ["int64", [-9223372036854775808, 9223372036854775807]],
+        ["uint64", [0, 9223372036854775807]],
+    ],
+)
+def test_tir_make_intimm(dtype, literals):
+    for l in literals:
+        imm = tirx.const(l, dtype)
+        assert imm.value == l, imm
+
+
+@pytest.mark.parametrize(
+    "dtype, literals",
+    [
+        ["int8", [-129, 128]],
+        ["uint8", [-1, 256]],
+        ["int32", [-2147483650, 2147483648]],
+        ["uint32", [-1, 4294967296]],
+        ["uint64", [-1, 18446744073709551616]],
+    ],
+)
+def test_tir_invalid_intimm(dtype, literals):
+    for l in literals:
+        # Out-of-range positive literals raise a builtin ValueError from
+        # the IntImm range check; negative-into-unsigned raises an
+        # InternalError ("cannot make uint from negative value") which is a
+        # RuntimeError subclass. Accept either.
+        with pytest.raises((RuntimeError, ValueError)):
+            tirx.const(l, dtype)
+
+
+@pytest.mark.parametrize(
+    "dtype, literals",
+    [
+        [
+            "uint64",
+            {
+                9223372036854775807: 9223372036854775807,
+                18446744073709551615: 18446744073709551615,
+            },
+        ],
+    ],
+)
+def test_tir_large_py_int_literals(dtype, literals):
+    """
+    For large uint value, use LargeUIntImm intrin,
+    """
+    for l in literals:
+        x = tirx.const(l, dtype)
+        if isinstance(x, tirx.IntImm | tirx.FloatImm):
+            assert x.value == literals[l]
+        else:
+            # LargeUIntImm(low32, hi32)
+            assert (int(x.args[1]) << 32) + int(x.args[0]) == literals[l]
+
+
+def test_tir_intimm_overflow():
+    assert int(tirx.const(255, "uint8") + tirx.const(1, "uint8")) == 0
+    assert int(tirx.const(2**31 - 1, "int32") + tirx.const(1, "int32")) == -(2**31)
+    assert int(tirx.const(2**32 - 1, "uint32") + tirx.const(1, "uint32")) == 0
+    assert int(tirx.const(2**63 - 1, "int64") + tirx.const(1, "int64")) == -(2**63)
+    assert int(tirx.const(2**32, "uint64") * tirx.const(2**32, "uint64")) == 0
+    # customized int types
+    assert int(tirx.const(7, "int4") + tirx.const(1, "int4")) == -8
+    assert int(tirx.const(2**39 - 1, "int40") + tirx.const(1, "int40")) == -(2**39)
+
+
+def compare_float_value(value, expect, msg):
+    if math.isfinite(value):
+        assert np.abs(value - expect) < 1e-5, f"{value} vs {expect}, {msg}"
+    elif math.isnan(value):
+        assert math.isnan(expect), f"{value} vs {expect}, {msg}"
+    elif math.isinf(value):
+        assert math.isinf(expect), f"{value} vs {expect}, {msg}"
+
+
+@pytest.mark.parametrize(
+    "dtype, literals",
+    [
+        ["float16", [-65504.0, 3.14, 65504.0, np.inf, np.nan]],
+        ["bfloat16", [-3.38953139e38, 3.38953139e38, 3.14]],
+        ["float32", [np.finfo("float32").min, 3.14, np.finfo("float32").max, np.inf, np.nan]],
+        ["float64", [np.finfo("float64").min, 3.14, np.finfo("float64").max, np.inf, np.nan]],
+    ],
+)
+def test_tir_make_floatimm(dtype, literals):
+    for l in literals:
+        imm = tirx.const(l, dtype)
+        compare_float_value(imm.value, l, "imm value should match feed value")
+
+
+@pytest.mark.parametrize(
+    "dtype, literals",
+    [
+        ["float16", [-65505.0, 65505.0]],
+        ["float32", [-3.402e39, 3.402e39]],
+    ],
+)
+def test_tir_invalid_floatimm(dtype, literals):
+    """Currently only fp16 and fp32 have range check."""
+    for l in literals:
+        # FloatImm out-of-range raises a builtin ValueError.
+        with pytest.raises(ValueError):
+            tirx.const(l, dtype)
+
+
+@pytest.mark.parametrize("dtype", ["float16", "float32", "float64"])
+@pytest.mark.parametrize("literal", [3.14, np.nan, np.inf])
+def test_tir_special_floatimms(dtype, literal):
+    x = tirx.const(literal, dtype)
+    compare_float_value(x.value, literal, "imm value should match feed value")
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_tir_too_large_literal_f64():
+    # Behavior check: if literal f64 value is out of dtype range, the
+    # object is still constructed, and eval to infinity.
+    @T.prim_func(s_tir=True)
+    def imm_overflow_fp64() -> T.float64:
+        T.evaluate(T.ret(T.float64(1.7976e309), dtype="float64"))
+
+    f = tvm.compile(imm_overflow_fp64, target="llvm")
+    assert math.isinf(f())
+
+
+@pytest.mark.parametrize(
+    "literal, expect_dtype",
+    [
+        (256, "int32"),
+        (2147483647, "int32"),
+        (-2147483648, "int32"),
+        (2147483648, "int64"),
+        (-2147483649, "int64"),
+        (3.14159, "float32"),
+        (np.finfo("float32").min, "float32"),
+        (np.finfo("float32").max, "float32"),
+        (-3.402e39, "float64"),
+        (3.402e39, "float64"),
+    ],
+)
+def test_tir_const_auto_dtype(literal, expect_dtype):
+    x = tirx.const(literal, dtype=None)
+    assert x.ty.dtype == expect_dtype
+    assert x.value == literal
+
+
+def check_tir_const_fold(
+    dtype, foldf, calcf, x_range=None, y_range=None, expect=None, skip_overflow=False
+):
+    """Helper to check constant folding behavior
+
+    Parameters
+    ----------
+    dtype: str
+        Datatype of constants
+
+    foldf: (x, y) -> z
+        Folding function to call
+
+    calcf: (x, y) -> z
+        Compiled calculation function to call
+
+    x_range: Union[int, float, tuple]
+        Single value or value range [min, max]
+
+    y_range: Union[int, float, tuple]
+        Single value or value range [min, max]
+
+    expect: Union[int, float]
+        Expected calculation result
+
+    skip_overflow: bool
+        Skip assertion if the overflow happens
+    """
+    seed = random.randint(0, 2147483648)
+    np.random.seed(seed)
+    ninfo = np.finfo(dtype) if dtype.startswith("float") else np.iinfo(dtype)
+
+    if x_range is None:
+        x_range = (ninfo.min, ninfo.max)
+    if isinstance(x_range, int | float):
+        x = x_range
+    elif dtype.startswith("int") or dtype.startswith("uint"):
+        x = np.random.randint(x_range[0], x_range[1] + 1, dtype=dtype)
+    else:
+        x = np.random.uniform(x_range[0], x_range[1])
+
+    if y_range is None:
+        y_range = (ninfo.min, ninfo.max)
+    if isinstance(y_range, int | float):
+        y = y_range
+    elif dtype.startswith("int") or dtype.startswith("uint"):
+        y = np.random.randint(y_range[0], y_range[1] + 1, dtype=dtype)
+    else:
+        y = np.random.uniform(y_range[0], y_range[1])
+
+    if skip_overflow:
+        py_res = foldf(x, y)
+        if isinstance(py_res, tirx.IntImm | tirx.FloatImm):
+            py_res = py_res.value
+        if not (ninfo.min <= py_res <= ninfo.max):
+            # If the result overflow, certain arithmetics is non-defined
+            # thus we intentionally do not make the test failed.
+            return
+
+    fold_res = foldf(tirx.const(x, dtype), tirx.const(y, dtype))
+    calc_res = calcf(x, y)
+
+    flaky_msg = (
+        f"{dtype} ({x}, {y}, {expect}) const folding check failed.\n"
+        + "This test is intentionally non-deterministic, "
+        + f"if it fails please report it in GitHub issue together with this seed {seed}\n"
+    )
+    if dtype.startswith("float"):
+        compare_float_value(calc_res, fold_res.value, flaky_msg)
+        if expect:
+            compare_float_value(expect, calc_res, flaky_msg)
+    else:
+        assert calc_res == fold_res.value, flaky_msg
+        if expect:
+            assert expect == calc_res, flaky_msg
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_tir_floatimm_const_fold():
+    """Behavior check: folding fp32 match platform f32 arithmetic"""
+
+    @T.prim_func(s_tir=True)
+    def float_imm_multiply(x: T.float32, y: T.float32, z: T.Buffer((), "float32")):
+        z[()] = x * y
+
+    @T.prim_func(s_tir=True)
+    def float_imm_add(x: T.float32, y: T.float32, z: T.Buffer((), "float32")):
+        z[()] = x + y
+
+    @T.prim_func(s_tir=True)
+    def float_imm_sub(x: T.float32, y: T.float32, z: T.Buffer((), "float32")):
+        z[()] = x - y
+
+    @T.prim_func(s_tir=True)
+    def float_imm_div(x: T.float32, y: T.float32, z: T.Buffer((), "float32")):
+        z[()] = x / y
+
+    def __wrap_build(f):
+        lib = tvm.compile(f, target="llvm")
+        z = tvm.runtime.tensor(np.zeros([]).astype("float32"))
+
+        def _func(x, y):
+            lib(x, y, z)
+            return z.numpy()
+
+        return _func
+
+    fmul = __wrap_build(float_imm_multiply)
+    fadd = __wrap_build(float_imm_add)
+    fsub = __wrap_build(float_imm_sub)
+    fdiv = __wrap_build(float_imm_div)
+
+    # overflow
+    check_tir_const_fold("float32", lambda x, y: x * y, fmul, 3.0e30, 3.0e30, np.inf)
+    check_tir_const_fold("float32", lambda x, y: x * y, fmul, 3.0e30, -3.0e30, -np.inf)
+    check_tir_const_fold("float32", lambda x, y: x / y, fdiv, 3.0e30, 3.0e-30, np.inf)
+
+    # divide by zero
+    with pytest.raises(RuntimeError):
+        check_tir_const_fold("float32", lambda x, y: x / y, fdiv, 1.0, 0.0)
+
+    # nan and inf
+    check_tir_const_fold("float32", lambda x, y: x + y, fadd, 1.0, np.nan, np.nan)
+    check_tir_const_fold("float32", lambda x, y: x + y, fadd, 1.0, np.inf, np.inf)
+    check_tir_const_fold("float32", lambda x, y: x + y, fadd, 1.0, -np.inf, -np.inf)
+
+    # randomized check
+    check_tir_const_fold("float32", lambda x, y: x * y, fmul)
+    check_tir_const_fold("float32", lambda x, y: x + y, fadd)
+    check_tir_const_fold("float32", lambda x, y: x - y, fsub)
+    check_tir_const_fold(
+        "float32", lambda x, y: x / y, fdiv, y_range=(0.01, np.finfo("float32").max)
+    )
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_tir_int8_const_fold():
+    """Behavior check: folding i8 operation match platform i8 arithmetic"""
+
+    @T.prim_func(s_tir=True)
+    def imm_multiply(x: T.int8, y: T.int8) -> T.int8:
+        T.evaluate(T.ret(x * y, dtype="int8"))
+
+    @T.prim_func(s_tir=True)
+    def imm_add(x: T.int8, y: T.int8) -> T.int8:
+        T.evaluate(T.ret(x + y, dtype="int8"))
+
+    @T.prim_func(s_tir=True)
+    def imm_sub(x: T.int8, y: T.int8) -> T.int8:
+        T.evaluate(T.ret(x - y, dtype="int8"))
+
+    @T.prim_func(s_tir=True)
+    def imm_truncdiv(x: T.int8, y: T.int8) -> T.int8:
+        T.evaluate(T.ret(T.truncdiv(x, y), dtype="int8"))
+
+    @T.prim_func(s_tir=True)
+    def imm_floordiv(x: T.int8, y: T.int8) -> T.int8:
+        T.evaluate(T.ret(T.floordiv(x, y), dtype="int8"))
+
+    fmul = tvm.compile(imm_multiply, target="llvm")
+    fadd = tvm.compile(imm_add, target="llvm")
+    fsub = tvm.compile(imm_sub, target="llvm")
+    ffloordiv = tvm.compile(imm_floordiv, target="llvm")
+    ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")
+
+    # overflow
+    check_tir_const_fold("int8", lambda x, y: x + y, fadd, 127, 1, -128)
+    check_tir_const_fold("int8", lambda x, y: x * y, fmul, 127, 127, 1)
+
+    # divide by zero
+    with pytest.raises(RuntimeError):
+        check_tir_const_fold("int8", lambda x, y: tirx.floordiv(x, y), ffloordiv, 1, 0)
+    with pytest.raises(RuntimeError):
+        check_tir_const_fold("int8", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, 1, 0)
+
+    # i8 mod folding is not implemented
+    assert not isinstance(tirx.floormod(tirx.const(7, "int8"), tirx.const(3, "int8")), tirx.IntImm)
+    assert not isinstance(tirx.truncmod(tirx.const(7, "int8"), tirx.const(3, "int8")), tirx.IntImm)
+
+    # randomized check
+    check_tir_const_fold("int8", lambda x, y: x * y, fmul)
+    check_tir_const_fold("int8", lambda x, y: x + y, fadd)
+    check_tir_const_fold("int8", lambda x, y: x - y, fsub)
+    check_tir_const_fold(
+        "int8", lambda x, y: tirx.floordiv(x, y), ffloordiv, y_range=(1, np.iinfo("int8").max)
+    )
+    check_tir_const_fold(
+        "int8", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, y_range=(1, np.iinfo("int8").max)
+    )
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_tir_uint8_const_fold():
+    """Behavior check: folding u8 operation match platform u8 arithmetic"""
+
+    @T.prim_func(s_tir=True)
+    def imm_multiply(x: T.uint8, y: T.uint8) -> T.uint8:
+        T.evaluate(T.ret(x * y, dtype="uint8"))
+
+    @T.prim_func(s_tir=True)
+    def imm_add(x: T.uint8, y: T.uint8) -> T.uint8:
+        T.evaluate(T.ret(x + y, dtype="uint8"))
+
+    @T.prim_func(s_tir=True)
+    def imm_sub(x: T.uint8, y: T.uint8) -> T.uint8:
+        T.evaluate(T.ret(x - y, dtype="uint8"))
+
+    @T.prim_func(s_tir=True)
+    def imm_truncdiv(x: T.uint8, y: T.uint8) -> T.uint8:
+        T.evaluate(T.ret(T.truncdiv(x, y), dtype="uint8"))
+
+    @T.prim_func(s_tir=True)
+    def imm_floordiv(x: T.uint8, y: T.uint8) -> T.uint8:
+        T.evaluate(T.ret(T.floordiv(x, y), dtype="uint8"))
+
+    fmul = tvm.compile(imm_multiply, target="llvm")
+    fadd = tvm.compile(imm_add, target="llvm")
+    fsub = tvm.compile(imm_sub, target="llvm")
+    ffloordiv = tvm.compile(imm_floordiv, target="llvm")
+    ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")
+
+    # overflow
+    check_tir_const_fold("uint8", lambda x, y: x + y, fadd, 255, 1, 0)
+
+    # zero sub
+    with pytest.raises(RuntimeError):
+        check_tir_const_fold("uint8", lambda x, y: x - y, fsub, 0, 10)
+
+    # divide by zero
+    with pytest.raises(RuntimeError):
+        check_tir_const_fold("uint8", lambda x, y: tirx.floordiv(x, y), ffloordiv, 1, 0)
+    with pytest.raises(RuntimeError):
+        check_tir_const_fold("uint8", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, 1, 0)
+
+    # u8 floormod folding is overflow-free and implemented
+    folded_floormod = tirx.floormod(tirx.const(7, "uint8"), tirx.const(3, "uint8"))
+    assert isinstance(folded_floormod, tirx.IntImm)
+    assert int(folded_floormod) == 1
+
+    # u8 truncmod folding is not implemented
+    assert not isinstance(
+        tirx.truncmod(tirx.const(7, "uint8"), tirx.const(3, "uint8")), tirx.IntImm
+    )
+
+    # randomized check
+    check_tir_const_fold("uint8", lambda x, y: x * y, fmul)
+    check_tir_const_fold("uint8", lambda x, y: x + y, fadd)
+    check_tir_const_fold("uint8", lambda x, y: x - y, fsub)
+    check_tir_const_fold(
+        "uint8", lambda x, y: tirx.floordiv(x, y), ffloordiv, y_range=(1, np.iinfo("uint8").max)
+    )
+    check_tir_const_fold(
+        "uint8", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, y_range=(1, np.iinfo("uint8").max)
+    )
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_tir_int32_const_fold():
+    """Behavior check: folding i32 operation match platform i32 arithmetic"""
+
+    @T.prim_func(s_tir=True)
+    def imm_multiply(x: T.int32, y: T.int32) -> T.int32:
+        T.evaluate(T.ret(x * y, dtype="int32"))
+
+    @T.prim_func(s_tir=True)
+    def imm_add(x: T.int32, y: T.int32) -> T.int32:
+        T.evaluate(T.ret(x + y, dtype="int32"))
+
+    @T.prim_func(s_tir=True)
+    def imm_sub(x: T.int32, y: T.int32) -> T.int32:
+        T.evaluate(T.ret(x - y, dtype="int32"))
+
+    @T.prim_func(s_tir=True)
+    def imm_truncdiv(x: T.int32, y: T.int32) -> T.int32:
+        T.evaluate(T.ret(T.truncdiv(x, y), dtype="int32"))
+
+    @T.prim_func(s_tir=True)
+    def imm_truncmod(x: T.int32, y: T.int32) -> T.int32:
+        T.evaluate(T.ret(T.truncmod(x, y), dtype="int32"))
+
+    @T.prim_func(s_tir=True)
+    def imm_floordiv(x: T.int32, y: T.int32) -> T.int32:
+        T.evaluate(T.ret(T.floordiv(x, y), dtype="int32"))
+
+    @T.prim_func(s_tir=True)
+    def imm_floormod(x: T.int32, y: T.int32) -> T.int32:
+        T.evaluate(T.ret(T.floormod(x, y), dtype="int32"))
+
+    fmul = tvm.compile(imm_multiply, target="llvm")
+    fadd = tvm.compile(imm_add, target="llvm")
+    fsub = tvm.compile(imm_sub, target="llvm")
+    ffloordiv = tvm.compile(imm_floordiv, target="llvm")
+    ffloormod = tvm.compile(imm_floormod, target="llvm")
+    ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")
+    ftruncmod = tvm.compile(imm_truncmod, target="llvm")
+
+    # i32 overflow is not specified, only check for range
+    assert -(2**31) <= int(tirx.const(2**31 - 1, "int32") + tirx.const(1, "int32")) < 2**31
+    assert -(2**31) <= int(tirx.const(-(2**31), "int32") - tirx.const(1, "int32")) < 2**31
+
+    # divide by zero
+    with pytest.raises(RuntimeError):
+        check_tir_const_fold("int32", lambda x, y: tirx.floordiv(x, y), ffloordiv, 1, 0)
+    with pytest.raises(RuntimeError):
+        check_tir_const_fold("int32", lambda x, y: tirx.floormod(x, y), ffloormod, 1, 0)
+    with pytest.raises(RuntimeError):
+        check_tir_const_fold("int32", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, 1, 0)
+    with pytest.raises(RuntimeError):
+        check_tir_const_fold("int32", lambda x, y: tirx.truncmod(x, y), ftruncmod, 1, 0)
+
+    # randomized check
+    check_tir_const_fold("int32", lambda x, y: x * y, fmul, skip_overflow=True)
+    check_tir_const_fold("int32", lambda x, y: x + y, fadd, skip_overflow=True)
+    check_tir_const_fold("int32", lambda x, y: x - y, fsub, skip_overflow=True)
+    check_tir_const_fold(
+        "int32",
+        lambda x, y: tirx.floordiv(x, y),
+        ffloordiv,
+        y_range=(1, np.iinfo("int32").max),
+        skip_overflow=True,
+    )
+    check_tir_const_fold(
+        "int32",
+        lambda x, y: tirx.truncdiv(x, y),
+        ftruncdiv,
+        y_range=(1, np.iinfo("int32").max),
+        skip_overflow=True,
+    )
+    check_tir_const_fold(
+        "int32",
+        lambda x, y: tirx.floormod(x, y),
+        ffloormod,
+        y_range=(1, np.iinfo("int32").max),
+        skip_overflow=False,
+    )
+    check_tir_const_fold(
+        "int32",
+        lambda x, y: tirx.truncmod(x, y),
+        ftruncmod,
+        y_range=(1, np.iinfo("int32").max),
+        skip_overflow=False,
+    )
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_tir_uint32_const_fold():
+    """Behavior check: folding u32 operation match platform u32 arithmetic"""
+
+    @T.prim_func(s_tir=True)
+    def imm_multiply(x: T.uint32, y: T.uint32) -> T.uint32:
+        T.evaluate(T.ret(x * y, dtype="uint32"))
+
+    @T.prim_func(s_tir=True)
+    def imm_add(x: T.uint32, y: T.uint32) -> T.uint32:
+        T.evaluate(T.ret(x + y, dtype="uint32"))
+
+    @T.prim_func(s_tir=True)
+    def imm_sub(x: T.uint32, y: T.uint32) -> T.uint32:
+        T.evaluate(T.ret(x - y, dtype="uint32"))
+
+    @T.prim_func(s_tir=True)
+    def imm_truncdiv(x: T.uint32, y: T.uint32) -> T.uint32:
+        T.evaluate(T.ret(T.truncdiv(x, y), dtype="uint32"))
+
+    @T.prim_func(s_tir=True)
+    def imm_floordiv(x: T.uint32, y: T.uint32) -> T.uint32:
+        T.evaluate(T.ret(T.floordiv(x, y), dtype="uint32"))
+
+    fmul = tvm.compile(imm_multiply, target="llvm")
+    fadd = tvm.compile(imm_add, target="llvm")
+    fsub = tvm.compile(imm_sub, target="llvm")
+    ffloordiv = tvm.compile(imm_floordiv, target="llvm")
+    ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")
+
+    # u32 overflow is not specified, only check for range
+    assert 0 <= int(tirx.const(2**32 - 1, "uint32") + tirx.const(1, "uint32")) < 2**32
+
+    # divide by zero
+    with pytest.raises(RuntimeError):
+        check_tir_const_fold("uint32", lambda x, y: tirx.floordiv(x, y), ffloordiv, 1, 0)
+    with pytest.raises(RuntimeError):
+        check_tir_const_fold("uint32", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, 1, 0)
+
+    # u32 floormod folding is overflow-free and implemented
+    folded_floormod = tirx.floormod(tirx.const(7, "uint32"), tirx.const(3, "uint32"))
+    assert isinstance(folded_floormod, tirx.IntImm)
+    assert int(folded_floormod) == 1
+
+    # u32 truncmod folding is not implemented
+    assert not isinstance(
+        tirx.truncmod(tirx.const(7, "uint32"), tirx.const(3, "uint32")), tirx.IntImm
+    )
+
+    # randomized check
+    check_tir_const_fold("uint32", lambda x, y: x * y, fmul, skip_overflow=True)
+    check_tir_const_fold("uint32", lambda x, y: x + y, fadd, skip_overflow=True)
+    check_tir_const_fold("uint32", lambda x, y: x - y, fsub, skip_overflow=True)
+    check_tir_const_fold(
+        "uint32",
+        lambda x, y: tirx.floordiv(x, y),
+        ffloordiv,
+        y_range=(1, np.iinfo("uint32").max),
+        skip_overflow=False,
+    )
+    check_tir_const_fold(
+        "uint32",
+        lambda x, y: tirx.truncdiv(x, y),
+        ftruncdiv,
+        y_range=(1, np.iinfo("uint32").max),
+        skip_overflow=False,
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-base/test_tir_index_map.py

@@ -0,0 +1,368 @@
+# 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, F401
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm.ir import assert_structural_equal
+from tvm.runtime import const
+from tvm.script import tirx as T
+from tvm.tirx import IndexMap, IntImm, floordiv, floormod, stmt_functor
+
+
+def assert_equal_index_map(map1: IndexMap, map2: IndexMap) -> None:
+    iters_1 = map1.map_indices(map2.initial_indices)
+    iters_2 = map2.final_indices
+    assert len(iters_1) == len(iters_2)
+
+    analyzer = tvm.arith.Analyzer()
+    for iter1, iter2 in zip(iters_1, iters_2):
+        assert analyzer.can_prove_equal(iter1, iter2)
+
+
+def test_index_mapping():
+    index_map = IndexMap.from_func(lambda i: [i // 4, i % 4], index_dtype="int32")
+
+    assert_structural_equal(index_map.map_indices([0]), [T.int32(0), T.int32(0)])
+    assert_structural_equal(index_map.map_indices([3]), [T.int32(0), T.int32(3)])
+    assert_structural_equal(index_map.map_indices([4]), [T.int32(1), T.int32(0)])
+    assert_structural_equal(index_map.map_indices([42]), [T.int32(10), T.int32(2)])
+    assert_structural_equal(index_map.map_indices([T.int64(42)]), [T.int64(10), T.int64(2)])
+
+
+def test_map_indices_accepts_external_analyzer():
+    tile = tvm.tirx.Var("tile", "int32")
+    index_map = IndexMap.from_func(lambda i: [i // tile], index_dtype="int32")
+    analyzer = tvm.arith.Analyzer()
+
+    unsimplified = index_map.map_indices([T.int32(32)])[0]
+    analyzer.bind(tile, T.int32(16))
+    simplified = index_map.map_indices([T.int32(32)], analyzer=analyzer)[0]
+
+    assert not tvm_ffi.structural_equal(unsimplified, T.int32(2))
+    assert_structural_equal(simplified, T.int32(2))
+
+
+def test_map_shape_accepts_external_analyzer():
+    tile = tvm.tirx.Var("tile", "int32")
+    index_map = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32")
+    analyzer = tvm.arith.Analyzer()
+
+    analyzer.bind(tile, T.int32(16))
+    mapped_shape = index_map.map_shape([T.int32(32)], analyzer=analyzer)
+
+    assert_structural_equal(mapped_shape, [T.int32(2), T.int32(16)])
+
+
+def test_is_equivalent_to_accepts_external_analyzer():
+    tile = tvm.tirx.Var("tile", "int32")
+    concrete = IndexMap.from_func(lambda i: [i // 4, i % 4], index_dtype="int32")
+    symbolic = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32")
+
+    # Without binding `tile`, the symbolic map cannot be proven equivalent.
+    assert not concrete.is_equivalent_to(symbolic)
+
+    analyzer = tvm.arith.Analyzer()
+    analyzer.bind(tile, T.int32(4))
+    assert concrete.is_equivalent_to(symbolic, analyzer=analyzer)
+
+
+def test_shape_mapping():
+    index_map = IndexMap.from_func(lambda i: [i // 4, i % 4], index_dtype="int32")
+
+    assert_structural_equal(index_map.map_shape([4]), [T.int32(1), T.int32(4)])
+    assert_structural_equal(index_map.map_shape([16]), [T.int32(4), T.int32(4)])
+
+    assert_structural_equal(index_map.map_shape([14]), [T.int32(4), T.int32(4)])
+    assert_structural_equal(index_map.map_shape([T.int64(16)]), [T.int64(4), T.int64(4)])
+    assert_structural_equal(index_map.map_shape([T.int64(14)]), [T.int64(4), T.int64(4)])
+
+
+def test_inverse():
+    index_map = IndexMap.from_func(lambda i: [i // 4, i % 4])
+    expected_inverse = IndexMap.from_func(lambda i, j: [4 * i + j])
+
+    assert index_map.inverse([16]).is_equivalent_to(expected_inverse)
+
+
+def test_inverse_preserves_passthrough_var_names():
+    index_map = IndexMap.from_func(lambda i, j: [j, i], index_dtype="int32")
+    inverse = index_map.inverse([8, 16])
+    assert [v.name for v in inverse.initial_indices] == ["j", "i"]
+
+
+def test_inverse_accepts_external_analyzer():
+    tile = tvm.tirx.Var("tile", "int32")
+    index_map = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32")
+    analyzer = tvm.arith.Analyzer()
+
+    analyzer.bind(tile, T.int32(16))
+    inverse = index_map.inverse([T.int32(32)], analyzer=analyzer)
+    mapped = inverse.map_indices([T.int32(1), T.int32(3)], analyzer=analyzer)
+
+    assert_structural_equal(mapped, [T.int32(19)])
+
+
+def test_nonbijective_inverse_gives_error():
+    index_map = IndexMap.from_func(lambda i: [i // 4, i % 4])
+
+    with pytest.raises(RuntimeError):
+        index_map.inverse([14])
+
+
+dynamic_N = tvm.tirx.Var("N", "int32")
+padding_test_case = tvm.testing.parameter(
+    by_dict={
+        "no_padding": dict(
+            forward=lambda i: [i // 4, i % 4],
+            inverse=lambda i, j: [4 * i + j],
+            pre_shape=[16],
+            post_shape=[T.int32(4), T.int32(4)],
+            padding=lambda i, j: tvm.runtime.convert(False),
+        ),
+        "right_padding": dict(
+            forward=lambda i: [i // 4, i % 4],
+            inverse=lambda i, j: [4 * i + j],
+            pre_shape=[15],
+            post_shape=[T.int32(4), T.int32(4)],
+            padding=lambda i, j: tvm.tirx.And(i == 3, tvm.runtime.convert(3) == j),
+        ),
+        "left_padding": dict(
+            forward=lambda i: [(i + 1) // 4, (i + 1) % 4],
+            inverse=lambda i, j: [4 * i + j - 1],
+            pre_shape=[15],
+            post_shape=[T.int32(4), T.int32(4)],
+            padding=lambda i, j: tvm.tirx.And(i == 0, j < 1),
+        ),
+        "left_and_right_padding": dict(
+            forward=lambda i: [(i + 1) // 4, (i + 1) % 4],
+            inverse=lambda i, j: [4 * i + j - 1],
+            pre_shape=[14],
+            post_shape=[T.int32(4), T.int32(4)],
+            padding=lambda i, j: tvm.tirx.Or(
+                tvm.tirx.And(i == 0, j < 1),
+                tvm.tirx.And(i == 3, tvm.runtime.convert(3) == j),
+            ),
+        ),
+        "dynamic_size": dict(
+            forward=lambda i: [i // 4, i % 4],
+            inverse=lambda i, j: [4 * i + j],
+            pre_shape=[dynamic_N],
+            post_shape=[(dynamic_N - dynamic_N % (-4)) // 4, T.int32(4)],
+            padding=lambda i, j: tvm.tirx.And(
+                dynamic_N % (-4) != 0,
+                tvm.tirx.And(i == dynamic_N // 4, j >= dynamic_N % 4),
+            ),
+        ),
+        "2d_padding": dict(
+            forward=lambda i, j: [(i + 1) // 4, (j + 5) // 8, (i + 1) % 4, (j + 5) % 8],
+            inverse=lambda i_outer, j_outer, i_inner, j_inner: [
+                4 * i_outer + i_inner - 1,
+                8 * j_outer + j_inner - 5,
+            ],
+            pre_shape=[14, 31],
+            post_shape=[
+                T.int32(4),  # ceildiv(left_pad + i.extent, 4) = ceildiv(1 + 14, 4) = 4
+                T.int32(5),  # ceildiv(left_pad + j.extent, 8) = ceildiv(5 + 31, 8) = 5
+                T.int32(4),  # Range of iter%4
+                T.int32(8),  # Range of iter%8
+            ],
+            padding=lambda i_outer, j_outer, i_inner, j_inner: tvm.tirx.Or(
+                tvm.tirx.Or(
+                    tvm.tirx.And(i_outer == 0, i_inner < 1),
+                    tvm.tirx.And(i_outer == 3, tvm.runtime.convert(3) == i_inner),
+                ),
+                tvm.tirx.Or(
+                    tvm.tirx.And(j_outer == 0, j_inner < 5),
+                    tvm.tirx.And(j_outer == 4, j_inner >= 4),
+                ),
+            ),
+        ),
+        "multiple_right_padding": dict(
+            forward=lambda i: [i // 32, (i // 4) % 8, i % 4],
+            inverse=lambda i, j, k: [32 * i + 4 * j + k],
+            pre_shape=[116],
+            post_shape=[T.int32(4), T.int32(8), T.int32(4)],
+            padding=lambda i, j, k: tvm.tirx.And(i == 3, 4 * j + k >= 20),
+        ),
+        "multiple_right_padding_transpose": dict(
+            forward=lambda i: [(i // 4) % 8, i // 32, i % 4],
+            inverse=lambda j, i, k: [32 * i + 4 * j + k],
+            pre_shape=[116],
+            post_shape=[T.int32(8), T.int32(4), T.int32(4)],
+            padding=lambda j, i, k: tvm.tirx.And(i == 3, 4 * j + k >= 20),
+        ),
+        "multiple_left_padding": dict(
+            forward=lambda i: [(i + 5) // 32, ((i + 5) // 4) % 8, (i + 5) % 4],
+            inverse=lambda i, j, k: [32 * i + 4 * j + k - 5],
+            pre_shape=[123],
+            post_shape=[T.int32(4), T.int32(8), T.int32(4)],
+            padding=lambda i, j, k: tvm.tirx.And(i == 0, j * 4 + k < 5),
+        ),
+        "multiple_left_padding_with_transpose": dict(
+            forward=lambda i: [((i + 5) // 4) % 8, (i + 5) // 32, (i + 5) % 4],
+            inverse=lambda j, i, k: [32 * i + 4 * j + k - 5],
+            pre_shape=[123],
+            post_shape=[T.int32(8), T.int32(4), T.int32(4)],
+            padding=lambda j, i, k: tvm.tirx.And(i == 0, j * 4 + k < 5),
+        ),
+        "outer_loop_extent_one": dict(
+            forward=lambda i: [i // 4, i % 4],
+            inverse=lambda i, j: [i * 4 + j],
+            pre_shape=[3],
+            post_shape=[T.int32(1), T.int32(4)],
+            padding=lambda i, j: tvm.runtime.convert(3) == j,
+        ),
+    }
+)
+
+
+def test_nonsurjective_inverse(padding_test_case):
+    index_map = IndexMap.from_func(padding_test_case["forward"], index_dtype="int32")
+
+    inverse, padding_predicate = index_map.non_surjective_inverse(padding_test_case["pre_shape"])
+    expected_inverse = IndexMap.from_func(padding_test_case["inverse"])
+    assert inverse.is_equivalent_to(expected_inverse)
+
+    post_shape = index_map.map_shape(padding_test_case["pre_shape"])
+    tvm.ir.assert_structural_equal(post_shape, padding_test_case["post_shape"])
+
+    expected_predicate = padding_test_case["padding"](*inverse.initial_indices)
+
+    # Can't use analyzer.can_prove_equal, because it can't simplify
+    # expressions like `(4*i+j >= 14) - (4*i+j >= 14)`.
+    analyzer = tvm.arith.Analyzer()
+    expected_predicate = analyzer.simplify(expected_predicate)
+    padding_predicate = analyzer.simplify(padding_predicate)
+    tvm.ir.assert_structural_equal(padding_predicate, expected_predicate)
+
+
+def test_non_surjective_inverse_accepts_external_analyzer():
+    tile = tvm.tirx.Var("tile", "int32")
+    index_map = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32")
+    analyzer = tvm.arith.Analyzer()
+
+    analyzer.bind(tile, T.int32(16))
+    inverse, padding_predicate = index_map.non_surjective_inverse([T.int32(31)], analyzer=analyzer)
+    mapped = inverse.map_indices([T.int32(1), T.int32(15)], analyzer=analyzer)
+
+    assert_structural_equal(mapped, [T.int32(31)])
+
+    padding_at_last_element = stmt_functor.substitute(
+        padding_predicate,
+        {inverse.initial_indices[0]: T.int32(1), inverse.initial_indices[1]: T.int32(15)},
+    )
+    padding_at_first_element = stmt_functor.substitute(
+        padding_predicate,
+        {inverse.initial_indices[0]: T.int32(0), inverse.initial_indices[1]: T.int32(0)},
+    )
+    assert_structural_equal(analyzer.simplify(padding_at_last_element), T.bool(True))
+    assert_structural_equal(analyzer.simplify(padding_at_first_element), T.bool(False))
+
+
+def test_non_surjective_inverse_does_not_bind_output_vars_to_external_analyzer():
+    tile = tvm.tirx.Var("tile", "int32")
+    index_map = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32")
+    analyzer = tvm.arith.Analyzer()
+
+    analyzer.bind(tile, T.int32(16))
+    inverse, _ = index_map.non_surjective_inverse([T.int32(31)], analyzer=analyzer)
+
+    analyzer.bind(inverse.initial_indices[0], T.int32(0))
+    analyzer.bind(inverse.initial_indices[1], T.int32(1))
+
+
+def test_index_map_inverse_no_iter():
+    def input_example(i0, i1, i2, i3):
+        j0 = floordiv(i3, 32)
+        j1 = floordiv(i2, 2)
+        j2 = floormod(i2, 2)
+        j3 = floormod(i3, 32)
+        return j0, j1, j2, j3
+
+    def expected_inverse(i0, i1, i2, i3):
+        return IntImm("int32", 0), IntImm("int32", 0), i2 + i1 * 2, i3 + i0 * 32
+
+    index_map = IndexMap.from_func(input_example)
+    inverse_map = index_map.inverse([1, 1, 64, 64])
+    expected_map = IndexMap.from_func(expected_inverse)
+    assert expected_map.is_equivalent_to(inverse_map)
+
+
+def test_map_tensor():
+    index_map = IndexMap.from_func(lambda i: [i // 4, i % 4])
+
+    inp = np.arange(16).astype("int8")
+
+    out = index_map.map_tensor(tvm.runtime.tensor(inp)).numpy()
+
+    ref = np.zeros(out.shape).astype("int8")
+
+    for i in range(16):
+        ref[i // 4, i % 4] = inp[i]
+
+    np.testing.assert_equal(ref, out)
+
+    index_map = IndexMap.from_func(lambda i0, i1, i2, i3: (i3, i0, i1, i2))
+
+    inp = np.random.randn(10, 10, 10, 10).astype("float16")
+
+    out = index_map.map_tensor(tvm.runtime.tensor(inp)).numpy()
+
+    ref = np.transpose(inp, (3, 0, 1, 2))
+
+    np.testing.assert_equal(ref, out)
+
+    index_map = IndexMap.from_func(
+        lambda i0, i1, i2, i3: (
+            floordiv(i3, 32),
+            i0,
+            floordiv(i2, 8),
+            floordiv(floormod(i3, 32), 16),
+            i1,
+            floormod(i2, 8),
+            floormod(i3, 16),
+        )
+    )
+
+    kH = kW = 3
+    I = 64
+    O = 64
+    inp = np.random.randn(kH, kW, I, O).astype("float32")
+    arr = tvm.runtime.tensor(inp)
+    out = index_map.map_tensor(arr).numpy()
+
+    ref = np.zeros(out.shape).astype("float32")
+
+    for i0 in range(kH):
+        for i1 in range(kW):
+            for i2 in range(I):
+                for i3 in range(O):
+                    v = inp[i0, i1, i2, i3]
+                    ref[i3 // 32, i0, i2 // 8, (i3 % 32) // 16, i1, i2 % 8, i3 % 16] = v
+
+    np.testing.assert_equal(ref, out)
+
+    inverse_map = index_map.inverse(inp.shape)
+    np.testing.assert_equal(inverse_map.map_tensor(index_map.map_tensor(arr)).numpy(), inp)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-base/test_tir_intrin.py

@@ -0,0 +1,413 @@
+# 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 ctypes
+import math
+
+import numpy as np
+import pytest
+
+pytest.importorskip("scipy")
+
+import scipy
+
+import tvm
+import tvm.testing
+from tvm import te, tirx, topi
+from tvm.script import tirx as T
+from tvm.support import clang, utils
+
+
+def test_nearbyint():
+    m = te.var(
+        "m",
+    )
+    A = te.placeholder((m,), name="A")
+    A_rounded = te.compute((m,), lambda *i: tvm.tirx.nearbyint(A(*i)), name="A")
+
+    # Convert to TIR and create schedule
+    mod = te.create_prim_func([A, A_rounded])
+    sch = tvm.s_tir.Schedule(mod)
+
+    # Build from scheduled TIR
+    func = tvm.compile(sch.mod, target="llvm")
+
+    dev = tvm.cpu(0)
+    n = 10
+    a = tvm.runtime.tensor(np.random.uniform(high=100, size=n).astype(A.dtype.dtype), dev)
+    a_rounded = tvm.runtime.tensor(np.random.uniform(size=n).astype(A_rounded.dtype.dtype), dev)
+    func(a, a_rounded)
+    # Note that numpys rint rounds to nearest integer with
+    # ties to halfway is broken by rounding to even.
+    # So that 1.5 and 2.5 will round 2.
+    # This is the default rounding mode with libc as well.
+    # However one can set a different rounding mode and in that
+    # case numpy result might differ.
+    tvm.testing.assert_allclose(a_rounded.numpy(), np.rint(a.numpy()))
+
+
+def test_round_ties_to_even():
+    """Test that tir.round uses ties-to-even (banker's rounding) semantics."""
+    m = te.var("m")
+    A = te.placeholder((m,), name="A")
+    A_rounded = te.compute((m,), lambda *i: tvm.tirx.round(A(*i)), name="A")
+
+    mod = te.create_prim_func([A, A_rounded])
+    sch = tvm.s_tir.Schedule(mod)
+    func = tvm.compile(sch.mod, target="llvm")
+
+    dev = tvm.cpu(0)
+    # Midpoint values where ties-to-even and ties-away differ
+    test_values = np.array([0.5, 1.5, 2.5, 3.5, -0.5, -1.5, -2.5, -3.5], dtype="float32")
+    expected = np.array([0.0, 2.0, 2.0, 4.0, 0.0, -2.0, -2.0, -4.0], dtype="float32")
+
+    a = tvm.runtime.tensor(test_values, dev)
+    a_rounded = tvm.runtime.tensor(np.zeros(len(test_values), dtype="float32"), dev)
+    func(a, a_rounded)
+    tvm.testing.assert_allclose(a_rounded.numpy(), expected)
+
+
+def test_round_intrinsics_on_int():
+    i = tvm.tirx.Var("i", "int32")
+    for op in [tvm.tirx.round, tvm.tirx.trunc, tvm.tirx.ceil, tvm.tirx.floor, tvm.tirx.nearbyint]:
+        assert op(tvm.tirx.const(10, "int32")).value == 10
+        assert op(tvm.tirx.const(True, "bool")).value == True
+        assert op(i).same_as(i)
+
+    assert tvm.tirx.isnan(tvm.tirx.const(10, "int32")).value == False
+
+
+def test_unary_intrin():
+    test_funcs = [
+        (tvm.tirx.exp, lambda x: np.exp(x)),
+        (tvm.tirx.exp10, lambda x: np.power(10, x)),
+        (tvm.tirx.log2, lambda x: np.log2(x)),
+        (tvm.tirx.log10, lambda x: np.log10(x)),
+        (tvm.tirx.sinh, lambda x: np.sinh(x)),
+        (tvm.tirx.cosh, lambda x: np.cosh(x)),
+        (tvm.tirx.log1p, lambda x: np.log1p(x)),
+        (tvm.tirx.asin, lambda x: np.arcsin(x)),
+        (tvm.tirx.acos, lambda x: np.arccos(x)),
+        (tvm.tirx.atan, lambda x: np.arctan(x)),
+        (tvm.tirx.asinh, lambda x: np.arcsinh(x)),
+        (tvm.tirx.acosh, lambda x: np.arccosh(x)),
+        (tvm.tirx.atanh, lambda x: np.arctanh(x)),
+        (tvm.tirx.erf, lambda x: scipy.special.erf(x)),
+    ]
+
+    def run_test(tvm_intrin, np_func, atol=1e-5, rtol=1e-5):
+        m = te.var(
+            "m",
+        )
+        A = te.placeholder((m,), name="A")
+        B = te.compute((m,), lambda *i: tvm_intrin(A(*i)), name="B")
+
+        # Convert to TIR and create schedule
+        mod = te.create_prim_func([A, B])
+        sch = tvm.s_tir.Schedule(mod)
+
+        # Build from scheduled TIR
+        func = tvm.compile(sch.mod, target="llvm")
+
+        dev = tvm.cpu(0)
+        n = 10
+        a = tvm.runtime.tensor(np.random.uniform(0.1, 0.5, size=n).astype(A.dtype.dtype), dev)
+        b = tvm.runtime.tensor(np.random.uniform(size=n).astype(A.dtype.dtype), dev)
+        func(a, b)
+        tvm.testing.assert_allclose(b.numpy(), np_func(a.numpy()), atol=atol, rtol=rtol)
+
+        # Out-of-bounds test for asin/acos
+        name = tvm_intrin.__name__
+        if name in ("asin", "acos"):
+            # generate some values outside [-1, 1]
+            n = 8
+            out_np = np.concatenate(
+                [
+                    np.random.uniform(1.1, 2.0, size=n // 2),
+                    np.random.uniform(-2.0, -1.1, size=n // 2),
+                ]
+            ).astype(A.dtype.dtype)
+            a2 = tvm.runtime.tensor(out_np, dev)
+            b2 = tvm.runtime.tensor(np.empty_like(out_np), dev)
+            func(a2, b2)
+            # all outputs should be NaN
+            assert np.all(np.isnan(b2.numpy()))
+        if name == "exp":
+            n = 8
+            out_np = np.random.randint(-20, 20, size=n).astype(A.dtype.dtype)
+            a2 = tvm.runtime.tensor(out_np, dev)
+            b2 = tvm.runtime.tensor(np.empty_like(out_np), dev)
+            func(a2, b2)
+            assert b2.numpy().dtype == np.float32
+            # Verify correctness against NumPy exp
+            expected = np.exp(out_np.astype(np.float32))
+            tvm.testing.assert_allclose(b2.numpy(), expected, rtol=1e-5, atol=1e-5)
+
+    for func in test_funcs:
+        atol = rtol = 1e-3 if func[0].__name__ in ["asin", "acos", "atan"] else 1e-5
+        run_test(*func, atol, rtol)
+
+
+def test_asin_acos_boundary_values():
+    """Test asin and acos with boundary values and threshold switching."""
+    test_funcs = [
+        (tvm.tirx.asin, lambda x: np.arcsin(x)),
+        (tvm.tirx.acos, lambda x: np.arccos(x)),
+    ]
+
+    def run_test(tvm_intrin, np_func):
+        m = te.var("m")
+        A = te.placeholder((m,), name="A")
+        B = te.compute((m,), lambda *i: tvm_intrin(A(*i)), name="B")
+
+        mod = te.create_prim_func([A, B])
+        sch = tvm.s_tir.Schedule(mod)
+        func = tvm.compile(sch.mod, target="llvm")
+
+        dev = tvm.cpu(0)
+
+        # Test boundary values: ±1.0 (should use system library)
+        boundary_values = np.array([1.0, -1.0], dtype=np.float32)
+        a1 = tvm.runtime.tensor(boundary_values, dev)
+        b1 = tvm.runtime.tensor(np.empty_like(boundary_values), dev)
+        func(a1, b1)
+        tvm.testing.assert_allclose(b1.numpy(), np_func(boundary_values), atol=1e-5, rtol=1e-5)
+
+        # Test values at threshold: ±0.5 (should use system library)
+        threshold_values = np.array([0.5, -0.5], dtype=np.float32)
+        a2 = tvm.runtime.tensor(threshold_values, dev)
+        b2 = tvm.runtime.tensor(np.empty_like(threshold_values), dev)
+        func(a2, b2)
+        tvm.testing.assert_allclose(b2.numpy(), np_func(threshold_values), atol=1e-4, rtol=1e-4)
+
+        # Test values just below threshold: ±0.49 (should use Taylor series)
+        below_threshold_values = np.array([0.49, -0.49, 0.3, -0.3, 0.0], dtype=np.float32)
+        a3 = tvm.runtime.tensor(below_threshold_values, dev)
+        b3 = tvm.runtime.tensor(np.empty_like(below_threshold_values), dev)
+        func(a3, b3)
+        tvm.testing.assert_allclose(
+            b3.numpy(), np_func(below_threshold_values), atol=1e-3, rtol=1e-3
+        )
+
+        # Test out-of-domain values: should return NaN
+        out_of_domain = np.array([1.1, -1.1, 2.0, -2.0], dtype=np.float32)
+        a4 = tvm.runtime.tensor(out_of_domain, dev)
+        b4 = tvm.runtime.tensor(np.empty_like(out_of_domain), dev)
+        func(a4, b4)
+        assert np.all(np.isnan(b4.numpy())), "Out-of-domain inputs should return NaN"
+
+    for func in test_funcs:
+        run_test(*func)
+
+
+def test_binary_intrin():
+    test_funcs = [
+        (tvm.tirx.atan2, lambda x1, x2: np.arctan2(x1, x2)),
+        (tvm.tirx.nextafter, lambda x1, x2: np.nextafter(x1, x2)),
+        (tvm.tirx.copysign, lambda x1, x2: np.copysign(x1, x2)),
+        (tvm.tirx.hypot, lambda x1, x2: np.hypot(x1, x2)),
+    ]
+
+    def run_test(tvm_intrin, np_func):
+        m = te.var(
+            "m",
+        )
+        A = te.placeholder((m,), name="A")
+        B = te.placeholder((m,), name="B")
+        C = te.compute((m,), lambda *i: tvm_intrin(A(*i), B(*i)), name="C")
+
+        # Convert to TIR and create schedule
+        mod = te.create_prim_func([A, B, C])
+        sch = tvm.s_tir.Schedule(mod)
+
+        # Build from scheduled TIR
+        func = tvm.compile(sch.mod, target="llvm")
+
+        dev = tvm.cpu(0)
+        n = 10
+        a = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(A.dtype.dtype), dev)
+        b = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(B.dtype.dtype), dev)
+        c = tvm.runtime.tensor(np.random.uniform(size=n).astype(A.dtype.dtype), dev)
+        func(a, b, c)
+        tvm.testing.assert_allclose(c.numpy(), np_func(a.numpy(), b.numpy()), atol=1e-5, rtol=1e-5)
+
+    for func in test_funcs:
+        run_test(*func)
+
+
+def test_ldexp():
+    m = te.var(
+        "m",
+    )
+    A = te.placeholder((m,), name="A")
+    B = te.placeholder((m,), name="B", dtype="int32")
+    C = te.compute((m,), lambda *i: tvm.tirx.ldexp(A(*i), B(*i)), name="C")
+
+    # Convert to TIR and create schedule
+    mod = te.create_prim_func([A, B, C])
+    sch = tvm.s_tir.Schedule(mod)
+
+    # Build from scheduled TIR
+    func = tvm.compile(sch.mod, target="llvm")
+
+    dev = tvm.cpu(0)
+    n = 10
+    a = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(A.dtype.dtype), dev)
+    b = tvm.runtime.tensor(np.random.randint(0, 5, size=n).astype(B.dtype.dtype), dev)
+    c = tvm.runtime.tensor(np.random.uniform(size=n).astype(A.dtype.dtype), dev)
+    func(a, b, c)
+    tvm.testing.assert_allclose(c.numpy(), np.ldexp(a.numpy(), b.numpy()), atol=1e-5, rtol=1e-5)
+
+
+dtype = tvm.testing.parameter("int32", "int64")
+
+
+@pytest.mark.parametrize(
+    "target",
+    ["llvm", pytest.param({"kind": "vulkan", "from_device": 0}, marks=pytest.mark.gpu)],
+)
+def test_clz(target, dtype):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+    target = tvm.target.Target(target)
+    if (
+        target.kind.name == "vulkan"
+        and dtype == "int64"
+        and not target.attrs.get("supports_int64", False)
+    ):
+        pytest.xfail("Vulkan target does not support Int64 types")
+
+    def clz_np(x, dtype):
+        ceil_log2 = np.ceil(np.log2(x)).astype(dtype)
+        bits = int(dtype[-2:])
+        clz = bits - ceil_log2
+        clz[np.bitwise_and(x, x - 1) == 0] -= 1
+        return clz
+
+    m = te.var("m")
+    A = te.placeholder((m,), name="A", dtype=dtype)
+    B = te.compute((m,), lambda *i: tvm.tirx.clz(A(*i)), name="B")
+
+    # Convert to TIR and create schedule
+    mod = te.create_prim_func([A, B])
+    sch = tvm.s_tir.Schedule(mod)
+
+    # Apply scheduling primitives if target is Vulkan
+    if target.kind.name == "vulkan":
+        block = sch.get_sblock("B")
+        loop = sch.get_loops(block)[0]
+        bx, tx = sch.split(loop, factors=[None, 64])
+        sch.bind(bx, "blockIdx.x")
+        sch.bind(tx, "threadIdx.x")
+
+    # Build from scheduled TIR
+    func = tvm.compile(sch.mod, target=target)
+
+    def run_and_check():
+        dev = tvm.device(target.kind.name)
+        n = 10
+        highs = [10, 100, 1000, 10000, 100000, 1000000]
+
+        if dtype == "int64":
+            highs.append((1 << 63) - 1)
+
+        for high in highs:
+            a_np = np.random.randint(1, high=high, size=(n,), dtype=dtype)
+            a = tvm.runtime.tensor(a_np, dev)
+            b = tvm.runtime.tensor(np.zeros((n,)).astype("int32"), dev)
+            func(a, b)
+            ref = clz_np(a_np, dtype)
+            np.testing.assert_equal(b.numpy(), ref)
+
+    if target.kind.name == "llvm":
+        run_and_check()
+    else:
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@tvm.script.ir_module
+class Module:
+    @T.prim_func(s_tir=True)
+    def test_tir_fma(A: T.handle, B: T.handle, C: T.handle, d: T.handle) -> None:
+        # function attr dict
+        T.func_attr({"global_symbol": "test_fma", "tirx.noalias": True})
+        n = T.int32()
+        stride = T.int32()
+        stride_1 = T.int32()
+        stride_2 = T.int32()
+        stride_3 = T.int32()
+        A_1 = T.match_buffer(
+            A,
+            [n],
+            strides=[stride],
+            elem_offset=0,
+            align=64,
+            offset_factor=1,
+            buffer_type="auto",
+        )
+        B_1 = T.match_buffer(
+            B,
+            [n],
+            strides=[stride_1],
+            elem_offset=0,
+            align=64,
+            offset_factor=1,
+            buffer_type="auto",
+        )
+        C_1 = T.match_buffer(
+            C,
+            [n],
+            strides=[stride_2],
+            elem_offset=0,
+            align=64,
+            offset_factor=1,
+            buffer_type="auto",
+        )
+        d_1 = T.match_buffer(
+            d,
+            [n],
+            strides=[stride_3],
+            elem_offset=0,
+            align=64,
+            offset_factor=1,
+            buffer_type="auto",
+        )
+        # body
+        for i in T.serial(0, n):
+            d_1[(i * stride_3)] = (A_1[(i * stride)] * B_1[(i * stride_1)]) + C_1[(i * stride_2)]
+
+
+def test_fma():
+    opt = tvm.transform.Sequential(
+        [
+            tvm.tirx.transform.Apply(lambda f: f.with_attr("target", tvm.target.Target("llvm"))),
+            tvm.tirx.transform.LowerIntrin(),
+        ]
+    )
+    mod = opt(Module)
+    assert mod["test_tir_fma"].body.body.value.op.name == "tirx.call_llvm_pure_intrin"
+
+
+if __name__ == "__main__":
+    test_nearbyint()
+    test_unary_intrin()
+    test_round_intrinsics_on_int()
+    test_asin_acos_boundary_values()
+    test_binary_intrin()
+    test_ldexp()
+    test_clz()
+    test_fma()

tests/python/tirx-base/test_tir_nodes.py

@@ -0,0 +1,505 @@
+# 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: F811, F841
+import numpy as np
+import pytest
+
+import tvm
+from tvm import ir
+
+
+def test_const():
+    x = tvm.tirx.const(1, "int32")
+    assert x.ty.dtype == "int32"
+    assert isinstance(x, tvm.tirx.IntImm)
+
+
+def test_te_const():
+    x = tvm.tirx.const(1, "int32")
+    assert x.ty.dtype == "int32"
+    assert isinstance(x, tvm.tirx.IntImm)
+
+
+def test_tir_const_dtype_inference():
+    for data in [
+        True,
+        bool(1),
+        np.uint8(1),
+        np.uint16(1),
+        np.uint32(1),
+        np.uint64(1),
+        np.int8(1),
+        np.int16(1),
+        np.int32(1),
+        np.int64(1),
+        np.float16(1),
+        np.float32(1),
+        np.float64(1),
+    ]:
+        assert tvm.tirx.const(data).ty.dtype == str(np.array(data).dtype)
+
+    assert tvm.tirx.const(True).ty.dtype == "bool"
+    assert tvm.tirx.const(1).ty.dtype == "int32"
+    assert tvm.tirx.const(1.0).ty.dtype == "float32"
+
+
+def test_make():
+    x = tvm.tirx.const(1, "int32")
+    y = tvm.tirx.Var("x", "int32")
+    z = x + y
+    assert isinstance(tvm.tirx.max(x, y), tvm.tirx.Max)
+    assert isinstance(tvm.tirx.min(x, y), tvm.tirx.Min)
+
+
+def test_ir():
+    x = tvm.tirx.const(1, "int32")
+    y = tvm.tirx.IntImm("int32", 1)
+    z = x + y
+    stmt = tvm.tirx.Evaluate(z)
+    assert isinstance(stmt, tvm.tirx.Evaluate)
+
+
+def test_ir2():
+    buf_size = tvm.tirx.Var("size", "int32")
+    x = tvm.tirx.Var("n", "int32")
+
+    storage_type = ir.PrimType("int32")
+    handle_type = ir.PointerType(storage_type)
+    array = tvm.tirx.Var("array", handle_type)
+    buf = tvm.tirx.decl_buffer([buf_size], "int32", data=array)
+
+    st = tvm.tirx.BufferStore(buf, x + 1, [1])
+    assert isinstance(st, tvm.tirx.BufferStore)
+    assert st.buffer == buf
+    assert st.buffer.data == array
+
+
+def test_let():
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    stmt = tvm.tirx.Bind(x, 10)
+
+
+def test_cast():
+    x = tvm.tirx.Var("x", "float32")
+    y = x.astype("int32")
+    z = x.astype("float32x4")
+    assert isinstance(y, tvm.tirx.Cast)
+    assert isinstance(z, tvm.tirx.Broadcast)
+    assert z.lanes == 4
+
+    s = tvm.tirx.StringImm("s")
+    with pytest.raises(TypeError, match="Cannot cast an expression with the void sentinel type"):
+        s.astype("int")
+
+
+def test_attr():
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    stmt = tvm.tirx.AttrStmt(y, "stride", 10, tvm.tirx.Evaluate(x + 1))
+    assert stmt.node == y
+
+    a = tvm.runtime.convert(1)
+    assert a == 1
+    try:
+        a.no_field
+        assert False
+    except AttributeError:
+        pass
+
+
+def test_basic():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    c = a + b
+    assert str(c) == f"{a.name} + {b.name}"
+
+
+def test_stmt():
+    x = tvm.tirx.Evaluate(0)
+    tvm.tirx.For(tvm.tirx.Var("i", "int32"), 0, 1, tvm.tirx.ForKind.SERIAL, x)
+    tvm.tirx.For(tvm.tirx.Var("i", "int32"), 0, 1, tvm.tirx.ForKind.UNROLLED, x, step=2)
+
+
+def test_dir():
+    x = tvm.tirx.Var("x", "int32")
+    dir(x)
+
+
+def test_dtype():
+    x = tvm.tirx.Var("x", "int32")
+    assert x.ty.dtype == "int32"
+    y = tvm.tirx.Var("y", "int32")
+    assert (x > y).ty.dtype == "bool"
+
+
+def test_any():
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    z = tvm.tirx.Var("z", "int32")
+    try:
+        t = x or x
+        assert False
+    except ValueError:
+        pass
+    try:
+        tvm.tirx.any()
+        assert False
+    except ValueError:
+        pass
+    assert str(tvm.tirx.any(x < y)) == f"{x.name} < {y.name}"
+    assert str(tvm.tirx.any(x < y, x > z)) == f"{x.name} < {y.name} or {x.name} > {z.name}"
+    assert (
+        str(tvm.tirx.any(x < y, y > z + 1, x < z * 2))
+        == f"{x.name} < {y.name} or {y.name} > {z.name} + 1 or {x.name} < {z.name} * 2"
+    )
+
+
+def test_all():
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    z = tvm.tirx.Var("z", "int32")
+    try:
+        t = x and x
+        assert False
+    except ValueError:
+        pass
+    try:
+        tvm.tirx.all()
+        assert False
+    except ValueError:
+        pass
+    assert str(tvm.tirx.all(x < y)) == f"{x.name} < {y.name}"
+    assert str(tvm.tirx.all(x < y, x > z)) == f"{x.name} < {y.name} and {x.name} > {z.name}"
+    assert (
+        str(tvm.tirx.all(x < y, y > z + 1, x < z * 2))
+        == f"{x.name} < {y.name} and {y.name} > {z.name} + 1 and {x.name} < {z.name} * 2"
+    )
+
+
+def test_bitwise():
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    assert str(x << y) == "T.shift_left(x, y)"
+    assert str(x >> y) == "T.shift_right(x, y)"
+    assert str(x & y) == "T.bitwise_and(x, y)"
+    assert str(x | y) == "T.bitwise_or(x, y)"
+    assert str(x ^ y) == "T.bitwise_xor(x, y)"
+    assert str(10 & x) == "T.bitwise_and(10, x)"
+    assert str(10 | x) == "T.bitwise_or(10, x)"
+    assert str(10 ^ x) == "T.bitwise_xor(10, x)"
+    assert str(10 >> x) == "T.shift_right(10, x)"
+    assert str(10 << x) == "T.shift_left(10, x)"
+    assert str(10 % x) == "10 % x"
+
+    assert str(~x) == "T.bitwise_not(x)"
+    assert (tvm.tirx.const(1, "int8x2") >> 1).ty.dtype == "int8x2"
+    assert (x >> tvm.tirx.const(1, "int32x2")).ty.dtype == "int32x2"
+    assert (tvm.tirx.Var("z", "int8x2") << tvm.tirx.const(1, "int8x2")).ty.dtype == "int8x2"
+
+
+def test_float_bitwise():
+    t = tvm.tirx.const(1.5, dtype="float32")
+    for test in [
+        lambda lhs, rhs: lhs << rhs,
+        lambda lhs, rhs: lhs >> rhs,
+        lambda lhs, rhs: lhs | rhs,
+        lambda lhs, rhs: lhs ^ rhs,
+        lambda lhs, rhs: lhs & rhs,
+    ]:
+        try:
+            test(t, 10.0)
+            assert False
+        except RuntimeError:
+            pass
+    try:
+        ~t
+        assert False
+    except RuntimeError:
+        pass
+
+
+def test_shift_bounds():
+    x = tvm.tirx.Var("x", "int32")
+    for test in [lambda lhs, rhs: lhs << rhs, lambda lhs, rhs: lhs >> rhs]:
+        # negative case
+        for testcase in [(x, -1), (x, 32)]:
+            try:
+                test(*testcase)
+                assert False
+            except RuntimeError:
+                pass
+
+        # positive case
+        for testcase in [(x, 0), (x, 16), (x, 31)]:
+            test(*testcase)
+
+
+def test_divide_by_zero():
+    for test in [
+        lambda lhs, rhs: tvm.tirx.floormod(lhs, rhs),
+        lambda lhs, rhs: tvm.tirx.floordiv(lhs, rhs),
+        lambda lhs, rhs: tvm.tirx.truncmod(lhs, rhs),
+        lambda lhs, rhs: tvm.tirx.truncdiv(lhs, rhs),
+        lambda lhs, rhs: tvm.tirx.div(lhs, rhs),
+    ]:
+        try:
+            test(tvm.tirx.const(5, "int32"), tvm.tirx.const(0, "int32"))
+            assert False
+        except RuntimeError:
+            pass
+
+
+def test_infinity():
+    assert str(tvm.tirx.infinity("float16")) == 'T.float16("inf")'
+    assert str(tvm.tirx.infinity("float32")) == 'T.float32("inf")'
+    assert str(tvm.tirx.infinity("float64")) == 'T.float64("inf")'
+
+
+def test_isnan():
+    x = tvm.tirx.Var("x", "float32")
+    assert str(tvm.tirx.isnan(x)) == "T.isnan(x)"
+    assert str(tvm.tirx.isnan(x).ty.dtype) == "bool"
+    y = tvm.tirx.Var("y", "float16")
+    assert str(tvm.tirx.isnan(y)) == 'T.isnan(T.Cast("float32", y))'
+    z = tvm.tirx.Var("z", "int32")
+    assert str(tvm.tirx.isnan(z)) == "T.bool(False)"
+    k = tvm.tirx.Var("k", "int8x2")
+    assert str(tvm.tirx.isnan(k).ty.dtype) == "boolx2"
+
+
+def test_equality():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    c = a == b
+    assert not c
+    d = c != c
+    assert not d
+
+
+def test_equality_string_imm():
+    x = "a"
+    y = tvm.tirx.StringImm(x)
+    x == y.value
+    x == y
+
+
+def test_prim_func():
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    b = tvm.tirx.decl_buffer((x,), "float32")
+    stmt = tvm.tirx.SeqStmt([tvm.tirx.Bind(x, 10), tvm.tirx.Evaluate(x + 1)])
+
+    func = tvm.tirx.PrimFunc([x, y, b], stmt)
+    # make sure we can print
+    assert func.buffer_map[func.params[2]].same_as(b)
+
+    assert len(func.buffer_map) == 1
+    f2 = func.with_attr({"calling_conv": 1, "tirx.noalias": True})
+    assert f2.attrs["calling_conv"] == 1
+    assert not func.attrs
+
+
+def test_vars():
+    x = tvm.tirx.Var("xyz", "int8")
+    assert x.ty.dtype == "int8"
+    ptype = tvm.ir.PointerType(tvm.ir.PrimType("float"))
+    x = tvm.tirx.Var("xyz", ptype)
+    assert x.ty == ptype
+    assert isinstance(ptype.element_type, tvm.ir.PrimType)
+
+
+def test_scoped_storage_vars():
+    dtype = "float"
+    storage_scope = "global.texture"
+    ptype = tvm.ir.PointerType(tvm.ir.PrimType(dtype), storage_scope)
+    x = tvm.tirx.Var("xyz", ptype)
+    assert x.ty == ptype
+    assert x.ty.storage_scope == storage_scope
+    assert isinstance(ptype.element_type, tvm.ir.PrimType)
+
+
+def test_buffer_load_store():
+    b = tvm.tirx.decl_buffer((10,), "float32")
+    x = tvm.tirx.BufferLoad(b, [0])
+    assert isinstance(x, tvm.tirx.BufferLoad)
+    assert x.ty.dtype == "float32"
+    assert x.buffer == b
+    s = tvm.tirx.BufferStore(b, 0.1, [0])
+    assert isinstance(s, tvm.tirx.BufferStore)
+
+
+def test_intimm_cond():
+    x = tvm.runtime.convert(1)
+    y = tvm.runtime.convert(1)
+    s = {x}
+    assert y in s
+    assert x == y
+    assert x < 20
+    assert not (x >= 20)
+    assert x < 10 and y < 10
+    assert not tvm.runtime.convert(x != 1)
+    assert x == 1
+
+
+def _create_ramp(lanes):
+    return tvm.tirx.Ramp(0, 1, lanes)
+
+
+def _create_broadcast(lanes):
+    return tvm.tirx.Broadcast(0, lanes)
+
+
+@pytest.mark.parametrize("lanes", [tvm.tirx.IntImm(dtype="int64", value=11)])
+@pytest.mark.parametrize("node_func", [_create_ramp, _create_broadcast])
+def test_lane_types(lanes, node_func):
+    def _check_dtype(node):
+        assert node.lanes.ty.dtype == "int32"
+        assert node.lanes == 11
+
+    _check_dtype(node_func(lanes))
+
+
+@pytest.mark.parametrize("lanes", [(11 * tvm.tirx.vscale()), (tvm.tirx.vscale() * 11)])
+@pytest.mark.parametrize("node_func", [_create_ramp, _create_broadcast])
+def test_scalable_vec(lanes, node_func):
+    def _check_dtype(node):
+        assert node.lanes.a.equal(tvm.tirx.vscale())
+        assert node.lanes.b == 11
+
+    _check_dtype(node_func(lanes))
+
+
+@pytest.mark.parametrize(
+    "lanes", [(tvm.tirx.vscale()), (tvm.tirx.vscale() + 3), (tvm.tirx.vscale() * 2 + 5)]
+)
+@pytest.mark.parametrize("node_func", [_create_ramp, _create_broadcast])
+def test_scalable_vec_error(lanes, node_func):
+    with pytest.raises(RuntimeError):
+        node_func(lanes)
+
+
+def test_broadcast_to_scalable_vec():
+    vec = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale()) + 3
+    broadcast = vec.b
+
+    assert isinstance(broadcast, tvm.tirx.expr.Broadcast)
+    assert broadcast.value == 3
+    assert broadcast.lanes.a.equal(tvm.tirx.vscale())
+    assert broadcast.lanes.b == 4
+
+
+def test_buffer_load_scalable_vec():
+    buf = tvm.tirx.decl_buffer((24,), "float32")
+    index = tvm.tirx.expr.Ramp(1, 1, 8 * tvm.tirx.vscale())
+    load = tvm.tirx.BufferLoad(buf, [index])
+
+    assert isinstance(load, tvm.tirx.BufferLoad)
+    assert load.ty.dtype == "float32xvscalex8"
+
+
+def test_buffer_store_scalable_vec():
+    b = tvm.tirx.decl_buffer((24,), "int32")
+    value = tvm.tirx.expr.Broadcast(1, 4 * tvm.tirx.vscale())
+    index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
+    store = tvm.tirx.BufferStore(b, value, [index])
+
+    assert isinstance(store, tvm.tirx.BufferStore)
+    assert store.value.ty.dtype == "int32xvscalex4"
+
+
+def test_buffer_store_predicate_invalid_scalability():
+    b = tvm.tirx.decl_buffer((24,), "int32")
+    value = tvm.tirx.expr.Broadcast(1, 4 * tvm.tirx.vscale())
+    index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
+    predicate = tvm.tirx.expr.Broadcast(tvm.tirx.IntImm("int1", 1), 4)
+
+    err_msg = "Predicate mask dtype and value dtype must both be scalable."
+    with pytest.raises(RuntimeError, match=err_msg):
+        tvm.tirx.BufferStore(b, value, [index], predicate)
+
+
+def test_buffer_store_predicate_invalid_lanes():
+    b = tvm.tirx.decl_buffer((24,), "int32")
+    value = tvm.tirx.expr.Broadcast(1, 4 * tvm.tirx.vscale())
+    index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
+    predicate = tvm.tirx.expr.Broadcast(tvm.tirx.IntImm("int1", 1), 8 * tvm.tirx.vscale())
+
+    err_msg = (
+        "Got a predicate mask with 8 lanes, but trying to store a "
+        "value with 4 lanes. The number of lanes must match."
+    )
+    with pytest.raises(RuntimeError, match=err_msg):
+        tvm.tirx.BufferStore(b, value, [index], predicate)
+
+
+def test_buffer_store_predicate_elements_invalid_type():
+    b = tvm.tirx.decl_buffer((24,), "int32")
+    value = tvm.tirx.expr.Broadcast(1, 4 * tvm.tirx.vscale())
+    index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
+    predicate = tvm.tirx.expr.Broadcast(1, 4 * tvm.tirx.vscale())
+
+    err_msg = "Predicate mask elements must be boolean values, but got int32."
+    with pytest.raises(RuntimeError, match=err_msg):
+        tvm.tirx.BufferStore(b, value, [index], predicate)
+
+
+def test_buffer_load_predicate_elements_invalid_type():
+    b = tvm.tirx.decl_buffer((24,), "int32")
+    index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
+    predicate = tvm.tirx.expr.Broadcast(1, 4 * tvm.tirx.vscale())
+
+    err_msg = "Predicate mask elements must be boolean values, but got int32."
+    with pytest.raises(RuntimeError, match=err_msg):
+        tvm.tirx.BufferLoad(b, [index], predicate)
+
+
+def test_buffer_store_predicate_invalid_scalability():
+    b = tvm.tirx.decl_buffer((24,), "int32")
+    index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
+    predicate = tvm.tirx.expr.Broadcast(tvm.tirx.IntImm("int1", 1), 4)
+
+    err_msg = "Predicate mask dtype and load indices must both be scalable."
+    with pytest.raises(RuntimeError, match=err_msg):
+        tvm.tirx.BufferLoad(b, [index], predicate)
+
+
+def test_buffer_store_predicate_invalid_lanes():
+    b = tvm.tirx.decl_buffer((24,), "int32")
+    index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
+    predicate = tvm.tirx.expr.Broadcast(tvm.tirx.IntImm("int1", 1), 8 * tvm.tirx.vscale())
+
+    err_msg = (
+        "Got a predicate mask with 8 lanes, but trying to load a "
+        "vector with 4 lanes. The number of lanes must match."
+    )
+    with pytest.raises(RuntimeError, match=err_msg):
+        tvm.tirx.BufferLoad(b, [index], predicate)
+
+
+def test_scalable_vec_cast():
+    b = tvm.tirx.decl_buffer((24,), "float32")
+    value = tvm.tirx.expr.Broadcast(1, 12 * tvm.tirx.vscale()).astype("float32xvscalex12")
+    index = tvm.tirx.expr.Ramp(0, 1, 12 * tvm.tirx.vscale())
+
+    store = tvm.tirx.BufferStore(b, value, [index])
+
+    assert isinstance(store.value.value, tvm.tirx.expr.FloatImm)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-base/test_tir_op_types.py

@@ -0,0 +1,347 @@
+# 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.
+# pylint: disable=missing-docstring
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.backend.cuda import op as _cuda_op
+
+
+def test_tir_op_tvm_tuple():
+    x = tirx.Var("x", dtype="float32")
+    y = tirx.Var("y", dtype="float32")
+    z = tirx.Var("z", dtype="float32")
+    expr = tirx.tvm_tuple(x, y, z, 1, 2, 3)
+    assert expr.op.name == "tirx.tvm_tuple"
+
+
+def test_tir_op_tvm_struct_get():
+    x = tirx.Var("x", dtype="handle")
+    expr = tirx.tvm_struct_get(x, 1, 2, dtype="int32")
+    assert expr.op.name == "tirx.tvm_struct_get"
+
+
+def test_tir_op_tvm_struct_set():
+    x = tirx.Var("x", dtype="handle")
+    expr = tirx.tvm_struct_set(x, 1, 2, 3)
+    assert expr.op.name == "tirx.tvm_struct_set"
+
+
+def test_tir_op_address_of():
+    buffer = tirx.decl_buffer((128), "float32")
+    expr = tirx.address_of(buffer[0])
+    assert expr.op.name == "tirx.address_of"
+    scalar_address = tirx.address_of(tirx.Var("value", "uint32"))
+    assert scalar_address.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint32"))
+
+
+def test_tir_op_trace_pointer():
+    pointer = tirx.Var("pointer", tvm.ir.PointerType(tvm.ir.PrimType("float32")))
+    traced = tirx.trace([pointer])
+    assert traced.ty == pointer.ty
+
+
+def test_tir_op_lookup_param():
+    expr = tirx.lookup_param("p0")
+    assert expr.op.name == "tirx.lookup_param"
+
+
+def test_tir_op_reinterpret():
+    x = tirx.Var("x", dtype="int32")
+    expr = tirx.reinterpret("float32", x)
+    assert expr.op.name == "tirx.reinterpret"
+    with pytest.raises(TypeError, match="scalar 64-bit integer source"):
+        tirx.reinterpret("handle", x)
+    pointer = tirx.reinterpret("handle", tirx.Var("address", dtype="uint64"))
+    assert pointer.ty == tvm.ir.PointerType(tvm.ir.PrimType("void"))
+
+
+def test_tir_op_isnullptr():
+    x = tirx.Var("x", dtype="int32")
+    expr = tirx.isnullptr(x)
+    assert expr.op.name == "tirx.isnullptr"
+
+
+def test_tir_op_call_assume():
+    x = tirx.Var("x", dtype="int32")
+    expr = tirx.assume(cond=x)
+    assert expr.op.name == "tirx.assume"
+
+
+def test_tir_op_call_undef():
+    expr = tirx.undef()
+    assert expr.op.name == "tirx.undef"
+
+
+def test_tir_op_call_likely():
+    x = tirx.Var("x", dtype="int32")
+    expr = tirx.likely(cond=x)
+    assert expr.op.name == "tirx.likely"
+
+
+def test_tir_op_tvm_thread_allreduce():
+    x = tirx.Var("x", "int32")
+    buffer = tirx.decl_buffer((128), "float32")
+    y = tirx.Var("y", "handle")
+    z = tirx.Var("z", "int32")
+    expr = tirx.tvm_thread_allreduce(x, buffer[0], True, y, z)
+    assert expr.op.name == "tirx.tvm_thread_allreduce"
+
+
+def test_tir_op_type_annotation():
+    expr = tirx.type_annotation("int32")
+    assert expr.op.name == "tirx.type_annotation"
+
+
+def test_tir_op_tvm_access_ptr():
+    buffer = tirx.decl_buffer((128), "float32")
+    expr = tirx.tvm_access_ptr("float32", buffer.data, 0, 1, 2)
+    assert expr.op.name == "tirx.tvm_access_ptr"
+    assert expr.ty == tvm.ir.PointerType(tvm.ir.PrimType("float32"))
+    offset_expr = tirx.ptr_byte_offset(buffer.data, 16, "uint8")
+    assert offset_expr.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint8"))
+
+
+def test_tir_op_tvm_throw_last_error():
+    expr = tirx.tvm_throw_last_error()
+    assert expr.op.name == "tirx.tvm_throw_last_error"
+
+
+def test_tir_op_tvm_load_matrix_sync():
+    buffer = tirx.decl_buffer((16, 16), "float32")
+    x = tirx.Var("x", "handle")
+    expr = tirx.tvm_load_matrix_sync(buffer.data, 16, 16, 16, 0, x, 128, "row_major")
+    assert expr.op.name == "tirx.tvm_load_matrix_sync"
+
+
+def test_tir_op_tvm_store_matrix_sync():
+    buffer = tirx.decl_buffer((16, 16), "float32")
+    x = tirx.Var("x", "handle")
+    expr = tirx.tvm_store_matrix_sync(buffer.data, 16, 16, 16, 0, x, 128, "row_major")
+    assert expr.op.name == "tirx.tvm_store_matrix_sync"
+
+
+def test_tir_op_tvm_mma_sync():
+    buffer_0 = tirx.decl_buffer((16, 16), "float32")
+    buffer_1 = tirx.decl_buffer((16, 16), "float32")
+    buffer_2 = tirx.decl_buffer((16, 16), "float32")
+    buffer_3 = tirx.decl_buffer((16, 16), "float32")
+    expr = tirx.tvm_mma_sync(buffer_0.data, 0, buffer_1.data, 0, buffer_2.data, 0, buffer_3.data, 0)
+    assert expr.op.name == "tirx.tvm_mma_sync"
+
+
+def test_tir_op_tvm_bmma_sync():
+    buffer_0 = tirx.decl_buffer((16, 16), "float32")
+    buffer_1 = tirx.decl_buffer((16, 16), "float32")
+    buffer_2 = tirx.decl_buffer((16, 16), "float32")
+    buffer_3 = tirx.decl_buffer((16, 16), "float32")
+    expr = tirx.tvm_bmma_sync(
+        buffer_0.data, 0, buffer_1.data, 0, buffer_2.data, 0, buffer_3.data, 0
+    )
+    assert expr.op.name == "tirx.tvm_bmma_sync"
+
+
+def test_tir_op_tvm_fill_fragment():
+    buffer = tirx.decl_buffer((16, 16), "float32")
+    expr = tirx.tvm_fill_fragment(buffer.data, 16, 16, 16, 0, 0)
+    assert expr.op.name == "tirx.tvm_fill_fragment"
+
+
+def test_tir_op_ptx_mma():
+    buffer_a = tirx.decl_buffer([32], "int4", scope="local")
+    buffer_b = tirx.decl_buffer([16], "uint4", scope="local")
+    buffer_c = tirx.decl_buffer([4], "int32", scope="local")
+    expr = _cuda_op.ptx_mma_legacy(
+        "m8n8k32",
+        "row",
+        "col",
+        "int4",
+        "uint4",
+        "int32",
+        buffer_a.data,
+        0,
+        buffer_b.data,
+        0,
+        buffer_c.data,
+        0,
+        False,
+    )
+    assert expr.op.name == "tirx.ptx.mma_legacy"
+
+
+def test_tir_op_ptx_mma_sp():
+    buffer_a = tirx.decl_buffer([32], "int4", scope="local")
+    buffer_b = tirx.decl_buffer([16], "uint4", scope="local")
+    buffer_c = tirx.decl_buffer([4], "int32", scope="local")
+    buffer_d = tirx.decl_buffer([1], "uint32", scope="local")
+    expr = _cuda_op.ptx_mma_sp_legacy(
+        "m8n8k32",
+        "row",
+        "col",
+        "int4",
+        "uint4",
+        "int32",
+        buffer_a.data,
+        0,
+        buffer_b.data,
+        0,
+        buffer_c.data,
+        0,
+        buffer_d.data,
+        0,
+        0,
+        False,
+    )
+    assert expr.op.name == "tirx.ptx.mma_sp"
+
+
+def test_tir_op_mma_store():
+    x = tirx.Var("x", dtype="int32")
+    y = tirx.Var("y", dtype="int32")
+    buffer_w = tirx.decl_buffer([16, 8], dtype="int32", scope="warp", offset_factor=1)
+    buffer = tirx.decl_buffer(
+        [16, 16], dtype="int32", scope="global", offset_factor=1, strides=[x, y]
+    )
+    expr = _cuda_op.mma_store(
+        "int32",
+        16,
+        16,
+        buffer.access_ptr("w"),
+        buffer_w.data,
+        buffer_w.elem_offset,
+        x,
+    )
+    assert expr.op.name == "tirx.mma_store"
+
+
+def test_tir_op_mma_fill():
+    buffer_w = tirx.decl_buffer([16, 8], dtype="int32", scope="warp", offset_factor=1)
+    expr = _cuda_op.mma_fill("int32", 8, buffer_w.data, buffer_w.elem_offset)
+    assert expr.op.name == "tirx.mma_fill"
+
+
+def test_op_ptx_ldmatrix():
+    buffer_shared = tirx.decl_buffer([16, 16], "float16", scope="shared")
+    buffer_local = tirx.decl_buffer([8], "float16", scope="local")
+    # New API: 4 scatter-form dst handles for .x4.b16 (one per output register).
+    expr = _cuda_op.ptx_ldmatrix(
+        False,
+        4,
+        ".b16",
+        buffer_shared.data,
+        buffer_local.data,
+        buffer_local.data,
+        buffer_local.data,
+        buffer_local.data,
+    )
+    assert expr.op.name == "tirx.ptx.ldmatrix"
+
+
+def test_op_ptx_cp_async():
+    buffer_shared = tirx.decl_buffer([16, 16], "float16", scope="shared")
+    buffer_local = tirx.decl_buffer([8], "float16", scope="local")
+    expr = _cuda_op.ptx_cp_async_legacy(buffer_shared.data, 0, buffer_local.data, 0, 16)
+    assert expr.op.name == "tirx.ptx.cp_async"
+
+    inner_dst = tirx.tvm_access_ptr("float16", buffer_shared.data, 2, 8, 1)
+    inner_src = tirx.tvm_access_ptr("float16", buffer_local.data, 4, 8, 1)
+    expr = _cuda_op.ptx_cp_async_legacy("float16", inner_dst, 3, inner_src, 5, 16)
+    for access_ptr, expected_offset in zip(expr.args[:2], [5, 9]):
+        assert access_ptr.op.name == "tirx.tvm_access_ptr"
+        assert isinstance(access_ptr.args[1], tirx.Var)
+        simplified_offset = tvm.arith.Analyzer().simplify(access_ptr.args[2])
+        assert int(simplified_offset) == expected_offset
+
+
+def test_op_ptx_cp_async_bulk():
+    buffer_shared = tirx.decl_buffer([16, 16], "float16", scope="shared")
+    buffer_local = tirx.decl_buffer([8], "float16", scope="local")
+    expr = _cuda_op.ptx_cp_async_bulk("float16", buffer_shared.data, 0, buffer_local.data, 0, 16, 0)
+    assert expr.op.name == "tirx.ptx.cp_async_bulk"
+
+
+def test_tir_op_vectorlow():
+    buffer = tirx.decl_buffer((4, 4), "int8", offset_factor=1)
+    vec = buffer.vload([0, 0], dtype="int8x16")
+    expr = tirx.vectorlow("int8x8", vec)
+    assert expr.op.name == "tirx.vectorlow"
+
+
+def test_tir_op_vectorhigh():
+    buffer = tirx.decl_buffer((4, 4), "int8", offset_factor=1)
+    vec = buffer.vload([0, 0], dtype="int8x16")
+    expr = tirx.vectorhigh("int8x8", vec)
+    assert expr.op.name == "tirx.vectorhigh"
+
+
+def test_tir_op_dp4a():
+    vec1 = tirx.Var("vec1", dtype="int8x4")
+    vec2 = tirx.Var("vec2", dtype="int8x4")
+    acc = tirx.Var("acc", dtype="int32")
+    expr = tirx.dp4a(vec1, vec2, acc)
+    assert expr.op.name == "tirx.dp4a"
+
+
+def test_tir_op_vectorcombine():
+    buffer = tirx.decl_buffer((4, 4), "int8", offset_factor=1)
+    vec = buffer.vload([0, 0], dtype="int8x16")
+    expr = tirx.vectorcombine("int8x8", vec, vec)
+    assert expr.op.name == "tirx.vectorcombine"
+
+
+def test_tir_op_shift_left():
+    x = tirx.Var("x", dtype="int32")
+    y = tirx.Var("x", dtype="int32")
+    expr = tirx.shift_left(x, y)
+    assert expr.op.name == "tirx.shift_left"
+
+
+def test_tir_op_shift_right():
+    x = tirx.Var("x", dtype="int32")
+    y = tirx.Var("x", dtype="int32")
+    expr = tirx.shift_right(x, y)
+    assert expr.op.name == "tirx.shift_right"
+
+
+def test_tir_op_bitwise():
+    x = tirx.Var("x", dtype="int32")
+    y = tirx.Var("y", dtype="int32")
+    expr = tirx.bitwise_and(x, y)
+    assert expr.op.name == "tirx.bitwise_and"
+    expr = tirx.bitwise_or(x, y)
+    assert expr.op.name == "tirx.bitwise_or"
+    expr = tirx.bitwise_not(x)
+    assert expr.op.name == "tirx.bitwise_not"
+    expr = tirx.bitwise_xor(x, y)
+    assert expr.op.name == "tirx.bitwise_xor"
+
+
+def test_tir_op_TVMBackendAllocWorkspace():
+    expr = tirx.TVMBackendAllocWorkspace(0, 1, 2, 3, 4)
+    assert expr.op.name == "tirx.TVMBackendAllocWorkspace"
+
+
+def test_tir_op_TVMBackendFreeWorkspace():
+    buffer = tirx.decl_buffer((128), "float32")
+    expr = tirx.TVMBackendFreeWorkspace(0, 1, buffer.data)
+    assert expr.op.name == "tirx.TVMBackendFreeWorkspace"
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-base/test_tir_ops.py

@@ -0,0 +1,245 @@
+# 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
+import tvm.testing
+
+
+def check_throws(f):
+    try:
+        f()
+    except RuntimeError:
+        pass
+    else:
+        raise AssertionError("Should have raised an exception but didn't.")
+
+
+def test_const_fold():
+    def check(f, *args):
+        x = f(*[tvm.tirx.const(x, "int32") for x in args])
+        y = f(*args)
+        if not isinstance(x, tvm.tirx.IntImm) or x.value != int(y):
+            raise ValueError(f"check error: {x} vs {y} ")
+
+    tmod = tvm.tirx.truncmod
+    check(lambda x, y: x + y, 3, 4)
+    check(lambda x, y: x * y, 3, 12)
+    check(lambda x, y: x * y - 10, 3, 12)
+    check(lambda x, y: x - tmod(y, 10), 3, 12)
+    check(lambda x, y: x // y + 10, 100, 12)
+    check(lambda x, y: x & y + 10, 112, 128)
+    check(lambda x, y: x > y, 112, 128)
+    check(lambda x, y: x < y, 112, 128)
+    check(lambda x, y: x <= y, 112, 128)
+    check(lambda x, y: x >= y, 112, 128)
+    check(lambda x, y: (x | y) ^ 10, 112, 128)
+
+
+def test_const_fold2():
+    x = tvm.tirx.Var("x", "int32")
+    tmod = tvm.tirx.truncmod
+    tdiv = tvm.tirx.truncdiv
+    assert (x + 0).same_as(x)
+    assert (0 + x).same_as(x)
+    assert (x - 0).same_as(x)
+    assert tmod(x, 1).value == 0
+    assert (x * 1).same_as(x)
+    assert (1 * x).same_as(x)
+    assert isinstance(tdiv(1, x), tvm.tirx.Div)
+
+
+def test_const_fold3():
+    # Test that using ints with logic operations is forbidden
+    x = tvm.tirx.Var("x", "int32")
+    for val in [0, 1]:
+        for func in [tvm.tirx.all, tvm.tirx.any]:
+            check_throws(lambda: func(tvm.tirx.const(val, "bool"), x))
+            check_throws(lambda: func(x, tvm.tirx.const(val, "bool")))
+
+    # Test const folding when both arguments are const
+    for tvm_func, py_func in [
+        (tvm.tirx.all, lambda a, b: a and b),
+        (tvm.tirx.any, lambda a, b: a or b),
+    ]:
+        for v1 in [0, 1]:
+            for v2 in [0, 1]:
+                tvm.ir.assert_structural_equal(
+                    tvm_func(tvm.tirx.const(v1, "bool"), tvm.tirx.const(v2, "bool")),
+                    tvm.tirx.const(py_func(v1, v2), "bool"),
+                )
+
+    x = tvm.tirx.Var("x", "bool")
+    true = tvm.tirx.const(1, "bool")
+    false = tvm.tirx.const(0, "bool")
+
+    assert tvm.tirx.all(x, true).same_as(x)
+    assert tvm.tirx.all(true, x).same_as(x)
+    assert tvm.tirx.any(x, false).same_as(x)
+    assert tvm.tirx.any(false, x).same_as(x)
+
+    assert tvm.tirx.all(x, false).same_as(false)
+    assert tvm.tirx.all(false, x).same_as(false)
+    assert tvm.tirx.any(x, true).same_as(true)
+    assert tvm.tirx.any(true, x).same_as(true)
+
+
+def test_const_fold4():
+    x1 = tvm.tirx.const(4, "int32")
+    x2 = x1 + 5
+    tdiv = tvm.tirx.truncdiv
+    assert isinstance(x2, tvm.tirx.IntImm) and x2.value == 9
+    x3 = tdiv(x2, 3)
+    assert isinstance(x3, tvm.tirx.IntImm) and x3.value == 3
+    x4 = x3 + 0.55
+    assert isinstance(x4, tvm.tirx.FloatImm) and abs(x4.value - 3.55) < 1e-6
+    x5 = tvm.tirx.ceil(x4)
+    assert isinstance(x5, tvm.tirx.FloatImm) and x5.value == 4
+    x6 = x5.astype("int")
+    assert isinstance(x6, tvm.tirx.IntImm) and x6.value == 4, f"x6={x6}"
+    y = (tvm.tirx.round((tvm.tirx.const(6.5, "float32") - 1) / 1.5) + 2).astype("int")
+    assert isinstance(y, tvm.tirx.IntImm) and y.value == 6
+
+
+def test_binary_dtype_match():
+    def verify_general_dtype_support(f, is_conditional=False):
+        rules = [
+            [("bool", "int32"), "int32"],
+            [("int32", "float32"), "float32"],
+            [("int32", "int64"), "int64"],
+            [("uint32", "int8"), "uint32"],
+            [("uint32", "int32"), "uint32"],
+        ]
+        for (lhs_dtype, rhs_dtype), out_dtype in rules:
+            lhs = tvm.tirx.Var("lhs", lhs_dtype)
+            rhs = tvm.tirx.Var("rhs", rhs_dtype)
+            out = f(lhs, rhs)
+            if not is_conditional:
+                assert out.ty.dtype == out_dtype
+            else:
+                assert out.ty.dtype == "bool"
+            if hasattr(out, "a"):
+                assert out.a.ty.dtype == out_dtype
+                assert out.b.ty.dtype == out_dtype
+            elif hasattr(out, "args"):
+                # CallOp
+                assert out.args[0].ty.dtype == out_dtype
+                assert out.args[1].ty.dtype == out_dtype
+            else:
+                raise ValueError("Unknown binary op format!")
+
+    def verify_callop_float_only(f):
+        for lhs_dtype in ["int32", "float32", "float64"]:
+            for rhs_dtype in ["int32", "float32", "float64"]:
+                lhs = tvm.tirx.Var("lhs", lhs_dtype)
+                rhs = tvm.tirx.Var("rhs", rhs_dtype)
+                if "float" not in lhs_dtype and "float" not in rhs_dtype:
+                    check_throws(lambda: f(lhs, rhs))
+                elif "float" in lhs_dtype:
+                    out = f(lhs, rhs)
+
+                    # Upcasting for floating point types
+                    dtypes = [lhs_dtype, rhs_dtype]
+                    if "float64" in dtypes:
+                        target_dtype = "float64"
+                    elif "float32" in dtypes:
+                        target_dtype = "float32"
+                    else:
+                        target_dtype = "int32"
+                    assert out.ty.dtype == target_dtype
+
+                    # Final inputs are the right type
+                    assert out.args[0].ty.dtype == target_dtype
+                    assert out.args[1].ty.dtype == target_dtype
+                else:
+                    out = f(lhs, rhs)
+                    assert out.ty.dtype == rhs_dtype
+                    assert out.args[0].ty.dtype == rhs_dtype
+                    assert out.args[1].ty.dtype == rhs_dtype
+
+    verify_general_dtype_support(lambda a, b: a + b)
+    verify_general_dtype_support(lambda a, b: a * b)
+    verify_general_dtype_support(lambda a, b: a >= b, is_conditional=True)
+    verify_general_dtype_support(lambda a, b: a <= b, is_conditional=True)
+    verify_callop_float_only(lambda a, b: tvm.tirx.power(a, b))
+
+    # verify bool & int32 constant folding
+    assert tvm.tirx.const(1) == tvm.tirx.const(True)
+    assert tvm.tirx.const(2) != tvm.tirx.const(True)
+
+
+def test_if_then_else():
+    cases = [
+        [(tvm.tirx.Var("cond", "bool"), "bool", "int32"), "int32"],
+        [(True, "int32", "float32"), "float32"],
+        [(False, "int32", "int64"), "int64"],
+        [(tvm.tirx.Var("cond", "bool"), "uint32", "int32"), "uint32"],
+        [(tvm.tirx.Var("cond", "int32"), "uint32", "int32"), "uint32"],
+    ]
+    for (cond, lhs_dtype, rhs_dtype), out_dtype in cases:
+        lhs = tvm.tirx.Var("lhs", lhs_dtype)
+        rhs = tvm.tirx.Var("rhs", rhs_dtype)
+        if cond is True or cond is False:
+            out = tvm.tirx.if_then_else(cond, lhs, rhs)
+            out2 = tvm.tirx.if_then_else(not cond, rhs, lhs)
+            out3 = tvm.tirx.if_then_else(not cond, lhs, rhs)
+            tvm.ir.assert_structural_equal(out, out2) == 1
+            if cond:
+                tvm.ir.assert_structural_equal(out, lhs.astype(out_dtype)) == 1
+                tvm.ir.assert_structural_equal(out3, rhs.astype(out_dtype)) == 1
+            else:
+                tvm.ir.assert_structural_equal(out, rhs.astype(out_dtype)) == 1
+                tvm.ir.assert_structural_equal(out3, lhs.astype(out_dtype)) == 1
+        elif cond.ty.dtype == "bool":
+            out = tvm.tirx.if_then_else(cond, lhs, rhs)
+            assert out.ty.dtype == out_dtype
+            assert out.args[1].ty.dtype == out_dtype
+            assert out.args[2].ty.dtype == out_dtype
+        elif cond.ty.dtype != "bool":
+            check_throws(lambda: tvm.tirx.if_then_else(cond, lhs, rhs))
+        else:
+            raise ValueError("Unknown combinations")
+
+
+@pytest.mark.parametrize("num_args", list(range(2, 10)))
+def test_comm_reducer(num_args):
+    """Handle all arguments in tirx comm_reducer
+
+    The `tirx.comm_reducer` API has two distinct usages.  It can reduce
+    a tensor along a specified axis, similar to numpy.max, or it can
+    reduce several arguments together, simililar to Python's built-in
+    max().  This choice is based on the type of the second argument.
+
+    If the `tirx.comm_reducer` is reducing all arguments, then all
+    arguments should be used.  In the past, the introduction of new
+    arguments intended for use when reducing along a tensor axis has
+    failed to forward these arguments when reducing along a list of
+    items.
+    """
+    assert tvm.tirx.max(*range(num_args)) == num_args - 1
+
+
+def test_llvm_intrin():
+    with pytest.raises(ValueError, match=r"Unknown llvm intrinsic function llvm.dummy"):
+        a = tvm.tirx.call_llvm_intrin("int32x4", "llvm.dummy")
+    with pytest.raises(ValueError, match=r"Unknown llvm intrinsic function llvm.dummy"):
+        a = tvm.tirx.call_llvm_pure_intrin("int32x4", "llvm.dummy")
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-base/test_tir_ptx_cp_async.py

@@ -0,0 +1,81 @@
+# 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 numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@T.prim_func(s_tir=True)
+def ptx_cp_async(A: T.Buffer((32, 128), "float16"), B: T.Buffer((32, 128), "float16")) -> None:
+    T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
+    bx = T.env_thread("blockIdx.x")
+    tx = T.env_thread("threadIdx.x")
+    T.launch_thread(bx, 1)
+    T.launch_thread(tx, 32)
+    with T.sblock():
+        A_shared = T.sblock_alloc_buffer([32, 128], "float16", scope="shared")
+        T.reads(A[0:32, 0:128])
+        T.writes(B[0:32, 0:128])
+
+        for i in range(16):
+            T.evaluate(
+                T.ptx.cp_async.legacy(
+                    A_shared.data, tx * 128 + 8 * i, A.data, tx * 128 + 8 * i, 16, dtype="float16"
+                )
+            )
+
+        # TODO(masahi): Remove dtype requirement from TVMScript parser
+        T.evaluate(T.ptx.cp_async.commit_group(dtype=""))
+        T.evaluate(T.ptx.cp_async.wait_group(0, dtype=""))
+
+        for i in range(128):
+            B[tx, i] = A_shared[tx, i]
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
+def test_ptx_cp_async():
+    f = ptx_cp_async
+
+    mod = tvm.compile(f, target="cuda")
+    A_np = np.random.rand(32, 128).astype("float16")
+    B_np = np.zeros((32, 128)).astype("float16")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A_nd = tvm.runtime.tensor(A_np, device=dev)
+        B_nd = tvm.runtime.tensor(B_np, device=dev)
+        mod(A_nd, B_nd)
+        tvm.testing.assert_allclose(B_nd.numpy(), A_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# Note: tests for the indexed barrier API (`create_barriers`,
+# `ptx_init_barrier_thread_count`, `ptx_arrive_barrier`, `ptx_wait_barrier`,
+# `ptx_cp_async_barrier`, `ptx_arrive_barrier_expect_tx`) were removed —
+# fork uses `ptx_mbarrier_*` instead and those intrinsics have no
+# users elsewhere in this codebase.
+
+
+if __name__ == "__main__":
+    test_ptx_cp_async()

tests/python/tirx-base/test_tir_ptx_griddepcontrol.py

@@ -0,0 +1,61 @@
+# 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 numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@T.prim_func(s_tir=True)
+def ptx_griddepcontrol(A: T.Buffer((32,), "float32"), B: T.Buffer((32,), "float32")) -> None:
+    T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
+    bx = T.env_thread("blockIdx.x")
+    tx = T.env_thread("threadIdx.x")
+    T.launch_thread(bx, 1)
+    T.launch_thread(tx, 32)
+    with T.sblock():
+        T.reads(A[0:32])
+        T.writes(B[0:32])
+        T.evaluate(T.ptx.griddepcontrol.wait(dtype=""))
+        B[tx] = A[tx]
+        T.evaluate(T.ptx.griddepcontrol.launch_dependents(dtype=""))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_ptx_griddepcontrol():
+    f = ptx_griddepcontrol
+    mod = tvm.compile(f, target="cuda")
+    A_np = np.random.default_rng(0).standard_normal(32).astype("float32")
+    B_np = np.zeros((32,), dtype="float32")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A_nd = tvm.runtime.tensor(A_np, device=dev)
+        B_nd = tvm.runtime.tensor(B_np, device=dev)
+        mod(A_nd, B_nd)
+        tvm.testing.assert_allclose(B_nd.numpy(), A_np, rtol=0, atol=0)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    test_ptx_griddepcontrol()

tests/python/tirx-base/test_tir_ptx_ldmatrix.py

@@ -0,0 +1,105 @@
+# 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 numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@T.prim_func(s_tir=True)
+def ptx_ldmatrix(
+    A: T.Buffer((16, 16), "float16"), B: T.Buffer((16, 16), "float16"), num: T.int32, trans: T.uint8
+) -> None:
+    T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
+    bx = T.env_thread("blockIdx.x")
+    tx = T.env_thread("threadIdx.x")
+    T.launch_thread(bx, 1)
+    T.launch_thread(tx, 32)
+    with T.sblock():
+        A_shared = T.sblock_alloc_buffer([16, 16], "float16", scope="shared")
+        A_local = T.sblock_alloc_buffer([8], "float16", scope="local")
+
+        for i in range(8):
+            A_shared[i * 2 + tx // 16, tx % 16] = A[i * 2 + tx // 16, tx % 16]
+
+        T.evaluate(
+            T.ptx.ldmatrix_legacy(
+                trans,
+                num,
+                ".b16",
+                A_local.data,
+                0,
+                A_shared.data,
+                16 * (tx % 16) + 8 * (tx // 16),
+                dtype="float16",
+            )
+        )
+
+        for k in range(2):
+            for j in range(2):
+                for i in range(2):
+                    B[8 * j + tx // 4, 8 * k + (tx % 4) * 2 + i] = A_local[4 * k + 2 * j + i]
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(7, 5), reason="need cuda compute >= 7.5")
+def test_ptx_ldmatrix():
+    f = ptx_ldmatrix
+    _, _, param_num, param_trans = f.params
+
+    for num in [1, 2, 4]:
+        for trans in [False, True]:
+            mod = tvm.compile(f.specialize({param_num: num, param_trans: trans}), target="cuda")
+            A_np = np.random.rand(16, 16).astype("float16")
+            A_mask_np = np.zeros_like(A_np)
+            if num == 1:
+                if trans:
+                    A_mask_np[:8, :8] = A_np[:8, :8].T
+                else:
+                    A_mask_np[:8, :8] = A_np[:8, :8]
+            elif num == 2:
+                if trans:
+                    A_mask_np[:8, :8] = A_np[:8, :8].T
+                    A_mask_np[8:16, :8] = A_np[8:16, :8].T
+                else:
+                    A_mask_np[:16, :8] = A_np[:16, :8]
+            else:  # num == 4
+                if trans:
+                    A_mask_np[:8, :8] = A_np[:8, :8].T
+                    A_mask_np[8:16, :8] = A_np[8:16, :8].T
+                    A_mask_np[:8, 8:16] = A_np[:8, 8:16].T
+                    A_mask_np[8:16, 8:16] = A_np[8:16, 8:16].T
+                else:
+                    A_mask_np[:16, :16] = A_np[:16, :16]
+            B_np = np.zeros((16, 16)).astype("float16")
+
+            def run_and_check():
+                dev = tvm.cuda(0)
+                A_nd = tvm.runtime.tensor(A_np, device=dev)
+                B_nd = tvm.runtime.tensor(B_np, device=dev)
+                mod(A_nd, B_nd)
+                tvm.testing.assert_allclose(B_nd.numpy(), A_mask_np)
+
+            tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    test_ptx_ldmatrix()

tests/python/tirx-base/test_tir_ptx_mma_sp.py

@@ -0,0 +1,345 @@
+# 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 numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+def gen_2in4_mask(m: int, n: int):
+    assert n % 4 == 0
+    return np.array(
+        [[np.sort(np.random.choice(4, 2, replace=False)) for _ in range(n // 4)] for _ in range(m)]
+    ).astype("uint8")
+
+
+def get_dense_mat_by_mask(val, mask):
+    m, n_chunks, _ = mask.shape
+    val = val.reshape(m, n_chunks, 2)
+    ret = np.zeros((m, n_chunks, 4)).astype(val.dtype)
+    for i in range(m):
+        for j in range(n_chunks):
+            for k in range(2):
+                ret[i, j, mask[i, j, k]] = val[i, j, k]
+    return ret.reshape(m, n_chunks * 4)
+
+
+@T.prim_func(s_tir=True)
+def mma_sp_m16n8k16_f16f16f16(a: T.handle, b: T.handle, c: T.handle, _metadata: T.handle):
+    T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
+    A = T.match_buffer(a, [16, 8], dtype="float16")
+    B = T.match_buffer(b, [16, 8], dtype="float16")
+    C = T.match_buffer(c, [16, 8], dtype="float16")
+    metadata = T.match_buffer(_metadata, [8], dtype="uint32")
+    brow = T.env_thread("blockIdx.y")
+    bcol = T.env_thread("blockIdx.x")
+    tx = T.env_thread("threadIdx.x")
+    T.launch_thread(brow, 1)
+    T.launch_thread(bcol, 1)
+    T.launch_thread(tx, 32)
+    multi_a = T.decl_buffer([4], "float16", scope="local")
+    multi_b = T.decl_buffer([4], "float16", scope="local")
+    accum = T.decl_buffer([4], "float16", scope="local")
+    meta_local = T.decl_buffer([1], "uint32", scope="local")
+    for i in range(4):
+        accum[i] = T.float16(0)
+
+    for i in range(4):
+        multi_a[i] = A[tx // 4 + i // 2 * 8, tx % 4 * 2 + i % 2]
+
+    for i in range(4):
+        multi_b[i] = B[tx % 4 * 2 + i % 2 + i // 2 * 8, tx // 4]
+
+    meta_local[0] = metadata[tx // 4]
+
+    T.evaluate(
+        T.ptx.mma.sp(
+            "m16n8k16",
+            "row",
+            "col",
+            "fp16",
+            "fp16",
+            "fp16",
+            multi_a.data,
+            0,
+            multi_b.data,
+            0,
+            accum.data,
+            0,
+            meta_local.data,
+            0,
+            0,
+            False,
+            dtype="float16",
+        )
+    )
+
+    for i in range(4):
+        C[i // 2 * 8 + tx // 4, tx % 4 * 2 + i % 2] = accum[i]
+
+
+@T.prim_func(s_tir=True)
+def mma_sp_m16n8k16_f16f16f32(a: T.handle, b: T.handle, c: T.handle, _metadata: T.handle):
+    T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
+    A = T.match_buffer(a, [16, 8], dtype="float16")
+    B = T.match_buffer(b, [16, 8], dtype="float16")
+    C = T.match_buffer(c, [16, 8], dtype="float32")
+    metadata = T.match_buffer(_metadata, [8], dtype="uint32")
+    brow = T.env_thread("blockIdx.y")
+    bcol = T.env_thread("blockIdx.x")
+    tx = T.env_thread("threadIdx.x")
+    T.launch_thread(brow, 1)
+    T.launch_thread(bcol, 1)
+    T.launch_thread(tx, 32)
+    multi_a = T.decl_buffer([4], "float16", scope="local")
+    multi_b = T.decl_buffer([4], "float16", scope="local")
+    accum = T.decl_buffer([4], "float32", scope="local")
+    meta_local = T.decl_buffer([1], "uint32", scope="local")
+    for i in range(4):
+        accum[i] = T.float16(0)
+
+    for i in range(4):
+        multi_a[i] = A[tx // 4 + i // 2 * 8, tx % 4 * 2 + i % 2]
+
+    for i in range(4):
+        multi_b[i] = B[tx % 4 * 2 + i % 2 + i // 2 * 8, tx // 4]
+
+    meta_local[0] = metadata[tx // 4]
+
+    T.evaluate(
+        T.ptx.mma.sp(
+            "m16n8k16",
+            "row",
+            "col",
+            "fp16",
+            "fp16",
+            "fp32",
+            multi_a.data,
+            0,
+            multi_b.data,
+            0,
+            accum.data,
+            0,
+            meta_local.data,
+            0,
+            0,
+            False,
+            dtype="float32",
+        )
+    )
+
+    for i in range(4):
+        C[i // 2 * 8 + tx // 4, tx % 4 * 2 + i % 2] = accum[i]
+
+
+@T.prim_func(s_tir=True)
+def mma_sp_m16n8k32_f16f16f16(a: T.handle, b: T.handle, c: T.handle, _metadata: T.handle):
+    T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
+    A = T.match_buffer(a, [16, 16], dtype="float16")
+    B = T.match_buffer(b, [32, 8], dtype="float16")
+    C = T.match_buffer(c, [16, 8], dtype="float16")
+    metadata = T.match_buffer(_metadata, [16], dtype="uint32")
+    brow = T.env_thread("blockIdx.y")
+    bcol = T.env_thread("blockIdx.x")
+    tx = T.env_thread("threadIdx.x")
+    T.launch_thread(brow, 1)
+    T.launch_thread(bcol, 1)
+    T.launch_thread(tx, 32)
+    multi_a = T.decl_buffer([8], "float16", scope="local")
+    multi_b = T.decl_buffer([8], "float16", scope="local")
+    accum = T.decl_buffer([4], "float16", scope="local")
+    meta_local = T.decl_buffer([1], "uint32", scope="local")
+    for i in range(4):
+        accum[i] = T.float16(0)
+
+    for i in range(8):
+        multi_a[i] = A[(i % 4) // 2 * 8 + tx // 4, i // 4 * 8 + tx % 4 * 2 + i % 2]
+
+    for i in range(8):
+        multi_b[i] = B[i // 2 * 8 + tx % 4 * 2 + i % 2, tx // 4]
+
+    meta_local[0] = metadata[tx // 4 * 2 + tx % 2]
+
+    T.evaluate(
+        T.ptx.mma.sp(
+            "m16n8k32",
+            "row",
+            "col",
+            "fp16",
+            "fp16",
+            "fp16",
+            multi_a.data,
+            0,
+            multi_b.data,
+            0,
+            accum.data,
+            0,
+            meta_local.data,
+            0,
+            0,
+            False,
+            dtype="float16",
+        )
+    )
+
+    for i in range(4):
+        C[i // 2 * 8 + tx // 4, tx % 4 * 2 + i % 2] = accum[i]
+
+
+@T.prim_func(s_tir=True)
+def mma_sp_m16n8k32_f16f16f32(a: T.handle, b: T.handle, c: T.handle, _metadata: T.handle):
+    T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
+    A = T.match_buffer(a, [16, 16], dtype="float16")
+    B = T.match_buffer(b, [32, 8], dtype="float16")
+    C = T.match_buffer(c, [16, 8], dtype="float32")
+    metadata = T.match_buffer(_metadata, [16], dtype="uint32")
+    brow = T.env_thread("blockIdx.y")
+    bcol = T.env_thread("blockIdx.x")
+    tx = T.env_thread("threadIdx.x")
+    T.launch_thread(brow, 1)
+    T.launch_thread(bcol, 1)
+    T.launch_thread(tx, 32)
+    multi_a = T.decl_buffer([8], "float16", scope="local")
+    multi_b = T.decl_buffer([8], "float16", scope="local")
+    accum = T.decl_buffer([4], "float32", scope="local")
+    meta_local = T.decl_buffer([1], "uint32", scope="local")
+    for i in range(4):
+        accum[i] = T.float16(0)
+
+    for i in range(8):
+        multi_a[i] = A[(i % 4) // 2 * 8 + tx // 4, i // 4 * 8 + tx % 4 * 2 + i % 2]
+
+    for i in range(8):
+        multi_b[i] = B[i // 2 * 8 + tx % 4 * 2 + i % 2, tx // 4]
+
+    meta_local[0] = metadata[tx // 4 * 2 + tx % 2]
+
+    T.evaluate(
+        T.ptx.mma.sp(
+            "m16n8k32",
+            "row",
+            "col",
+            "fp16",
+            "fp16",
+            "fp32",
+            multi_a.data,
+            0,
+            multi_b.data,
+            0,
+            accum.data,
+            0,
+            meta_local.data,
+            0,
+            0,
+            False,
+            dtype="float32",
+        )
+    )
+
+    for i in range(4):
+        C[i // 2 * 8 + tx // 4, tx % 4 * 2 + i % 2] = accum[i]
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
+def test_mma_sp_m16n8k16_f16():
+    def get_meta_m16n8k16_half(mask):
+        assert mask.shape == (16, 4, 2)
+        mask = mask.reshape(16, 8)
+        ret = np.zeros((8,)).astype("uint32")
+
+        for i in range(8):
+            base = 1
+            for blk in range(2):
+                for j in range(8):
+                    ret[i] |= int(mask[blk * 8 + i, j]) * base
+                    base = base << 2
+        return ret
+
+    for out_dtype in ["float16", "float32"]:
+        func = mma_sp_m16n8k16_f16f16f16 if out_dtype == "float16" else mma_sp_m16n8k16_f16f16f32
+        sch = tvm.s_tir.Schedule(func)
+        cuda_mod = tvm.compile(sch.mod, target="cuda")
+
+        A_np = np.random.uniform(-1, 1, [16, 8]).astype("float16")
+        B_np = np.random.uniform(-1, 1, [16, 8]).astype("float16")
+        mask = gen_2in4_mask(16, 16)
+        A_dense_np = get_dense_mat_by_mask(A_np, mask)
+        C_np = np.matmul(A_dense_np, B_np).astype(out_dtype)
+        meta = get_meta_m16n8k16_half(mask)
+
+        def run_and_check():
+            ctx = tvm.cuda()
+            A_tvm = tvm.runtime.tensor(A_np, ctx)
+            B_tvm = tvm.runtime.tensor(B_np, ctx)
+            C_tvm = tvm.runtime.tensor(np.zeros_like(C_np), ctx)
+            meta_tvm = tvm.runtime.tensor(meta, ctx)
+            cuda_mod(A_tvm, B_tvm, C_tvm, meta_tvm)
+            tvm.testing.assert_allclose(C_tvm.numpy(), C_np, atol=1e-3, rtol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
+def test_mma_sp_m16n8k32_f16():
+    def get_meta_m16n8k32_half(mask):
+        assert mask.shape == (16, 8, 2)
+        mask = mask.reshape(16, 2, 8)
+        ret = np.zeros((8, 2)).astype("uint32")
+
+        for i in range(8):
+            for k in range(2):
+                base = 1
+                for blk in range(2):
+                    for j in range(8):
+                        ret[i, k] |= int(mask[blk * 8 + i, k, j]) * base
+                        base = base << 2
+
+        return ret.reshape(16)
+
+    for out_dtype in ["float16", "float32"]:
+        func = mma_sp_m16n8k32_f16f16f16 if out_dtype == "float16" else mma_sp_m16n8k32_f16f16f32
+        sch = tvm.s_tir.Schedule(func)
+        cuda_mod = tvm.compile(sch.mod, target="cuda")
+
+        A_np = np.random.uniform(-1, 1, [16, 16]).astype("float16")
+        B_np = np.random.uniform(-1, 1, [32, 8]).astype("float16")
+        mask = gen_2in4_mask(16, 32)
+        A_dense_np = get_dense_mat_by_mask(A_np, mask)
+        C_np = np.matmul(A_dense_np, B_np).astype(out_dtype)
+        meta = get_meta_m16n8k32_half(mask)
+
+        def run_and_check():
+            ctx = tvm.cuda()
+            A_tvm = tvm.runtime.tensor(A_np, ctx)
+            B_tvm = tvm.runtime.tensor(B_np, ctx)
+            C_tvm = tvm.runtime.tensor(np.zeros_like(C_np), ctx)
+            meta_tvm = tvm.runtime.tensor(meta, ctx)
+            cuda_mod(A_tvm, B_tvm, C_tvm, meta_tvm)
+            tvm.testing.assert_allclose(C_tvm.numpy(), C_np, atol=1e-3, rtol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    test_mma_sp_m16n8k16_f16()
+    test_mma_sp_m16n8k32_f16()

tests/python/tirx-base/test_tir_ptx_scalar_f32_math.py

@@ -0,0 +1,74 @@
+# 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 numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@T.prim_func(s_tir=True)
+def ptx_scalar_f32_math(
+    A: T.Buffer((32,), "float32"),
+    B: T.Buffer((32,), "float32"),
+    C_add: T.Buffer((32,), "float32"),
+    C_mul: T.Buffer((32,), "float32"),
+    C_max: T.Buffer((32,), "float32"),
+) -> None:
+    T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
+    bx = T.env_thread("blockIdx.x")
+    tx = T.env_thread("threadIdx.x")
+    T.launch_thread(bx, 1)
+    T.launch_thread(tx, 32)
+    with T.sblock():
+        T.reads(A[0:32], B[0:32])
+        T.writes(C_add[0:32], C_mul[0:32], C_max[0:32])
+        T.evaluate(T.ptx.add_f32(T.address_of(C_add[tx]), A[tx], B[tx]))
+        T.evaluate(T.ptx.mul_f32(T.address_of(C_mul[tx]), A[tx], B[tx]))
+        C_max[tx] = T.ptx.max_f32(A[tx], B[tx])
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(7), reason="need cuda compute >= 7.0")
+def test_ptx_scalar_f32_math():
+    f = ptx_scalar_f32_math
+    mod = tvm.compile(f, target="cuda")
+    rng = np.random.default_rng(0)
+    A_np = rng.standard_normal(32).astype("float32")
+    B_np = rng.standard_normal(32).astype("float32")
+    Z = np.zeros((32,), dtype="float32")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A_nd = tvm.runtime.tensor(A_np, device=dev)
+        B_nd = tvm.runtime.tensor(B_np, device=dev)
+        Cadd = tvm.runtime.tensor(Z.copy(), device=dev)
+        Cmul = tvm.runtime.tensor(Z.copy(), device=dev)
+        Cmax = tvm.runtime.tensor(Z.copy(), device=dev)
+        mod(A_nd, B_nd, Cadd, Cmul, Cmax)
+        tvm.testing.assert_allclose(Cadd.numpy(), A_np + B_np, rtol=0, atol=0)
+        tvm.testing.assert_allclose(Cmul.numpy(), A_np * B_np, rtol=0, atol=0)
+        tvm.testing.assert_allclose(Cmax.numpy(), np.maximum(A_np, B_np), rtol=0, atol=0)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    test_ptx_scalar_f32_math()

tests/python/tirx-base/test_tir_scalable_datatype.py

@@ -0,0 +1,65 @@
+# 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 pytest
+
+import tvm
+from tvm import tirx
+from tvm.script import tirx as T
+from tvm.target.codegen import llvm_version_major
+
+"""
+Tests for scalable data types.
+"""
+
+
+def test_create_scalable_data_type_python_api():
+    dtype = tvm.DataType("float32xvscalex4")
+    assert str(dtype) == "float32xvscalex4"
+
+
+# LLVM 20 renamed llvm.experimental.stepvector to llvm.stepvector and dropped
+# the old name from the intrinsic table:
+# https://releases.llvm.org/20.1.0/docs/ReleaseNotes.html
+_STEPVECTOR_NAME = (
+    "llvm.stepvector" if llvm_version_major() >= 20 else "llvm.experimental.stepvector"
+)
+
+
+@pytest.mark.skipif(llvm_version_major() < 13, reason="Stepvector intrinsic was added in LLVM 13.")
+def test_create_scalable_tir_intrin():
+    intrin = tirx.call_llvm_intrin("int32xvscalex4", _STEPVECTOR_NAME)
+    assert intrin.ty.dtype == "int32xvscalex4"
+    assert str(intrin) == f'T.call_llvm_intrin("int32xvscalex4", "{_STEPVECTOR_NAME}")'
+
+
+@pytest.mark.skipif(llvm_version_major() < 13, reason="Stepvector intrinsic was added in LLVM 13.")
+def test_tvm_script_create_scalable_tir_intrin():
+    @T.prim_func(s_tir=True)
+    def my_func():
+        T.call_llvm_intrin("int32xvscalex4", _STEPVECTOR_NAME)
+
+    assert f'T.call_llvm_intrin("int32xvscalex4", "{_STEPVECTOR_NAME}")' in my_func.script()
+
+
+def test_invalid_data_type():
+    with pytest.raises(ValueError):
+        tvm.DataType("float32x4xvscale")
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-base/test_tir_specialize.py

@@ -0,0 +1,365 @@
+# 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.
+# pylint: disable=missing-function-docstring, missing-module-docstring
+# ruff: noqa: F401, F841
+
+import pytest
+
+import tvm
+from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
+from tvm.script import tirx as T
+
+
+@T.prim_func(s_tir=True)
+def matmul(a: T.handle, b: T.handle, c: T.handle, n: T.int32) -> None:
+    m = T.int32()
+    A = T.match_buffer(a, [m, n])
+    B = T.match_buffer(b, [m, n])
+    C = T.match_buffer(c, [m, m])
+
+    for i, j, k in T.grid(m, m, n):
+        with T.sblock("update"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def matmul_128(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("update"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def matmul_m_128(a: T.handle, b: T.handle, c: T.handle) -> None:
+    m = T.int32()
+    A = T.match_buffer(a, [m, 128])
+    B = T.match_buffer(b, [m, 128])
+    C = T.match_buffer(c, [m, m])
+
+    for i, j, k in T.grid(m, m, 128):
+        with T.sblock("update"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+# x is considered undefined because it appears as part of x*8,
+# but not on its own
+@T.prim_func(check_well_formed=False, s_tir=True)
+def matmul_m_8x(a: T.handle, b: T.handle, c: T.handle) -> None:
+    x = T.int32()
+    m = T.int32()
+    A = T.match_buffer(a, [m, x * 8])
+    B = T.match_buffer(b, [m, x * 8])
+    C = T.match_buffer(c, [m, m])
+
+    for i, j, k in T.grid(m, m, x * 8):
+        with T.sblock("update"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func(s_tir=True)
+def element_wise(a: T.handle, c: T.handle) -> None:
+    m = T.int32()
+    n = T.int32()
+    A = T.match_buffer(a, (m, n), "float32")
+    C = T.match_buffer(c, (m, n), "float32")
+
+    B = T.sblock_alloc_buffer((m, n), "float32")
+
+    for i, j in T.grid(m, n):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+
+    for i, j in T.grid(m, n):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_128_64(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 64), "float32")
+    C = T.match_buffer(c, (128, 64), "float32")
+    B = T.sblock_alloc_buffer((128, 64), "float32")
+
+    for i, j in T.grid(128, 64):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+
+    for i, j in T.grid(128, 64):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def element_wise_128_n(a: T.handle, c: T.handle) -> None:
+    n = T.int32()
+    A = T.match_buffer(a, (128, n), "float32")
+    C = T.match_buffer(c, (128, n), "float32")
+    B = T.sblock_alloc_buffer((128, n), "float32")
+
+    for i, j in T.grid(128, n):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+
+    for i, j in T.grid(128, n):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def mem_copy(a: T.handle, b: T.handle, m: T.int32, n: T.int32, p: T.int32, q: T.int32) -> None:
+    A = T.match_buffer(a, (m, n), "float32", strides=[p, 1], elem_offset=q)
+    B = T.match_buffer(b, (m, n), "float32", strides=[p, 1], elem_offset=q)
+
+    for i, j in T.grid(m, n):
+        with T.sblock():
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def mem_copy_16_16_8_4(a: T.handle, b: T.handle) -> None:
+    A = T.match_buffer(a, (16, 16), "float32", strides=[8, 1], elem_offset=4)
+    B = T.match_buffer(b, (16, 16), "float32", strides=[8, 1], elem_offset=4)
+
+    for i, j in T.grid(16, 16):
+        with T.sblock():
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj]
+
+
+@T.prim_func(s_tir=True)
+def mem_copy_m_n_p_n(a: T.handle, b: T.handle, m: T.int32, n: T.int32, p: T.int32) -> None:
+    A = T.match_buffer(a, (m, n), "float32", strides=[p, 1], elem_offset=n)
+    B = T.match_buffer(b, (m, n), "float32", strides=[p, 1], elem_offset=n)
+
+    for i, j in T.grid(m, n):
+        with T.sblock():
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj]
+
+
+def test_specialize_nothing():
+    func = matmul.specialize({})
+    assert func.same_as(matmul)  # Pointer the same
+
+
+def test_specialize_matmul():
+    a, _, _, n = matmul.params
+    # fully specialized
+    func = matmul.specialize({a: tvm.tirx.decl_buffer((128, 128))})
+    assert_structural_equal_ignore_global_symbol(func, matmul_128)
+    # partially specialized
+    func = matmul.specialize({n: 128})
+    assert_structural_equal_ignore_global_symbol(func, matmul_m_128)
+    # symbolic specialized
+    func = matmul.specialize({n: tvm.tirx.Var("x", "int32") * 8})
+    assert_structural_equal_ignore_global_symbol(func, matmul_m_8x)
+
+
+def test_specialize_elemwise():
+    a, c = element_wise.params
+    C = element_wise.buffer_map[c]
+    # fully specialized
+    func = element_wise.specialize({a: tvm.tirx.decl_buffer((128, 64))})
+    assert_structural_equal_ignore_global_symbol(func, element_wise_128_64)
+    # partially specialized
+    func = element_wise.specialize({c: tvm.tirx.decl_buffer((128, C.shape[1]))})
+    assert_structural_equal_ignore_global_symbol(func, element_wise_128_n)
+
+
+def test_specialize_mem_copy():
+    a, _, m, n, p, q = mem_copy.params
+    # fully specialized
+    func = mem_copy.specialize({a: tvm.tirx.decl_buffer((16, 16), strides=[8, 1], elem_offset=4)})
+    assert_structural_equal_ignore_global_symbol(func, mem_copy_16_16_8_4)
+    func = mem_copy.specialize({n: 16, m: 16, p: 8, q: 4})
+    assert_structural_equal_ignore_global_symbol(func, mem_copy_16_16_8_4)
+    # partially specialized
+    func = mem_copy.specialize({q: n})
+    assert_structural_equal_ignore_global_symbol(func, mem_copy_m_n_p_n)
+
+
+def test_specialize_recursive_load():
+    # TODO(Siyuan): add recursive Load testcase, e.g. A[C[i]]
+    pass
+
+
+def test_specialize_with_const_folding():
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle):
+        n = T.int32()
+        A = T.match_buffer(a, [n // 8, 8], "int32")
+        B = T.match_buffer(b, [n], "int32")
+        for i in range(n - 1):
+            with T.sblock():
+                vi = T.axis.S(n - 1, i)
+                B[vi] = A[vi // 8, vi % 8] + (n + 1) * 42
+
+    @T.prim_func(s_tir=True)
+    def expected(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, [2, 8], "int32")
+        B = T.match_buffer(b, [16], "int32")
+        for i in range(15):
+            with T.sblock():
+                vi = T.axis.S(15, i)
+                B[vi] = A[vi // 8, vi % 8] + 714
+
+    b = before.params[1]
+    after = before.specialize({b: tvm.tirx.decl_buffer([16], dtype="int32")})
+    assert_structural_equal_ignore_global_symbol(expected, after)
+
+
+def test_specialize_decl_buffer():
+    """Buffers occurring in a DeclBuffer statement should be updated"""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A_data: T.handle("float32"), A_size: T.int32):
+        A_buf = T.decl_buffer(A_size, "float32", data=A_data)
+        for i in range(A_size):
+            A_buf[i] = A_buf[i] * 2.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A_data: T.handle("float32")):
+        A_buf = T.decl_buffer(16, "float32", data=A_data)
+        for i in range(16):
+            A_buf[i] = A_buf[i] * 2.0
+
+    param_map = {before.params[1]: T.int32(16)}
+    after = before.specialize(param_map)
+
+    tvm.ir.assert_structural_equal(expected, after)
+
+
+def test_specialize_buffer_var_to_var():
+    """A buffer var may be remapped by specialization
+
+    If a buffer variable is replaced by a specialization, then other
+    buffers using the same buffer var should also be updated.
+    """
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer([16, 16], "float32"), B: T.Buffer([16, 16], "float32")):
+        A_flat = T.decl_buffer([256], "float32", data=A.data)
+        B_flat = T.decl_buffer([256], "float32", data=B.data)
+        for i in range(256):
+            B_flat[i] = A_flat[i] * 2.0
+
+    # well-formed checker complains about multiple nested definitions of B_flat
+    # since it appears in the buffer map twice
+    @T.prim_func(private=True, check_well_formed=False, s_tir=True)
+    def expected(A: T.Buffer([16, 16], "float32"), B_handle: T.handle):
+        B = T.match_buffer(B_handle, [16, 16], "float32", data=A.data)
+        A_flat = T.decl_buffer([256], "float32", data=A.data)
+        B_flat = T.decl_buffer([256], "float32", data=A.data)
+        for i in range(256):
+            B_flat[i] = A_flat[i] * 2.0
+
+    A = before.buffer_map[before.params[0]]
+    B_handle = before.params[1]
+    param_map = {B_handle: A}
+    after = before.specialize(param_map)
+
+    tvm.ir.assert_structural_equal(expected, after)
+
+
+def test_specialize_buffer_var_to_expr():
+    """Handle specialization of buffer var
+
+    The `tirx::Buffer::data` field must be an explicit `tirx::Var`, and
+    cannot be replaced with a handle-typed `tirx::Expr`.  However,
+    these substitutions are useful
+    when lowering.  If these occur, a binding of the `tirx::Var` is
+    included in the specialized function.
+    """
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A_data: T.handle("float32"), B_data: T.handle("float32")):
+        A_buf = T.decl_buffer(32, "float32", data=A_data)
+        B_buf = T.decl_buffer(16, "float32", data=B_data)
+        for i in range(16):
+            B_buf[i] = A_buf[i] * 2.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A_data: T.handle("float32")):
+        A_buf = T.decl_buffer(32, "float32", data=A_data)
+        B_data: T.let[T.Ptr[T.float32]] = T.address_of(A_buf[16])
+        B_buf = T.decl_buffer(16, "float32", data=B_data)
+        for i in range(16):
+            B_buf[i] = A_buf[i] * 2.0
+
+    B_data = before.params[1]
+    # body is a SeqStmt; the first statement is DeclBuffer for A_buf
+    A_buf = before.body[0].buffer
+    param_map = {B_data: tvm.tirx.address_of(A_buf[16])}
+    after = before.specialize(param_map)
+
+    tvm.ir.assert_structural_equal(expected, after)
+
+
+def test_specialization_updates_ty():
+    """Update type in specialization
+
+    A PrimFunc may have a `relax.Type`.  If that PrimFunc is
+    specialized, the type should be updated.
+    """
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(n: T.int32) -> T.int32:
+        T.ret(n * 10)
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected() -> T.int32:
+        T.ret(50)
+
+    ty_before = tvm.relax.FuncType([tvm.ir.PrimType("int32")], tvm.ir.PrimType("int32"))
+    tvm.ir.assert_structural_equal(before.ty, ty_before)
+
+    ty_expected = tvm.relax.FuncType([], tvm.ir.PrimType("int32"))
+    tvm.ir.assert_structural_equal(expected.ty, ty_expected)
+
+    n = before.params[0]
+    param_map = {n: 5}
+    after = before.specialize(param_map)
+
+    tvm.ir.assert_structural_equal(after, expected)
+    tvm.ir.assert_structural_equal(after.ty, ty_expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-base/test_tir_stmt_functor_ir_transform.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.
+import tvm
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_ir_transform():
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(n: T.int32):
+            for i in T.serial(n):
+                for j in T.serial(10):
+                    # Inline call_extern to avoid Let binding (x must be the Call node itself)
+                    T.evaluate(
+                        T.call_extern(
+                            "int32", "TestB", T.call_extern("int32", "TestA", i * 3 + j * 1)
+                        )
+                    )
+                    T.evaluate(
+                        T.call_extern(
+                            "int32", "TestC", T.call_extern("int32", "TestA", i * 3 + j * 1)
+                        )
+                    )
+
+    body = Module["main"].body
+    builtin_call_extern = tvm.ir.Op.get("tirx.call_extern")
+
+    def preorder(op):
+        if op.op.same_as(builtin_call_extern) and op.args[0].value == "TestC":
+            return tvm.tirx.const(42, "int32")
+        return None
+
+    def postorder(op):
+        assert isinstance(op, tvm.ir.Call)
+        assert tvm.ir.is_prim_expr(op)
+        if op.op.same_as(builtin_call_extern) and op.args[0].value == "TestA":
+            return tvm.tirx.call_extern("int32", "TestB", op.args[1] + 1)
+        return op
+
+    body = tvm.tirx.stmt_functor.ir_transform(body, preorder, postorder, ["ir.Call"])
+    stmt_list = tvm.tirx.stmt_list(body.body.body)
+    assert stmt_list[0].value.args[1].args[0].value == "TestB"
+    assert stmt_list[1].value.value == 42
+
+
+if __name__ == "__main__":
+    test_ir_transform()

tests/python/tirx-base/test_tir_stmt_functor_substitute.py

@@ -0,0 +1,116 @@
+# 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.script import ir as I
+from tvm.script import tirx as T
+from tvm.tirx.stmt_functor import substitute
+
+
+def _apply_substitute(mod):
+    """Apply substitute transform to replace the first parameter with 16."""
+    func = mod["main"]
+    vmap = {func.params[0]: 16}
+    new_func = (
+        tvm.tirx.PrimFunc(params=[], body=substitute(func.body, vmap))
+        .with_attr("global_symbol", func.attrs["global_symbol"])
+        .with_attr("s_tir", True)
+    )
+    return tvm.IRModule.from_expr(new_func)
+
+
+def test_basic_substitute():
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(n: T.int32):
+            for i in range(n):
+                T.evaluate(i)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            for i in range(16):
+                T.evaluate(i)
+
+    After = _apply_substitute(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_substitute_allocate():
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(n: T.int32):
+            A = T.alloc_buffer((n,), "float32")
+            T.evaluate(A.data)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            A = T.alloc_buffer((16,), "float32")
+            T.evaluate(A.data)
+
+    After = _apply_substitute(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_substitute_buffer_load():
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(n: T.int32):
+            A = T.alloc_buffer((n,), "float32")
+            for i in range(n):
+                T.evaluate(A[i])
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            A = T.alloc_buffer((16,), "float32")
+            for i in range(16):
+                T.evaluate(A[i])
+
+    After = _apply_substitute(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_substitute_decl_buffer():
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(n: T.int32):
+            A = T.alloc_buffer((n,), "float32")
+            T.evaluate(A.data)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            A = T.alloc_buffer((16,), "float32")
+            T.evaluate(A.data)
+
+    After = _apply_substitute(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-base/test_tir_texture_scope.py

@@ -0,0 +1,64 @@
+# 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, F841
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.ir.module import IRModule
+from tvm.script import tirx as T
+
+
+def test_texture_scope():
+    @tvm.script.ir_module
+    class PlusOneMultTwo:
+        @T.prim_func(s_tir=True)
+        def main(a: T.handle, b: T.handle) -> None:
+            T.func_attr({"tirx.noalias": True})
+            A = T.match_buffer(a, (128, 128, 4), dtype="float32", scope="global.texture")
+            B = T.sblock_alloc_buffer((128, 128, 4), dtype="float32", scope="global.texture")
+            C = T.match_buffer(b, (128, 128, 4), dtype="float32", scope="global.texture")
+            for block_idx in T.thread_binding(0, 128, thread="blockIdx.x"):
+                for thread_idx in T.thread_binding(0, 128, thread="threadIdx.x"):
+                    for k in T.serial(4):
+                        with T.sblock("B"):
+                            vb, vt, vk = T.axis.remap("SSS", [block_idx, thread_idx, k])
+                            B[vb, vt, vk] = A[vb, vt, vk] + T.float32(1)
+            for block_idx in T.thread_binding(0, 128, thread="blockIdx.x"):
+                for thread_idx in T.thread_binding(0, 128, thread="threadIdx.x"):
+                    for k in T.serial(4):
+                        with T.sblock("C"):
+                            vb, vt, vk = T.axis.remap("SSS", [block_idx, thread_idx, k])
+                            C[vb, vt, vk] = B[vb, vt, vk] * T.float32(2)
+
+    sch = tvm.s_tir.Schedule(PlusOneMultTwo, debug_mask="all")
+
+    def schedule_block(block):
+        _, _, inner = sch.get_loops(block)
+        sch.vectorize(inner)
+
+    schedule_block(sch.get_sblock("B"))
+    schedule_block(sch.get_sblock("C"))
+
+    target = tvm.target.Target("opencl")
+    mod = tvm.compile(sch.mod["main"], target=target)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-base/test_tir_unsafe_hide_buffer_access.py

@@ -0,0 +1,107 @@
+# 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.
+# pylint: disable=missing-function-docstring,missing-module-docstring
+# ruff: noqa: F401
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.s_tir.schedule.testing import (
+    assert_structural_equal_ignore_global_symbol,
+    verify_trace_roundtrip,
+)
+from tvm.script import tirx as T
+
+
+@T.prim_func(s_tir=True)
+def indirect_mem_access(a: T.handle, idx_a: T.handle, b: T.handle, idx_b: T.handle) -> None:
+    A = T.match_buffer(a, [128], dtype="float32")
+    IA = T.match_buffer(idx_a, [10], dtype="int32")
+    B = T.match_buffer(b, [128], dtype="float32")
+    IB = T.match_buffer(idx_b, [10], dtype="int32")
+
+    for i in range(10):
+        with T.sblock("B"):
+            vi = T.axis.spatial(10, i)
+            T.reads(A[IA[vi]], IA[vi])
+            T.writes(B[IB[vi]], IB[vi])
+            B[IB[vi]] = A[IA[vi]]
+
+
+@T.prim_func(s_tir=True)
+def indirect_mem_access_hide_ia(a: T.handle, idx_a: T.handle, b: T.handle, idx_b: T.handle) -> None:
+    A = T.match_buffer(a, [128], dtype="float32")
+    IA = T.match_buffer(idx_a, [10], dtype="int32")
+    B = T.match_buffer(b, [128], dtype="float32")
+    IB = T.match_buffer(idx_b, [10], dtype="int32")
+
+    for i in range(10):
+        with T.sblock("B"):
+            vi = T.axis.spatial(10, i)
+            T.reads(A[IA[vi]])
+            T.writes(B[IB[vi]], IB[vi])
+            B[IB[vi]] = A[IA[vi]]
+
+
+@T.prim_func(s_tir=True)
+def indirect_mem_access_hide_ib(a: T.handle, idx_a: T.handle, b: T.handle, idx_b: T.handle) -> None:
+    A = T.match_buffer(a, [128], dtype="float32")
+    IA = T.match_buffer(idx_a, [10], dtype="int32")
+    B = T.match_buffer(b, [128], dtype="float32")
+    IB = T.match_buffer(idx_b, [10], dtype="int32")
+
+    for i in range(10):
+        with T.sblock("B"):
+            vi = T.axis.spatial(10, i)
+            T.reads(A[IA[vi]], IA[vi])
+            T.writes(B[IB[vi]])
+            B[IB[vi]] = A[IA[vi]]
+
+
+def test_hide_buffer_access_read():
+    sch = tvm.s_tir.Schedule(indirect_mem_access, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    sch.unsafe_hide_buffer_access(block_b, "read", [1])
+    assert_structural_equal_ignore_global_symbol(indirect_mem_access_hide_ia, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=indirect_mem_access)
+
+
+def test_hide_buffer_access_write():
+    sch = tvm.s_tir.Schedule(indirect_mem_access, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    sch.unsafe_hide_buffer_access(block_b, "write", [1])
+    assert_structural_equal_ignore_global_symbol(indirect_mem_access_hide_ib, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=indirect_mem_access)
+
+
+def test_hide_buffer_access_fail_buffer_type():
+    sch = tvm.s_tir.Schedule(indirect_mem_access, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    with pytest.raises(RuntimeError):
+        sch.unsafe_hide_buffer_access(block_b, "opaque", [0])
+
+
+def test_hide_buffer_access_fail_buffer_index():
+    sch = tvm.s_tir.Schedule(indirect_mem_access, debug_mask="all")
+    block_b = sch.get_sblock("B")
+    with pytest.raises(RuntimeError):
+        sch.unsafe_hide_buffer_access(block_b, "read", [2])
+
+
+if __name__ == "__main__":
+    tvm.testing.main()