apache--tvm/tests/python/tirx-base/test_tir_expr_functor.py

# 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()