chore: import upstream snapshot with attribution

python/tvm/script/parser/__init__.py

@@ -0,0 +1,51 @@
+# 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 parser subpackage of TVMScript.
+
+Per-dialect parser submodules (``tvm.script.parser.tirx``, etc.) are
+resolved lazily via :data:`tvm.script._DIALECT_REGISTRY`.  When a dialect
+is accessed (e.g. ``tvm.script.parser.tirx``), this subpackage's
+``__getattr__`` looks up the dialect in ``_DIALECT_REGISTRY`` and imports
+``<dialect_module_path>.parser`` (e.g. ``tvm.tirx.script.parser``),
+caching the result so subsequent accesses skip ``__getattr__``.
+
+The IR layer is foundational and is NOT registered as a dialect — its
+parser lives as a real submodule ``tvm.script.parser.ir``, with
+``ir_module`` re-exported at this level for convenience.
+
+See :mod:`tvm.script` for a full description of the dialect resolution
+mechanism, including the ``_DialectRedirectFinder`` that handles
+deep statement-form imports.
+"""
+
+import importlib
+from typing import Any
+
+from . import _core, ir, tirx
+from ._core import parse
+from .ir import ir_module
+
+
+def __getattr__(name: str) -> Any:
+    # Lazy import to avoid loading tvm.script during dialect bootstrap.
+    from tvm.script import _DIALECT_REGISTRY  # pylint: disable=import-outside-toplevel
+
+    if name in _DIALECT_REGISTRY:
+        module = importlib.import_module(f"{_DIALECT_REGISTRY[name]}.parser")
+        globals()[name] = module
+        return module
+    raise AttributeError(f"module 'tvm.script.parser' has no attribute {name!r}")

python/tvm/script/parser/_core.py

@@ -0,0 +1,24 @@
+# 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 Licens.
+# ruff: noqa: F401
+"""The core parser infra"""
+
+# pylint: disable=unused-import
+from .core import dispatch, doc, utils
+from .core.dispatch import OpMethod, register_op
+from .core.entry import parse, scan_macro
+from .core.parser import Parser

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

@@ -0,0 +1,20 @@
+# isort: skip_file
+# 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 core parser infra"""
+
+from . import diagnostics, dispatch, doc, doc_core, entry, evaluator, parser, utils

python/tvm/script/parser/core/diagnostics.py

@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E741
+"""TVM Script Parser Source and diagnostics"""
+
+import inspect
+import sys
+
+from tvm.error import DiagnosticError
+
+from . import doc
+
+
+class Source:
+    """Source code class for TVMScript.
+
+    It is constructed by source code str or doc AST tree.
+
+    Parameters
+    ----------
+    source_name : str
+        The filename of the file where the source code locates.
+
+    start_line : int
+        The first line number of the source code.
+
+    start_column : int
+        The first column number of the first line of the source code.
+
+    source : str
+        The source code str of source code.
+
+    full_source : str
+        The complete source code of the file where the source code locates.
+    """
+
+    source_name: str
+    start_line: int
+    start_column: int
+    source: str
+    full_source: str
+
+    def __init__(self, program: str | doc.AST):
+        if isinstance(program, str):
+            self.source_name = "<str>"
+            self.start_line = 1
+            self.start_column = 0
+            self.source = program
+            self.full_source = program
+            return
+
+        self.source_name = inspect.getsourcefile(program)  # type: ignore
+        lines, self.start_line = getsourcelines(program)  # type: ignore
+        if lines:
+            self.start_column = len(lines[0]) - len(lines[0].lstrip())
+        else:
+            self.start_column = 0
+        if self.start_column and lines:
+            self.source = "\n".join([l[self.start_column :].rstrip() for l in lines])
+        else:
+            self.source = "".join(lines)
+        try:
+            # It will cause a problem when running in Jupyter Notebook.
+            # `mod` will be <module '__main__'>, which is a built-in module
+            # and `getsource` will throw a TypeError
+            mod = inspect.getmodule(program)
+            if mod:
+                self.full_source = inspect.getsource(mod)
+            else:
+                self.full_source = self.source
+        except TypeError:
+            # It's a work around for Jupyter problem.
+            # Since `findsource` is an internal API of inspect, we just use it
+            # as a fallback method.
+            src, _ = inspect.findsource(program)  # type: ignore
+            self.full_source = "".join(src)
+
+    def as_ast(self) -> doc.AST:
+        """Parse the source code into AST.
+
+        Returns
+        -------
+        res : doc.AST
+            The AST of source code.
+        """
+        return doc.parse(self.source)
+
+
+_getfile = inspect.getfile  # pylint: disable=invalid-name
+_findsource = inspect.findsource  # pylint: disable=invalid-name
+
+
+def _patched_inspect_getfile(obj):
+    """Work out which source or compiled file an object was defined in."""
+    if not inspect.isclass(obj):
+        return _getfile(obj)
+    mod = getattr(obj, "__module__", None)
+    if mod is not None:
+        file = getattr(sys.modules[mod], "__file__", None)
+        if file is not None:
+            return file
+    for _, member in inspect.getmembers(obj):
+        if inspect.isfunction(member):
+            if obj.__qualname__ + "." + member.__name__ == member.__qualname__:
+                return inspect.getfile(member)
+    raise TypeError(f"Source for {obj!r} not found")
+
+
+def findsource(obj):
+    """Return the entire source file and starting line number for an object."""
+    import linecache  # pylint: disable=import-outside-toplevel
+
+    if not inspect.isclass(obj):
+        return _findsource(obj)
+
+    file = inspect.getsourcefile(obj)
+    if file:
+        linecache.checkcache(file)
+    else:
+        file = inspect.getfile(obj)
+        if not (file.startswith("<") and file.endswith(">")):
+            raise OSError("source code not available")
+
+    module = inspect.getmodule(obj, file)
+    if module:
+        lines = linecache.getlines(file, module.__dict__)
+    else:
+        lines = linecache.getlines(file)
+    if not lines:
+        raise OSError("could not get source code")
+    qual_names = obj.__qualname__.replace(".<locals>", "<locals>").split(".")
+    in_comment = 0
+    scope_stack = []
+    indent_info = {}
+    for i, line in enumerate(lines):
+        n_comment = line.count('"""')
+        if n_comment:
+            # update multi-line comments status
+            in_comment = in_comment ^ (n_comment & 1)
+            continue
+        if in_comment:
+            # skip lines within multi-line comments
+            continue
+        indent = len(line) - len(line.lstrip())
+        tokens = line.split()
+        if len(tokens) > 1:
+            name = None
+            if tokens[0] == "def":
+                name = tokens[1].split(":")[0].split("(")[0] + "<locals>"
+            elif tokens[0] == "class":
+                name = tokens[1].split(":")[0].split("(")[0]
+            # pop scope if we are less indented
+            while scope_stack and indent_info[scope_stack[-1]] >= indent:
+                scope_stack.pop()
+            if name:
+                scope_stack.append(name)
+                indent_info[name] = indent
+                if scope_stack == qual_names:
+                    return lines, i
+
+    raise OSError("could not find class definition")
+
+
+def getsourcelines(obj):
+    """Extract the block of code at the top of the given list of lines."""
+    obj = inspect.unwrap(obj)
+    lines, l_num = findsource(obj)
+    return inspect.getblock(lines[l_num:]), l_num + 1
+
+
+inspect.getfile = _patched_inspect_getfile
+
+
+def _format_source_snippet(
+    source_lines: list,
+    lineno: int,
+    col_offset: int,
+    end_lineno: int,
+    end_col_offset: int,
+) -> str:
+    """Format a source code snippet with a column/span marker.
+
+    Renders every source line spanned by the diagnostic (``lineno`` through
+    ``end_lineno``, inclusive) with a per-line gutter, followed by a caret
+    underline covering the offending span. For a single-line span
+    (``end_lineno == lineno``) only the columns ``col_offset`` (inclusive)
+    through ``end_col_offset`` (exclusive) are underlined; for a multi-line
+    span the underline covers the start column to end-of-line on the first
+    line, the full text of interior lines, and the start of the final line up
+    to ``end_col_offset``.
+
+    Parameters
+    ----------
+    source_lines : list of str
+        Lines of the source code.
+
+    lineno : int
+        1-based starting line number in the source.
+
+    col_offset : int
+        1-based starting column (inclusive) on ``lineno``.
+
+    end_lineno : int
+        1-based ending line number in the source (>= ``lineno``).
+
+    end_col_offset : int
+        1-based ending column (exclusive) on ``end_lineno``.
+
+    Returns
+    -------
+    snippet : str
+        Formatted source snippet with caret-marker line(s).
+    """
+    if end_lineno < lineno:
+        end_lineno = lineno
+
+    # Determine the gutter width so that all line numbers line up.
+    header_width = len(f" {end_lineno} ")
+    no_line_header = " " * header_width
+
+    parts = [f"{no_line_header}|  "]
+    for cur_lineno in range(lineno, end_lineno + 1):
+        idx = cur_lineno - 1
+        if not 0 <= idx < len(source_lines):
+            continue
+        line_text = source_lines[idx].rstrip("\n")
+        line_header = f" {cur_lineno} ".rjust(header_width)
+
+        # Compute the underline span [start_col, stop_col) for this line.
+        start_col = col_offset if cur_lineno == lineno else 1
+        stop_col = end_col_offset if cur_lineno == end_lineno else len(line_text) + 1
+
+        marker = ""
+        for i in range(1, len(line_text) + 1):
+            if start_col <= i < stop_col:
+                marker += "^"
+            else:
+                marker += " "
+        parts.append(f"{line_header}|  {line_text}")
+        parts.append(f"{no_line_header}|  {marker}")
+
+    if len(parts) == 1:
+        return ""
+    return "\n".join(parts)
+
+
+class Diagnostics:
+    """Diagnostics class for error reporting in parser.
+
+    Formats parse errors with source location context and raises directly,
+    without going through DiagnosticContext.
+
+    Parameters
+    ----------
+    source : Source
+        The source code.
+    """
+
+    source: Source
+
+    def __init__(self, source: Source):
+        self.source = source
+
+    def error(self, node: doc.AST, message: str) -> None:
+        """Emit a diagnostic error by raising with source location context.
+
+        Parameters
+        ----------
+        node : doc.AST
+            The node with diagnostic error.
+
+        message : str
+            The diagnostic message.
+        """
+        lineno = getattr(node, "lineno", 1)
+        col_offset = getattr(node, "col_offset", self.source.start_column)
+        end_lineno = getattr(node, "end_lineno", lineno)
+        end_col_offset = getattr(node, "end_col_offset", col_offset)
+        lineno += self.source.start_line - 1
+        end_lineno += self.source.start_line - 1
+        col_offset += self.source.start_column + 1
+        end_col_offset += self.source.start_column + 1
+
+        source_lines = self.source.full_source.splitlines(keepends=True)
+        snippet = _format_source_snippet(
+            source_lines, lineno, col_offset, end_lineno, end_col_offset
+        )
+
+        location = f"{self.source.source_name}:{lineno}:{col_offset}"
+        formatted = f"error: {message}\n --> {location}\n{snippet}"
+        raise DiagnosticError(formatted)

python/tvm/script/parser/core/dispatch.py

