chore: import upstream snapshot with attribution

tests/python/tirx-transform/test_tir_functor.py

@@ -0,0 +1,437 @@
+# 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 tvm
+import tvm.testing
+from tvm import tirx
+from tvm.tirx import (
+    EQ,
+    LT,
+    Add,
+    Cast,
+    Evaluate,
+    FloatImm,
+    For,
+    IfThenElse,
+    IntImm,
+    Max,
+    Min,
+    Mul,
+    PyStmtExprMutator,
+    PyStmtExprVisitor,
+    StringImm,
+    Sub,
+    Var,
+)
+
+
+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)
+
+
+@tirx.functor.visitor
+class ASTPrinter(PyStmtExprVisitor):
+    """Print tirx AST in structured format. The shape of Node is ignored."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.log = ASTLog()
+
+    def visit_var_(self, op: Var) -> None:
+        self.log.add("Stmt: Var")
+        super().visit_var_(op)
+
+    def visit_add_(self, op: Add) -> None:
+        self.log.add("Stmt: Add")
+        super().visit_add_(op)
+
+
+@tirx.functor.visitor
+class SimpleExprCounter(PyStmtExprVisitor):
+    """Count expressions without recursion"""
+
+    def __init__(self):
+        super().__init__()
+        self.var_count = 0
+        self.add_count = 0
+        self.mul_count = 0
+
+    def visit_var_(self, op: Var):
+        self.var_count += 1
+        # Don't recursively visit children to avoid infinite recursion
+
+    def visit_add_(self, op: Add):
+        self.add_count += 1
+        # Visit children manually
+        super().visit_add_(op)
+
+    def visit_mul_(self, op: Mul):
+        self.mul_count += 1
+        # Visit children manually
+        super().visit_mul_(op)
+
+
+@tirx.functor.mutator
+class VariableReplacer(PyStmtExprMutator):
+    """Replace variables with constants"""
+
+    def __init__(self, replacements):
+        super().__init__()
+        self.replacements = replacements
+
+    def visit_var_(self, op: Var):
+        if op.name in self.replacements:
+            return IntImm("int32", self.replacements[op.name])
+        return op
+
+
+@tirx.functor.mutator
+class AddToSubMutator(PyStmtExprMutator):
+    """Convert Add operations to Sub operations"""
+
+    def visit_add_(self, op: Add):
+        # First mutate the operands
+        a = self.visit_expr(op.a)
+        b = self.visit_expr(op.b)
+        # Convert Add to Sub
+        return Sub(a, b)
+
+
+@tirx.functor.visitor
+class SimpleStmtCounter(PyStmtExprVisitor):
+    """Count statements without recursion"""
+
+    def __init__(self):
+        super().__init__()
+        self.for_count = 0
+        self.if_count = 0
+        self.evaluate_count = 0
+
+    def visit_for_(self, op: For):
+        self.for_count += 1
+        super().visit_for_(op)
+
+    def visit_if_then_else_(self, op: IfThenElse):
+        self.if_count += 1
+        super().visit_if_then_else_(op)
+
+    def visit_evaluate_(self, op: Evaluate):
+        self.evaluate_count += 1
+        super().visit_evaluate_(op)
+
+
+@tirx.functor.mutator
+class ForLoopUnroller(PyStmtExprMutator):
+    """Simple loop unroller for demonstration"""
+
+    def __init__(self, unroll_factor=2):
+        super().__init__()
+        self.unroll_factor = unroll_factor
+
+    def visit_for_(self, op: For):
+        # For demonstration, just return the original for now
+        # In a real implementation, we would unroll small loops
+        return super().visit_for_(op)
+
+
+@tirx.functor.visitor
+class SimpleStmtExprVisitor(PyStmtExprVisitor):
+    """Visitor that handles both statements and expressions"""
+
+    def __init__(self):
+        super().__init__()
+        self.expr_count = 0
+        self.stmt_count = 0
+        self.var_names = set()
+
+    def visit_var_(self, op: Var):
+        self.var_names.add(op.name)
+        self.expr_count += 1
+
+    def visit_evaluate_(self, op: Evaluate):
+        self.stmt_count += 1
+        # Visit the expression
+        self.visit_expr(op.value)
+
+
+@tirx.functor.mutator
+class ComplexMutator(PyStmtExprMutator):
+    """Mutator that handles both statements and expressions"""
+
+    def __init__(self):
+        super().__init__()
+        self.modifications = 0
+
+    def visit_add_(self, op: Add):
+        self.modifications += 1
+        # Convert a + b to a * 2 + b for demonstration
+        a = self.visit_expr(op.a)
+        b = self.visit_expr(op.b)
+        return Add(Mul(a, IntImm("int32", 2)), b)
+
+
+def test_basic_visitor():
+    """Test the basic AST printer visitor"""
+    expr = Add(Var("x", dtype="int32"), Var("y", dtype="int32"))
+    printer = ASTPrinter()
+    printer.visit_expr(expr)
+    assert str(printer.log) == "\n".join(["Stmt: Add", "Stmt: Var", "Stmt: Var"])
+
+
+def test_simple_expr_counter():
+    """Test simple expression counting visitor"""
+    x = Var("x", dtype="int32")
+    y = Var("y", dtype="int32")
+
+    # Create simple expression: x + y
+    expr = Add(x, y)
+
+    counter = SimpleExprCounter()
+    counter.visit_expr(expr)
+
+    assert counter.var_count == 2  # x and y
+    assert counter.add_count == 1  # one add
+
+
+def test_variable_replacer():
+    """Test expression mutator that replaces variables"""
+    x = Var("x", dtype="int32")
+    y = Var("y", dtype="int32")
+    expr = Add(x, Mul(y, IntImm("int32", 3)))
+
+    replacer = VariableReplacer({"x": 10, "y": 5})
+    result = replacer.visit_expr(expr)
+
+    # Should be Add(IntImm(10), Mul(IntImm(5), IntImm(3)))
+    assert isinstance(result, Add)
+    assert isinstance(result.a, IntImm)
+    assert result.a.value == 10
+    assert isinstance(result.b, Mul)
+    assert isinstance(result.b.a, IntImm)
+    assert result.b.a.value == 5
+
+
+def test_add_to_sub_mutator():
+    """Test mutator that converts Add to Sub"""
+    x = Var("x", dtype="int32")
+    y = Var("y", dtype="int32")
+    expr = Add(x, y)
+
+    mutator = AddToSubMutator()
+    result = mutator.visit_expr(expr)
+
+    assert isinstance(result, Sub)
+    assert isinstance(result.a, Var)
+    assert isinstance(result.b, Var)
+    assert result.a.name == "x"
+    assert result.b.name == "y"
+
+
+def test_simple_stmt_counter():
+    """Test statement visitor that counts statements"""
+    i = Var("i", dtype="int32")
+
+    # Create a simple for loop
+    loop_body = Evaluate(IntImm("int32", 0))
+    for_stmt = For(i, IntImm("int32", 0), IntImm("int32", 10), tirx.ForKind.SERIAL, loop_body)
+
+    counter = SimpleStmtCounter()
+    counter.visit_stmt(for_stmt)
+
+    assert counter.for_count == 1  # One for loop
+    assert counter.evaluate_count == 1  # One evaluate in the body
+
+
+def test_if_then_else_visitor():
+    """Test visitor with if-then-else statements"""
+    x = Var("x", dtype="int32")
+    condition = EQ(x, IntImm("int32", 0))
+    then_stmt = Evaluate(IntImm("int32", 1))
+    else_stmt = Evaluate(IntImm("int32", 2))
+
+    if_stmt = IfThenElse(condition, then_stmt, else_stmt)
+
+    counter = SimpleStmtCounter()
+    counter.visit_stmt(if_stmt)
+
+    assert counter.if_count == 1
+    assert counter.for_count == 0
+
+
+def test_simple_stmt_expr_visitor():
+    """Test stmt_expr_visitor with mixed statements and expressions"""
+    x = Var("x", dtype="int32")
+    y = Var("y", dtype="int32")
+
+    # Create an evaluate statement with an expression
+    expr = Add(x, y)
+    stmt = Evaluate(expr)
+
+    visitor = SimpleStmtExprVisitor()
+    visitor.visit_stmt(stmt)
+
+    assert visitor.stmt_count == 1  # One Evaluate statement
+    assert visitor.expr_count == 2  # Two variables
+    assert "x" in visitor.var_names
+    assert "y" in visitor.var_names
+
+
+def test_complex_mutator():
+    """Test stmt_expr_mutator"""
+    x = Var("x", dtype="int32")
+    y = Var("y", dtype="int32")
+
+    # Expression with Add operations
+    expr = Add(x, y)
+    stmt = Evaluate(expr)
+
+    mutator = ComplexMutator()
+    result = mutator.visit_stmt(stmt)
+    print(type(mutator))
+
+    assert mutator.modifications == 1  # One Add operation modified
+    assert isinstance(result, Evaluate)
+
+    # Check that the expression was modified
+    modified_expr = result.value
+    assert isinstance(modified_expr, Add)
+    assert isinstance(modified_expr.a, Mul)  # First operand should be multiplied by 2
+
+
+def test_different_expr_types():
+    """Test visitor with various expression types"""
+    x = Var("x", dtype="int32")
+
+    # Test different expression types individually
+    exprs = [
+        IntImm("int32", 42),
+        FloatImm("float32", 3.14),
+        StringImm("hello"),
+        Cast("float32", x),
+        Min(x, IntImm("int32", 10)),
+        Max(x, IntImm("int32", 0)),
+        LT(x, IntImm("int32", 5)),
+    ]
+
+    # Just test that we can create and visit each type
+    counter = SimpleExprCounter()
+    for expr in exprs:
+        try:
+            counter.visit_expr(expr)
+        except Exception as e:
+            # Some expressions might not be supported, that's ok
+            pass
+
+
+def test_decorator_functionality():
+    """Test that decorators work correctly"""
+
+    # Test that decorated classes are properly wrapped
+    visitor = SimpleExprCounter()
+    assert hasattr(visitor, "_outer")  # Should have the wrapper functionality
+
+    mutator = VariableReplacer({})
+    assert hasattr(mutator, "_outer")
+
+
+def test_empty_expressions():
+    """Test handling of simple expressions"""
+    counter = SimpleExprCounter()
+
+    # Test with just a variable
+    x = Var("x", dtype="int32")
+    counter.visit_expr(x)
+
+    assert counter.var_count == 1
+
+    # Test with just a constant
+    counter = SimpleExprCounter()
+    const = IntImm("int32", 5)
+    counter.visit_expr(const)
+
+    # Constants don't increase var_count
+    assert counter.var_count == 0
+
+
+def test_stmt_mutator():
+    """Test basic statement mutator functionality"""
+    x = Var("x", dtype="int32")
+    stmt = Evaluate(Add(x, IntImm("int32", 1)))
+
+    unroller = ForLoopUnroller()
+    result = unroller.visit_stmt(stmt)
+
+    # Should return the same statement (no actual unrolling implemented)
+    assert isinstance(result, Evaluate)
+
+
+def test_nested_expressions():
+    """Test with nested expressions"""
+    x = Var("x", dtype="int32")
+    y = Var("y", dtype="int32")
+    z = Var("z", dtype="int32")
+
+    # Create nested expression: (x + y) * z
+    inner_add = Add(x, y)
+    expr = Mul(inner_add, z)
+
+    counter = SimpleExprCounter()
+    counter.visit_expr(expr)
+
+    assert counter.var_count == 3  # x, y, z
+    assert counter.add_count == 1  # one add
+    assert counter.mul_count == 1  # one mul
+
+
+def test_simple_mutations():
+    """Test simple expression mutations"""
+    x = Var("x", dtype="int32")
+    y = Var("y", dtype="int32")
+
+    # Test multiple replacements
+    expr = Add(x, y)
+    replacer = VariableReplacer({"x": 1, "y": 2})
+    result = replacer.visit_expr(expr)
+
+    assert isinstance(result, Add)
+    assert isinstance(result.a, IntImm)
+    assert isinstance(result.b, IntImm)
+    assert result.a.value == 1
+    assert result.b.value == 2
+
+
+if __name__ == "__main__":
+    test_basic_visitor()
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_inline_private_functions.py

@@ -0,0 +1,255 @@
+# 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: F821
+
+import pytest
+
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+class BaseTestCase:
+    def test_well_formed(self):
+        After = tvm.tirx.transform.InlinePrivateFunctions()(self.Before)
+        tvm.tirx.analysis.verify_well_formed(After)
+
+    def test_produces_expected(self):
+        After = tvm.tirx.transform.InlinePrivateFunctions()(self.Before)
+        tvm.ir.assert_structural_equal(self.Expected, After)
+
+
+class TestSimple(BaseTestCase):
+    """Simple case directly acting on PrimFunc"""
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer([80, 16], "float32"), B: T.Buffer([64, 16], "float32")):
+            for i in range(64):
+                Before.subroutine(T.address_of(A[i, 0]), T.address_of(B[i, 0]))
+
+        @T.prim_func(private=True, s_tir=True)
+        def subroutine(A_data: T.handle("float32"), B_data: T.handle("float32")):
+            A = T.decl_buffer([16, 16], "float32", data=A_data)
+            B = T.decl_buffer([16], "float32", data=B_data)
+            for i in range(16):
+                B[i] = 0.0
+                for j in range(16):
+                    B[i] = B[i] + A[i, j]
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer([80, 16], "float32"), B: T.Buffer([64, 16], "float32")):
+            for i in range(64):
+                A_view_data: T.let[T.handle("float32")] = T.address_of(A[i, 0])
+                Aview = T.decl_buffer([16, 16], "float32", data=A_view_data)
+                B_view_data: T.let[T.handle("float32")] = T.address_of(B[i, 0])
+                Bview = T.decl_buffer([16], "float32", data=B_view_data)
+                for j in range(16):
+                    Bview[j] = 0.0
+                    for k in range(16):
+                        Bview[j] = Bview[j] + Aview[j, k]
+
+
+class TestRetainCrossFunctionSubroutines(BaseTestCase):
+    """Do not inline functions that cross device boundaries
+
+    When lowering TIR, calls for which the callsite and callee have
+    different targets are used at some stages, before being further
+    lowered to explicit device kernel launches.  Since inlining the
+    function would remove this cross-device information,
+    InlinePrivateSubroutines should not inline these cases.
+    """
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer([80, 16], "float32"), B: T.Buffer([64, 16], "float32")):
+            T.func_attr({"target": T.target("llvm")})
+            for i in range(64):
+                Before.subroutine(T.address_of(A[i, 0]), T.address_of(B[i, 0]))
+
+        @T.prim_func(private=True, s_tir=True)
+        def subroutine(A_data: T.handle("float32"), B_data: T.handle("float32")):
+            T.func_attr({"target": T.target("cuda")})
+            A = T.decl_buffer([16, 16], "float32", data=A_data)
+            B = T.decl_buffer([16], "float32", data=B_data)
+            for i in range(16):
+                B[i] = 0.0
+                for j in range(16):
+                    B[i] = B[i] + A[i, j]
+
+    Expected = Before
+
+
+class TestRetainRecursiveSubroutines(BaseTestCase):
+    """Do not inline recursive functions
+
+    To avoid potentially infinite loops at compile-time, disable
+    inlining of recursive functions.  If inlining of these functions
+    would be useful, this restriction may be relaxed with improved
+    analysis of the subroutine.
+    """
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32")):
+            Before.subroutine(T.address_of(A[0]), 16)
+
+        @T.prim_func(private=True, s_tir=True)
+        def subroutine(A_data: T.handle("float32"), A_size: T.int32):
+            A = T.decl_buffer(A_size, "float32", data=A_data)
+            A[1] = A[0] + A[1]
+
+            if A_size > 1:
+                Before.subroutine(T.address_of(A[1]), A_size - 1)
+
+    Expected = Before
+
+
+class TestDeduplicateBlockName(BaseTestCase):
+    """Block names must be de-duplicated after inlining"""
+
+    @pytest.mark.xfail(reason="Inlining of schedulable TIR not yet supported")
+    def test_produces_expected(self):
+        super().test_produces_expected(self)
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer([2, 16], "float32"), B: T.Buffer([2, 16], "float32")):
+            Before.subroutine(T.address_of(A[0, 0]), T.address_of(B[0, 0]))
+            Before.subroutine(T.address_of(A[1, 0]), T.address_of(B[1, 0]))
+
+        @T.prim_func(private=True, s_tir=True)
+        def subroutine(A_data: T.handle("float32"), B_data: T.handle("float32")):
+            A = T.decl_buffer(16, "float32", data=A_data)
+            B = T.decl_buffer(16, "float32", data=B_data)
+            for i in range(16):
+                with T.sblock("scalar_mul"):
+                    B[i] = A[i] * 2.0
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer([80, 16], "float32"), B: T.Buffer([64, 16], "float32")):
+            A_data_1 = T.bind(T.address_of(A[0, 0]), T.handle("float32"))
+            A_1 = T.decl_buffer(16, "float32", data=A_data_1)
+            B_data_1: T.let[T.handle("float32")] = T.address_of(B[0, 0])
+            B_1 = T.decl_buffer(16, "float32", data=B_data_1)
+            for i in range(16):
+                with T.sblock("scalar_mul_1"):
+                    B_1[i] = A_1[i] * 2.0
+
+            A_data_2 = T.bind(T.address_of(A[1, 0]), T.handle("float32"))
+            A_2 = T.decl_buffer(16, "float32", data=A_data_2)
+            B_data_2: T.let[T.handle("float32")] = T.address_of(B[1, 0])
+            B_2 = T.decl_buffer(16, "float32", data=B_data_2)
+            for i in range(16):
+                with T.sblock("scalar_mul_2"):
+                    B_2[i] = A_2[i] * 2.0
+
+
+class TestInlineCallOccurringInExpression(BaseTestCase):
+    """Inline a Call node that is used in a function
+
+    The current implementation only replaces `ir.Call` instances that
+    occur in a `tirx.Evaluate` context.  This is the primary use case,
+    used in destination-passing style.
+
+    This unit test is marked as xfail.  If/when the implementation
+    supports inlining of function calls occurring as part of an
+    expression, the annotation can be removed.
+    """
+
+    @pytest.mark.xfail(reason="Inlining of PrimFuncs outside of tirx.Evaluate is not yet supported")
+    def test_produces_expected(self):
+        super().test_produces_expected(self)
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32")):
+            for i in range(16):
+                A[i] = Before.subroutine(i)
+
+        @T.prim_func(private=True, s_tir=True)
+        def subroutine(i: T.int32) -> T.float32:
+            cos = T.cos(T.cast(i, "float32"))
+            sin = T.sin(T.cast(i, "float32"))
+            retval = cos * cos + sin * sin
+            T.ret(retval)
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32")):
+            for i in range(16):
+                cos = T.cos(T.cast(i, "float32"))
+                sin = T.sin(T.cast(i, "float32"))
+                retval = cos * cos + sin * sin
+                A[i] = retval
+
+
+class TestInlineFunctionWithBufferArguments(BaseTestCase):
+    """Inline a function that accepts buffer arguments
+
+    The current implementation does not support this usage.  This unit
+    test is provided to display a possible user interaction, and is
+    marked with `@pytest.mark.xfail`.  If/when the implementation
+    supports inlining of function calls with buffer arguments, the
+    annotation can be removed.
+    """
+
+    @pytest.mark.xfail(reason="Inlining of PrimFuncs with buffer arguments")
+    def test_produces_expected(self):
+        super().test_produces_expected(self)
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32")):
+            Before.subroutine(
+                T.tvm_stack_make_array(
+                    A.data,
+                    T.tvm_stack_make_shape(*A.shape, dtype="handle"),
+                    0,
+                    len(A.shape),
+                    0.0,
+                    A.elem_offset,
+                    dtype="handle",
+                )
+            )
+
+        @T.prim_func(private=True, s_tir=True)
+        def subroutine(A: T.Buffer(16, "float32")):
+            for i in range(16):
+                A[i] = A[i] * 2.0
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32")):
+            for i in range(16):
+                A[i] = A[i] * 2.0
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_bf16_legalize.py

@@ -0,0 +1,455 @@
+# 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.script
+from tvm.script import tirx as T
+from tvm.target import Target
+from tvm.tirx.transform.transform import BindTarget
+
+
+def u16tof32(v):
+    uint32_v = v.astype("uint32")
+    uint32_v = uint32_v << tvm.tirx.const(16, "uint32")
+    return T.reinterpret("float32", uint32_v)
+
+
+def bf16tof32(v):
+    return u16tof32(T.reinterpret("uint16", v))
+
+
+def f32tou16(v):
+    uint32_v = T.reinterpret("uint32", v)
+    rounding_bias = (uint32_v >> tvm.tirx.const(16, "uint32")) & tvm.tirx.const(1, "uint32")
+    rounding_bias += tvm.tirx.const(0x7FFF, "uint32")
+    uint32_v = uint32_v + rounding_bias
+    return (uint32_v >> tvm.tirx.const(16, "uint32")).astype("uint16")
+
+
+def f32tobf16(v):
+    return T.reinterpret("bfloat16", f32tou16(v))
+
+
+def test_bf16_simple_store_will_legalize():
+    def get_before():
+        @tvm.script.ir_module
+        class Before:
+            @T.prim_func(s_tir=True)
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+                Cptr: T.handle("bfloat16"),
+            ):
+                T.func_attr({"global_symbol": "main"})
+                A = T.decl_buffer((100,), "bfloat16", data=Aptr)
+                B = T.decl_buffer((100,), "bfloat16")
+                C = T.decl_buffer((100,), "bfloat16", data=Cptr)
+                for i in T.grid(100):
+                    B[i] = A[i]
+                    C[i] = T.exp(B[i])
+
+        return Before
+
+    def after_compute_legalize():
+        @tvm.script.ir_module
+        class After:
+            @T.prim_func(s_tir=True)
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+                Cptr: T.handle("bfloat16"),
+            ):
+                T.func_attr({"global_symbol": "main"})
+                A = T.decl_buffer((100,), "bfloat16", data=Aptr)
+                B = T.decl_buffer((100,), "float32")
+                C = T.decl_buffer((100,), "bfloat16", data=Cptr)
+                for i in T.grid(100):
+                    B[i] = bf16tof32(A[i])
+                    C[i] = f32tobf16(T.exp(B[i]))
+
+        return After
+
+    def after_storage_legalize():
+        @tvm.script.ir_module
+        class After:
+            @T.prim_func(s_tir=True)
+            def main(
+                Aptr: T.handle("uint16", storage_scope="shared"),
+                Cptr: T.handle("uint16"),
+            ):
+                T.func_attr({"global_symbol": "main"})
+                A = T.decl_buffer((100,), "uint16", data=Aptr)
+                B = T.decl_buffer((100,), "float32")
+                C = T.decl_buffer((100,), "uint16", data=Cptr)
+                for i in T.grid(100):
+                    B[i] = u16tof32(A[i])
+                    C[i] = f32tou16(T.exp(B[i]))
+
+        return After
+
+    target = Target("nvidia/geforce-rtx-2080-ti")
+    before = BindTarget(target)(get_before())
+    after_compute = tvm.tirx.transform.BF16ComputeLegalize()(before)
+    after_storage = tvm.tirx.transform.BF16StorageLegalize()(after_compute)
+    tvm.ir.assert_structural_equal(after_compute, BindTarget(target)(after_compute_legalize()))
+    tvm.ir.assert_structural_equal(after_storage, BindTarget(target)(after_storage_legalize()))
+
+
+def test_bf16_storage_compute_scope_will_legalize():
+    def get_before():
+        @tvm.script.ir_module
+        class Before:
+            @T.prim_func(s_tir=True)
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+                Bptr: T.handle("bfloat16", storage_scope="local"),
+                Dptr: T.handle("bfloat16"),
+            ):
+                T.func_attr({"global_symbol": "main"})
+                A = T.decl_buffer((100,), "bfloat16", data=Aptr)
+                B = T.decl_buffer((100,), "bfloat16", data=Bptr)
+                D = T.decl_buffer((100,), "bfloat16", data=Dptr)
+                C = T.decl_buffer((100,), "bfloat16")
+                for i in T.grid(100):
+                    C[i] = A[i] + B[i]
+                    D[i] = T.exp(C[i])
+
+        return Before
+
+    def after_compute_legalize():
+        @tvm.script.ir_module
+        class After:
+            @T.prim_func(s_tir=True)
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+                Bptr: T.handle("bfloat16", storage_scope="local"),
+                Dptr: T.handle("bfloat16"),
+            ):
+                T.func_attr({"global_symbol": "main"})
+                A = T.decl_buffer((100,), "bfloat16", data=Aptr)
+                B = T.decl_buffer((100,), "bfloat16", data=Bptr)
+                D = T.decl_buffer((100,), "bfloat16", data=Dptr)
+                C = T.decl_buffer((100,), "float32")
+                for i in T.grid(100):
+                    C[i] = bf16tof32(A[i]) + bf16tof32(B[i])
+                    D[i] = f32tobf16(T.exp(C[i]))
+
+        return After
+
+    def after_storage_legalize():
+        @tvm.script.ir_module
+        class After:
+            @T.prim_func(s_tir=True)
+            def main(
+                Aptr: T.handle("uint16", storage_scope="shared"),
+                Bptr: T.handle("uint16", storage_scope="local"),
+                Dptr: T.handle("uint16"),
+            ):
+                T.func_attr({"global_symbol": "main"})
+                A = T.decl_buffer((100,), "uint16", data=Aptr)
+                B = T.decl_buffer((100,), "uint16", data=Bptr)
+                D = T.decl_buffer((100,), "uint16", data=Dptr)
+                C = T.decl_buffer((100,), "float32")
+                for i in T.grid(100):
+                    C[i] = u16tof32(A[i]) + u16tof32(B[i])
+                    D[i] = f32tou16(T.exp(C[i]))
+
+        return After
+
+    target = Target("nvidia/geforce-rtx-2080-ti")
+    before = BindTarget(target)(get_before())
+    after_compute = tvm.tirx.transform.BF16ComputeLegalize()(before)
+    after_storage = tvm.tirx.transform.BF16StorageLegalize()(after_compute)
+    tvm.ir.assert_structural_equal(after_compute, BindTarget(target)(after_compute_legalize()))
+    tvm.ir.assert_structural_equal(after_storage, BindTarget(target)(after_storage_legalize()))
+
+
+def test_bf16_storage_compute_scope_wont_legalize():
+    def get_before():
+        @tvm.script.ir_module
+        class Before:
+            @T.prim_func(s_tir=True)
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+                Bptr: T.handle("bfloat16", storage_scope="local"),
+                Dptr: T.handle("bfloat16"),
+            ):
+                T.func_attr({"global_symbol": "main"})
+                A = T.decl_buffer((100,), "bfloat16", data=Aptr)
+                B = T.decl_buffer((100,), "bfloat16", data=Bptr)
+                D = T.decl_buffer((100,), "bfloat16", data=Dptr)
+                C = T.decl_buffer((100,), "bfloat16")
+                for i in T.grid(100):
+                    C[i] = A[i] + B[i]
+                    D[i] = T.exp(C[i])
+
+        return Before
+
+    def after_compute_legalize():
+        @tvm.script.ir_module
+        class After:
+            @T.prim_func(s_tir=True)
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+                Bptr: T.handle("bfloat16", storage_scope="local"),
+                Dptr: T.handle("bfloat16"),
+            ):
+                T.func_attr({"global_symbol": "main"})
+                A = T.decl_buffer((100,), "bfloat16", data=Aptr)
+                B = T.decl_buffer((100,), "bfloat16", data=Bptr)
+                D = T.decl_buffer((100,), "bfloat16", data=Dptr)
+                C = T.decl_buffer((100,), "bfloat16")
+                for i in T.grid(100):
+                    C[i] = A[i] + B[i]
+                    D[i] = T.exp(C[i])
+
+        return After
+
+    def after_storage_legalize():
+        @tvm.script.ir_module
+        class After:
+            @T.prim_func(s_tir=True)
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+                Bptr: T.handle("bfloat16", storage_scope="local"),
+                Dptr: T.handle("bfloat16"),
+            ):
+                T.func_attr({"global_symbol": "main"})
+                A = T.decl_buffer((100,), "bfloat16", data=Aptr)
+                B = T.decl_buffer((100,), "bfloat16", data=Bptr)
+                D = T.decl_buffer((100,), "bfloat16", data=Dptr)
+                C = T.decl_buffer((100,), "bfloat16")
+                for i in T.grid(100):
+                    C[i] = A[i] + B[i]
+                    D[i] = T.exp(C[i])
+
+        return After
+
+    target = Target("nvidia/geforce-rtx-3090-ti")
+    before = BindTarget(target)(get_before())
+    after_compute = tvm.tirx.transform.BF16ComputeLegalize()(before)
+    after_storage = tvm.tirx.transform.BF16StorageLegalize()(after_compute)
+    tvm.ir.assert_structural_equal(after_compute, BindTarget(target)(after_compute_legalize()))
+    tvm.ir.assert_structural_equal(after_storage, BindTarget(target)(after_storage_legalize()))
+
+
+def test_bf16_reduce_will_legalize():
+    def get_before():
+        @tvm.script.ir_module
+        class Before:
+            @T.prim_func(private=True, s_tir=True)
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+            ):
+                A_flat = T.decl_buffer(4096, "bfloat16", data=Aptr)
+
+                for i in range(128):
+                    threadIdx_x = T.launch_thread("threadIdx.x", 32)
+
+                    reduce = T.decl_buffer(1, dtype="bfloat16", scope="local")
+
+                    with T.attr(
+                        T.comm_reducer(lambda x, y: x + y, [T.bfloat16(0)]),
+                        "reduce_scope",
+                        T.int32(0),
+                    ):
+                        T.tvm_thread_allreduce(
+                            T.uint32(1),
+                            A_flat[0],
+                            T.bool(True),
+                            reduce[0],
+                            threadIdx_x,
+                        )
+
+        return Before
+
+    def after_compute_legalize():
+        @tvm.script.ir_module
+        class After:
+            @T.prim_func(private=True, s_tir=True)
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+            ):
+                A_flat_1 = T.decl_buffer(4096, "bfloat16", data=Aptr)
+
+                for i in range(128):
+                    threadIdx_x = T.launch_thread("threadIdx.x", 32)
+
+                    reduce = T.decl_buffer(1, dtype="float32", scope="local")
+
+                    with T.attr(
+                        T.comm_reducer(lambda x, y: x + y, [T.float32(0)]),
+                        "reduce_scope",
+                        T.int32(0),
+                    ):
+                        T.tvm_thread_allreduce(
+                            T.uint32(1),
+                            T.reinterpret(
+                                "float32",
+                                T.shift_left(
+                                    T.Cast("uint32", T.reinterpret("uint16", A_flat_1[0])),
+                                    T.uint32(16),
+                                ),
+                            ),
+                            T.bool(True),
+                            reduce[0],
+                            threadIdx_x,
+                        )
+
+        return After
+
+    def after_storage_legalize():
+        @tvm.script.ir_module
+        class After:
+            @T.prim_func(private=True, s_tir=True)
+            def main(
+                Aptr: T.handle("uint16", storage_scope="shared"),
+            ):
+                A_flat_1 = T.decl_buffer(4096, "uint16", data=Aptr)
+
+                for i in range(128):
+                    threadIdx_x = T.launch_thread("threadIdx.x", 32)
+
+                    reduce = T.decl_buffer(1, dtype="float32", scope="local")
+
+                    with T.attr(
+                        T.comm_reducer(lambda x, y: x + y, [T.float32(0)]),
+                        "reduce_scope",
+                        T.int32(0),
+                    ):
+                        T.tvm_thread_allreduce(
+                            T.uint32(1),
+                            T.reinterpret(
+                                "float32",
+                                T.shift_left(
+                                    T.Cast("uint32", T.reinterpret("uint16", A_flat_1[0])),
+                                    T.uint32(16),
+                                ),
+                            ),
+                            T.bool(True),
+                            reduce[0],
+                            threadIdx_x,
+                        )
+
+        return After
+
+    target = Target("nvidia/geforce-rtx-2080-ti")
+    before = BindTarget(target)(get_before())
+    after_compute = tvm.tirx.transform.BF16ComputeLegalize()(before)
+    after_storage = tvm.tirx.transform.BF16StorageLegalize()(after_compute)
+    tvm.ir.assert_structural_equal(after_compute, BindTarget(target)(after_compute_legalize()))
+    tvm.ir.assert_structural_equal(after_storage, BindTarget(target)(after_storage_legalize()))
+
+
+def test_bf16_reduce_wont_legalize():
+    def get_before():
+        @tvm.script.ir_module
+        class Before:
+            @T.prim_func(private=True, s_tir=True)
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+            ):
+                A_flat = T.decl_buffer(4096, "bfloat16", data=Aptr)
+
+                for i in range(128):
+                    threadIdx_x = T.launch_thread("threadIdx.x", 32)
+
+                    reduce = T.decl_buffer(1, dtype="bfloat16", scope="local")
+
+                    with T.attr(
+                        T.comm_reducer(lambda x, y: x + y, [T.bfloat16(0)]),
+                        "reduce_scope",
+                        T.int32(0),
+                    ):
+                        T.tvm_thread_allreduce(
+                            T.uint32(1),
+                            A_flat[0],
+                            T.bool(True),
+                            reduce[0],
+                            threadIdx_x,
+                        )
+
+        return Before
+
+    def after_compute_legalize():
+        @tvm.script.ir_module
+        class After:
+            @T.prim_func(private=True, s_tir=True)
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+            ):
+                A_flat = T.decl_buffer(4096, "bfloat16", data=Aptr)
+
+                for i in range(128):
+                    threadIdx_x = T.launch_thread("threadIdx.x", 32)
+
+                    reduce = T.decl_buffer(1, dtype="bfloat16", scope="local")
+
+                    with T.attr(
+                        T.comm_reducer(lambda x, y: x + y, [T.bfloat16(0)]),
+                        "reduce_scope",
+                        T.int32(0),
+                    ):
+                        T.tvm_thread_allreduce(
+                            T.uint32(1),
+                            A_flat[0],
+                            T.bool(True),
+                            reduce[0],
+                            threadIdx_x,
+                        )
+
+        return After
+
+    def after_storage_legalize():
+        @tvm.script.ir_module
+        class After:
+            @T.prim_func(private=True, s_tir=True)
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+            ):
+                A_flat = T.decl_buffer(4096, "bfloat16", data=Aptr)
+
+                for i in range(128):
+                    threadIdx_x = T.launch_thread("threadIdx.x", 32)
+
+                    reduce = T.decl_buffer(1, dtype="bfloat16", scope="local")
+
+                    with T.attr(
+                        T.comm_reducer(lambda x, y: x + y, [T.bfloat16(0)]),
+                        "reduce_scope",
+                        T.int32(0),
+                    ):
+                        T.tvm_thread_allreduce(
+                            T.uint32(1),
+                            A_flat[0],
+                            T.bool(True),
+                            reduce[0],
+                            threadIdx_x,
+                        )
+
+        return After
+
+    target = Target("nvidia/geforce-rtx-3090-ti")
+    before = BindTarget(target)(get_before())
+    after_compute = tvm.tirx.transform.BF16ComputeLegalize()(before)
+    after_storage = tvm.tirx.transform.BF16StorageLegalize()(after_compute)
+    tvm.ir.assert_structural_equal(after_compute, BindTarget(target)(after_compute_legalize()))
+    tvm.ir.assert_structural_equal(after_storage, BindTarget(target)(after_storage_legalize()))
+
+
+if __name__ == "__main__":
+    test_bf16_storage_compute_scope_will_legalize()
+    test_bf16_storage_compute_scope_wont_legalize()
+    test_bf16_reduce_will_legalize()
+    test_bf16_reduce_wont_legalize()

tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py

@@ -0,0 +1,780 @@
+# 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 hashlib
+
+import tvm
+from tvm.ir.base import save_json
+from tvm.script import tirx as T
+
+
+# =====================================================================
+# T1: Basic multi-level CSE
+# Two common subexpressions at different scoping levels.
+# =====================================================================
+def test_basic():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), i1: T.int32, i2: T.int32, z3: T.int32):
+            z1 = T.bind(1)
+            z2 = T.bind(2)
+            B[i1] = z1 + z2
+            x = T.bind(1)
+            y = T.bind(1)
+            a = T.bind((x + y) + (z1 + z2))
+            b = T.bind((x + y) + z3)
+            B[i2] = a + b
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), i1: T.int32, i2: T.int32, z3: T.int32):
+            z1 = T.bind(1)
+            z2 = T.bind(2)
+            cse_v1 = T.bind(z1 + z2)
+            B[i1] = cse_v1
+            x = T.bind(1)
+            y = T.bind(1)
+            cse_v2 = T.bind(x + y)
+            a = T.bind(cse_v2 + cse_v1)
+            b = T.bind(cse_v2 + z3)
+            B[i2] = a + b
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T2: If -- single-branch CSE
+# Duplicated expression only in then-branch stays inside then-branch.
+# =====================================================================
+def test_if_single_branch():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(
+            B: T.Buffer((50,), "int32"),
+            i1: T.int32,
+            i2: T.int32,
+            i3: T.int32,
+            y: T.int32,
+            z: T.int32,
+        ):
+            b = T.bind(1)
+            if b:
+                B[i1] = y + z
+                B[i2] = y + z
+            else:
+                B[i3] = y
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(
+            B: T.Buffer((50,), "int32"),
+            i1: T.int32,
+            i2: T.int32,
+            i3: T.int32,
+            y: T.int32,
+            z: T.int32,
+        ):
+            b = T.bind(1)
+            if b:
+                cse_v1 = T.bind(y + z)
+                B[i1] = cse_v1
+                B[i2] = cse_v1
+            else:
+                B[i3] = y
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T3: If -- both-branch CSE
+# Duplicated expression in both branches is hoisted before the if.
+# =====================================================================
+def test_if_both_branches():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(
+            B: T.Buffer((50,), "int32"),
+            i1: T.int32,
+            i2: T.int32,
+            i3: T.int32,
+            y: T.int32,
+            z: T.int32,
+        ):
+            b = T.bind(1)
+            if b:
+                B[i1] = y + z
+                B[i2] = y
+            else:
+                B[i3] = y + z
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(
+            B: T.Buffer((50,), "int32"),
+            i1: T.int32,
+            i2: T.int32,
+            i3: T.int32,
+            y: T.int32,
+            z: T.int32,
+        ):
+            b = T.bind(1)
+            cse_v1 = T.bind(y + z)
+            if b:
+                B[i1] = cse_v1
+                B[i2] = y
+            else:
+                B[i3] = cse_v1
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T4: Cascade CSE
+# Introducing (x+y)+z creates opportunity for x+y.
+# =====================================================================
+def test_cascade():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(
+            B: T.Buffer((50,), "int32"),
+            i1: T.int32,
+            i2: T.int32,
+            i3: T.int32,
+            x: T.int32,
+            y: T.int32,
+            z: T.int32,
+        ):
+            B[i1] = (x + y) + z
+            B[i2] = (x + y) + z
+            B[i3] = x + y
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(
+            B: T.Buffer((50,), "int32"),
+            i1: T.int32,
+            i2: T.int32,
+            i3: T.int32,
+            x: T.int32,
+            y: T.int32,
+            z: T.int32,
+        ):
+            cse_v2 = T.bind(x + y)
+            cse_v1 = T.bind(cse_v2 + z)
+            B[i1] = cse_v1
+            B[i2] = cse_v1
+            B[i3] = cse_v2
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T5: No change when no duplication
+# Single occurrence — pass is identity.
+# =====================================================================
+def test_no_duplication():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(x: T.int32, y: T.int32, z: T.int32):
+            a = T.bind(x + (y + z))
+            T.evaluate(a)
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(x: T.int32, y: T.int32, z: T.int32):
+            a = T.bind(x + (y + z))
+            T.evaluate(a)
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T6: Deterministic output
+# Multiple runs on same input produce identical output (same var names,
+# ordering). Verified via JSON serialization hash.
+# =====================================================================
+def test_deterministic():
+    NUM_TERMS = 10
+    REPEATS = 10
+
+    x = tvm.tirx.Var("x", "int32")
+    result = tvm.tirx.Var("result", "int32")
+
+    offsets = sorted([i + 1 for i in range(NUM_TERMS)])
+    inc1 = [(x + offsets[i]) for i in range(NUM_TERMS)]
+    inc2 = [(x + offsets[i]) for i in range(NUM_TERMS)]
+
+    expression = x
+    for add in inc1 + inc2:
+        expression = expression + add
+    body = tvm.tirx.SeqStmt([tvm.tirx.Bind(result, expression), tvm.tirx.Evaluate(result)])
+    mod = tvm.IRModule.from_expr(tvm.tirx.PrimFunc([x], body))
+
+    initial_hash = None
+    for _ in range(REPEATS):
+        out = tvm.tirx.transform.CommonSubexprElim()(mod)
+        func = out["main"]
+        json_val = save_json(func)
+        json_hash = hashlib.sha256(json_val.encode()).hexdigest()
+        if initial_hash is None:
+            initial_hash = json_hash
+        assert json_hash == initial_hash
+
+
+# =====================================================================
+# T7: CSE inside for-loop
+# Common sub-expression inside a for-loop body.
+# =====================================================================
+def test_for_loop():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), y: T.int32, z: T.int32):
+            for i in range(10):
+                B[i] = y + z
+                B[i + 10] = y + z
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), y: T.int32, z: T.int32):
+            for i in range(10):
+                cse_v1 = T.bind(y + z)
+                B[i] = cse_v1
+                B[i + 10] = cse_v1
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T8: CSE across for-loop and outer scope
+# Expression appears both outside and inside a for-loop. Binding placed
+# in outer scope before the for.
+# =====================================================================
+def test_for_hoist():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), y: T.int32, z: T.int32):
+            B[0] = y + z
+            for i in range(10):
+                B[i + 1] = y + z
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), y: T.int32, z: T.int32):
+            cse_v1 = T.bind(y + z)
+            B[0] = cse_v1
+            for i in range(10):
+                B[i + 1] = cse_v1
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T9: Cannot-lift -- expressions containing BufferLoad
+# Expressions containing BufferLoad are ineligible even when duplicated,
+# because lifting them would change semantics (buffer may alias).
+# =====================================================================
+def test_cannot_lift_bufferload():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((50,), "int32"), B: T.Buffer((50,), "int32")):
+            B[0] = A[0] + A[0]
+            B[1] = A[0] + A[0]
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((50,), "int32"), B: T.Buffer((50,), "int32")):
+            B[0] = A[0] + A[0]
+            B[1] = A[0] + A[0]
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T10: Nested if -- multi-level scope LCA
+# Expression in both branches of an inner if, but not the outer else.
+# Binding placed at the inner if's parent scope.
+# =====================================================================
+def test_nested_if():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(
+            B: T.Buffer((50,), "int32"),
+            c1: T.int32,
+            c2: T.int32,
+            y: T.int32,
+            z: T.int32,
+        ):
+            if c1:
+                if c2:
+                    B[0] = y + z
+                else:
+                    B[1] = y + z
+            else:
+                B[2] = y
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(
+            B: T.Buffer((50,), "int32"),
+            c1: T.int32,
+            c2: T.int32,
+            y: T.int32,
+            z: T.int32,
+        ):
+            if c1:
+                cse_v1 = T.bind(y + z)
+                if c2:
+                    B[0] = cse_v1
+                else:
+                    B[1] = cse_v1
+            else:
+                B[2] = y
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T11: Multiple independent CSE candidates
+# Several independent expressions, each duplicated.
+# =====================================================================
+def test_multi_independent():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(
+            B: T.Buffer((50,), "int32"),
+            a: T.int32,
+            b: T.int32,
+            c: T.int32,
+            d: T.int32,
+        ):
+            B[0] = a + b
+            B[1] = c + d
+            B[2] = a + b
+            B[3] = c + d
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(
+            B: T.Buffer((50,), "int32"),
+            a: T.int32,
+            b: T.int32,
+            c: T.int32,
+            d: T.int32,
+        ):
+            cse_v1 = T.bind(a + b)
+            B[0] = cse_v1
+            cse_v2 = T.bind(c + d)
+            B[1] = cse_v2
+            B[2] = cse_v1
+            B[3] = cse_v2
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T12: Expression in if-condition and branch
+# One occurrence in the if-condition (parent scope), one in the
+# then-branch. Binding hoisted before the if.
+# =====================================================================
+def test_if_condition():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), y: T.int32, z: T.int32):
+            if y + z > 0:
+                B[0] = y + z
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), y: T.int32, z: T.int32):
+            cse_v1 = T.bind(y + z)
+            if cse_v1 > 0:
+                B[0] = cse_v1
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T13: Cannot-lift -- expression containing Call
+# Function calls cannot be lifted.
+# =====================================================================
+def test_cannot_lift_call():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), x: T.int32):
+            B[0] = T.call_extern("my_func", x, dtype="int32") + 1
+            B[1] = T.call_extern("my_func", x, dtype="int32") + 1
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), x: T.int32):
+            B[0] = T.call_extern("my_func", x, dtype="int32") + 1
+            B[1] = T.call_extern("my_func", x, dtype="int32") + 1
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T14: No single-use binding
+# When all occurrences of a sub-expression are inside a deeper
+# expression that is also CSE'd, the sub-expression should NOT get
+# its own binding (it would be used only once in the parent's value).
+# =====================================================================
+def test_no_single_use_binding():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(
+            B: T.Buffer((50,), "int32"),
+            x: T.int32,
+            y: T.int32,
+            z: T.int32,
+        ):
+            B[0] = (x + y) + z
+            B[1] = (x + y) + z
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(
+            B: T.Buffer((50,), "int32"),
+            x: T.int32,
+            y: T.int32,
+            z: T.int32,
+        ):
+            cse_v1 = T.bind((x + y) + z)
+            B[0] = cse_v1
+            B[1] = cse_v1
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T15: For-loop extent lifting
+# Expression shared between loop extent and body is hoisted outside the
+# loop rather than being re-evaluated every iteration.
+# =====================================================================
+def test_for_extent_lift():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), y: T.int32, z: T.int32):
+            for i in range(y + z):
+                B[i] = y + z
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), y: T.int32, z: T.int32):
+            cse_v1 = T.bind(y + z)
+            for i in range(cse_v1):
+                B[i] = cse_v1
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T16: Loop-variable expression stays inside loop
+# An expression using the loop variable (i*4) that appears multiple
+# times inside the body must NOT be hoisted outside the loop.
+# =====================================================================
+def test_loop_var_expr_stays_inside():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((50,), "int32"),
+            B: T.Buffer((50,), "int32"),
+        ):
+            for i in range(10):
+                A[i * 4] = B[i * 4]
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((50,), "int32"),
+            B: T.Buffer((50,), "int32"),
+        ):
+            for i in range(10):
+                cse_v1 = T.bind(i * 4)
+                A[cse_v1] = B[cse_v1]
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T17: No normalization without commoning
+# Single-occurrence expression is not normalized or modified.
+# =====================================================================
+def test_no_normalization_without_commoning():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(x: T.int32, y: T.int32, z: T.int32):
+            a = T.bind(x + (y + z))
+            T.evaluate(a)
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(x: T.int32, y: T.int32, z: T.int32):
+            a = T.bind(x + (y + z))
+            T.evaluate(a)
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+# =====================================================================
+# T18: Let-bound variable -- no extraction from Let body
+# Expressions inside a Let body that reference the Let-bound variable
+# must NOT be extracted, because the variable is only in scope inside
+# the Let body, not at the statement level where Bind would be placed.
+# =====================================================================
+def test_let_body_no_extraction():
+    """CSE must not extract expressions from Let bodies that use Let-bound vars."""
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    # Let(x, 1, (x+y) + (x+y)) -- x+y appears twice but x is Let-bound
+    let_expr = tvm.tirx.Let(x, tvm.tirx.IntImm("int32", 1), (x + y) + (x + y))
+    buf = tvm.tirx.decl_buffer((10,), "int32", name="B")
+    i = tvm.tirx.Var("i", "int32")
+    store = tvm.tirx.BufferStore(buf, let_expr, [i])
+    loop = tvm.tirx.For(
+        i,
+        tvm.tirx.const(0, "int32"),
+        tvm.tirx.const(10, "int32"),
+        tvm.tirx.ForKind.SERIAL,
+        store,
+    )
+    func = tvm.tirx.PrimFunc([buf, y], loop)
+    mod = tvm.IRModule({"main": func})
+    mod_after = tvm.tirx.transform.CommonSubexprElim()(mod)
+    # No CSE variables should be introduced
+    script = mod_after["main"].script()
+    assert "cse_v" not in script, f"CSE incorrectly extracted from Let body:\n{script}"
+
+
+# =====================================================================
+# T19: Let value -- CSE works for shared Let value expressions
+# The Let value is evaluated before the binding takes effect, so
+# expressions in the Let value CAN be CSE'd with expressions outside.
+# =====================================================================
+def test_let_value_cse():
+    """CSE can extract from Let values (computed before binding)."""
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    z = tvm.tirx.Var("z", "int32")
+    # Let(x, y+z, x+1) with y+z also appearing outside the Let
+    let_expr = tvm.tirx.Let(x, y + z, x + 1)
+    buf = tvm.tirx.decl_buffer((10,), "int32", name="B")
+    i = tvm.tirx.Var("i", "int32")
+    store = tvm.tirx.BufferStore(buf, (y + z) + let_expr, [i])
+    loop = tvm.tirx.For(
+        i,
+        tvm.tirx.const(0, "int32"),
+        tvm.tirx.const(10, "int32"),
+        tvm.tirx.ForKind.SERIAL,
+        store,
+    )
+    func = tvm.tirx.PrimFunc([buf, y, z], loop)
+    mod = tvm.IRModule({"main": func})
+    mod_after = tvm.tirx.transform.CommonSubexprElim()(mod)
+    # y+z should be extracted (appears in Let value AND outside)
+    script = mod_after["main"].script()
+    assert "cse_v" in script, f"CSE should extract y+z from Let value:\n{script}"
+
+
+# =====================================================================
+# T20: Nested Let -- both levels respected
+# Multiple nested Let expressions with common sub-expressions in the
+# innermost body. None should be extracted.
+# =====================================================================
+def test_nested_let_no_extraction():
+    """CSE must not extract from nested Let bodies."""
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    z = tvm.tirx.Var("z", "int32")
+    # Let(x, 1, Let(y, 2, (x+y+z) + (x+y+z)))
+    inner = (x + y + z) + (x + y + z)
+    nested_let = tvm.tirx.Let(
+        x, tvm.tirx.IntImm("int32", 1), tvm.tirx.Let(y, tvm.tirx.IntImm("int32", 2), inner)
+    )
+    buf = tvm.tirx.decl_buffer((10,), "int32", name="B")
+    i = tvm.tirx.Var("i", "int32")
+    store = tvm.tirx.BufferStore(buf, nested_let, [i])
+    loop = tvm.tirx.For(
+        i,
+        tvm.tirx.const(0, "int32"),
+        tvm.tirx.const(10, "int32"),
+        tvm.tirx.ForKind.SERIAL,
+        store,
+    )
+    func = tvm.tirx.PrimFunc([buf, z], loop)
+    mod = tvm.IRModule({"main": func})
+    mod_after = tvm.tirx.transform.CommonSubexprElim()(mod)
+    script = mod_after["main"].script()
+    assert "cse_v" not in script, f"CSE incorrectly extracted from nested Let body:\n{script}"
+
+
+# =====================================================================
+# T21: Let with lowered floordiv pattern
+# Simulates the pattern produced by LowerIntrin for floordiv:
+# Let(rmod, truncmod(a,b), Let(rdiv, div(a,b), Select(...)))
+# wrapped in Let(x, load, Let(y, load, ...))
+# This is the regression test for the CI failure in lower_intrin tests.
+# =====================================================================
+def test_let_floordiv_pattern():
+    """CSE must handle the Let pattern from LowerIntrin's floordiv lowering."""
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    rmod = tvm.tirx.Var("rmod", "int32")
+    rdiv = tvm.tirx.Var("rdiv", "int32")
+    # Simulate lowered floordiv: Let(rmod, x%y, Let(rdiv, x/y, Select(...)))
+    select_cond = tvm.tirx.And(y >= 0, rmod >= 0) | tvm.tirx.And(y < 0, rmod <= 0)
+    select_expr = tvm.tirx.Select(select_cond, rdiv, rdiv - 1)
+    inner_let = tvm.tirx.Let(rdiv, tvm.tirx.Div(x, y), select_expr)
+    outer_let = tvm.tirx.Let(rmod, tvm.tirx.Mod(x, y), inner_let)
+    # Wrap in Let(x, load, Let(y, load, ...))
+    buf_a = tvm.tirx.decl_buffer((10,), "int32", name="A")
+    buf_b = tvm.tirx.decl_buffer((10,), "int32", name="B")
+    buf_c = tvm.tirx.decl_buffer((10,), "int32", name="C")
+    i = tvm.tirx.Var("i", "int32")
+    full_expr = tvm.tirx.Let(
+        x,
+        tvm.tirx.BufferLoad(buf_a, [i]),
+        tvm.tirx.Let(y, tvm.tirx.BufferLoad(buf_b, [i]), outer_let),
+    )
+    store = tvm.tirx.BufferStore(buf_c, full_expr, [i])
+    loop = tvm.tirx.For(
+        i,
+        tvm.tirx.const(0, "int32"),
+        tvm.tirx.const(10, "int32"),
+        tvm.tirx.ForKind.SERIAL,
+        store,
+    )
+    func = tvm.tirx.PrimFunc([buf_a, buf_b, buf_c], loop)
+    mod = tvm.IRModule({"main": func})
+    # Should not crash and should not extract Let-bound vars
+    mod_after = tvm.tirx.transform.CommonSubexprElim()(mod)
+    script = mod_after["main"].script()
+    assert "cse_v" not in script, f"CSE incorrectly extracted from Let body:\n{script}"
+
+
+# =====================================================================
+# T22: No lifting of bool predicate (comparison expression)
+# A duplicated `i < n` feeds two if-statements.  CSE must leave it
+# inline rather than hoisting a `cse_v: bool = (i < n)` binding.
+# =====================================================================
+def test_no_lift_bool_predicate():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), n: T.int32, x: T.int32):
+            for i in range(50):
+                if i < n:
+                    B[i] = x
+                if i < n:
+                    B[i] = x + 1
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Before)
+    assert "cse_v" not in after["main"].script()
+
+
+# =====================================================================
+# T23: No lifting of bool logical expression (And)
+# A duplicated `a && b` feeds two if-statements.  CSE must leave it
+# inline rather than hoisting a `cse_v: bool = T.And(a, b)` binding.
+# =====================================================================
+def test_no_lift_bool_logical():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((50,), "int32"), a: T.bool, b: T.bool, x: T.int32):
+            if T.And(a, b):
+                B[0] = x
+            if T.And(a, b):
+                B[1] = x + 1
+
+    after = tvm.tirx.transform.CommonSubexprElim()(Before)
+    tvm.ir.assert_structural_equal(after, Before)
+    assert "cse_v" not in after["main"].script()
+
+
+if __name__ == "__main__":
+    test_basic()
+    test_if_single_branch()
+    test_if_both_branches()
+    test_cascade()
+    test_no_duplication()
+    test_deterministic()
+    test_for_loop()
+    test_for_hoist()
+    test_cannot_lift_bufferload()
+    test_nested_if()
+    test_multi_independent()
+    test_if_condition()
+    test_cannot_lift_call()
+    test_no_single_use_binding()
+    test_for_extent_lift()
+    test_loop_var_expr_stays_inside()
+    test_no_normalization_without_commoning()
+    test_let_body_no_extraction()
+    test_let_value_cse()
+    test_nested_let_no_extraction()
+    test_let_floordiv_pattern()
+    test_no_lift_bool_predicate()
+    test_no_lift_bool_logical()

