chore: import upstream snapshot with attribution

python/tvm/tirx/script/parser/__init__.py

@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: RUF005
+"""The tirx parser"""
+
+from typing import TYPE_CHECKING
+
+from tvm.tirx.script.builder import *  # pylint: disable=redefined-builtin
+from tvm.tirx.script.builder import ir as _tir
+
+from . import operation as _operation
+from . import parser as _parser
+from .entry import Buffer, Ptr, constexpr
+
+if TYPE_CHECKING:
+    # pylint: disable=invalid-name
+    # Define prim_func and make it type check as static method
+    # so most tvmscript won't trigger pylint error here.
+    prim_func = staticmethod
+    jit = staticmethod
+else:
+    from .entry import inline, jit, macro, prim_func
+
+__all__ = _tir.__all__ + [
+    "Buffer",
+    "Ptr",
+    "SMEMPool",
+    "TMEMPool",
+    "TMEMStages",
+    "bool",
+    "constexpr",
+    "inline",
+    "jit",
+    "macro",
+    "prim_func",
+]

python/tvm/tirx/script/parser/entry.py

@@ -0,0 +1,520 @@
+# 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.
+"""The entry point of TVM parser for tirx."""
+
+import inspect
+from collections.abc import Callable
+from typing import Any
+
+from tvm.ir.base import deprecated
+from tvm.script.parser._core import parse, scan_macro, utils
+from tvm.script.parser.core.parser import Parser, ScriptMacro, VarTable
+from tvm.tirx import Buffer, PrimFunc
+from tvm.tirx.script.builder import block_name_suffix_context, buffer, ptr
+
+
+def prim_func(
+    func: Callable | None = None,
+    private: bool = False,
+    check_well_formed=True,
+    s_tir: bool = False,
+    persistent: bool = False,
+) -> PrimFunc | Callable:
+    """The parsing method for tirx prim func, by using `@prim_func` as decorator.
+
+    Parameters
+    ----------
+    func : Callable
+        The function to be parsed as prim func.
+        (Listed as optional to allow the decorator to be used
+        without arguments, like `@prim_func`,
+        or with an argument, `@prim_func(private=True)`)
+
+    private : bool, optional
+        Whether the function should be treated as private.
+        A private function has no global symbol attribute;
+        if the function is not private, it will have a global symbol
+        matching the function name.
+
+    Returns
+    -------
+    res : Union[PrimFunc, Callable]
+        The parsed tirx prim func.
+    """
+    # pylint: disable=unused-argument
+    # (private will be used in the parser, but not immediately)
+
+    # need to capture this var outside the wrapper because the wrapper
+    # adds to the stack
+    outer_stack = inspect.stack()
+
+    def decorator_wrapper(func):
+        if not inspect.isfunction(func):
+            raise TypeError(f"Expect a function, but got: {func}")
+        if utils.is_defined_in_class(outer_stack, func):
+            return func
+        extra_vars = utils.inspect_function_capture(func)
+        utils.resolve_closure_vars(func, extra_vars, outer_stack)
+        f = parse(func, extra_vars, check_well_formed=check_well_formed, s_tir=s_tir)
+        setattr(f, "__name__", func.__name__)
+        return f
+
+    if func is not None:
+        # no optional args given => use wrapper directly
+        return decorator_wrapper(func)
+    else:
+        # if there is an optional arg given, return a new decorator
+        # that will then be invoked
+        setattr(decorator_wrapper, "dispatch_token", "tirx")
+        return decorator_wrapper
+
+
+setattr(prim_func, "dispatch_token", "tirx")
+
+
+class TIRInline(ScriptMacro):
+    """Specialization of ScriptMacro for TIR with Python LEGB scoping.
+
+    Two definition paths:
+    1. Outside @T.prim_func (standalone @T.inline): definition_depth is None,
+       closure_vars captured at definition time are used (module globals are
+       effectively late-bound since they don't change during parsing).
+    2. Inside @T.prim_func (inline def in parsed body): definition_depth is set
+       to the VarTable frame depth at definition time, and defining_var_table
+       stores a reference to the VarTable that was active. At call time,
+       defining_var_table.get_at_depth(definition_depth) reads current values
+       from the lexically enclosing frames.
+
+    Attributes
+    ----------
+    definition_depth : Optional[int]
+        VarTable frame depth at definition time, or None for outside-prim_func.
+    defining_var_table : Optional[VarTable]
+        Reference to the VarTable that was active at definition time.
+    call_count : int
+        Counter for unique block name suffixes.
+    """
+
+    def __init__(
+        self,
+        source,
+        closure_vars: dict[str, Any],
+        func: Callable,
+        definition_depth: int | None = None,
+        defining_var_table: VarTable | None = None,
+    ) -> None:
+        # hygienic=True for the base class (field kept for compat but not used in dispatch)
+        super().__init__(source, closure_vars, func, hygienic=True)
+        self.definition_depth = definition_depth
+        self.defining_var_table = defining_var_table
+        self.call_count = 0
+
+    def parse_macro(self, parser: Parser) -> None:
+        macro_def = self.get_macro_def()
+        suffix = f"_{self.call_count}" if self.call_count > 0 else ""
+        self.call_count += 1
+        with block_name_suffix_context(suffix):
+            parser.visit_body(macro_def.body)
+
+    def __call__(self, *args, **kwargs):
+        param_binding = inspect.signature(self.func).bind(*args, **kwargs)
+        param_binding.apply_defaults()
+        local_vars = param_binding.arguments
+        parser = self._find_parser_def()
+
+        with parser.with_diag_source(self.source):
+            if self.defining_var_table is not None:
+                # Inside-prim_func path: LEGB late binding from the defining scope
+                enclosing_vars = self.defining_var_table.get_at_depth(self.definition_depth)
+            else:
+                # Outside-prim_func path: use captured closure vars
+                enclosing_vars = self.closure_vars
+
+            saved_var_table = parser.var_table
+            parser.var_table = VarTable()
+
+            with parser.var_table.with_frame():
+                for k, v in enclosing_vars.items():
+                    parser.var_table.add(k, v)
+                with parser.var_table.with_frame():
+                    for k, v in local_vars.items():
+                        parser.var_table.add(k, v)
+
+                    parse_result = self.parse_macro(parser)
+
+            parser.var_table = saved_var_table
+
+        return parse_result
+
+
+def inline(*args, definition_depth: int | None = None, defining_var_table=None) -> Callable:
+    """Decorator for inline function definitions with Python LEGB scoping.
+
+    @T.inline follows Python's lexical scoping with late binding:
+    - At definition time, record which scopes are visible.
+    - At call time, read current values from those scopes.
+
+    Example::
+
+        import tvm
+        from tvm.script import tirx as T
+
+        x_value = 128
+
+        @T.inline
+        def capture(A, B):
+            B[()] = A[x_value]          # x_value resolved from enclosing scope
+
+        @T.prim_func(s_tir=True)
+        def use(A: T.Buffer((1024,), "int32"), B: T.Buffer((), "int32")) -> None:
+            capture(A, B)               # Produces B[()] = A[128]
+    """
+
+    def _decorator(func: Callable) -> Callable:
+        source, closure_vars = scan_macro(func, utils.inspect_function_capture(func))
+        obj = TIRInline(
+            source,
+            closure_vars,
+            func,
+            definition_depth=definition_depth,
+            defining_var_table=defining_var_table,
+        )
+
+        def wrapper(*args, **kwargs):
+            return obj(*args, **kwargs)
+
+        return wrapper
+
+    if len(args) == 0:
+        setattr(_decorator, "dispatch_token", "tir.inline")
+        return _decorator
+    if len(args) == 1 and inspect.isfunction(args[0]):
+        return _decorator(args[0])
+
+    raise ValueError("Invalid use of T.inline. Usage: @T.inline or @T.inline()")
+
+
+setattr(inline, "dispatch_token", "tir.inline")
+
+
+class TIRJit:
+    """Top-level kernel decorator with constexpr params + ``.specialize()``.
+
+    Parses the function body lazily: parsing is deferred until ``.specialize()``
+    supplies concrete values for the params annotated as ``T.constexpr``. The
+    return type of ``.specialize()`` is a ``tvm.tirx.PrimFunc``, identical in
+    type to what ``@T.prim_func`` produces today.
+
+    Constexpr params are removed from the resulting PrimFunc's parameter list;
+    their values are baked into the IR (e.g. into ``T.Buffer((M, K), ...)``
+    shape annotations and into the body).
+    """
+
+    def __init__(
+        self,
+        func: Callable,
+        check_well_formed: bool = True,
+        is_stir: bool = False,
+        persistent: bool = False,
+        private: bool = False,
+    ) -> None:
+        self.func = func
+        self.check_well_formed = check_well_formed
+        self.is_stir = is_stir
+        self.persistent = persistent  # pylint: disable=unused-private-member
+        self.private = private  # pylint: disable=unused-private-member
+        # Resolved closure vars (computed once; the function itself is the
+        # capture point, so this never changes between specializations).
+        self._closure_vars: dict[str, Any] = utils.inspect_function_capture(func)
+        # Detect which params are marked T.constexpr. With PEP 563
+        # (``from __future__ import annotations``), each annotation is a
+        # string; we eval them one-by-one so a constexpr probe is not
+        # blocked by sibling annotations that reference yet-undefined names
+        # (e.g. ``A: T.Buffer((N,), ...)`` referencing constexpr ``N``).
+        raw_anns = getattr(func, "__annotations__", {}) or {}
+        eval_globals = {**func.__globals__, **self._closure_vars}
+        sig = inspect.signature(func)
+        constexpr_names: set[str] = set()
+        constexpr_defaults: dict[str, Any] = {}
+        for name, param in sig.parameters.items():
+            ann = raw_anns.get(name)
+            if isinstance(ann, str):
+                try:
+                    ann = eval(ann, eval_globals)  # pylint: disable=eval-used
+                except Exception:  # pylint: disable=broad-except
+                    ann = None
+            if ann is constexpr:
+                constexpr_names.add(name)
+                if param.default is not inspect.Parameter.empty:
+                    constexpr_defaults[name] = param.default
+        self.constexpr_names: frozenset[str] = frozenset(constexpr_names)
+        self.constexpr_defaults: dict[str, Any] = constexpr_defaults
+        self._cache: dict[tuple, PrimFunc] = {}
+
+    def specialize(self, **constexpr_kwargs) -> PrimFunc:
+        """Build a concrete PrimFunc by binding the constexpr params.
+
+        Parameters
+        ----------
+        **constexpr_kwargs
+            One value per ``T.constexpr``-annotated parameter. All such
+            parameters must be supplied; passing names that are not
+            constexpr-annotated is an error.
+
+        Returns
+        -------
+        PrimFunc
+            A concrete TIRx PrimFunc, identical in type to the output of
+            ``@T.prim_func``.
+        """
+        extra = constexpr_kwargs.keys() - self.constexpr_names
+        if extra:
+            raise TypeError(
+                f"{self.func.__name__}.specialize() got unexpected arg(s): "
+                f"{sorted(extra)} (constexpr params are: {sorted(self.constexpr_names)})"
+            )
+        effective = {**self.constexpr_defaults, **constexpr_kwargs}
+        missing = self.constexpr_names - effective.keys()
+        if missing:
+            raise TypeError(
+                f"{self.func.__name__}.specialize() missing constexpr arg(s) "
+                f"(no default provided): {sorted(missing)}"
+            )
+
+        try:
+            cache_key = tuple(sorted(effective.items()))
+            cached = self._cache.get(cache_key)
+        except TypeError as err:
+            raise TypeError(
+                f"{self.func.__name__}.specialize(): all constexpr values must "
+                f"be hashable (got: {effective!r})"
+            ) from err
+        if cached is not None:
+            return cached
+
+        extra_vars = {**self._closure_vars, **effective}
+        prim_func = parse(
+            self.func,
+            extra_vars,
+            check_well_formed=self.check_well_formed,
+            s_tir=self.is_stir,
+        )
+        setattr(prim_func, "__name__", self.func.__name__)
+        self._cache[cache_key] = prim_func
+        return prim_func
+
+
+def jit(
+    func: Callable | None = None,
+    private: bool = False,
+    check_well_formed: bool = True,
+    is_stir: bool = False,
+    persistent: bool = False,
+) -> "TIRJit | Callable":
+    """Decorator: capture the kernel and defer parsing until ``.specialize()``.
+
+    Use ``@T.jit`` (instead of ``@T.prim_func``) when the kernel takes
+    compile-time parameters annotated with ``T.constexpr``. The resulting
+    object exposes ``.specialize(**constexpr_kwargs)``, which returns a
+    ``tvm.tirx.PrimFunc``.
+
+    Example::
+
+        from tvm.script import tirx as T
+
+        @T.jit
+        def add(
+            A: T.Buffer((N,), "float32"),
+            B: T.Buffer((N,), "float32"),
+            *,
+            N: T.constexpr,
+        ):
+            ...
+
+        kernel = add.specialize(N=1024)  # returns a PrimFunc
+    """
+
+    def decorator_wrapper(func: Callable) -> TIRJit:
+        if not inspect.isfunction(func):
+            raise TypeError(f"Expect a function, but got: {func}")
+        return TIRJit(
+            func,
+            check_well_formed=check_well_formed,
+            is_stir=is_stir,
+            persistent=persistent,
+            private=private,
+        )
+
+    if func is not None:
+        return decorator_wrapper(func)
+    setattr(decorator_wrapper, "dispatch_token", "tirx")
+    return decorator_wrapper
+
+
+setattr(jit, "dispatch_token", "tirx")
+
+
+class TIRMacro(ScriptMacro):
+    """Specialization of the ScriptMacro class for TIR.
+
+    Apache-compatible hygienic macro. Distinct from ``TIRInline`` (which
+    uses Python LEGB late binding) so upstream code that relies on
+    capture-at-definition-time semantics keeps working.
+
+    Attributes
+    ----------
+    call_count : int
+        Counter for the number of times this macro has been invoked.
+        Used to generate unique block name suffixes.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.call_count = 0
+
+    def parse_macro(self, parser: Parser) -> None:
+        macro_def = self.get_macro_def()
+        suffix = f"_{self.call_count}" if self.call_count > 0 else ""
+        self.call_count += 1
+        with block_name_suffix_context(suffix):
+            parser.visit_body(macro_def.body)
+
+
+def macro(*args, hygienic: bool = True) -> Callable:
+    """Decorator for macro definitions with hygienic capture.
+
+    Parameters
+    ----------
+    hygienic: bool
+        Specifies whether the macro is hygienic or not. A hygienic macro
+        resolves symbols at definition time; a non-hygienic macro at use
+        time. Defaults to ``True``.
+    """
+
+    def _decorator(func: Callable) -> TIRMacro:
+        source, closure_vars = scan_macro(func, utils.inspect_function_capture(func))
+        obj = TIRMacro(source, closure_vars, func, hygienic)
+
+        def wrapper(*args, **kwargs):
+            return obj(*args, **kwargs)
+
+        return wrapper
+
+    if len(args) == 0:
+        return _decorator
+    if len(args) == 1 and inspect.isfunction(args[0]):
+        return _decorator(args[0])
+
+    raise ValueError("Invalid use of T.macro. Usage: @T.macro or @T.macro()")
+
+
+setattr(macro, "dispatch_token", "tir.macro")
+
+
+class BufferProxy:
+    """Buffer proxy class for constructing tirx buffer."""
+
+    def __or__(self, other):
+        """Support ``T.Buffer | None`` union syntax in annotations."""
+        return self
+
+    def __ror__(self, other):
+        """Support ``None | T.Buffer`` union syntax in annotations."""
+        return self
+
+    def __call__(
+        self,
+        shape,
+        dtype="float32",
+        data=None,
+        strides=None,
+        elem_offset=None,
+        byte_offset=None,
+        scope="global",
+        align=0,
+        offset_factor=0,
+        buffer_type="",
+        axis_separators=None,
+        layout="default",
+    ) -> Buffer:
+        return buffer(
+            shape,
+            dtype=dtype,
+            data=data,
+            strides=strides,
+            elem_offset=elem_offset,
+            byte_offset=byte_offset,
+            scope=scope,
+            align=align,
+            offset_factor=offset_factor,
+            buffer_type=buffer_type,
+            axis_separators=axis_separators,
+            layout=layout,
+        )
+
+    @deprecated("T.Buffer[...]", "T.Buffer(...)")
+    def __getitem__(self, keys) -> Buffer:
+        if not isinstance(keys, tuple):
+            return self(keys)
+        if len(keys) >= 2 and not isinstance(keys[1], str):
+            return self(keys)
+        return self(*keys)  # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+class PtrProxy:
+    """Ptr proxy class for constructing tirx pointer."""
+
+    def __or__(self, other):
+        """Support union syntax in annotations."""
+        return self
+
+    def __ror__(self, other):
+        """Support union syntax in annotations."""
+        return self
+
+    @deprecated("T.Ptr(...)", "T.handle(...)")
+    def __call__(self, dtype, storage_scope="global"):
+        if callable(dtype):
+            dtype = dtype().ty.dtype
+        return ptr(dtype, storage_scope)  # type: ignore[attr-defined] # pylint: disable=no-member
+
+    @deprecated("T.Ptr[...]", "T.handle(...)")
+    def __getitem__(self, keys):
+        if not isinstance(keys, tuple):
+            return self(keys)
+        return self(*keys)
+
+
+class _ConstexprProxy:
+    """Sentinel marker for compile-time (specialization-time) parameters.
+
+    Used as a parameter annotation in ``@T.jit`` decorated functions to mark
+    a parameter as constexpr — its value is supplied to ``.specialize(**kwargs)``
+    rather than at call time, and it is removed from the generated PrimFunc's
+    runtime parameter list.
+    """
+
+    def __or__(self, other):
+        return self
+
+    def __ror__(self, other):
+        return self
+
+
+Buffer = BufferProxy()  # pylint: disable=invalid-name
+Ptr = PtrProxy()  # pylint: disable=invalid-name
+constexpr = _ConstexprProxy()  # pylint: disable=invalid-name

python/tvm/tirx/script/parser/operation.py

@@ -0,0 +1,167 @@
+# 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.
+"""The tirx expression operation registration"""
+
+import tvm
+from tvm import tirx
+from tvm.ir import PrimType
+from tvm.runtime import DataTypeCode
+from tvm.script.parser._core import OpMethod, doc, register_op
+from tvm.tirx import IntImm
+from tvm.tirx.expr import FloatImm
+
+
+def _register_expr_op(ty: type):  # pylint: disable=invalid-name
+    ty._dispatch_type = ty  # pylint: disable=protected-access
+
+    def _expr_ty(expr):
+        ty = expr.ty if tvm.ir.is_prim_expr(expr) else None
+        if not isinstance(ty, PrimType):
+            ty = expr.expr_ty()
+        if not isinstance(ty, PrimType):
+            raise TypeError(f"Expected a PrimType expression, but got {ty}")
+        return ty
+
+    def _and(a, b):
+        if isinstance(a, bool):
+            a = IntImm("bool", a)
+        if isinstance(b, bool):
+            b = IntImm("bool", b)
+        if not _expr_ty(a).is_scalar() or not _expr_ty(b).is_scalar():
+            return a & b
+        else:
+            return tirx.And(a, b)
+
+    def _or(a, b):
+        if isinstance(a, bool):
+            a = IntImm("bool", a)
+        if isinstance(b, bool):
+            b = IntImm("bool", b)
+        if not _expr_ty(a).is_scalar() or not _expr_ty(b).is_scalar():
+            return a | b
+        else:
+            return tirx.Or(a, b)
+
+    def _get_type_str(ty: PrimType):
+        dtype_str = str(ty.dtype)
+        if ty.is_scalar():
+            return dtype_str
+        index = dtype_str.find("x")
+        return dtype_str[0:index]
+
+    def _auto_broadcast(a, b, op):
+        if isinstance(a, int):
+            if tvm.ir.is_prim_expr(b) or hasattr(b, "expr_ty"):
+                b_ty = _expr_ty(b)
+                if b_ty.matches_code(DataTypeCode.INT, DataTypeCode.UINT, DataTypeCode.BOOL):
+                    a = IntImm(_get_type_str(b_ty), a)
+                elif b_ty.matches_code(DataTypeCode.FLOAT):
+                    a = FloatImm(_get_type_str(b_ty), a)
+            elif isinstance(b, float):
+                a = FloatImm("float32", a)
+            else:
+                a = IntImm("int32", a)
+        elif isinstance(a, float):
+            b_ty = _expr_ty(b)
+            if b_ty.matches_code(DataTypeCode.FLOAT):
+                a = FloatImm(_get_type_str(b_ty), a)
+            else:
+                a = FloatImm("float32", a)
+
+        assert tvm.ir.is_prim_expr(a), "Operand should be a Expr."
+        if isinstance(b, int):
+            a_ty = _expr_ty(a)
+            if a_ty.matches_code(DataTypeCode.INT, DataTypeCode.UINT, DataTypeCode.BOOL):
+                b = IntImm(_get_type_str(a_ty), b)
+            elif a_ty.matches_code(DataTypeCode.FLOAT):
+                b = FloatImm(_get_type_str(a_ty), b)
+        elif isinstance(b, float):
+            b = FloatImm(_get_type_str(_expr_ty(a)), b)
+
+        a_ty = _expr_ty(a)
+        b_ty = _expr_ty(b)
+        if a_ty.dtype.lanes == b_ty.dtype.lanes:
+            return op(a, b)
+        elif a_ty.is_scalar() and a_ty.dtype.lanes != b_ty.dtype.lanes:
+            broadcast_a = tirx.Broadcast(a, b_ty.dtype.lanes)
+            return op(broadcast_a, b)
+        elif b_ty.is_scalar() and a_ty.dtype.lanes != b_ty.dtype.lanes:
+            broadcast_b = tirx.Broadcast(b, a_ty.dtype.lanes)
+            return op(a, broadcast_b)
+        else:
+            raise TypeError("do not know how to deal with it.")
+
+    def _eq(a, b):
+        return _auto_broadcast(a, b, tirx.EQ)
+
+    def _ne(a, b):
+        return _auto_broadcast(a, b, tirx.NE)
+
+    def _lt(a, b):
+        return _auto_broadcast(a, b, tirx.LT)
+
+    def _le(a, b):
+        return _auto_broadcast(a, b, tirx.LE)
+
+    def _gt(a, b):
+        return _auto_broadcast(a, b, tirx.GT)
+
+    def _ge(a, b):
+        return _auto_broadcast(a, b, tirx.GE)
+
+    def r(op: type, i: int, m: OpMethod):  # pylint: disable=invalid-name
+        register_op(ty, op, i)(m)
+
+    for i in [0, 1]:
+        # Case 1. binop
+        # doc.Add <-- is overloaded
+        # doc.Sub <-- is overloaded
+        # doc.Mult <-- is overloaded
+        # doc.Div <-- is overloaded
+        # doc.FloorDiv <-- is overloaded
+        # doc.Mod <-- is overloaded
+        # doc.LShift <-- is overloaded
+        # doc.RShift <-- is overloaded
+        # doc.BitOr <-- is overloaded
+        # doc.BitXor <-- is overloaded
+        # doc.BitAnd <-- is overloaded
+        # doc.MatMult <-- not implemented
+        # doc.Pow <-- not implemented
+        # Case 2. cmpop
+        r(doc.Eq, i, _eq)
+        r(doc.NotEq, i, _ne)
+        r(doc.Lt, i, _lt)
+        r(doc.LtE, i, _le)
+        r(doc.Gt, i, _gt)
+        r(doc.GtE, i, _ge)
+        # doc.Is <-- not implemented
+        # doc.IsNot <-- not implemented
+        # doc.In <-- not implemented
+        # doc.NotIn <-- not implemented
+        # Case 3. boolop
+        r(doc.And, i, _and)
+        r(doc.Or, i, _or)
+    for i in [0]:
+        #  Case 4. unaryop
+        # doc.Invert <-- is overloaded
+        r(doc.Not, i, tirx.Not)
+        # doc.UAdd <-- is overloaded
+        # doc.USub <-- is overloaded
+
+
+_register_expr_op(tirx.Expr)
+_register_expr_op(tirx.IterVar)

python/tvm/tirx/script/parser/parser.py

@@ -0,0 +1,914 @@
+# 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.
+"""The base parser for tirx"""
+
+import ast
+import contextlib
+from copy import deepcopy
+from functools import partial
+from typing import Any
+
+import tvm
+from tvm.ir import Expr, GlobalVar, PrimType
+from tvm.script.ir_builder import ir as I
+from tvm.script.ir_builder.base import IRBuilder
+from tvm.script.ir_builder.base import IRBuilderFrame as Frame
+from tvm.script.parser._core import Parser, dispatch, doc
+from tvm.script.parser.core.doc import from_doc
+from tvm.tirx import Buffer, IterVar, Layout, Var
+from tvm.tirx.script import builder as T
+from tvm.tirx.script.builder.ir import name_meta_class_value
+from tvm.tirx.stmt import BufferRegion
+
+from .entry import constexpr as _constexpr_sentinel
+from .entry import inline
+
+
+def slice_buffer_from_region(br: BufferRegion) -> Buffer:
+    """Create a matched DeclBuffer from a BufferRegion.
+
+    Slices the layout (if present) or computes elem_offset for the sub-region,
+    producing a DeclBuffer that views the same underlying data.
+    """
+    import functools  # pylint: disable=import-outside-toplevel
+
+    buf = br.buffer
+    region = br.region
+    new_shape = [r.extent for r in region]
+    sliced_layout = None
+    if buf.layout is not None:
+        range_pairs = [(r.min, r.min + r.extent) for r in region]
+        sliced_layout = buf.layout.slice(list(buf.shape), range_pairs)
+    if sliced_layout is not None:
+        return T.decl_buffer(
+            new_shape,
+            buf.dtype,
+            buf.data,
+            buf.strides,
+            buf.elem_offset,
+            None,
+            buf.scope(),
+            buf.data_alignment,
+            buf.offset_factor,
+            "",
+            buf.axis_separators,
+            sliced_layout,
+        )
+    # Fallback: compute elem_offset for default/no layout
+    strides = []
+    for i in range(len(buf.shape)):
+        stride = functools.reduce(
+            lambda x, y: x * y, buf.shape[i + 1 :], tvm.tirx.const(1, "int32")
+        )
+        strides.append(stride)
+    offset = tvm.tirx.const(0, "int32")
+    for i, r in enumerate(region):
+        offset = offset + r.min * strides[i]
+    new_elem_offset = buf.elem_offset + offset
+    return T.decl_buffer(
+        new_shape,
+        buf.dtype,
+        buf.data,
+        buf.strides,
+        new_elem_offset,
+        None,
+        buf.scope(),
+        buf.data_alignment,
+        buf.offset_factor,
+        "",
+        buf.axis_separators,
+        buf.layout,
+    )
+
+
+def bind_with_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
+    """Value binding methods when parsing with statement.
+    e.g. binding i, j, k with T.grid(128, 128, 128), when parsing
+        with T.grid(128, 128, 18) as i, j, k.
+
+    Parameters
+    ----------
+    self : Parser
+        The current parser.
+
+    node : doc.expr
+        The doc AST expression node for error reporting.
+
+    var_name : str
+        The variable name.
+
+    value : Any
+        The value to be bound with.
+
+    Returns
+    -------
+    res : Any
+        The bound value.
+    """
+    if isinstance(value, list | tuple):
+        for i, v in enumerate(value):
+            bind_with_value(self, node, f"{var_name}_{i}", v)
+        return value
+    elif isinstance(value, Buffer | Var):
+        IRBuilder.name(var_name, value)
+        return value
+    else:
+        self.report_error(node, f"Do not know how to bind type: {type(value)} in with statement")
+        raise NotImplementedError
+
+
+def bind_for_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
+    """Value binding methods when parsing for statement.
+    e.g. binding i, j, k with T.grid(128, 128, 128), when parsing
+        for i, j, k in T.grid(128, 128, 128).
+
+    Parameters
+    ----------
+    self : Parser
+        The current parser.
+
+    node : doc.expr
+        The doc AST expression node for error reporting.
+
+    var_name : str
+        The variable name.
+
+    value : Any
+        The value to be bound with.
+
+    Returns
+    -------
+    res : Any
+        The bound value.
+    """
+    if isinstance(value, list | tuple | tvm.ir.Array):
+        for i, v in enumerate(value):
+            bind_for_value(self, node, f"{var_name}_{i}", v)
+        return value
+    elif isinstance(value, Var):
+        IRBuilder.name(var_name, value)
+        return value
+    else:
+        self.report_error(node, f"Do not know how to bind type: {type(value)} in for statement")
+        raise NotImplementedError
+
+
+def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
+    """Value binding methods when parsing assign statement.
+    e.g. binding vi, vj, vk with T.axis.remap("SSR", [i, j, k]), when parsing
+        vi, vj, vk = T.axis.remap("SSR", [i, j, k]).
+
+    Parameters
+    ----------
+    self : Parser
+        The current parser.
+
+    node : doc.expr
+        The doc AST expression node for error reporting.
+
+    var_name : str
+        The variable name.
+
+    value : Any
+        The value to be bound with.
+
+    Returns
+    -------
+    res : Any
+        The bound value.
+    """
+    if isinstance(value, T.scalar_wrapper):  # pylint: disable=protected-access
+        # special case for scalar, name the buffer, but the var is used as BufferLoad
+        assert isinstance(value.scalar, T.BufferLoad)
+        IRBuilder.name(var_name, value.scalar.buffer)
+        return value.scalar
+    if isinstance(value, T.meta_var):
+        return value.value
+    elif getattr(type(value), "_is_meta_class", False):
+        name_meta_class_value(var_name, value)
+        return value
+    elif isinstance(value, list | tuple):
+        # Tuple-unpacking with a starred target (e.g. ``vi, *vs = T.axis.remap(...)``)
+        # collects multiple elements into a single list bound here. Recurse so each
+        # element gets a per-index name; this matches apache's behavior.
+        for i, v in enumerate(value):
+            bind_assign_value(self, node, f"{var_name}_{i}", v)
+        return value
+    elif isinstance(value, BufferRegion):
+        return value
+    elif isinstance(value, Frame):
+        value.add_callback(partial(value.__exit__, None, None, None))
+        res = value.__enter__()
+        IRBuilder.name(var_name, res)
+        return res
+    elif isinstance(value, Buffer | IterVar | Layout) or (
+        isinstance(value, Var) and not self.var_table.exist(value)
+    ):
+        IRBuilder.name(var_name, value)
+        return value
+    else:
+        if not tvm.ir.is_prim_expr(value):
+            value = tvm.tirx.const(value)
+        if not isinstance(value, tvm.tirx.StringImm):
+            # x = expr -> scalar (auto-typed from value)
+            scalar = T.local_scalar(dtype=str(value.ty.dtype))
+            IRBuilder.name(var_name, scalar.scalar.buffer)
+            T.buffer_store(scalar.scalar.buffer, value, [0])
+            return scalar.scalar
+        else:
+            # StringImm: x = expr -> immutable Bind var
+            ann_var = tvm.tirx.Var(var_name, value.ty)
+            IRBuilder.name(var_name, ann_var)
+            T.Bind(value, var=ann_var)
+            return ann_var
+
+
+def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: bool = True) -> bool:
+    """
+    Check the value of given annotation (argument name) in the prim_func decorator.
+    Returns the value of the annotation if present, otherwise giving the default value.
+    """
+    # look for the named argument in the prim_func / jit decorator
+    for dec in node.decorator_list:
+        if not isinstance(dec, doc.Call) or dec.func.attr not in ("prim_func", "jit"):
+            continue
+        for keyword in dec.keywords:
+            if keyword.arg == annotation:
+                return keyword.value.value
+    return default
+
+
+@dispatch.register(token="tirx", type_name="For")
+def visit_for(self: Parser, node: doc.For) -> None:
+    """The for visiting method for tirx.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.For
+        The doc AST for node.
+    """
+    # Intercept range() at AST level so it works with both Python ints and PrimExprs.
+    # In other contexts (e.g. list comprehensions), range remains Python's builtin.
+    if (
+        isinstance(node.iter, doc.Call)
+        and isinstance(node.iter.func, doc.Name)
+        and node.iter.func.id == "range"
+    ):
+        args = [self.eval_expr(a) for a in node.iter.args]
+        kwargs = {kw.arg: self.eval_expr(kw.value) for kw in node.iter.keywords}
+        if len(args) == 1:
+            for_frame = T.serial(0, args[0], **kwargs)
+        elif len(args) == 2:
+            for_frame = T.serial(args[0], args[1], **kwargs)
+        elif len(args) == 3:
+            for_frame = T.serial(args[0], args[1], step=args[2], **kwargs)
+        else:
+            self.report_error(node.iter, "range() takes 1 to 3 arguments")
+    else:
+        for_frame = self.eval_expr(node.iter)
+    if not isinstance(for_frame, T.frame.ForFrame):
+        self.report_error(
+            node.iter,
+            "Expect the for loop to be one of the following: "
+            "range, T.serial, T.grid, T.parallel, T.vectorized, T.unroll, T.thread_binding",
+        )
+    with self.var_table.with_frame():
+        with for_frame as iters:
+            self.eval_assign(target=node.target, source=iters, bind_value=bind_for_value)
+            self.visit_body(node.body)
+
+
+@dispatch.register(token="tirx", type_name="While")
+def visit_while(self: Parser, node: doc.While) -> None:
+    """The while visiting method for tirx.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.While
+        The doc AST while node.
+    """
+    with self.var_table.with_frame():
+        cond = self.eval_expr(node.test)
+        with T.While(cond):
+            self.visit_body(node.body)
+
+
+@dispatch.register(token="tirx", type_name="Break")
+def visit_break(self: Parser, node: doc.Break) -> None:
+    """The break visiting method for tir.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.Break
+        The doc AST break node.
+    """
+    T.evaluate(T.break_loop())
+
+
+@dispatch.register(token="tirx", type_name="Continue")
+def visit_continue(self: Parser, node: doc.Continue) -> None:
+    """The continue visiting method for tir.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.Continue
+        The doc AST continue node.
+    """
+    T.evaluate(T.continue_loop())
+
+
+@dispatch.register(token="tirx", type_name="Assign")
+def visit_assign(self: Parser, node: doc.Assign) -> None:
+    """The assign visiting method for tirx.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.Assign
+        The doc AST assign node.
+    """
+    if len(node.targets) != 1:
+        self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
+    lhs = node.targets[0]
+
+    if isinstance(node.value, doc.Subscript):
+        check_slices = []
+        if isinstance(node.value.slice, doc.Slice):
+            check_slices = [node.value.slice]
+        elif isinstance(node.value.slice, doc.Tuple):
+            for p in node.value.slice.elts:
+                if isinstance(p, doc.Slice):
+                    check_slices.append(p)
+        for s in check_slices:
+            if not s.step and s.upper and s.lower:
+                s.step = doc.Constant(
+                    1,
+                    None,
+                    s.upper.lineno,
+                    s.upper.end_col_offset + 1,
+                    s.upper.lineno,
+                    s.upper.end_col_offset + 2,
+                )
+
+    rhs = self.eval_expr(node.value)
+    if isinstance(lhs, doc.Subscript):
+        if isinstance(lhs.slice, doc.Tuple):
+            indices = []
+            for index in lhs.slice.elts:
+                if isinstance(index, doc.Starred):
+                    # x[*y]
+                    indices.extend(self.eval_expr(index.value))
+                else:
+                    indices.append(self.eval_expr(index))
+        else:
+            indices = self.eval_expr(lhs.slice)
+        T.buffer_store(self.eval_expr(lhs.value), rhs, indices)
+    else:
+        # special case for scalar buffers
+        # scalar = xxx <=> scalar.buffer[()] = xxx
+        # or for a normal 1-dim buffer with shape (1,)
+        # buffer = xxx <=> buffer[()] = xxx
+        # Try to resolve lhs as a buffer/scalar variable. eval_expr may raise
+        # if the name is not yet defined (i.e. this is a new variable binding),
+        # which is the expected fallthrough case.
+        lhs_value = None
+        try:
+            lhs_copy = deepcopy(lhs)
+            if hasattr(lhs_copy, "ctx"):
+                lhs_copy.ctx = doc.Load()
+            lhs_value = self.eval_expr(lhs_copy)
+        except Exception:  # pylint: disable=broad-except
+            pass
+        # Buffer check and store are intentionally outside the try/except so
+        # that genuine errors (e.g. wrong shape, bad store) are not swallowed.
+        # Only TypeError from FFI type mismatch (e.g. rhs is a meta_var, not
+        # a Expr or auto-convertible scalar) triggers fallthrough.
+        if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad | tvm.tirx.Buffer):
+            if isinstance(lhs_value, T.scalar_wrapper):
+                buffer = lhs_value.scalar.buffer
+            else:
+                buffer = lhs_value.buffer if isinstance(lhs_value, T.BufferLoad) else lhs_value
+            if len(buffer.shape) == 1 and bool(buffer.shape[0] == 1):
+                # only 1-dim buffer with shape (1,) can be assigned directly
+                # Note that shape can be a Expr, so we only judge by
+                # bool(shape[0] == 1) rather than int(shape[0]) == 1.
+                try:
+                    T.buffer_store(buffer, rhs, [0])
+                    return
+                except TypeError:
+                    pass  # rhs not compatible with buffer_store, fall through
+        # otherwise
+        self.eval_assign(target=lhs, source=rhs, bind_value=bind_assign_value)
+
+
+@dispatch.register(token="tirx", type_name="AugAssign")
+def visit_aug_assign(self: Parser, node: doc.AugAssign) -> None:
+    """The augmented assign visiting method for tirx.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.AugAssign
+        The doc AST augmented assign node.
+    """
+    lhs_pos = (
+        node.target.lineno,
+        node.target.col_offset,
+        node.target.end_lineno,
+        node.target.end_col_offset,
+    )
+    rhs_pos = (
+        node.value.lineno,
+        node.value.col_offset,
+        node.value.end_lineno,
+        node.value.end_col_offset,
+    )
+    node.target.ctx = doc.Load()
+    with self.var_table.with_frame():
+        lhs_name = "__tvm_tmp_value_aug_assign_lhs"
+        rhs_name = "__tvm_tmp_value_aug_assign_rhs"
+        lhs_expr = self.eval_expr(node.target)
+        rhs_expr = self.eval_expr(node.value)
+        self.var_table.add(lhs_name, lhs_expr)
+        self.var_table.add(rhs_name, rhs_expr)
+        op = doc.BinOp(
+            doc.Name(lhs_name, doc.Load(), *lhs_pos),
+            node.op,
+            doc.Name(rhs_name, doc.Load(), *rhs_pos),
+            *lhs_pos,
+        )
+        rhs = self.eval_expr(op)
+    lhs = node.target
+    lhs.ctx = doc.Store()
+    if isinstance(lhs, doc.Subscript):
+        if isinstance(lhs.slice, doc.Tuple):
+            indices = []
+            for index in lhs.slice.elts:
+                if isinstance(index, doc.Starred):
+                    # x[*y]
+                    indices.extend(self.eval_expr(index.value))
+                else:
+                    indices.append(self.eval_expr(index))
+        else:
+            indices = [self.eval_expr(lhs.slice)]
+        T.buffer_store(self.eval_expr(lhs.value), rhs, indices)
+    else:
+        lhs_value = None
+        try:
+            lhs_copy = deepcopy(lhs)
+            if hasattr(lhs_copy, "ctx"):
+                lhs_copy.ctx = doc.Load()
+            lhs_value = self.eval_expr(lhs_copy)
+        except Exception:  # pylint: disable=broad-except
+            pass
+        if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad | tvm.tirx.Buffer):
+            if isinstance(lhs_value, T.scalar_wrapper):
+                buffer = lhs_value.scalar.buffer
+            else:
+                buffer = lhs_value.buffer if isinstance(lhs_value, T.BufferLoad) else lhs_value
+            if len(buffer.shape) == 1 and bool(buffer.shape[0] == 1):
+                try:
+                    T.buffer_store(buffer, rhs, [0])
+                    return
+                except TypeError:
+                    pass
+        self.eval_assign(target=lhs, source=rhs, bind_value=bind_assign_value)
+
+
+@dispatch.register(token="tirx", type_name="AnnAssign")
+def visit_ann_assign(self: Parser, node: doc.AnnAssign) -> None:
+    """The annotated assign visiting method for tirx.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.AnnAssign
+        The doc AST annotated assign node.
+    """
+    lhs = node.target
+    rhs = self.eval_expr(node.value) if node.value is not None else None
+    raw_ann = self.eval_expr(node.annotation)
+
+    if isinstance(raw_ann, T.LocalVectorAnnotation):
+        # x: T.float32[N] or x: T.f32[M, N] -> local buffer allocation
+        if rhs is not None:
+            self.report_error(node, "Vector annotation does not support initial value")
+        buf = T.alloc_local(shape=raw_ann.shape, dtype=raw_ann.dtype)
+        self.eval_assign(target=lhs, source=buf, bind_value=bind_assign_value)
+    elif isinstance(raw_ann, T.LetAnnotation):
+        # T.let or T.let[type] -> immutable Bind var
+        if rhs is None:
+            self.report_error(node, "T.let annotation requires a value")
+        if not isinstance(rhs, Expr):
+            if isinstance(rhs, str):
+                rhs = tvm.tirx.StringImm(rhs)
+            else:
+                rhs = tvm.tirx.const(rhs)
+        if raw_ann.type_spec is not None:
+            ann_var = raw_ann.as_var()
+        else:
+            ann_var = raw_ann.as_var(rhs_dtype=rhs.ty)
+        if not isinstance(ann_var, Var):
+            self.report_error(node.annotation, "Annotation should resolve to Var")
+        self.eval_assign(target=lhs, source=ann_var, bind_value=bind_assign_value)
+        T.Bind(rhs, var=ann_var)
+    else:
+        ann_var = raw_ann() if callable(raw_ann) else raw_ann
+        if not isinstance(ann_var, Var):
+            self.report_error(node.annotation, "Annotation should resolve to Var")
+        if not isinstance(ann_var.ty, PrimType):
+            self.report_error(
+                node.annotation,
+                "Use T.let[...] for non-PrimType annotations (e.g. PointerType, handle)",
+            )
+        if str(ann_var.ty) == "handle":
+            self.report_error(
+                node.annotation,
+                "handle type cannot be used as scalar annotation; use T.let[T.handle] instead",
+            )
+        # x: T.int32 = expr -> scalar (mutable scalar buffer)
+        scalar = T.local_scalar(dtype=str(ann_var.ty))
+        self.eval_assign(target=lhs, source=scalar, bind_value=bind_assign_value)
+        if rhs is not None:
+            T.buffer_store(scalar.scalar.buffer, rhs, [0])
+
+
+@dispatch.register(token="tirx", type_name="With")
+def visit_with(self: Parser, node: doc.With) -> None:
+    """The with visiting method for tirx.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.With
+        The doc AST with node.
+    """
+    with contextlib.ExitStack() as stack:
+        stack.enter_context(self.var_table.with_frame())
+        for item in node.items:
+            frame = self.eval_expr(item.context_expr)
+            if not isinstance(frame, Frame) and not (
+                hasattr(frame, "__enter__") and hasattr(frame, "__exit__")
+            ):
+                self.report_error(
+                    item.context_expr,
+                    "Invalid context expression in the with-statement.",
+                )
+            rhs = stack.enter_context(frame)
+            if item.optional_vars is not None:
+                self.eval_assign(target=item.optional_vars, source=rhs, bind_value=bind_with_value)
+        self.visit_body(node.body)
+
+
+@dispatch.register(token="tirx", type_name="FunctionDef")
+def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
+    """The function definition visiting method for tirx.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.FunctionDef
+        The doc AST function definition node.
+    """
+    supplied_annotation = self.function_annotations
+    func_annotation = supplied_annotation.get(node.name, {})
+    privacy = find_decorator_annotation(node, "private", default=False)
+    s_tir = find_decorator_annotation(node, "s_tir", default=False)
+    persistent = find_decorator_annotation(node, "persistent", default=False)
+    self.function_annotations = None
+    with self.var_table.with_frame():
+        prim_func_ctx = T.prim_func(is_private=privacy, s_tir=s_tir, persistent=persistent)
+        with prim_func_ctx:
+            T.func_name(node.name)
+            if node.returns is not None:
+                ret_type = self.eval_expr(node.returns)
+                if callable(ret_type):
+                    ret_type = ret_type().ty
+                T.func_ret(ret_type)
+            with self.with_dispatch_token("tirx"):
+                # TODO: handle different types of arguments:
+                # - vararg: arg | None
+                # - kwonlyargs: list[arg]
+                # - kw_defaults: list[expr | None]
+                # - kwarg: arg | None
+                # - defaults: list[expr]
+                # - posonlyargs: list[arg]
+                for arg in node.args.args:
+                    if arg.annotation is None:
+                        self.report_error(arg, "Type annotation required for function parameters.")
+                    try:
+                        ann = self.eval_expr(arg.annotation)
+                        if callable(ann) and ann is not _constexpr_sentinel:
+                            ann = ann()
+                    except Exception:  # pylint: disable=broad-except
+                        ann = func_annotation.get(arg.arg, None)
+                        if ann is None:
+                            raise
+                    if ann is _constexpr_sentinel:
+                        # T.constexpr param: value was bound in extra_vars by
+                        # TIRJit.specialize() and lives in an outer var_table
+                        # frame; do not register a runtime PrimFunc param.
+                        continue
+                    param = T.arg(arg.arg, ann)
+                    self.var_table.add(arg.arg, param)
+                self.visit_body(node.body)
+    self.function_annotations = supplied_annotation
+
+
+@dispatch.register(token="tir.inline", type_name="FunctionDef")
+def visit_inline_function_def(self: Parser, node: doc.FunctionDef) -> None:
+    """The function definition visiting method for inline functions in tir.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.FunctionDef
+        The doc AST function definition node.
+    """
+    # remove the inline decorator
+    node.decorator_list.pop()
+    # adjust the node location to the source code location
+    node.lineno += self.diag.source.start_line - 1
+    node.col_offset += self.diag.source.start_column + 1
+    node.end_lineno += self.diag.source.start_line - 1
+    node.end_col_offset += self.diag.source.start_column + 1
+
+    # Record definition depth for LEGB late binding
+    definition_depth = len(self.var_table.frames)
+
+    def get_func():
+        func_ast = from_doc(node)
+        module_ast = ast.Module(body=[func_ast], type_ignores=[])
+        ast.fix_missing_locations(module_ast)
+        # set the filename to the source name, so that the error message can be reported correctly
+        code_obj = compile(module_ast, filename=self.diag.source.source_name, mode="exec")
+        namespace = self.var_table.get()
+        exec(code_obj, namespace)  # pylint: disable=exec-used
+        func_name = func_ast.name
+        func = namespace[func_name]
+        return func, func_name
+
+    func, func_name = get_func()
+    wrapper = inline(func, definition_depth=definition_depth, defining_var_table=self.var_table)
+
+    self.var_table.add(func_name, wrapper, allow_shadowing=False)
+    return None
+
+
+@dispatch.register(token="tirx", type_name="tvm_annotation")
+def visit_tvm_annotation(self: Parser, node: doc.expr):
+    """The TVM annotation visiting method for tirx.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.expr
+        The doc AST expr node.
+    """
+    annotation = self.eval_expr(node)
+    if callable(annotation):
+        annotation = annotation()
+    return annotation
+
+
+@dispatch.register(token="tirx", type_name="Expr")
+def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
+    """The expr statement visiting method for tirx.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.Expr
+        The doc AST Expr node.
+    """
+
+    res = self.eval_expr(node.value)
+    if res is None:
+        pass
+    elif isinstance(res, Frame):
+        res.add_callback(partial(res.__exit__, None, None, None))
+        res.__enter__()
+    elif hasattr(res, "frames") and hasattr(res, "__enter__"):
+        # _FrameScope from T.attr({...}) — enter each inner frame for concise scoping
+        for f in res.frames:
+            f.add_callback(partial(f.__exit__, None, None, None))
+            f.__enter__()
+    elif isinstance(res, Var):
+        # Standalone Var expression (e.g. from T.bind(value, var=v)) --
+        # the Bind statement was already emitted to the parent frame by the FFI call,
+        # so just discard the returned Var.
+        pass
+    elif tvm.ir.is_prim_expr(res):
+        T.evaluate(res)
+    elif isinstance(res, int | bool):
+        T.evaluate(tvm.tirx.const(res))
+    elif isinstance(res, tvm.ir.Call) and not tvm.ir.is_prim_expr(res):
+        if isinstance(res.op, tvm.ir.GlobalVar) and res.ty.is_missing():
+            # GlobalVar calls with a missing return type are ambiguous, as each IR has a
+            # different function Call representation. Convert to the TIR representation.
+            T.evaluate(tvm.tirx.call_tir(res.op, *res.args))
+        else:
+            # Pointer-valued TIR calls are general Expr rather than PrimExpr,
+            # but are still valid standalone Evaluate statements.
+            T.evaluate(res)
+    elif isinstance(res, str):
+        # Ignore docstrings
+        pass
+    elif isinstance(res, tvm.tirx.stmt.BufferStore):
+        T.buffer_store(res.buffer, res.value, res.indices, res.predicate)
+    elif isinstance(res, tvm.tirx.Buffer):
+        # ``T.match_buffer(...)`` used as a bare statement (no LHS) — the
+        # buffer object is discarded; the underlying side effect (the
+        # match_buffer node) has already been emitted into the frame.
+        pass
+    else:
+        self.report_error(node, f"Parsing resulted in unexpected type {type(res)}")
+
+
+@dispatch.register(token="tirx", type_name="If")
+def visit_if(self: Parser, node: doc.If) -> None:
+    """The if visiting method for tirx.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.If
+        The doc AST if node.
+    """
+    with self.var_table.with_frame():
+        predicate = self.eval_expr(node.test)
+        if tvm.ir.is_prim_expr(predicate) or isinstance(predicate, tvm.tirx.expr.ExprOp):
+            with T.If(self.eval_expr(node.test)):
+                with T.Then():
+                    with self.var_table.with_frame():
+                        self.visit_body(node.body)
+                if node.orelse:
+                    with T.Else():
+                        with self.var_table.with_frame():
+                            self.visit_body(node.orelse)
+        elif isinstance(predicate, bool):
+            if predicate:
+                with self.var_table.with_frame():
+                    self.visit_body(node.body)
+            elif node.orelse:
+                with self.var_table.with_frame():
+                    self.visit_body(node.orelse)
+        else:
+            self.report_error(
+                node.test,
+                f"If condition must be a boolean expression, but got {predicate}",
+            )
+
+
+@dispatch.register(token="tirx", type_name="Assert")
+def visit_assert(self: Parser, node: doc.Assert) -> None:
+    """The assert visiting method for tirx.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.Assert
+        The doc AST assert node.
+
+    The assert message can be either:
+    - A plain string: ``assert cond, "message"``
+    - A tuple of (kind, [parts...]): ``assert cond, ("ValueError", ["part0", "part1"])``
+    """
+    cond = self.eval_expr(node.test)
+    msg = self.eval_expr(node.msg)
+
+    kind = "RuntimeError"
+    message = msg
+
+    if isinstance(msg, tuple):
+        if len(msg) != 2:
+            self.report_error(
+                node,
+                f"Assert message tuple must have exactly 2 elements (kind, [parts...]), "
+                f"got {len(msg)} elements",
+            )
+        kind_str, parts = msg
+        if isinstance(kind_str, tvm.tirx.StringImm):
+            kind_str = kind_str.value
+        if not isinstance(kind_str, str):
+            self.report_error(
+                node,
+                f"Assert message tuple first element must be a string (error kind like "
+                f'"ValueError"), got {type(kind_str).__name__}',
+            )
+        kind = kind_str
+        message = parts
+
+    if isinstance(message, list | tuple):
+        message = [p.value if isinstance(p, tvm.tirx.StringImm) else str(p) for p in message]
+
+    frame = T.Assert(cond, message, error_kind=kind)
+    frame.add_callback(partial(frame.__exit__, None, None, None))
+    frame.__enter__()
+
+
+@dispatch.register(token="tirx", type_name="Return")
+def visit_return(self: Parser, node: doc.Return) -> None:
+    """The return visiting method for tirx.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.Return
+        The doc AST return node.
+    """
+    value = self.eval_expr(node.value)
+    if value is None:
+        self.report_error(node, "Expression to be returned must be a Expr")
+    T.evaluate(tvm.tirx.ret(value))
+
+
+@dispatch.register(token="tirx", type_name="tvm_declare_function")
+def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
+    """The function declaration step for tirx
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.Return
+        The doc AST return node.
+    """
+
+    supplied_annotation = self.function_annotations
+    func_annotation = supplied_annotation.get(node.name, {})
+
+    ret_type = None
+    with self.var_table.with_frame():
+        if node.returns is not None:
+            ret_type = self.eval_expr(node.returns)
+            if callable(ret_type):
+                ret_type = ret_type().ty
+
+        arg_annotations = []
+        for arg in node.args.args:
+            if arg.annotation is None:
+                self.report_error(arg, "Type annotation required for function parameters.")
+            try:
+                ann = self.eval_expr(arg.annotation)
+                if callable(ann):
+                    ann = ann()
+            except Exception:  # pylint: disable=broad-except
+                ann = func_annotation.get(arg.arg, None)
+                if ann is None:
+                    raise
+
+            IRBuilder.name(arg.arg, ann)
+            arg_annotations.append(ann)
+
+    func_signature = tvm.tirx.PrimFunc(arg_annotations, None, ret_type=ret_type)
+    return I.decl_function(node.name, func_signature)