@@ -0,0 +1,157 @@
+# 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.
+"""Parser dispatching infrastructure"""
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+from .doc import AST
+
+if TYPE_CHECKING:
+    from .parser import Parser
+
+
+ParseMethod = Callable[["Parser", AST], None]
+ParseVTable: dict[tuple[str, str], ParseMethod] = {}
+
+OpMethod = Callable[..., Any]
+OpVTable: dict[tuple[type, AST, int], OpMethod] = {}
+
+
+def register(token: str, type_name: str):
+    """Register a method for a dispatch token and type name.
+
+    Parameters
+    ----------
+    token : str
+        The token for IR, e.g., T for TIR and R for Relax.
+
+    type_name : str
+        The type name of AST node, e.g., FunctionDef, With, For.
+
+    Returns
+    -------
+    func : callable
+        The function to register dispatched method of parsing
+        corresponding token and AST node type.
+    """
+
+    def func(method: ParseMethod):
+        """Register a method in parser virtual table.
+
+        Parameters
+        ----------
+        method : ParseMethod
+            The dispatched method to be registered in parser virtual table.
+        """
+        ParseVTable[(token, type_name)] = method
+
+    return func
+
+
+def get(
+    token: str,
+    type_name: str,
+    default: ParseMethod | None = None,
+) -> ParseMethod | None:
+    """Get a registered method for a dispatch token and type name,
+    or return a default method if no registered methods with this dispatch token and type name.
+
+    Parameters
+    ----------
+    token : str
+        The token for IR, e.g., T for TIR and R for Relax.
+
+    type_name : str
+        The type name of AST node, e.g., FunctionDef, With, For.
+
+    default : Optional[ParseMethod]
+        The default method when no registered methods with this dispatch token and type name.
+
+    Returns
+    -------
+    func : Optional[ParseMethod]
+        The dispatched method of parsing corresponding token and AST node type.
+    """
+    return ParseVTable.get((token, type_name), default)
+
+
+def register_op(operand_type: type, op_node_type: AST, operand_index: int):
+    """Register a method for a operand type, AST operator node and operand index.
+
+    Parameters
+    ----------
+    operand_type : Type
+        The type of operands, e.g., tirx.Expr, tirx.IterVar.
+
+    op_node_type : AST
+        The doc AST operator node type, e.g., doc.Add, doc.Eq.
+
+    operand_index : int
+        The operand index, i.e., 0 for left operand and 1 for right operand.
+
+    Returns
+    -------
+    func : callable
+        The function to register dispatched method of parsing
+        corresponding a operand type, AST operator node and operand index.
+    """
+
+    def func(method: OpMethod):
+        """Register a method in parser operator virtual table.
+
+        Parameters
+        ----------
+        method : ParseMethod
+            The dispatched method to be registered in parser operator virtual table.
+        """
+        OpVTable[(operand_type, op_node_type, operand_index)] = method
+
+    return func
+
+
+def get_op(
+    operand_type: type,
+    op_node_type: type,
+    operand_index: int,
+    default: OpMethod | None = None,
+) -> OpMethod | None:
+    """Register a method for a operand type, AST operator node and operand index.
+
+    Parameters
+    ----------
+    operand_type : Type
+        The type of operands, e.g., tirx.Expr, tirx.IterVar.
+
+    op_node_type : AST
+        The doc AST operator node type, e.g., doc.Add, doc.Eq.
+
+    operand_index : int
+        The operand index, i.e., 0 for left operand and 1 for right operand.
+
+
+    default : Optional[OpMethod]
+        The default method when no registered methods with this operand type,
+        AST operator node and operand index.
+
+    Returns
+    -------
+    func : Optional[OpMethod]
+        The function to register dispatched method of parsing
+        corresponding a operand type, AST operator node and operand index.
+    """
+    return OpVTable.get((operand_type, op_node_type, operand_index), default)

python/tvm/script/parser/core/doc.py

@@ -0,0 +1,315 @@
+# 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: E722, F403
+"""TVM Script Parser doc AST"""
+
+import ast
+import inspect
+import typing
+from collections import defaultdict
+
+from . import doc_core as doc
+from .doc_core import *  # pylint: disable=unused-import,wildcard-import,redefined-builtin,W0614
+
+FnToDoc = typing.Callable[[ast.AST], doc.AST]
+FnFromDoc = typing.Callable[[doc.AST], ast.AST]
+
+
+class Entry:
+    """Mapping entry between python AST node type str and doc AST.
+
+    Parameters
+    ----------
+    to_doc : typing.Optional[FnToDoc]
+        The callable methods for converting python AST node to doc AST.
+
+    from_doc : typing.Optional[FnFromDoc]
+        The callable methods for converting doc AST to python AST node.
+    """
+
+    to_doc: FnToDoc | None
+    from_doc: FnFromDoc | None
+
+    def __init__(self):
+        self.to_doc = None
+        self.from_doc = None
+
+
+class Registry:
+    """Registration map from python AST node type str to methods of conversion
+    between python AST node and doc AST node.
+
+    Parameters
+    ----------
+    _inst : typing.Optional["Registry"]
+        The instance of Registry.
+
+    table : typing.Dict[str, Entry]
+        The registration map from python AST node type str to methods of conversion
+        between python AST node and doc AST node.
+    """
+
+    _inst: typing.Optional["Registry"] = None
+    table: dict[str, Entry]
+
+    def __init__(self):
+        self.table = defaultdict(Entry)
+
+
+def register_to_doc(name: str):
+    """Register the to_doc method for python AST node type.
+
+    Parameters
+    ----------
+    name : str
+        The type of python AST node.
+
+    Returns
+    -------
+    f : Callable[[FnToDoc], None]
+        The function of registering the to_doc method for python AST node type.
+    """
+
+    def f(to_doc: FnToDoc):  # pylint: disable=redefined-outer-name
+        reg = Registry._inst  # pylint: disable=protected-access
+        reg.table[name].to_doc = to_doc
+
+    return f
+
+
+def register_from_doc(name: str):
+    """Register the from_doc method for python AST node type.
+
+    Parameters
+    ----------
+    name : str
+        The type of python AST node.
+
+    Returns
+    -------
+    f : Callable[[FnFromDoc], None]
+        The function of registering the from_doc method for python AST node type.
+    """
+
+    def f(to_doc: FnFromDoc):  # pylint: disable=redefined-outer-name
+        reg = Registry._inst  # pylint: disable=protected-access
+        reg.table[name].from_doc = to_doc
+
+    return f
+
+
+def _is_atomic_type(node):
+    return (
+        node is None
+        or node in [..., True, False]
+        or isinstance(
+            node,
+            int | float | str | bool | bytes | complex,
+        )
+    )
+
+
+def _get_registry_entry(cls_name, attr):
+    cls_name = cls_name.split(".")[-1]
+    reg = Registry._inst  # pylint: disable=protected-access
+    if cls_name in reg.table:
+        entry = reg.table[cls_name]
+        return getattr(entry, attr, None)
+    return None
+
+
+def from_doc(node):
+    """Get original python AST node from doc AST node.
+
+    Parameters
+    ----------
+    node : doc.AST
+        The doc AST node.
+
+    Returns
+    -------
+    res : ast.AST
+        The corresponding AST node.
+    """
+    if _is_atomic_type(node):
+        return node
+    if isinstance(node, tuple):
+        return tuple(from_doc(n) for n in node)
+    if isinstance(node, list):
+        return [from_doc(n) for n in node]
+    func = _get_registry_entry(node.__class__.__name__, "from_doc")
+    if not func:
+        raise NotImplementedError(f"from_doc is not implemented for: {node.__class__.__name__}")
+    return func(node)
+
+
+def to_doc(node):
+    """Get doc AST node from python AST node.
+
+    Parameters
+    ----------
+    node : ast.AST
+        The AST node.
+
+    Returns
+    -------
+    res : doc.AST
+        The corresponding doc AST node.
+    """
+    if _is_atomic_type(node):
+        return node
+    if isinstance(node, tuple):
+        return tuple(to_doc(n) for n in node)
+    if isinstance(node, list):
+        return [to_doc(n) for n in node]
+    func = _get_registry_entry(node.__class__.__name__, "to_doc")
+    if not func:
+        raise NotImplementedError(f"to_doc is not implemented for: {node.__class__.__name__}")
+    return func(node)
+
+
+def parse(
+    source: str,
+    filename: str = "<unknown>",
+    mode: str = "exec",
+) -> doc.AST:
+    """Parse TVMScript source code str to doc AST.
+
+    Its interface is consistent with python built-in ast.parse.
+    And it will parse by python 3.8 first if possible,
+    or it will parse with python version in current environment.
+
+    Parameters
+    ----------
+    source : str
+        The TVMScript source code.
+
+    filename : str
+        The optional filename of the file where source code locates.
+
+    mode : str
+        The parsing mode for ast.parse.
+
+    Returns
+    -------
+    res : doc.AST
+        The parsed doc AST.
+    """
+    try:
+        program = ast.parse(  # pylint: disable=unexpected-keyword-arg
+            source=source,
+            filename=filename,
+            mode=mode,
+            feature_version=(3, 8),
+        )
+    except:  # pylint: disable=bare-except
+        program = ast.parse(
+            source=source,
+            filename=filename,
+            mode=mode,
+        )
+    return to_doc(program)
+
+
+class NodeVisitor:
+    """Node visitor for doc AST"""
+
+    def visit(self, node: doc.AST) -> None:
+        if isinstance(node, list | tuple):
+            for item in node:
+                self.visit(item)
+            return
+        if not isinstance(node, doc.AST):
+            return
+        getattr(
+            self,
+            "visit_" + node.__class__.__name__.split(".")[-1],
+            self.generic_visit,
+        )(node)
+
+    def generic_visit(self, node: doc.AST) -> None:
+        for field in node.__class__._FIELDS:  # pylint: disable=protected-access
+            value = getattr(node, field, None)
+            if value is None:
+                pass
+            elif isinstance(value, doc.AST | list | tuple):
+                self.visit(value)
+
+
+class NodeTransformer:
+    """Node transformer for doc AST"""
+
+    def visit(self, node: doc.AST) -> doc.AST:
+        if isinstance(node, list):
+            return [self.visit(item) for item in node]
+        if isinstance(node, tuple):
+            return tuple(self.visit(item) for item in node)
+        if not isinstance(node, doc.AST):
+            return node
+        return getattr(
+            self,
+            "visit_" + node.__class__.__name__.split(".")[-1],
+            self.generic_visit,
+        )(node)
+
+    def generic_visit(self, node: doc.AST) -> doc.AST:
+        kv: dict[str, typing.Any] = {}
+        for field in node.__class__._FIELDS:  # pylint: disable=protected-access
+            value = getattr(node, field, None)
+            if value is None:
+                pass
+            elif isinstance(value, doc.AST | list | tuple):
+                value = self.visit(value)
+            kv[field] = value
+        return node.__class__(**kv)
+
+
+def _register_default():
+    class DefaultTranslator:
+        def __init__(self, doc_cls, func, fields):
+            self.doc_cls = doc_cls  # getattr(doc, name)
+            self.func = func
+            self.fields = fields
+
+        def __call__(self, node):
+            kv = {attr: self.func(getattr(node, attr, None)) for attr in self.fields}
+            return self.doc_cls(**kv)
+
+    Registry._inst = Registry()  # pylint: disable=protected-access
+    for cls_name in dir(doc):
+        doc_cls = getattr(doc, cls_name)
+        if not hasattr(ast, cls_name):
+            continue
+        if inspect.isclass(doc_cls) and issubclass(doc_cls, doc.AST):
+            assert "." not in cls_name
+            register_to_doc(cls_name)(
+                DefaultTranslator(
+                    getattr(doc, cls_name),
+                    to_doc,
+                    doc_cls._FIELDS,  # pylint: disable=protected-access
+                )
+            )
+            register_from_doc(cls_name)(
+                DefaultTranslator(
+                    getattr(ast, cls_name),
+                    from_doc,
+                    doc_cls._FIELDS,  # pylint: disable=protected-access
+                )
+            )
+
+
+_register_default()

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