tests/python/tirx-transform/test_tir_transform_convert_ssa.py

@@ -0,0 +1,539 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401
+
+import tvm
+import tvm.testing
+from tvm import ir, tirx
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_reuse_in_sequential_bind():
+    """De-dup sequential variable bindings"""
+
+    # Manually construct the PrimFunc body, as SSA violations are
+    # not valid TIR, and may not be expressible in future versions
+    # of TVMSCript.
+    var = tirx.Var("var", "int32")
+    sequential_bindings = tirx.SeqStmt(
+        [
+            tirx.Bind(var, 16),
+            tirx.Evaluate(var),
+            tirx.Bind(var, 32),
+            tirx.Evaluate(var),
+        ]
+    )
+    before = tirx.PrimFunc([], sequential_bindings).with_attr("s_tir", True)
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected():
+        var1 = T.bind(T.int32(16))
+        T.evaluate(var1)
+        var2 = T.bind(T.int32(32))
+        T.evaluate(var2)
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = tvm.tirx.transform.ConvertSSA()(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_reuse_in_nested_bind():
+    """De-dup sequential bindings of the same variable.
+
+    In the flat Bind model, all Binds are siblings in a SeqStmt. A second
+    Bind of the same variable redefines it for all subsequent siblings.
+    ConvertSSA should create a new variable for the second binding and
+    update all subsequent uses to refer to the new variable.
+    """
+
+    # Manually construct the PrimFunc body, as SSA violations are
+    # not valid TIR, and may not be expressible in future versions
+    # of TVMScript.
+    var = tirx.Var("var", "int32")
+    # Note: nested SeqStmt is flattened by the IR builder, so the input
+    # is actually a flat SeqStmt with 5 elements.
+    inner_seq = tirx.SeqStmt(
+        [
+            tirx.Bind(var, 16),
+            tirx.Evaluate(var),
+        ]
+    )
+    outer_seq = tirx.SeqStmt(
+        [
+            tirx.Bind(var, 32),
+            tirx.Evaluate(var),
+            inner_seq,
+            tirx.Evaluate(var),
+        ]
+    )
+    before = tirx.PrimFunc([], outer_seq)
+
+    # In the flat model, the second Bind(var, 16) redefines var for
+    # ALL subsequent siblings including the last Evaluate.
+    var1 = tirx.Var("var", "int32")
+    var2 = tirx.Var("var", "int32")
+    expected_body = tirx.SeqStmt(
+        [
+            tirx.Bind(var1, 32),
+            tirx.Evaluate(var1),
+            tirx.Bind(var2, 16),
+            tirx.Evaluate(var2),
+            tirx.Evaluate(var2),
+        ]
+    )
+    expected = tirx.PrimFunc([], expected_body)
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = tvm.tirx.transform.ConvertSSA()(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_reused_var_across_module():
+    """De-duplicate Var bindings across entire module"""
+
+    @T.prim_func(private=True, s_tir=True)
+    def func():
+        var = T.bind(10)
+        T.evaluate(var)
+
+    before = tvm.IRModule(
+        {
+            "func_a": func.with_attr("global_symbol", "func_a"),
+            "func_b": func.with_attr("global_symbol", "func_b"),
+        }
+    )
+
+    @I.ir_module
+    class expected:
+        @T.prim_func(s_tir=True)
+        def func_a():
+            var: T.let = T.int32(10)
+            T.evaluate(var)
+
+        @T.prim_func(s_tir=True)
+        def func_b():
+            var: T.let = T.int32(10)
+            T.evaluate(var)
+
+    after = tvm.tirx.transform.ConvertSSA()(before)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_reused_parameter():
+    """De-duplicate Var usage in parameters
+
+    In this test, the same `tirx.Var` instance is used for the
+    parameter `n` in both functions.
+    """
+
+    @T.prim_func(private=True, s_tir=True)
+    def func(n: T.int32):
+        T.evaluate(n)
+
+    before = tvm.IRModule(
+        {
+            "func_a": func.with_attr("global_symbol", "func_a"),
+            "func_b": func.with_attr("global_symbol", "func_b"),
+        }
+    )
+
+    @I.ir_module
+    class expected:
+        @T.prim_func(s_tir=True)
+        def func_a(n: T.int32):
+            T.evaluate(n)
+
+        @T.prim_func(s_tir=True)
+        def func_b(n: T.int32):
+            T.evaluate(n)
+
+    after = tvm.tirx.transform.ConvertSSA()(before)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_reused_buffer_obj():
+    """De-duplicate buffer usage across entire module"""
+
+    @T.prim_func(private=True, s_tir=True)
+    def func(a: T.handle("float32")):
+        A = T.decl_buffer(shape=1, dtype="float32", data=a)
+        T.evaluate(A[0])
+
+    before = tvm.IRModule(
+        {
+            "func_a": func.with_attr("global_symbol", "func_a"),
+            "func_b": func.with_attr("global_symbol", "func_b"),
+        }
+    )
+
+    @I.ir_module
+    class expected:
+        @T.prim_func(s_tir=True)
+        def func_a(a: T.handle("float32")):
+            A = T.decl_buffer(shape=1, dtype="float32", data=a)
+            T.evaluate(A[0])
+
+        @T.prim_func(s_tir=True)
+        def func_b(a: T.handle("float32")):
+            A = T.decl_buffer(shape=1, dtype="float32", data=a)
+            T.evaluate(A[0])
+
+    after = tvm.tirx.transform.ConvertSSA()(before)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_reused_buffer_parameter():
+    """De-duplicate buffer_map across entire module"""
+
+    @T.prim_func(private=True, s_tir=True)
+    def func(A: T.Buffer(1, "float32")):
+        T.evaluate(A[0])
+
+    before = tvm.IRModule(
+        {
+            "func_a": func.with_attr("global_symbol", "func_a"),
+            "func_b": func.with_attr("global_symbol", "func_b"),
+        }
+    )
+
+    @I.ir_module
+    class expected:
+        @T.prim_func(s_tir=True)
+        def func_a(A: T.Buffer(1, "float32")):
+            T.evaluate(A[0])
+
+        @T.prim_func(s_tir=True)
+        def func_b(A: T.Buffer(1, "float32")):
+            T.evaluate(A[0])
+
+    after = tvm.tirx.transform.ConvertSSA()(before)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_no_change_if_already_ssa():
+    """A module that is already SSA should be unchanged"""
+
+    @I.ir_module
+    class before:
+        @T.prim_func(s_tir=True)
+        def func(A: T.Buffer(1, "float32")):
+            T.evaluate(A[0])
+
+    after = tvm.tirx.transform.ConvertSSA()(before)
+    tvm.ir.assert_structural_equal(before, after)
+    assert before.same_as(after)
+
+
+def test_keep_duplicate_thread_idx_in_same_function():
+    """Environment threads are treated as being at function scope
+
+    The `"thread_extent"` attribute has some unique semantics.  It
+    serves as the definition of the `tirx::Var` representing the
+    environment thread (e.g. `threadIdx.x` in CUDA).  However,
+    multiple `"thread_extent"` attributes may co-exist in the same
+    PrimFunc.  For the purpose of variable scope, use of the
+    `tirx::Var` is only allowed within the body of the `AttrStmt`.
+    However, for the purpose of well-formed-ness, all
+    `"thread_extent"` attributes must use the same IterVar instance
+    (e.g. `WarpIndexFinder` in `lower_warp_memory.cc` may throw an
+    error if multiple IterVar instances occur).
+
+    If there are multiple `AttrStmt` with key `"thread_extent"` in a
+    single function (represented in TVMScript as `T.launch_thread`),
+    these should be treated as a definition of a single variable at
+    function scope, and should not be de-duplicated.
+    """
+
+    @I.ir_module
+    class before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer([256], "float32")):
+            threadIdx_x = T.env_thread("threadIdx.x")
+            with T.launch_thread(threadIdx_x, 256):
+                A[threadIdx_x] = A[threadIdx_x] + 1.0
+
+            with T.launch_thread(threadIdx_x, 256):
+                A[threadIdx_x] = A[threadIdx_x] + 2.0
+
+    after = tvm.tirx.transform.ConvertSSA()(before)
+    tvm.ir.assert_structural_equal(after, before)
+
+
+def test_de_duplicate_thread_idx_across_multiple_functions():
+    """Environment threads are treated as being at function scope
+
+    See `test_keep_duplicate_thread_idx_in_same_function` for background
+    information.
+
+    If there are multiple functions in an IRModule, the `AttrStmt`
+    with key `"thread_extent"` in a single function (represented in
+    TVMScript as `T.launch_thread`), these should be treated as a
+    definition of a single variable at function scope, and should not
+    be de-duplicated.
+
+    For this test case, the `AttrStmt` for `"thread_extent"` are
+    written explicitly, without using the usual `T.env_thread` and
+    `T.launch_thread`, as they cannot represent the duplciate
+    Var/IterVar usage across the two PrimFuncs.
+    """
+
+    threadIdx_x = tvm.tirx.Var("threadIdx_x", "int32")
+
+    # threadIdx_x is defined outside
+    @I.ir_module(check_well_formed=False)
+    class before:
+        @T.prim_func(s_tir=True)
+        def kernel_1(A: T.Buffer([256], "float32")):
+            T.attr(
+                T.iter_var(threadIdx_x, T.Range(0, 256), "ThreadIndex", "threadIdx.x"),
+                "thread_extent",
+                256,
+            )
+            A[threadIdx_x] = A[threadIdx_x] + T.float32(1)
+
+        @T.prim_func(s_tir=True)
+        def kernel_2(A: T.Buffer([256], "float32")):
+            T.attr(
+                T.iter_var(threadIdx_x, T.Range(0, 256), "ThreadIndex", "threadIdx.x"),
+                "thread_extent",
+                256,
+            )
+            A[threadIdx_x] = A[threadIdx_x] + T.float32(1)
+
+    @I.ir_module
+    class expected:
+        @T.prim_func(s_tir=True)
+        def kernel_1(A: T.Buffer([256], "float32")):
+            threadIdx_x = T.int32()
+            T.attr(
+                T.iter_var(threadIdx_x, T.Range(0, 256), "ThreadIndex", "threadIdx.x"),
+                "thread_extent",
+                256,
+            )
+            A[threadIdx_x] = A[threadIdx_x] + T.float32(1)
+
+        @T.prim_func(s_tir=True)
+        def kernel_2(A: T.Buffer([256], "float32")):
+            threadIdx_x = T.int32()
+            T.attr(
+                T.iter_var(threadIdx_x, T.Range(0, 256), "ThreadIndex", "threadIdx.x"),
+                "thread_extent",
+                256,
+            )
+            A[threadIdx_x] = A[threadIdx_x] + T.float32(1)
+
+    after = tvm.tirx.transform.ConvertSSA()(before)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_de_duplicate_thread_idx_iter_var_across_multiple_functions():
+    """Environment threads are treated as being at function scope
+
+    Like `test_de_duplicate_thread_idx_across_multiple_functions`, except the
+    `IterVar` for the environment thread is duplicated across multiple
+    PrimFuncs, not just the `tirx.Var` inside the `IterVar`.
+    """
+
+    threadIdx_x = tvm.tirx.Var("threadIdx_x", "int32")
+    iter_var = tvm.tirx.IterVar(
+        tvm.ir.Range(0, 256), threadIdx_x, tvm.tirx.IterVar.ThreadIndex, "threadIdx.x"
+    )
+
+    # complaints of multiple definitions for threadIdx_x
+    @I.ir_module(check_well_formed=False)
+    class before:
+        @T.prim_func(s_tir=True)
+        def kernel_1(A: T.Buffer([256], "float32")):
+            T.attr(iter_var, "thread_extent", 256)
+            A[threadIdx_x] = A[threadIdx_x] + T.float32(1)
+
+        @T.prim_func(s_tir=True)
+        def kernel_2(A: T.Buffer([256], "float32")):
+            T.attr(iter_var, "thread_extent", 256)
+            A[threadIdx_x] = A[threadIdx_x] + T.float32(1)
+
+    @I.ir_module(check_well_formed=False)
+    class expected:
+        @T.prim_func(s_tir=True)
+        def kernel_1(A: T.Buffer([256], "float32")):
+            threadIdx_x = T.int32()
+            T.attr(
+                T.iter_var(threadIdx_x, T.Range(0, 256), "ThreadIndex", "threadIdx.x"),
+                "thread_extent",
+                256,
+            )
+            A[threadIdx_x] = A[threadIdx_x] + T.float32(1)
+
+        @T.prim_func(s_tir=True)
+        def kernel_2(A: T.Buffer([256], "float32")):
+            threadIdx_x = T.int32()
+            T.attr(
+                T.iter_var(threadIdx_x, T.Range(0, 256), "ThreadIndex", "threadIdx.x"),
+                "thread_extent",
+                256,
+            )
+            A[threadIdx_x] = A[threadIdx_x] + T.float32(1)
+
+    after = tvm.tirx.transform.ConvertSSA()(before)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_thread_idx_reused_within_and_across_functions():
+    """Environment threads are treated as being at function scope
+
+    A combination of
+    test_de_duplicate_thread_idx_iter_var_across_multiple_functions and
+    test_keep_duplicate_thread_idx_in_same_function.  The re-use within a
+    function should be maintained, while re-use across functions is
+    de-duplicated.
+    """
+
+    threadIdx_x = tvm.tirx.Var("threadIdx_x", "int32")
+    iter_var = tvm.tirx.IterVar(
+        tvm.ir.Range(0, 256), threadIdx_x, tvm.tirx.IterVar.ThreadIndex, "threadIdx.x"
+    )
+
+    # complaints of multiple definitions of threadIdx_x
+    @I.ir_module(check_well_formed=False)
+    class before:
+        @T.prim_func(s_tir=True)
+        def kernel_1(A: T.Buffer([256], "float32")):
+            with T.attr(iter_var, "thread_extent", 256):
+                A[threadIdx_x] = A[threadIdx_x] + 1.0
+            with T.attr(iter_var, "thread_extent", 256):
+                A[threadIdx_x] = A[threadIdx_x] + 2.0
+
+        @T.prim_func(s_tir=True)
+        def kernel_2(A: T.Buffer([256], "float32")):
+            with T.attr(iter_var, "thread_extent", 256):
+                A[threadIdx_x] = A[threadIdx_x] + 1.0
+            with T.attr(iter_var, "thread_extent", 256):
+                A[threadIdx_x] = A[threadIdx_x] + 2.0
+
+    @I.ir_module
+    class expected:
+        @T.prim_func(s_tir=True)
+        def kernel_1(A: T.Buffer([256], "float32")):
+            threadIdx_x = T.env_thread("threadIdx.x")
+            with T.launch_thread(threadIdx_x, 256):
+                A[threadIdx_x] = A[threadIdx_x] + 1.0
+            with T.launch_thread(threadIdx_x, 256):
+                A[threadIdx_x] = A[threadIdx_x] + 2.0
+
+        @T.prim_func(s_tir=True)
+        def kernel_2(A: T.Buffer([256], "float32")):
+            threadIdx_x = T.env_thread("threadIdx.x")
+            with T.launch_thread(threadIdx_x, 256):
+                A[threadIdx_x] = A[threadIdx_x] + 1.0
+            with T.launch_thread(threadIdx_x, 256):
+                A[threadIdx_x] = A[threadIdx_x] + 2.0
+
+    after = tvm.tirx.transform.ConvertSSA()(before)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_track_forward_declarations_in_attr_stmt():
+    """T.attr statements may refer to a about-to-be-defined tirx.Var"""
+
+    # Generate the PrimFunc, which is already SSA
+    #
+    # This is constructed directly, rather than using TVMScript.
+    # This test case requires a `tirx.AttrStmt` that references a
+    # variable, followed by the `tirx.For` defining that variable.
+    # This is not expressible in TVMScript, as it only provides the
+    # loop iterator within the body of the loop.
+    i0_outer_outer = tirx.Var("i0_outer_outer", "int32")
+    i0_outer_inner = tirx.Var("i0_outer_inner", "int32")
+    i0_inner = tirx.Var("i0_inner", "int32")
+
+    A = tirx.decl_buffer(1024, "float32", "A")
+    B = tirx.decl_buffer(1024, "float32", "B")
+
+    index = i0_outer_outer * 52 + i0_outer_inner * 4 + i0_inner
+
+    stmt = tirx.BufferStore(B, tirx.BufferLoad(A, [index]), [index])
+    stmt = tirx.IfThenElse(i0_outer_outer * 13 + i0_outer_inner < 256, stmt, None)
+    stmt = tirx.For(i0_inner, 0, 4, tirx.ForKind.VECTORIZED, stmt)
+    stmt = tirx.For(i0_outer_inner, 0, 13, tirx.ForKind.PARALLEL, stmt)
+    stmt = tirx.AttrStmt(
+        T.iter_var(i0_outer_inner, None, "DataPar", ""),
+        "pragma_parallal_barrier_when_finish",
+        1,
+        stmt,
+    )
+    stmt = tirx.AttrStmt(
+        T.iter_var(i0_outer_inner, None, "DataPar", ""),
+        "pragma_parallal_stride_pattern",
+        1,
+        stmt,
+    )
+    stmt = tirx.For(i0_outer_outer, 0, 20, tirx.ForKind.SERIAL, stmt)
+    stmt = tirx.AttrStmt(
+        T.iter_var(i0_outer_outer, None, "DataPar", ""),
+        "pragma_parallal_launch_point",
+        1,
+        stmt,
+    )
+
+    A_handle = tirx.Var("A_handle", "handle")
+    B_handle = tirx.Var("B_handle", "handle")
+
+    before = tirx.PrimFunc(
+        [A_handle, B_handle],
+        stmt,
+        buffer_map={A_handle: A, B_handle: B},
+    )
+
+    mod = tvm.IRModule.from_expr(before)
+    after = tvm.tirx.transform.ConvertSSA()(mod)
+    tvm.ir.assert_structural_equal(after["main"], before)
+
+
+def test_shared_shape_var_in_buffer_map_and_alloc_buffer():
+    """Shape var shared across buffer_map entries and AllocBuffer should not be renamed.
+
+    When the same Var (e.g., `n`) appears in multiple buffer_map
+    entries (A and B both have shape [n]), ConvertSSA should not treat
+    the second occurrence as a redefinition.  All uses of `n` in the
+    function body (including AllocBuffer shapes) must remain the same
+    Var object so that MakePackedAPI can bind it from the DLTensor shape.
+    """
+    n = tirx.Var("n", "int32")
+    A_handle = tirx.Var("A_handle", "handle")
+    B_handle = tirx.Var("B_handle", "handle")
+    A = tirx.decl_buffer((n,), "float32", "A")
+    B = tirx.decl_buffer((n,), "float32", "B")
+
+    # AllocBuffer with shape [n] in the body (flat, no body)
+    C = tirx.decl_buffer((n,), "float32", "C")
+    body = tirx.SeqStmt([tirx.AllocBuffer(C), tirx.Evaluate(1)])
+
+    before = tirx.PrimFunc(
+        [A_handle, B_handle],
+        body,
+        buffer_map={A_handle: A, B_handle: B},
+    )
+
+    mod = tvm.IRModule.from_expr(before)
+    after = tvm.tirx.transform.ConvertSSA()(mod)
+    # The function is already SSA — ConvertSSA should not change it.
+    tvm.ir.assert_structural_equal(after["main"], before)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_flatten_buffer.py

@@ -0,0 +1,440 @@
+# 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
+
+
+def _transform():
+    return tvm.transform.Sequential(
+        [
+            tvm.tirx.transform.FlattenBuffer(),
+            tvm.tirx.transform.StmtSimplify(),
+        ]
+    )
+
+
+def test_elementwise():
+    """2-d buffers are flattened to 1-d"""
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16, 16), "float32"), C: T.Buffer((16, 16), "float32")):
+            for i in T.serial(0, 16):
+                B_new = T.decl_buffer([1, 16], "float32")
+                for j in T.serial(0, 16):
+                    B_new[0, j] = A[i, j] + 1.0
+                for j in T.serial(0, 16):
+                    C[i, j] = B_new[0, j] * 2.0
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16, 16), "float32"), C: T.Buffer((16, 16), "float32")):
+            A_1 = T.decl_buffer(256, dtype="float32", data=A.data)
+            C_1 = T.decl_buffer(256, dtype="float32", data=C.data)
+            for i in T.serial(0, 16):
+                B_new = T.decl_buffer([16], "float32")
+                for j in T.serial(0, 16):
+                    B_new[j] = A_1[((i * 16) + j)] + 1.0
+                for j in T.serial(0, 16):
+                    C_1[((i * 16) + j)] = B_new[j] * 2.0
+
+    After = _transform()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_elementwise_without_decl_buffer():
+    """2-d buffers are flattened to 1-d
+
+    Like test_elementwise, but the TIR doesn't have the DeclBuffer
+    node.  The T.Buffer declaration applies only during the
+    parsing the TVMScript, and doesn't occur in the TIR itself.  In
+    this case, the allocation should be assumed to be targeting flat
+    memory, and should be flattened to a 1-d allocation.
+    """
+
+    @I.ir_module(check_well_formed=False, s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16, 16), "float32"), C: T.Buffer((16, 16), "float32")):
+            for i in T.serial(0, 16):
+                B_new_buf = T.alloc_buffer((1, 16), "float32")
+                B_new = T.Buffer([1, 16], "float32", data=B_new_buf.data)
+                for j in T.serial(0, 16):
+                    B_new[0, j] = A[i, j] + 1.0
+                for j in T.serial(0, 16):
+                    C[i, j] = B_new[0, j] * 2.0
+
+    @I.ir_module(check_well_formed=False, s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(input_A: T.Buffer((16, 16), "float32"), input_C: T.Buffer((16, 16), "float32")):
+            A = T.decl_buffer(256, dtype="float32", data=input_A.data)
+            C = T.decl_buffer(256, dtype="float32", data=input_C.data)
+            for i in T.serial(0, 16):
+                B_new_buf = T.alloc_buffer((16,), "float32")
+                B_new = T.Buffer(16, "float32", data=B_new_buf.data)
+                for j in T.serial(0, 16):
+                    B_new[j] = A[((i * 16) + j)] + 1.0
+                for j in T.serial(0, 16):
+                    C[((i * 16) + j)] = B_new[j] * 2.0
+
+    After = _transform()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_gpu():
+    """Buffer flattening may have indices based on GPU thread vars"""
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16, 16), "float32"), C: T.Buffer((16, 16), "float32")):
+            i0 = T.env_thread("blockIdx.x")
+            i1 = T.env_thread("threadIdx.x")
+            i2 = T.env_thread("vthread")
+
+            T.launch_thread(i0, 4)
+            T.launch_thread(i1, 2)
+            T.launch_thread(i2, 2)
+            B = T.decl_buffer([1, 16], "float32", scope="local")
+            for j in range(0, 16):
+                B[0, j] = A[i0 * 4 + i1 * 2 + i2, j] + 1.0
+            for j in range(0, 16):
+                C[i0 * 4 + i1 * 2 + i2, j] = B[0, j] * 2.0
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16, 16), "float32"), C: T.Buffer((16, 16), "float32")):
+            A_1 = T.decl_buffer(256, dtype="float32", data=A.data)
+            C_1 = T.decl_buffer(256, dtype="float32", data=C.data)
+
+            i0 = T.env_thread("blockIdx.x")
+            i1 = T.env_thread("threadIdx.x")
+            i2 = T.env_thread("vthread")
+
+            T.launch_thread(i0, 4)
+            T.launch_thread(i1, 2)
+            T.launch_thread(i2, 2)
+            B = T.decl_buffer([16], "float32", scope="local")
+            for j in range(0, 16):
+                B[j] = A_1[i0 * 64 + i1 * 32 + i2 * 16 + j] + 1.0
+            for j in range(0, 16):
+                C_1[i0 * 64 + i1 * 32 + i2 * 16 + j] = B[j] * 2.0
+
+    After = _transform()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_symbolic():
+    """Dynamically-sized arrrays are flattened"""
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(a: T.handle, c: T.handle, n: T.int32, m: T.int32) -> None:
+            A = T.match_buffer(a, (n, m), "float32")
+            C = T.match_buffer(c, (n, m), "float32")
+
+            for i in range(0, n):
+                B = T.decl_buffer([m], "float32")
+                for j in range(0, m):
+                    B[j] = A[i, j] + 1.0
+                for j in range(0, m):
+                    C[i, j] = B[j] * 2.0
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(a: T.handle, c: T.handle, n: T.int32, m: T.int32) -> None:
+            A = T.match_buffer(a, (n, m), "float32")
+            C = T.match_buffer(c, (n, m), "float32")
+            A_1 = T.decl_buffer(n * m, "float32", data=A.data)
+            C_1 = T.decl_buffer(n * m, "float32", data=C.data)
+
+            for i in range(0, n):
+                B = T.decl_buffer([m], "float32")
+                for j in range(0, m):
+                    B[j] = A_1[i * m + j] + 1.0
+                for j in range(0, m):
+                    C_1[i * m + j] = B[j] * 2.0
+
+    After = _transform()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_fused_symbolic():
+    """Dynamically-sized arrrays with fused iterator which can be flattened"""
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(a: T.handle, b: T.handle, n: T.int32) -> None:
+            A = T.match_buffer(a, (32, n, n), "float32")
+            B = T.match_buffer(b, (32, n, n), "float32")
+
+            for i in range(0, n * n * 32):
+                B[i // (n * n), (i % (n * n)) // n, i % n] = A[
+                    i // (n * n), (i % (n * n)) // n, i % n
+                ]
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(a: T.handle, b: T.handle, n: T.int32) -> None:
+            input_A = T.match_buffer(a, (32, n, n), "float32")
+            input_B = T.match_buffer(b, (32, n, n), "float32")
+            A = T.decl_buffer(n * n * 32, "float32", data=input_A.data)
+            B = T.decl_buffer(n * n * 32, "float32", data=input_B.data)
+
+            for i in range(0, n * n * 32):
+                B[i] = A[i]
+
+    After = _transform()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_fused_symbolic_with_predicate():
+    """Dynamically-sized arrrays with fused iterator which can be flattened with extra predicate"""
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(a: T.handle, b: T.handle, n: T.int32) -> None:
+            A = T.match_buffer(a, (32, n, n), "float32")
+            B = T.match_buffer(b, (32, n, n), "float32")
+            for bx, tx in T.grid((n * n + 1) // 2, 64):
+                if bx * 64 + tx < n * n * 32:
+                    B[
+                        (bx * 64 + tx) // (n * n),
+                        ((bx * 64 + tx) % (n * n)) // n,
+                        (bx * 64 + tx) % n,
+                    ] = A[
+                        (bx * 64 + tx) // (n * n),
+                        ((bx * 64 + tx) % (n * n)) // n,
+                        (bx * 64 + tx) % n,
+                    ]
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(a: T.handle, b: T.handle, n: T.int32) -> None:
+            input_A = T.match_buffer(a, (32, n, n), "float32")
+            input_B = T.match_buffer(b, (32, n, n), "float32")
+            A = T.decl_buffer(n * n * 32, "float32", data=input_A.data)
+            B = T.decl_buffer(n * n * 32, "float32", data=input_B.data)
+
+            for bx, tx in T.grid((n * n + 1) // 2, 64):
+                if bx * 64 + tx < n * n * 32:
+                    B[bx * 64 + tx] = A[bx * 64 + tx]
+
+    After = _transform()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_multi_alloc():
+    """If multiple allocations occur, all are flattened."""
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((4, 32), "float32"), D: T.Buffer((4, 32), "float32")):
+            for i, j in T.grid(4, 32):
+                B = T.decl_buffer((4, 32), "float32", scope="global")
+                C = T.decl_buffer((4, 32), "float32", scope="global")
+                B[i, j] = A[i, j] + 1.0
+                C[i, j] = A[i, j] + B[i, j]
+                D[i, j] = C[i, j] * 2.0
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((4, 32), "float32"), D: T.Buffer((4, 32), "float32")):
+            A_1 = T.decl_buffer(128, "float32", data=A.data)
+            D_1 = T.decl_buffer(128, "float32", data=D.data)
+
+            for i, j in T.grid(4, 32):
+                B = T.decl_buffer([128], "float32")
+                C = T.decl_buffer([128], "float32")
+                B[i * 32 + j] = A_1[i * 32 + j] + 1.0
+                C[i * 32 + j] = A_1[i * 32 + j] + B[i * 32 + j]
+                D_1[i * 32 + j] = C[i * 32 + j] * 2.0
+
+    After = _transform()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_strided():
+    """Indices for flattened buffers use the specified striding."""
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16, 16), "float32"), C: T.Buffer((16, 16), "float32")):
+            for i0 in T.serial(4):
+                B = T.decl_buffer([4, 17], "float32")
+                B_1 = T.decl_buffer([4, 16], dtype="float32", data=B.data, strides=[17, 1])
+                for i1, j in T.grid(4, 16):
+                    B_1[i1, j] = A[i0 * 4 + i1, j] + 1.0
+                for i1, j in T.grid(4, 16):
+                    C[i0 * 4 + i1, j] = B_1[i1, j] * 2.0
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16, 16), "float32"), C: T.Buffer((16, 16), "float32")):
+            A_1 = T.decl_buffer(256, dtype="float32", data=A.data)
+            C_1 = T.decl_buffer(256, dtype="float32", data=C.data)
+            for i0 in T.serial(0, 4):
+                B = T.decl_buffer([68], "float32")
+                B_1 = T.decl_buffer([68], "float32", data=B.data)
+                for i1 in T.serial(0, 4):
+                    for j in T.serial(0, 16):
+                        B_1[i1 * 17 + j] = A_1[i0 * 64 + i1 * 16 + j] + 1.0
+                for i1 in T.serial(0, 4):
+                    for j in T.serial(0, 16):
+                        C_1[i0 * 64 + i1 * 16 + j] = B_1[i1 * 17 + j] * 2.0
+
+    After = _transform()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_boolean():
+    """Boolean buffers are flattened but kept as bool (no int8 backing array)"""
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(10, "bool"), B: T.Buffer(10, "bool")) -> None:
+            for i0 in T.serial(10):
+                B[i0] = A[i0]
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(input_A: T.Buffer(10, "bool"), input_B: T.Buffer(10, "bool")) -> None:
+            A = T.decl_buffer(10, dtype="bool", data=input_A.data)
+            B = T.decl_buffer(10, dtype="bool", data=input_B.data)
+            # body
+            for i0 in T.serial(10):
+                B[i0] = A[i0]
+
+    After = _transform()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_flatten_inside_block():
+    """Flattening access inside a block flattens the accessed region."""
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            A = T.sblock_alloc_buffer([32, 32])
+            for i, j in T.grid(32, 32):
+                with T.sblock("block"):
+                    T.reads(A[i, j])
+                    T.evaluate(A[i, j])
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            A = T.sblock_alloc_buffer([1024])
+            for i, j in T.grid(32, 32):
+                with T.sblock("block"):
+                    T.reads(A[i * 32 + j])
+                    T.evaluate(A[i * 32 + j])
+
+    After = _transform()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_no_change_to_2d_physical_buffer():
+    """Flattening preserves axis separators."""
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            A = T.sblock_alloc_buffer([32, 32], axis_separators=[1])
+            for i, j in T.grid(32, 32):
+                T.evaluate(A[i, j])
+
+    Expected = Before
+
+    After = _transform()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_flatten_alloc_buffer_with_axis_separators():
+    """Flattening preserves axis separators"""
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            A = T.sblock_alloc_buffer([2, 3, 5, 7, 11, 13], axis_separators=[3])
+            for i0, i1, i2, i3, i4, i5 in T.grid(2, 3, 5, 7, 11, 13):
+                T.evaluate(A[i0, i1, i2, i3, i4, i5])
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            A = T.sblock_alloc_buffer([30, 1001], axis_separators=[1])
+            for i0, i1, i2, i3, i4, i5 in T.grid(2, 3, 5, 7, 11, 13):
+                T.evaluate(A[i0 * 15 + i1 * 5 + i2, i3 * 143 + i4 * 13 + i5])
+
+    After = _transform()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_flatten_decl_buffer_with_axis_separators():
+    """Flattening preserves axis separators
+
+    Like test_flatten_alloc_buffer_with_axis_separators, but the allocations
+    is done using Allocate/DeclBuffer, rather than through
+    BlockNode::alloc_buffers.
+    """
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            A = T.decl_buffer([2, 3, 5, 7, 11, 13], axis_separators=[3])
+            for i0, i1, i2, i3, i4, i5 in T.grid(2, 3, 5, 7, 11, 13):
+                T.evaluate(A[i0, i1, i2, i3, i4, i5])
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            A = T.decl_buffer([30, 1001], axis_separators=[1])
+            for i0, i1, i2, i3, i4, i5 in T.grid(2, 3, 5, 7, 11, 13):
+                T.evaluate(A[i0 * 15 + i1 * 5 + i2, i3 * 143 + i4 * 13 + i5])
+
+    After = _transform()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py

