chore: import upstream snapshot with attribution

tests/python/tirx/transform/test_tirx_expr_functor.py

@@ -0,0 +1,829 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import tvm
+import tvm.testing
+from tvm import tirx as tir
+from tvm.ir import Call, Op
+from tvm.ir.base import assert_structural_equal
+from tvm.tirx.expr import (
+    EQ,
+    GE,
+    GT,
+    LE,
+    LT,
+    NE,
+    Add,
+    And,
+    Broadcast,
+    BufferLoad,
+    Cast,
+    Div,
+    FloatImm,
+    FloorDiv,
+    FloorMod,
+    IntImm,
+    Let,
+    Max,
+    Min,
+    Mod,
+    Mul,
+    Not,
+    Or,
+    ProducerLoad,
+    Ramp,
+    Reduce,
+    Select,
+    Shuffle,
+    StringImm,
+    Sub,
+    Var,
+)
+from tvm.tirx.expr_functor import ExprMutator, ExprVisitor
+
+# Basic example variables for testing
+n = tir.Var("n", "int32")
+m = tir.Var("m", "int32")
+x = tir.Var("x", "float32")
+y = tir.Var("y", "float32")
+
+
+class BasicVisitor(ExprVisitor):
+    """Default ExprVisitor"""
+
+
+class ASTLog:
+    """Helper class to log AST"""
+
+    def __init__(self) -> None:
+        self.log = []
+        self.indent = "\t"
+        self.level = 0
+
+    def push_scope(self):
+        self.level += 1
+
+    def pop_scope(self):
+        self.level -= 1
+
+    def add(self, s: str):
+        self.log.append(self.indent * self.level + s)
+
+    def __str__(self) -> str:
+        return "\n".join(self.log)
+
+
+class ASTPrinter(ExprVisitor):
+    """Print TIR AST in structured format."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.log = ASTLog()
+
+    def visit_var_(self, op: Var) -> None:
+        self.log.add("Var")
+
+    def visit_buffer_load_(self, op: BufferLoad) -> None:
+        self.log.add("BufferLoad")
+        self.log.push_scope()
+        for idx in op.indices:
+            self.visit_expr(idx)
+        self.log.pop_scope()
+
+    def visit_producer_load_(self, op: ProducerLoad) -> None:
+        self.log.add("ProducerLoad")
+        self.log.push_scope()
+        for idx in op.indices:
+            self.visit_expr(idx)
+        self.log.pop_scope()
+
+    def visit_let_(self, op: Let) -> None:
+        self.log.add("Let")
+        self.log.push_scope()
+        self.visit_expr(op.var)
+        self.visit_expr(op.value)
+        self.visit_expr(op.body)
+        self.log.pop_scope()
+
+    def visit_call_(self, op: Call) -> None:
+        self.log.add("Call")
+        self.log.push_scope()
+        if isinstance(op.op, Op):
+            self.log.add("Op")
+        else:
+            self.visit_expr(op.op)
+        for arg in op.args:
+            self.visit_expr(arg)
+        self.log.pop_scope()
+
+    def visit_add_(self, op: Add) -> None:
+        self.log.add("Add")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_sub_(self, op: Sub) -> None:
+        self.log.add("Sub")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_mul_(self, op: Mul) -> None:
+        self.log.add("Mul")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_div_(self, op: Div) -> None:
+        self.log.add("Div")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_mod_(self, op: Mod) -> None:
+        self.log.add("Mod")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_floordiv_(self, op: FloorDiv) -> None:
+        self.log.add("FloorDiv")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_floormod_(self, op: FloorMod) -> None:
+        self.log.add("FloorMod")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_min_(self, op: Min) -> None:
+        self.log.add("Min")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_max_(self, op: Max) -> None:
+        self.log.add("Max")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_eq_(self, op: EQ) -> None:
+        self.log.add("EQ")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_ne_(self, op: NE) -> None:
+        self.log.add("NE")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_lt_(self, op: LT) -> None:
+        self.log.add("LT")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_le_(self, op: LE) -> None:
+        self.log.add("LE")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_gt_(self, op: GT) -> None:
+        self.log.add("GT")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_ge_(self, op: GE) -> None:
+        self.log.add("GE")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_and_(self, op: And) -> None:
+        self.log.add("And")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_or_(self, op: Or) -> None:
+        self.log.add("Or")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_reduce_(self, op: Reduce) -> None:
+        self.log.add("Reduce")
+        self.log.push_scope()
+        for source in op.source:
+            self.visit_expr(source)
+        for axis in op.axis:
+            self.visit_expr(axis.var)
+        self.visit_expr(op.condition)
+        self.log.pop_scope()
+
+    def visit_cast_(self, op: Cast) -> None:
+        self.log.add("Cast")
+        self.log.push_scope()
+        self.visit_expr(op.value)
+        self.log.pop_scope()
+
+    def visit_not_(self, op: Not) -> None:
+        self.log.add("Not")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.log.pop_scope()
+
+    def visit_select_(self, op: Select) -> None:
+        self.log.add("Select")
+        self.log.push_scope()
+        self.visit_expr(op.condition)
+        self.visit_expr(op.true_value)
+        self.visit_expr(op.false_value)
+        self.log.pop_scope()
+
+    def visit_ramp_(self, op: Ramp) -> None:
+        self.log.add("Ramp")
+        self.log.push_scope()
+        self.visit_expr(op.base)
+        self.visit_expr(op.stride)
+        self.visit_expr(op.lanes)
+        self.log.pop_scope()
+
+    def visit_broadcast_(self, op: Broadcast) -> None:
+        self.log.add("Broadcast")
+        self.log.push_scope()
+        self.visit_expr(op.value)
+        self.visit_expr(op.lanes)
+        self.log.pop_scope()
+
+    def visit_shuffle_(self, op: Shuffle) -> None:
+        self.log.add("Shuffle")
+        self.log.push_scope()
+        for vec in op.vectors:
+            self.visit_expr(vec)
+        for idx in op.indices:
+            self.visit_expr(idx)
+        self.log.pop_scope()
+
+    def visit_int_imm_(self, op: IntImm) -> None:
+        self.log.add("IntImm")
+
+    def visit_float_imm_(self, op: FloatImm) -> None:
+        self.log.add("FloatImm")
+
+    def visit_string_imm_(self, op: StringImm) -> None:
+        self.log.add("StringImm")
+
+
+class BasicMutator(ExprMutator):
+    """Default ExprMutator"""
+
+
+class ASTPostPrinterMutator(ExprMutator):
+    """Print TIR AST in the post order format."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.log = ASTLog()
+
+    def visit_var_(self, op: Var) -> tir.Expr:
+        result = super().visit_var_(op)
+        self.log.add("Var")
+        return result
+
+    def visit_buffer_load_(self, op: BufferLoad) -> tir.Expr:
+        result = super().visit_buffer_load_(op)
+        self.log.add("BufferLoad")
+        return result
+
+    def visit_producer_load_(self, op: ProducerLoad) -> tir.Expr:
+        result = super().visit_producer_load_(op)
+        self.log.add("ProducerLoad")
+        return result
+
+    def visit_let_(self, op: Let) -> tir.Expr:
+        result = super().visit_let_(op)
+        self.log.add("Let")
+        return result
+
+    def visit_call_(self, op: Call) -> tir.Expr:
+        result = super().visit_call_(op)
+        self.log.add("Call")
+        return result
+
+    def visit_add_(self, op: Add) -> tir.Expr:
+        result = super().visit_add_(op)
+        self.log.add("Add")
+        return result
+
+    def visit_sub_(self, op: Sub) -> tir.Expr:
+        result = super().visit_sub_(op)
+        self.log.add("Sub")
+        return result
+
+    def visit_mul_(self, op: Mul) -> tir.Expr:
+        result = super().visit_mul_(op)
+        self.log.add("Mul")
+        return result
+
+    def visit_div_(self, op: Div) -> tir.Expr:
+        result = super().visit_div_(op)
+        self.log.add("Div")
+        return result
+
+    def visit_mod_(self, op: Mod) -> tir.Expr:
+        result = super().visit_mod_(op)
+        self.log.add("Mod")
+        return result
+
+    def visit_floordiv_(self, op: FloorDiv) -> tir.Expr:
+        result = super().visit_floordiv_(op)
+        self.log.add("FloorDiv")
+        return result
+
+    def visit_floormod_(self, op: FloorMod) -> tir.Expr:
+        result = super().visit_floormod_(op)
+        self.log.add("FloorMod")
+        return result
+
+    def visit_min_(self, op: Min) -> tir.Expr:
+        result = super().visit_min_(op)
+        self.log.add("Min")
+        return result
+
+    def visit_max_(self, op: Max) -> tir.Expr:
+        result = super().visit_max_(op)
+        self.log.add("Max")
+        return result
+
+    def visit_eq_(self, op: EQ) -> tir.Expr:
+        result = super().visit_eq_(op)
+        self.log.add("EQ")
+        return result
+
+    def visit_ne_(self, op: NE) -> tir.Expr:
+        result = super().visit_ne_(op)
+        self.log.add("NE")
+        return result
+
+    def visit_lt_(self, op: LT) -> tir.Expr:
+        result = super().visit_lt_(op)
+        self.log.add("LT")
+        return result
+
+    def visit_le_(self, op: LE) -> tir.Expr:
+        result = super().visit_le_(op)
+        self.log.add("LE")
+        return result
+
+    def visit_gt_(self, op: GT) -> tir.Expr:
+        result = super().visit_gt_(op)
+        self.log.add("GT")
+        return result
+
+    def visit_ge_(self, op: GE) -> tir.Expr:
+        result = super().visit_ge_(op)
+        self.log.add("GE")
+        return result
+
+    def visit_and_(self, op: And) -> tir.Expr:
+        result = super().visit_and_(op)
+        self.log.add("And")
+        return result
+
+    def visit_or_(self, op: Or) -> tir.Expr:
+        result = super().visit_or_(op)
+        self.log.add("Or")
+        return result
+
+    def visit_reduce_(self, op: Reduce) -> tir.Expr:
+        result = super().visit_reduce_(op)
+        self.log.add("Reduce")
+        return result
+
+    def visit_cast_(self, op: Cast) -> tir.Expr:
+        result = super().visit_cast_(op)
+        self.log.add("Cast")
+        return result
+
+    def visit_not_(self, op: Not) -> tir.Expr:
+        result = super().visit_not_(op)
+        self.log.add("Not")
+        return result
+
+    def visit_select_(self, op: Select) -> tir.Expr:
+        result = super().visit_select_(op)
+        self.log.add("Select")
+        return result
+
+    def visit_ramp_(self, op: Ramp) -> tir.Expr:
+        result = super().visit_ramp_(op)
+        self.log.add("Ramp")
+        return result
+
+    def visit_broadcast_(self, op: Broadcast) -> tir.Expr:
+        result = super().visit_broadcast_(op)
+        self.log.add("Broadcast")
+        return result
+
+    def visit_shuffle_(self, op: Shuffle) -> tir.Expr:
+        result = super().visit_shuffle_(op)
+        self.log.add("Shuffle")
+        return result
+
+    def visit_int_imm_(self, op: IntImm) -> tir.Expr:
+        result = super().visit_int_imm_(op)
+        self.log.add("IntImm")
+        return result
+
+    def visit_float_imm_(self, op: FloatImm) -> tir.Expr:
+        result = super().visit_float_imm_(op)
+        self.log.add("FloatImm")
+        return result
+
+    def visit_string_imm_(self, op: StringImm) -> tir.Expr:
+        result = super().visit_string_imm_(op)
+        self.log.add("StringImm")
+        return result
+
+
+def basic_check(expr, visitor_str, mutator_str):
+    """Helper function to check visitor and mutator on an expression"""
+
+    # Check visitor
+    basic_visitor = BasicVisitor()
+    basic_visitor.visit_expr(expr)
+    # Check AST printer visitor
+    log_visitor = ASTPrinter()
+    log_visitor.visit_expr(expr)
+    assert str(log_visitor.log) == visitor_str
+
+    # Check basic mutator
+    basic_mutator = BasicMutator()
+    mutated_expr = basic_mutator.visit_expr(expr)
+    assert_structural_equal(mutated_expr, expr)
+
+    # Check post-order printer mutator
+    post_log_mutator = ASTPostPrinterMutator()
+    mutated_expr = post_log_mutator.visit_expr(expr)
+    assert_structural_equal(mutated_expr, expr)
+    assert str(post_log_mutator.log) == mutator_str
+
+
+def test_var():
+    basic_check(n, "Var", "Var")
+
+
+def test_int_imm():
+    basic_check(tir.IntImm("int32", 10), "IntImm", "IntImm")
+
+
+def test_float_imm():
+    basic_check(tir.FloatImm("float32", 1.5), "FloatImm", "FloatImm")
+
+
+def test_string_imm():
+    basic_check(tir.StringImm("hello"), "StringImm", "StringImm")
+
+
+def test_add():
+    add_node = tir.Add(n, m)
+    basic_check(add_node, "\n".join(["Add", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Add"]))
+
+
+def test_sub():
+    sub_node = tir.Sub(n, m)
+    basic_check(sub_node, "\n".join(["Sub", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Sub"]))
+
+
+def test_mul():
+    mul_node = tir.Mul(n, m)
+    basic_check(mul_node, "\n".join(["Mul", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Mul"]))
+
+
+def test_div():
+    div_node = tir.Div(n, m)
+    basic_check(div_node, "\n".join(["Div", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Div"]))
+
+
+def test_floor_div():
+    floor_div_node = tir.FloorDiv(n, m)
+    basic_check(
+        floor_div_node,
+        "\n".join(["FloorDiv", "\tVar", "\tVar"]),
+        "\n".join(["Var", "Var", "FloorDiv"]),
+    )
+
+
+def test_floor_mod():
+    floor_mod_node = tir.FloorMod(n, m)
+    basic_check(
+        floor_mod_node,
+        "\n".join(["FloorMod", "\tVar", "\tVar"]),
+        "\n".join(["Var", "Var", "FloorMod"]),
+    )
+
+
+def test_min():
+    min_node = tir.Min(n, m)
+    basic_check(min_node, "\n".join(["Min", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Min"]))
+
+
+def test_max():
+    max_node = tir.Max(n, m)
+    basic_check(max_node, "\n".join(["Max", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Max"]))
+
+
+def test_eq():
+    eq_node = tir.EQ(n, m)
+    basic_check(eq_node, "\n".join(["EQ", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "EQ"]))
+
+
+def test_ne():
+    ne_node = tir.NE(n, m)
+    basic_check(ne_node, "\n".join(["NE", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "NE"]))
+
+
+def test_lt():
+    lt_node = tir.LT(n, m)
+    basic_check(lt_node, "\n".join(["LT", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "LT"]))
+
+
+def test_le():
+    le_node = tir.LE(n, m)
+    basic_check(le_node, "\n".join(["LE", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "LE"]))
+
+
+def test_gt():
+    gt_node = tir.GT(n, m)
+    basic_check(gt_node, "\n".join(["GT", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "GT"]))
+
+
+def test_ge():
+    ge_node = tir.GE(n, m)
+    basic_check(ge_node, "\n".join(["GE", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "GE"]))
+
+
+def test_and():
+    and_node = tir.And(tir.EQ(n, m), tir.LT(n, 10))
+    basic_check(
+        and_node,
+        "\n".join(["And", "\tEQ", "\t\tVar", "\t\tVar", "\tLT", "\t\tVar", "\t\tIntImm"]),
+        "\n".join(["Var", "Var", "EQ", "Var", "IntImm", "LT", "And"]),
+    )
+
+
+def test_or():
+    or_node = tir.Or(tir.EQ(n, m), tir.LT(n, 10))
+    basic_check(
+        or_node,
+        "\n".join(["Or", "\tEQ", "\t\tVar", "\t\tVar", "\tLT", "\t\tVar", "\t\tIntImm"]),
+        "\n".join(["Var", "Var", "EQ", "Var", "IntImm", "LT", "Or"]),
+    )
+
+
+def test_not():
+    not_node = tir.Not(tir.EQ(n, m))
+    basic_check(
+        not_node,
+        "\n".join(["Not", "\tEQ", "\t\tVar", "\t\tVar"]),
+        "\n".join(["Var", "Var", "EQ", "Not"]),
+    )
+
+
+def test_select():
+    select_node = tir.Select(tir.EQ(n, m), n, m)
+    basic_check(
+        select_node,
+        "\n".join(["Select", "\tEQ", "\t\tVar", "\t\tVar", "\tVar", "\tVar"]),
+        "\n".join(["Var", "Var", "EQ", "Var", "Var", "Select"]),
+    )
+
+
+def test_cast():
+    cast_node = tir.Cast("float32", n)
+    basic_check(cast_node, "\n".join(["Cast", "\tVar"]), "\n".join(["Var", "Cast"]))
+
+
+def test_let():
+    let_node = tir.Let(n, tir.IntImm("int32", 10), n + 1)
+    basic_check(
+        let_node,
+        "\n".join(["Let", "\tVar", "\tIntImm", "\tAdd", "\t\tVar", "\t\tIntImm"]),
+        "\n".join(["Var", "IntImm", "Var", "IntImm", "Add", "Let"]),
+    )
+
+
+def test_ramp():
+    ramp_node = tir.Ramp(n, 1, 4)
+    basic_check(
+        ramp_node,
+        "\n".join(["Ramp", "\tVar", "\tIntImm", "\tIntImm"]),
+        "\n".join(["Var", "IntImm", "IntImm", "Ramp"]),
+    )
+
+
+def test_broadcast():
+    broadcast_node = tir.Broadcast(n, 4)
+    basic_check(
+        broadcast_node,
+        "\n".join(["Broadcast", "\tVar", "\tIntImm"]),
+        "\n".join(["Var", "IntImm", "Broadcast"]),
+    )
+
+
+def test_inherit():
+    # The internal class is not instantiated.
+    class InternalVisitor(ExprVisitor):
+        def __init__(self) -> None:
+            super().__init__()
+            self.log = ASTLog()
+
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("InternalAdd")
+            self.log.push_scope()
+            self.visit_expr(op.a)
+            self.visit_expr(op.b)
+            self.log.pop_scope()
+
+        def visit_var_(self, op: Var) -> None:
+            self.log.add("InternalVar")
+
+    class LeafVisitor(InternalVisitor):
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("LeafAdd")
+            self.log.push_scope()
+            self.visit_expr(op.a)
+            self.visit_expr(op.b)
+            self.log.pop_scope()
+
+    add_node = tir.Add(n, m)
+    lv = LeafVisitor()
+    lv.visit_expr(add_node)
+    assert str(lv.log) == "\n".join(["LeafAdd", "\tInternalVar", "\tInternalVar"])
+
+
+def test_inherit_with_cls():
+    class InternalVisitor(ExprVisitor):
+        def __init__(self) -> None:
+            super().__init__()
+            self.log = ASTLog()
+
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("InternalAdd")
+            self.log.push_scope()
+            self.visit_expr(op.a)
+            self.visit_expr(op.b)
+            self.log.pop_scope()
+
+        def visit_var_(self, op: Var) -> None:
+            self.log.add("InternalVar")
+
+    class LeafVisitor(InternalVisitor):
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("LeafAdd")
+            self.log.push_scope()
+            self.visit_expr(op.a)
+            self.visit_expr(op.b)
+            self.log.pop_scope()
+
+    add_node = tir.Add(n, m)
+    iv = InternalVisitor()
+    iv.visit_expr(add_node)
+    assert str(iv.log) == "\n".join(["InternalAdd", "\tInternalVar", "\tInternalVar"])
+
+    lv = LeafVisitor()
+    lv.visit_expr(add_node)
+    assert str(lv.log) == "\n".join(["LeafAdd", "\tInternalVar", "\tInternalVar"])
+
+
+def test_call_visitor_super():
+    class InternalVisitor(ExprVisitor):
+        def __init__(self) -> None:
+            super().__init__()
+            self.log = ASTLog()
+
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("InternalAdd")
+            super().visit_add_(op)  # call ExprVisitor.visit_add_
+
+        def visit_var_(self, op: Var) -> None:
+            self.log.add("InternalVar")
+
+        def visit_int_imm_(self, op: IntImm) -> None:
+            self.log.add("InternalIntImm")
+
+    class LeafVisitor(InternalVisitor):
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("LeafAdd")
+            super().visit_add_(op)  # call InternalVisitor.visit_add_
+
+    add_node = tir.Add(n, tir.IntImm("int32", 10))
+    iv = InternalVisitor()
+    iv.visit_expr(add_node)
+    assert str(iv.log) == "\n".join(["InternalAdd", "InternalVar", "InternalIntImm"])
+
+    lv = LeafVisitor()
+    lv.visit_expr(add_node)
+    assert str(lv.log) == "\n".join(["LeafAdd", "InternalAdd", "InternalVar", "InternalIntImm"])
+
+
+def test_call_mutator_super():
+    class InternalMutator(ExprMutator):
+        def __init__(self) -> None:
+            super().__init__()
+            self.log = ASTLog()
+
+        def visit_add_(self, op: Add) -> tir.Expr:
+            self.log.add("InternalAdd")
+            return super().visit_add_(op)  # call ExprMutator.visit_add_
+
+        def visit_var_(self, op: Var) -> tir.Expr:
+            self.log.add("InternalVar")
+            return super().visit_var_(op)  # call ExprMutator.visit_var_
+
+        def visit_int_imm_(self, op: IntImm) -> tir.Expr:
+            self.log.add("InternalIntImm")
+            return super().visit_int_imm_(op)  # call ExprMutator.visit_int_imm_
+
+    class LeafMutator(InternalMutator):
+        def visit_add_(self, op: Add) -> tir.Expr:
+            self.log.add("LeafAdd")
+            return super().visit_add_(op)  # call InternalMutator.visit_add_
+
+    add_node = tir.Add(n, tir.IntImm("int32", 10))
+    im = InternalMutator()
+    im.visit_expr(add_node)
+    assert str(im.log) == "\n".join(["InternalAdd", "InternalVar", "InternalIntImm"])
+
+    lm = LeafMutator()
+    lm.visit_expr(add_node)
+    assert str(lm.log) == "\n".join(["LeafAdd", "InternalAdd", "InternalVar", "InternalIntImm"])
+
+
+def test_var_mutation():
+    """Test mutating variables in a TIR expression"""
+
+    class VarMutator(ExprMutator):
+        def __init__(self, var_map):
+            super().__init__()
+            self.var_map = var_map
+
+        def visit_var_(self, op: Var) -> tir.Expr:
+            if op.name in self.var_map:
+                return self.var_map[op.name]
+            return op
+
+    # Create a simple expression
+    expr = n + m
+
+    # Create a mutator that replaces 'n' with a constant
+    var_map = {"n": tir.IntImm("int32", 42)}
+    mutator = VarMutator(var_map)
+    result = mutator.visit_expr(expr)
+
+    # The result should be 42 + m
+    expected = tir.Add(tir.IntImm("int32", 42), m)
+    assert_structural_equal(result, expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/transform/test_transform_naive_allocator.py

@@ -0,0 +1,177 @@
+# 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.ir import assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.trn.transform import TrnNaiveAllocator
+
+
+def test_one_alloc():
+    src_shape = [128, 512]
+    src_layout = TileLayout(S[(128, 512) : (512, 1)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(A_sbuf, A)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle) -> None:
+        T.func_attr({"global_symbol": "copy"})
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout, allocated_addr=[0])  # noqa: E501
+        Tx.copy(A_sbuf, A)
+        # fmt: on
+
+    mod = tvm.IRModule({"copy": copy})
+    mod = TrnNaiveAllocator()(mod)
+    assert_structural_equal(mod["copy"], expected)
+
+
+def test_two_alloc():
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF")
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF")
+        Tx.copy(B_sbuf[0:256, :], A_sbuf)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle) -> None:
+        T.func_attr({"global_symbol": "copy"})
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[0])  # noqa: E501
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[2*512*4])  # noqa: E501
+        Tx.copy(B_sbuf[0:256, :], A_sbuf)
+        # fmt: on
+
+    mod = tvm.IRModule({"copy": copy})
+    mod = TrnNaiveAllocator()(mod)
+    assert_structural_equal(mod["copy"], expected)
+
+
+def test_existing_alloc():
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF")
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[1])  # noqa: E501
+        Tx.copy(B_sbuf[0:256, :], A_sbuf)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle) -> None:
+        T.func_attr({"global_symbol": "copy"})
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[4*512*4+1])  # noqa: E501
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[1])  # noqa: E501
+        Tx.copy(B_sbuf[0:256, :], A_sbuf)
+        # fmt: on
+
+    mod = tvm.IRModule({"copy": copy})
+    mod = TrnNaiveAllocator()(mod)
+    assert_structural_equal(mod["copy"], expected)
+
+
+def test_workspace():
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF")
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF")
+        C_sbuf = T.alloc_buffer([128, 1024], "float32", scope="trn.sbuf")
+        Tx.copy(B_sbuf[0:256, :], A_sbuf, workspace={"C": C_sbuf})
+
+    @T.prim_func
+    def expected(A_ptr: T.handle) -> None:
+        T.func_attr({"global_symbol": "copy"})
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[0])  # noqa: E501
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[2*512*4])  # noqa: E501
+        C_sbuf = T.alloc_buffer([128, 1024], "float32", scope="trn.sbuf", allocated_addr=[2*512*4+4*512*4])  # noqa: E501
+        Tx.copy(B_sbuf[0:256, :], A_sbuf, workspace={"C": C_sbuf})
+        # fmt: on
+
+    mod = tvm.IRModule({"copy": copy})
+    mod = TrnNaiveAllocator()(mod)
+    assert_structural_equal(mod["copy"], expected)
+
+
+def test_other_scope_alloc():
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF")
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF")
+        C_sbuf = T.alloc_buffer([8, 128, 512], "float32", scope="global")
+        Tx.copy(B_sbuf[0:256, :], A_sbuf, workspace={"C": C_sbuf})
+
+    @T.prim_func
+    def expected(A_ptr: T.handle) -> None:
+        T.func_attr({"global_symbol": "copy"})
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[0])  # noqa: E501
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[2*512*4])  # noqa: E501
+        C_sbuf = T.alloc_buffer([8, 128, 512], "float32", scope="global")
+        Tx.copy(B_sbuf[0:256, :], A_sbuf, workspace={"C": C_sbuf})
+        # fmt: on
+
+    mod = tvm.IRModule({"copy": copy})
+    mod = TrnNaiveAllocator()(mod)
+    assert_structural_equal(mod["copy"], expected)
+
+
+def test_buffer_views():
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF")
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF")
+        B_view = B_sbuf.view(2, 256, 512)
+        Tx.copy(B_view[0], A_sbuf)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle) -> None:
+        T.func_attr({"global_symbol": "copy"})
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[0])  # noqa: E501
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[2*512*4])  # noqa: E501
+        B_view = B_sbuf.view(2, 256, 512)
+        Tx.copy(B_view[0], A_sbuf)
+        # fmt: on
+
+    mod = tvm.IRModule({"copy": copy})
+    mod = TrnNaiveAllocator()(mod)
+    assert_structural_equal(mod["copy"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/transform/test_transform_static_horizontal_fusion.py

@@ -0,0 +1,20 @@
+# 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.
+
+
+SM_CNT = 148
+NUM_THREADS = 256