@@ -0,0 +1,210 @@
+# 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."""
+
+import inspect
+from typing import Any
+
+import tvm
+
+from ....ir.module import IRModule
+from ...ir_builder import IRBuilder
+from . import doc
+from .diagnostics import Source
+from .error import ParserError
+from .parser import Parser
+
+WELL_FORMED_ERROR_MESSAGE = (
+    "Program is not well-formed. If this is deliberate, consider "
+    "setting check_well_formed in the top-level decorator to False "
+    "(e.g., @I.ir_module(check_well_formed=False) or "
+    "@R.function(check_well_formed=False))."
+)
+
+
+def _default_globals() -> dict[str, Any]:
+    # lazy import here to avoid circular deps
+    from tvm.script import tirx as _tirx_dsl  # pylint: disable=import-outside-toplevel
+    from tvm.script.parser import (
+        ir,  # pylint: disable=import-outside-toplevel
+        relax,  # pylint: disable=import-outside-toplevel
+    )
+    from tvm.script.parser import tirx as _tirx_parser  # pylint: disable=import-outside-toplevel
+    from tvm.script.tirx import tile as _tirx_tile  # pylint: disable=import-outside-toplevel
+    from tvm.tirx import layout as _tirx_layout  # pylint: disable=import-outside-toplevel
+
+    # Expose the layout `Axis` class so printed layout sugar like
+    # `4 @ Axis.laneid` round-trips without per-script imports. Injecting just
+    # `Axis` (one short symbol) avoids name collisions with common user shape
+    # vars like `m`, `P`, `F` that registered axes happen to share names with.
+    return {
+        "tvm": tvm,
+        "I": ir,
+        "ir": ir,
+        "T": _tirx_parser,
+        "tir": _tirx_parser,
+        "R": relax,
+        "relax": relax,
+        "Tx": _tirx_tile,
+        "tirx": _tirx_dsl,
+        "Axis": _tirx_layout.Axis,
+    }
+
+
+def scan_macro(program: Any | str, extra_vars: dict[str, Any] | None = None) -> Any:
+    """Generate the AST, and the source code for __repr__."""
+    # The AST will be converted into TIR at the time of expansion.
+    source = Source(program)
+    closure_vars = extra_vars or _default_globals()
+    return source, closure_vars
+
+
+def parse(
+    program: doc.AST | Any | str,
+    extra_vars: dict[str, Any] | None = None,
+    check_well_formed: bool = True,
+    s_tir: bool = False,
+) -> Any:
+    """Register a method for a operand type, AST operator node and operand index.
+
+    Parameters
+    ----------
+    program : Union[doc.AST, Any, str]
+        The TVMScript code to parse.
+
+    extra_vars : Dict[str, Any]
+        The extra variable table for parsing.
+
+    check_well_formed : bool
+        Whether to check well-formedness after parsing.
+
+    Returns
+    -------
+    func : Any
+        The parsed TVMScript program.
+    """
+    if extra_vars is None:
+        extra_vars = _default_globals()
+
+    ann = {}
+    all_pyfuncs = {}
+    if inspect.isfunction(program):
+        ann = {program.__name__: program.__annotations__}
+    elif inspect.isclass(program):
+        for name, func in program.__dict__.items():
+            if inspect.isfunction(func):
+                ann[name] = func.__annotations__
+                all_pyfuncs[name] = func
+
+    source = Source(program)
+    parser = Parser(source, ann)
+    with IRBuilder() as builder:
+        try:
+            parser.parse(extra_vars=extra_vars)
+        except ParserError as err:
+            parser.report_error(err.node, err.args[0])
+    ret = builder.get()
+    # Attach pyfuncs to the IRModule
+    if inspect.isclass(program) and isinstance(ret, IRModule):
+        _attach_pyfuncs_to_irmodule(ret, all_pyfuncs)
+
+    # check well-formedness in both Relax and TIR
+    if check_well_formed:
+        check_ret = ret
+        if not isinstance(check_ret, IRModule):
+            check_ret = IRModule.from_expr(ret)
+
+        source_ast = source.as_ast()
+
+        if isinstance(
+            ret, IRModule | tvm.relax.Function
+        ) and not tvm.relax.analysis.check_well_formed(ret):
+            parser.report_error(source_ast, err=WELL_FORMED_ERROR_MESSAGE)
+
+        try:
+            if s_tir:
+                tvm.tirx.analysis.verify_well_formed(check_ret)
+            else:
+                tvm.tirx.analysis.verify_tirx_well_formed(check_ret)
+        except Exception as err:  # pylint: disable=broad-exception-caught
+            parser.report_error(
+                source_ast,
+                err=f"{WELL_FORMED_ERROR_MESSAGE}\n\nTraceback: {err!s}",
+            )
+    return ret
+
+
+def _create_python_packed_func(pyfunc):
+    """Create a PackedFunc wrapper for a Python function.
+
+    This function creates a PackedFunc that can be called from TVM runtime
+    and will execute the original Python function.
+
+    Parameters
+    ----------
+    pyfunc : Callable
+        The Python function to wrap.
+
+    Returns
+    -------
+    PackedFunc
+        A PackedFunc that wraps the Python function.
+    """
+
+    def packed_func_wrapper(*args, **kwargs):
+        """Wrapper function that calls the original Python function."""
+        try:
+            result = pyfunc(*args, **kwargs)
+            return result
+        except Exception as error:
+            print(f"Error calling Python function {pyfunc.__name__}: {error}")
+            raise
+
+    return packed_func_wrapper
+
+
+def _attach_pyfuncs_to_irmodule(irmodule, all_pyfuncs):
+    """Attach Python functions to IRModule with reduced nesting."""
+    if not all_pyfuncs:
+        return
+
+    if not hasattr(irmodule, "pyfuncs"):
+        irmodule.pyfuncs = {}
+
+    for global_var, func in irmodule.functions_items():
+        if not isinstance(func, tvm.relax.ExternFunc):
+            continue
+        if not func.attrs.get("is_pyfunc", False):
+            continue
+
+        pyfunc_name = global_var.name_hint
+        if pyfunc_name not in all_pyfuncs:
+            continue
+
+        pyfunc = all_pyfuncs[pyfunc_name]
+        irmodule.pyfuncs[pyfunc_name] = pyfunc
+
+        try:
+            source_code = inspect.getsource(pyfunc)
+            func = func.with_attr("python_source", source_code)
+        except (OSError, TypeError):
+            func = func.with_attr("python_source", f"# Source unavailable for {pyfunc_name}")
+
+        packed_func = _create_python_packed_func(pyfunc)
+        func = func.with_attr("python_packed_func", packed_func)
+
+        irmodule[global_var] = func

python/tvm/script/parser/core/error.py

@@ -0,0 +1,27 @@
+# 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.
+"""Error classes for diagnostics."""
+
+from . import doc
+
+
+class ParserError(Exception):
+    """Error class for diagnostics."""
+
+    def __init__(self, node: doc.AST, msg: str):
+        super().__init__(msg)
+        self.node = node

python/tvm/script/parser/core/evaluator.py