@@ -0,0 +1,311 @@
+# 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 pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+
+
+def test_thread_axis1():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((T.int64(64),), "float32"), B: T.Buffer((T.int64(64),), "float32")):
+        blockIdx_x = T.env_thread("blockIdx.x")
+        T.launch_thread(blockIdx_x, T.int64(2))
+        threadIdx_x = T.env_thread("threadIdx.x")
+        T.launch_thread(threadIdx_x, T.int64(32))
+        B[T.Cast("int64", blockIdx_x) * T.int64(32) + T.Cast("int64", threadIdx_x)] = A[
+            T.Cast("int64", blockIdx_x) * T.int64(32) + T.Cast("int64", threadIdx_x)
+        ] + T.float32(1)
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((64,), "float32"), B: T.Buffer((64,), "float32")):
+        blockIdx_x = T.env_thread("blockIdx.x")
+        T.launch_thread(blockIdx_x, 2)
+        threadIdx_x = T.env_thread("threadIdx.x")
+        T.launch_thread(threadIdx_x, 32)
+        B[blockIdx_x * 32 + threadIdx_x] = A[blockIdx_x * 32 + threadIdx_x] + T.float32(1)
+
+    mod = tvm.IRModule.from_expr(before)
+    func = tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"]
+    tvm.ir.assert_structural_equal(func, expected)
+
+
+def test_thread_axis2():
+    @T.prim_func(s_tir=True)
+    def before(
+        T_reshape: T.Buffer((1, 12, 384, 384), "float32"),
+        placeholder_1: T.Buffer((T.int64(1), T.int64(12), T.int64(384), 384), "bool"),
+        T_where: T.Buffer((T.int64(1), T.int64(12), T.int64(384), 384), "float32"),
+    ) -> None:
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i0_i1_i2_i3_fused_1 in T.thread_binding(T.int64(256), thread="blockIdx.x"):
+            for i0_i1_i2_i3_fused_2 in T.thread_binding(T.int64(1024), thread="threadIdx.x"):
+                for i0_i1_i2_i3_fused_0 in T.serial(T.int64(7)):
+                    with T.sblock("T_where"):
+                        ax0 = T.axis.spatial(T.int64(1), T.int64(0))
+                        ax1 = T.axis.spatial(
+                            T.int64(12),
+                            (
+                                (i0_i1_i2_i3_fused_0 * T.int64(256) + i0_i1_i2_i3_fused_1)
+                                * T.int64(1024)
+                                + i0_i1_i2_i3_fused_2
+                            )
+                            % T.int64(1769472)
+                            // T.int64(147456),
+                        )
+                        ax2 = T.axis.spatial(
+                            T.int64(384),
+                            (
+                                (i0_i1_i2_i3_fused_0 * T.int64(256) + i0_i1_i2_i3_fused_1)
+                                * T.int64(1024)
+                                + i0_i1_i2_i3_fused_2
+                            )
+                            % T.int64(147456)
+                            // T.int64(384),
+                        )
+                        ax3 = T.axis.spatial(
+                            384,
+                            T.cast(
+                                (
+                                    (i0_i1_i2_i3_fused_0 * T.int64(256) + i0_i1_i2_i3_fused_1)
+                                    * T.int64(1024)
+                                    + i0_i1_i2_i3_fused_2
+                                )
+                                % T.int64(384),
+                                "int32",
+                            ),
+                        )
+                        T.where(
+                            (i0_i1_i2_i3_fused_0 * T.int64(256) + i0_i1_i2_i3_fused_1)
+                            * T.int64(1024)
+                            + i0_i1_i2_i3_fused_2
+                            < T.int64(1769472)
+                        )
+                        T.reads(placeholder_1[ax0, ax1, ax2, ax3], T_reshape[ax0, ax1, ax2, ax3])
+                        T.writes(T_where[ax0, ax1, ax2, ax3])
+                        T_where[ax0, ax1, ax2, ax3] = T.Select(
+                            T.cast(placeholder_1[ax0, ax1, ax2, ax3], "int32") != 0,
+                            T.float32(-1000000000),
+                            T_reshape[ax0, ax1, ax2, ax3],
+                        )
+
+    @T.prim_func(s_tir=True)
+    def expected(
+        T_reshape: T.Buffer((1, 12, 384, 384), "float32"),
+        placeholder_1: T.Buffer((1, 12, 384, 384), "bool"),
+        T_where: T.Buffer((1, 12, 384, 384), "float32"),
+    ):
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i0_i1_i2_i3_fused_1 in T.thread_binding(256, thread="blockIdx.x"):
+            for i0_i1_i2_i3_fused_2 in T.thread_binding(1024, thread="threadIdx.x"):
+                for i0_i1_i2_i3_fused_0 in range(7):
+                    with T.sblock("T_where"):
+                        ax0 = T.axis.spatial(1, 0)
+                        ax1 = T.axis.spatial(
+                            12,
+                            (
+                                (i0_i1_i2_i3_fused_0 * 256 + i0_i1_i2_i3_fused_1) * 1024
+                                + i0_i1_i2_i3_fused_2
+                            )
+                            % 1769472
+                            // 147456,
+                        )
+                        ax2 = T.axis.spatial(
+                            384,
+                            (
+                                (i0_i1_i2_i3_fused_0 * 256 + i0_i1_i2_i3_fused_1) * 1024
+                                + i0_i1_i2_i3_fused_2
+                            )
+                            % 147456
+                            // 384,
+                        )
+                        ax3 = T.axis.spatial(
+                            384,
+                            (
+                                (i0_i1_i2_i3_fused_0 * 256 + i0_i1_i2_i3_fused_1) * 1024
+                                + i0_i1_i2_i3_fused_2
+                            )
+                            % 384,
+                        )
+                        T.where(
+                            (i0_i1_i2_i3_fused_0 * 256 + i0_i1_i2_i3_fused_1) * 1024
+                            + i0_i1_i2_i3_fused_2
+                            < 1769472
+                        )
+                        T.reads(placeholder_1[ax0, ax1, ax2, ax3], T_reshape[ax0, ax1, ax2, ax3])
+                        T.writes(T_where[ax0, ax1, ax2, ax3])
+                        T_where[ax0, ax1, ax2, ax3] = T.Select(
+                            T.Cast("int32", placeholder_1[ax0, ax1, ax2, ax3]) != 0,
+                            T.float32(-1000000000),
+                            T_reshape[ax0, ax1, ax2, ax3],
+                        )
+
+    mod = tvm.IRModule.from_expr(before)
+    func = tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"]
+    tvm.ir.assert_structural_equal(func, expected)
+
+
+def test_block():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
+        for i in T.serial(0, T.int64(16)):
+            for j in T.serial(0, T.int64(8)):
+                with T.sblock():
+                    vi = T.axis.spatial(T.int64(128), i * T.int64(8) + j)
+                    B[vi] = A[vi] + T.float32(1)
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
+        for i in T.serial(0, T.int32(16)):
+            for j in T.serial(0, T.int32(8)):
+                with T.sblock():
+                    vi = T.axis.spatial(T.int32(128), i * T.int32(8) + j)
+                    B[vi] = A[vi] + T.float32(1)
+
+    mod = tvm.IRModule.from_expr(before)
+    func = tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"]
+    tvm.ir.assert_structural_equal(func, expected)
+
+
+def test_i16_buffer():
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer((128,), "int16"), B: T.Buffer((128,), "int16")):
+        for i in T.serial(0, T.int64(16)):
+            for j in T.serial(0, T.int64(16)):
+                with T.sblock():
+                    vi = T.axis.spatial(T.int64(128), i * 8 + j)
+                    B[vi] = A[vi] + T.int16(1)
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer((128,), "int16"), B: T.Buffer((128,), "int16")):
+        for i in T.serial(0, 16):
+            for j in T.serial(0, 16):
+                with T.sblock():
+                    vi = T.axis.spatial(128, i * 8 + j)
+                    B[vi] = A[vi] + T.int16(1)
+
+    mod = tvm.IRModule.from_expr(before)
+    after = tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"]
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_fail_on_buffer_map():
+    @T.prim_func(private=True, s_tir=True)
+    def func(A: T.Buffer((128,), "int64"), B: T.Buffer((128,), "int64")):
+        for i in T.serial(0, 16):
+            for j in T.serial(0, 8):
+                with T.sblock():
+                    vi = T.axis.spatial(128, i * 8 + j)
+                    B[vi] = A[vi] + T.int64(1)
+
+    mod = tvm.IRModule.from_expr(func)
+    with pytest.raises(RuntimeError):
+        tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"]
+
+
+def test_fail_on_buffer_map():
+    @T.prim_func(private=True, s_tir=True)
+    def func(A: T.Buffer((128,), "int32"), B: T.Buffer((128,), "int32")):
+        C = T.sblock_alloc_buffer((128,), "int64")
+        for i in T.serial(0, 16):
+            for j in T.serial(0, 8):
+                with T.sblock():
+                    vi = T.axis.spatial(128, i * 8 + j)
+                    C[vi] = T.cast(A[vi], "int64") + T.int64(1)
+        for i in T.serial(0, 16):
+            for j in T.serial(0, 8):
+                with T.sblock():
+                    vi = T.axis.spatial(128, i * 8 + j)
+                    B[vi] = T.cast(C[vi] + T.int64(1), "int32")
+
+    mod = tvm.IRModule.from_expr(func)
+    with pytest.raises(RuntimeError):
+        tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"]
+
+
+def test_pod_params_and_select():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((T.int64(4),), "float32"), B: T.Buffer((T.int64(4),), "float32"), n: T.int64
+        ):
+            for i in T.serial(T.int64(4)):
+                B[i] = T.Select(T.int64(1) <= i, A[i + n], T.Cast("float32", i))
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32"), n: T.int32):
+            for i in range(4):
+                B[i] = T.Select(1 <= i, A[i + n], T.Cast("float32", i))
+
+    after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before)
+    tvm.ir.assert_structural_equal(Expected, after)
+
+
+def test_clz():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((T.int64(4),), "int32")):
+            for i in T.serial(T.int64(4)):
+                B[i] = T.clz(i)
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((4,), "int32")):
+            for i in range(4):
+                B[i] = T.clz(i) - 32 + 64
+
+    after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before)
+    tvm.ir.assert_structural_equal(Expected, after)
+
+
+def test_let_binding():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(buf: T.handle):
+            n = T.int64()
+            Buf = T.match_buffer(buf, [n], "int32")
+            ceil_log2 = T.Cast("int64", T.ceil(T.log2(T.Cast("float32", n))))
+            for i in T.serial(ceil_log2):
+                T.evaluate(0)
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(buf: T.handle):
+            n = T.int32()
+            Buf = T.match_buffer(buf, [n], "int32")
+            # The pass narrows indexing variables (n, the For extent) but leaves
+            # an explicitly-typed `T.Cast("int64", ...)` storage alone; a Cast to
+            # int32 is inserted at the use site (the For iter) instead.
+            ceil_log2 = T.Cast("int64", T.ceil(T.log2(T.Cast("float32", n))))
+            for i in range(T.Cast("int32", ceil_log2)):
+                T.evaluate(0)
+
+    after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before)
+    tvm.ir.assert_structural_equal(Expected, after)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_fp8_legalize.py

