tvm.script

tvm.script

TVMScript public namespace.

Dialect resolution mechanism

tvm.script is a virtual namespace: dialect names like tirx and relax are not bound as static attributes here. Instead:

  • register_dialect(name, module_path) writes an entry to _DIALECT_REGISTRY: dict[str, str]. Each in-tree dialect’s __init__.py calls this on import (e.g., tvm.tirx.__init__.py calls tvm.script.register_dialect("tirx", "tvm.tirx.script")). Out-of-tree dialects can register themselves the same way.

  • __getattr__(name) (PEP 562) fires on missing attribute access. If name is in _DIALECT_REGISTRY, the listed module is imported and cached as a normal module attribute. Subsequent accesses skip __getattr__ (cached in globals()).

  • Subpackages tvm.script.parser, tvm.script.ir_builder, etc. each define their own __getattr__ that consults the SAME _DIALECT_REGISTRY and appends their suffix. So tvm.script.parser.tirx resolves to tvm.tirx.script.parser via the dialect registry + .parser suffix.

  • For deep statement-form imports like from tvm.script.parser.tirx.entry import ObjectProxy, PEP 562’s __getattr__ is not enough — it only handles one-level from X import Y. A sys.meta_path finder (see _DialectRedirectFinder) intercepts the import machinery to register the real module under the legacy name in sys.modules, so subsequent attribute walks resolve correctly.

Each dialect’s tvm.<dialect>.script package MUST expose parser, ir_builder, and (where applicable) printer as submodules. This convention is what makes the suffix-append redirect work uniformly. IR is foundational (script depends on ir) and is NOT a dialect; its script handlers live in the shared core, not via this registry.

Bootstrap order

python/tvm/__init__.py imports tvm.script BEFORE importing any dialect package (tvm.tirx, tvm.relax, …). This guarantees that tvm.script.register_dialect is reachable the moment a dialect’s own __init__.py runs and calls it. The tvm.script module itself stays dialect-agnostic at load time (no dialect submodules are eagerly imported here), so there is no circular dependency.

tvm.script.register_dialect(name: str, module_path: str) → None

Register a dialect’s script package path.

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 module_path + ".parser" / module_path + ".builder" etc. via each subpackage’s own __getattr__. Deep statement-form imports (e.g., from tvm.script.parser.<name>.entry import X) are handled by _DialectRedirectFinder on sys.meta_path.

This function is idempotent — re-registering the same name with the same path is harmless.

Each in-tree dialect calls this from its own __init__.py:

import tvm.script
tvm.script.register_dialect("tirx", "tvm.tirx.script")

Out-of-tree dialects do the same in their own package init without editing any in-tree file.

Parameters:

  • name (str) – The short name exposed under tvm.script.<name> (e.g. "tirx").

  • module_path (str) – The full dotted module path of the dialect’s script package, e.g. "tvm.tirx.script". That package must expose parser and ir_builder as submodules (and printer where applicable) so that the suffix-append redirect works uniformly.

tvm.script.ir_module(mod: type | None = None, check_well_formed: bool = True) → IRModule

The parsing method for ir module, by using @ir_module as decorator.

Parameters:

  • mod (Type) – The class to be parsed as ir module.

  • check_well_formed (bool) – Whether to check well-formedness during parsing.

Returns:

ir_module – The parsed ir module.

Return type:

IRModule

tvm.script.relax

Relax-layer TVMScript pieces (parser, builder).

After the per-dialect TVMScript restructure, the Relax layer owns its own script/{parser,builder} subpackages. tvm.script.relax resolves to this module via the dialect registry, so the public parser surface (function, Tensor, match_cast, etc.) is re-exported here.

tvm.script.tirx

TIRX-layer TVMScript pieces (parser, builder).

After the per-dialect TVMScript restructure, the TIRX layer owns its own script/{parser,builder} subpackages. tvm.script.tirx resolves to this module via the dialect registry, so the public parser surface (prim_func, Buffer, Ptr, etc.) is re-exported here.