@@ -0,0 +1,604 @@
+# 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.
+"""AST Evaluation"""
+
+import ast
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+import tvm
+
+from . import dispatch, doc
+from .error import ParserError
+
+if TYPE_CHECKING:
+    from .parser import Parser
+
+DEFAULT_OP: dict[type, Callable[..., Any]] = {
+    doc.Add: lambda a, b: a + b,
+    doc.Sub: lambda a, b: a - b,
+    doc.Mult: lambda a, b: a * b,
+    doc.Div: lambda a, b: a / b,
+    doc.FloorDiv: lambda a, b: a // b,
+    doc.Mod: lambda a, b: a % b,
+    doc.LShift: lambda a, b: a << b,
+    doc.RShift: lambda a, b: a >> b,
+    doc.BitOr: lambda a, b: a | b,
+    doc.BitXor: lambda a, b: a ^ b,
+    doc.BitAnd: lambda a, b: a & b,
+    doc.MatMult: lambda a, b: a @ b,
+    doc.Pow: lambda a, b: a**b,
+    doc.Eq: lambda a, b: a == b,
+    doc.NotEq: lambda a, b: a != b,
+    doc.Lt: lambda a, b: a < b,
+    doc.LtE: lambda a, b: a <= b,
+    doc.Gt: lambda a, b: a > b,
+    doc.GtE: lambda a, b: a >= b,
+    doc.Is: lambda a, b: a is b,
+    doc.IsNot: lambda a, b: a is not b,
+    doc.In: lambda a, b: a in b,
+    doc.NotIn: lambda a, b: a not in b,
+    doc.And: lambda a, b: a and b,
+    doc.Or: lambda a, b: a or b,
+    doc.Invert: lambda a: ~a,
+    doc.Not: lambda a: not a,
+    doc.UAdd: lambda a: +a,
+    doc.USub: lambda a: -a,
+}
+
+
+def _get_builtin_or_none(name: str):
+    builtins = globals().get("__builtins__")
+    if not builtins:
+        return None
+    if not isinstance(builtins, dict) and hasattr(builtins, "__dict__"):
+        builtins = builtins.__dict__
+    if isinstance(builtins, dict):
+        return builtins.get(name)
+    return None
+
+
+class ExprEvaluator:
+    """Expression evaluator for TVMScript parser.
+
+    Parameters
+    ----------
+    parser : Parser
+        The parser bound with the evaluator.
+
+    value_table : Dict[str, Any]
+        The value table for expression evaluation.
+
+    new_value_count : int
+        The count for intermediate result added during evaluation.
+    """
+
+    parser: "Parser"
+    value_table: dict[str, Any]
+    new_value_count: int
+
+    def __init__(self, parser: "Parser", value_table: dict[str, Any]) -> None:
+        super().__init__()
+        self.parser = parser
+        self.value_table = value_table
+        self.new_value_count = 0
+
+    @staticmethod
+    def eval(parser: "Parser", value_table: dict[str, Any], node: doc.AST) -> Any:
+        """Expression evaluation for TVMScript parser.
+
+        Parameters
+        ----------
+        parser : Parser
+            The parser bound with the evaluator.
+
+        value_table : Dict[str, Any]
+            The value table for expression evaluation.
+
+        node : doc.AST
+            The root node of AST tree node of expression to evaluate.
+
+        Returns
+        -------
+        res : Any
+            The evaluation result.
+        """
+        self = ExprEvaluator(parser, value_table)
+        result = self._visit(node)  # pylint: disable=protected-access
+        if isinstance(result, doc.Name):
+            if result.id in self.value_table:
+                return self.value_table[result.id]
+            else:
+                builtin = _get_builtin_or_none(result.id)
+                if builtin:
+                    return builtin
+                raise ParserError(result, f"Undefined variable: {result.id}")
+        if isinstance(result, doc.Constant):
+            return result.value
+        raise TypeError(f"Unexpected result type: {type(result)}")
+
+    def _add_intermediate_result(self, value: Any) -> doc.Name:
+        """Add intermediate result during evaluation into value table.
+
+        Parameters
+        ----------
+        value : Any
+            The intermediate result.
+
+        Returns
+        -------
+        name : doc.Name
+            The doc AST name node with intermediate name for intermediate result.
+        """
+        name = f"__tvm_tmp_value_{self.new_value_count}"
+        self.new_value_count += 1
+        self.value_table[name] = value
+        lineno = 0
+        col_offset = 0
+        return doc.Name(
+            id=name,
+            ctx=doc.Load(),
+            lineno=lineno,
+            col_offset=col_offset,
+            end_lineno=None,
+            end_col_offset=None,
+        )
+
+    def _visit(self, node: doc.AST) -> Any:
+        """General doc AST node visiting method for expression evaluation.
+
+        Parameters
+        ----------
+        node : doc.AST
+            The root node of AST tree node of expression to evaluate.
+
+        Returns
+        -------
+        res : Any
+            The evaluation result.
+        """
+        args = []
+        if (
+            isinstance(node, doc.Call)
+            and hasattr(node.func, "attr")
+            and node.func.attr not in ["reads", "writes", "match_buffer"]
+        ) or isinstance(node, doc.BinOp | doc.UnaryOp | doc.Compare | doc.BoolOp | doc.IfExp):
+            if isinstance(node, doc.BinOp):
+                args = [node.left, node.right]
+            elif isinstance(node, doc.UnaryOp):
+                args = [node.operand]
+            elif isinstance(node, doc.Compare):
+                args = [node.left, *node.comparators]
+            elif isinstance(node, doc.IfExp):
+                args = [node.test, node.body, node.orelse]
+            elif isinstance(node, doc.Call):
+                args = node.args
+            elif isinstance(node, doc.BoolOp):
+                args = node.values
+        for arg in args:
+            if isinstance(arg, doc.Subscript) and isinstance(arg.slice, doc.Slice | doc.Tuple):
+                if isinstance(arg.slice, doc.Slice):
+                    check_slices = [arg.slice]
+                else:
+                    check_slices = []
+                    for p in arg.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,
+                        )
+        if isinstance(node, list):
+            return [self._visit(n) for n in node]
+        if isinstance(node, tuple):
+            return tuple(self._visit(n) for n in node)
+        assert isinstance(node, doc.AST)
+        if isinstance(node, doc.Name):
+            if node.id not in self.value_table and not _get_builtin_or_none(node.id):
+                raise ParserError(node, f"Undefined variable: {node.id}")
+            return node
+        if isinstance(
+            node,
+            doc.Constant | doc.expr_context | doc.operator | doc.boolop | doc.unaryop | doc.cmpop,
+        ):
+            return node
+        if isinstance(node, doc.keyword):
+            return doc.keyword(arg=node.arg, value=self._visit(node.value))
+        if not isinstance(node, doc.expr | doc.Slice):
+            return node
+        if isinstance(node, doc.Lambda):
+            return self._eval_lambda(node)
+        if isinstance(node, doc.Starred):
+            value = self._visit(node.value)
+            return doc.Starred(
+                value=value,
+                ctx=node.ctx,
+                lineno=node.lineno,
+                col_offset=node.col_offset,
+                end_lineno=node.end_lineno,
+                end_col_offset=node.end_col_offset,
+            )
+
+        if isinstance(node, doc.ListComp | doc.SetComp | doc.DictComp):
+            value = self._eval_expr(node)
+            return self._add_intermediate_result(value)
+
+        fields = {}
+        for field in node.__class__._FIELDS:  # pylint: disable=protected-access
+            attr = getattr(node, field)
+            if isinstance(attr, doc.AST | tuple | list):
+                fields[field] = self._visit(attr)
+            else:
+                fields[field] = attr
+        try:
+            if isinstance(node, doc.BoolOp):
+                value = self._eval_bool_op(fields)
+            elif isinstance(node, doc.Compare):
+                value = self._eval_compare(fields)
+            elif isinstance(node, doc.UnaryOp):
+                value = self._eval_unary_op(fields)
+            elif isinstance(node, doc.BinOp):
+                value = self._eval_bin_op(fields)
+            elif isinstance(node, doc.IfExp):
+                value = self._eval_if_exp(fields)
+            elif isinstance(node, doc.Slice):
+                value = self._eval_slice(fields)
+            else:
+                value = self._eval_expr(node.__class__(**fields))
+        except Exception as err:  # pylint: disable=broad-except
+            self.parser.report_error(node, err)
+        return self._add_intermediate_result(value)
+
+    def _eval_lambda(self, node: doc.Lambda) -> Any:
+        """The doc AST lambda node evaluating method.
+
+        Parameters
+        ----------
+        node : doc.Lambda
+            The root node of AST tree node of expression to evaluate.
+
+        Returns
+        -------
+        res : Any
+            The evaluation result.
+        """
+        try:
+            value = self._eval_expr(node)
+        except Exception as err:  # pylint: disable=broad-except
+            self.parser.report_error(node, err)
+        return self._add_intermediate_result(value)
+
+    def _eval_bool_op(self, fields: dict[str, Any]) -> Any:
+        """The doc AST boolean operator node evaluating method.
+
+        Parameters
+        ----------
+        fields : Dict[str, Any]
+            The dictionary of boolean operation information,
+            e.g., operator types, operand values.
+
+        Returns
+        -------
+        res : Any
+            The evaluation result.
+        """
+        op = fields["op"]
+        if not isinstance(op, doc.And | doc.Or):
+            raise TypeError(f"Unexpected operator: {op}")
+        value = self._eval_expr(fields["values"][0])
+        for rhs in fields["values"][1:]:
+            value = _eval_op(op, values=[value, self._eval_expr(rhs)])
+        return value
+
+    def _eval_compare(self, fields: dict[str, Any]) -> Any:
+        """The doc AST comparison operation node evaluating method.
+
+        Parameters
+        ----------
+        fields : Dict[str, Any]
+            The dictionary of comparison operation information,
+            e.g., operator types, operand values.
+
+        Returns
+        -------
+        res : Any
+            The evaluation result.
+        """
+        values = [self._eval_expr(fields["left"])]
+        values.extend([self._eval_expr(rhs) for rhs in fields["comparators"]])
+        result = None
+        assert len(fields["ops"]) == len(values) - 1
+
+        for index, op in enumerate(fields["ops"]):
+            sub_result = _eval_op(op, values=[values[index], values[index + 1]])
+            if result is None:
+                result = sub_result
+            else:
+                result = _eval_op(doc.And(), values=[result, sub_result])
+        return result
+
+    def _eval_unary_op(self, fields: dict[str, Any]) -> Any:
+        """The doc AST unary operation node evaluating method.
+
+        Parameters
+        ----------
+        fields : Dict[str, Any]
+            The dictionary of unary operation information,
+            e.g., operator types, operand values.
+
+        Returns
+        -------
+        res : Any
+            The evaluation result.
+        """
+        value = self._eval_expr(fields["operand"])
+        value = _eval_op(fields["op"], values=[value])
+        return value
+
+    def _eval_bin_op(self, fields: dict[str, Any]) -> Any:
+        """The doc AST binary operation node evaluating method.
+
+        Parameters
+        ----------
+        fields : Dict[str, Any]
+            The dictionary of binary operation information,
+            e.g., operator types, operand values.
+
+        Returns
+        -------
+        res : Any
+            The evaluation result.
+        """
+        return _eval_op(
+            fields["op"],
+            values=[
+                self._eval_expr(fields["left"]),
+                self._eval_expr(fields["right"]),
+            ],
+        )
+
+    def _eval_if_exp(self, fields: dict[str, Any]) -> Any:
+        """The doc AST if-else expression node evaluating method.
+
+        Parameters
+        ----------
+        fields : Dict[str, Any]
+            The dictionary of if-else expression information,
+            e.g., test, body, orelse.
+
+        Returns
+        -------
+        res : Any
+            The evaluation result.
+        """
+        test = self._eval_expr(fields["test"])
+        body = self._eval_expr(fields["body"])
+        orelse = self._eval_expr(fields["orelse"])
+        if isinstance(test, bool):
+            return body if test else orelse
+        elif tvm.ir.is_prim_expr(test) and test.ty.matches_code(tvm.DataTypeCode.BOOL):
+            return tvm.tirx.op.if_then_else(test, body, orelse)
+        else:
+            raise TypeError(f"Expected Python bool or TIR bool, but got {type(test)}")
+
+    def _eval_slice(self, fields: dict[str, Any]) -> slice:
+        """The doc AST slice node evaluating method.
+
+        Parameters
+        ----------
+        fields : Dict[str, Any]
+            The dictionary of slice information,
+            e.g., lower bound, upper bound, step.
+
+        Returns
+        -------
+        res : slice
+            The evaluation result.
+        """
+        lower, upper, step = fields["lower"], fields["upper"], fields["step"]
+
+        lower = self._eval_expr(lower) if lower is not None else None
+        upper = self._eval_expr(upper) if upper is not None else None
+        step = self._eval_expr(step) if step is not None else None
+
+        return slice(lower, upper, step)
+
+    def _eval_expr(self, v: Any) -> Any:
+        """The doc AST expression node evaluating method.
+
+        Parameters
+        ----------
+        v : Any
+            The root node of AST tree node of expression to evaluate.
+
+        Returns
+        -------
+        res : Any
+            The evaluation result.
+        """
+        return _eval_expr(v, self.value_table)
+
+
+def eval_expr(
+    parser: "Parser",
+    node: doc.expr | doc.Expression,
+    dict_globals: dict[str, Any] | None,
+) -> Any:
+    """Expression evaluation for TVMScript parser.
+
+    Parameters
+    ----------
+    parser : Parser
+        The parser bound with the evaluator.
+
+    node : Union[doc.expr, doc.Expression]
+        The root node of AST tree node of expression to evaluate.
+
+    dict_globals : Optional[Dict[str, Any]]
+        The optional global value table for expression evaluation.
+
+    Returns
+    -------
+    res : Any
+        The evaluation result.
+    """
+    value_table = {}
+    if dict_globals is not None:
+        value_table.update(dict_globals)
+    return ExprEvaluator.eval(parser, value_table, node)
+
+
+def eval_assign(
+    parser: "Parser",
+    target: doc.expr,
+    source: Any,
+) -> dict[str, Any]:
+    """Expression assignment evaluation for TVMScript parser.
+
+    Parameters
+    ----------
+    parser : Parser
+        The parser bound with the evaluator.
+
+    target : doc.expr
+        The root node of AST tree node of assigned expression to evaluate.
+
+    source : Any
+        The source to be assigned with evaluated expression.
+
+    Returns
+    -------
+    res : Any
+        The evaluation result.
+    """
+    try:
+        return _eval_assign(target, source)
+    except Exception as err:  # pylint: disable=broad-except
+        parser.report_error(target, err)
+        raise
+
+
+def _eval_expr(
+    node: doc.expr | doc.Expression,
+    dict_globals: dict[str, Any] | None,
+) -> Any:
+    """Expression evaluation implementation for TVMScript parser.
+
+    Parameters
+    ----------
+    node : Union[doc.expr, doc.Expression]
+        The root node of AST tree node of expression to evaluate.
+
+    dict_globals : Optional[Dict[str, Any]]
+        The optional global value table for expression evaluation.
+
+    Returns
+    -------
+    res : Any
+        The evaluation result.
+    """
+    node = doc.from_doc(node)
+    if isinstance(node, ast.expr):
+        node = ast.Expression(body=node)
+    assert isinstance(node, ast.Expression), "Expects an ast.Expression, but gets: " + str(node)
+    if dict_globals is None:
+        dict_globals = {}
+    node = ast.fix_missing_locations(node)
+    exe = compile(node, filename="<ast>", mode="eval")
+    return eval(exe, dict_globals)  # pylint: disable=eval-used
+
+
+def _eval_op(
+    op: doc.AST,
+    values: list[Any],
+):
+    """Operation expression evaluation implementation for TVMScript parser.
+
+    Parameters
+    ----------
+    op : doc.AST
+        The root node of AST tree node of operation expression to evaluate.
+
+    values : List[Any]
+        The list of values of operands.
+
+    Returns
+    -------
+    res : Any
+        The evaluation result.
+    """
+    op_type = type(op)  # pylint: disable=protected-access
+    for i, v in enumerate(values):
+        v_type = getattr(type(v), "_dispatch_type", None)
+        if v_type is None:
+            continue
+        f = dispatch.get_op(
+            operand_type=v_type, op_node_type=op_type, operand_index=i, default=None
+        )
+        if f is not None:
+            return f(*values)
+    return DEFAULT_OP[op_type](*values)
+
+
+def _eval_assign(
+    target: doc.expr,
+    source: Any,
+) -> dict[str, Any]:
+    """Expression assignment evaluation implementation for TVMScript parser.
+
+    Parameters
+    ----------
+    target : doc.expr
+        The root node of AST tree node of assigned expression to evaluate.
+
+    source : Any
+        The source to be assigned with evaluated expression.
+
+    Returns
+    -------
+    res : Any
+        The evaluation result.
+    """
+    target = doc.from_doc(target)
+    assert isinstance(target, ast.expr)
+    RHS_VAR_NAME = "__tvm_rhs_var__"  # pylint: disable=invalid-name
+    rhs_var_name = RHS_VAR_NAME
+    dict_locals = {rhs_var_name: source}
+    mod = ast.fix_missing_locations(
+        ast.Module(
+            body=[
+                ast.Assign(
+                    targets=[target],
+                    value=ast.Name(
+                        id=rhs_var_name,
+                        ctx=ast.Load(),
+                    ),
+                )
+            ],
+            type_ignores=[],
+        )
+    )
+    exe = compile(mod, filename="<ast>", mode="exec")
+    exec(exe, {}, dict_locals)  # pylint: disable=exec-used
+    del dict_locals[rhs_var_name]
+    return dict_locals

