This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 951cad0e77 [REFACTOR][TVMScript] Consolidate Python IR builder
packages (#20433)
951cad0e77 is described below
commit 951cad0e77c4c6ffa221da91414ccfda091a23a5
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Sep 24 17:08:39 2026 -0400
[REFACTOR][TVMScript] Consolidate Python IR builder packages (#20433)
Consolidate Python construction APIs under `tvm.script.ir_builder`:
shared module, frame, and protocol helpers live at the root, with real
TIRx and Relax builder subpackages. Dialect packages import their
canonical builders, while the shared script layer retains generic
registration and lazy lookup for external dialects.
Complete public builder, registration, and decorator parameter
documentation, and expose the shared protocol in the API reference so
dialect hooks link to their construction contracts.
---
docs/arch/tvmscript.rst | 10 +-
docs/reference/api/python/script/ir_builder.rst | 16 +-
docs/reference/api/python/script/parser.rst | 10 +
docs/tirx/api/script.rst | 4 +-
docs/tirx/api/tile_dispatch.rst | 2 +-
docs/tirx/arch/tile_dispatch.rst | 2 +-
python/tvm/backend/cuda/__init__.py | 4 +-
python/tvm/backend/cuda/lang/alloc_pool.py | 2 +-
python/tvm/backend/cuda/script.py | 2 +-
python/tvm/backend/metal/__init__.py | 4 +-
python/tvm/backend/metal/script.py | 2 +-
python/tvm/backend/trn/__init__.py | 4 +-
python/tvm/backend/trn/script.py | 2 +-
python/tvm/relax/__init__.py | 2 +-
python/tvm/relax/script/__init__.py | 7 +-
python/tvm/relax/transform/legalize_ops/grad.py | 2 +-
python/tvm/script/__init__.py | 29 ++-
python/tvm/script/ir_builder/__init__.py | 66 +++++-
python/tvm/script/ir_builder/_ffi_api.py | 3 +
python/tvm/script/ir_builder/base.py | 107 ++++++++-
python/tvm/script/ir_builder/{ir => }/frame.py | 4 +-
python/tvm/script/ir_builder/{ir => }/ir.py | 4 +-
python/tvm/script/ir_builder/ir/__init__.py | 36 ---
python/tvm/script/ir_builder/ir/_ffi_api.py | 21 --
.../script/ir_builder/{ir => }/parser_protocol.py | 38 ++--
.../ir_builder/relax}/__init__.py | 249 ++++++++++++++++++++-
.../ir_builder/relax}/_ffi_api.py | 0
.../ir_builder/relax}/distributed/__init__.py | 0
.../ir_builder/relax}/distributed/_ffi_api.py | 0
.../ir_builder/relax}/distributed/ir.py | 7 +-
.../builder => script/ir_builder/relax}/frame.py | 0
.../builder => script/ir_builder/relax}/ir.py | 5 +-
.../ir_builder/relax}/parser_protocol.py | 131 +++++++----
.../builder => script/ir_builder/tirx}/__init__.py | 171 +++++++++++++-
.../builder => script/ir_builder/tirx}/_ffi_api.py | 0
.../ir_builder/tirx}/external_kernel.py | 2 +-
.../builder => script/ir_builder/tirx}/frame.py | 0
.../builder => script/ir_builder/tirx}/ir.py | 38 +++-
.../ir_builder/tirx}/parser_protocol.py | 113 ++++++----
.../builder => script/ir_builder/tirx}/tirx.py | 0
.../builder => script/ir_builder/tirx}/triton.py | 0
.../builder => script/ir_builder/tirx}/utils.py | 2 +-
python/tvm/script/parser/__init__.py | 40 +++-
python/tvm/script/parser/entry.py | 28 ++-
python/tvm/script/parser/ir.py | 2 +-
python/tvm/script/parser/protocol_registry.py | 128 +++++++++--
python/tvm/tirx/__init__.py | 2 +-
python/tvm/tirx/_buffer_view.py | 4 +-
python/tvm/tirx/script/__init__.py | 9 +-
python/tvm/tirx/script/jit.py | 2 +-
python/tvm/tirx/script/tile.py | 3 +-
tests/python/codegen/test_target_codegen_vulkan.py | 2 +-
tests/python/tirx/test_op.py | 2 +-
tests/python/tirx/test_op_namespace_cleanup.py | 6 +-
tests/python/tirx/test_parser_printer.py | 4 +-
.../test_tvmscript_ir_builder_irmodule.py | 2 +-
.../tvmscript/test_tvmscript_parser_evaluator.py | 2 +-
.../python/tvmscript/test_tvmscript_printer_ir.py | 2 +-
58 files changed, 1034 insertions(+), 305 deletions(-)
diff --git a/docs/arch/tvmscript.rst b/docs/arch/tvmscript.rst
index ce3f3f86f4..48cba1dd33 100644
--- a/docs/arch/tvmscript.rst
+++ b/docs/arch/tvmscript.rst
@@ -282,11 +282,11 @@ or emission receipt, allowing result attachment without a
construction context.
helpers still require ``.ctx`` to locate internal effects. At an ordinary
assignment boundary,
``.ctx(..., attach_result=False)`` preserves construction context while
leaving the returned
value's attribution to ``bind_`` through its separate ``value_span``.
-The parser-owned ``parser.protocol_registry`` module defines registration
APIs, metadata and
-persistent registration state. Dialects apply these policies beside concrete
definitions,
+The parser-owned :mod:`tvm.script.parser.protocol_registry` module defines
registration APIs,
+metadata and persistent registration state. Dialects apply these policies
beside concrete definitions,
constructor creation or necessary exposure sites, without namespace scans or
bulk
registration inventories. Dialect ``parser_protocol`` modules implement
construction hooks. The shared
-``ir_builder.ir.parser_protocol`` documents the detailed builder-hook contract
and links
+:mod:`tvm.script.ir_builder.parser_protocol` documents the detailed
builder-hook contract and links
to the registry for special syntax policies. Dialect hooks retain typed
signatures and
document only their specific behavior. For example, dialects import
``args_policy``
from ``tvm.script.parser.protocol_registry`` to register literal argument
policies.
@@ -384,7 +384,7 @@ construction policy and finalization; the syntax transpiler
does not maintain an
TIR builder
~~~~~~~~~~~
-The TIR builder (``python/tvm/tirx/script/builder/``) implements the
construction operations
+The TIR builder (``python/tvm/script/ir_builder/tirx/``) implements the
construction operations
behind public ``T`` syntax. Builder modules also support direct programmatic
construction;
public decorators and internal frame-opening operations have distinct roles.
Key categories:
@@ -416,7 +416,7 @@ public decorators and internal frame-opening operations
have distinct roles. Key
Relax builder
~~~~~~~~~~~~~
-The Relax builder (``python/tvm/relax/script/builder/``) implements:
+The Relax builder (``python/tvm/script/ir_builder/relax/``) implements:
**Function and dataflow**:
diff --git a/docs/reference/api/python/script/ir_builder.rst
b/docs/reference/api/python/script/ir_builder.rst
index 150d4a43de..e1f7f0c6ba 100644
--- a/docs/reference/api/python/script/ir_builder.rst
+++ b/docs/reference/api/python/script/ir_builder.rst
@@ -23,17 +23,17 @@ tvm.script.ir_builder
.. automodule:: tvm.script.ir_builder
:members:
:imported-members:
- :exclude-members: Any
+ :exclude-members: GenericConst, Range, StringImm, StringType
-tvm.script.ir_builder.ir
-************************
-The :class:`~tvm.script.ir_builder.IRBuilder` class is also available as
-``tvm.script.ir_builder.ir.IRBuilder``.
+tvm.script.ir_builder.parser_protocol
+*************************************
+Shared module construction and location helpers are exported directly from
+:mod:`tvm.script.ir_builder`. The protocol below describes the hooks
implemented
+by each language variant, including
:func:`~tvm.script.ir_builder.parser_protocol.if_`
+and :func:`~tvm.script.ir_builder.parser_protocol.function_`.
-.. automodule:: tvm.script.ir_builder.ir
+.. automodule:: tvm.script.ir_builder.parser_protocol
:members:
- :imported-members:
- :exclude-members: GenericConst, IRBuilder, Range, StringImm, StringType
tvm.script.ir_builder.relax
***************************
diff --git a/docs/reference/api/python/script/parser.rst
b/docs/reference/api/python/script/parser.rst
index 716738c747..02e5304d60 100644
--- a/docs/reference/api/python/script/parser.rst
+++ b/docs/reference/api/python/script/parser.rst
@@ -24,6 +24,16 @@ tvm.script.parser
:members:
:imported-members:
+tvm.script.parser.protocol_registry
+***********************************
+Language variants register source syntax policies independently of their
+construction hooks. Namespace objects and lazy initialization callbacks are
+registered through :func:`tvm.script.parser.register_namespace` and
+:func:`tvm.script.parser.register_namespace_initializer`.
+
+.. automodule:: tvm.script.parser.protocol_registry
+ :members: constexpr, args_policy, register_type_var_decl,
mutable_cell_decl, result_span, module_decorator, declaration_kind
+
The language variant aliases below share the public construction namespaces
documented
in :doc:`script`. Parser entry points above use the canonical frontend.
diff --git a/docs/tirx/api/script.rst b/docs/tirx/api/script.rst
index 8d905e11f5..dccc4e21e8 100644
--- a/docs/tirx/api/script.rst
+++ b/docs/tirx/api/script.rst
@@ -37,12 +37,12 @@ Parser entry points
Core IR builder
---------------
-.. automodule:: tvm.tirx.script.builder.ir
+.. automodule:: tvm.script.ir_builder.tirx.ir
:members:
:no-index:
:exclude-members: LetAnnotation, alloc_tcgen05_ldst_frag, anylist_getitem,
anylist_resetitem, anylist_setitem_call_packed, anylist_setitem_call_cpacked,
match_buffer
-.. currentmodule:: tvm.tirx.script.builder.ir
+.. currentmodule:: tvm.script.ir_builder.tirx.ir
.. py:function:: match_buffer(param, shape=None, dtype="float32", data=None,
strides=None, elem_offset=None, scope="global", align=-1, offset_factor=0,
layout="default")
diff --git a/docs/tirx/api/tile_dispatch.rst b/docs/tirx/api/tile_dispatch.rst
index 41c5cf3764..8e13d46f5d 100644
--- a/docs/tirx/api/tile_dispatch.rst
+++ b/docs/tirx/api/tile_dispatch.rst
@@ -21,7 +21,7 @@ Tile Dispatch Extension API
Raw node constructors
---------------------
-``tvm.tirx.script.builder.tirx`` constructs the same
+``tvm.script.ir_builder.tirx.tirx`` constructs the same
``TilePrimitiveCall`` nodes as the validated ``Tx.tile`` facade. This surface
is intended for parser, builder, and extension authors; kernel code should
normally use the :doc:`tile authoring API <tile>`. Its public builder spelling
diff --git a/docs/tirx/arch/tile_dispatch.rst b/docs/tirx/arch/tile_dispatch.rst
index fecf2476f0..fe99b0a947 100644
--- a/docs/tirx/arch/tile_dispatch.rst
+++ b/docs/tirx/arch/tile_dispatch.rst
@@ -32,7 +32,7 @@ The authoritative operation list is the C++ registry
wrapper classes live in
``python/tvm/tirx/operator/tile_primitive/ops.py``. Raw
``TilePrimitiveCall`` constructors live in
-``python/tvm/tirx/script/builder/tirx.py``, while the validated authoring
facade
+``python/tvm/script/ir_builder/tirx/tirx.py``, while the validated authoring
facade
is in ``python/tvm/tirx/script/tile.py``. Both Python construction surfaces
produce the same IR node type.
diff --git a/python/tvm/backend/cuda/__init__.py
b/python/tvm/backend/cuda/__init__.py
index 1b334c6f45..b14688ca23 100644
--- a/python/tvm/backend/cuda/__init__.py
+++ b/python/tvm/backend/cuda/__init__.py
@@ -55,8 +55,10 @@ def _detect_target_from_device(dev):
def register_backend():
"""Register CUDA-owned Python semantics."""
+ from tvm.script.ir_builder.tirx import (
+ ir as builder_ir, # pylint: disable=import-outside-toplevel
+ )
from tvm.target.detect_target import register_device_target_detector
- from tvm.tirx.script.builder import ir as builder_ir # pylint:
disable=import-outside-toplevel
runtime_dir = Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent
try:
diff --git a/python/tvm/backend/cuda/lang/alloc_pool.py
b/python/tvm/backend/cuda/lang/alloc_pool.py
index 8887b56835..0120f71ffe 100644
--- a/python/tvm/backend/cuda/lang/alloc_pool.py
+++ b/python/tvm/backend/cuda/lang/alloc_pool.py
@@ -34,7 +34,7 @@ _ir = None
def _get_ir():
global _ir
if _ir is None:
- from tvm.tirx.script.builder import ir as _mod
+ from tvm.script.ir_builder.tirx import ir as _mod
_ir = _mod
return _ir
diff --git a/python/tvm/backend/cuda/script.py
b/python/tvm/backend/cuda/script.py
index e5c1c846af..753a7bbe13 100644
--- a/python/tvm/backend/cuda/script.py
+++ b/python/tvm/backend/cuda/script.py
@@ -22,9 +22,9 @@ from collections.abc import Callable
from typing import Any
from tvm.backend.cuda import op as _cuda_op
+from tvm.script.ir_builder.tirx.ir import _dtype_forward, _op_wrapper
from tvm.tirx import is_buffer_var
from tvm.tirx import op as _tir_op
-from tvm.tirx.script.builder.ir import _dtype_forward, _op_wrapper
# pylint: disable=protected-access
diff --git a/python/tvm/backend/metal/__init__.py
b/python/tvm/backend/metal/__init__.py
index 54f99fbd8e..4954c19500 100644
--- a/python/tvm/backend/metal/__init__.py
+++ b/python/tvm/backend/metal/__init__.py
@@ -41,8 +41,10 @@ def _detect_target_from_device(dev):
def register_backend():
"""Register Metal-owned Python semantics."""
+ from tvm.script.ir_builder.tirx import (
+ ir as builder_ir, # pylint: disable=import-outside-toplevel
+ )
from tvm.target.detect_target import register_device_target_detector
- from tvm.tirx.script.builder import ir as builder_ir # pylint:
disable=import-outside-toplevel
runtime_dir = Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent
try:
diff --git a/python/tvm/backend/metal/script.py
b/python/tvm/backend/metal/script.py
index e9913dac18..1fa677bd39 100644
--- a/python/tvm/backend/metal/script.py
+++ b/python/tvm/backend/metal/script.py
@@ -19,9 +19,9 @@
from __future__ import annotations
from tvm.backend.metal import op as _metal_op
+from tvm.script.ir_builder.tirx.ir import _op_wrapper
from tvm.tirx import is_buffer_var
from tvm.tirx import op as _tir_op
-from tvm.tirx.script.builder.ir import _op_wrapper
class MetalNamespace:
diff --git a/python/tvm/backend/trn/__init__.py
b/python/tvm/backend/trn/__init__.py
index bb75669376..bc8e45dc8d 100644
--- a/python/tvm/backend/trn/__init__.py
+++ b/python/tvm/backend/trn/__init__.py
@@ -31,8 +31,10 @@ _LAZY_SUBMODULES = {
def register_backend():
"""Register Trainium-owned Python semantics."""
+ from tvm.script.ir_builder.tirx import (
+ ir as builder_ir, # pylint: disable=import-outside-toplevel
+ )
from tvm.tirx import compilation_pipeline # pylint:
disable=import-outside-toplevel
- from tvm.tirx.script.builder import ir as builder_ir # pylint:
disable=import-outside-toplevel
for name, namespace in script_namespaces().items():
builder_ir.register_script_namespace(name, namespace)
diff --git a/python/tvm/backend/trn/script.py b/python/tvm/backend/trn/script.py
index 4f9d9e436d..4a6bd01398 100644
--- a/python/tvm/backend/trn/script.py
+++ b/python/tvm/backend/trn/script.py
@@ -30,7 +30,7 @@ OpWrapper = Callable[[Callable[..., Any]], Callable[..., Any]]
def _default_op_wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
"""Create the default facade for a returned-node NKI intrinsic."""
- from tvm.tirx.script.builder.ir import _op_wrapper # pylint:
disable=import-outside-toplevel
+ from tvm.script.ir_builder.tirx.ir import _op_wrapper # pylint:
disable=import-outside-toplevel
# NKI producers return a single call_intrin node and never emit separately.
return
_result_span(f"T.nki.{func.__name__.removeprefix('nki_')}")(_op_wrapper(func))
diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py
index 955c988906..0b71d0a08d 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -117,7 +117,7 @@ tvm.script.register_dialect("relax", "tvm.relax.script")
def _check_script_module(module: "tvm.ir.IRModule") -> None:
# Delay builder imports until validation, keeping dialect/runtime
bootstrap safe.
- from .script.builder.parser_protocol import _check_module_well_formed
+ from tvm.script.ir_builder.relax.parser_protocol import
_check_module_well_formed
_check_module_well_formed(module)
diff --git a/python/tvm/relax/script/__init__.py
b/python/tvm/relax/script/__init__.py
index 172af4eece..beff701106 100644
--- a/python/tvm/relax/script/__init__.py
+++ b/python/tvm/relax/script/__init__.py
@@ -45,11 +45,10 @@ def _initialize() -> None:
return
_initializing = True
try:
+ from tvm.script.ir_builder import relax as builder
from tvm.script.parser import entry, register_namespace
from tvm.script.parser.protocol_registry import declaration_kind
- from . import builder
-
globals().update(
(name, value) for name, value in vars(builder).items() if not
name.startswith("_")
)
@@ -70,14 +69,12 @@ def _initialize() -> None:
def __getattr__(name: str) -> _Any:
- if name in ("builder", "tile"):
- return _importlib.import_module(__name__ + "." + name)
if name.startswith("_") and name != "__all__":
raise AttributeError(name)
_initialize()
if name in globals():
return globals()[name]
- return getattr(_importlib.import_module(__name__ + ".builder"), name)
+ return getattr(_importlib.import_module("tvm.script.ir_builder.relax"),
name)
class _PyModuleFactory:
diff --git a/python/tvm/relax/transform/legalize_ops/grad.py
b/python/tvm/relax/transform/legalize_ops/grad.py
index 951deb5666..96e32ac2c4 100644
--- a/python/tvm/relax/transform/legalize_ops/grad.py
+++ b/python/tvm/relax/transform/legalize_ops/grad.py
@@ -23,7 +23,7 @@ from tvm import te, tirx, topi
from tvm.ir import Call
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import tirx as T
-from tvm.tirx.script.builder.utils import buffer_proxy
+from tvm.script.ir_builder.tirx.utils import buffer_proxy
from ...block_builder import BlockBuilder
from ...expr import Expr
diff --git a/python/tvm/script/__init__.py b/python/tvm/script/__init__.py
index a71ce59180..6c846b2079 100644
--- a/python/tvm/script/__init__.py
+++ b/python/tvm/script/__init__.py
@@ -22,6 +22,7 @@ The registry permits out-of-tree language variants to expose
the same public spe
"""
import importlib
+import importlib.machinery
import importlib.util
import sys
from collections.abc import Callable
@@ -32,9 +33,8 @@ _DIALECT_REGISTRY: dict[str, str] = {}
# Callbacks own concrete eligibility; this list retains no source or
construction state.
_MODULE_VALIDATORS: list[Callable[[Any], None]] = []
-# Subpackages of `tvm.script` whose per-language variant children are
redirected to a
-# canonical package under `tvm.<language variant>.script`. An empty suffix
names
-# the public language variant namespace itself.
+# An empty suffix names the public language variant namespace itself.
+# Builder redirection applies only when no real package owns that name.
_REDIRECTED_SUBPACKAGES = {
"tvm.script.parser": "",
"tvm.script.ir_builder": "builder",
@@ -79,9 +79,10 @@ def register_dialect(name: str, module_path: str) -> None:
Writes ``name -> module_path`` into ``_DIALECT_REGISTRY``. After
registration, ``tvm.script.<name>`` resolves to ``module_path`` via
``__getattr__``, and ``tvm.script.parser.<name>`` /
``tvm.script.ir_builder.<name>``
- resolve to the same script namespace and ``module_path + ".builder"``
- respectively. Deep builder imports are handled by
- ``_DialectRedirectFinder`` on ``sys.meta_path``.
+ resolve to the same script namespace and the corresponding real builder
+ package, respectively. When no real builder package exists, external
language
+ variants resolve through ``module_path + ".builder"``. Registered namespace
+ imports are handled by ``_DialectRedirectFinder`` on ``sys.meta_path``.
This function is idempotent — re-registering the same name with the same
path is harmless.
@@ -101,7 +102,9 @@ def register_dialect(name: str, module_path: str) -> None:
module_path : str
The full dotted module path of the language variant's script package,
e.g.
``"tvm.tirx.script"``. That package exposes its public construction
- entry points directly and its imperative APIs in a ``builder``
submodule.
+ entry points directly. Out-of-tree language variants provide imperative
+ APIs in a ``builder`` submodule; in-tree builders live in
+ ``tvm.script.ir_builder.<name>``.
"""
_DIALECT_REGISTRY[name] = module_path
@@ -153,6 +156,14 @@ class _DialectRedirectFinder:
@classmethod
def find_spec(cls, fullname, path, target=None):
+ # A real builder package owns its canonical name. Only missing packages
+ # use the out-of-tree language variant registry fallback.
+ prefix = "tvm.script.ir_builder."
+ if fullname.startswith(prefix):
+ package_name = prefix + fullname[len(prefix) :].partition(".")[0]
+ parent_path = sys.modules["tvm.script.ir_builder"].__path__
+ if importlib.machinery.PathFinder.find_spec(package_name,
parent_path) is not None:
+ return None
redirected = _redirect_target(fullname)
if redirected is None:
return None
@@ -196,7 +207,9 @@ class _AliasLoader:
# Install the redirect finder once. Re-importing tvm.script (e.g. during a
# pytest reload) must not stack duplicates.
if not any(isinstance(f, _DialectRedirectFinder) for f in sys.meta_path):
- sys.meta_path.append(_DialectRedirectFinder())
+ # Handle deep aliases before PathFinder can execute a canonical source file
+ # again under its alias package's __path__.
+ sys.meta_path.insert(0, _DialectRedirectFinder())
def __getattr__(name: str) -> Any:
diff --git a/python/tvm/script/ir_builder/__init__.py
b/python/tvm/script/ir_builder/__init__.py
index 2f1769459c..5438e851e8 100644
--- a/python/tvm/script/ir_builder/__init__.py
+++ b/python/tvm/script/ir_builder/__init__.py
@@ -14,20 +14,72 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""The ir_builder subpackage of TVMScript."""
+"""Shared TVMScript construction APIs and lazy language variant builders."""
-import importlib
-from typing import Any
+import importlib as _importlib
+from typing import Any as _Any
-from .base import AlreadyEmitted, IRBuilder
+from tvm.ir import GenericConst, Range, StringImm, StringType
+from tvm.script.parser.protocol_registry import constexpr
+from .base import (
+ MISSING,
+ AlreadyEmitted,
+ IRBuilder,
+ annotation_value_,
+ at_,
+ require_defined,
+ with_at_group_,
+)
+from .frame import IRModuleFrame
+from .ir import (
+ decl_function,
+ def_function,
+ ir_module,
+ lookup_name,
+ meta_var,
+ module_attrs,
+ module_get_attr,
+ module_global_infos,
+ module_set_attr,
+)
+from .parser_protocol import check_well_formed_, module_member_
-def __getattr__(name: str) -> Any:
- # Lazy import to avoid loading tvm.script during language variant
bootstrap.
+# Keep source namespaces independent of imported helper modules and lazy
builders.
+__all__ = [
+ "MISSING",
+ "AlreadyEmitted",
+ "GenericConst",
+ "IRBuilder",
+ "IRModuleFrame",
+ "Range",
+ "StringImm",
+ "StringType",
+ "annotation_value_",
+ "at_",
+ "check_well_formed_",
+ "constexpr",
+ "decl_function",
+ "def_function",
+ "ir_module",
+ "lookup_name",
+ "meta_var",
+ "module_attrs",
+ "module_get_attr",
+ "module_global_infos",
+ "module_member_",
+ "module_set_attr",
+ "require_defined",
+ "with_at_group_",
+]
+
+
+def __getattr__(name: str) -> _Any:
+ # Real packages resolve directly; the finder supplies registered external
builders.
from tvm.script import _DIALECT_REGISTRY # pylint:
disable=import-outside-toplevel
if name in _DIALECT_REGISTRY:
- module = importlib.import_module(f"{_DIALECT_REGISTRY[name]}.builder")
+ module = _importlib.import_module(f"{__name__}.{name}")
globals()[name] = module
return module
raise AttributeError(f"module 'tvm.script.ir_builder' has no attribute
{name!r}")
diff --git a/python/tvm/script/ir_builder/_ffi_api.py
b/python/tvm/script/ir_builder/_ffi_api.py
index dbaad5b791..e3caa415f5 100644
--- a/python/tvm/script/ir_builder/_ffi_api.py
+++ b/python/tvm/script/ir_builder/_ffi_api.py
@@ -19,3 +19,6 @@
import tvm_ffi
tvm_ffi.init_ffi_api("script.ir_builder", __name__) # pylint:
disable=protected-access
+
+# Shared module construction retains its native registry namespace.
+tvm_ffi.init_ffi_api("script.ir_builder.ir", __name__) # pylint:
disable=protected-access
diff --git a/python/tvm/script/ir_builder/base.py
b/python/tvm/script/ir_builder/base.py
index 815bee7be2..6411c821ad 100644
--- a/python/tvm/script/ir_builder/base.py
+++ b/python/tvm/script/ir_builder/base.py
@@ -251,7 +251,21 @@ def source_span(
class AlreadyEmitted(Generic[_T]):
- """Hold the exact emitted value; location handling preserves this
receipt."""
+ """Hold the exact emitted value; location handling preserves this receipt.
+
+ Parameters
+ ----------
+ value : Any
+ Object already emitted by a language variant's builder. The receipt
+ retains this object without copying it. Binding or emitting the receipt
+ must not emit the object again.
+
+ Attributes
+ ----------
+ value : Any
+ The same emitted object, available for identity checks and source-span
+ attachment.
+ """
__slots__ = ("value",)
@@ -264,6 +278,24 @@ def at(
) -> _T:
"""Attach source context to the same IR node, emission receipt, or frame.
+ Parameters
+ ----------
+ span : SpanEntry, Span, tuple or None
+ Source location to compose with the active construction context. A
tuple
+ contains ``(source_name, line, end_line, column, end_column)``; its
source
+ name may be a string or SourceName. None leaves the value unchanged.
+ value : Any
+ Native object, :class:`AlreadyEmitted` receipt, or list/tuple of native
+ objects to annotate. A receipt's contained object receives the span.
+
+ Returns
+ -------
+ Any
+ The exact ``value`` object, including its original receipt or
container.
+ With no active builder, the value is returned without modification.
+
+ Notes
+ -----
Native mutation annotates the statement held by the builder itself. Keep
the original Python facade as well, including callable objects and frames.
Unsupported objects and ordinary Python values pass through unchanged.
@@ -279,7 +311,26 @@ def at(
def require_defined(value, name):
- """Report a source name whose designated region output was not produced."""
+ """Report a source name whose designated region output was not produced.
+
+ Parameters
+ ----------
+ value : Any
+ Candidate binding. Only the canonical ``MISSING`` singleton denotes an
+ absent value; explicit None is a defined value.
+ name : str
+ Source identifier to include in the missing-name diagnostic.
+
+ Returns
+ -------
+ Any
+ The exact ``value`` when it is defined.
+
+ Raises
+ ------
+ NameError
+ If ``value`` is ``MISSING``.
+ """
if value is MISSING:
raise NameError(f"name {name!r} is not defined")
return value
@@ -291,7 +342,32 @@ def with_at_group_(
*,
attach_result: bool = True,
) -> _T:
- """Evaluate once under a location, optionally attaching it to the same
result."""
+ """Evaluate once under a location, optionally attaching it to the same
result.
+
+ Parameters
+ ----------
+ location : SpanEntry, Span, tuple or None
+ Source context for the call. A tuple contains
+ ``(source_name, line, end_line, column, end_column)``. None, or the
+ absence of an active builder, leaves construction context unchanged.
+ thunk : Callable[[], Any]
+ Zero-argument callable evaluated exactly once inside that context.
+ attach_result : bool, optional
+ Attach the location to the returned object with :func:`at_`. Defaults
+ to True. False supplies construction context only, leaving explicit
+ result attachment to a later operation.
+
+ Returns
+ -------
+ Any
+ The exact result of ``thunk``, with its receipt or container preserved.
+
+ Notes
+ -----
+ The prior source context is restored even if the callable raises; its
+ exception propagates unchanged. Frames created during the call retain their
+ construction spans regardless of ``attach_result``.
+ """
span = source_span(location)
context = (
IRBuilder.current().with_source_span(span)
@@ -402,7 +478,30 @@ def _return_annotation(annotation):
def annotation_value_(name, value):
- """Adapt a real definition-context symbol using the native function map."""
+ """Adapt a real definition-context symbol using the native function map.
+
+ Parameters
+ ----------
+ name : str
+ Source spelling used to resolve the symbol in the nearest active
+ function frame.
+ value : TypeVar, Var or Any
+ Captured definition-context value. An unconstrained ``typing.TypeVar``
+ resolves a symbolic variable; a primitive IR Var supplies its existing
+ value to the resolver. Other values pass through unchanged.
+
+ Returns
+ -------
+ Any
+ The function frame's resolved symbol, or the unchanged nonsymbolic
value.
+
+ Raises
+ ------
+ TypeError
+ If a TypeVar has a bound or constraints.
+ ValueError
+ If a symbolic value requires resolution without an active function
frame.
+ """
if isinstance(value, TypeVar):
if value.__bound__ is not None or value.__constraints__:
raise TypeError("A symbolic TypeVar cannot have constraints or a
bound")
diff --git a/python/tvm/script/ir_builder/ir/frame.py
b/python/tvm/script/ir_builder/frame.py
similarity index 94%
rename from python/tvm/script/ir_builder/ir/frame.py
rename to python/tvm/script/ir_builder/frame.py
index 346ec00c62..778fe2fed3 100644
--- a/python/tvm/script/ir_builder/ir/frame.py
+++ b/python/tvm/script/ir_builder/frame.py
@@ -14,11 +14,11 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Package tvm.script.ir_builder.ir.frame"""
+"""Package tvm.script.ir_builder.frame"""
from tvm_ffi import register_object as _register_object
-from ..base import IRBuilderFrame
+from .base import IRBuilderFrame
@_register_object("script.ir_builder.IRModuleFrame")
diff --git a/python/tvm/script/ir_builder/ir/ir.py
b/python/tvm/script/ir_builder/ir.py
similarity index 99%
rename from python/tvm/script/ir_builder/ir/ir.py
rename to python/tvm/script/ir_builder/ir.py
index efe38fc5a7..5b6e0f3b59 100644
--- a/python/tvm/script/ir_builder/ir/ir.py
+++ b/python/tvm/script/ir_builder/ir.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Package tvm.script.ir_builder.ir.ir"""
+"""Package tvm.script.ir_builder.ir"""
import inspect
from typing import TypeVar
@@ -22,8 +22,8 @@ from typing import TypeVar
from tvm.ir import BaseFunc, GlobalInfo, GlobalVar
from tvm.runtime import Object as tvm_Object
-from ..base import IRBuilder
from . import _ffi_api
+from .base import IRBuilder
from .frame import IRModuleFrame
T = TypeVar("T")
diff --git a/python/tvm/script/ir_builder/ir/__init__.py
b/python/tvm/script/ir_builder/ir/__init__.py
deleted file mode 100644
index 59ccf8e395..0000000000
--- a/python/tvm/script/ir_builder/ir/__init__.py
+++ /dev/null
@@ -1,36 +0,0 @@
-# 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.
-"""Package tvm.script.ir_builder.ir"""
-
-from tvm.ir import GenericConst, Range, StringImm, StringType
-from tvm.script.parser.protocol_registry import constexpr
-
-from ..base import IRBuilder, MISSING, annotation_value_, at_, with_at_group_,
require_defined
-from .parser_protocol import check_well_formed_, module_member_
-from .frame import IRModuleFrame
-from .ir import (
- decl_function,
- def_function,
- ir_module,
- module_attrs,
- module_get_attr,
- module_set_attr,
- module_global_infos,
- lookup_name,
- meta_var,
-)
diff --git a/python/tvm/script/ir_builder/ir/_ffi_api.py
b/python/tvm/script/ir_builder/ir/_ffi_api.py
deleted file mode 100644
index 9d0e22b6bf..0000000000
--- a/python/tvm/script/ir_builder/ir/_ffi_api.py
+++ /dev/null
@@ -1,21 +0,0 @@
-# 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.
-"""FFI APIs"""
-
-import tvm_ffi
-
-tvm_ffi.init_ffi_api("script.ir_builder.ir", __name__) # pylint:
disable=protected-access
diff --git a/python/tvm/script/ir_builder/ir/parser_protocol.py
b/python/tvm/script/ir_builder/parser_protocol.py
similarity index 97%
rename from python/tvm/script/ir_builder/ir/parser_protocol.py
rename to python/tvm/script/ir_builder/parser_protocol.py
index 48c2f1594d..4f97019d4f 100644
--- a/python/tvm/script/ir_builder/ir/parser_protocol.py
+++ b/python/tvm/script/ir_builder/parser_protocol.py
@@ -33,14 +33,14 @@ from typing import Any
from tvm import ir as _ir
-from ..base import MISSING, AlreadyEmitted, IRBuilderFrame, SpanEntry
+from .base import MISSING, AlreadyEmitted, IRBuilderFrame, SpanEntry
_Span = SpanEntry | _ir.Span | tuple[_ir.SourceName, int, int, int, int] | None
-# ------------
+# --------------------------------------
# Section: control flow
-# ------------
+# --------------------------------------
#
# ``with X.if_(cond):`` opens a conditional frame.
# ``with X.then_():`` enters its true branch.
@@ -73,7 +73,7 @@ def if_(condition: Any, *, span: _Span = None) ->
IRBuilderFrame:
Notes
-----
- Requires an active function region. Enter then_/else_ frames inside it;
each branch's
+ Requires an active function region. Enter ``then_``/``else_`` frames
inside it; each branch's
lexical helper is defined and called inside that branch frame. TIRx
permits a missing
else; value-producing Relax branches must agree on their final binding
name/type.
Invalid context/condition/outputs produce builder diagnostics. Explicit
constexpr
@@ -345,9 +345,9 @@ def continue_(*, span: _Span = None) -> AlreadyEmitted[Any]:
raise NotImplementedError
-# ------------
+# --------------------------------------
# Section: operator overloading
-# ------------
+# --------------------------------------
#
# ``X.if_then_else_(c, a, b)`` selects an expression.
# ``X.and_(a, b)`` constructs conjunction.
@@ -702,9 +702,9 @@ def ne_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
raise NotImplementedError
-# ------------
+# --------------------------------------
# Section: context lookup and resolution
-# ------------
+# --------------------------------------
#
# ``X.resolve_global_info_(key)`` resolves module metadata.
# ``X.resolve_type_var_("n")`` resolves a symbolic dimension.
@@ -862,14 +862,14 @@ def module_member_(name: str, value: Any) -> Any:
return value
-# ------------
+# --------------------------------------
# Section: special protocol
-# ------------
+# --------------------------------------
#
# Syntax markers live in tvm.script.parser.protocol_registry.
# ``constexpr(value)`` selects host evaluation in marked control flow.
# ``args_policy(path, fields)`` marks expression-string or global-info
arguments.
-# ``register_type_var_decl(path, dtype)`` marks symbolic declarations.
+# ``register_type_var_decl(path, constructor, dtype=...)`` marks symbolic
declarations.
# ``mutable_cell_decl(path)`` marks mutable storage declarations.
# ``result_span(path)`` permits attaching a call's result span without a call
context.
# ``module_decorator(path)`` marks module declaration decorators.
@@ -894,9 +894,11 @@ def function_(*, decl: bool = False, span: _Span = None,
**options: Any) -> IRBu
unspecified; active source-call provenance is composed by the builder.
Frames
retain their location until finalization.
options : Any
- Language variant options: TIRx private, s_tir and persistent default
to False. Relax
- is_pure defaults to True, is_private/local to False; local=True
requires a
- declared reference when building its body.
+ Language variant options: TIRx ``private``, ``s_tir`` and
``persistent``
+ default to False. Relax ``pure`` defaults to True and
``private``/``local``
+ to False; ``local=True`` requires a declared reference when building
its
+ body. See :func:`tvm.script.ir_builder.tirx.function_` and
+ :func:`tvm.script.ir_builder.relax.function_` for their concrete
options.
Returns
-------
@@ -1074,9 +1076,9 @@ def check_well_formed_(module: _ir.IRModule) -> None:
validator(module)
-# ------------
+# --------------------------------------
# Section: binding
-# ------------
+# --------------------------------------
#
# ``a = X.bind_(value, name="a")`` binds a source name.
# ``X.decl_mutable_cell_(value, ty=ty)`` declares mutable storage.
@@ -1283,9 +1285,9 @@ def unpack(value: Any) -> Any:
raise NotImplementedError
-# ------------
+# --------------------------------------
# Section: statement
-# ------------
+# --------------------------------------
#
# ``X.emit_(value)`` emits an expression statement.
# ``X.return_(value)`` emits a function return.
diff --git a/python/tvm/relax/script/builder/__init__.py
b/python/tvm/script/ir_builder/relax/__init__.py
similarity index 51%
rename from python/tvm/relax/script/builder/__init__.py
rename to python/tvm/script/ir_builder/relax/__init__.py
index b94c54cce9..f381cb31c5 100644
--- a/python/tvm/relax/script/builder/__init__.py
+++ b/python/tvm/script/ir_builder/relax/__init__.py
@@ -85,7 +85,31 @@ supports_mutable_declarations = False
@_args_policy("R.Tensor", {"shape": "expr_str", "vdevice": "global_info"},
scalar_strings=False)
def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *, span=None):
- """Construct a tensor type from concrete shape dimensions."""
+ """Construct a Relax tensor type.
+
+ Parameters
+ ----------
+ shape : Expr or sequence of Expr, optional
+ Tensor shape, or None when unknown. A string supplied without dtype
+ is shorthand for the dtype. Script dimension strings follow the
+ registered shape-expression argument policy.
+ dtype : str or PrimType, optional
+ Element type; None leaves the element type unknown.
+ vdevice : VDevice or str, optional
+ Concrete virtual device or a module metadata selector, such as
"cuda:0".
+ None leaves the virtual device unspecified.
+ ndim : int, optional
+ Rank when shape is unknown; -1 means unknown rank. Do not supply
+ an explicit rank together with a known shape.
+ span : Span or source location, optional
+ Source location attached to the constructed IR; None leaves it
unspecified.
+
+ Returns
+ -------
+ result : TensorType or Type
+ The tensor type. A string device selector evaluated before the module
+ builder opens returns a missing type for the generated declaration to
resolve.
+ """
if isinstance(shape, _python.str) and dtype is None:
dtype, shape = shape, None
if isinstance(vdevice, _python.str) and not _IRBuilder.is_in_scope():
@@ -100,7 +124,32 @@ def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1,
*, span=None):
"R.DTensor", {"shape": "expr_str", "device_mesh": "global_info"},
scalar_strings=False
)
def DTensor(shape=None, dtype=None, device_mesh=None, placement="", *,
ndim=-1, span=None):
- """Construct a distributed tensor type from concrete dimensions."""
+ """Construct a Relax distributed tensor type.
+
+ Parameters
+ ----------
+ shape : Expr or sequence of Expr, optional
+ Global tensor shape, or None when unknown; dimension strings in script
+ follow the registered shape-expression argument policy.
+ dtype : str or PrimType, optional
+ Element type; None leaves the element type unknown.
+ device_mesh : DeviceMesh or str, optional
+ Concrete mesh or module metadata selector. None creates an empty mesh
+ placeholder. A selector is resolved in the enclosing module builder.
+ placement : Placement or str, optional
+ Distribution placement. Text, including the default empty string, is
+ parsed with Placement.from_text.
+ ndim : int, optional
+ Global rank when shape is unknown; -1 means unknown rank.
+ span : Span or source location, optional
+ Source location attached to the constructed IR; None leaves it
unspecified.
+
+ Returns
+ -------
+ result : DTensorType or Type
+ The distributed type. A string mesh selector evaluated before the
+ module builder opens returns a missing type for later resolution.
+ """
if isinstance(device_mesh, _python.str) and not _IRBuilder.is_in_scope():
return _ir.Type.missing()
if device_mesh is None:
@@ -125,7 +174,24 @@ __tvm_value_if__ = True
@_args_policy("R.Shape", {"values": "expr_str"}, dtype="int64")
def Shape(values=None, ndim=-1, *, span=None):
- """Construct a shape type from concrete dimensions."""
+ """Construct a Relax shape type.
+
+ Parameters
+ ----------
+ values : sequence of Expr, optional
+ Known dimensions, or None for an unknown shape value. Script dimension
+ strings use the registered int64 shape-expression policy.
+ ndim : int, optional
+ Number of dimensions when values is None; -1 leaves it unknown.
+ Do not supply an explicit count together with known values.
+ span : Span or source location, optional
+ Source location attached to the constructed IR; None leaves it
unspecified.
+
+ Returns
+ -------
+ result : ShapeType
+ The constructed shape type.
+ """
return _relax.ShapeType(values, ndim, _source_span(span))
@@ -142,7 +208,36 @@ def _type(value):
def Callable(params=None, ret=None, purity=None, derive_func=None, *,
span=None):
- """Construct a concrete or opaque Relax function type."""
+ """Construct a concrete or opaque Relax function type.
+
+ Parameters
+ ----------
+ params : Type, callable, or sequence of annotations, optional
+ Parameter annotations. A single annotation is accepted; None creates
+ an opaque callable with an unspecified parameter list.
+ ret : Type or callable, optional
+ Return annotation. None means an empty tuple for a concrete callable
+ and an unspecified result for an opaque callable.
+ purity : bool, optional
+ Whether the callable is pure. None selects True for a concrete
+ parameter list and False for an opaque callable.
+ derive_func : str or EnvFunc, optional
+ Custom result-type derivation for an opaque callable. It is not
+ accepted when params supplies a concrete parameter list.
+ span : Span or source location, optional
+ Source location attached to the constructed IR; None leaves it
unspecified.
+
+ Returns
+ -------
+ result : FuncType
+ The constructed function type.
+
+ Notes
+ -----
+ Annotations may be types, primitive expressions supplying their types, or
+ zero-argument factories returning either. Opaque result and derivation
rules
+ follow :meth:`tvm.relax.FuncType.opaque_func`.
+ """
if purity is None:
purity = params is not None
if params is None:
@@ -162,14 +257,39 @@ def Callable(params=None, ret=None, purity=None,
derive_func=None, *, span=None)
def Tuple(*fields, span=None):
- """Construct a tuple type while preserving missing component types."""
+ """Construct a Relax tuple type.
+
+ Parameters
+ ----------
+ fields : Type or callable
+ Field annotations as positional arguments, or one list or tuple.
+ Each annotation is normalized to a type; None denotes an empty tuple.
+ span : Span or source location, optional
+ Source location attached to the constructed IR; None leaves it
unspecified.
+
+ Returns
+ -------
+ result : TupleType
+ The tuple type with fields in the supplied order.
+ """
if len(fields) == 1 and isinstance(fields[0], list | _python.tuple):
fields = fields[0]
return _ir.TupleType([_type(field) for field in fields],
_source_span(span))
def Object(*, span=None):
- """Construct the unconstrained Relax value type."""
+ """Construct the unconstrained Relax value type.
+
+ Parameters
+ ----------
+ span : Span or source location, optional
+ Source location attached to the constructed IR; None leaves it
unspecified.
+
+ Returns
+ -------
+ result : AnyType
+ A type accepting any Relax value.
+ """
return _relax.AnyType(_source_span(span))
@@ -180,7 +300,27 @@ is_type_var = _ir.is_prim_var
def type_var(name, *, dtype=None, span=None):
- """Construct a fresh standalone primitive symbol."""
+ """Construct a fresh standalone primitive symbol.
+
+ Parameters
+ ----------
+ name : str
+ Name of the symbol.
+ dtype : str or PrimType, optional
+ Primitive type of the symbol; None selects "int64".
+ span : Span or source location, optional
+ Source location attached to the constructed IR; None leaves it
unspecified.
+
+ Returns
+ -------
+ result : Var
+ The newly constructed primitive variable.
+
+ Notes
+ -----
+ This constructor creates a new symbol on each call. Use the language
variant
+ resolver for symbols shared by name within a function signature.
+ """
return _ir.Var(name, "int64" if dtype is None else dtype,
_source_span(span))
@@ -214,7 +354,24 @@ def _value(value, ty=None):
def match_cast(value, ty, *, span=None):
- """Construct a match-cast descriptor for ``bind_`` to consume."""
+ """Construct a match-cast descriptor for the binding hook.
+
+ Parameters
+ ----------
+ value : Expr or Python value
+ Value to match against the asserted type. Numbers and Python tuples
+ are converted to Relax expressions; None is not accepted.
+ ty : Type or callable
+ Asserted type, or a zero-argument factory producing its annotation.
+ span : Span or source location, optional
+ Source location attached to the constructed IR; None leaves it
unspecified.
+
+ Returns
+ -------
+ result : MatchCast
+ An unbound match-cast descriptor, consumed by the language variant
+ binding hook to emit and name the result.
+ """
if value is None:
raise ValueError("The match-cast value cannot be None")
ty = _type(ty)
@@ -276,7 +433,25 @@ def _logical_pair(lhs, rhs, operation, primitive,
python_operation):
def logical_and(*values):
- """Construct conjunction of concrete host, primitive, or tensor values."""
+ """Construct conjunction of host, primitive, or tensor values.
+
+ Parameters
+ ----------
+ values : Expr or Python value
+ One or more eagerly evaluated operands. Host pairs follow Python
+ logical operations; primitive pairs use primitive IR and tensor pairs
+ use the corresponding Relax operation.
+
+ Returns
+ -------
+ result : Expr or Python value
+ The conjunction reduced from left to right.
+
+ Notes
+ -----
+ All arguments are evaluated before this call; it does not provide Python
+ short-circuit evaluation of the argument expressions.
+ """
if not values:
raise TypeError("logical_and requires at least one operand")
result = values[0]
@@ -286,7 +461,25 @@ def logical_and(*values):
def logical_or(*values):
- """Construct disjunction of concrete host, primitive, or tensor values."""
+ """Construct disjunction of host, primitive, or tensor values.
+
+ Parameters
+ ----------
+ values : Expr or Python value
+ One or more eagerly evaluated operands. Host pairs follow Python
+ logical operations; primitive pairs use primitive IR and tensor pairs
+ use the corresponding Relax operation.
+
+ Returns
+ -------
+ result : Expr or Python value
+ The disjunction reduced from left to right.
+
+ Notes
+ -----
+ All arguments are evaluated before this call; it does not provide Python
+ short-circuit evaluation of the argument expressions.
+ """
if not values:
raise TypeError("logical_or requires at least one operand")
result = values[0]
@@ -296,7 +489,19 @@ def logical_or(*values):
def logical_not(value):
- """Negate a host or IR boolean without coercing IR to Python bool."""
+ """Negate a host, primitive, or tensor condition.
+
+ Parameters
+ ----------
+ value : Expr or Python value
+ Operand to negate. Primitive and tensor expressions use their IR
+ logical operation; host values use Python truth testing.
+
+ Returns
+ -------
+ result : Expr or bool
+ The logical negation without testing an IR expression as a Python bool.
+ """
if _ir.is_prim_expr(value):
return _tir.Not(value)
if isinstance(value, _ir.Expr):
@@ -305,7 +510,27 @@ def logical_not(value):
def select(condition, true_value, false_value):
- """Construct elementwise selection or select already-built host values."""
+ """Select between already-constructed values.
+
+ Parameters
+ ----------
+ condition : Expr or Python value
+ Primitive condition, tensor condition, or host truth value.
+ true_value : Expr or Python value
+ Value selected when the condition is true.
+ false_value : Expr or Python value
+ Value selected when the condition is false.
+
+ Returns
+ -------
+ result : Expr or Python value
+ A primitive Select, Relax elementwise where, or the selected host
value.
+
+ Notes
+ -----
+ Both value arguments are evaluated before this call. Host selection returns
+ the selected object; tensor selection converts Python numbers and tuples.
+ """
if _ir.is_prim_expr(condition):
return _tir.Select(condition, true_value, false_value)
if isinstance(condition, _ir.Expr):
diff --git a/python/tvm/relax/script/builder/_ffi_api.py
b/python/tvm/script/ir_builder/relax/_ffi_api.py
similarity index 100%
rename from python/tvm/relax/script/builder/_ffi_api.py
rename to python/tvm/script/ir_builder/relax/_ffi_api.py
diff --git a/python/tvm/relax/script/builder/distributed/__init__.py
b/python/tvm/script/ir_builder/relax/distributed/__init__.py
similarity index 100%
rename from python/tvm/relax/script/builder/distributed/__init__.py
rename to python/tvm/script/ir_builder/relax/distributed/__init__.py
diff --git a/python/tvm/relax/script/builder/distributed/_ffi_api.py
b/python/tvm/script/ir_builder/relax/distributed/_ffi_api.py
similarity index 100%
rename from python/tvm/relax/script/builder/distributed/_ffi_api.py
rename to python/tvm/script/ir_builder/relax/distributed/_ffi_api.py
diff --git a/python/tvm/relax/script/builder/distributed/ir.py
b/python/tvm/script/ir_builder/relax/distributed/ir.py
similarity index 96%
rename from python/tvm/relax/script/builder/distributed/ir.py
rename to python/tvm/script/ir_builder/relax/distributed/ir.py
index 827dea67fd..cb9e6391c4 100644
--- a/python/tvm/relax/script/builder/distributed/ir.py
+++ b/python/tvm/script/ir_builder/relax/distributed/ir.py
@@ -40,12 +40,11 @@ from tvm.relax.op.distributed import (
from tvm.relax.op.distributed import (
redistribute as _redistribute,
)
-from tvm.relax.script.builder.ir import py_str
from tvm.relax.utils import convert_to_expr
from tvm.runtime import _tensor
-from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder.ir import IRModuleFrame
-from tvm.script.ir_builder.ir.ir import lookup_global_info
+from tvm.script.ir_builder import IRBuilder, IRModuleFrame
+from tvm.script.ir_builder.ir import lookup_global_info
+from tvm.script.ir_builder.relax.ir import py_str
from . import _ffi_api
diff --git a/python/tvm/relax/script/builder/frame.py
b/python/tvm/script/ir_builder/relax/frame.py
similarity index 100%
rename from python/tvm/relax/script/builder/frame.py
rename to python/tvm/script/ir_builder/relax/frame.py
diff --git a/python/tvm/relax/script/builder/ir.py
b/python/tvm/script/ir_builder/relax/ir.py
similarity index 99%
rename from python/tvm/relax/script/builder/ir.py
rename to python/tvm/script/ir_builder/relax/ir.py
index a406c402de..4996b000e6 100644
--- a/python/tvm/relax/script/builder/ir.py
+++ b/python/tvm/script/ir_builder/relax/ir.py
@@ -207,9 +207,8 @@ from tvm.runtime._tensor import (
vulkan,
webgpu,
)
-from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder.ir import decl_function
-from tvm.script.ir_builder.ir.ir import _class_global_infos
+from tvm.script.ir_builder import IRBuilder, decl_function
+from tvm.script.ir_builder.ir import _class_global_infos
from tvm.script.parser.protocol_registry import module_decorator as
_module_decorator
from . import _ffi_api, frame
diff --git a/python/tvm/relax/script/builder/parser_protocol.py
b/python/tvm/script/ir_builder/relax/parser_protocol.py
similarity index 81%
rename from python/tvm/relax/script/builder/parser_protocol.py
rename to python/tvm/script/ir_builder/relax/parser_protocol.py
index e759767195..c71fae3a19 100644
--- a/python/tvm/relax/script/builder/parser_protocol.py
+++ b/python/tvm/script/ir_builder/relax/parser_protocol.py
@@ -19,7 +19,7 @@
Hooks delegate construction to this language variant's native builder and IR
APIs.
For example, generated ``X.if_(condition)`` creates the native conditional
frame;
``X.then_()`` and ``X.else_()`` enter its branches. See the corresponding
shared
-``tvm.script.ir_builder.ir.parser_protocol`` hooks for operand and span
contracts.
+``tvm.script.ir_builder.parser_protocol`` hooks for operand and span contracts.
"""
from __future__ import annotations
@@ -32,15 +32,15 @@ from typing import Any, NoReturn
import tvm_ffi as _ffi
+import tvm.script.ir_builder.relax as _builder
from tvm import ir as _ir
from tvm import relax as _relax
from tvm import tirx as _tir
from tvm.ir.prim import _ffi_api as _prim_ffi
from tvm.script.ir_builder import IRBuilder as _IRBuilder
from tvm.script.ir_builder import base as _base
-from tvm.script.ir_builder.ir.frame import IRModuleFrame as _IRModuleFrame
+from tvm.script.ir_builder.frame import IRModuleFrame as _IRModuleFrame
-from .. import builder as _builder
from . import _ffi_api
from . import frame as _frame
from . import ir as _native
@@ -48,9 +48,9 @@ from . import ir as _native
_Span = _base.SpanEntry | _ir.Span | tuple[_ir.SourceName, int, int, int, int]
| None
-# ------------
+# --------------------------------------
# Section: control flow
-# ------------
+# --------------------------------------
#
# ``with X.if_(cond):`` opens a conditional frame.
# ``with X.then_():`` enters its true branch.
@@ -63,24 +63,24 @@ _Span = _base.SpanEntry | _ir.Span | tuple[_ir.SourceName,
int, int, int, int] |
def if_(condition: Any, *, span: _Span = None) -> _frame.IfFrame:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.if_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.if_`."""
return _base.at_(span, _native.If(condition))
def then_(*, span: _Span = None) -> _frame.ThenFrame:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.then_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.then_`."""
return _base.at_(span, _native.Then())
def else_(*, span: _Span = None) -> _frame.ElseFrame:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.else_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.else_`."""
return _base.at_(span, _native.Else())
def for_(
iterable: Any, *, names: str | Sequence[str] | None = None, span: _Span =
None
) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.for_`.
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.for_`.
Relax rejects source loops; use supported functional control flow.
"""
@@ -88,17 +88,17 @@ def for_(
def while_(condition: Any, *, span: _Span = None) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.while_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.while_`."""
raise TypeError("Relax does not support imperative while loops")
def range_(*args: Any, annotations: dict[str, Any] | None = None) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.range_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.range_`."""
raise TypeError("Relax does not support imperative for loops")
def break_(*, span: _Span = None) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.break_`.
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.break_`.
Relax rejects loop-control statements.
"""
@@ -106,16 +106,16 @@ def break_(*, span: _Span = None) -> NoReturn:
def continue_(*, span: _Span = None) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.continue_`.
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.continue_`.
Relax rejects loop-control statements.
"""
raise TypeError("Relax does not support continue")
-# ------------
+# --------------------------------------
# Section: operator overloading
-# ------------
+# --------------------------------------
#
# ``X.if_then_else_(c, a, b)`` selects an expression.
# ``X.and_(a, b)`` constructs conjunction.
@@ -130,7 +130,7 @@ def continue_(*, span: _Span = None) -> NoReturn:
def if_then_else_(condition: Any, true_value: Any, false_value: Any) -> Any:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.if_then_else_`."""
+ """Implements
:func:`tvm.script.ir_builder.parser_protocol.if_then_else_`."""
if isinstance(condition, _ffi.ObjectConvertible):
condition = condition.asobject()
if not isinstance(condition, _ir.Expr):
@@ -152,22 +152,22 @@ def if_then_else_(condition: Any, true_value: Any,
false_value: Any) -> Any:
def and_(*values: Any) -> Any:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.and_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.and_`."""
return _builder.logical_and(*values)
def or_(*values: Any) -> Any:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.or_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.or_`."""
return _builder.logical_or(*values)
def not_(value: Any) -> Any:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.not_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.not_`."""
return _builder.logical_not(value)
def lt_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.lt_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.lt_`."""
if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for
value in (lhs, rhs)):
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
@@ -176,7 +176,7 @@ def lt_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
def le_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.le_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.le_`."""
if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for
value in (lhs, rhs)):
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
@@ -185,7 +185,7 @@ def le_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
def gt_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.gt_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.gt_`."""
if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for
value in (lhs, rhs)):
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
@@ -194,7 +194,7 @@ def gt_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
def ge_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.ge_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.ge_`."""
if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for
value in (lhs, rhs)):
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
@@ -203,7 +203,7 @@ def ge_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
def eq_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.eq_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.eq_`."""
if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for
value in (lhs, rhs)):
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
@@ -212,7 +212,7 @@ def eq_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
def ne_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.ne_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.ne_`."""
if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for
value in (lhs, rhs)):
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
@@ -220,9 +220,9 @@ def ne_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
return _prim_ffi._OpNE(lhs, rhs, _base.source_span(span))
-# ------------
+# --------------------------------------
# Section: context lookup and resolution
-# ------------
+# --------------------------------------
#
# ``X.resolve_global_info_(key)`` resolves module metadata.
# ``X.resolve_type_var_("n")`` resolves a symbolic dimension.
@@ -294,18 +294,18 @@ def resolve_type_var_(
value: _ir.Var | None = None,
span: _Span = None,
) -> _ir.Var:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.resolve_type_var_`."""
+ """Implements
:func:`tvm.script.ir_builder.parser_protocol.resolve_type_var_`."""
return _base._current_function_frame().resolve_type_var(name, dtype,
value=value, span=span)
def call_global_var_(function: _ir.GlobalVar, args: Sequence[Any]) -> _ir.Expr:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.call_global_var_`."""
+ """Implements
:func:`tvm.script.ir_builder.parser_protocol.call_global_var_`."""
return _relax.Call(function, [_relax.utils.convert_to_expr(value) for
value in args])
-# ------------
+# --------------------------------------
# Section: special protocol
-# ------------
+# --------------------------------------
#
# Syntax markers live in tvm.script.parser.protocol_registry.
# ``with X.function_(...) as fn:`` builds a function frame.
@@ -324,7 +324,7 @@ def function_(
reference: _ir.Var | None = None,
span: _Span = None,
) -> _frame.FunctionFrame:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.function_`.
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.function_`.
Public pure/private options pass to native purity and visibility controls.
Local function bodies require their previously declared reference.
@@ -339,7 +339,25 @@ def function_(
def arg(name: str, ty: Any, *, span: _Span = None) -> _ir.Var:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.arg`."""
+ """Declare a Relax parameter using the shared argument-hook contract.
+
+ Parameters
+ ----------
+ name : str
+ Source parameter name.
+ ty : Type, Var, or callable
+ Parameter annotation; corresponds to ``annotation`` in the shared
+ :func:`tvm.script.ir_builder.parser_protocol.arg` contract. Primitive
+ annotations resolve through the active signature symbol context; an
+ existing variable retains its identity.
+ span : Span or source location, optional
+ Source location attached to the constructed IR; None leaves it
unspecified.
+
+ Returns
+ -------
+ result : Var
+ The parameter registered in the current function frame.
+ """
if not isinstance(ty, _ir.Var):
ty = _builder._type(ty)
if isinstance(ty, _ir.PrimType) or _ir.is_prim_var(ty):
@@ -350,17 +368,32 @@ def arg(name: str, ty: Any, *, span: _Span = None) ->
_ir.Var:
def func_name(name: str) -> None:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.func_name`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.func_name`."""
return _native.func_name(name)
def func_ret_type(annotation: Any, *, span: _Span = None) -> None:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.func_ret_type`."""
+ """Implements
:func:`tvm.script.ir_builder.parser_protocol.func_ret_type`."""
return
_native.func_ret_type(_builder._type(_base._return_annotation(annotation)))
def check_well_formed_(function: _relax.Function) -> None:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.check_well_formed_`.
+ """Validate a completed Relax function.
+
+ Parameters
+ ----------
+ function : tvm.relax.Function
+ Completed function to validate, without changing its IR.
+
+ Raises
+ ------
+ ValueError
+ If the function fails language variant validation.
+
+ Notes
+ -----
+ See :func:`tvm.script.ir_builder.parser_protocol.check_well_formed_`
+ for the shared whole-module validation coordinator.
Validate the completed Relax function. Whole-module cross-function
validation is registered separately with the root coordinator.
@@ -388,9 +421,9 @@ def _check_module_well_formed(module: _ir.IRModule) -> None:
)
-# ------------
+# --------------------------------------
# Section: binding
-# ------------
+# --------------------------------------
#
# ``a = X.bind_(value, name="a")`` binds a source name.
# ``X.decl_mutable_cell_(value, ty=ty)`` declares mutable storage.
@@ -408,7 +441,7 @@ def bind_(
name_span: _Span = None,
frame_value: bool = False,
) -> Any:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.bind_`.
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.bind_`.
Relax values emit a binding or MatchCast; value_span belongs to the actual
RHS and span to the binding target. Metadata passes through unchanged.
@@ -460,7 +493,7 @@ def decl_mutable_cell_(
span: _Span = None,
name_span: _Span = None,
) -> NoReturn:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.decl_mutable_cell_`.
+ """Implements
:func:`tvm.script.ir_builder.parser_protocol.decl_mutable_cell_`.
Relax has immutable bindings and rejects mutable storage declarations.
"""
@@ -468,7 +501,7 @@ def decl_mutable_cell_(
def set_mutable_cell_(target: Any, value: Any, *, span: _Span = None) ->
NoReturn:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.set_mutable_cell_`.
+ """Implements
:func:`tvm.script.ir_builder.parser_protocol.set_mutable_cell_`.
Relax has immutable bindings and rejects mutable storage updates.
"""
@@ -476,7 +509,7 @@ def set_mutable_cell_(target: Any, value: Any, *, span:
_Span = None) -> NoRetur
def unpack(value: Any) -> Any:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.unpack`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.unpack`."""
if isinstance(value, _relax.Tuple):
return _python.tuple(value.fields)
if isinstance(value, _relax.Expr) and isinstance(value.ty, _ir.TupleType):
@@ -484,9 +517,9 @@ def unpack(value: Any) -> Any:
return value
-# ------------
+# --------------------------------------
# Section: statement
-# ------------
+# --------------------------------------
#
# ``X.emit_(value)`` emits an expression statement.
# ``X.return_(value)`` emits a function return.
@@ -496,7 +529,7 @@ def unpack(value: Any) -> Any:
def emit_(value: Any, *, span: _Span = None) -> None:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.emit_`.
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.emit_`.
Only void expressions may be discarded. Receipts and None emit nothing;
non-void expressions raise ValueError and unsupported host values raise
TypeError.
@@ -517,7 +550,7 @@ def emit_(value: Any, *, span: _Span = None) -> None:
def return_(value: Any = None, *, span: _Span = None) -> None:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.return_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.return_`."""
if value is None:
value = _relax.Tuple([])
# Normalization may emit bindings, but an existing result keeps its own
span.
@@ -525,12 +558,12 @@ def return_(value: Any = None, *, span: _Span = None) ->
None:
def setitem_(target: Any, key: Any, value: Any, *, span: _Span = None) ->
NoReturn:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.setitem_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.setitem_`."""
raise TypeError("Relax does not support indexed assignment")
def setattr_(target: Any, name: str, value: Any, *, span: _Span = None) ->
NoReturn:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.setattr_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.setattr_`."""
raise TypeError("Relax does not support attribute assignment")
@@ -540,7 +573,7 @@ def assert_(
*,
span: _Span = None,
) -> None:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.assert_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.assert_`."""
if not isinstance(message, _python.str):
raise TypeError("An assertion message must be construction-time text")
emit_(_base.at_(span, _native.assert_op(condition, format=message)),
span=span)
diff --git a/python/tvm/tirx/script/builder/__init__.py
b/python/tvm/script/ir_builder/tirx/__init__.py
similarity index 71%
rename from python/tvm/tirx/script/builder/__init__.py
rename to python/tvm/script/ir_builder/tirx/__init__.py
index 3cfcc6a6b1..babeaf151d 100644
--- a/python/tvm/tirx/script/builder/__init__.py
+++ b/python/tvm/script/ir_builder/tirx/__init__.py
@@ -97,7 +97,27 @@ is_type_var = _ir.is_prim_var
def type_var(name, *, dtype=None, span=None):
- """Construct an explicit standalone primitive symbol."""
+ """Construct a fresh standalone primitive symbol.
+
+ Parameters
+ ----------
+ name : str
+ Name of the symbol.
+ dtype : str or PrimType, optional
+ Primitive type of the symbol; None selects "int64".
+ span : Span or source location, optional
+ Source location attached to the constructed IR; None leaves it
unspecified.
+
+ Returns
+ -------
+ result : Var
+ The newly constructed primitive variable.
+
+ Notes
+ -----
+ This constructor creates a new symbol on each call. Use the language
variant
+ resolver for symbols shared by name within a function signature.
+ """
return _ir.Var(name, "int64" if dtype is None else dtype,
_source_span(span))
@@ -242,7 +262,24 @@ def _as_expr(value):
def emit(value):
- """Emit a standalone value, preserving already-emitted statement
receipts."""
+ """Emit a standalone value without repeating an earlier emission.
+
+ Parameters
+ ----------
+ value : Expr, Stmt, IRBuilderFrame, AlreadyEmitted, or sequence
+ Value consumed by the statement hook. AlreadyEmitted receipts and None
+ add no statement; other values follow the language variant ``emit_``
contract.
+
+ Returns
+ -------
+ None
+ No source-visible result.
+
+ Notes
+ -----
+ See :func:`tvm.script.ir_builder.parser_protocol.emit_` for construction
+ context, sequence handling, and supported inert values.
+ """
from tvm.script.ir_builder.base import AlreadyEmitted
if isinstance(value, AlreadyEmitted):
@@ -274,20 +311,67 @@ def grid(*extents, dtype=None):
@_mutable_cell_decl("T.alloc_scalar")
def alloc_scalar(dtype="float32", scope="global"):
- """Allocate scalar storage and return its load expression."""
+ """Allocate scalar storage and return its load expression.
+
+ Parameters
+ ----------
+ dtype : str, optional
+ Element dtype; defaults to "float32".
+ scope : str, optional
+ Storage scope for the one-element allocation; defaults to "global".
+
+ Returns
+ -------
+ result : TensorLoad
+ A load from the allocated scalar storage, usable as a mutable-cell
target.
+
+ Notes
+ -----
+ Requires an active allocation scope in the TIRx builder.
+ """
value = _native.alloc_scalar(dtype, scope)
return value.scalar if isinstance(value, _native.scalar_wrapper) else value
@_mutable_cell_decl("T.local_scalar")
def local_scalar(dtype="float32"):
- """Allocate scalar storage in local memory."""
+ """Allocate scalar storage in local memory.
+
+ Parameters
+ ----------
+ dtype : str, optional
+ Element dtype; defaults to "float32".
+
+ Returns
+ -------
+ result : TensorLoad
+ A load from the allocated scalar storage, usable as a mutable-cell
target.
+
+ Notes
+ -----
+ Equivalent to ``alloc_scalar(dtype, "local")`` in an active allocation
scope.
+ """
return alloc_scalar(dtype, "local")
@_mutable_cell_decl("T.shared_scalar")
def shared_scalar(dtype="float32"):
- """Allocate scalar storage in shared memory."""
+ """Allocate scalar storage in shared memory.
+
+ Parameters
+ ----------
+ dtype : str, optional
+ Element dtype; defaults to "float32".
+
+ Returns
+ -------
+ result : TensorLoad
+ A load from the allocated scalar storage, usable as a mutable-cell
target.
+
+ Notes
+ -----
+ Equivalent to ``alloc_scalar(dtype, "shared")`` in an active allocation
scope.
+ """
return alloc_scalar(dtype, "shared")
@@ -374,7 +458,25 @@ def match_buffer(*args, **kwargs):
def logical_and(*values):
- """Construct scalar or vector conjunction from eager operands."""
+ """Construct scalar or vector conjunction from eager operands.
+
+ Parameters
+ ----------
+ values : Expr or Python value
+ One or more operands. Object-convertible values are normalized first.
+ Host pairs follow Python logical operations; IR pairs use scalar
logical
+ or vector bitwise operations according to their types.
+
+ Returns
+ -------
+ result : Expr or Python value
+ The conjunction reduced from left to right.
+
+ Notes
+ -----
+ All arguments are evaluated before this call; it does not provide Python
+ short-circuit evaluation of the argument expressions.
+ """
if not values:
raise TypeError("logical_and requires at least one operand")
values = [
@@ -391,7 +493,25 @@ def logical_and(*values):
def logical_or(*values):
- """Construct scalar or vector disjunction from eager operands."""
+ """Construct scalar or vector disjunction from eager operands.
+
+ Parameters
+ ----------
+ values : Expr or Python value
+ One or more operands. Object-convertible values are normalized first.
+ Host pairs follow Python logical operations; IR pairs use scalar
logical
+ or vector bitwise operations according to their types.
+
+ Returns
+ -------
+ result : Expr or Python value
+ The disjunction reduced from left to right.
+
+ Notes
+ -----
+ All arguments are evaluated before this call; it does not provide Python
+ short-circuit evaluation of the argument expressions.
+ """
if not values:
raise TypeError("logical_or requires at least one operand")
values = [
@@ -408,14 +528,47 @@ def logical_or(*values):
def logical_not(value):
- """Negate a host or IR value without testing IR truth in Python."""
+ """Negate a host or IR value.
+
+ Parameters
+ ----------
+ value : Expr or Python value
+ Operand to negate. Object-convertible values are normalized first;
+ IR expressions use primitive Not and host values use Python not.
+
+ Returns
+ -------
+ result : Expr or bool
+ The logical negation without testing an IR expression as a Python bool.
+ """
if isinstance(value, _ffi.ObjectConvertible):
value = value.asobject()
return _tir.Not(value) if isinstance(value, _ir.Expr) else not value
def select(condition, true_value, false_value):
- """Construct a scalar conditional whose runtime evaluates one arm."""
+ """Construct a scalar conditional or select a host value.
+
+ Parameters
+ ----------
+ condition : Expr or Python value
+ Scalar IR condition or host truth value. Object-convertible conditions
+ are normalized before selection.
+ true_value : Expr or Python value
+ Value selected when the condition is true.
+ false_value : Expr or Python value
+ Value selected when the condition is false.
+
+ Returns
+ -------
+ result : Expr or Python value
+ A primitive if_then_else expression, or the selected host object.
+
+ Notes
+ -----
+ Both Python value arguments are constructed before this call. The generated
+ IR conditional evaluates only its selected arm at runtime.
+ """
if isinstance(condition, _ffi.ObjectConvertible):
condition = condition.asobject()
if not isinstance(condition, _ir.Expr):
diff --git a/python/tvm/tirx/script/builder/_ffi_api.py
b/python/tvm/script/ir_builder/tirx/_ffi_api.py
similarity index 100%
rename from python/tvm/tirx/script/builder/_ffi_api.py
rename to python/tvm/script/ir_builder/tirx/_ffi_api.py
diff --git a/python/tvm/tirx/script/builder/external_kernel.py
b/python/tvm/script/ir_builder/tirx/external_kernel.py
similarity index 99%
rename from python/tvm/tirx/script/builder/external_kernel.py
rename to python/tvm/script/ir_builder/tirx/external_kernel.py
index 6f308c7ef8..9bf3d52464 100644
--- a/python/tvm/tirx/script/builder/external_kernel.py
+++ b/python/tvm/script/ir_builder/tirx/external_kernel.py
@@ -213,7 +213,7 @@ def call_kernel(
kwargs : Dict[str, Any]
Additional keyword arguments to pass to the kernel or compilation.
"""
- from tvm.script.ir_builder.ir import ( # pylint:
disable=import-outside-toplevel
+ from tvm.script.ir_builder import ( # pylint:
disable=import-outside-toplevel
module_get_attr,
module_set_attr,
)
diff --git a/python/tvm/tirx/script/builder/frame.py
b/python/tvm/script/ir_builder/tirx/frame.py
similarity index 100%
rename from python/tvm/tirx/script/builder/frame.py
rename to python/tvm/script/ir_builder/tirx/frame.py
diff --git a/python/tvm/tirx/script/builder/ir.py
b/python/tvm/script/ir_builder/tirx/ir.py
similarity index 99%
rename from python/tvm/tirx/script/builder/ir.py
rename to python/tvm/script/ir_builder/tirx/ir.py
index 2b303d4a75..63a5e39606 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/script/ir_builder/tirx/ir.py
@@ -42,9 +42,9 @@ from tvm.ir import register_op_attr as _register_op_attr
from tvm.ir.base import deprecated
from tvm.ir.prim import _ffi_api as _prim_ffi_api
from tvm.runtime import convert
+from tvm.script.ir_builder import meta_var
from tvm.script.ir_builder.base import AlreadyEmitted, IRBuilder
-from tvm.script.ir_builder.ir import meta_var
-from tvm.script.ir_builder.ir.frame import IRModuleFrame
+from tvm.script.ir_builder.frame import IRModuleFrame
from tvm.script.parser.protocol_registry import (
mutable_cell_decl as _mutable_cell_decl,
)
@@ -1608,7 +1608,23 @@ def bind(
*,
var: Var | None = None, # pylint: disable=redefined-outer-name
) -> Var:
- """Create an immutable binding and return its native variable directly."""
+ """Create an immutable binding and return its native variable.
+
+ Parameters
+ ----------
+ value : Expr
+ Expression bound in the current statement frame.
+ type_annotation : Type or callable, optional
+ Explicit binding type, or a zero-argument annotation factory. A
variable
+ annotation supplies its type; None lets the native builder infer it.
+ var : Var, optional
+ Existing variable to bind. None asks the native builder to create one.
+
+ Returns
+ -------
+ result : Var
+ The same variable stored in the emitted Bind statement.
+ """
return Bind(value, type_annotation, var=var)
@@ -1785,7 +1801,19 @@ def While(condition: Expr) -> frame.WhileFrame: #
pylint: disable=invalid-name
def Return(value: Expr) -> AlreadyEmitted[tir.Stmt]: # pylint:
disable=invalid-name
- """Emit a return and retain the stored statement in an emission receipt."""
+ """Emit a return and retain the stored statement in an emission receipt.
+
+ Parameters
+ ----------
+ value : Expr
+ Expression returned by the active primitive function.
+
+ Returns
+ -------
+ result : AlreadyEmitted[Stmt]
+ Receipt holding the exact return statement emitted into the current
frame;
+ consuming it does not emit that statement again.
+ """
return AlreadyEmitted(_ffi_api.Return(value)) # type:
ignore[attr-defined] # pylint: disable=no-member
@@ -3084,7 +3112,7 @@ def register_script_namespace(name: str, namespace:
object) -> object:
import sys # pylint: disable=import-outside-toplevel
for module_name in [
- "tvm.tirx.script.builder",
+ "tvm.script.ir_builder.tirx",
"tvm.tirx.script",
"tvm.script.tirx",
]:
diff --git a/python/tvm/tirx/script/builder/parser_protocol.py
b/python/tvm/script/ir_builder/tirx/parser_protocol.py
similarity index 85%
rename from python/tvm/tirx/script/builder/parser_protocol.py
rename to python/tvm/script/ir_builder/tirx/parser_protocol.py
index e2e5bac325..32fa88528e 100644
--- a/python/tvm/tirx/script/builder/parser_protocol.py
+++ b/python/tvm/script/ir_builder/tirx/parser_protocol.py
@@ -19,7 +19,7 @@
Hooks delegate construction to this language variant's native builder and IR
APIs.
For example, generated ``X.if_(condition)`` creates the native conditional
frame;
``X.then_()`` and ``X.else_()`` enter its branches. See the corresponding
shared
-``tvm.script.ir_builder.ir.parser_protocol`` hooks for operand and span
contracts.
+``tvm.script.ir_builder.parser_protocol`` hooks for operand and span contracts.
"""
from __future__ import annotations
@@ -29,13 +29,13 @@ from collections.abc import Sequence
from functools import partial as _partial
from typing import Any
+import tvm.script.ir_builder.tirx as _builder
from tvm import ir as _ir
from tvm import tirx as _tir
from tvm.ir.prim import _ffi_api as _prim_ffi
from tvm.script.ir_builder import IRBuilder as _IRBuilder
from tvm.script.ir_builder import base as _base
-from .. import builder as _builder
from . import _ffi_api
from . import frame as _frame
from . import ir as _native
@@ -43,9 +43,9 @@ from . import ir as _native
_Span = _base.SpanEntry | _ir.Span | tuple[_ir.SourceName, int, int, int, int]
| None
-# ------------
+# --------------------------------------
# Section: control flow
-# ------------
+# --------------------------------------
#
# ``with X.if_(cond):`` opens a conditional frame.
# ``with X.then_():`` enters its true branch.
@@ -58,24 +58,24 @@ _Span = _base.SpanEntry | _ir.Span | tuple[_ir.SourceName,
int, int, int, int] |
def if_(condition: Any, *, span: _Span = None) -> _frame.IfFrame:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.if_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.if_`."""
return _base.at_(span, _native.If(condition))
def then_(*, span: _Span = None) -> _frame.ThenFrame:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.then_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.then_`."""
return _base.at_(span, _native.Then())
def else_(*, span: _Span = None) -> _frame.ElseFrame:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.else_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.else_`."""
return _base.at_(span, _native.Else())
def for_(
iterable: Any, *, names: str | Sequence[str] | None = None, span: _Span =
None
) -> _frame.ForFrame:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.for_`.
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.for_`.
A single native loop returns its scalar Var; multiple loops return their
sequence. frame.vars remains the stable sequence for source unpacking.
@@ -89,12 +89,12 @@ def for_(
def while_(condition: Any, *, span: _Span = None) -> _frame.WhileFrame:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.while_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.while_`."""
return _base.at_(span, _native.While(condition))
def range_(*args: Any, annotations: dict[str, Any] | None = None) ->
_frame.ForFrame:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.range_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.range_`."""
if len(args) == 1:
args = (0, args[0], None)
elif len(args) == 2:
@@ -107,7 +107,7 @@ def range_(*args: Any, annotations: dict[str, Any] | None =
None) -> _frame.ForF
def break_(*, span: _Span = None) -> _base.AlreadyEmitted[_tir.Stmt]:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.break_`.
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.break_`.
Legality is checked on the completed function, across loop and function
boundaries.
"""
@@ -115,16 +115,16 @@ def break_(*, span: _Span = None) ->
_base.AlreadyEmitted[_tir.Stmt]:
def continue_(*, span: _Span = None) -> _base.AlreadyEmitted[_tir.Stmt]:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.continue_`.
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.continue_`.
Legality is checked on the completed function, across loop and function
boundaries.
"""
return _base.at_(span, _native.evaluate(_native.continue_loop()))
-# ------------
+# --------------------------------------
# Section: operator overloading
-# ------------
+# --------------------------------------
#
# ``X.if_then_else_(c, a, b)`` selects an expression.
# ``X.and_(a, b)`` constructs conjunction.
@@ -139,58 +139,58 @@ def continue_(*, span: _Span = None) ->
_base.AlreadyEmitted[_tir.Stmt]:
def if_then_else_(condition: Any, true_value: Any, false_value: Any) -> Any:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.if_then_else_`."""
+ """Implements
:func:`tvm.script.ir_builder.parser_protocol.if_then_else_`."""
return _builder.select(condition, true_value, false_value)
def and_(*values: Any) -> Any:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.and_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.and_`."""
return _builder.logical_and(*values)
def or_(*values: Any) -> Any:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.or_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.or_`."""
return _builder.logical_or(*values)
def not_(value: Any) -> Any:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.not_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.not_`."""
return _builder.logical_not(value)
def lt_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.lt_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.lt_`."""
return _prim_ffi._OpLT(lhs, rhs, _base.source_span(span))
def le_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.le_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.le_`."""
return _prim_ffi._OpLE(lhs, rhs, _base.source_span(span))
def gt_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.gt_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.gt_`."""
return _prim_ffi._OpGT(lhs, rhs, _base.source_span(span))
def ge_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.ge_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.ge_`."""
return _prim_ffi._OpGE(lhs, rhs, _base.source_span(span))
def eq_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.eq_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.eq_`."""
return _prim_ffi._OpEQ(lhs, rhs, _base.source_span(span))
def ne_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.ne_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.ne_`."""
return _prim_ffi._OpNE(lhs, rhs, _base.source_span(span))
-# ------------
+# --------------------------------------
# Section: context lookup and resolution
-# ------------
+# --------------------------------------
#
# ``X.resolve_global_info_(key)`` resolves module metadata.
# ``X.resolve_type_var_("n")`` resolves a symbolic dimension.
@@ -198,7 +198,7 @@ def ne_(lhs: Any, rhs: Any, *, span: _Span = None) ->
_ir.Expr:
def resolve_global_info_(content: Any) -> Any:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.resolve_global_info_`.
+ """Implements
:func:`tvm.script.ir_builder.parser_protocol.resolve_global_info_`.
TIRx does not define global-info selectors.
"""
@@ -212,18 +212,18 @@ def resolve_type_var_(
value: _ir.Var | None = None,
span: _Span = None,
) -> _ir.Var:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.resolve_type_var_`."""
+ """Implements
:func:`tvm.script.ir_builder.parser_protocol.resolve_type_var_`."""
return _base._current_function_frame().resolve_type_var(name, dtype,
value=value, span=span)
def call_global_var_(function: _ir.GlobalVar, args: Sequence[Any]) -> _ir.Expr:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.call_global_var_`."""
+ """Implements
:func:`tvm.script.ir_builder.parser_protocol.call_global_var_`."""
return _native._call_global(function, *args)
-# ------------
+# --------------------------------------
# Section: special protocol
-# ------------
+# --------------------------------------
#
# Syntax markers live in tvm.script.parser.protocol_registry.
# ``with X.function_(...) as fn:`` builds a function frame.
@@ -242,7 +242,7 @@ def function_(
decl: bool = False,
span: _Span = None,
) -> _frame.PrimFuncFrame:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.function_`.
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.function_`.
Public private/s_tir/persistent options pass to the native function frame;
is_stir supplies the JIT spelling. The same frame supports declaration and
body entry.
@@ -258,7 +258,7 @@ def function_(
def arg(name: str, annotation: Any, *, span: _Span = None) -> _ir.Var:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.arg`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.arg`."""
if getattr(annotation, "__tvm_optional_annotation__", None) is not None:
raise TypeError("T.Optional is only supported by @T.jit")
if callable(annotation) and not isinstance(annotation, _ir.Expr):
@@ -287,12 +287,12 @@ def arg(name: str, annotation: Any, *, span: _Span =
None) -> _ir.Var:
def func_name(name: str) -> None:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.func_name`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.func_name`."""
return _native.func_name(name)
def func_ret_type(annotation: Any, *, span: _Span = None) -> None:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.func_ret_type`."""
+ """Implements
:func:`tvm.script.ir_builder.parser_protocol.func_ret_type`."""
annotation = _base._return_annotation(annotation)
if callable(annotation) and not isinstance(annotation, _ir.Expr |
_ir.Type):
annotation = annotation()
@@ -302,7 +302,22 @@ def func_ret_type(annotation: Any, *, span: _Span = None)
-> None:
def check_well_formed_(function: _tir.PrimFunc) -> None:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.check_well_formed_`.
+ """Validate a completed TIRx function.
+
+ Parameters
+ ----------
+ function : tvm.tirx.PrimFunc
+ Completed function to validate, without changing its IR.
+
+ Raises
+ ------
+ ValueError
+ If the function fails language variant validation.
+
+ Notes
+ -----
+ See :func:`tvm.script.ir_builder.parser_protocol.check_well_formed_`
+ for the shared whole-module validation coordinator.
The completed native function owns semantic loop and IR validation.
Module-wide S-TIR and TIRx checks are supplied separately to the root
coordinator.
@@ -338,9 +353,9 @@ def _check_module_well_formed(module: _ir.IRModule) -> None:
raise ValueError(f"{message}\n{error}") from error
-# ------------
+# --------------------------------------
# Section: binding
-# ------------
+# --------------------------------------
#
# ``a = X.bind_(value, name="a")`` binds a source name.
# ``X.decl_mutable_cell_(value, ty=ty)`` declares mutable storage.
@@ -371,7 +386,7 @@ def bind_(
name_span: _Span = None,
frame_value: bool = False,
) -> Any:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.bind_`.
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.bind_`.
Returned Vars, including buffers, and metadata retain identity, names and
spans.
Other expressions create native Bind nodes; value_span belongs to the RHS.
@@ -441,7 +456,7 @@ def decl_mutable_cell_(
span: _Span = None,
name_span: _Span = None,
) -> Any:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.decl_mutable_cell_`.
+ """Implements
:func:`tvm.script.ir_builder.parser_protocol.decl_mutable_cell_`.
Primitive annotations allocate scalar local storage; vector annotations
allocate their declared shape. Buffer declaration producers retain their
own effects.
@@ -477,7 +492,7 @@ def decl_mutable_cell_(
def set_mutable_cell_(
target: _ir.TensorLoad | _ir.Var | _native.scalar_wrapper, value: Any, *,
span: _Span = None
) -> _base.AlreadyEmitted[_tir.Stmt]:
- """Implements
:func:`tvm.script.ir_builder.ir.parser_protocol.set_mutable_cell_`.
+ """Implements
:func:`tvm.script.ir_builder.parser_protocol.set_mutable_cell_`.
Updates emit a scalar buffer store. Targets must denote scalar storage.
"""
@@ -497,7 +512,7 @@ def set_mutable_cell_(
def unpack(value: Any) -> Any:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.unpack`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.unpack`."""
if isinstance(value, _ir.Tuple):
return _python.tuple(value.fields)
if isinstance(value, _ir.Expr) and isinstance(value.ty, _ir.TupleType):
@@ -505,9 +520,9 @@ def unpack(value: Any) -> Any:
return value
-# ------------
+# --------------------------------------
# Section: statement
-# ------------
+# --------------------------------------
#
# ``X.emit_(value)`` emits an expression statement.
# ``X.return_(value)`` emits a function return.
@@ -517,7 +532,7 @@ def unpack(value: Any) -> Any:
def emit_(value: Any, *, span: _Span = None) -> None:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.emit_`.
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.emit_`.
Native statements emit once; receipts are already emitted. Vars, layouts
and meta_class instances are inert. Sequences are consumed
@@ -553,7 +568,7 @@ def emit_(value: Any, *, span: _Span = None) -> None:
def return_(value: Any = None, *, span: _Span = None) ->
_base.AlreadyEmitted[_tir.Stmt]:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.return_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.return_`."""
if value is None:
raise TypeError("A primitive function return requires an expression")
return _base.with_at_group_(span, lambda:
_native.Return(_builder._as_expr(value)))
@@ -562,14 +577,14 @@ def return_(value: Any = None, *, span: _Span = None) ->
_base.AlreadyEmitted[_t
def setitem_(
target: Any, key: Any, value: Any, *, span: _Span = None
) -> _base.AlreadyEmitted[_tir.Stmt]:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.setitem_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.setitem_`."""
return _base.at_(span, _builder.buffer_store(target, value, key))
def setattr_(
target: Any, name: str, value: Any, *, span: _Span = None
) -> _base.AlreadyEmitted[_tir.Stmt] | None:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.setattr_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.setattr_`."""
previous = getattr(target, name, _base.MISSING)
if isinstance(previous, _native.scalar_wrapper):
previous = previous.scalar
@@ -587,7 +602,7 @@ def assert_(
*,
span: _Span = None,
) -> None:
- """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.assert_`."""
+ """Implements :func:`tvm.script.ir_builder.parser_protocol.assert_`."""
kind = "RuntimeError"
if isinstance(message, tuple):
if len(message) != 2 or not isinstance(message[0], str):
diff --git a/python/tvm/tirx/script/builder/tirx.py
b/python/tvm/script/ir_builder/tirx/tirx.py
similarity index 100%
rename from python/tvm/tirx/script/builder/tirx.py
rename to python/tvm/script/ir_builder/tirx/tirx.py
diff --git a/python/tvm/tirx/script/builder/triton.py
b/python/tvm/script/ir_builder/tirx/triton.py
similarity index 100%
rename from python/tvm/tirx/script/builder/triton.py
rename to python/tvm/script/ir_builder/tirx/triton.py
diff --git a/python/tvm/tirx/script/builder/utils.py
b/python/tvm/script/ir_builder/tirx/utils.py
similarity index 98%
rename from python/tvm/tirx/script/builder/utils.py
rename to python/tvm/script/ir_builder/tirx/utils.py
index 5b60b16cf9..7f72c46008 100644
--- a/python/tvm/tirx/script/builder/utils.py
+++ b/python/tvm/script/ir_builder/tirx/utils.py
@@ -213,7 +213,7 @@ def buffer_proxy(buf: Buffer) -> _BufferProxy:
--------
.. code-block:: python
- from tvm.tirx.script.builder.utils import buffer_proxy
+ from tvm.script.ir_builder.tirx.utils import buffer_proxy
buf = tvm.tirx.decl_buffer([2, 3], "float32")
ptr = buffer_proxy(buf)
diff --git a/python/tvm/script/parser/__init__.py
b/python/tvm/script/parser/__init__.py
index c67be91e9d..46c85f1685 100644
--- a/python/tvm/script/parser/__init__.py
+++ b/python/tvm/script/parser/__init__.py
@@ -41,6 +41,23 @@ _initializing = False
def register_namespace(alias: str, namespace: object) -> None:
"""Register an opaque fixed namespace; its first alias is the syntax key
root.
+ Parameters
+ ----------
+ alias : str
+ Source name for the namespace, such as ``"I"`` or ``"M"``. The first
+ registered alias for an object is its canonical syntax-policy prefix;
+ additional aliases for that object resolve to the same prefix.
+ namespace : object
+ Namespace object exposing construction operations or source decorators.
+ It is borrowed without copying or inspecting its members.
+
+ Returns
+ -------
+ None
+ Registration updates the namespaces available to subsequent parses.
+
+ Notes
+ -----
Each parse borrows these namespaces. Replacing an alias affects future
parses
only; registration enters no builder frame and inspects no namespace
members.
"""
@@ -50,7 +67,28 @@ def register_namespace(alias: str, namespace: object) ->
None:
def register_namespace_initializer(
initializer: Callable[[], None], *, aliases: tuple[str, ...] = ()
) -> None:
- """Register a lazy bootstrap callback without importing the language
variant."""
+ """Register a lazy bootstrap callback without importing the language
variant.
+
+ Parameters
+ ----------
+ initializer : Callable[[], None]
+ Callback that registers the language variant's namespaces. It receives
+ no arguments and is called during the parser's first initialization.
+ If initialization has already completed, a newly registered callback
+ runs immediately. Re-registering the same callback object does not
+ enqueue or invoke it again.
+ aliases : tuple[str, ...], optional
+ Names whose attribute lookup should trigger initialization. Defaults
+ to an empty tuple. These names advertise lazy namespaces; the callback
+ must still register their actual objects with
:func:`register_namespace`.
+ Aliases are added even when the callback was previously registered.
+
+ Returns
+ -------
+ None
+ The callback is registered, and invoked immediately when the parser
+ is already initialized.
+ """
_NAMESPACE_ALIASES.update(aliases)
if any(existing is initializer for existing in _NAMESPACE_INITIALIZERS):
return
diff --git a/python/tvm/script/parser/entry.py
b/python/tvm/script/parser/entry.py
index 0d80724773..44aae05d9d 100644
--- a/python/tvm/script/parser/entry.py
+++ b/python/tvm/script/parser/entry.py
@@ -30,8 +30,8 @@ from types import FrameType, FunctionType
from typing import TYPE_CHECKING, Any, TypeVar
from tvm.ir import SourceName, Span
+from tvm.script import ir_builder as builder_ir
from tvm.script.ir_builder import base
-from tvm.script.ir_builder import ir as builder_ir
from tvm.script.ir_builder.base import SpanEntry
from . import _NAMESPACES, _initialize, jit_support
@@ -258,9 +258,25 @@ def make_decorator(
check_well_formed : bool, optional
Whether to check that the constructed function is well formed.
Defaults to True.
+ s_tir : bool, optional
+ For ``T.prim_func``, select scheduled-TIR construction semantics:
+ buffers have no default layout and the body is wrapped in a root
+ block. Defaults to False, selecting TIRx construction. See
+ :func:`tvm.script.ir_builder.tirx.prim_func`.
+ persistent : bool, optional
+ For ``T.prim_func``, mark the resulting function as a persistent
+ kernel. Defaults to False. See
+ :func:`tvm.script.ir_builder.tirx.prim_func`.
+ pure : bool, optional
+ For ``R.function``, declare whether the function is pure, meaning
+ that it has no observable side effects. Defaults to True. See
+ :func:`tvm.script.ir_builder.relax.function_`.
**options
- Additional language variant options. ``T.prim_func`` accepts
``s_tir`` and
- ``persistent``; ``R.function`` accepts ``pure``.
+ Keyword options are forwarded to the selected language variant's
+ :func:`~tvm.script.ir_builder.parser_protocol.function_` hook,
+ except ``check_well_formed``, which controls parser validation.
+ Options supported by only one language variant are not shared
+ between ``T.prim_func`` and ``R.function``.
Returns
-------
@@ -321,7 +337,7 @@ def make_macro_decorator(
Accepts a function directly or keyword options. The ``hygienic``
option defaults to True and snapshots the definition environment;
False captures the calling environment on each invocation. Other
- options remain metadata for the namespace consumer.
+ keyword options are accepted but do not affect helper construction.
Raises
------
@@ -350,6 +366,10 @@ def make_macro_decorator(
instead of its calling environment. Defaults to True. ``T.macro``
and ``R.macro`` capture values at definition time; ``T.inline``
refreshes captured closure cells when called.
+ **options
+ Keyword configuration for this helper decorator. ``hygienic``
+ controls name lookup as described above; other options are
+ accepted but are not consumed or forwarded to builder hooks.
Returns
-------
diff --git a/python/tvm/script/parser/ir.py b/python/tvm/script/parser/ir.py
index 5e32059eb6..86541c85f4 100644
--- a/python/tvm/script/parser/ir.py
+++ b/python/tvm/script/parser/ir.py
@@ -18,7 +18,7 @@
import sys as _sys
-from tvm.script.ir_builder.ir import * # noqa: F403
+from tvm.script.ir_builder import * # noqa: F403
from . import register_namespace
from .entry import ir_module as ir_module
diff --git a/python/tvm/script/parser/protocol_registry.py
b/python/tvm/script/parser/protocol_registry.py
index 0b4094ae8a..c336a17781 100644
--- a/python/tvm/script/parser/protocol_registry.py
+++ b/python/tvm/script/parser/protocol_registry.py
@@ -28,7 +28,7 @@ Tables retain static syntax facts only, never source
functions, captures, frames
or constructed results. Registration decorators return the same callable,
except
that ``args_policy`` preserves its existing builder-owned eager annotation
adapter.
Shared builder operations are documented in
-``tvm.script.ir_builder.ir.parser_protocol``; concrete language variants own
registration
+``tvm.script.ir_builder.parser_protocol``; concrete language variants own
registration
and namespace initialization.
"""
@@ -81,6 +81,20 @@ DECLARATION_KIND: dict[str, Literal["function", "helper"]] =
{}
def constexpr(value: object) -> NoReturn:
"""Mark host control syntax or a JIT specialization annotation.
+ Parameters
+ ----------
+ value : object
+ Source expression to evaluate with ordinary Python semantics in a
+ supported control-flow position. The marker itself can also appear
+ as an annotation identifying a JIT specialization parameter.
+
+ Raises
+ ------
+ TypeError
+ If invoked directly instead of being recognized in parsed source.
+
+ Notes
+ -----
The parser recognizes this marker through its fixed namespace path and
removes it before execution. Host operators retain ordinary Python
behavior.
Direct invocation raises TypeError; no builder frame or IR is constructed.
@@ -113,12 +127,14 @@ def args_policy(
Registered namespace alias followed by the exported member path.
fields : Mapping[str, str]
Parameter names mapped to ``expr_str`` or ``global_info``.
- scalar_strings : bool
- Whether bare strings in expression fields denote expressions.
- dtype : object
- Opaque dtype passed to the language variant's symbol resolver.
- as_type : bool
+ scalar_strings : bool, optional
+ Whether bare strings in expression fields denote expressions. Defaults
to True.
+ dtype : object, optional
+ Opaque dtype passed to the language variant's symbol resolver. None
(the default)
+ leaves dtype selection to that resolver.
+ as_type : bool, optional
Preserve the eager annotation-class surface, including Python type
unions.
+ Defaults to False.
Returns
-------
@@ -200,9 +216,27 @@ def register_type_var_decl(
) -> _Callable:
"""Register symbolic declaration syntax and return the unchanged
constructor.
- ``namespace_path`` identifies the fixed exported callable. Zero-argument
- calls denote declarations; ``dtype`` supplies static scalar type facts.
- Dictionary membership distinguishes an unspecified dtype from an
unregistered
+ Parameters
+ ----------
+ namespace_path : str
+ Canonical registered namespace alias and exported callable path, such
+ as ``"T.int32"``. A later registration at this path replaces its dtype.
+ constructor : Callable
+ Callable providing the eager construction operation. Registration
+ records its syntax path without invoking or wrapping this callable.
+ dtype : object, optional
+ Static scalar dtype forwarded to the language variant's symbol
resolver.
+ None (the default) leaves the dtype unspecified for that resolver.
+
+ Returns
+ -------
+ Callable
+ The exact ``constructor`` object.
+
+ Notes
+ -----
+ Zero-argument calls denote declarations. Dictionary membership
distinguishes
+ an unspecified dtype from an unregistered
constructor. The parser predeclares symbols needed by signatures while the
native function owns their identity. No constructor executes at
registration.
@@ -221,9 +255,29 @@ def mutable_cell_decl(
) -> Callable[[_Callable], _Callable]:
"""Register persistent mutable storage syntax at a fixed namespace path.
- ``syntax`` is ``call``, ``annotation`` or ``parameter``. Repeated
registration
- adds a syntax kind for the same path. The decorator returns its callable
- unchanged. Explicit declarations take precedence over mutable target
updates;
+ Parameters
+ ----------
+ namespace_path : str
+ Canonical registered namespace alias and exported constructor path.
+ syntax : str, optional
+ Declaration form: ``"call"`` (the default) for a constructor call,
+ ``"annotation"`` for an annotated local declaration, or ``"parameter"``
+ for a function parameter annotation. Repeated registration adds forms
+ for the same path without removing existing ones.
+
+ Returns
+ -------
+ Callable
+ Registration decorator that returns its constructor unchanged.
+
+ Raises
+ ------
+ ValueError
+ If ``syntax`` is not one of the three supported forms.
+
+ Notes
+ -----
+ Explicit declarations take precedence over mutable target updates;
storage creation and updates stay outside ordinary ``bind_``.
.. code:: python
@@ -247,10 +301,23 @@ def mutable_cell_decl(
def result_span(namespace_path: str) -> Callable[[_Callable], _Callable]:
"""Declare that a call's complete IR effect is represented by its result.
+ Parameters
+ ----------
+ namespace_path : str
+ Canonical registered namespace alias and exported callable path whose
+ returned object carries the call's complete IR effect.
+
+ Returns
+ -------
+ Callable
+ Registration decorator that returns its callable unchanged.
+
+ Notes
+ -----
The unchanged callable owns no unrelated emitted statements needing caller
context. This permits result attachment instead of an opaque-call context;
- argument instrumentation and ordinary binding/emission remain. D64 passes
- assignment RHS locations separately to the language variant's ``bind_``.
+ argument instrumentation and ordinary binding/emission remain. Assignment
+ RHS locations pass separately to the language variant's ``bind_``.
.. code:: python
@@ -271,6 +338,19 @@ def result_span(namespace_path: str) ->
Callable[[_Callable], _Callable]:
def module_decorator(namespace_path: str) -> Callable[[_Callable], _Callable]:
"""Mark a module entry point and return the same callable.
+ Parameters
+ ----------
+ namespace_path : str
+ Canonical registered namespace alias and exported module-decorator
path,
+ such as ``"I.ir_module"``.
+
+ Returns
+ -------
+ Callable
+ Registration decorator that returns the module entry point unchanged.
+
+ Notes
+ -----
Member function decorators use this fixed namespace syntax to defer parsing
until the complete enclosing module is available. No source-function
records,
captured scopes or constructed module results are retained in the table.
@@ -294,9 +374,23 @@ def declaration_kind(
) -> Callable[[_Callable], _Callable]:
"""Classify a fixed namespace decorator's declaration syntax.
- ``function`` constructs an IR function; ``helper`` keeps ordinary Python
- helper semantics, including macro expansion or retained ``I.pyfunc``
- callables. Unregistered decorators have neither kind.
+ Parameters
+ ----------
+ namespace_path : str
+ Canonical registered namespace alias and exported source-decorator
path.
+ kind : {"function", "helper"}
+ ``"function"`` declares an IR function. ``"helper"`` retains ordinary
+ Python helper semantics, including macro expansion and ``I.pyfunc``.
+ A later registration at the same path replaces its declaration kind.
+
+ Returns
+ -------
+ Callable
+ Registration decorator that returns its source decorator unchanged.
+
+ Notes
+ -----
+ Unregistered decorators have neither declaration kind.
Concrete namespaces register their entry points where they expose them::
prim_func = declaration_kind("T.prim_func",
"function")(make_decorator(builder))
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 924ae415d1..38158bff41 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -122,7 +122,7 @@ tvm.script.register_dialect("tirx", "tvm.tirx.script")
def _check_script_module(module: "tvm.ir.IRModule") -> None:
# Delay builder imports until validation, keeping dialect/runtime
bootstrap safe.
- from .script.builder.parser_protocol import _check_module_well_formed
+ from tvm.script.ir_builder.tirx.parser_protocol import
_check_module_well_formed
_check_module_well_formed(module)
diff --git a/python/tvm/tirx/_buffer_view.py b/python/tvm/tirx/_buffer_view.py
index 5a384d3ee0..091ae9ea41 100644
--- a/python/tvm/tirx/_buffer_view.py
+++ b/python/tvm/tirx/_buffer_view.py
@@ -40,7 +40,7 @@ def _redecl(buf: Buffer, shape, layout, *, dtype=None,
elem_offset=None, addr_of
addr = buf.allocated_addr[0]
if addr_offset is not None:
addr = addr + addr_offset
- return tvm.tirx.script.builder.decl_buffer(
+ return tvm.script.ir_builder.tirx.decl_buffer(
shape,
buf.dtype if dtype is None else dtype,
None,
@@ -53,7 +53,7 @@ def _redecl(buf: Buffer, shape, layout, *, dtype=None,
elem_offset=None, addr_of
layout,
allocated_addr=addr,
)
- return tvm.tirx.script.builder.decl_buffer(
+ return tvm.script.ir_builder.tirx.decl_buffer(
shape,
buf.dtype if dtype is None else dtype,
buf.data,
diff --git a/python/tvm/tirx/script/__init__.py
b/python/tvm/tirx/script/__init__.py
index ad42c61a48..2ab9680e1d 100644
--- a/python/tvm/tirx/script/__init__.py
+++ b/python/tvm/tirx/script/__init__.py
@@ -31,11 +31,12 @@ def _initialize() -> None:
return
_initializing = True
try:
+ from tvm.script.ir_builder import tirx as builder
from tvm.script.parser import entry, register_namespace
from tvm.script.parser.protocol_registry import declaration_kind
from tvm.tirx.layout import Axis
- from . import builder, tile
+ from . import tile
from .jit import make_jit
globals().update(
@@ -67,14 +68,14 @@ def _initialize() -> None:
def __getattr__(name: str) -> _Any:
- if name in ("builder", "tile"):
- return _importlib.import_module(__name__ + "." + name)
+ if name == "tile":
+ return _importlib.import_module(__name__ + ".tile")
if name.startswith("_") and name != "__all__":
raise AttributeError(name)
_initialize()
if name in globals():
return globals()[name]
- return getattr(_importlib.import_module(__name__ + ".builder"), name)
+ return getattr(_importlib.import_module("tvm.script.ir_builder.tirx"),
name)
from .jit import OptionalAnnotation as Optional
diff --git a/python/tvm/tirx/script/jit.py b/python/tvm/tirx/script/jit.py
index 9cc507300b..b93a5b9bb1 100644
--- a/python/tvm/tirx/script/jit.py
+++ b/python/tvm/tirx/script/jit.py
@@ -257,7 +257,7 @@ class TIRJit:
return cached
if self.builder is None:
- from tvm.tirx.script import builder
+ from tvm.script.ir_builder import tirx as builder
else:
builder = self.builder
prim_func = parse(
diff --git a/python/tvm/tirx/script/tile.py b/python/tvm/tirx/script/tile.py
index e891c495ec..a9d4921f03 100644
--- a/python/tvm/tirx/script/tile.py
+++ b/python/tvm/tirx/script/tile.py
@@ -19,10 +19,9 @@
import functools
from tvm.ir import TensorRegion
+from tvm.script.ir_builder.tirx import tirx as _builder
from tvm.tirx import is_buffer_var
-from .builder import tirx as _builder
-
def _get_arg(args, kwargs, index, name):
if len(args) > index:
diff --git a/tests/python/codegen/test_target_codegen_vulkan.py
b/tests/python/codegen/test_target_codegen_vulkan.py
index 9faa2eb250..9ac20f0236 100644
--- a/tests/python/codegen/test_target_codegen_vulkan.py
+++ b/tests/python/codegen/test_target_codegen_vulkan.py
@@ -24,9 +24,9 @@ import pytest
import tvm
import tvm.testing
from tvm.script import ir as I
+from tvm.script import ir_builder as I_builder
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import ir as I_builder
from tvm.script.ir_builder import tirx as T_builder
from tvm.testing import env
diff --git a/tests/python/tirx/test_op.py b/tests/python/tirx/test_op.py
index 8d6e6326e1..bc728d7863 100644
--- a/tests/python/tirx/test_op.py
+++ b/tests/python/tirx/test_op.py
@@ -103,7 +103,7 @@ def test_buffer_replacer_replaces_strides_and_elem_offset():
def test_gemm_async_partial_scale_factor():
"""Regression test for F7: gemm_async must reject partial scale factors."""
- from tvm.tirx.script.builder.tirx import gemm_async
+ from tvm.script.ir_builder.tirx.tirx import gemm_async
A = decl_buffer((64, 64), "float16", scope="shared")
B = decl_buffer((64, 64), "float16", scope="shared")
diff --git a/tests/python/tirx/test_op_namespace_cleanup.py
b/tests/python/tirx/test_op_namespace_cleanup.py
index 2f2d4ff862..6f23681a3c 100644
--- a/tests/python/tirx/test_op_namespace_cleanup.py
+++ b/tests/python/tirx/test_op_namespace_cleanup.py
@@ -160,7 +160,7 @@ def
test_device_intrinsic_namespaces_are_canonical_and_classified():
)
from tvm.backend.metal.script import MetalNamespace as
BackendMetalNamespace
from tvm.backend.trn.script import NKINamespace as BackendNKINamespace
- from tvm.tirx.script.builder import ir as builder_ir
+ from tvm.script.ir_builder.tirx import ir as builder_ir
assert isinstance(builder_ir.cuda, BackendCUDANamespace)
assert isinstance(builder_ir.s_tir, BackendSTIRNamespace)
@@ -229,9 +229,9 @@ def test_backend_specific_wrappers_are_not_root_exports():
def test_backend_load_updates_tirx_alias_and_script_facades(monkeypatch):
+ from tvm.script.ir_builder import tirx as builder
+ from tvm.script.ir_builder.tirx import ir as builder_ir
from tvm.script.parser import tirx as parser
- from tvm.tirx.script import builder
- from tvm.tirx.script.builder import ir as builder_ir
backend_name = "unit_test_backend"
backend_module_name = f"tvm.backend.{backend_name}"
diff --git a/tests/python/tirx/test_parser_printer.py
b/tests/python/tirx/test_parser_printer.py
index 0d2713252f..d507cad6bc 100644
--- a/tests/python/tirx/test_parser_printer.py
+++ b/tests/python/tirx/test_parser_printer.py
@@ -1088,7 +1088,7 @@ def test_scalar_assign_error_not_swallowed():
swallowed and the assignment would silently fall through to eval_assign."""
from unittest.mock import patch
- original = tvm.tirx.script.builder.buffer_store
+ original = tvm.script.ir_builder.tirx.buffer_store
def bomb(*args, **kwargs):
# Intercept only the scalar-assignment path (indices == [0])
@@ -1107,7 +1107,7 @@ def func():
"""
# The ValueError propagates unchanged. A broad ``except Exception`` here
# previously swallowed it and fell through to eval_assign.
- with patch("tvm.tirx.script.builder.buffer_store", side_effect=bomb):
+ with patch("tvm.script.ir_builder.tirx.buffer_store", side_effect=bomb):
with pytest.raises(ValueError, match="boom"):
from_source(src)
diff --git a/tests/python/tvmscript/test_tvmscript_ir_builder_irmodule.py
b/tests/python/tvmscript/test_tvmscript_ir_builder_irmodule.py
index ca46b6a0f4..85c9359795 100644
--- a/tests/python/tvmscript/test_tvmscript_ir_builder_irmodule.py
+++ b/tests/python/tvmscript/test_tvmscript_ir_builder_irmodule.py
@@ -22,8 +22,8 @@ import pytest
import tvm.testing
from tvm import ir
from tvm.ir.base import assert_structural_equal
+from tvm.script import ir_builder as I
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import ir as I
def test_ir_builder_irmodule():
diff --git a/tests/python/tvmscript/test_tvmscript_parser_evaluator.py
b/tests/python/tvmscript/test_tvmscript_parser_evaluator.py
index b198691ea0..684b191e4d 100644
--- a/tests/python/tvmscript/test_tvmscript_parser_evaluator.py
+++ b/tests/python/tvmscript/test_tvmscript_parser_evaluator.py
@@ -20,8 +20,8 @@
import pytest
import tvm.testing
+from tvm.script.ir_builder import tirx as T
from tvm.script.parser.entry import _run_statements
-from tvm.tirx.script import builder as T
def _calc(expr, extra_vars=None):
diff --git a/tests/python/tvmscript/test_tvmscript_printer_ir.py
b/tests/python/tvmscript/test_tvmscript_printer_ir.py
index 027317497c..a18bdb6a7b 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_ir.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_ir.py
@@ -19,8 +19,8 @@
import pytest
from tvm import IRModule
+from tvm.script import ir_builder as I
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import ir as I
from tvm.script.ir_builder import tirx as T