@@ -0,0 +1,228 @@
+# 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.script
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.target import Target
+from tvm.tirx.transform.transform import BindTarget
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+def get_before(dtype: str):
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(Aptr: T.handle(dtype), Bptr: T.handle(dtype), Dptr: T.handle(dtype)):
+            T.func_attr({"global_symbol": "main"})
+            A = T.decl_buffer((100,), dtype, data=Aptr)
+            B = T.decl_buffer((100,), dtype, data=Bptr)
+            D = T.decl_buffer((100,), dtype, data=Dptr)
+            C = T.decl_buffer((100,), dtype)
+            for i in T.grid(100):
+                C[i] = A[i] + B[i]
+                D[i] = T.exp(C[i])
+
+    return Before
+
+
+def promote_f8(f8_dtype: str, promote_dtype: str, v):
+    return promote_uint8(f8_dtype, promote_dtype, T.reinterpret("uint8", v))
+
+
+def cast_to_f8(f8_dtype: str, promote_dtype: str, v):
+    return T.reinterpret(f8_dtype, cast_to_uint8(f8_dtype, promote_dtype, v))
+
+
+def get_after_compute_legalize(dtype: str, promote_dtype: str):
+    @tvm.script.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(Aptr: T.handle(dtype), Bptr: T.handle(dtype), Dptr: T.handle(dtype)):
+            T.func_attr({"global_symbol": "main"})
+            A = T.decl_buffer((100,), dtype, data=Aptr)
+            B = T.decl_buffer((100,), dtype, data=Bptr)
+            D = T.decl_buffer((100,), dtype, data=Dptr)
+            C = T.decl_buffer((100,), promote_dtype)
+            for i in T.grid(100):
+                C[i] = promote_f8(dtype, promote_dtype, A[i]) + promote_f8(
+                    dtype, promote_dtype, B[i]
+                )
+                D[i] = cast_to_f8(dtype, promote_dtype, T.exp(C[i]))
+
+    return After
+
+
+def promote_uint8(f8_dtype: str, promote_dtype: str, v):
+    if f8_dtype == "float8_e4m3fn":
+        if promote_dtype == "float16":
+            mantissa = T.bitwise_and(
+                T.shift_left(T.Cast("uint16", v), T.uint16(7)), T.uint16(0x3FF)
+            )
+            exponent = T.shift_left(
+                T.Cast(
+                    "uint16",
+                    T.shift_right(T.shift_left(v, T.uint8(1)), T.uint8(4)) + T.uint8(8),
+                ),
+                T.uint16(10),
+            )
+            sign = T.shift_left(T.Cast("uint16", T.shift_right(v, T.uint8(7))), T.uint16(15))
+            return T.reinterpret("float16", T.bitwise_or(T.bitwise_or(mantissa, exponent), sign))
+        else:  # promote_dtype == "float32"
+            mantissa = T.bitwise_and(
+                T.shift_left(T.Cast("uint32", v), T.uint32(20)), T.uint32(0x7FFFFF)
+            )
+            exponent = T.shift_left(
+                T.Cast(
+                    "uint32",
+                    T.shift_right(T.shift_left(v, T.uint8(1)), T.uint8(4)) + T.uint8(120),
+                ),
+                T.uint32(23),
+            )
+            sign = T.shift_left(T.Cast("uint32", T.shift_right(v, T.uint8(7))), T.uint32(31))
+            return T.reinterpret("float32", T.bitwise_or(T.bitwise_or(mantissa, exponent), sign))
+    else:  # f8_dtype == "float8_e5m2"
+        if promote_dtype == "float16":
+            return T.reinterpret("float16", T.shift_left(T.Cast("uint16", v), T.uint16(8)))
+        else:  # promote_dtype == "float32"
+            mantissa = T.bitwise_and(
+                T.shift_left(T.Cast("uint32", v), T.uint32(21)), T.uint32(0x7FFFFF)
+            )
+            exponent = T.shift_left(
+                T.Cast(
+                    "uint32",
+                    T.shift_right(T.shift_left(v, T.uint8(1)), T.uint8(3)) + T.uint8(112),
+                ),
+                T.uint32(23),
+            )
+            sign = T.shift_left(T.Cast("uint32", T.shift_right(v, T.uint8(7))), T.uint32(31))
+            return T.reinterpret("float32", T.bitwise_or(T.bitwise_or(mantissa, exponent), sign))
+
+
+def cast_to_uint8(f8_dtype: str, promote_dtype: str, v):
+    if f8_dtype == "float8_e4m3fn":
+        if promote_dtype == "float16":
+            uint16_v = T.reinterpret("uint16", v)
+            rounding_bias = T.bitwise_and(
+                T.shift_right(uint16_v, T.uint16(7)),
+                T.uint16(1),
+            ) + T.uint16(0x3F)
+            uint16_v = uint16_v + rounding_bias
+            mantissa = T.bitwise_and(
+                T.Cast("uint8", T.shift_right(uint16_v, T.uint8(7))), T.uint8(0x7)
+            )
+            exponent_before_delta = T.shift_right(T.shift_left(uint16_v, T.uint16(1)), T.uint16(11))
+            round_to_zero = exponent_before_delta < T.uint16(8)
+            exponent = T.shift_left(
+                T.Cast("uint8", exponent_before_delta - T.uint16(8)),
+                T.uint8(3),
+            )
+            sign = T.shift_left(T.Cast("uint8", T.shift_right(uint16_v, T.uint16(15))), T.uint8(7))
+            return T.if_then_else(
+                round_to_zero, T.uint8(0), T.bitwise_or(T.bitwise_or(mantissa, exponent), sign)
+            )
+        else:  # promote_dtype == "float32"
+            uint32_v = T.reinterpret("uint32", v)
+            rounding_bias = T.bitwise_and(
+                T.shift_right(uint32_v, T.uint32(20)), T.uint32(1)
+            ) + T.uint32(0x7FFFF)
+            uint32_v = uint32_v + rounding_bias
+            mantissa = T.bitwise_and(
+                T.Cast("uint8", T.shift_right(uint32_v, T.uint8(20))), T.uint8(0x7)
+            )
+            exponent_before_delta = T.shift_right(T.shift_left(uint32_v, T.uint32(1)), T.uint32(24))
+            round_to_zero = exponent_before_delta < T.uint32(120)
+            exponent = T.shift_left(
+                T.Cast("uint8", exponent_before_delta - T.uint32(120)), T.uint8(3)
+            )
+            sign = T.shift_left(T.Cast("uint8", T.shift_right(uint32_v, T.uint32(31))), T.uint8(7))
+            return T.if_then_else(
+                round_to_zero, T.uint8(0), T.bitwise_or(T.bitwise_or(mantissa, exponent), sign)
+            )
+    else:  # f8_dtype == "float8_e5m2"
+        if promote_dtype == "float16":
+            uint16_v = T.reinterpret("uint16", v)
+            rounding_bias = T.bitwise_and(
+                T.shift_right(uint16_v, T.uint16(8)), T.uint16(1)
+            ) + T.uint16(0x7F)
+            uint16_v = uint16_v + rounding_bias
+            return T.Cast("uint8", T.shift_right(uint16_v, T.uint16(8)))
+        else:  # promote_dtype == "float32"
+            uint32_v = T.reinterpret("uint32", v)
+            rounding_bias = T.bitwise_and(
+                T.shift_right(uint32_v, T.uint32(21)), T.uint32(1)
+            ) + T.uint32(0xFFFFF)
+            uint32_v = uint32_v + rounding_bias
+            mantissa = T.bitwise_and(
+                T.Cast("uint8", T.shift_right(uint32_v, T.uint8(21))), T.uint8(0x3)
+            )
+            exponent_before_delta = T.shift_right(T.shift_left(uint32_v, T.uint32(1)), T.uint32(24))
+            round_to_zero = exponent_before_delta < T.uint32(112)
+            exponent = T.shift_left(
+                T.Cast("uint8", exponent_before_delta - T.uint32(112)), T.uint8(2)
+            )
+            sign = T.shift_left(T.Cast("uint8", T.shift_right(uint32_v, T.uint32(31))), T.uint8(7))
+            return T.if_then_else(
+                round_to_zero, T.uint8(0), T.bitwise_or(T.bitwise_or(mantissa, exponent), sign)
+            )
+
+
+def get_after_storage_legalize(dtype: str, promote_dtype: str):
+    @tvm.script.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(Aptr: T.handle("uint8"), Bptr: T.handle("uint8"), Dptr: T.handle("uint8")):
+            T.func_attr({"global_symbol": "main"})
+            A = T.decl_buffer((100,), "uint8", data=Aptr)
+            B = T.decl_buffer((100,), "uint8", data=Bptr)
+            D = T.decl_buffer((100,), "uint8", data=Dptr)
+            C = T.decl_buffer((100,), promote_dtype)
+            for i in T.grid(100):
+                C[i] = promote_uint8(dtype, promote_dtype, A[i]) + promote_uint8(
+                    dtype, promote_dtype, B[i]
+                )
+                D[i] = cast_to_uint8(dtype, promote_dtype, T.exp(C[i]))
+
+    return After
+
+
+dtype = tvm.testing.parameter("float8_e4m3fn", "float8_e5m2")
+promote_dtype = tvm.testing.parameter("float16", "float32")
+
+
+def test_fp8_compute_legalize(dtype, promote_dtype):
+    target = Target("nvidia/nvidia-a100")
+    before = BindTarget(target)(get_before(dtype))
+    expected = BindTarget(target)(get_after_compute_legalize(dtype, promote_dtype))
+    # run the transform twice to ensure we can afford to deal
+    # with this repeative optimizations
+    after = tvm.tirx.transform.FP8ComputeLegalize(promote_dtype)(before)
+    after = tvm.tirx.transform.FP8ComputeLegalize(promote_dtype)(after)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_fp8_storage_legalize(dtype, promote_dtype):
+    target = Target("nvidia/nvidia-a100")
+    before = BindTarget(target)(get_after_compute_legalize(dtype, promote_dtype))
+    after = tvm.tirx.transform.FP8StorageLegalize()(before)
+    expected = BindTarget(target)(get_after_storage_legalize(dtype, promote_dtype))
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_helpers.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.
+# ruff: noqa: F841
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_annotate_entry_func_single_primfunc():
+    @tvm.script.ir_module
+    class MockModule:
+        @T.prim_func(private=True, s_tir=True)
+        def func1(A: T.Buffer((16,), "float32")):
+            for i in T.serial(16):
+                if i == 5:
+                    if i == 5:
+                        A[i] = 0.0
+
+    mod = MockModule
+    assert mod
+    assert "tirx.is_entry_func" not in (mod["func1"].attrs or {})
+    after = tvm.tirx.transform.AnnotateEntryFunc()(mod)
+    assert (
+        after["func1"].attrs
+        and "tirx.is_entry_func" in after["func1"].attrs
+        and after["func1"].attrs["tirx.is_entry_func"]
+    )
+
+
+# Test module
+@tvm.script.ir_module
+class MockModule:
+    @T.prim_func(private=True, s_tir=True)
+    def func1(A: T.Buffer((16,), "float32")):
+        for i in T.serial(16):
+            if i == 5:
+                if i == 5:
+                    A[i] = 0.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def func2(A: T.Buffer((32,), "float32")):
+        for i in T.serial(32):
+            if i == 15:
+                if i == 15:
+                    A[i] = 0.0
+
+
+@pytest.mark.xfail
+def test_annotate_entry_func_multiple_primfunc():
+    mod = MockModule
+    assert mod
+    assert "target" not in (mod["func1"].attrs or {})
+    assert "target" not in (mod["func2"].attrs or {})
+    # This should fail
+    after = tvm.tirx.transform.AnnotateEntryFunc()(mod)
+
+
+def test_bind_target():
+    mod = MockModule
+    assert mod
+
+    target = tvm.target.Target("cuda")
+    assert "target" not in (mod["func1"].attrs or {})
+    assert "target" not in (mod["func2"].attrs or {})
+    after = tvm.tirx.transform.BindTarget(target)(mod)
+
+    assert "target" in after["func1"].attrs
+    assert after["func1"].attrs["target"] == target
+    assert "target" in after["func2"].attrs
+    assert after["func2"].attrs["target"] == target
+
+
+def test_bind_target_adds_attribute():
+    """BindTarget adds the "target" attribute"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.evaluate(0)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"target": T.target("cuda")})
+            T.evaluate(0)
+
+    After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda"))(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bind_target_with_host_to_exposed_function():
+    """BindTarget adds the host target to externally-exposed functions"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"global_symbol": "main"})
+            T.evaluate(0)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"global_symbol": "main", "target": T.target("cuda", host="llvm")})
+            T.evaluate(0)
+
+    After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda", host="llvm"))(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bind_target_with_host_to_internal_function():
+    """Internal functions have a target annotation, but without the host
+
+    The host portion of the target annotation provides host
+    parameters, and is used to expose a function externally as part of
+    `MakePackedAPI` and `MakeUnpackedAPI`.  For internal functions, no
+    external exposure is required, so the host attribute should not be
+    used.
+    """
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main():
+            T.evaluate(0)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(private=True, s_tir=True)
+        def main():
+            T.func_attr({"target": T.target("cuda")})
+            T.evaluate(0)
+
+    After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda", host="llvm"))(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bind_target_ignores_existing():
+    """BindTarget should not replace existing annotations"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"target": T.target("nvptx")})
+            T.evaluate(0)
+
+    Expected = Before
+
+    After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda"))(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bind_target_updates_host():
+    """BindTarget should update host for existing annotations"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"global_symbol": "func", "target": T.target("nvptx")})
+            T.evaluate(0)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr(
+                {
+                    "global_symbol": "func",
+                    "target": T.target("nvptx", host={"kind": "llvm", "opt-level": 0}),
+                }
+            )
+            T.evaluate(0)
+
+    After = tvm.tirx.transform.BindTarget(
+        tvm.target.Target("cuda", host={"kind": "llvm", "opt-level": 0})
+    )(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bind_target_multiple_functions():
+    """BindTarget may apply to multiple functions in a module"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def func1():
+            T.evaluate(0)
+
+        @T.prim_func(s_tir=True)
+        def func2():
+            T.evaluate(0)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def func1():
+            T.func_attr({"target": T.target("cuda")})
+            T.evaluate(0)
+
+        @T.prim_func(s_tir=True)
+        def func2():
+            T.func_attr({"target": T.target("cuda")})
+            T.evaluate(0)
+
+    After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda"))(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bind_target_with_device_host_call_same_func():
+    """BindTarget should bind the device target to the function if it is called from device"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def add(a: T.int32, b: T.int32) -> T.int32:
+            return a + b
+
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((128, 128), "int32"),
+            B: T.Buffer((128, 128), "int32"),
+            C: T.Buffer((128, 128), "int32"),
+        ):
+            T.func_attr({"global_symbol": "main"})
+            length: T.let[T.int32] = Before.add(64, 64)  # Call from host
+            for bx in T.thread_binding(length, "blockIdx.x"):
+                for tx in T.thread_binding(length, "threadIdx.x"):
+                    C[bx, tx] = Before.add(A[bx, tx], B[bx, tx])  # Call from device
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(private=True, s_tir=True)
+        def add(a: T.int32, b: T.int32) -> T.int32:
+            T.func_attr({"target": T.target("cuda")})
+            return a + b
+
+        @T.prim_func(private=True, s_tir=True)
+        def add_host(a: T.int32, b: T.int32) -> T.int32:
+            T.func_attr({"target": T.target({"kind": "llvm", "opt-level": 0})})
+            return a + b
+
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((128, 128), "int32"),
+            B: T.Buffer((128, 128), "int32"),
+            C: T.Buffer((128, 128), "int32"),
+        ):
+            T.func_attr(
+                {
+                    "global_symbol": "main",
+                    "target": T.target("cuda", host={"kind": "llvm", "opt-level": 0}),
+                }
+            )
+            length: T.let[T.int32] = Expected.add_host(64, 64)  # Call from host
+            for bx in T.thread_binding(length, "blockIdx.x"):
+                for tx in T.thread_binding(length, "threadIdx.x"):
+                    C[bx, tx] = Expected.add(A[bx, tx], B[bx, tx])  # Call from device
+
+    After = tvm.tirx.transform.BindTarget(
+        tvm.target.Target("cuda", host={"kind": "llvm", "opt-level": 0})
+    )(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_filter_primfunc():
+    mod = MockModule
+    assert mod
+    # Annotate each function for testing
+    mod["func1"] = mod["func1"].with_attr("temp", "test1")
+    mod["func2"] = mod["func2"].with_attr("temp", "test2")
+
+    # Test condition that does not filter out anything
+    def checker_filter_out_none(func: tvm.tirx.PrimFunc):
+        return "temp" in func.attrs
+
+    after = tvm.tirx.transform.Filter(checker_filter_out_none)(mod)
+    assert len(after.functions) == 2
+    # Filtered functions should satisfy the given condition.
+    assert checker_filter_out_none(after["func1"])
+    assert checker_filter_out_none(after["func2"])
+
+    # Test condition that selectively filters out primfuncs
+    def checker_filter_out_one(func: tvm.tirx.PrimFunc):
+        return ("temp" in func.attrs) and func.attrs["temp"] == "test1"
+
+    after = tvm.tirx.transform.Filter(checker_filter_out_one)(mod)
+    assert len(after.functions) == 1
+    # Filtered functions should satisfy the given condition.
+    assert checker_filter_out_one(after["func1"])
+
+    # Test condition that filters out everything
+    def checker_filter_out_both(func: tvm.tirx.PrimFunc):
+        return "invalid_attr" in func.attrs
+
+    after = tvm.tirx.transform.Filter(checker_filter_out_both)(mod)
+    assert len(after.functions) == 0
+
+
+def test_filter_removes_global_var_map():
+    """Filtering out a function should be identical to never adding it
+
+    This test is to guard against hidden state in the IRModule that
+    remains after filtering.  Previously, this was observed in the
+    `IRModuleNode::global_var_map_`, which retained entries of
+    filtered-out functions.
+    """
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def func():
+            T.evaluate(0)
+
+    @I.ir_module
+    class Expected:
+        pass
+
+    After = tvm.tirx.transform.Filter(lambda prim_func: False)(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_lower_intrin.py

@@ -0,0 +1,226 @@
+# 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: RUF005
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.testing import env
+
+
+def lower_intrin(params, stmt):
+    """wrapper to call transformation in stmt"""
+    lower_expr = tvm.ir.is_prim_expr(stmt)
+    stmt = tvm.tirx.Evaluate(stmt) if lower_expr else stmt
+    mod = tvm.IRModule.from_expr(
+        tvm.tirx.PrimFunc(params, stmt).with_attr("target", tvm.target.Target("llvm"))
+    )
+    mod = tvm.transform.Sequential(
+        [tvm.tirx.transform.StmtSimplify(), tvm.tirx.transform.LowerIntrin()]
+    )(mod)
+    func = mod["main"]
+    stmt = func.body
+    return stmt.value if lower_expr else stmt.body
+
+
+def check_value(expr, variables, data, fref):
+    """
+    Check that expr evaluates to fref(*row) for each row in data.
+    variables: list of TIR vars [x] or [x, y] bound to the columns of data.
+    data: list of tuples, each tuple has len(variables) elements.
+    """
+    n = len(data)
+    num_vars = len(variables)
+    assert num_vars >= 1 and all(len(row) == num_vars for row in data)
+
+    # Build input and output buffers
+    input_bufs = [
+        tvm.tirx.decl_buffer((n,), dtype=variables[i].ty, name=f"v{i}") for i in range(num_vars)
+    ]
+    out_buf = tvm.tirx.decl_buffer((n,), dtype=expr.ty, name="C")
+
+    # Build loop body: for each i, bind variables[j] = input_bufs[j][i], then store expr to out
+    loop_var = tvm.tirx.Var("i", "int32")
+
+    def make_store(i_var):
+        # Build the expression with each variable bound to the corresponding buffer load
+        result = expr
+        for j in range(num_vars - 1, -1, -1):
+            result = tvm.tirx.Let(variables[j], tvm.tirx.BufferLoad(input_bufs[j], [i_var]), result)
+        return tvm.tirx.BufferStore(out_buf, result, [i_var])
+
+    loop = tvm.tirx.For(
+        loop_var,
+        tvm.tirx.const(0, "int32"),
+        tvm.tirx.const(n, "int32"),
+        tvm.tirx.ForKind.SERIAL,
+        make_store(loop_var),
+    )
+
+    prim_func = tvm.tirx.PrimFunc(input_bufs + [out_buf], loop)
+    prim_func = prim_func.with_attr({"tirx.noalias": True, "global_symbol": "main"})
+    f = tvm.compile(prim_func, "llvm")
+
+    arrays = [
+        tvm.runtime.tensor(np.array([row[j] for row in data], dtype=str(variables[j].ty)))
+        for j in range(num_vars)
+    ]
+    c = tvm.runtime.tensor(np.zeros(n, dtype=str(expr.ty)))
+    f(*arrays, c)
+    cref = np.array([fref(*row) for row in data])
+    np.testing.assert_equal(c.numpy(), cref)
+
+
+def test_lower_nested_access_ptr():
+    data = tvm.tirx.Var("data", tvm.ir.PointerType(tvm.ir.PrimType("float32")))
+    inner = tvm.tirx.tvm_access_ptr("float32", data, 2, 16, 1)
+    outer = tvm.tirx.tvm_access_ptr("float32", inner, 3, 8, 1)
+    body = tvm.tirx.Evaluate(tvm.tirx.call_extern("void", "consume", outer))
+    mod = tvm.IRModule.from_expr(
+        tvm.tirx.PrimFunc([data], body).with_attr("target", tvm.target.Target("llvm"))
+    )
+
+    lowered = tvm.tirx.transform.LowerIntrin()(mod)["main"]
+    access_ptr_calls = []
+    address_calls = []
+
+    def collect(node):
+        if isinstance(node, tvm.ir.Call):
+            if node.op.name == "tirx.tvm_access_ptr":
+                access_ptr_calls.append(node)
+            elif node.op.name == "tirx.address_of":
+                address_calls.append(node)
+
+    tvm.tirx.stmt_functor.post_order_visit(lowered.body, collect)
+    assert not access_ptr_calls
+    assert len(address_calls) == 1
+    load = address_calls[0].args[0]
+    assert isinstance(load, tvm.tirx.BufferLoad)
+    assert int(tvm.arith.Analyzer().simplify(load.indices[0])) == 5
+
+
+def get_ref_data():
+    """Get reference data for every pairs"""
+    import itertools
+
+    x = range(-10, 10)
+    y = list(range(-10, 10))
+    y.remove(0)
+    return list(itertools.product(x, y))
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_lower_floordiv():
+    data = get_ref_data()
+    for dtype in ["int32", "int64", "int16"]:
+        x = tvm.tirx.Var("x", dtype)
+        y = tvm.tirx.Var("y", dtype)
+        zero = tvm.tirx.const(0, dtype)
+        # no constraints
+        res = lower_intrin([x, y], tvm.tirx.floordiv(x, y))
+        check_value(res, [x, y], data, lambda a, b: a // b)
+        # rhs >= 0
+        res = lower_intrin([x, y], tvm.tirx.Select(y >= 0, tvm.tirx.floordiv(x, y), zero))
+        check_value(res, [x, y], data, lambda a, b: a // b if b > 0 else 0)
+        # involves max
+        res = lower_intrin(
+            [x, y], tvm.tirx.Select(y >= 0, tvm.tirx.max(tvm.tirx.floordiv(x, y), zero), zero)
+        )
+        check_value(res, [x, y], data, lambda a, b: max(a // b, 0) if b > 0 else 0)
+        # lhs >= 0
+        res = lower_intrin(
+            [x, y], tvm.tirx.Select(tvm.tirx.all(y >= 0, x >= 0), tvm.tirx.floordiv(x, y), zero)
+        )
+        check_value(res, [x, y], data, lambda a, b: a // b if b > 0 and a >= 0 else 0)
+        # const power of two
+        res = lower_intrin([x, y], tvm.tirx.floordiv(x, tvm.tirx.const(8, dtype=dtype)))
+        check_value(res, [x, y], [(a, b) for a, b in data if b == 8], lambda a, b: a // b)
+        # floordiv(x + m, k), m and k are positive constant. 2 <= m <= k-1.
+        res = lower_intrin(
+            [x, y],
+            tvm.tirx.floordiv(x + tvm.tirx.const(4, dtype=dtype), tvm.tirx.const(5, dtype=dtype)),
+        )
+        check_value(res, [x, y], [(a, b) for a, b in data if b == 5], lambda a, b: (a + 4) // b)
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_lower_floormod():
+    data = get_ref_data()
+    for dtype in ["int32", "int64", "int16"]:
+        x = tvm.tirx.Var("x", dtype)
+        y = tvm.tirx.Var("y", dtype)
+        zero = tvm.tirx.const(0, dtype)
+        # no constraints
+        res = lower_intrin([x, y], tvm.tirx.floormod(x, y))
+        check_value(res, [x, y], data, lambda a, b: a % b)
+        # rhs >= 0
+        res = lower_intrin([x, y], tvm.tirx.Select(y >= 0, tvm.tirx.floormod(x, y), zero))
+        check_value(res, [x, y], data, lambda a, b: a % b if b > 0 else 0)
+        # lhs >= 0
+        res = lower_intrin(
+            [x, y], tvm.tirx.Select(tvm.tirx.all(y >= 0, x >= 0), tvm.tirx.floormod(x, y), zero)
+        )
+        check_value(res, [x, y], data, lambda a, b: a % b if b > 0 and a >= 0 else 0)
+        # const power of two
+        res = lower_intrin([x, y], tvm.tirx.floormod(x, tvm.tirx.const(8, dtype=dtype)))
+        check_value(res, [x, y], [(a, b) for a, b in data if b == 8], lambda a, b: a % b)
+        # floormod(x + m, k), m and k are positive constant. 2 <= m <= k-1.
+        res = lower_intrin(
+            [x, y],
+            tvm.tirx.floormod(x + tvm.tirx.const(4, dtype=dtype), tvm.tirx.const(5, dtype=dtype)),
+        )
+        check_value(res, [x, y], [(a, b) for a, b in data if b == 5], lambda a, b: (a + 4) % b)
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_lower_floordiv_overflow_checks():
+    """
+    Regression tests for overflow checks in TryFindShiftCoefficientForPositiveRange.
+    Divisor is constant 3 (not 1 to avoid CSE, not power-of-two so we don't take the shift path).
+    Reuses lower_intrin and check_value; overflow tests use one var [x].
+    """
+    # Check 3: (b-1) - a_min must not overflow (numerator and C++ int64).
+    # x (int64) full range -> min_value = -2^63. With b = 3: numerator = 2 - (-2^63) > LLONG_MAX.
+    x = tvm.tirx.Var("x", "int64")
+    res = lower_intrin([x], tvm.tirx.floordiv(x, tvm.tirx.const(3, "int64")))
+    data_check3 = [(-(2**63),), (0,), (100,)]
+    check_value(res, [x], data_check3, lambda a: a // 3)
+
+    # Check 4: c_value * b_value must not overflow dtype.
+    # x (int16) full range -> min_value = -32768, c = ceil(32770/3) = 10923; 10923*3 > 32767.
+    x = tvm.tirx.Var("x", "int16")
+    res = lower_intrin([x], tvm.tirx.floordiv(x, tvm.tirx.const(3, "int16")))
+    data_check4 = [(-32768,), (0,), (100,)]
+    check_value(res, [x], data_check4, lambda a: a // 3)
+
+    # Check 5: a_max + b*c must not overflow (offset numerator).
+    # tirx.min(tirx.max(x, -10), 32758) can give bounds [-10, 32758]; b=3, c=4; a_max + 12 > 32767.
+    # In practice this path may not be triggered. This test still validates correct lowering.
+    x = tvm.tirx.Var("x", "int16")
+    clamped = tvm.tirx.min(
+        tvm.tirx.max(x, tvm.tirx.const(-10, "int16")), tvm.tirx.const(32758, "int16")
+    )
+    res = lower_intrin([x], tvm.tirx.floordiv(clamped, tvm.tirx.const(3, "int16")))
+    data_check5 = [(-10,), (0,), (32758,), (32757,)]
+    check_value(res, [x], data_check5, lambda a: (min(max(a, -10), 32758)) // 3)
+
+
+if __name__ == "__main__":
+    test_lower_floordiv()
+    test_lower_floormod()
+    test_lower_floordiv_overflow_checks()

tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py

@@ -0,0 +1,327 @@
+# 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 ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@tvm.register_global_func("tvm.test_matmul")
+def my_matmul(a, b, c):
+    c.copyfrom(np.dot(a.numpy(), b.numpy()))
+
+
+def test_lower_call_packed():
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((64, 64), "float32"),
+            B: T.Buffer((64, 64), "float32"),
+            C: T.Buffer((64, 64), "float32"),
+        ):
+            T.func_attr({"target": tvm.target.Target("llvm")})
+            T.attr("", "device_id", T.int32(0))
+            T.call_packed("tvm.test_matmul", A, B, C)
+
+    @I.ir_module(check_well_formed=False)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((64, 64), "float32"),
+            B: T.Buffer((64, 64), "float32"),
+            C: T.Buffer((64, 64), "float32"),
+        ):
+            T.func_attr({"target": tvm.target.Target("llvm")})
+            stack_ffi_any: T.let[T.handle] = T.tvm_stack_alloca("tvm_ffi_any", 4)
+            stack_array: T.let[T.handle] = T.tvm_stack_alloca("array", 3)
+            stack_shape: T.let[T.handle("int64")] = T.tvm_stack_alloca("shape", 6)
+            stack_shape_1 = T.decl_buffer((T.int64(6),), "int64", data=stack_shape)
+            stack_shape_1[0] = T.int64(64)
+            stack_shape_1[1] = T.int64(64)
+            T.tvm_struct_set(stack_array, 0, 1, A.data)
+            stack_shape_2 = T.Buffer((1,), "int64", data=stack_shape)
+            T.tvm_struct_set(stack_array, 0, 2, T.address_of(stack_shape_2[0]))
+            T.tvm_struct_set(stack_array, 0, 3, T.reinterpret("handle", T.uint64(0)))
+            T.tvm_struct_set(stack_array, 0, 4, 2)
+            T.tvm_struct_set(stack_array, 0, 5, T.uint8(2))
+            T.tvm_struct_set(stack_array, 0, 6, T.uint8(32))
+            T.tvm_struct_set(stack_array, 0, 7, T.uint16(1))
+            T.tvm_struct_set(stack_array, 0, 8, T.uint64(0))
+            T.tvm_struct_set(stack_array, 0, 9, 0)
+            T.tvm_struct_set(stack_array, 0, 10, 1)
+            stack_shape_1[2] = T.int64(64)
+            stack_shape_1[3] = T.int64(64)
+            T.tvm_struct_set(stack_array, 1, 1, B.data)
+            stack_shape_3 = T.Buffer((3,), "int64", data=stack_shape)
+            T.tvm_struct_set(stack_array, 1, 2, T.address_of(stack_shape_3[2]))
+            T.tvm_struct_set(stack_array, 1, 3, T.reinterpret("handle", T.uint64(0)))
+            T.tvm_struct_set(stack_array, 1, 4, 2)
+            T.tvm_struct_set(stack_array, 1, 5, T.uint8(2))
+            T.tvm_struct_set(stack_array, 1, 6, T.uint8(32))
+            T.tvm_struct_set(stack_array, 1, 7, T.uint16(1))
+            T.tvm_struct_set(stack_array, 1, 8, T.uint64(0))
+            T.tvm_struct_set(stack_array, 1, 9, 0)
+            T.tvm_struct_set(stack_array, 1, 10, 1)
+            stack_shape_1[4] = T.int64(64)
+            stack_shape_1[5] = T.int64(64)
+            T.tvm_struct_set(stack_array, 2, 1, C.data)
+            stack_shape_4 = T.Buffer((5,), "int64", data=stack_shape)
+            T.tvm_struct_set(stack_array, 2, 2, T.address_of(stack_shape_4[4]))
+            T.tvm_struct_set(stack_array, 2, 3, T.reinterpret("handle", T.uint64(0)))
+            T.tvm_struct_set(stack_array, 2, 4, 2)
+            T.tvm_struct_set(stack_array, 2, 5, T.uint8(2))
+            T.tvm_struct_set(stack_array, 2, 6, T.uint8(32))
+            T.tvm_struct_set(stack_array, 2, 7, T.uint16(1))
+            T.tvm_struct_set(stack_array, 2, 8, T.uint64(0))
+            T.tvm_struct_set(stack_array, 2, 9, 0)
+            T.tvm_struct_set(stack_array, 2, 10, 1)
+            T.tvm_struct_set(stack_ffi_any, 0, 13, 7)
+            T.tvm_struct_set(stack_ffi_any, 0, 14, 0)
+            T.tvm_struct_set(stack_ffi_any, 0, 15, T.tvm_struct_get(stack_array, 0, 0, "handle"))
+            T.tvm_struct_set(stack_ffi_any, 1, 13, 7)
+            T.tvm_struct_set(stack_ffi_any, 1, 14, 0)
+            T.tvm_struct_set(stack_ffi_any, 1, 15, T.tvm_struct_get(stack_array, 1, 0, "handle"))
+            T.tvm_struct_set(stack_ffi_any, 2, 13, 7)
+            T.tvm_struct_set(stack_ffi_any, 2, 14, 0)
+            T.tvm_struct_set(stack_ffi_any, 2, 15, T.tvm_struct_get(stack_array, 2, 0, "handle"))
+            T.tvm_struct_set(stack_ffi_any, 3, 13, 0)
+            T.tvm_struct_set(stack_ffi_any, 3, 14, 0)
+            T.tvm_struct_set(stack_ffi_any, 3, 15, T.int64(0))
+            T.call_packed_lowered("tvm.test_matmul", stack_ffi_any, 0, 3)
+
+    After = tvm.tirx.transform.LowerTVMBuiltin()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_lower_call_packed_raw_string():
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"target": tvm.target.Target("llvm")})
+            T.call_packed("testing.echo", "payload")
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"target": tvm.target.Target("llvm")})
+            stack_ffi_any: T.let[T.handle] = T.tvm_stack_alloca("tvm_ffi_any", 2)
+            T.tvm_struct_set(stack_ffi_any, 0, 13, 8)
+            T.tvm_struct_set(stack_ffi_any, 0, 14, 0)
+            T.tvm_struct_set(stack_ffi_any, 0, 15, T.reinterpret(T.handle().ty, "payload"))
+            T.tvm_struct_set(stack_ffi_any, 1, 13, 0)
+            T.tvm_struct_set(stack_ffi_any, 1, 14, 0)
+            T.tvm_struct_set(stack_ffi_any, 1, 15, T.int64(0))
+            T.call_packed_lowered("testing.echo", stack_ffi_any, 0, 1)
+
+    After = tvm.tirx.transform.LowerTVMBuiltin()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+    # The typed pointer is required by the LLVM TVMFFIAny lowering.
+    tvm.compile(Before, target="llvm")
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_call_packed_return_non_i32():
+    # This call packed that return non i32 types
+    expected_value = np.array([1.2, 1.4], dtype="float32")
+
+    def packed_echo(value):
+        return tvm.tirx.call_intrin(
+            value.ty, tvm.ir.Op.get("tirx.tvm_call_packed"), "testing.echo", value
+        )
+
+    def build_tir():
+        Ab = tvm.tirx.decl_buffer((2,), "float32")
+
+        # Build statements using direct TIR construction (no ir_builder)
+        # 1. Store packed_echo(const) result into Ab[0]
+        store0 = tvm.tirx.BufferStore(
+            Ab, packed_echo(tvm.tirx.const(expected_value[0], "float32")), [0]
+        )
+
+        # 2. Let binding: Aptr_dup = packed_echo(Ab.data), then store const into Ab[1]
+        Aptr_dup = tvm.tirx.Var("Aptr_dup", Ab.data.ty)
+        store1 = tvm.tirx.BufferStore(Ab, tvm.tirx.const(expected_value[1], "float32"), [1])
+        bind_stmt = tvm.tirx.Bind(Aptr_dup, packed_echo(Ab.data))
+
+        # Combine into sequence
+        stmt = tvm.tirx.SeqStmt([store0, bind_stmt, store1])
+
+        return tvm.IRModule.from_expr(
+            tvm.tirx.PrimFunc([Ab], stmt).with_attr("global_symbol", "packed_test")
+        )
+
+    mod = build_tir()
+    f = tvm.compile(mod, None)
+    a = tvm.runtime.tensor(np.zeros(2, dtype="float32"))
+    f(a)
+    tvm.testing.assert_allclose(a.numpy(), expected_value)
+
+
+def test_lower_overflow_int32():
+    @T.prim_func(check_well_formed=False, s_tir=True)
+    def variance4(rxplaceholder: T.Buffer((T.int64(1), T.int64(32), T.int64(25690112)), "float32")):
+        T.func_attr({"global_symbol": "variance4", "tirx.noalias": True})
+        rxplaceholder_red = T.alloc_buffer((32,), "float32")
+        T_subtract = T.alloc_buffer((822083584,), "float32")
+        rxplaceholder_red_1 = T.Buffer((T.int64(32),), data=rxplaceholder_red.data)
+        rxplaceholder_1 = T.Buffer((T.int64(822083584),), data=rxplaceholder.data)
+        T_subtract_1 = T.Buffer((T.int64(822083584),), data=T_subtract.data)
+        for ax1, ax2 in T.grid(32, 25690112):
+            cse_v1: T.let[T.int32] = ax1 * 25690112 + ax2
+            T_subtract_1[cse_v1] = rxplaceholder_1[cse_v1] - rxplaceholder_red_1[ax1]
+
+    func = variance4
+    tvm.compile(func, target="llvm")  # should not crash
+
+
+def test_lower_device_allocate():
+    """Device allocations are lowered to TVMBackend* calls
+
+    This test validates the current behavior of LowerTVMBuiltin.  This
+    unit test may be improved in the future by addressing:
+
+    - TVMScript always produces "handle" dtype for
+      `T.tvm_throw_last_error`, while LowerTVMBuiltin outputs "int32"
+      dtype.
+    """
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"target": T.target("llvm")})
+            T.attr("dummy", "device_type", 2)  # kDLCuda
+            T.attr("dummy", "device_id", 0)
+            ptr = T.alloc_buffer((16,), "float32")
+            buf = T.decl_buffer(16, "float32", data=ptr.data)
+            buf[0] = 0.0
+
+    After = tvm.tirx.transform.LowerTVMBuiltin()(Before)
+    # Verify the lowered module can be printed (no crash)
+    script_output = After.script()
+    # Should contain TVMBackendAllocWorkspace and TVMBackendFreeWorkspace
+    assert "TVMBackendAllocWorkspace" in script_output
+    assert "TVMBackendFreeWorkspace" in script_output
+    # DeclBuffer should appear as a flat statement
+    assert "T.decl_buffer" in script_output
+
+
+def test_lower_cpu_allocation():
+    """CPU allocations can be handled at codegen time"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"target": T.target("llvm")})
+            T.attr("dummy", "device_type", 1)  # kDLCPU
+            T.attr("dummy", "device_id", 0)
+            ptr = T.alloc_buffer((16,), "float32")
+            buf = T.decl_buffer(16, "float32", data=ptr.data)
+            buf[0] = 0.0
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"target": T.target("llvm")})
+            ptr = T.alloc_buffer((16,), "float32")
+            buf = T.decl_buffer(16, "float32", data=ptr.data)
+            buf[0] = 0.0
+
+    After = tvm.tirx.transform.LowerTVMBuiltin()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_lower_allocate_requires_device_id():
+    """If device id is missing, error."""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"target": T.target("llvm")})
+            T.attr("dummy", "device_type", 2)  # kDLCuda
+            ptr = T.alloc_buffer((16,), "float32")
+            buf = T.decl_buffer(16, "float32", data=ptr.data)
+            buf[0] = 0.0
+
+    with pytest.raises(RuntimeError):
+        tvm.tirx.transform.LowerTVMBuiltin()(Before)
+
+
+def test_lower_allocate_requires_device_type():
+    """If device type is missing, error.
+
+    The device type can be inferred either from the `"device_type"`
+    statement attribute, or from the `"target"` function attribute.
+    Here, we provide neither.  The `"tirx.is_host_func"` attribute is
+    provided as otherwise the function would be skipped altogether by
+    LowerTVMBuiltin.
+    """
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"tirx.is_host_func": True})
+            T.attr("dummy", "device_id", 0)
+            ptr = T.alloc_buffer((1024 * 1024,), "float32")
+            buf = T.decl_buffer(1024 * 1024, "float32", data=ptr.data)
+            buf[0] = 0.0
+
+    with pytest.raises(RuntimeError):
+        tvm.tirx.transform.LowerTVMBuiltin()(Before)
+
+
+def test_lower_cpu_alloc_with_function_attr():
+    """CPU allocations can be handled at codegen time
+
+    Like `test_lower_cpu_allocation`, but the device type is taken from
+    the function attribute.  The `AttrStmt` can override the device
+    type for allocations within its scope, but it defaults to the
+    function's target.
+    """
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"target": T.target("llvm")})
+            ptr = T.alloc_buffer((16,), "float32")
+            buf = T.decl_buffer(16, "float32", data=ptr.data)
+            buf[0] = 0.0
+
+    # Expected is same as before for this transform
+    Expected = Before
+
+    After = tvm.tirx.transform.LowerTVMBuiltin()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_make_packed_api.py

@@ -0,0 +1,453 @@
+# 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.
+"""Tests for tirx.transform.MakePackedAPI TIR transform.
+
+Tests verify the transform output using TVMScript before/after patterns.
+Runtime error tests are in tests/python/codegen/test_codegen_error_handling.py.
+"""
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def _find_compute_scope(func):
+    result = None
+
+    def _visitor(stmt):
+        if isinstance(stmt, tirx.AttrStmt) and stmt.attr_key == "compute_scope":
+            nonlocal result
+            result = stmt
+
+    tirx.stmt_functor.post_order_visit(func.body, _visitor)
+
+    return result
+
+
+@pytest.mark.parametrize("use_global_symbol", [True, False])
+def test_no_op_when_global_symbol_is_absent(use_global_symbol):
+    func_attr = {"target": tvm.target.Target("llvm", host="llvm")}
+
+    @T.prim_func(private=True, s_tir=True)
+    def before():
+        T.func_attr(func_attr)
+        T.evaluate(0)
+
+    if use_global_symbol:
+        before = before.with_attr("global_symbol", "main")
+
+    after = tvm.tirx.transform.MakePackedAPI()(tvm.IRModule.from_expr(before))["main"]
+    if use_global_symbol:
+        assert len(after.params) == 4
+    else:
+        tvm.ir.assert_structural_equal(before, after)
+
+
+def test_target_host_removed():
+    """After MakePackedAPI, host-side target should be the host
+
+    MakePackedAPI is the last transform that requires both the device
+    and the host.  After MakePackedAPI, the target attribute should
+    only contain the host-side target.
+    """
+
+    host = tvm.target.Target("llvm")
+
+    @I.ir_module
+    class before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "float32")):
+            T.func_attr({"global_symbol": "main", "target": T.target("cuda", host=host)})
+            T.evaluate(0)
+
+    after = tvm.tirx.transform.MakePackedAPI()(before)
+    target_attr = after["main"].attrs["target"]
+    assert str(host) == str(target_attr)
+
+
+def test_internal_subroutine_call():
+    """Internal subroutines should not use the PackedFunc API
+
+    A subroutine without the "global_symbol" attribute is an internal
+    subroutine, and is not directly exposed to a user of the generated
+    `runtime.Module`.  Therefore, it doesn't need to follow the
+    PackedFunc API.
+    """
+
+    @I.ir_module
+    class before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "float32")):
+            T.func_attr({"target": T.target("llvm", host="llvm")})
+            before.subroutine(A.data)
+
+        # this test fails if it's made public
+        @T.prim_func(private=True, s_tir=True)
+        def subroutine(A_data: T.handle("float32")):
+            T.func_attr({"target": T.target("llvm")})
+            T.evaluate(A_data)
+
+    after = tvm.tirx.transform.MakePackedAPI()(before)
+    tvm.ir.assert_structural_equal(before["subroutine"], after["subroutine"])
+
+    compute_scope = _find_compute_scope(after["main"])
+    subroutine_call_op = compute_scope.body.value.op
+    assert isinstance(subroutine_call_op, tvm.ir.GlobalVar), (
+        f"The main function's CallNode should use the subroutine's GlobalVar as the operation, "
+        f"but instead has an operation of type {subroutine_call_op}"
+    )
+
+
+def test_subroutine_call_to_externally_visible_subroutine():
+    """Externally-visible subroutines should use the PackedFunc API
+
+    Because the subroutine may be called directly by a user, it must
+    use the PackedFunc API.  Its signature should be updated to the
+    PackedFunc signature, and call sites should be updated to use
+    `T.tvm_call_cpacked`.
+    """
+
+    @I.ir_module
+    class before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "float32")):
+            T.func_attr({"global_symbol": "main", "target": T.target("llvm", host="llvm")})
+            before.subroutine(A.data)
+
+        @T.prim_func(s_tir=True)
+        def subroutine(A_data: T.handle("float32")):
+            T.func_attr({"global_symbol": "subroutine", "target": T.target("llvm", host="llvm")})
+            T.evaluate(A_data)
+
+    after = tvm.tirx.transform.MakePackedAPI()(before)
+
+    main_compute_scope = _find_compute_scope(after["main"])
+    assert main_compute_scope is not None
+    subroutine_compute_scope = _find_compute_scope(after["subroutine"])
+    assert subroutine_compute_scope is not None
+
+    subroutine_call_op = main_compute_scope.body.value.op
+    assert (
+        isinstance(subroutine_call_op, tvm.ir.Op)
+        and subroutine_call_op.name == "tirx.tvm_call_cpacked"
+    ), (
+        f"The main function's CallNode should be lowered to the builtin 'tirx.tvm_call_cpacked', "
+        f"but instead has an operation of type {subroutine_call_op}"
+    )
+
+
+def test_zero_arg_function():
+    """Zero-arg function emits num_args check but no null-pointer check."""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def func_without_arg() -> T.int64:
+            T.func_attr({"target": T.target("llvm", host="llvm")})
+            return T.int64(42)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def func_without_arg(
+            self_handle: T.handle,
+            args: T.handle,
+            num_args: T.int32,
+            result: T.handle("void", "global"),
+        ) -> T.int32:
+            T.func_attr(
+                {
+                    "calling_conv": 1,
+                    "global_symbol": "__tvm_ffi_func_without_arg",
+                    "target": T.target("llvm"),
+                }
+            )
+            assert num_args == 0, (
+                "TypeError",
+                ["Expected ", "0", " arguments", " when calling:\n  `", "func_without_arg()", "`"],
+            )
+            with T.attr(0, "compute_scope", "func_without_arg_compute_"):
+                T.tvm_struct_set(result, 0, 13, 1)
+                T.tvm_struct_set(result, 0, 14, 0)
+                T.tvm_struct_set(result, 0, 15, T.Cast("int64", T.int64(42)))
+                return 0
+            return 0
+
+    After = tvm.tirx.transform.MakePackedAPI()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_pointer_return():
+    """Pointer returns are stored as an opaque pointer in TVMFFIAny."""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(arg: T.handle) -> T.handle:
+            T.func_attr({"target": T.target("llvm", host="llvm")})
+            return arg
+
+    after = tvm.tirx.transform.MakePackedAPI()(Before)["main"]
+    return_type_indices = []
+
+    def collect(node):
+        if not isinstance(node, tvm.ir.Call) or node.op.name != "tirx.tvm_struct_set":
+            return
+        field = node.args[2]
+        if isinstance(field, tvm.tirx.IntImm) and int(field) == 13:
+            return_type_indices.append(int(node.args[3]))
+
+    tvm.tirx.stmt_functor.post_order_visit(after.body, collect)
+    assert 4 in return_type_indices  # ffi::TypeIndex::kTVMFFIOpaquePtr
+
+
+def test_int_parameter():
+    """Int parameter emits type check accepting int or bool."""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(arg: T.int32) -> T.int32:
+            T.func_attr({"target": T.target("llvm", host="llvm")})
+            if arg > 0:
+                return 10
+            else:
+                return 20
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(
+            self_handle: T.handle,
+            args: T.handle,
+            num_args: T.int32,
+            result: T.handle("void", "global"),
+        ) -> T.int32:
+            T.func_attr(
+                {
+                    "calling_conv": 1,
+                    "global_symbol": "__tvm_ffi_main",
+                    "target": T.target("llvm"),
+                }
+            )
+            assert num_args == 1, (
+                "TypeError",
+                ["Expected ", "1", " arguments", " when calling:\n  `", "main(arg: int32)", "`"],
+            )
+            assert not T.isnullptr(args), (
+                "TypeError",
+                ["args pointer is NULL", " when calling:\n  `", "main(arg: int32)", "`"],
+            )
+            arg_type_index: T.let[T.int32] = T.tvm_struct_get(args, 0, 13, "int32")
+            assert arg_type_index == 1 or arg_type_index == 2, (
+                "TypeError",
+                [
+                    "Mismatched type on argument #",
+                    "0",
+                    " when calling:\n  `",
+                    "main(arg: int32)",
+                    "`,\n  expected ",
+                    "int",
+                ],
+            )
+            arg: T.let[T.int32] = T.Cast("int32", T.tvm_struct_get(args, 0, 15, "int64"))
+            with T.attr(0, "compute_scope", "main_compute_"):
+                if arg > 0:
+                    T.tvm_struct_set(result, 0, 13, 1)
+                    T.tvm_struct_set(result, 0, 14, 0)
+                    T.tvm_struct_set(result, 0, 15, T.Cast("int64", 10))
+                    return 0
+                else:
+                    T.tvm_struct_set(result, 0, 13, 1)
+                    T.tvm_struct_set(result, 0, 14, 0)
+                    T.tvm_struct_set(result, 0, 15, T.Cast("int64", 20))
+                    return 0
+            return 0
+
+    After = tvm.tirx.transform.MakePackedAPI()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bool_parameter():
+    """Bool parameter emits type check accepting bool or int."""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(arg: T.bool) -> T.int32:
+            T.func_attr({"target": T.target("llvm", host="llvm")})
+            if arg:
+                return 10
+            else:
+                return 20
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(
+            self_handle: T.handle,
+            args: T.handle,
+            num_args: T.int32,
+            result: T.handle("void", "global"),
+        ) -> T.int32:
+            T.func_attr(
+                {
+                    "calling_conv": 1,
+                    "global_symbol": "__tvm_ffi_main",
+                    "target": T.target("llvm"),
+                }
+            )
+            assert num_args == 1, (
+                "TypeError",
+                ["Expected ", "1", " arguments", " when calling:\n  `", "main(arg: bool)", "`"],
+            )
+            assert not T.isnullptr(args), (
+                "TypeError",
+                ["args pointer is NULL", " when calling:\n  `", "main(arg: bool)", "`"],
+            )
+            arg_type_index: T.let[T.int32] = T.tvm_struct_get(args, 0, 13, "int32")
+            assert arg_type_index == 2 or arg_type_index == 1, (
+                "TypeError",
+                [
+                    "Mismatched type on argument #",
+                    "0",
+                    " when calling:\n  `",
+                    "main(arg: bool)",
+                    "`,\n  expected ",
+                    "boolean",
+                ],
+            )
+            arg: T.let[T.bool] = T.Cast("bool", T.tvm_struct_get(args, 0, 15, "int64"))
+            with T.attr(0, "compute_scope", "main_compute_"):
+                if arg:
+                    T.tvm_struct_set(result, 0, 13, 1)
+                    T.tvm_struct_set(result, 0, 14, 0)
+                    T.tvm_struct_set(result, 0, 15, T.Cast("int64", 10))
+                    return 0
+                else:
+                    T.tvm_struct_set(result, 0, 13, 1)
+                    T.tvm_struct_set(result, 0, 14, 0)
+                    T.tvm_struct_set(result, 0, 15, T.Cast("int64", 20))
+                    return 0
+            return 0
+
+    After = tvm.tirx.transform.MakePackedAPI()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_float_parameter():
+    """Float parameter emits type check accepting float, int, or bool."""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(arg: T.float32) -> T.int32:
+            T.func_attr({"target": T.target("llvm", host="llvm")})
+            if arg > T.float32(0):
+                return 10
+            else:
+                return 20
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(
+            self_handle: T.handle,
+            args: T.handle,
+            num_args: T.int32,
+            result: T.handle("void", "global"),
+        ) -> T.int32:
+            T.func_attr(
+                {
+                    "calling_conv": 1,
+                    "global_symbol": "__tvm_ffi_main",
+                    "target": T.target("llvm"),
+                }
+            )
+            assert num_args == 1, (
+                "TypeError",
+                ["Expected ", "1", " arguments", " when calling:\n  `", "main(arg: float32)", "`"],
+            )
+            assert not T.isnullptr(args), (
+                "TypeError",
+                ["args pointer is NULL", " when calling:\n  `", "main(arg: float32)", "`"],
+            )
+            arg_type_index: T.let[T.int32] = T.tvm_struct_get(args, 0, 13, "int32")
+            assert arg_type_index == 3 or arg_type_index == 1 or arg_type_index == 2, (
+                "TypeError",
+                [
+                    "Mismatched type on argument #",
+                    "0",
+                    " when calling:\n  `",
+                    "main(arg: float32)",
+                    "`,\n  expected ",
+                    "float",
+                ],
+            )
+            arg: T.let[T.float32] = T.Select(
+                arg_type_index == 3,
+                T.Cast("float32", T.tvm_struct_get(args, 0, 15, "float64")),
+                T.Cast("float32", T.tvm_struct_get(args, 0, 15, "int64")),
+            )
+            with T.attr(0, "compute_scope", "main_compute_"):
+                if arg > T.float32(0.0):
+                    T.tvm_struct_set(result, 0, 13, 1)
+                    T.tvm_struct_set(result, 0, 14, 0)
+                    T.tvm_struct_set(result, 0, 15, T.Cast("int64", 10))
+                    return 0
+                else:
+                    T.tvm_struct_set(result, 0, 13, 1)
+                    T.tvm_struct_set(result, 0, 14, 0)
+                    T.tvm_struct_set(result, 0, 15, T.Cast("int64", 20))
+                    return 0
+            return 0
+
+    After = tvm.tirx.transform.MakePackedAPI()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_forward_reference_symbolic_variable():
+    """MakePackedAPI succeeds when a symbolic variable is used before it is defined.
+
+    When buffer A has shape (batch_size+1,) and buffer B has shape (batch_size,),
+    batch_size is referenced (in A's shape check) before it is defined (from B's
+    shape). The three-sequence separation (init_nest, asserts, decl_buffers)
+    ensures all variable definitions precede all assertions.
+    """
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(a: T.handle, b: T.handle):
+            T.func_attr({"target": T.target("llvm", host="llvm")})
+            batch_size = T.int64()
+            A = T.match_buffer(a, (batch_size + 1,), "int32")
+            B = T.match_buffer(b, (batch_size,), "int32")
+            for i in range(batch_size):
+                B[i] = A[i] + A[i + 1]
+
+    # Should not raise "variable batch_size has been used before definition"
+    After = tvm.tirx.transform.MakePackedAPI()(Before)
+    assert len(After["main"].params) == 4
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_narrow_datatype.py

@@ -0,0 +1,322 @@
+# 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 tirx as T
+from tvm.tirx import const
+
+
+def lower_stmt(params, stmt, target_bits):
+    func = tvm.tirx.PrimFunc(params, stmt)
+    func = tvm.tirx.transform.NarrowDataType(target_bits)(tvm.IRModule.from_expr(func))["main"]
+    stmt = func.body
+    return stmt
+
+
+def lower_func_body(func, target_bits):
+    """Lower a TVMScript function and return the first For loop in the body."""
+    mod = tvm.IRModule.from_expr(func)
+    gvar = next(iter(mod.functions.keys()))
+    func = tvm.tirx.transform.NarrowDataType(target_bits)(mod)[gvar]
+    body = func.body
+    # With flat buffer semantics, navigate to the first For node
+    if isinstance(body, tvm.tirx.SeqStmt):
+        for stmt in body:
+            if isinstance(stmt, tvm.tirx.For):
+                return stmt
+    while hasattr(body, "body") and not isinstance(body, tvm.tirx.For):
+        body = body.body
+    return body
+
+
+def test_basic():
+    def check_const(m, n, target_bits, target_dtype):
+        """Check with constant values using TVMScript closure."""
+
+        @T.prim_func(s_tir=True)
+        def func(A: T.Buffer((m * n,), "float32"), B: T.Buffer((m * n,), "float32")):
+            for i in T.serial(m):
+                for j in T.serial(n):
+                    B[i * n + j] = A[i * n + j] + T.float32(1)
+
+        stmt = lower_func_body(func, target_bits)
+        assert stmt.loop_var.ty.dtype == target_dtype
+        assert stmt.body.loop_var.ty.dtype == target_dtype
+
+    def check_symbolic(m_dtype, n_dtype, target_bits, target_dtype):
+        """Check with symbolic shapes as function parameters."""
+        if m_dtype == "int32":
+
+            @T.prim_func(s_tir=True)
+            def func(A: T.handle("float32"), B: T.handle("float32"), m: T.int32, n: T.int32):
+                A_buf = T.decl_buffer((m * n,), "float32", data=A)
+                B_buf = T.decl_buffer((m * n,), "float32", data=B)
+                for i in T.serial(m):
+                    for j in T.serial(n):
+                        B_buf[i * n + j] = A_buf[i * n + j] + T.float32(1)
+
+        else:
+
+            @T.prim_func(s_tir=True)
+            def func(A: T.handle("float32"), B: T.handle("float32"), m: T.int64, n: T.int64):
+                A_buf = T.decl_buffer((m * n,), "float32", data=A)
+                B_buf = T.decl_buffer((m * n,), "float32", data=B)
+                for i in T.serial(m):
+                    for j in T.serial(n):
+                        B_buf[i * n + j] = A_buf[i * n + j] + T.float32(1)
+
+        stmt = lower_func_body(func, target_bits)
+        assert stmt.loop_var.ty.dtype == target_dtype
+        assert stmt.body.loop_var.ty.dtype == target_dtype
+
+    # const shape
+    # i32 -> i32
+    check_const(2, 2, 32, "int32")
+    # i64 -> i32
+    check_const(const(2, dtype="int64"), const(2, dtype="int64"), 32, "int32")
+    check_const(const(2**16, dtype="int64"), const(2**16, dtype="int64"), 32, "int64")
+    # i32 -> i16
+    check_const(2, 2, 16, "int16")
+    check_const(2**10, 2**10, 16, "int32")
+
+    # symbolic shape
+    check_symbolic("int32", "int32", 32, "int32")
+    check_symbolic("int64", "int64", 32, "int64")
+
+
+def test_thread_axis():
+    # This test uses launch_thread to create AttrStmt nodes with "thread_extent"
+    # and checks the dtype of thread axis variables after narrowing.
+    def check_const(m, n, target_bits, target_dtype):
+        @T.prim_func(s_tir=True)
+        def func(A: T.Buffer((m * n,), "float32"), B: T.Buffer((m * n,), "float32")):
+            bx = T.launch_thread("blockIdx.x", m)
+            tx = T.launch_thread("threadIdx.x", n)
+            B[bx * n + tx] = A[bx * n + tx] + T.float32(1)
+
+        mod = tvm.IRModule.from_expr(func)
+        gvar = next(iter(mod.functions.keys()))
+        func_narrowed = tvm.tirx.transform.NarrowDataType(target_bits)(mod)[gvar]
+        stmt = func_narrowed.body
+        assert stmt.node.var.ty.dtype == target_dtype
+        assert stmt.body.node.var.ty.dtype == target_dtype
+
+    # i32 -> i32
+    check_const(2, 32, target_bits=32, target_dtype="int32")
+    # i64 -> i32
+    check_const(
+        const(2, dtype="int64"), const(32, dtype="int64"), target_bits=32, target_dtype="int32"
+    )
+    check_const(
+        const(2**30, dtype="int64"),
+        const(32, dtype="int64"),
+        target_bits=32,
+        target_dtype="int64",
+    )
+    # i32 -> i16
+    check_const(2, 32, target_bits=16, target_dtype="int16")
+    check_const(2**14, 32, target_bits=16, target_dtype="int32")
+
+
+def test_multilanes():
+    # Test narrowing with vector types. Uses closure capture for m and lanes.
+    def check(m, lanes, target_bits, target_dtype):
+        vec_dtype = f"float32x{lanes}"
+
+        @T.prim_func(s_tir=True)
+        def func(
+            A: T.Buffer((m,), vec_dtype),
+            B: T.Buffer((m,), vec_dtype),
+        ):
+            for i in T.serial(m):
+                B[i] = A[i] + T.Broadcast(T.float32(1), lanes)
+            A[0] = B[1]
+
+        mod = tvm.IRModule.from_expr(func)
+        gvar = next(iter(mod.functions.keys()))
+        func_narrowed = tvm.tirx.transform.NarrowDataType(target_bits)(mod)[gvar]
+        stmt = func_narrowed.body
+        assert stmt.seq[0].loop_var.ty.dtype == target_dtype
+
+    # i32 -> i32
+    check(const(2**10, dtype="int32"), 2, target_bits=32, target_dtype="int32")
+    # i64 -> i32
+    check(const(2**10, dtype="int64"), 2, target_bits=32, target_dtype="int32")
+    check(const(2**32, dtype="int64"), 2, target_bits=32, target_dtype="int64")
+    # i32 -> i16
+    check(const(2**10, dtype="int32"), 2, target_bits=16, target_dtype="int16")
+    check(const(2**16, dtype="int32"), 2, target_bits=16, target_dtype="int32")
+
+
+def test_slice():
+    # Test narrowing with slice indexing where buffer B has different index ranges.
+    def check(m, n, target_bits, target_dtype):
+        # The index may overflow in B, while not in A
+        @T.prim_func(s_tir=True)
+        def func(
+            A: T.Buffer((m * n,), "float32"),
+            B: T.Buffer((m * n * 2,), "float32"),
+        ):
+            for i in T.serial(m):
+                for j in T.serial(n):
+                    A[i * n + j] = B[i * 2 * n + 2 * j] + T.float32(1)
+
+        stmt = lower_func_body(func, target_bits)
+        assert stmt.loop_var.ty.dtype == target_dtype
+        assert stmt.body.loop_var.ty.dtype == target_dtype
+
+    # The maximum index is (2**15 * 2**15 - 1) * 2 <= 2**31 - 1
+    check(const(2**15, "int64"), const(2**15, "int64"), target_bits=32, target_dtype="int32")
+    # The maximum index is (2**15 * 2**15 - 1 + 2**15) * 2 > 2**31 - 1
+    check(const(2**15, "int64"), const((2**15 + 1), "int64"), target_bits=32, target_dtype="int64")
+
+
+def test_condition():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((128,), "float32"), B: T.Buffer((130,), "float32")):
+        for i, j in T.grid(T.int64(2), T.int64(65)):
+            if i * T.int64(65) + j >= T.int64(0) and i * T.int64(65) + j < T.int64(128):
+                A[i * T.int64(65) + j] = 0.0
+        for i, j in T.grid(T.int64(2), T.int64(65)):
+            B[i * T.int64(65) + j] = T.if_then_else(
+                i * T.int64(65) + j >= T.int64(0) and i * T.int64(65) + j < T.int64(128),
+                A[i * T.int64(65) + j],
+                0.0,
+                dtype="float32",
+            )
+
+    @T.prim_func(s_tir=True)
+    def expected_after(A: T.Buffer(128, "float32"), B: T.Buffer(130, "float32")):
+        for i, j in T.grid(2, 65):
+            if i * 65 + j >= 0 and i * 65 + j < 128:
+                A[i * 65 + j] = T.float32(0)
+        for i, j in T.grid(2, 65):
+            B[i * 65 + j] = T.if_then_else(
+                i * 65 + j >= 0 and i * 65 + j < 128, A[i * 65 + j], T.float32(0), dtype="float32"
+            )
+
+    after = tvm.tirx.transform.NarrowDataType(32)(
+        tvm.IRModule.from_expr(before.with_attr("global_symbol", "main"))
+    )["main"]
+    tvm.ir.assert_structural_equal(after, expected_after.with_attr("global_symbol", "main"))
+
+
+def test_block():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
+        for i in T.serial(0, T.int64(16)):
+            for j in T.serial(0, T.int64(8)):
+                with T.sblock():
+                    vi = T.axis.spatial(T.int64(128), i * T.int64(8) + j)
+                    B[vi] = A[vi] + T.float32(1)
+
+    @T.prim_func(s_tir=True)
+    def expected_after(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
+        for i in T.serial(0, T.int32(16)):
+            for j in T.serial(0, T.int32(8)):
+                with T.sblock():
+                    vi = T.axis.spatial(T.int32(128), i * T.int32(8) + j)
+                    B[vi] = A[vi] + T.float32(1)
+
+    after = tvm.tirx.transform.NarrowDataType(32)(
+        tvm.IRModule.from_expr(before.with_attr("global_symbol", "main"))
+    )["main"]
+    tvm.ir.assert_structural_equal(after, expected_after.with_attr("global_symbol", "main"))
+
+
+def test_avg_pool2d():
+    @T.prim_func(s_tir=True)
+    def before(PSUM: T.Buffer((313600,), "int32"), PAVG: T.Buffer((313600,), "int32")):
+        for j in T.parallel(T.int64(0), T.int64(280)):
+            for i in T.serial(T.int64(0), T.int64(35)):
+                for vi in T.vectorized(T.int64(0), T.int64(32)):
+                    PAVG[(((j * T.int64(1120)) + (i * T.int64(32))) + vi)] = T.cast(
+                        T.Div(
+                            T.cast(PSUM[(((j * T.int64(1120)) + (i * T.int64(32))) + vi)], "int64"),
+                            T.max(
+                                (
+                                    (
+                                        (
+                                            T.min(
+                                                T.int64(1),
+                                                (T.int64(34) - T.floormod(j, T.int64(35))),
+                                            )
+                                            + T.int64(2)
+                                        )
+                                        - T.max(
+                                            (T.int64(1) - T.floormod(j, T.int64(35))), T.int64(0)
+                                        )
+                                    )
+                                    * (
+                                        (T.min(T.int64(1), (T.int64(34) - i)) + T.int64(2))
+                                        - T.max((T.int64(1) - i), T.int64(0))
+                                    )
+                                ),
+                                T.int64(1),
+                            ),
+                        ),
+                        "int32",
+                    )
+
+    @T.prim_func(s_tir=True)
+    def expected_after(PSUM: T.Buffer((313600,), "int32"), PAVG: T.Buffer((313600,), "int32")):
+        for j in T.parallel(T.int32(0), T.int32(280)):
+            for i in T.serial(T.int32(0), T.int32(35)):
+                for vi in T.vectorized(T.int32(0), T.int32(32)):
+                    PAVG[(((j * T.int32(1120)) + (i * T.int32(32))) + vi)] = T.Div(
+                        PSUM[(((j * T.int32(1120)) + (i * T.int32(32))) + vi)],
+                        (
+                            (
+                                (
+                                    T.min(T.int32(1), (T.int32(34) - T.floormod(j, T.int32(35))))
+                                    + T.int32(2)
+                                )
+                                - T.max((T.int32(1) - T.floormod(j, T.int32(35))), T.int32(0))
+                            )
+                            * (
+                                (T.min(T.int32(1), (T.int32(34) - i)) + T.int32(2))
+                                - T.max((T.int32(1) - i), T.int32(0))
+                            )
+                        ),
+                    )
+
+    after = tvm.tirx.transform.NarrowDataType(32)(
+        tvm.IRModule.from_expr(before.with_attr("global_symbol", "main"))
+    )
+    after = tvm.tirx.transform.StmtSimplify()(after)
+    tvm.ir.assert_structural_equal(after["main"], expected_after.with_attr("global_symbol", "main"))
+
+
+def test_narrow_i64_valued_bufferload_index_to_i32():
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer((16,), "int64")):
+        for i in range(T.int64(15)):
+            A[i + T.int64(1)] = A[i] + T.int64(1)
+
+    @T.prim_func(s_tir=True)
+    def expect(A: T.Buffer((16,), "int64")):
+        for i in range(15):
+            A[i + 1] = A[i] + T.int64(1)
+
+    after = tvm.tirx.transform.NarrowDataType(32)(
+        tvm.IRModule.from_expr(before.with_attr("global_symbol", "main"))
+    )["main"]
+    tvm.ir.assert_structural_equal(after, expect.with_attr("global_symbol", "main"))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py

@@ -0,0 +1,146 @@
+# 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=invalid-name, missing-docstring
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_rewrite_to_shuffle_0():
+    transform = tvm.tirx.transform.PointerValueTypeRewrite()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16,), "float32"), B: T.Buffer((4,), "float32")):
+            A_local = T.alloc_buffer((16,), scope="local")
+            for i in range(4):
+                A_local[i * 4 : i * 4 + 4] = A[i * 4 : i * 4 + 4]
+            for i in range(4):
+                B[i] = A_local[i * 4] + A_local[i * 4 + 1] + A_local[i * 4 + 2] + A_local[i * 4 + 3]
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((4,), "float32x4"), B: T.Buffer((4,), "float32")):
+            A_local = T.alloc_buffer((4,), "float32x4", scope="local")
+            for i in range(4):
+                A_local[T.Div(i * 4, 4)] = A[T.Div(i * 4, 4)]
+            for i in range(4):
+                B[i] = (
+                    T.Shuffle([A_local[T.Div(i * 4, 4)]], [0])
+                    + T.Shuffle([A_local[T.Div(i * 4 + 1, 4)]], [1])
+                    + T.Shuffle([A_local[T.Div(i * 4 + 2, 4)]], [2])
+                    + T.Shuffle([A_local[T.Div(i * 4 + 3, 4)]], [3])
+                )
+
+    After = transform(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_rewrite_to_shuffle_1():
+    transform = tvm.tirx.transform.PointerValueTypeRewrite()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((8,), "float32"), B: T.Buffer((1,), "float32")):
+            A_local = T.alloc_buffer((8,), scope="local")
+            A_local[0:4] = A[0:4]
+            A_local[4:8] = A[4:8]
+            B[0] = (
+                A_local[0]
+                + A_local[1]
+                + A_local[2]
+                + A_local[3]
+                + A_local[4]
+                + A_local[5]
+                + A_local[6]
+                + A_local[7]
+            )
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((2,), "float32x4"), B: T.Buffer((1,), "float32")):
+            A_local = T.alloc_buffer((2,), "float32x4", scope="local")
+            A_local[0] = A[0]
+            A_local[1] = A[1]
+            B[0] = (
+                T.Shuffle([A_local[0]], [0])
+                + T.Shuffle([A_local[0]], [1])
+                + T.Shuffle([A_local[0]], [2])
+                + T.Shuffle([A_local[0]], [3])
+                + T.Shuffle([A_local[1]], [0])
+                + T.Shuffle([A_local[1]], [1])
+                + T.Shuffle([A_local[1]], [2])
+                + T.Shuffle([A_local[1]], [3])
+            )
+
+    After = transform(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_address_of():
+    transform = tvm.tirx.transform.PointerValueTypeRewrite()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")):
+            for i in range(4):
+                T.evaluate(T.address_of(A[i * 4]))
+                B[i * 4 : i * 4 + 4] = A[i * 4 : i * 4 + 4]
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16,), "float32"), B: T.Buffer((4,), "float32x4")):
+            for i in range(4):
+                T.evaluate(T.address_of(A[i * 4]))
+                B[T.Div(i * 4, 4)] = A[i * 4 : i * 4 + 4]
+
+    After = transform(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_scalar_read_without_write():
+    transform = tvm.tirx.transform.PointerValueTypeRewrite()
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16,), "float32")):
+            for i in range(4):
+                T.evaluate(A[i * 4])
+
+    # Expected is the same as Before - no transformation
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16,), "float32")):
+            for i in range(4):
+                T.evaluate(A[i * 4])
+
+    After = transform(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_prim_func_pass.py

@@ -0,0 +1,67 @@
+# 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
+
+
+def test_prim_func_pass():
+    @tvm.tirx.transform.prim_func_pass(opt_level=1)
+    class TestReplaceFunc:
+        """Simple test function to replace one argument to another."""
+
+        def __init__(self, new_func):
+            self.new_func = new_func
+
+        def transform_function(self, func, mod, ctx):
+            return self.new_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)
+
+    new_func = tvm.tirx.PrimFunc([x, y, b], tvm.tirx.Evaluate(0))
+
+    mod = tvm.IRModule({"main": func})
+    mod = TestReplaceFunc(new_func)(mod)
+
+    tvm.ir.assert_structural_equal(mod["main"].body, new_func.body)
+
+
+def test_cow_pass():
+    def fapply(f):
+        assert tvm.testing.object_use_count(f) == 1
+        return f
+
+    pidentity = tvm.tirx.transform.Apply(fapply)
+    x = tvm.tirx.Var("x", "int32")
+    func = tvm.tirx.PrimFunc([x], tvm.tirx.Evaluate(x)).with_attr("target_bits", 32)
+    func_hash = func.__hash__()
+    mod = tvm.IRModule({"main": func})
+    del func
+    # copy on write
+    mod_hash = mod.__hash__()
+    mod = tvm.transform.Sequential([pidentity, tvm.tirx.transform.NarrowDataType(32)])(mod._move())
+    assert mod_hash == mod.__hash__()
+    assert func_hash == mod["main"].__hash__()
+
+
+if __name__ == "__main__":
+    test_cow_pass()
+    test_prim_func_pass()

tests/python/tirx-transform/test_tir_transform_remove_assume.py

@@ -0,0 +1,69 @@
+# 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
+
+
+def test_remove_assume():
+    """Remove any instance of T.assume"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "int32")):
+            T.evaluate(T.assume(A[0] == 5))
+            A[0] = 10
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "int32")):
+            A[0] = 10
+
+    After = tvm.tirx.transform.RemoveAssume()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_remove_assume_loop():
+    """Loops containing only T.assume should be removed"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "int32")):
+            for i in T.serial(16):
+                T.evaluate(T.assume(A[i] == 0))
+
+            for i in T.serial(16):
+                A[i] = 10
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "int32")):
+            for i in T.serial(16):
+                A[i] = 10
+
+    After = tvm.tirx.transform.RemoveAssume()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_remove_no_op.py

@@ -0,0 +1,528 @@
+# 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
+from tvm.script import tirx as T
+
+
+def nop():
+    return tvm.tirx.Evaluate(1)
+
+
+def test_remove_no_op():
+    i = tvm.tirx.Var("i", "int32")
+    j = tvm.tirx.Var("j", "int32")
+    k = tvm.tirx.Var("k", "int32")
+    m = tvm.tirx.Var("m", "int32")
+    n = tvm.tirx.Var("n", "int32")
+    dtype = "int64"
+    Ab = tvm.tirx.decl_buffer((n,), dtype)
+    stmt = tvm.tirx.For(
+        i,
+        0,
+        4,
+        tvm.tirx.ForKind.SERIAL,
+        tvm.tirx.For(
+            j,
+            0,
+            n,
+            tvm.tirx.ForKind.SERIAL,
+            tvm.tirx.For(
+                k,
+                0,
+                m,
+                tvm.tirx.ForKind.SERIAL,
+                tvm.tirx.IfThenElse(
+                    (i * m + j + k < n), tvm.tirx.Evaluate(m), tvm.tirx.Evaluate(n)
+                ),
+            ),
+        ),
+    )
+
+    mod = tvm.IRModule.from_expr(tvm.tirx.PrimFunc([Ab], stmt))
+    ret = tvm.tirx.transform.RemoveNoOp()(mod)["main"].body
+
+    assert isinstance(ret, tvm.tirx.Evaluate)
+    store = tvm.tirx.BufferStore(Ab, tvm.tirx.BufferLoad(Ab, [i]) + 1, [i + 1])
+    stmt2 = tvm.tirx.SeqStmt([nop(), tvm.tirx.SeqStmt([store, nop()])])
+
+    mod = tvm.IRModule.from_expr(tvm.tirx.PrimFunc([Ab], stmt2))
+    ret = tvm.tirx.transform.RemoveNoOp()(mod)["main"].body
+    assert ret == store
+
+    # remove zero extent loop
+    stmt3 = tvm.tirx.For(i, 0, 0, tvm.tirx.ForKind.SERIAL, store)
+    mod = tvm.IRModule.from_expr(tvm.tirx.PrimFunc([Ab], stmt3))
+    ret = tvm.tirx.transform.RemoveNoOp()(mod)["main"].body
+    assert isinstance(ret, tvm.tirx.Evaluate)
+
+
+def test_remove_no_op_with_invalid_extent():
+    @T.prim_func(s_tir=True)
+    def main(A: T.Buffer((16), "int32"), B: T.Buffer((16), "int32")) -> None:
+        for i in T.serial(16):
+            for j in T.serial(i - 20):
+                B[i] = A[i] + j
+
+    mod = tvm.ir.module.IRModule.from_expr(main)
+    ret = tvm.tirx.transform.RemoveNoOp()(mod)["main"].body
+    assert isinstance(ret, tvm.tirx.Evaluate)
+
+
+def _apply_remove_no_op(mod, max_simplification_steps=0):
+    """Helper function to apply RemoveNoOp transform with config."""
+    config = {
+        "tirx.RemoveNoOp": {
+            "max_simplification_steps": max_simplification_steps,
+        }
+    }
+    with tvm.transform.PassContext(config=config):
+        mod = tvm.tirx.transform.RemoveNoOp()(mod)
+    return mod
+
+
+def test_remove_empty_for_loop():
+    """A for-loop whose body is a no-op is itself a no-op."""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before():
+        for i in T.serial(16):
+            T.evaluate(0)
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected():
+        T.evaluate(0)
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_remove_zero_extent_loop():
+    """A for-loop with no extent is a no-op."""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(16, "int32")):
+        for i in T.serial(0):
+            A[i] = 42
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer(16, "int32")):
+        T.evaluate(0)
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_remove_unused_let():
+    """With flat Bind, unused bindings are preserved by remove_no_op.
+
+    Unused Bind elimination requires a separate two-pass approach
+    and is not handled by the current remove_no_op pass.
+    """
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(16, "int32")):
+        x = 5
+        for i in T.serial(16):
+            A[i] = 0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer(16, "int32")):
+        x = 5
+        for i in T.serial(16):
+            A[i] = 0
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_remove_let_used_only_in_no_op():
+    """With flat Bind, unused bindings are preserved by remove_no_op.
+
+    The zero-extent for loop is removed, but the Bind for x remains
+    since unused Bind elimination is not handled by remove_no_op.
+    """
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(16, "int32")):
+        x = 5
+        for i in T.serial(0):
+            A[i] = x
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer(16, "int32")):
+        x = 5
+        T.evaluate(0)
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_keep_side_effects_of_let():
+    """Side-effect Bind is preserved as-is by remove_no_op."""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before():
+        x = T.call_extern("extern_func", dtype="int32")
+        T.evaluate(0)
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected():
+        x = T.call_extern("extern_func", dtype="int32")
+        T.evaluate(0)
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_remove_empty_then_case():
+    """A no-op then_case can be removed."""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(16, "int32")):
+        for i in T.serial(16):
+            if i < 8:
+                T.evaluate(0)
+            else:
+                A[i] = 42
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer(16, "int32")):
+        for i in T.serial(16):
+            if not (i < 8):
+                A[i] = 42
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_remove_empty_else_case():
+    """A no-op else_case can be removed."""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(16, "int32")):
+        for i in T.serial(16):
+            if i < 8:
+                A[i] = 42
+            else:
+                T.evaluate(0)
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer(16, "int32")):
+        for i in T.serial(16):
+            if i < 8:
+                A[i] = 42
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_suppress_removal_of_unused_write():
+    """Sequential writes to the same location are not removed.
+
+    Dataflow analysis is no longer supported.
+    """
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(16, "int32")):
+        for i in T.serial(16):
+            A[i] = 100
+            A[i] = 42
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], before)
+
+
+def test_keep_first_write_when_used():
+    """For two sequential writes, keep the first if it is used"""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(16, "int32")):
+        for i in T.serial(16):
+            A[i] = 100
+            A[i] = A[i] + 1
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], before)
+
+
+def test_keep_partially_overwritten_loop():
+    """Keep partially overwritten regions
+
+    If the second loop doesn't entirely overwrite the first, the first
+    may not be removed be kept.
+    """
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(16, "int32")):
+        for i in T.serial(16):
+            A[i] = 100
+
+        for i in T.serial(16):
+            if i < 12:
+                A[i] = 42
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], before)
+
+
+def test_remove_read_write():
+    """Writing a value to the same location as was just read is a no-op."""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(1, "int32")):
+        A[0] = A[0]
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer(1, "int32")):
+        T.evaluate(0)
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_keep_read_write_to_different_indices():
+    """Writing a value to a different index should not be removed"""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(16, "int32")):
+        for i in T.serial(15):
+            A[i] = A[i + 1]
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], before)
+
+
+def test_remove_read_write_same_index_different_expression():
+    """Writing a value to the same location as the read is a no-op.
+
+    If the value of the index can be proven to be the same, then the
+    store can be removed. With flat Bind, the Bind for i and the
+    enclosing loops remain since unused Bind elimination is not
+    handled by remove_no_op.
+    """
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(16, "int32")):
+        for io, ii in T.grid(4, 4):
+            i: T.let[T.int32] = 4 * io + ii
+            A[4 * io + ii] = A[i]
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer(16, "int32")):
+        for io in range(4):
+            for ii in range(4):
+                i: T.let[T.int32] = 4 * io + ii
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_remove_read_write_same_index_using_constraint():
+    """Writing a value to the same location as the read is a no-op.
+
+    If the value of the index can be proven to be the same, then the
+    no-op can be removed.  This may require using the a constraint
+    that is known from a conditional containing the read/write.
+    """
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(16, "int32")):
+        for i in T.serial(16):
+            if i != 0:
+                A[i] = A[i - 1]
+            else:
+                A[i] = A[0]
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer(16, "int32")):
+        for i in T.serial(16):
+            if i != 0:
+                A[i] = A[i - 1]
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+@pytest.mark.xfail(reason="Dead alloc removal not yet implemented for flat AllocBuffer")
+def test_remove_empty_temporary():
+    """An allocation with a no-op body is a no-op."""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before():
+        A = T.alloc_buffer((16,), "int32", scope="local")
+        T.evaluate(0)
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected():
+        T.evaluate(0)
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+@pytest.mark.xfail(reason="Dead alloc removal not yet implemented for flat AllocBuffer")
+def test_remove_empty_temporary_with_decl_buffer():
+    """Remove DeclBuffer alongside Allocate
+
+    If an unused allocation is removed, any DeclBuffer instances that
+    refer to it should also be removed.
+    """
+
+    @T.prim_func(private=True, s_tir=True)
+    def before():
+        A = T.decl_buffer([4, 4], "int32", scope="local")
+        A_flat = T.decl_buffer(16, "int32", scope="local", data=A.data)
+        T.evaluate(0)
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected():
+        T.evaluate(0)
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+@pytest.mark.xfail(reason="Not implemented yet")
+def test_remove_unused_temporary():
+    """An unused allocation is a no-op."""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(16, "int32")):
+        B = T.alloc_buffer((16,), "int32", scope="local")
+        for i in T.serial(16):
+            A[i] = 1
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer(16, "int32")):
+        for i in T.serial(16):
+            A[i] = 1
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+@pytest.mark.xfail(reason="Not implemented yet")
+def test_remove_unused_write_into_temporary():
+    """A write that only impacts a temporary allocation is a no-op."""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before():
+        A = T.decl_buffer([16], "int32", scope="local")
+        for i in T.serial(16):
+            A[i] = 0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected():
+        T.evaluate(0)
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_keep_used_write_into_temporary():
+    """A write into a temporary that is used later must be kept."""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(B: T.Buffer(16, "int32")):
+        A = T.decl_buffer([16], "int32", scope="local")
+        for i in T.serial(16):
+            A[i] = 0
+
+        for i in T.serial(16):
+            B[i] = A[i]
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], before)
+
+
+@pytest.mark.xfail(reason="Not implemented yet")
+def test_remove_write_into_temporary():
+    """A write that only impacts a temporary allocation is a no-op."""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before(A: T.Buffer(16, "int32"), C: T.Buffer(1, "int32")):
+        B = T.decl_buffer([16], "int32", scope="local")
+        for i in T.serial(16):
+            B[i] = A[i]
+
+        C[0] = 0
+        for i in T.serial(16):
+            C[0] = C[0] + B[i]
+
+        for i in T.serial(16):
+            B[i] = 0
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected(A: T.Buffer(16, "int32"), C: T.Buffer(1, "int32")):
+        B = T.decl_buffer([16], "int32", scope="local")
+        for i in T.serial(16):
+            B[i] = A[i]
+
+        C[0] = 0
+        for i in T.serial(16):
+            C[0] = C[0] + B[i]
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+def test_certain_condition():
+    """The conditon of the If-Else node is certain.
+    This would cause `Segmentation fault` error before."""
+
+    @T.prim_func(private=True, s_tir=True)
+    def before():
+        if True:
+            T.evaluate(0)
+        else:
+            T.evaluate(0)
+
+    @T.prim_func(private=True, s_tir=True)
+    def expected():
+        T.evaluate(0)
+
+    mod = tvm.IRModule.from_expr(before)
+    mod = _apply_remove_no_op(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_split_host_device.py

@@ -0,0 +1,618 @@
+# 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
+
+
+def test_public_api_surface():
+    assert hasattr(tvm.tirx.transform, "SplitHostDevice")
+    assert not hasattr(tvm.tirx.transform, "AnnotateDeviceRegions")
+    assert not hasattr(tvm.tirx.transform, "LowerDeviceKernelLaunch")
+
+
+def test_ssa_across_entire_module():
+    """The host and device functions should not share TIR vars
+
+    Any arguments that are passed from the host to the device should
+    be in terms of independent TIR variables.
+    """
+
+    @I.ir_module
+    class before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"global_symbol": "main", "target": T.target("cuda", host="llvm")})
+            for i in range(16):
+                T.attr(0, "device_scope", 0)
+                for j in range(16):
+                    T.evaluate(i)
+
+    after = tvm.tirx.transform.SplitHostDevice()(before)
+    loop_var = after["main"].body.loop_var
+    param_var = after["main_kernel"].params[0]
+
+    assert not loop_var.same_as(param_var)
+
+
+def test_split_host_device():
+    """SplitHostDevice divides a function at the "target" attribute"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(n: T.int32):
+            T.func_attr({"target": T.target("cuda", host={"kind": "llvm", "opt-level": 0})})
+            T.attr(T.target("cuda"), "target", 0)
+            T.evaluate(n)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(n: T.int32):
+            T.func_attr({"target": T.target("cuda", host={"kind": "llvm", "opt-level": 0})})
+            T.call_packed("main_kernel", n)
+
+        @T.prim_func(s_tir=True)
+        def main_kernel(n: T.int32):
+            T.func_attr(
+                {
+                    "target": T.target("cuda"),
+                    "calling_conv": 2,
+                    "tirx.kernel_launch_params": [],
+                    "global_symbol": "main_kernel",
+                    "tirx.noalias": True,
+                    "tirx.is_global_func": True,
+                }
+            )
+            T.evaluate(n)
+
+    After = tvm.tirx.transform.SplitHostDevice()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_split_host_device_on_cpu():
+    """A kernel running on the CPU may return an error code"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(n: T.int32):
+            T.func_attr({"target": T.target("cuda", host={"kind": "llvm", "opt-level": 0})})
+            T.attr(T.target("llvm"), "target", 0)
+            T.evaluate(n)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(n: T.int32):
+            T.func_attr({"target": T.target("cuda", host={"kind": "llvm", "opt-level": 0})})
+            kernel_error_code: T.let[T.int32] = T.call_extern("int32", "main_kernel", n)
+            assert kernel_error_code == 0, "Error executing compute kernel"
+
+        @T.prim_func(s_tir=True)
+        def main_kernel(n: T.int32) -> T.int32:
+            T.func_attr(
+                {
+                    "target": T.target("llvm"),
+                    "tirx.noalias": True,
+                    "tirx.is_global_func": True,
+                }
+            )
+            T.evaluate(n)
+            T.ret(0)
+
+    After = tvm.tirx.transform.SplitHostDevice()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_split_host_device_without_func_host_attribute():
+    """Like test_split_host_device, but no host specified in the host's target
+
+    The `T.attr` specifying the device still requires splitting out
+    the kernel.
+    """
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(n: T.int32):
+            T.func_attr({"target": T.target("llvm")})
+            T.attr(T.target("cuda"), "target", 0)
+            T.evaluate(n)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(n: T.int32):
+            T.func_attr({"target": T.target("llvm")})
+            T.call_packed("main_kernel", n)
+
+        @T.prim_func(s_tir=True)
+        def main_kernel(n: T.int32):
+            T.func_attr(
+                {
+                    "target": T.target("cuda"),
+                    "calling_conv": 2,
+                    "tirx.kernel_launch_params": [],
+                    "global_symbol": "main_kernel",
+                    "tirx.noalias": True,
+                    "tirx.is_global_func": True,
+                }
+            )
+            T.evaluate(n)
+
+    After = tvm.tirx.transform.SplitHostDevice()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_split_host_device_without_device_region():
+    """Like test_split_host_device, but no device regions to extract
+
+    Because MakePackedAPI/MakeUnpackedAPI still require both the
+    device and host, SplitHostDevice does not modify the "target"
+    attribute.
+    """
+
+    @T.prim_func(s_tir=True)
+    def Before():
+        T.func_attr({"target": T.target("ext_dev", host="llvm")})
+        T.evaluate(0)
+
+    Expected = Before
+
+    After = tvm.tirx.transform.SplitHostDevice()(tvm.IRModule.from_expr(Before))
+    tvm.ir.assert_structural_equal(After["Before"], Expected)
+
+
+def test_split_host_device_name_collision():
+    """Like test_split_host_device, but with the default name already taken
+
+    The default name is generated as `func.name + "_kernel"`.  If this
+    name is already taken by another function in the IRModule, then
+    SplitHostDevice should select a different name.
+    """
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(n: T.int32):
+            T.func_attr({"target": T.target("cuda", host={"kind": "llvm", "opt-level": 0})})
+            T.attr(T.target("cuda"), "target", 0)
+            T.evaluate(n)
+
+        @T.prim_func(s_tir=True)
+        def main_kernel():
+            T.func_attr({"target": T.target("llvm")})
+            T.evaluate(0)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(n: T.int32):
+            T.func_attr({"target": T.target("cuda", host={"kind": "llvm", "opt-level": 0})})
+            T.call_packed("main_kernel_1", n)
+
+        @T.prim_func(s_tir=True)
+        def main_kernel_1(n: T.int32):
+            T.func_attr(
+                {
+                    "target": T.target("cuda"),
+                    "calling_conv": 2,
+                    "tirx.kernel_launch_params": [],
+                    "global_symbol": "main_kernel_1",
+                    "tirx.noalias": True,
+                    "tirx.is_global_func": True,
+                }
+            )
+            T.evaluate(n)
+
+        @T.prim_func(s_tir=True)
+        def main_kernel():
+            T.func_attr({"target": T.target("llvm")})
+            T.evaluate(0)
+
+    After = tvm.tirx.transform.SplitHostDevice()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_dynamic_launch_thread():
+    """Dynamic T.launch_thread may depend on host-side variable
+
+    A dynamic parameter for `T.launch_thread` may have an extent that
+    is computed using variables outside of the `T.target` section.
+
+    This is a regression test to catch a previous failure mode, in
+    which SplitHostDevice generated output with undefined variables,
+    if the only use of a variable occurred in the extent of a
+    `T.launch_thread` statement.
+
+    While the launch-lowering stage will hoist the
+    computation of the extent from the device kernel to the host
+    function, the IRModule must be well-defined at all stages of
+    lowering.  Even if a variable is only used as part of a thread
+    extent, `SplitHostDevice` should treat it as a kernel parameter, to
+    provide a definition of the variable within the TIR device kernel.
+    """
+
+    @I.ir_module
+    class before:
+        @T.prim_func(s_tir=True)
+        def default_function(var_A: T.handle, var_B: T.handle, seq_len: T.int32):
+            T.func_attr({"target": T.target("cuda")})
+
+            A = T.match_buffer(var_A, [seq_len], "int32")
+            B = T.match_buffer(var_B, [seq_len], "int32")
+
+            num_blocks: T.let[T.int32] = (seq_len + 127) // 128
+            with T.attr(T.target("cuda"), "target", 0):
+                blockIdx_x = T.launch_thread("blockIdx.x", num_blocks)
+                threadIdx_x = T.launch_thread("threadIdx.x", 128)
+                if blockIdx_x * 128 + threadIdx_x < seq_len:
+                    B[blockIdx_x * 128 + threadIdx_x] = A[blockIdx_x * 128 + threadIdx_x]
+
+    @I.ir_module
+    class expected:
+        @T.prim_func(s_tir=True)
+        def default_function(var_A: T.handle, var_B: T.handle, seq_len: T.int32):
+            T.func_attr({"target": T.target("cuda")})
+            A = T.match_buffer(var_A, (seq_len,), "int32")
+            B = T.match_buffer(var_B, (seq_len,), "int32")
+            num_blocks: T.let[T.int32] = (seq_len + 127) // 128
+            expected.default_function_kernel(A.data, B.data, num_blocks, seq_len)
+
+        @T.prim_func(private=True, s_tir=True)
+        def default_function_kernel(
+            A_data: T.handle("int32"),
+            B_data: T.handle("int32"),
+            num_blocks: T.int32,
+            seq_len: T.int32,
+        ):
+            T.func_attr(
+                {
+                    "target": T.target("cuda"),
+                    "tirx.is_global_func": True,
+                    "tirx.noalias": True,
+                }
+            )
+            A = T.decl_buffer(seq_len, "int32", data=A_data)
+            B = T.decl_buffer(seq_len, "int32", data=B_data)
+            blockIdx_x = T.launch_thread("blockIdx.x", num_blocks)
+            threadIdx_x = T.launch_thread("threadIdx.x", 128)
+            if blockIdx_x * 128 + threadIdx_x < seq_len:
+                B[blockIdx_x * 128 + threadIdx_x] = A[blockIdx_x * 128 + threadIdx_x]
+
+    after = tvm.tirx.transform.SplitHostDevice()(before)
+
+    tvm.tirx.analysis.verify_well_formed(after)
+    tvm.ir.assert_structural_equal(expected, after)
+
+
+def test_symbolic_var_parameter():
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle):
+            T.func_attr({"target": T.target("cuda")})
+            m = T.int64()
+            A = T.match_buffer(var_A, (m,))
+            B = T.match_buffer(var_B, (m,))
+            T.attr(T.target("cuda"), "target", 0)
+            blockIdx_x = T.launch_thread("blockIdx.x", m)
+            B_1 = T.decl_buffer((m,), data=B.data)
+            A_1 = T.decl_buffer((m,), data=A.data)
+            B_1[blockIdx_x] = A_1[blockIdx_x]
+
+    after = tvm.tirx.transform.SplitHostDevice()(Module)
+    assert len(after["main_kernel"].params) == 3
+    assert isinstance(after["main_kernel"].params[2], tvm.tirx.Var)
+
+
+def test_thread_extent_region_extracted_as_device_kernel():
+    """A bare thread_extent is annotated and extracted as a device kernel."""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32")):
+            T.func_attr({"target": T.target("cuda", host="llvm")})
+            i = T.launch_thread("threadIdx.x", 16)
+            A[i] = 0.0
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32")):
+            T.func_attr({"target": T.target("cuda", host="llvm")})
+            T.call_packed("main_kernel", A.data, 16)
+
+        @T.prim_func(s_tir=True)
+        def main_kernel(A_data: T.handle("float32")):
+            T.func_attr(
+                {
+                    "target": T.target("cuda"),
+                    "calling_conv": 2,
+                    "tirx.kernel_launch_params": ["threadIdx.x"],
+                    "global_symbol": "main_kernel",
+                    "tirx.noalias": True,
+                    "tirx.is_global_func": True,
+                }
+            )
+            A = T.decl_buffer(16, dtype="float32", data=A_data)
+            i = T.launch_thread("threadIdx.x", 16)
+            A[i] = 0.0
+
+    After = tvm.tirx.transform.SplitHostDevice()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_device_scope_region_extracted_as_device_kernel():
+    """A bare device_scope is annotated and extracted as a device kernel."""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "float32")):
+            T.func_attr({"target": T.target("cuda", host="llvm")})
+            T.attr(0, "device_scope", 0)
+            A[0] = 0.0
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "float32")):
+            T.func_attr({"target": T.target("cuda", host="llvm")})
+            T.call_packed("main_kernel", A.data)
+
+        @T.prim_func(s_tir=True)
+        def main_kernel(A_data: T.handle("float32")):
+            T.func_attr(
+                {
+                    "target": T.target("cuda"),
+                    "calling_conv": 2,
+                    "tirx.kernel_launch_params": [],
+                    "global_symbol": "main_kernel",
+                    "tirx.noalias": True,
+                    "tirx.is_global_func": True,
+                }
+            )
+            A = T.decl_buffer(1, dtype="float32", data=A_data)
+            T.attr(0, "device_scope", 0)
+            A[0] = 0.0
+
+    After = tvm.tirx.transform.SplitHostDevice()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_lower_device_kernel_launch():
+    """Kernel calls are lowered using the public SplitHostDevice pass."""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "float32")):
+            T.func_attr({"target": T.target("llvm")})
+            Before.kernel(A.data)
+
+        @T.prim_func(s_tir=True)
+        def kernel(A_data: T.handle("float32")):
+            T.func_attr({"target": T.target("cuda")})
+            A = T.decl_buffer(1, dtype="float32", data=A_data)
+            A[0] = 0.0
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "float32")):
+            T.func_attr({"target": T.target("llvm")})
+            T.call_packed("kernel", A.data)
+
+        @T.prim_func(s_tir=True)
+        def kernel(A_data: T.handle("float32")):
+            T.func_attr(
+                {
+                    "target": T.target("cuda"),
+                    "calling_conv": 2,
+                    "tirx.kernel_launch_params": [],
+                    "global_symbol": "kernel",
+                    "tirx.is_global_func": True,
+                }
+            )
+            A = T.decl_buffer(1, dtype="float32", data=A_data)
+            A[0] = 0.0
+
+    After = tvm.tirx.transform.SplitHostDevice()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_externally_visible_kernel_launch():
+    """Kernel launch lowering preserves a pre-defined global_symbol."""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "float32")):
+            T.func_attr({"target": T.target("llvm")})
+            Before.kernel(A.data)
+
+        @T.prim_func(s_tir=True)
+        def kernel(A_data: T.handle("float32")):
+            T.func_attr({"target": T.target("cuda"), "global_symbol": "kernel_by_another_name"})
+            A = T.decl_buffer(1, dtype="float32", data=A_data)
+            A[0] = 0.0
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "float32")):
+            T.func_attr({"target": T.target("llvm")})
+            T.call_packed("kernel_by_another_name", A.data)
+
+        @T.prim_func(s_tir=True)
+        def kernel(A_data: T.handle("float32")):
+            T.func_attr(
+                {
+                    "target": T.target("cuda"),
+                    "calling_conv": 2,
+                    "tirx.kernel_launch_params": [],
+                    "global_symbol": "kernel_by_another_name",
+                    "tirx.is_global_func": True,
+                }
+            )
+            A = T.decl_buffer(1, dtype="float32", data=A_data)
+            A[0] = 0.0
+
+    After = tvm.tirx.transform.SplitHostDevice()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_collect_launch_parameter():
+    """Thread launch extents are appended to the host launch call."""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32")):
+            T.func_attr({"target": T.target("llvm")})
+            Before.kernel(A.data)
+
+        @T.prim_func(s_tir=True)
+        def kernel(A_data: T.handle("float32")):
+            T.func_attr(
+                {
+                    "target": T.target("cuda"),
+                    "global_symbol": "kernel",
+                }
+            )
+            A = T.decl_buffer(16, dtype="float32", data=A_data)
+            i = T.launch_thread("threadIdx.x", 16)
+            A[i] = 0.0
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32")):
+            T.func_attr({"target": T.target("llvm")})
+            T.call_packed("kernel", A.data, 16)
+
+        @T.prim_func(s_tir=True)
+        def kernel(A_data: T.handle("float32")):
+            T.func_attr(
+                {
+                    "target": T.target("cuda"),
+                    "calling_conv": 2,
+                    "tirx.kernel_launch_params": ["threadIdx.x"],
+                    "global_symbol": "kernel",
+                    "tirx.is_global_func": True,
+                }
+            )
+            A = T.decl_buffer(16, dtype="float32", data=A_data)
+            i = T.launch_thread("threadIdx.x", 16)
+            A[i] = 0.0
+
+    After = tvm.tirx.transform.SplitHostDevice()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_same_device_different_target():
+    """Same-device calls with different codegen are lowered to extern calls."""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "float32")):
+            T.func_attr({"target": T.target("llvm")})
+            Before.kernel(A.data)
+
+        @T.prim_func(s_tir=True)
+        def kernel(A_data: T.handle("float32")):
+            T.func_attr({"target": T.target("c")})
+            A = T.decl_buffer(16, dtype="float32", data=A_data)
+            A[0] = 0.0
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, "float32")):
+            T.func_attr({"target": T.target("llvm")})
+            T.call_extern("kernel", A.data, dtype="void")
+
+        @T.prim_func(s_tir=True)
+        def kernel(A_data: T.handle("float32")):
+            T.func_attr(
+                {
+                    "target": T.target("c"),
+                    "global_symbol": "kernel",
+                    "tirx.is_global_func": True,
+                }
+            )
+            A = T.decl_buffer(16, dtype="float32", data=A_data)
+            A[0] = 0.0
+
+    After = tvm.tirx.transform.SplitHostDevice()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bind_before_thread_extent():
+    """Bind-defined thread extents are inlined into launch arguments."""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32"), n: T.int32):
+            T.func_attr({"target": T.target("llvm")})
+            Before.kernel(A.data, n)
+
+        @T.prim_func(s_tir=True)
+        def kernel(A_data: T.handle("float32"), n: T.int32):
+            T.func_attr({"target": T.target("cuda"), "global_symbol": "kernel"})
+            A = T.decl_buffer(16, dtype="float32", data=A_data)
+            v: T.let[T.int32] = n + 1
+            i = T.launch_thread("threadIdx.x", v)
+            A[i] = 0.0
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32"), n: T.int32):
+            T.func_attr({"target": T.target("llvm")})
+            T.call_packed("kernel", A.data, n, n + 1)
+
+        @T.prim_func(s_tir=True)
+        def kernel(A_data: T.handle("float32"), n: T.int32):
+            T.func_attr(
+                {
+                    "target": T.target("cuda"),
+                    "calling_conv": 2,
+                    "tirx.kernel_launch_params": ["threadIdx.x"],
+                    "global_symbol": "kernel",
+                    "tirx.is_global_func": True,
+                }
+            )
+            A = T.decl_buffer(16, dtype="float32", data=A_data)
+            v: T.let[T.int32] = n + 1
+            i = T.launch_thread("threadIdx.x", v)
+            A[i] = 0.0
+
+    After = tvm.tirx.transform.SplitHostDevice()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_storage_rewrite.py

@@ -0,0 +1,587 @@
+# 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 sys
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_alloc_seq():
+    scope_tb = "local.L0A"
+
+    @T.prim_func(s_tir=True)
+    def func(n: T.int32):
+        for i in T.serial(n):
+            for j in range(10):
+                A = T.alloc_buffer((200,), scope=scope_tb)
+                A[j] = T.float32(1.2)
+            for j in range(10):
+                B = T.alloc_buffer((200,), scope=scope_tb)
+                B[j] = T.float32(1.3)
+
+    mod = tvm.IRModule.from_expr(func)
+    body = tvm.tirx.transform.StorageRewrite()(mod)["func"].body
+
+    num_alloc = [0]
+
+    def verify(n):
+        if isinstance(n, tvm.tirx.AllocBuffer):
+            num_alloc[0] += 1
+            assert n.buffer.shape[0].value == 200
+
+    tvm.tirx.stmt_functor.post_order_visit(body, verify)
+    assert num_alloc[0] == 1
+
+
+def test_alloc_different_dtypes():
+    # Test cross-loop buffer access with buffers allocated in parent scope
+    def make_mod(dtype_list, length):
+        assert len(dtype_list) == 4
+
+        @T.prim_func(s_tir=True)
+        def func():
+            # Allocate all buffers in parent scope (before any loops)
+            A = T.alloc_buffer((length,), dtype_list[0], scope="local.L0A")
+            B = T.alloc_buffer((length,), dtype_list[1], scope="local.L0A")
+            C = T.alloc_buffer((length,), dtype_list[2], scope="local.L0A")
+            D = T.alloc_buffer((length,), dtype_list[3], scope="local.L0A")
+            E = T.alloc_buffer((length,), "int8", scope="local.L0A")
+
+            for j in range(length):
+                A[j] = T.Cast(dtype_list[0], 1)
+            for j in range(length):
+                B[j] = T.Cast(dtype_list[1], 1)
+            for j in range(length):
+                C[j] = T.Cast(dtype_list[2], 1)
+            for j in range(length):
+                D[j] = T.Cast(dtype_list[3], 1)
+            for j in range(length):
+                E[j] = (
+                    T.Cast("int8", A[j])
+                    + T.Cast("int8", B[j])
+                    + T.Cast("int8", C[j])
+                    + T.Cast("int8", D[j])
+                )
+
+        return tvm.IRModule.from_expr(func)
+
+    def dtype_bit_len(dtype):
+        index = 0
+        for i in dtype:
+            if i.isdigit():
+                break
+            index += 1
+        return int(dtype[index:])
+
+    def offset_generater(dtype_list, length):
+        dtype_len_list = [dtype_bit_len(i) for i in dtype_list]
+        base_len = dtype_len_list[0]
+        return sum([i * length / base_len for i in dtype_len_list])
+
+    def dtype_test(dtype_list, length):
+        def verify(n):
+            if isinstance(n, tvm.tirx.AllocBuffer):
+                assert n.buffer.shape[0].value == offset
+
+        mod = make_mod(dtype_list, length)
+        offset = offset_generater(dtype_list, length)
+
+        body = tvm.tirx.transform.StorageRewrite()(mod)["func"].body
+        tvm.tirx.stmt_functor.post_order_visit(body, verify)
+
+    length = 1024
+    dtype_list = ["float16", "int32", "uint16", "int8"]
+    dtype_test(dtype_list, length)
+
+    dtype_list = ["float32", "int32", "uint16", "int8"]
+    dtype_test(dtype_list, length)
+
+    dtype_list = ["float64", "int32", "uint16", "int8"]
+    dtype_test(dtype_list, length)
+
+    dtype_list = ["int8", "int32", "uint16", "uint8"]
+    dtype_test(dtype_list, length)
+
+
+def test_address_of():
+    # In this test, the storage rewrite pass is allowed to
+    # combine buffers B and D, but not C
+    @T.prim_func(s_tir=True)
+    def before(A: T.Buffer(8, "float32"), E: T.Buffer(8, "float32")):
+        B = T.alloc_buffer((8,))
+        for i in range(8):
+            B[i] = (
+                T.call_extern("deref", T.address_of(A[i]), dtype="float32")
+                + T.call_extern("deref", T.address_of(A[0]), dtype="float32")
+                + T.float32(1)
+            )
+        C = T.alloc_buffer((8,))
+        for i in range(8):
+            C[i] = (
+                T.call_extern("deref", T.address_of(B[i]), dtype="float32")
+                + T.call_extern("deref", T.address_of(B[0]), dtype="float32")
+                + T.float32(2)
+            )
+        D = T.alloc_buffer((8,))
+        for i in range(8):
+            D[i] = (
+                T.call_extern("deref", T.address_of(C[i]), dtype="float32")
+                + T.call_extern("deref", T.address_of(C[0]), dtype="float32")
+                + T.float32(2)
+            )
+        for i in range(8):
+            E[i] = (
+                T.call_extern("deref", T.address_of(D[i]), dtype="float32")
+                + T.call_extern("deref", T.address_of(D[0]), dtype="float32")
+                + T.float32(3)
+            )
+
+    def verify(n):
+        if isinstance(n, tvm.tirx.AllocBuffer):
+            total_alloc[0] += n.buffer.shape[0].value
+
+    total_alloc = [0]
+    mod = tvm.IRModule.from_expr(before.with_attr("global_symbol", "main"))
+    tvm.tirx.stmt_functor.post_order_visit(mod["main"].body, verify)
+    assert total_alloc[0] == 24
+
+    total_alloc[0] = 0
+    mod = tvm.tirx.transform.StorageRewrite()(mod)
+    tvm.tirx.stmt_functor.post_order_visit(mod["main"].body, verify)
+    assert total_alloc[0] == 16
+
+
+def test_parallel_alloc():
+    @T.prim_func(s_tir=True)
+    def func1(n: T.int32):
+        for i in T.parallel(n):
+            for j in range(10):
+                A = T.alloc_buffer((n,))
+                A[j] = A[j] + T.float32(2)
+
+    mod = tvm.IRModule.from_expr(func1)
+    body = tvm.tirx.transform.StorageRewrite()(mod)["func1"]
+
+    # With flat AllocBuffer, the for body is a SeqStmt; first element is AllocBuffer
+    assert isinstance(body.body.body[0], tvm.tirx.AllocBuffer)
+
+    @T.prim_func(s_tir=True)
+    def func2(n: T.int32):
+        for t in T.serial(n):
+            with T.attr(T.int32(1), "pragma_scope", "parallel_launch_point"):
+                for i in T.parallel(n):
+                    for j in range(10):
+                        A = T.alloc_buffer((n,))
+                        A[j] = A[j] + T.float32(2)
+
+    mod = tvm.IRModule.from_expr(func2)
+    body = tvm.tirx.transform.StorageRewrite()(mod)["func2"]
+
+    assert isinstance(body.body.body.body.body[0], tvm.tirx.AllocBuffer)
+
+
+def test_while_alloc():
+    @T.prim_func(s_tir=True)
+    def func_parallel(n: T.int32):
+        for i in T.parallel(n):
+            j = T.alloc_buffer((1,), "int32")
+            j[0] = 0
+            while j[0] < 10:
+                A = T.alloc_buffer((n,))
+                A[j[0]] = A[j[0]] + T.float32(2)
+                j[0] = j[0] + j[0] + 1
+
+    @T.prim_func(s_tir=True)
+    def func_serial(n: T.int32):
+        for i in T.serial(n):
+            j = T.alloc_buffer((1,), "int32")
+            j[0] = 0
+            while j[0] < 10:
+                A = T.alloc_buffer((n,))
+                A[j[0]] = A[j[0]] + T.float32(2)
+                j[0] = j[0] + j[0] + 1
+
+    mod = tvm.IRModule.from_expr(func_parallel)
+    # parallel (i, 0, n) {
+    #   allocate j[int32 * 1]
+    #   j[0] = 0
+    #   while((j[0] < 10)){
+    #     // attr [A] storage_scope = "global"
+    #     allocate A[float32 * n]
+    #     A[j[0]] = (A[j[0]] + 2f)
+    #     j[0] = (j[0] + (j[0] + 1))
+    #   }
+    # }
+    body = tvm.tirx.transform.StorageRewrite()(mod)["func_parallel"]
+    # Navigate to inside the for loop, then check that allocations exist
+    # The structure with DeclBuffer is:
+    #   parallel (i, 0, n) { DeclBuffer(j, DeclBuffer(A, ...)) }
+    # or with Allocate+DeclBuffer pairs
+    inner = body.body.body  # inside For
+    # Skip DeclBuffer nodes to find Allocate
+    num_alloc = [0]
+
+    def count_alloc(n):
+        if isinstance(n, tvm.tirx.AllocBuffer):
+            num_alloc[0] += 1
+
+    tvm.tirx.stmt_functor.post_order_visit(inner, count_alloc)
+    assert num_alloc[0] == 2  # j and A allocations
+
+    mod = tvm.IRModule.from_expr(func_serial)
+    body = tvm.tirx.transform.StorageRewrite()(mod)["func_serial"]
+    num_alloc[0] = 0
+    tvm.tirx.stmt_functor.post_order_visit(body.body, count_alloc)
+    assert num_alloc[0] == 2  # j and A allocations
+
+
+def test_alloc_seq_type():
+    @T.prim_func(s_tir=True)
+    def func(n: T.int32):
+        for i in T.serial(n):
+            for j in range(10):
+                A = T.alloc_buffer((200,), scope="local.L0A")
+                A1 = T.alloc_buffer((200,), scope="local.L0A")
+                A[j] = T.float32(1.2)
+                A1[j] = T.float32(1.3)
+                B = T.alloc_buffer((200,), "int16", scope="local.L0A")
+                B[j] = T.int16(1)
+                C = T.alloc_buffer((200,), "int16", scope="local.L0A")
+                C[j] = T.int16(1)
+                D = T.alloc_buffer((200,), "int16", scope="local.L0A")
+                D[j] = B[j] + C[j]
+                A2 = T.alloc_buffer((200,), scope="local.L0A")
+                A2[j] = A[j]
+
+    mod = tvm.IRModule.from_expr(func)
+    body = tvm.tirx.transform.StorageRewrite()(mod)["func"].body
+
+    num_alloc = [0]
+
+    def verify(n):
+        if isinstance(n, tvm.tirx.AllocBuffer):
+            num_alloc[0] += 1
+            assert n.buffer.shape[0].value == 500
+
+    tvm.tirx.stmt_functor.post_order_visit(body, verify)
+    assert num_alloc[0] == 1
+
+
+def test_alloc_seq_type2():
+    scope_tb = "local.L0A2"
+
+    @T.prim_func(s_tir=True)
+    def func(n: T.int32):
+        for i in T.serial(n):
+            for j in range(10):
+                A = T.alloc_buffer((200,), scope=scope_tb)
+                A[j] = T.float32(1.2)
+            for j in range(20):
+                B = T.alloc_buffer((400,), "int16", scope=scope_tb)
+                B[j] = T.int16(1)
+            for j in range(10):
+                C = T.alloc_buffer((200,), scope=scope_tb)
+                C[j] = T.float32(1.2)
+
+    mod = tvm.IRModule.from_expr(func)
+    body = tvm.tirx.transform.StorageRewrite()(mod)["func"].body
+
+    num_alloc = [0]
+
+    def verify(n):
+        if isinstance(n, tvm.tirx.AllocBuffer):
+            num_alloc[0] += 1
+            assert n.buffer.shape[0].value == 200
+
+    tvm.tirx.stmt_functor.post_order_visit(body, verify)
+    assert num_alloc[0] == 1
+
+
+def test_reuse_small_buffer():
+    @T.prim_func(s_tir=True)
+    def func(n: T.int32):
+        for i in T.serial(n):
+            for j in range(10):
+                A = T.alloc_buffer((200,), "int16", scope="local.L0A")
+                A[j] = T.int16(1)
+                B = T.alloc_buffer((200,), "int16", scope="local.L0A")
+                B[j] = T.int16(1)
+                B1 = T.alloc_buffer((200,), "int16", scope="local.L0A")
+                B1[j] = A[j] + B[j]
+                C = T.alloc_buffer((400,), "int16", scope="local.L0A")
+                C[j] = T.int16(1)
+                D = T.alloc_buffer((400,), "int16", scope="local.L0A")
+                D[j] = T.int16(1)
+                E = T.alloc_buffer((400,), "int16", scope="local.L0A")
+                E[j] = C[j]
+
+    mod = tvm.IRModule.from_expr(func)
+    body = tvm.tirx.transform.StorageRewrite()(mod)["func"].body
+
+    num_alloc = [0]
+
+    def verify(n):
+        if isinstance(n, tvm.tirx.AllocBuffer):
+            num_alloc[0] += 1
+            assert n.buffer.shape[0].value == 800
+
+    tvm.tirx.stmt_functor.post_order_visit(body, verify)
+    assert num_alloc[0] == 1
+
+
+def test_access_in_let_value():
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((8,), "float32")):
+        for i in range(8):
+            B = T.alloc_buffer((1,))
+            B[0] = 3.14
+            x: T.let[T.float32] = T.exp(B[0], dtype="float32")
+            A[i] = (x + 1.0) / (x - 1.0)
+
+    @T.prim_func(s_tir=True)
+    def func_rewritten(A: T.Buffer((8,), "float32")) -> None:
+        B = T.alloc_buffer((1,))
+        for i in range(8):
+            B[0] = 3.14
+            x: T.let[T.float32] = T.exp(B[0], dtype="float32")
+            A[i] = (x + 1.0) / (x - 1.0)
+
+    mod = tvm.tirx.transform.StorageRewrite()(
+        tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+    )
+    tvm.ir.assert_structural_equal(mod["main"], func_rewritten.with_attr("global_symbol", "main"))
+
+
+def test_let_buffer_rewrite():
+    """StorageRewrite replaces the bound var of backing allocations
+
+    If StorageRewrite replaces the backing variable of an array, such
+    as when vectorizing the storage type, the variable must be
+    replaced in the Bind that defines it.  Currently, StmtMutator
+    only visits usage of variables, and does not visit definitions of
+    variables, so the definition in a Bind must be explicitly
+    handled.
+    """
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main() -> None:
+            A_data: T.let[T.handle("int32")] = T.call_extern(
+                "dummy_func", dtype=T.handle("int32").ty
+            )
+            A = T.decl_buffer([8], "int32", data=A_data)
+            A[0:8] = T.broadcast(42, 8)
+
+    @I.ir_module(check_well_formed=False)
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main() -> None:
+            A_data: T.let[T.handle("int32x8")] = T.call_extern(
+                "dummy_func", dtype=T.handle("int32x8").ty
+            )
+            A = T.decl_buffer([8], "int32", data=A_data)
+            A_1 = T.Buffer([1], "int32x8", data=A_data)
+            A_1[0] = T.broadcast(42, 8)
+
+    After = tvm.tirx.transform.StorageRewrite()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_rewrite_in_place_use_of_non_flat_buffer():
+    """A non-flat buffer may be re-used for in-place operations"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16, 16), "float32"), D: T.Buffer((16, 16), "float32")):
+            B = T.decl_buffer(
+                [16, 16],
+                dtype="float32",
+                axis_separators=[1],
+            )
+            C = T.decl_buffer(
+                [16, 16],
+                dtype="float32",
+                axis_separators=[1],
+            )
+
+            for i, j in T.grid(16, 16):
+                B[i, j] = A[i, j]
+
+            for i, j in T.grid(16, 16):
+                C[i, j] = 2.0 * B[i, j]
+
+            for i, j in T.grid(16, 16):
+                D[i, j] = C[i, j]
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16, 16), "float32"), D: T.Buffer((16, 16), "float32")):
+            B = T.decl_buffer([16, 16], dtype="float32", axis_separators=[1])
+            C = T.decl_buffer(
+                [16, 16],
+                dtype="float32",
+                axis_separators=[1],
+                data=B.data,
+            )
+
+            for i, j in T.grid(16, 16):
+                B[i, j] = A[i, j]
+
+            for i, j in T.grid(16, 16):
+                C[i, j] = 2.0 * B[i, j]
+
+            for i, j in T.grid(16, 16):
+                D[i, j] = C[i, j]
+
+    After = tvm.tirx.transform.StorageRewrite()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_no_rewrite_of_shared_non_flat_buffer():
+    """In general, sharing of non-flat buffer isn't supported
+
+    The current packing algorithms in StorageRewrite assume a flat
+    memory space, and do not support packing of N-d buffers.  For
+    buffers with axis separators, normal buffer sharing should be
+    disabled.
+
+    Like test_rewrite_in_place_use_of_non_flat_buffer, except that B and C do
+    not have matching shapes.
+    """
+
+    @T.prim_func(s_tir=True)
+    def Before(A: T.Buffer((16, 16), "float32"), D: T.Buffer((16, 16), "float32")):
+        B = T.decl_buffer(
+            [16, 16],
+            dtype="float32",
+            axis_separators=[1],
+        )
+        C = T.decl_buffer(
+            [20, 20],
+            dtype="float32",
+            axis_separators=[1],
+        )
+
+        for i, j in T.grid(16, 16):
+            B[i, j] = A[i, j]
+
+        for i, j in T.grid(16, 16):
+            C[i, j] = 2.0 * B[i, j]
+
+        for i, j in T.grid(16, 16):
+            D[i, j] = C[i, j]
+
+    Expected = Before
+
+    After = tvm.tirx.transform.StorageRewrite()(tvm.IRModule.from_expr(Before))
+    tvm.ir.assert_structural_equal(After["Before"], Expected)
+
+
+def test_rewrite_decl_buffer():
+    """A DeclBuffer node may appear in StorageRewrite's input"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32"), D: T.Buffer(16, "float32")):
+            B = T.decl_buffer(16, dtype="float32")
+            C = T.decl_buffer(16, dtype="float32")
+
+            for i in range(16):
+                B[i] = A[i]
+
+            for i in range(16):
+                C[i] = 2.0 * B[i]
+
+            for i in range(16):
+                D[i] = C[i]
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32"), D: T.Buffer(16, "float32")):
+            B = T.decl_buffer(16, dtype="float32")
+            C = T.decl_buffer(16, dtype="float32", data=B.data)
+
+            for i in range(16):
+                B[i] = A[i]
+
+            for i in range(16):
+                C[i] = 2.0 * B[i]
+
+            for i in range(16):
+                D[i] = C[i]
+
+    After = tvm.tirx.transform.StorageRewrite()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_no_orphaned_decl_buffer():
+    """A DeclBuffer of an unused Allocate should be removed
+
+    StorageRewrite removes any allocations that are unused.  When it
+    does so, any DeclBuffer that refers to that allocation should also
+    be removed.
+    """
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32"), D: T.Buffer(16, "float32")):
+            B = T.decl_buffer(16, dtype="float32")
+            C = T.decl_buffer(16, dtype="float32")
+            Unused = T.decl_buffer(16, dtype="float32")
+
+            for i in range(16):
+                B[i] = A[i]
+
+            for i in range(16):
+                C[i] = 2.0 * B[i]
+
+            for i in range(16):
+                D[i] = C[i]
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(16, "float32"), D: T.Buffer(16, "float32")):
+            B = T.decl_buffer(16, dtype="float32")
+            C = T.decl_buffer(16, dtype="float32", data=B.data)
+
+            for i in range(16):
+                B[i] = A[i]
+
+            for i in range(16):
+                C[i] = 2.0 * B[i]
+
+            for i in range(16):
+                D[i] = C[i]
+
+    After = tvm.tirx.transform.StorageRewrite()(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_unroll_loop.py

@@ -0,0 +1,161 @@
+# 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_unroll_loop():
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.handle, n: T.int64):
+            Ab = T.match_buffer(A, (n,), "int64")
+            for i in T.serial(n, n + 2):
+                for j in T.unroll(8):
+                    Ab[j + 1] = Ab[i] + T.int64(1)
+
+    mod = Module
+    stmt = mod["main"].body
+
+    assert isinstance(stmt, tvm.tirx.For)
+
+    with tvm.transform.PassContext(config={"tirx.UnrollLoop": {"auto_max_step": 16}}):
+        ret = tvm.tirx.transform.UnrollLoop()(mod)["main"].body
+        assert not isinstance(ret, tvm.tirx.For)
+
+    with tvm.transform.PassContext(config={"tirx.UnrollLoop": {"auto_max_step": 15}}):
+        ret = tvm.tirx.transform.UnrollLoop()(mod)["main"].body
+        assert isinstance(ret, tvm.tirx.For)
+
+    with tvm.transform.PassContext(
+        config={"tirx.UnrollLoop": {"auto_max_step": 16, "explicit_unroll": False}}
+    ):
+        ret = tvm.tirx.transform.UnrollLoop()(mod)["main"].body
+        assert isinstance(ret, tvm.tirx.For)
+        assert ret.kind == tvm.tirx.ForKind.UNROLLED
+
+    @I.ir_module
+    class ModuleWithPragma:
+        @T.prim_func(s_tir=True)
+        def main(A: T.handle, n: T.int64):
+            Ab = T.match_buffer(A, (n,), "int64")
+            with T.attr(T.int32(0), "pragma_auto_unroll_max_step", 16):
+                for i in T.serial(n, n + 2):
+                    for j in T.unroll(8):
+                        Ab[j + 1] = Ab[i] + T.int64(1)
+            for i in T.serial(n, n + 2):
+                for j in T.unroll(8):
+                    Ab[j + 1] = Ab[i] + T.int64(1)
+
+    with tvm.transform.PassContext(
+        config={"tirx.UnrollLoop": {"auto_max_depth": 8, "explicit_unroll": False}}
+    ):
+        ret = tvm.tirx.transform.UnrollLoop()(ModuleWithPragma)["main"].body
+        assert isinstance(ret[0], tvm.tirx.For)
+        assert ret[0].kind == tvm.tirx.ForKind.UNROLLED
+        assert isinstance(ret[1], tvm.tirx.For)
+        assert ret[1].kind != tvm.tirx.ForKind.UNROLLED
+
+
+def test_unroll_fake_loop():
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.handle, n: T.int64):
+            Ab = T.match_buffer(A, (n,), "int32")
+            for i in T.serial(1):
+                Ab[i * 2] = 3
+                for j in T.serial(10):
+                    Ab[j + 1] = Ab[i] + 1
+
+    with tvm.transform.PassContext(
+        config={
+            "tirx.UnrollLoop": {"auto_max_depth": 8, "auto_max_extent": 1, "explicit_unroll": False}
+        }
+    ):
+        ret = tvm.tirx.transform.UnrollLoop()(Module)["main"].body
+        assert isinstance(ret[0], tvm.tirx.BufferStore)
+
+
+def test_unroll_allocations():
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            for i in T.unroll(2):
+                buf = T.alloc_buffer([16], "float32")
+                buf[0] = 0.0
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            buf1 = T.alloc_buffer([16], "float32")
+            buf1[0] = 0.0
+            buf2 = T.alloc_buffer([16], "float32")
+            buf2[0] = 0.0
+
+    after = tvm.tirx.transform.UnrollLoop()(Before)
+
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_unroll_local_access():
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((64,), "float32")):
+            for bx in T.thread_binding(4, thread="blockIdx.x"):
+                for tx in T.thread_binding(4, thread="threadIdx.x"):
+                    A_local = T.alloc_buffer((4,), scope="local")
+                    for i in T.serial(4):
+                        A_local[i] = T.float32(i)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main(B: T.Buffer((64,), "float32")):
+            for bx in T.thread_binding(4, thread="blockIdx.x"):
+                for tx in T.thread_binding(4, thread="threadIdx.x"):
+                    A_local = T.alloc_buffer((4,), scope="local")
+                    A_local[0] = T.float32(0)
+                    A_local[1] = T.float32(1)
+                    A_local[2] = T.float32(2)
+                    A_local[3] = T.float32(3)
+
+    with tvm.transform.PassContext(
+        config={
+            "tirx.UnrollLoop": {
+                "auto_max_depth": 0,
+                "auto_max_extent": 1,
+                "explicit_unroll": True,
+                "unroll_local_access": True,
+            }
+        }
+    ):
+        after = tvm.tirx.transform.UnrollLoop()(Before)
+        after = tvm.tirx.transform.StmtSimplify()(after)
+
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+if __name__ == "__main__":
+    test_unroll_local_access()
+    test_unroll_loop()
+    test_unroll_fake_loop()
+    test_unroll_allocations()

tests/python/tirx-transform/test_tir_transform_vectorize.py

@@ -0,0 +1,820 @@
+# 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
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.target.codegen import llvm_version_major
+
+simple_target = tvm.target.Target({"kind": "llvm", "mtriple": "x86_64-linux-gnu"})
+sve_target = tvm.target.Target(
+    {
+        "kind": "llvm",
+        "device": "arm_cpu",
+        "mtriple": "aarch64-linux-gnu",
+        "mattr": ["+v8.2a", "+sve"],
+    }
+)
+
+
+@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() * 4, sve_target)])
+def test_vectorize_loop(extent, target):
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16,), "float32")):
+            for j in T.vectorized(0, extent):
+                A[j] = 1
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16,), "float32")):
+            A[T.Ramp(0, 1, extent)] = T.Broadcast(1, extent)
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+
+
+def test_vectorize_vector():
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((4,), "float32x4"), n: T.int32):
+            for i in range(n):
+                for j in T.vectorized(4):
+                    A[j] = T.Broadcast(T.float32(1), 4)
+
+    mod = tvm.tirx.transform.VectorizeLoop()(Module)
+    stmt = mod["main"].body
+
+    assert isinstance(stmt, tvm.tirx.For)
+    assert not isinstance(stmt.body, tvm.tirx.For)
+    assert len(stmt.body.indices) == 1
+    assert isinstance(stmt.body.indices[0], tvm.tirx.Ramp)
+    assert isinstance(stmt.body.value, tvm.tirx.Broadcast)
+
+
+def test_vectorize_vector_scalable_error():
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32")):
+            for j in T.vectorized(T.vscale() * 4):
+                A[j * 4 : j * 4 + 4] = T.Broadcast(T.float32(1), 4)
+
+    error_msg = "Creating scalable vectors from existing vectors is not supported."
+    with tvm.target.Target(sve_target):
+        with pytest.raises(tvm.error.InternalError, match=error_msg):
+            tvm.tirx.transform.VectorizeLoop()(Module)
+
+
+def test_vectorize_vector_scalable_error2():
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32xvscalex4")):
+            for j in T.vectorized(4):
+                A[j] = T.Broadcast(T.float32(1), T.vscale() * 4)
+
+    error_msg = "Vectorizing over scalable buffer elements is not supported in vectorizer."
+    with pytest.raises(tvm.error.InternalError, match=error_msg):
+        tvm.tirx.transform.VectorizeLoop()(Module)
+
+
+def test_vectorize_vector_scalable_error3():
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32")):
+            for j in T.vectorized(4):
+                A[j * T.vscale() * 4 : j * T.vscale() * 4 + T.vscale() * 4] = T.Broadcast(
+                    T.float32(1), T.vscale() * 4
+                )
+
+    error_msg = "Vectorizing over existing scalable vectors is not supported."
+    with pytest.raises(tvm.error.InternalError, match=error_msg):
+        with tvm.target.Target(sve_target):
+            tvm.tirx.transform.VectorizeLoop()(Module)
+
+
+def test_vectorize_vector_scalable_error4():
+    @I.ir_module
+    class Module:
+        @T.prim_func(private=True, s_tir=True)
+        def main(A: T.Buffer((25,), "float32")):
+            for j in T.vectorized(T.vscale() * 4):
+                A[j * T.vscale() * 4 : j * T.vscale() * 4 + T.vscale() * 4] = T.Broadcast(
+                    T.float32(1), T.vscale() * 4
+                )
+
+    error_msg = "Creating scalable vectors from existing vectors is not supported."
+    with pytest.raises(tvm.error.InternalError, match=error_msg):
+        with tvm.target.Target(sve_target):
+            tvm.tirx.transform.VectorizeLoop()(Module)
+
+
+def test_vectorize_with_if():
+    extent = 4
+    target = simple_target
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(a: T.handle, n: T.int32, x: T.int32):
+            A = T.match_buffer(a, (25,), "float32")
+            for i in T.vectorized(extent):
+                if x < n:
+                    A[i] = A[i] + T.float32(1)
+                else:
+                    if i < n:
+                        A[i] = T.float32(2)
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(a: T.handle, n: T.int32, x: T.int32):
+            A = T.match_buffer(a, (25,), "float32")
+            if x < n:
+                A[T.Ramp(0, 1, extent)] = A[T.Ramp(0, 1, extent)] + T.Broadcast(
+                    T.float32(1), extent
+                )
+            else:
+                for i_s in range(extent):
+                    if i_s < n:
+                        A[i_s] = T.float32(2)
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+
+
+def test_vectorize_if_scalable_extent():
+    extent = T.vscale() * 4
+    target = sve_target
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(a: T.handle, n: T.int32, x: T.int32):
+            A = T.match_buffer(a, (25,), "float32")
+            for i in T.vectorized(extent):
+                if x < n:
+                    A[i] = A[i] + T.float32(1)
+                else:
+                    if i < n:
+                        A[i] = T.float32(2)
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(a: T.handle, n: T.int32, x: T.int32):
+            A = T.match_buffer(a, (25,), "float32")
+            if x < n:
+                A[T.Ramp(0, 1, extent)] = A[T.Ramp(0, 1, extent)] + T.Broadcast(
+                    T.float32(1), extent
+                )
+            else:
+                A.vstore(
+                    [T.Ramp(0, 1, T.vscale() * 4)],
+                    T.Broadcast(T.float32(2), T.vscale() * 4),
+                    predicate=T.get_active_lane_mask("uint1xvscalex4", 0, n),
+                )
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+
+
+@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() * 4, sve_target)])
+def test_vectorize_let(extent, target):
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32")):
+            for i in T.vectorized(extent):
+                v: T.let = A[i] + T.float32(1)
+                A[i] = v + T.float32(2)
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32")):
+            v: T.let = A[T.Ramp(0, 1, extent)] + T.Broadcast(T.float32(1), extent)
+            A[T.Ramp(0, 1, extent)] = v + T.Broadcast(T.float32(2), extent)
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+
+
+@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() * 4, sve_target)])
+def test_vectorize_with_le_cond(extent, target):
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16,), "float32"), n: T.int32):
+            for i in T.vectorized(extent):
+                if i <= n:
+                    A[i] = A[i] + T.float32(1)
+
+    with tvm.target.Target(target):
+        stmt = tvm.tirx.transform.VectorizeLoop()(Module)["main"].body
+
+        # Check that the loop wasn't vectorised
+        assert isinstance(stmt, tvm.tirx.For)
+
+
+@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() * 4, sve_target)])
+def test_vectorize_with_ge_cond(extent, target):
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16,), "float32"), n: T.int32):
+            for i in T.vectorized(extent):
+                if i >= n:
+                    A[i] = A[i] + T.float32(1)
+
+    with tvm.target.Target(target):
+        stmt = tvm.tirx.transform.VectorizeLoop()(Module)["main"].body
+
+        # Check that the loop wasn't vectorised
+        assert isinstance(stmt, tvm.tirx.For)
+
+
+@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() * 4, sve_target)])
+def test_vectorize_if_then_else_scalarize(extent, target):
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32")):
+            for i in T.vectorized(extent):
+                A[i] = T.if_then_else(i > 0, A[i] + T.float32(1), A[i])
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32")):
+            for i_s in range(extent):
+                A[i_s] = T.if_then_else(i_s > 0, A[i_s] + T.float32(1), A[i_s])
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+
+
+@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() * 4, sve_target)])
+def test_vectorize_if_then_else_vector(extent, target):
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32"), n: T.int32):
+            for i in range(n):
+                for j in T.vectorized(extent):
+                    A[i * extent + j] = T.if_then_else(i > 0, A[i * extent + j], 0)
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32"), n: T.int32):
+            for i in range(n):
+                A[T.Ramp(i * extent, 1, extent)] = T.if_then_else(
+                    i > 0, A[T.Ramp(i * extent, 1, extent)], T.Broadcast(0, extent)
+                )
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+
+
+def test_vectorize_let_if_then_else():
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            for i in T.vectorized(4):
+                if i < 2:
+                    result: T.let[T.int32] = T.if_then_else(i < 1, 1, 2)
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main():
+            for i_s in range(4):
+                if i_s < 2:
+                    result: T.let[T.int32] = T.if_then_else(i_s < 1, 1, 2)
+                    T.evaluate(0)
+
+    with tvm.target.Target(simple_target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+
+
+def test_vectorize_while_fail():
+    """A while loop inside a vectorized loop should fail."""
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((64,), "float32"),
+            B: T.Buffer((64,), "float32"),
+            C: T.Buffer((64,), "float32"),
+        ):
+            # Initialize C to 0
+            for j in range(64):
+                C[j] = T.float32(0)
+
+            # While loop inside vectorized loop (should fail)
+            i = T.decl_buffer((1,), "int32", scope="local")
+            i[0] = 0
+            for j in T.vectorized(64):
+                while i[0] < 10:
+                    C[j] = C[j] + A[j] + B[j]
+                    i[0] = i[0] + 1
+
+    try:
+        tvm.compile(Module, target="llvm")
+        assert False
+    except RuntimeError as e:
+        error_msg = str(e).split("\n")[-1]
+        expected = "A while loop inside a vectorized loop not supported"
+        assert expected in error_msg
+
+
+@pytest.mark.parametrize(
+    "extent, vec_str, target",
+    [(16, "float32x16", simple_target), (T.vscale() * 8, "float32xvscalex8", sve_target)],
+)
+def test_vectorize_with_reinterpret(extent, vec_str, target):
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16,), "int32"), B: T.Buffer((16,), "float32")):
+            for i in T.vectorized(0, extent):
+                B[i] = T.reinterpret("float32", A[i])
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16,), "int32"), B: T.Buffer((16,), "float32")):
+            B[T.Ramp(0, 1, extent)] = T.reinterpret(vec_str, A[T.Ramp(0, 1, extent)])
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+
+
+@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() * 4, sve_target)])
+@pytest.mark.parametrize(
+    "op",
+    (
+        T.Mul,
+        T.Add,
+        T.Sub,
+        T.Div,
+        T.Mod,
+        T.FloorDiv,
+        T.FloorMod,
+        T.Min,
+        T.Max,
+        T.EQ,
+        T.LT,
+        T.LE,
+        T.GE,
+        T.GT,
+        T.NE,
+    ),
+)
+def test_vectorize_binary(op, extent, target):
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32"), B: T.Buffer((25,), "float32")):
+            for j in T.vectorized(extent):
+                A[j] = op(T.float32(3), B[j])
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32"), B: T.Buffer((25,), "float32")):
+            A[T.Ramp(0, 1, extent)] = op(T.Broadcast(T.float32(3), extent), B[T.Ramp(0, 1, extent)])
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+
+
+@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() * 4, sve_target)])
+@pytest.mark.parametrize("op", (T.And, T.Or))
+def test_vectorize_logical(op, extent, target):
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "bool"), B: T.Buffer((25,), "bool")):
+            for j in T.vectorized(extent):
+                A[j] = op(T.bool(1), B[j])
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "bool"), B: T.Buffer((25,), "bool")):
+            A[T.Ramp(0, 1, extent)] = op(T.Broadcast(T.bool(1), extent), B[T.Ramp(0, 1, extent)])
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+
+
+@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() * 4, sve_target)])
+def test_vectorize_select(extent, target):
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32"), B: T.Buffer((25,), "float32")):
+            for j in T.vectorized(extent):
+                A[j] = T.Select(T.bool(True), A[j], B[j])
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32"), B: T.Buffer((25,), "float32")):
+            A[T.Ramp(0, 1, extent)] = T.Select(
+                T.Broadcast(T.bool(True), extent),
+                A[T.Ramp(0, 1, extent)],
+                B[T.Ramp(0, 1, extent)],
+            )
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+
+
+@pytest.mark.parametrize(
+    "extent, vec_str, target",
+    [(4, "int32x4", simple_target), (T.vscale() * 4, "int32xvscalex4", sve_target)],
+)
+def test_vectorize_cast(extent, vec_str, target):
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "int32"), B: T.Buffer((25,), "float32")):
+            for j in T.vectorized(extent):
+                A[j] = T.Cast("int32", B[j])
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "int32"), B: T.Buffer((25,), "float32")):
+            A[T.Ramp(0, 1, extent)] = T.Cast(vec_str, B[T.Ramp(0, 1, extent)])
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+
+
+def test_illegal_extent():
+    @I.ir_module(check_well_formed=False)
+    class Mod:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "int32")):
+            n = T.Var("n", dtype="int32")
+            for j in T.vectorized(n):
+                A[j] = 3
+
+    error_msg = "Failed to vectorize loop with extent n for target None"
+    with pytest.raises(tvm.error.InternalError, match=error_msg):
+        tvm.tirx.transform.VectorizeLoop()(Mod)
+
+
+def test_illegal_vscale_in_non_sve_compilation():
+    @I.ir_module
+    class Mod:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16,), "float32")):
+            for j in T.vectorized(0, 4 * T.vscale()):
+                A[j] = 13
+
+    msg = "Failed to vectorize loop with extent T.vscale\\(\\) \\* 4 for target"
+    with tvm.target.Target(simple_target):
+        with pytest.raises(tvm.error.InternalError, match=msg):
+            tvm.tirx.transform.VectorizeLoop()(Mod)
+
+
+def test_vectorize_and_predicate_all_buffer_loads_stores():
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i_0 in T.serial(T.ceildiv(14, 4)):
+            for i_1 in T.vectorized(4):
+                if i_0 * 4 + i_1 < 14:
+                    B[i_0 * 4 + i_1] = A[i_0 * 4 + i_1] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def expected(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i_0 in range(4):
+            load_a = T.meta_var(
+                A.vload(
+                    [T.Ramp(i_0 * 4, 1, 4)],
+                    predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+                )
+            )
+            add_1 = T.meta_var(load_a + T.Broadcast(T.float32(1), 4))
+            B.vstore(
+                [T.Ramp(i_0 * 4, 1, 4)],
+                add_1,
+                predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+            )
+
+    mod = tvm.IRModule.from_expr(before)
+    with tvm.transform.PassContext(config={"tirx.enable_buffer_level_predication": True}):
+        after = tvm.tirx.transform.VectorizeLoop()(mod)["main"]
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_vectorize_and_predicate_some_buffer_loads_stores():
+    # Currently revert to scalarizing the block if not all accesses
+    # have been predicated, otherwise incorrect code is generated.
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i_0 in T.serial(T.ceildiv(14, 4)):
+            for i_1 in T.vectorized(4):
+                if i_0 * 4 + i_1 < 14:
+                    B[i_0 * 4 + i_1] = A[i_0] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def expected(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i_0, i_1_s in T.grid(4, 4):
+            if i_0 * 4 + i_1_s < 14:
+                B[i_0 * 4 + i_1_s] = A[i_0] + T.float32(1)
+
+    mod = tvm.IRModule.from_expr(before)
+    with tvm.transform.PassContext(config={"tirx.enable_buffer_level_predication": True}):
+        after = tvm.tirx.transform.VectorizeLoop()(mod)["main"]
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_vectorize_and_predicate_multiple_access_statements():
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i_0 in T.serial(T.ceildiv(14, 4)):
+            for i_1 in T.vectorized(4):
+                if i_0 * 4 + i_1 < 14:
+                    A[i_0 * 4 + i_1] = 2.0
+                    B[i_0 * 4 + i_1] = 1.0
+
+    @T.prim_func(s_tir=True)
+    def expected(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i_0 in range(4):
+            A.vstore(
+                [T.Ramp(i_0 * 4, 1, 4)],
+                T.Broadcast(T.float32(2), 4),
+                predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+            )
+            B.vstore(
+                [T.Ramp(i_0 * 4, 1, 4)],
+                T.Broadcast(T.float32(1), 4),
+                predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+            )
+
+    before_mod = tvm.IRModule.from_expr(before)
+    with tvm.transform.PassContext(config={"tirx.enable_buffer_level_predication": True}):
+        after = tvm.tirx.transform.VectorizeLoop()(before_mod)["main"]
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_vectorize_and_predicate_invalid_conditions():
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i_0 in T.serial(T.ceildiv(14, 4)):
+            for i_1 in T.vectorized(4):
+                if i_0 * 4 + i_1 > 14:
+                    A[i_0 * 4 + i_1] = 2.0
+                if 14 < i_0 * 4 + i_1:
+                    A[i_0 * 4 + i_1] = 2.0
+                if i_0 * 4 + i_1 < i_0 * 4 + i_1:
+                    A[i_0 * 4 + i_1] = 2.0
+
+    @T.prim_func(s_tir=True)
+    def expected(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i_0 in range(4):
+            for i_1_s in range(4):
+                if i_0 * 4 + i_1_s > 14:
+                    A[i_0 * 4 + i_1_s] = T.float32(2)
+            for i_1_s in range(4):
+                if 14 < i_0 * 4 + i_1_s:
+                    A[i_0 * 4 + i_1_s] = T.float32(2)
+            for i_1_s in range(4):
+                if i_0 * 4 + i_1_s < i_0 * 4 + i_1_s:
+                    A[i_0 * 4 + i_1_s] = T.float32(2)
+
+    before_mod = tvm.IRModule.from_expr(before)
+    with tvm.transform.PassContext(config={"tirx.enable_buffer_level_predication": True}):
+        after = tvm.tirx.transform.VectorizeLoop()(before_mod)["main"]
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_vectorize_with_explicitly_disabled_buffer_level_predication():
+    # Since the target has the VLA feature, buffer level predication is enabled
+    # by default. However, it has been explicitly disabled by the pass context
+    # option, so no buffer-level predicates should be added.
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i_0 in T.serial(T.ceildiv(14, 4)):
+            for i_1 in T.vectorized(4):
+                if i_0 * 4 + i_1 < 14:
+                    B[i_0 * 4 + i_1] = A[i_0 * 4 + i_1] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def expected(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i_0, i_1_s in T.grid(4, 4):
+            if i_0 * 4 + i_1_s < 14:
+                B[i_0 * 4 + i_1_s] = A[i_0 * 4 + i_1_s] + T.float32(1)
+
+    mod = tvm.IRModule.from_expr(before)
+    with tvm.transform.PassContext(config={"tirx.enable_buffer_level_predication": False}):
+        with tvm.target.Target(sve_target):
+            after = tvm.tirx.transform.VectorizeLoop()(mod)["main"]
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_vectorize_and_predicate_buffer_load_stores_with_sve_func_attr_target():
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True, "target": sve_target})
+        for i_0 in T.serial(T.ceildiv(14, 4)):
+            for i_1 in T.vectorized(4):
+                if i_0 * 4 + i_1 < 14:
+                    B[i_0 * 4 + i_1] = A[i_0 * 4 + i_1] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def expected(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True, "target": sve_target})
+        for i_0 in range(4):
+            load_a = T.meta_var(
+                A.vload(
+                    [T.Ramp(i_0 * 4, 1, 4)],
+                    predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+                )
+            )
+            add_1 = T.meta_var(load_a + T.Broadcast(T.float32(1), 4))
+            B.vstore(
+                [T.Ramp(i_0 * 4, 1, 4)],
+                add_1,
+                predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+            )
+
+    mod = tvm.IRModule.from_expr(before)
+    after = tvm.tirx.transform.VectorizeLoop()(mod)["main"]
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_vectorize_and_predicate_buffer_load_stores_with_sve_attr_scope_target():
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        with T.attr(sve_target, "target", 0):
+            for i_0 in T.serial(T.ceildiv(14, 4)):
+                for i_1 in T.vectorized(4):
+                    if i_0 * 4 + i_1 < 14:
+                        B[i_0 * 4 + i_1] = A[i_0 * 4 + i_1] + 1.0
+
+    @T.prim_func(s_tir=True)
+    def expected(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        with T.attr(sve_target, "target", 0):
+            for i_0 in range(4):
+                load_a = T.meta_var(
+                    A.vload(
+                        [T.Ramp(i_0 * 4, 1, 4)],
+                        predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+                    )
+                )
+                add_1 = T.meta_var(load_a + T.Broadcast(T.float32(1), 4))
+                B.vstore(
+                    [T.Ramp(i_0 * 4, 1, 4)],
+                    add_1,
+                    predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+                )
+
+    mod = tvm.IRModule.from_expr(before)
+    after = tvm.tirx.transform.VectorizeLoop()(mod)["main"]
+    tvm.ir.assert_structural_equal(after, expected)
+
+
+@pytest.mark.parametrize(
+    "extent, vec_str, target",
+    [(4, "float32x4", simple_target)],
+)
+def test_vectorize_llvm_pure_intrin(extent, vec_str, target):
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32"), B: T.Buffer((25,), "float32")):
+            for j in T.vectorized(extent):
+                A[j] = T.call_llvm_pure_intrin("float32", "llvm.sqrt", B[j])
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "float32"), B: T.Buffer((25,), "float32")):
+            A[T.Ramp(0, 1, extent)] = T.call_llvm_pure_intrin(
+                vec_str, "llvm.sqrt", B[T.Ramp(0, 1, extent)]
+            )
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+        mod = tvm.compile(mod, target=target)
+
+
+@pytest.mark.parametrize(
+    "extent, vec_str, target",
+    [(4, "int32x4", simple_target)],
+)
+def test_vectorize_llvm_pure_intrin_fail(extent, vec_str, target):
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "int32"), B: T.Buffer((25,), "float32")):
+            for j in T.vectorized(extent):
+                A[j] = T.call_llvm_pure_intrin("int32", "llvm.lround", B[j])
+
+    @I.ir_module
+    class After:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((25,), "int32"), B: T.Buffer((25,), "float32")):
+            A[T.Ramp(0, 1, extent)] = T.call_llvm_pure_intrin(
+                vec_str, "llvm.lround", B[T.Ramp(0, 1, extent)]
+            )
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.transform.VectorizeLoop()(Before)
+        tvm.ir.assert_structural_equal(mod, After)
+        # LLVM 20 added vector support for llvm.lround/llvm.llround.  The IR Verifier's
+        # "Intrinsic does not support vectors" check was removed in release/20.x, so
+        # compilation only fails on LLVM <= 19.
+        if llvm_version_major() >= 20:
+            tvm.compile(mod, target=target)
+        else:
+            with pytest.raises(Exception, match="Intrinsic does not support vectors"):
+                tvm.compile(mod, target=target)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()