python/tvm/script/parser/core/parser.py

@@ -0,0 +1,952 @@
+# 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 core parser"""
+
+import abc
+import inspect
+from collections import defaultdict
+from collections.abc import Callable
+from contextlib import contextmanager
+from typing import Any
+
+import numpy as np
+
+from tvm.error import DiagnosticError
+from tvm.ir import GlobalVar
+
+from . import dispatch, doc
+from .diagnostics import Diagnostics, Source
+from .evaluator import eval_assign, eval_expr
+
+DEFAULT_VISIT = {
+    "Interactive",
+    "Module",
+    "Expression",
+    "Pass",
+}
+
+
+def _deferred(exit_f: Callable[[], None]):
+    """Created context with certain exit function.
+
+    Parameters
+    ----------
+    exit_f : Callable[[], None]
+        The function to call when exiting the context.
+
+    Returns
+    -------
+    res : Any
+        The created context.
+    """
+
+    @contextmanager
+    def context():
+        try:
+            yield
+        finally:
+            exit_f()
+
+    return context()
+
+
+def _do_nothing(*args, **kwargs):  # pylint: disable=unused-argument
+    pass
+
+
+class ScriptMacro(abc.ABC):
+    """Representation of a script macro.
+
+    This is a callable object, intended to be called from the expression evaluator.
+    The evaluator is expected to insert the current parser into the environment
+    undef the name given by "parser_object_name".
+
+    Once called, the ScriptMacro object will locate the current parser, and use it
+    to parse the macro's body and produce the result.
+
+    There were two major considerations for this design:
+    1. Implementing hygienic and non-hygienic macros.
+    2. Implementing macros that return values.
+
+    Macro uses in TIR are only allowed at a statement-level, and they don't produce
+    any values. Parsing of such macros could easily be done by intercepting doc.Call
+    nodes in the TIR parser. If a macro is a value-producing expression, then there
+    may not be a direct way to intercept calls to it if it's embedded in a complex
+    expression. Because macros use function-call syntax, the evaluator will try to
+    call the macro object, which this design relies on to parse and evaluate the macro.
+    """
+
+    parser_object_name = "__current_script_parser__"
+
+    def __init__(
+        self,
+        source: Source,
+        closure_vars: dict[str, Any],
+        func: Callable,
+        hygienic: bool,
+    ) -> None:
+        self.source = source
+        self.closure_vars = closure_vars
+        self.func = func
+        self.hygienic = hygienic
+
+    def __repr__(self):
+        return self.source.source
+
+    @abc.abstractmethod
+    def parse_macro(self, parser: "Parser") -> Any:
+        """The main macro parsing function. Different scripts may have different
+        ways to parse a macro, and to return a value to the evaluator.
+
+        Parameters
+        ----------
+        parser : Parser
+            The parser with the appropriate frame already created and populated depending
+            macro's hygiene settings,
+
+        Returns
+        -------
+            The return value depends on the specifics of the particular script. It can be
+            "None" or any other value or any type.
+        """
+
+    def _find_parser_def(self):
+        outer_frame_infos = inspect.getouterframes(inspect.currentframe())
+        for finfo in outer_frame_infos:
+            parser = finfo.frame.f_globals.get(ScriptMacro.parser_object_name)
+            if parser is not None:
+                return parser
+        raise RuntimeError(f"{ScriptMacro.parser_object_name} not available")
+
+    def get_macro_def(self):
+        ast_module = self.source.as_ast()
+        for decl in ast_module.body:
+            if isinstance(decl, doc.FunctionDef) and decl.name == self.func.__name__:
+                return decl
+        raise RuntimeError(f"cannot find macro definition for {self.func.__name__}")
+
+    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.hygienic:
+                saved_var_table = parser.var_table
+                parser.var_table = VarTable()
+
+                with parser.var_table.with_frame():
+                    for k, v in self.closure_vars.items():
+                        parser.var_table.add(k, v)
+                    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
+
+            else:
+                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)
+
+        return parse_result
+
+
+class VarTableFrame:
+    """The variable table frame.
+    A frame of variable table stores the variables created in one block or scope.
+
+    Parameters
+    ----------
+    vars : Set[str]
+        The set of variable names in the variable table frame.
+    """
+
+    vars: set[str]
+
+    def __init__(self):
+        self.vars = set()
+
+    def add(self, var: str):
+        """Add a new variable into variable table frame.
+
+        Parameters
+        ----------
+        var : str
+            The name of new variable.
+        """
+        if var in self.vars:
+            raise ValueError(f"Variable {var} already defined in current scope")
+        self.vars.add(var)
+
+    def pop_all(self, fn_pop: Callable[[str], None]):
+        """Pop out all variable in variable table frame.
+
+        Parameters
+        ----------
+        fn_pop : Callable[[str], None]
+            The methods to call when popping each variable.
+        """
+        for var in self.vars:
+            fn_pop(var)
+        self.vars.clear()
+
+
+class VarTable:
+    """The variable table.
+    A variable table stores the all variables when parsing TVMScript.
+
+    Parameters
+    ----------
+    frames : List[VarTableFrame]
+        The list or stack of variable table frame.
+
+    name2value : Dict[str, List[Any]]
+        The dictionary for variable table name-based query.
+    """
+
+    frames: list[VarTableFrame]
+    name2value: dict[str, list[Any]]
+
+    def __init__(self):
+        self.frames = []
+        self.name2value = defaultdict(list)
+
+    def with_frame(self):
+        """Create a new variable table frame as with statement.
+
+        Returns
+        -------
+        res : Any
+            The context with new variable table frame.
+        """
+
+        def pop_frame():
+            frame = self.frames.pop()
+            frame.pop_all(lambda name: self.name2value[name].pop())
+
+        self.frames.append(VarTableFrame())
+        return _deferred(pop_frame)
+
+    def add(self, var: str, value: Any, allow_shadowing: bool = False):
+        """Add a new variable to variable table.
+
+        Parameters
+        ----------
+        var : str
+            The name of variable.
+
+        value : Any
+            The value of variable.
+
+        allow_shadowing : bool
+            The options of whether variable shadowing allowed for this variable.
+        """
+        # Skip if the key and value are equal to those in the var_table
+        if self.name2value[var] and isinstance(self.name2value[var][-1], type(value)):
+            if isinstance(value, np.ndarray) and (self.name2value[var][-1] == value).all():
+                return
+            elif self.name2value[var][-1] == value:
+                return
+        if allow_shadowing and var in self.frames[-1].vars:
+            # Shadowing
+            self.name2value[var][-1] = value
+        else:
+            self.frames[-1].add(var)
+            self.name2value[var].append(value)
+
+    def get(self) -> dict[str, Any]:
+        """Get a variable dictionary of latest variables.
+
+        Returns
+        -------
+        res : Any
+            The variable dictionary copy of latest variables.
+        """
+        return {key: values[-1] for key, values in self.name2value.items() if values}
+
+    def get_at_depth(self, depth: int) -> dict[str, Any]:
+        """Get variables visible at the given frame depth, using current values.
+
+        For each variable name that appears in frames 0..depth-1, count how many
+        times it was pushed (to handle shadowing), then index into name2value at
+        count-1 to retrieve the latest value visible at that depth.
+
+        Parameters
+        ----------
+        depth : int
+            The frame depth (number of frames visible).
+
+        Returns
+        -------
+        res : dict[str, Any]
+            Variable dictionary of values visible at the given depth.
+        """
+        result: dict[str, Any] = {}
+        name_count: dict[str, int] = defaultdict(int)
+        for frame_idx in range(min(depth, len(self.frames))):
+            for name in self.frames[frame_idx].vars:
+                name_count[name] += 1
+        for name, count in name_count.items():
+            if self.name2value[name]:
+                result[name] = self.name2value[name][count - 1]
+        return result
+
+    def exist(self, value: Any) -> bool:
+        """Check if any value exists in variable table.
+
+        Parameters
+        ----------
+        value : Any
+            The value of variable.
+
+        Returns
+        -------
+        res : bool
+            The existence of the value.
+        """
+        return any(
+            value.same_as(known_value)
+            for known_value_stack in self.name2value.values()
+            for known_value in known_value_stack
+        )
+
+
+def _dispatch_wrapper(func: dispatch.ParseMethod) -> dispatch.ParseMethod:
+    def _wrapper(self: "Parser", node: doc.AST) -> None:
+        try:
+            return func(self, node)
+        except Exception as err:  # pylint: disable=broad-except
+            self.report_error(node, err)
+            raise
+
+    return _wrapper
+
+
+def _dispatch(self: "Parser", type_name: str) -> dispatch.ParseMethod:
+    for token in [self.dispatch_tokens[-1], "default"]:
+        func = dispatch.get(token=token, type_name=type_name, default=None)
+        if func is not None:
+            return _dispatch_wrapper(func)
+    return _dispatch_wrapper(lambda self, node: self.generic_visit(node))
+
+
+class Parser(doc.NodeVisitor):
+    """The TVMScript parser
+
+    Parameters
+    ----------
+    diag : Diagnostics
+        The diagnostics for error reporting.
+
+    dispatch_tokens : List[str]
+        The list of dispatching tokens to dispatching parsing method
+        of different IRs and different doc AST structure.
+
+    var_table : VarTable
+        The variable table for parsing.
+    """
+
+    diag: Diagnostics
+    dispatch_tokens: list[str]
+    function_annotations: dict[str, dict[str, Any]] | None
+    var_table: VarTable
+    inside_function: bool  # whether we are within a function
+    current_class: str | None = None  # current class being parsed
+    base_py_module_context: bool = False  # whether current class inherits from BasePyModule
+
+    def __init__(
+        self,
+        source: Source,
+        function_annotations: dict[str, dict[str, Any]],
+    ) -> None:
+        self.diag = Diagnostics(source)
+        self.dispatch_tokens = ["default"]
+        self.function_annotations = function_annotations
+        self.var_table = VarTable()
+        self.inside_function = False
+
+    def parse(self, extra_vars: dict[str, Any] | None = None) -> Any:
+        """The main parse method for parser.
+
+        Parameters
+        ----------
+        extra_vars : Optional[Dict[str, Any]]
+            The optional global value table for parsing.
+
+        Returns
+        -------
+        res : Any
+            The doc AST node visiting result.
+        """
+        if extra_vars is None:
+            extra_vars = {}
+        with self.var_table.with_frame():
+            for k, v in extra_vars.items():
+                self.var_table.add(k, v)
+            node = self.diag.source.as_ast()
+            self.visit(node)
+
+    def get_dispatch_token(self, node: doc.FunctionDef) -> str:
+        if not isinstance(node, doc.FunctionDef):
+            self.report_error(node, "Only can get dispatch token for function.")
+        if not node.decorator_list:
+            self.report_error(node, "Function must be decorated")
+        # TODO: only the last decorator is parsed
+        decorator = self.eval_expr(node.decorator_list[-1])
+        if not hasattr(decorator, "dispatch_token"):
+            self.report_error(node, "The parser does not understand the decorator")
+        return decorator.dispatch_token
+
+    def with_dispatch_token(self, token: str):
+        """Add a new dispatching token as with statement.
+
+        Parameters
+        ----------
+        token : str
+            The dispatching token.
+
+        Returns
+        -------
+        res : Any
+            The context with new dispatching token.
+        """
+
+        self.dispatch_tokens.append(token)
+        enter_func = dispatch.get(token=token, type_name="enter_token", default=lambda *args: None)
+        context = enter_func(self)
+
+        def pop_token():
+            exit_func = dispatch.get(
+                token=token, type_name="exit_token", default=lambda *args: None
+            )
+            exit_func(self, context)
+            self.dispatch_tokens.pop()
+
+        return _deferred(pop_token)
+
+    def set_class_context(self, class_name: str, is_base_py_module: bool = False):
+        """Set the current class context for parsing.
+
+        Parameters
+        ----------
+        class_name : str
+            The name of the current class being parsed.
+        is_base_py_module : bool
+            Whether the current class inherits from BasePyModule.
+        """
+        self.current_class = class_name
+        self.base_py_module_context = is_base_py_module
+
+    def _get_current_class_context(self) -> str | None:
+        """Get the current class context.
+
+        Returns
+        -------
+        Optional[str]
+            The name of the current class, or None if not in a class context.
+        """
+        return self.current_class
+
+    def _is_base_py_module_context(self) -> bool:
+        """Check if the current class context allows Python functions.
+
+        Returns
+        -------
+        bool
+            True if Python functions are allowed in the current context.
+        """
+        return self.base_py_module_context
+
+    def with_diag_source(self, source: Source):
+        """Add a new source as with statement.
+
+        Parameters
+        ----------
+        source : Source
+            The source for diagnostics.
+
+        Returns
+        -------
+        res : Any
+            The context with new source.
+        """
+
+        last_diag = self.diag
+        self.diag = Diagnostics(source)
+
+        def pop_source():
+            self.diag = last_diag
+
+        return _deferred(pop_source)
+
+    def eval_expr(
+        self,
+        node: doc.Expression | doc.expr,
+        extra_vars: dict[str, Any] | None = None,
+    ) -> Any:
+        """Expression evaluation when parsing.
+
+        Parameters
+        ----------
+        node : Union[doc.expr, doc.Expression]
+            The root node of AST tree node of expression to evaluate.
+
+        extra_vars : Optional[Dict[str, Any]]
+            The optional global value table for expression evaluation.
+
+        Returns
+        -------
+        res : Any
+            The evaluation result.
+        """
+        var_values = self.var_table.get()
+        if extra_vars is not None:
+            for k, v in extra_vars.items():
+                var_values[k] = v
+        var_values[ScriptMacro.parser_object_name] = self
+        return eval_expr(self, node, var_values)
+
+    def _duplicate_lhs_check(self, target: doc.expr) -> bool | set[str]:
+        """Check whether duplicate lhs exists in assignment.
+
+        Parameters
+        ----------
+        target : doc.expr
+            The doc AST expr node for lhs.
+
+        Returns
+        -------
+        res : Union[bool, Set[str]]
+            The result of true if duplicate lhs exists,
+            or the set of lhs names if no duplicate lhs exists.
+        """
+        if isinstance(target, doc.Tuple | doc.List):
+            vars: set[str] = set()  # pylint: disable=redefined-builtin
+            for i in target.elts:
+                res = self._duplicate_lhs_check(i)
+                if isinstance(res, bool) and res:
+                    return True
+                assert isinstance(res, set)
+                if vars & res:
+                    return True
+                vars = vars.union(res)
+            return vars
+        elif isinstance(target, doc.Name):
+            return {target.id}
+        elif isinstance(target, doc.Starred):
+            return self._duplicate_lhs_check(target.value)
+        elif isinstance(target, doc.Attribute):
+            # Attribute assignment like packedB.data = ..., treated as rebinding.
+            return {target.attr}
+        else:
+            self.report_error(target, "Invalid type in assign statement")
+            raise NotImplementedError
+
+    def eval_assign(
+        self,
+        target: doc.expr,
+        source: Any,
+        bind_value: Callable[["Parser", doc.expr, str, Any], Any],
+        allow_shadowing: bool = False,
+    ) -> dict[str, Any]:
+        """Expression assignment evaluation when parsing.
+
+        Parameters
+        ----------
+        target : doc.expr
+            The root node of AST tree node of assigned expression to evaluate.
+
+        source : Any
+            The source to be assigned with evaluated expression.
+
+        bind_value : Callable[["Parser", doc.expr, str, Any], Any]
+            The value binding method when assigning the values to variables.
+
+        allow_shadowing : bool
+            The options of whether variable shadowing allowed for assignment.
+
+        Returns
+        -------
+        res : Dict[str, Any]
+            The dictionary of assignment result.
+        """
+        if self._duplicate_lhs_check(target) is True:
+            self.report_error(target, "Duplicate vars assigned.")
+        var_values = eval_assign(self, target, source)
+        for k, v in var_values.items():
+            var = bind_value(self, target, k, v)
+            self.var_table.add(k, var, allow_shadowing)
+        return var_values
+
+    def report_error(self, node: doc.AST, err: Exception | str) -> None:  # pylint: disable=no-self-use
+        """The error reporting when parsing.
+
+        Parameters
+        ----------
+        node : doc.AST
+            The doc AST node with errors.
+
+        err: Union[Exception, str]
+            The error to report.
+        """
+
+        # If the error is already being raised as a DiagnosticError,
+        # re-raise it without wrapping it in a DiagnosticContext.
+        if isinstance(err, DiagnosticError):
+            raise err
+
+        # Only take the last line of the error message
+        if isinstance(err, RuntimeError):
+            lines = list(filter(None, str(err).split("\n")))
+            msg = lines[-1] if lines else (str(err) or type(err).__name__)
+        elif isinstance(err, KeyError):
+            msg = "KeyError: " + str(err)
+        else:
+            msg = str(err)
+
+        try:
+            self.diag.error(node, msg)
+        except Exception as diag_err:
+            # Calling self.diag.error is guaranteed to throw an
+            # exception.  When shown to a user, this error should
+            # reference the point of error within the provided
+            # TVMScript.  However, when caught in pdb, the full
+            # traceback should be available for debugging.
+            if isinstance(err, Exception):
+                diag_err = diag_err.with_traceback(err.__traceback__)
+            raise diag_err
+
+    def visit(self, node: doc.AST) -> None:
+        """The general visiting method.
+
+        Parameters
+        ----------
+        node : doc.AST
+            The doc AST node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        if isinstance(node, list | tuple):
+            for item in node:
+                self.visit(item)
+            return
+        if not isinstance(node, doc.AST):
+            return
+        name = node.__class__.__name__.split(".")[-1]
+        if name in DEFAULT_VISIT:
+            func = self.generic_visit
+        else:
+            func = getattr(self, "visit_" + name, None)
+        if func is None:
+            raise NotImplementedError(f"Visitor of AST node is not implemented: {name}")
+        try:
+            func(node)
+        except Exception as err:  # pylint: disable=broad-except
+            self.report_error(node, err)
+
+    def visit_body(self, node: list[doc.stmt]) -> Any:
+        """The general body visiting method.
+
+        Parameters
+        ----------
+        node : List[doc.stmt]
+            The list of statements in body.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        for stmt in node:
+            self.visit(stmt)
+
+    def visit_tvm_annotation(self, node: doc.expr) -> Any:
+        """The general TVM annotation visiting method.
+
+        Parameters
+        ----------
+        node : doc.expr
+            The doc AST expr node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "tvm_annotation")(self, node)
+
+    def visit_FunctionDef(self, node: doc.FunctionDef) -> None:  # pylint: disable=invalid-name
+        """The general function definition visit method.
+
+        Parameters
+        ----------
+        node : doc.FunctionDef
+            The doc FunctionDef node.
+        """
+        token = self.get_dispatch_token(node)
+        func = dispatch.get(token=token, type_name="FunctionDef", default=None)
+        if func is None:
+            self.report_error(
+                node,
+                """The parser does not understand the decorator,
+                or visit_FunctionDef is not implemented for the decorator with token: """
+                + token,
+            )
+        _dispatch(self, "pre_visit_local_function")(self, node)
+        _dispatch_wrapper(func)(self, node)
+        _dispatch(self, "post_visit_local_function")(self, node)
+
+    def visit_tvm_declare_function(self, node: doc.FunctionDef) -> GlobalVar:
+        token = self.get_dispatch_token(node)
+        with self.with_dispatch_token(token):
+            return _dispatch(self, "tvm_declare_function")(self, node)
+
+    def visit_ClassDef(self, node: doc.ClassDef) -> Any:  # pylint: disable=invalid-name
+        """The general class definition visiting method.
+
+        Parameters
+        ----------
+        node : doc.ClassDef
+            The doc AST class definition node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        func = dispatch.get(token="ir", type_name="ClassDef", default=None)
+        if func is None:
+            self.report_error(node, "The parser does not understand the decorator")
+        _dispatch_wrapper(func)(self, node)
+
+    def visit_arguments(self, node: doc.arguments) -> Any:
+        """The general arguments visiting method.
+
+        Parameters
+        ----------
+        node : doc.arguments
+            The doc AST arguments node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "arguments")(self, node)
+
+    def visit_For(self, node: doc.For) -> Any:  # pylint: disable=invalid-name
+        """The general for visiting method.
+
+        Parameters
+        ----------
+        node : doc.For
+            The doc AST for node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "For")(self, node)
+
+    def visit_While(self, node: doc.While) -> Any:  # pylint: disable=invalid-name
+        """The general while visiting method.
+
+        Parameters
+        ----------
+        node : doc.While
+            The doc AST while node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "While")(self, node)
+
+    def visit_With(self, node: doc.With) -> Any:  # pylint: disable=invalid-name
+        """The general with visiting method.
+
+        Parameters
+        ----------
+        node : doc.With
+            The doc AST with node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "With")(self, node)
+
+    def visit_Assign(self, node: doc.Assign) -> Any:  # pylint: disable=invalid-name
+        """The general assign visiting method.
+
+        Parameters
+        ----------
+        node : doc.Assign
+            The doc AST assign node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "Assign")(self, node)
+
+    def visit_AnnAssign(self, node: doc.AnnAssign) -> Any:  # pylint: disable=invalid-name
+        """The general annotated assign visiting method.
+
+        Parameters
+        ----------
+        node : doc.Assign
+            The doc AST annotated assign node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "AnnAssign")(self, node)
+
+    def visit_Expr(self, node: doc.Expr) -> Any:  # pylint: disable=invalid-name
+        """The general expression visiting method.
+
+        Parameters
+        ----------
+        node : doc.Expr
+            The doc AST expression node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "Expr")(self, node)
+
+    def visit_If(self, node: doc.If) -> Any:  # pylint: disable=invalid-name
+        """The general if visiting method.
+
+        Parameters
+        ----------
+        node : doc.If
+            The doc AST if node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "If")(self, node)
+
+    def visit_AugAssign(self, node: doc.AugAssign) -> Any:  # pylint: disable=invalid-name
+        """The general augmented assignment visiting method.
+
+        Parameters
+        ----------
+        node : doc.AugAssign
+            The doc AST augmented assignment node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "AugAssign")(self, node)
+
+    def visit_Assert(self, node: doc.Assert) -> Any:  # pylint: disable=invalid-name
+        """The general assert visiting method.
+
+        Parameters
+        ----------
+        node : doc.Assert
+            The doc AST assert node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "Assert")(self, node)
+
+    def visit_Return(self, node: doc.Return) -> Any:  # pylint: disable=invalid-name
+        """The general return visiting method.
+
+        Parameters
+        ----------
+        node : doc.Return
+            The doc AST return node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "Return")(self, node)
+
+    def visit_Continue(self, node: doc.Continue) -> Any:  # pylint: disable=invalid-name
+        """The general continue visiting method.
+
+        Parameters
+        ----------
+        node : doc.Continue
+            The doc AST continue node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "Continue")(self, node)
+
+    def visit_Break(self, node: doc.Break) -> Any:  # pylint: disable=invalid-name
+        """The general break visiting method.
+
+        Parameters
+        ----------
+        node : doc.Break
+            The doc AST break node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "Break")(self, node)
+
+    def visit_Nonlocal(self, node: doc.Nonlocal) -> Any:  # pylint: disable=invalid-name
+        """The general nonlocal visiting method.
+
+        Parameters
+        ----------
+        node : doc.Nonlocal
+            The doc AST nonlocal node.
+
+        Returns
+        -------
+        res : Any
+            The visiting result.
+        """
+        return _dispatch(self, "Nonlocal")(self, node)

python/tvm/script/parser/core/utils.py

@@ -0,0 +1,212 @@
+# 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.
+"""TVM Script Parser utils"""
+
+import inspect
+from collections.abc import Callable
+from types import FrameType
+from typing import Any
+
+from .diagnostics import findsource
+
+
+def get_func_nonlocals(func):
+    """A modified version of `inspect.getclosurevars`"""
+
+    if inspect.ismethod(func):
+        func = func.__func__
+
+    if not inspect.isfunction(func):
+        raise TypeError(f"{func!r} is not a Python function")
+
+    code = func.__code__
+    # Nonlocal references are named in co_freevars and resolved
+    # by looking them up in __closure__ by positional index
+    nonlocal_vars = {}
+    if func.__closure__ is not None:
+        for var, cell in zip(code.co_freevars, func.__closure__):
+            try:
+                nonlocal_vars[var] = cell.cell_contents
+            except ValueError as err:
+                # cell_contents may raise ValueError if the cell is empty.
+                if "empty" not in str(err):
+                    raise
+    return nonlocal_vars
+
+
+def inspect_function_capture(func: Callable) -> dict[str, Any]:
+    """Capture function non-locals and global variables.
+
+    Parameters
+    ----------
+    func : Callable
+        The function to inspect.
+
+    Returns
+    -------
+    res : Dict[str, Any]
+        The function variables map with non-local or global variables.
+    """
+    captured = {
+        **func.__globals__,  # type: ignore
+        **get_func_nonlocals(func),
+    }
+    return captured
+
+
+def inspect_class_capture(cls: type) -> dict[str, Any]:
+    """Capture class non-locals and global variables.
+
+    Parameters
+    ----------
+    cls : type
+        The class to inspect.
+
+    Returns
+    -------
+    res : Dict[str, Any]
+        The class variables map with non-local or global variables.
+    """
+    result: dict[str, Any] = {}
+    for _, v in cls.__dict__.items():
+        if inspect.isfunction(v):
+            func_vars = inspect_function_capture(v)
+            result.update(**func_vars)
+    return result
+
+
+def _collect_annotation_names(source_obj: type | Callable) -> set[str]:
+    """Parse source AST to find names used in function annotations.
+
+    Returns the set of ``ast.Name`` identifiers found inside argument
+    annotations and return annotations of any function definitions in
+    *source_obj*.
+    """
+    import ast
+    import textwrap
+
+    try:
+        source = textwrap.dedent(inspect.getsource(source_obj))
+        tree = ast.parse(source)
+    except (OSError, TypeError):
+        return set()
+
+    names: set[str] = set()
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
+            for arg in node.args.args + node.args.posonlyargs + node.args.kwonlyargs:
+                if arg.annotation:
+                    for n in ast.walk(arg.annotation):
+                        if isinstance(n, ast.Name):
+                            names.add(n.id)
+            if node.returns:
+                for n in ast.walk(node.returns):
+                    if isinstance(n, ast.Name):
+                        names.add(n.id)
+    return names
+
+
+def _has_string_annotations(source_obj: type | Callable) -> bool:
+    """Check if *source_obj* has stringified annotations (PEP 563)."""
+    if inspect.isclass(source_obj):
+        return any(
+            isinstance(a, str)
+            for v in source_obj.__dict__.values()
+            if inspect.isfunction(v)
+            for a in v.__annotations__.values()
+        )
+    return any(isinstance(a, str) for a in getattr(source_obj, "__annotations__", {}).values())
+
+
+def _get_enclosing_scope_names(qualname: str) -> set[str]:
+    """Extract lexically enclosing scope names from ``__qualname__``.
+
+    For ``outer.<locals>.inner.<locals>.func`` this returns ``{"outer", "inner"}``.
+    """
+    parts = qualname.split(".")
+    return {p for p in parts[:-1] if p != "<locals>"}
+
+
+def resolve_closure_vars(
+    source_obj: type | Callable, extra_vars: dict[str, Any], outer_stack: list
+) -> None:
+    """Resolve closure variables hidden by PEP 563.
+
+    With ``from __future__ import annotations``, variables used only in
+    annotations are not captured in ``__closure__``.  This function parses
+    the source AST to find names used in function annotations, then looks
+    them up in lexically enclosing scope frames identified via
+    ``__qualname__``.
+
+    Only triggered when annotations are actually strings (PEP 563 active).
+    Only annotation-referenced names are added, and only from enclosing
+    scopes — not from arbitrary caller frames.
+
+    Works for both classes (``@I.ir_module``) and functions (``@T.prim_func``).
+    """
+    if not _has_string_annotations(source_obj):
+        return
+    ann_names = _collect_annotation_names(source_obj)
+    enclosing = _get_enclosing_scope_names(source_obj.__qualname__)
+    for name in ann_names:
+        if name not in extra_vars:
+            for frame_info in outer_stack[1:]:
+                if frame_info.frame.f_code.co_name in enclosing:
+                    if name in frame_info.frame.f_locals:
+                        extra_vars[name] = frame_info.frame.f_locals[name]
+                        break
+
+
+def is_defined_in_class(frames: list[FrameType], obj: Any) -> bool:
+    """Check whether a object is defined in a class scope.
+
+    Parameters
+    ----------
+    frames : List[FrameType]
+        The frame stack of the object, obtained by `inspect.stack()`.
+
+    Returns
+    -------
+    res : bool
+        The result if the object is defined in a class scope.
+    """
+
+    def _is_tvmscript_class_annotator(line: str) -> bool:
+        """Checks if the line contains a TVMScript annotator for a class
+
+        These match either `@I.ir_module` or `@R.rewriter`, or their
+        imported names `@ir_module` or `@rewriter`.
+        """
+
+        return line.startswith("@") and ("ir_module" in line or "rewriter" in line)
+
+    if len(frames) > 2:
+        frame_info = frames[2]
+        code_context = frame_info.code_context
+        if code_context is None:
+            return False
+        line = code_context[0].strip()
+        if _is_tvmscript_class_annotator(line):
+            return True
+        if line.startswith("class"):
+            lineno = frame_info.lineno
+            if lineno >= 2:
+                source, _ = findsource(obj)
+                line = source[lineno - 2].strip()
+                if _is_tvmscript_class_annotator(line):
+                    return True
+    return False

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

@@ -0,0 +1,35 @@
+# isort: skip_file
+# 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 ir module parser"""
+
+from tvm.ir import Range
+from ...ir_builder.ir import *  # pylint: disable=redefined-builtin
+from . import parser as _parser
+from .entry import ir_module, pyfunc
+
+
+__all__ = [
+    "Range",
+    "dummy_global_info",
+    "ir_module",
+    "lookup_vdevice",
+    "module_attrs",
+    "module_global_infos",
+    "pyfunc",
+    "vdevice",
+]

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

@@ -0,0 +1,163 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=import-outside-toplevel
+"""The entry point of TVM parser for ir module."""
+
+import inspect
+from collections.abc import Callable
+
+from tvm import cpu, ir
+from tvm.ir import GlobalVar, IRModule
+
+from .._core import parse, utils
+
+
+# this formulation allows us to support having @I.ir_module
+# appear as a decorator by itself or to have optional arguments
+# like @I.ir_module(check_well_formed=False)
+def ir_module(
+    mod: type | None = None, check_well_formed: bool = True, s_tir: bool = False
+) -> IRModule:
+    """The parsing method for ir module, by using `@ir_module` as decorator.
+
+    Parameters
+    ----------
+    mod : Type
+        The class to be parsed as ir module.
+
+    check_well_formed : bool
+        Whether to check well-formedness during parsing.
+
+    Returns
+    -------
+    ir_module : IRModule
+        The parsed ir module.
+    """
+
+    # Capture stack outside wrapper (wrapper adds to the stack)
+    outer_stack = inspect.stack()
+
+    def decorator_wrapper(mod):
+        if not inspect.isclass(mod):
+            raise TypeError(f"Expect a class, but got: {mod}")
+
+        # Check BasePyModule inheritance
+        base_py_module_inherited = any(base.__name__ == "BasePyModule" for base in mod.__bases__)
+
+        extra_vars = utils.inspect_class_capture(mod)
+        # Resolve closure variables hidden by PEP 563 (annotation-only names)
+        utils.resolve_closure_vars(mod, extra_vars, outer_stack)
+        m = parse(mod, extra_vars, check_well_formed=check_well_formed, s_tir=s_tir)
+
+        if base_py_module_inherited:
+            # Lazy import: tvm.relax cannot be imported at module level in tvm.script.parser
+            # because tvm.script is loaded before tvm.relax during tvm initialization.
+            from tvm.relax.base_py_module import BasePyModule
+            from tvm.relax.expr import ExternFunc  # pylint: disable=import-outside-toplevel
+
+            # Collect pyfunc methods
+            pyfunc_methods = [
+                name
+                for name, attr in mod.__dict__.items()
+                if hasattr(attr, "dispatch_token") and attr.dispatch_token == "pyfunc"
+            ]
+
+            mod._pyfunc_methods = pyfunc_methods
+
+            # Create ExternFunc nodes
+
+            for method_name in pyfunc_methods:
+                try:
+                    existing_gvars = [
+                        global_var
+                        for global_var in m.get_global_vars()
+                        if global_var.name_hint == method_name
+                    ]
+
+                    extern_func = ExternFunc(method_name)
+                    extern_func = extern_func.with_attr("is_pyfunc", True)
+                    extern_func = extern_func.with_attr("function_type", "python")
+                    extern_func = extern_func.with_attr("python_function_name", method_name)
+                    extern_func = extern_func.with_attr(
+                        "python_source", f"# Source for {method_name}"
+                    )
+                    extern_func = extern_func.with_attr("python_packed_func", None)
+
+                    if existing_gvars:
+                        m[existing_gvars[0]] = extern_func
+                    else:
+                        m[GlobalVar(method_name)] = extern_func
+
+                except Exception:  # pylint: disable=broad-exception-caught
+                    continue
+
+            class ModuleFactory:
+                """Factory class for creating BasePyModule instances with Python functions."""
+
+                def __init__(self, module, pyfunc_methods, original_class):
+                    self.ir_module = module
+                    self.pyfunc_methods = pyfunc_methods
+                    self.original_class = original_class
+
+                def __call__(self, device=None, target=None):
+                    if device is None:
+                        device = cpu(0)
+
+                    instance_ir_mod = ir.IRModule()
+                    for global_var, func in self.ir_module.functions_items():
+                        instance_ir_mod[global_var] = func
+
+                    instance = BasePyModule(instance_ir_mod, device, target)
+
+                    for method_name in self.pyfunc_methods:
+                        if hasattr(self.original_class, method_name):
+                            method = getattr(self.original_class, method_name)
+                            instance.add_python_function(method_name, method)
+
+                    return instance
+
+                def __getattr__(self, name):
+                    if hasattr(self.ir_module, name):
+                        return getattr(self.ir_module, name)
+                    raise AttributeError(
+                        f"'{self.__class__.__name__}' object has no attribute '{name}'"
+                    )
+
+            factory = ModuleFactory(m, pyfunc_methods, mod)
+            setattr(factory, "__name__", mod.__name__)
+            return factory
+
+        setattr(m, "__name__", mod.__name__)
+        return m
+
+    if mod is not None:
+        # if there are no optional args given, this will directly invoke the wrapper
+        return decorator_wrapper(mod)
+    else:
+        # if there is a optional arg given, it returns the wrapper function
+        # as a new decorator and applies it
+        setattr(decorator_wrapper, "dispatch_token", "ir")
+        return decorator_wrapper
+
+
+def pyfunc(func: Callable):
+    # Set the dispatch_token on the decorated function
+    setattr(func, "dispatch_token", "pyfunc")
+    return func
+
+
+setattr(pyfunc, "dispatch_token", "pyfunc")

python/tvm/script/parser/ir/parser.py

@@ -0,0 +1,212 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-argument
+"""The base parser for ir module"""
+
+from tvm.ir import GlobalVar
+
+from ...ir_builder import ir as I
+from .._core import Parser, dispatch, doc
+
+
+class ModuleWithGlobalVars:
+    """A Module that can add global vars during parsing, to support `Module.function` syntax."""
+
+    def __getattr__(self, attr):
+        # Customize the error message.
+        # NOTE: `__getattr__` is only called when the attribute access fails with an AttributeError
+        raise AttributeError(f"Cannot find the function `{attr}` in the current IRModule")
+
+
+@dispatch.register(token="ir", type_name="ClassDef")
+def _visit_class_def(self: Parser, node: doc.ClassDef) -> None:
+    """The class definition visiting method for ir module.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.ClassDef
+        The doc AST class definition node.
+    """
+
+    with self.var_table.with_frame():
+        with I.ir_module():
+            # Step 0. Add the class name to the var table
+            fake_module = ModuleWithGlobalVars()
+            self.var_table.add(node.name, fake_module)
+
+            # Step 1: Check if this class inherits from BasePyModule
+            is_base_py_module = _check_base_py_module_inheritance(node)
+            if is_base_py_module:
+                # Store this information in the IRModule for later use
+                I.module_attrs({"base_py_module": True})
+                # Set the parser context to allow Python functions
+                self.set_class_context(node.name, True)
+            else:
+                # Set the parser context to disallow Python functions
+                self.set_class_context(node.name, False)
+
+            # Step 2. Visit non-function stmts, including but not limited to
+            # 1. `I.module_attrs`
+            # 2. `I.module_global_infos`
+            with self.with_dispatch_token("ir"):
+                for stmt in node.body:
+                    if not isinstance(stmt, doc.FunctionDef):
+                        self.visit(stmt)
+
+            # Step 3. Visit function stmts to declare the global vars
+            for stmt in node.body:
+                if isinstance(stmt, doc.FunctionDef):
+                    global_var = self.visit_tvm_declare_function(stmt)
+                    fake_module.__setattr__(stmt.name, global_var)
+
+            # Step 4. Visit and parse the functions
+            with self.with_dispatch_token("ir"):
+                for stmt in node.body:
+                    if isinstance(stmt, doc.FunctionDef):
+                        self.visit(stmt)
+
+
+@dispatch.register(token="ir", type_name="Assign")
+def _visit_assign(self: Parser, node: doc.Assign) -> None:
+    """The assign visiting method for ir module.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.ClassDef
+        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].id
+    rhs = self.eval_expr(node.value)
+
+    I.decl_function(lhs, rhs)
+    I.def_function(lhs, rhs)
+
+
+@dispatch.register(token="ir", type_name="Expr")
+def _visit_expr(self: Parser, node: doc.Expr) -> None:
+    """The expression visiting method for ir module.
+
+    Parameters
+    ----------
+    self : Parser
+        The visiting parser.
+
+    node : doc.ClassDef
+        The doc AST expression node.
+    """
+    self.eval_expr(node.value)
+
+
+@dispatch.register(token="default", type_name="Assign")
+def visit_assign(self: Parser, node: doc.Assign) -> None:
+    if len(node.targets) != 1:
+        self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
+    lhs = node.targets[0]
+    rhs = self.eval_expr(node.value)
+    self.eval_assign(
+        target=lhs, source=rhs, bind_value=lambda _a, _b, _c, value: value, allow_shadowing=True
+    )
+
+
+@dispatch.register(token="default", type_name="pre_visit_local_function")
+def pre_visit_local_function(self: Parser, node: doc.Expr) -> None:
+    pass
+
+
+@dispatch.register(token="default", type_name="post_visit_local_function")
+def post_visit_local_function(self: Parser, node: doc.Expr) -> None:
+    pass
+
+
+@dispatch.register(token="pyfunc", type_name="tvm_declare_function")
+def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
+    """Declare a Python function as an ExternFunc in the IRModule."""
+    # Check if Python functions are allowed in this context
+    # We need to check if we're in a class that inherits from BasePyModule
+    current_class = self._get_current_class_context()
+    if current_class and not self._is_base_py_module_context():
+        self.report_error(
+            node,
+            "@I.pyfunc are only allowed in classes that inherit from BasePyModule. "
+            f"Class '{current_class}' does not inherit from BasePyModule.",
+        )
+
+    # Lazy import: tvm.relax cannot be imported at module level in tvm.script.parser
+    # because tvm.script is loaded before tvm.relax during tvm initialization.
+    from tvm.relax import ExternFunc  # pylint: disable=import-outside-toplevel
+
+    # Create ExternFunc with proper attributes for Python functions
+    func = ExternFunc(node.name)
+    func = func.with_attr("is_pyfunc", True)
+    func = func.with_attr("function_type", "python")
+    func = func.with_attr("python_function_name", node.name)
+
+    # Add placeholder attributes that will be filled in later
+    func = func.with_attr("python_source", f"# Source will be filled for {node.name}")
+    func = func.with_attr("python_packed_func", None)  # Will be filled in entry.py
+
+    # Store the function name for later retrieval
+    return I.decl_function(node.name, func)
+
+
+@dispatch.register(token="pyfunc", type_name="FunctionDef")
+def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
+    """Visit Python function definition - no need to parse the body."""
+    # Python function body is not parsed in TVMScript
+
+
+def _check_base_py_module_inheritance(node: doc.ClassDef) -> bool:
+    """Check if a class inherits from BasePyModule.
+
+    Parameters
+    ----------
+    node : doc.ClassDef
+        The class definition node to check.
+
+    Returns
+    -------
+    bool
+        True if the class inherits from BasePyModule, False otherwise.
+    """
+    if not node.bases:
+        return False
+
+    # Check each base class
+    for base in node.bases:
+        if hasattr(base, "id"):
+            if base.id == "BasePyModule":
+                return True
+        elif hasattr(base, "attr"):
+            if base.attr == "BasePyModule":
+                return True
+        elif hasattr(base, "value") and hasattr(base.value, "id"):
+            if (
+                base.value.id in ["BasePyModule", "tvm", "relax"]
+                and hasattr(base, "attr")
+                and base.attr == "BasePyModule"
+            ):
+                return True
+
+    return False