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 78ade2f649 [REFACTOR][TVMScript] Restore dialect-owned IR builder
packages (#20440)
78ade2f649 is described below
commit 78ade2f649385213814e2f7e6ecd01b392d93716
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Sep 24 18:55:32 2026 -0400
[REFACTOR][TVMScript] Restore dialect-owned IR builder packages (#20440)
Move the TIRx and Relax builder implementations into
`tvm.tirx.script.ir_builder` and `tvm.relax.script.ir_builder`,
including Relax distributed helpers. Each dialect registers its
canonical builder path with the generic TVMScript namespace, so existing
`tvm.script.ir_builder.tirx` and `.relax` imports resolve to those same
modules.
Keep shared construction APIs in `tvm.script.ir_builder`, retain lazy
external dialect registration and deep import identity, and update
internal imports and documentation to the owning packages.
---
docs/arch/tvmscript.rst | 4 +-
docs/reference/api/python/script/ir_builder.rst | 15 +++--
docs/tirx/api/script.rst | 4 +-
docs/tirx/api/tile_dispatch.rst | 6 +-
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 | 4 +-
python/tvm/relax/script/__init__.py | 7 ++-
.../relax => relax/script/ir_builder}/__init__.py | 0
.../relax => relax/script/ir_builder}/_ffi_api.py | 2 +-
.../script/ir_builder}/distributed/__init__.py | 2 +-
.../script/ir_builder}/distributed/_ffi_api.py | 2 +-
.../script/ir_builder}/distributed/ir.py | 2 +-
.../relax => relax/script/ir_builder}/frame.py | 0
.../relax => relax/script/ir_builder}/ir.py | 0
.../script/ir_builder}/parser_protocol.py | 2 +-
python/tvm/relax/testing/attention.py | 2 +-
python/tvm/relax/testing/matmul.py | 2 +-
python/tvm/relax/transform/legalize_ops/grad.py | 4 +-
python/tvm/s_tir/tensor_intrin/arm_cpu.py | 2 +-
python/tvm/script/__init__.py | 65 +++++++++++-----------
python/tvm/script/ir_builder/__init__.py | 2 +-
python/tvm/script/ir_builder/parser_protocol.py | 4 +-
python/tvm/script/parser/entry.py | 6 +-
python/tvm/tirx/__init__.py | 6 +-
python/tvm/tirx/_buffer_view.py | 4 +-
python/tvm/tirx/script/__init__.py | 6 +-
.../tirx => tirx/script/ir_builder}/__init__.py | 0
.../tirx => tirx/script/ir_builder}/_ffi_api.py | 0
.../script/ir_builder}/external_kernel.py | 0
.../tirx => tirx/script/ir_builder}/frame.py | 0
.../tirx => tirx/script/ir_builder}/ir.py | 4 +-
.../script/ir_builder}/parser_protocol.py | 2 +-
.../tirx => tirx/script/ir_builder}/tirx.py | 0
.../tirx => tirx/script/ir_builder}/triton.py | 0
.../tirx => tirx/script/ir_builder}/utils.py | 2 +-
python/tvm/tirx/script/jit.py | 2 +-
python/tvm/tirx/script/tile.py | 2 +-
python/tvm/topi/gpu/scan.py | 2 +-
python/tvm/topi/gpu/scatter_elements.py | 2 +-
python/tvm/topi/gpu/scatter_nd.py | 2 +-
python/tvm/topi/gpu/sort.py | 2 +-
python/tvm/topi/index_put.py | 2 +-
python/tvm/topi/scan.py | 2 +-
python/tvm/topi/scatter.py | 2 +-
python/tvm/topi/scatter_elements.py | 2 +-
python/tvm/topi/searchsorted.py | 2 +-
python/tvm/topi/signal.py | 2 +-
python/tvm/topi/sparse_reshape.py | 2 +-
python/tvm/topi/unique.py | 2 +-
python/tvm/topi/vision/nms.py | 2 +-
python/tvm/topi/vision/nms_util.py | 2 +-
tests/python/codegen/test_target_codegen_vulkan.py | 2 +-
tests/python/relax/backend/adreno/mod_utils.py | 2 +-
tests/python/relax/backend/adreno/test_clml_ops.py | 2 +-
.../relax/backend/adreno/test_texture_network.py | 2 +-
tests/python/relax/test_codegen_cublas.py | 2 +-
tests/python/relax/test_codegen_cudnn.py | 2 +-
tests/python/relax/test_codegen_cutlass.py | 6 +-
...st_transform_attach_attr_layout_free_buffers.py | 2 +-
.../test_transform_combine_parallel_matmul.py | 2 +-
.../test_transform_fold_batch_norm_to_conv2d.py | 2 +-
tests/python/relax/test_tvmscript_ir_builder.py | 4 +-
tests/python/tirx/test_layout.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 +-
.../tvmscript/test_tvmscript_ir_builder_tir.py | 4 +-
.../tvmscript/test_tvmscript_parser_evaluator.py | 2 +-
.../python/tvmscript/test_tvmscript_printer_ir.py | 2 +-
.../python/tvmscript/test_tvmscript_printer_tir.py | 2 +-
77 files changed, 139 insertions(+), 128 deletions(-)
diff --git a/docs/arch/tvmscript.rst b/docs/arch/tvmscript.rst
index 48cba1dd33..2b1e1f9d6d 100644
--- a/docs/arch/tvmscript.rst
+++ b/docs/arch/tvmscript.rst
@@ -384,7 +384,7 @@ construction policy and finalization; the syntax transpiler
does not maintain an
TIR builder
~~~~~~~~~~~
-The TIR builder (``python/tvm/script/ir_builder/tirx/``) implements the
construction operations
+The TIR builder (``python/tvm/tirx/script/ir_builder/``) 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/script/ir_builder/relax/``) implements:
+The Relax builder (``python/tvm/relax/script/ir_builder/``) 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 e1f7f0c6ba..c89862a9b7 100644
--- a/docs/reference/api/python/script/ir_builder.rst
+++ b/docs/reference/api/python/script/ir_builder.rst
@@ -35,18 +35,21 @@ and
:func:`~tvm.script.ir_builder.parser_protocol.function_`.
.. automodule:: tvm.script.ir_builder.parser_protocol
:members:
-tvm.script.ir_builder.relax
+Dialect builders are owned by their respective script packages. Registered
+compatibility paths under ``tvm.script.ir_builder`` resolve to the same
modules.
+
+tvm.relax.script.ir_builder
***************************
-.. automodule:: tvm.script.ir_builder.relax
+.. automodule:: tvm.relax.script.ir_builder
:members:
:exclude-members: ExternFunc, ShapeExpr, TupleGetItem, Range
-tvm.script.ir_builder.relax.distributed
+tvm.relax.script.ir_builder.distributed
***************************************
-.. automodule:: tvm.script.ir_builder.relax.distributed
+.. automodule:: tvm.relax.script.ir_builder.distributed
:members:
-tvm.script.ir_builder.tirx
+tvm.tirx.script.ir_builder
**************************
-.. automodule:: tvm.script.ir_builder.tirx
+.. automodule:: tvm.tirx.script.ir_builder
:members:
diff --git a/docs/tirx/api/script.rst b/docs/tirx/api/script.rst
index dccc4e21e8..0744f09299 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.script.ir_builder.tirx.ir
+.. automodule:: tvm.tirx.script.ir_builder.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.script.ir_builder.tirx.ir
+.. currentmodule:: tvm.tirx.script.ir_builder.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 8e13d46f5d..d8ddf8d833 100644
--- a/docs/tirx/api/tile_dispatch.rst
+++ b/docs/tirx/api/tile_dispatch.rst
@@ -21,13 +21,13 @@ Tile Dispatch Extension API
Raw node constructors
---------------------
-``tvm.script.ir_builder.tirx.tirx`` constructs the same
+``tvm.tirx.script.ir_builder.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
-is ``tvm.script.ir_builder.tirx.tile``::
+is ``tvm.tirx.script.ir_builder.tile``::
- from tvm.script.ir_builder import tirx as Tx_builder
+ from tvm.tirx.script import ir_builder as Tx_builder
Tx_builder.tile.copy(dst, src)
diff --git a/docs/tirx/arch/tile_dispatch.rst b/docs/tirx/arch/tile_dispatch.rst
index fe99b0a947..a22be597bf 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/script/ir_builder/tirx/tirx.py``, while the validated authoring
facade
+``python/tvm/tirx/script/ir_builder/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 b14688ca23..d604f85d02 100644
--- a/python/tvm/backend/cuda/__init__.py
+++ b/python/tvm/backend/cuda/__init__.py
@@ -55,10 +55,10 @@ def _detect_target_from_device(dev):
def register_backend():
"""Register CUDA-owned Python semantics."""
- from tvm.script.ir_builder.tirx import (
+ from tvm.target.detect_target import register_device_target_detector
+ from tvm.tirx.script.ir_builder import (
ir as builder_ir, # pylint: disable=import-outside-toplevel
)
- from tvm.target.detect_target import register_device_target_detector
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 0120f71ffe..225166b80f 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.script.ir_builder.tirx import ir as _mod
+ from tvm.tirx.script.ir_builder 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 753a7bbe13..2f2fc2ea65 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.ir_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 4954c19500..25a8252cae 100644
--- a/python/tvm/backend/metal/__init__.py
+++ b/python/tvm/backend/metal/__init__.py
@@ -41,10 +41,10 @@ def _detect_target_from_device(dev):
def register_backend():
"""Register Metal-owned Python semantics."""
- from tvm.script.ir_builder.tirx import (
+ from tvm.target.detect_target import register_device_target_detector
+ from tvm.tirx.script.ir_builder import (
ir as builder_ir, # pylint: disable=import-outside-toplevel
)
- from tvm.target.detect_target import register_device_target_detector
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 1fa677bd39..2cd4de49d4 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.ir_builder.ir import _op_wrapper
class MetalNamespace:
diff --git a/python/tvm/backend/trn/__init__.py
b/python/tvm/backend/trn/__init__.py
index bc8e45dc8d..b45c411af1 100644
--- a/python/tvm/backend/trn/__init__.py
+++ b/python/tvm/backend/trn/__init__.py
@@ -31,10 +31,10 @@ _LAZY_SUBMODULES = {
def register_backend():
"""Register Trainium-owned Python semantics."""
- from tvm.script.ir_builder.tirx import (
+ from tvm.tirx import compilation_pipeline # pylint:
disable=import-outside-toplevel
+ from tvm.tirx.script.ir_builder import (
ir as builder_ir, # pylint: disable=import-outside-toplevel
)
- from tvm.tirx import compilation_pipeline # 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 4a6bd01398..701523f273 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.script.ir_builder.tirx.ir import _op_wrapper # pylint:
disable=import-outside-toplevel
+ from tvm.tirx.script.ir_builder.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 0b71d0a08d..c07375dbe8 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -112,12 +112,12 @@ from .binding_rewrite import DataflowBlockRewrite
import tvm.script
-tvm.script.register_dialect("relax", "tvm.relax.script")
+tvm.script.register_dialect("relax", "tvm.relax.script",
builder_path="tvm.relax.script.ir_builder")
def _check_script_module(module: "tvm.ir.IRModule") -> None:
# Delay builder imports until validation, keeping dialect/runtime
bootstrap safe.
- from tvm.script.ir_builder.relax.parser_protocol import
_check_module_well_formed
+ from tvm.relax.script.ir_builder.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 beff701106..b1bdfe93e8 100644
--- a/python/tvm/relax/script/__init__.py
+++ b/python/tvm/relax/script/__init__.py
@@ -45,10 +45,11 @@ 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 ir_builder as builder
+
globals().update(
(name, value) for name, value in vars(builder).items() if not
name.startswith("_")
)
@@ -69,12 +70,14 @@ def _initialize() -> None:
def __getattr__(name: str) -> _Any:
+ if name == "ir_builder":
+ return _importlib.import_module(__name__ + ".ir_builder")
if name.startswith("_") and name != "__all__":
raise AttributeError(name)
_initialize()
if name in globals():
return globals()[name]
- return getattr(_importlib.import_module("tvm.script.ir_builder.relax"),
name)
+ return getattr(_importlib.import_module("tvm.relax.script.ir_builder"),
name)
class _PyModuleFactory:
diff --git a/python/tvm/script/ir_builder/relax/__init__.py
b/python/tvm/relax/script/ir_builder/__init__.py
similarity index 100%
rename from python/tvm/script/ir_builder/relax/__init__.py
rename to python/tvm/relax/script/ir_builder/__init__.py
diff --git a/python/tvm/script/ir_builder/relax/_ffi_api.py
b/python/tvm/relax/script/ir_builder/_ffi_api.py
similarity index 95%
rename from python/tvm/script/ir_builder/relax/_ffi_api.py
rename to python/tvm/relax/script/ir_builder/_ffi_api.py
index ac4bd324af..8b14e9dd7b 100644
--- a/python/tvm/script/ir_builder/relax/_ffi_api.py
+++ b/python/tvm/relax/script/ir_builder/_ffi_api.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""FFI APIs for tvm.script.ir_builder.relax"""
+"""FFI APIs for tvm.relax.script.ir_builder"""
import tvm_ffi
diff --git a/python/tvm/script/ir_builder/relax/distributed/__init__.py
b/python/tvm/relax/script/ir_builder/distributed/__init__.py
similarity index 94%
rename from python/tvm/script/ir_builder/relax/distributed/__init__.py
rename to python/tvm/relax/script/ir_builder/distributed/__init__.py
index 9c81ce63b9..8b0bf3b40f 100644
--- a/python/tvm/script/ir_builder/relax/distributed/__init__.py
+++ b/python/tvm/relax/script/ir_builder/distributed/__init__.py
@@ -16,6 +16,6 @@
# specific language governing permissions and limitations
# under the License.
# pylint: disable=unused-import
-"""Package tvm.script.ir_builder.relax.distributed"""
+"""Package tvm.relax.script.ir_builder.distributed"""
from .ir import * # pylint: disable=wildcard-import,redefined-builtin
diff --git a/python/tvm/script/ir_builder/relax/distributed/_ffi_api.py
b/python/tvm/relax/script/ir_builder/distributed/_ffi_api.py
similarity index 93%
rename from python/tvm/script/ir_builder/relax/distributed/_ffi_api.py
rename to python/tvm/relax/script/ir_builder/distributed/_ffi_api.py
index 2ebc4935b8..e8df67e7df 100644
--- a/python/tvm/script/ir_builder/relax/distributed/_ffi_api.py
+++ b/python/tvm/relax/script/ir_builder/distributed/_ffi_api.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""FFI APIs for tvm.script.ir_builder.relax.distributed"""
+"""FFI APIs for tvm.relax.script.ir_builder.distributed"""
import tvm_ffi
diff --git a/python/tvm/script/ir_builder/relax/distributed/ir.py
b/python/tvm/relax/script/ir_builder/distributed/ir.py
similarity index 99%
rename from python/tvm/script/ir_builder/relax/distributed/ir.py
rename to python/tvm/relax/script/ir_builder/distributed/ir.py
index cb9e6391c4..f805af0cf4 100644
--- a/python/tvm/script/ir_builder/relax/distributed/ir.py
+++ b/python/tvm/relax/script/ir_builder/distributed/ir.py
@@ -40,11 +40,11 @@ from tvm.relax.op.distributed import (
from tvm.relax.op.distributed import (
redistribute as _redistribute,
)
+from tvm.relax.script.ir_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, 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/script/ir_builder/relax/frame.py
b/python/tvm/relax/script/ir_builder/frame.py
similarity index 100%
rename from python/tvm/script/ir_builder/relax/frame.py
rename to python/tvm/relax/script/ir_builder/frame.py
diff --git a/python/tvm/script/ir_builder/relax/ir.py
b/python/tvm/relax/script/ir_builder/ir.py
similarity index 100%
rename from python/tvm/script/ir_builder/relax/ir.py
rename to python/tvm/relax/script/ir_builder/ir.py
diff --git a/python/tvm/script/ir_builder/relax/parser_protocol.py
b/python/tvm/relax/script/ir_builder/parser_protocol.py
similarity index 99%
rename from python/tvm/script/ir_builder/relax/parser_protocol.py
rename to python/tvm/relax/script/ir_builder/parser_protocol.py
index c71fae3a19..539d9506de 100644
--- a/python/tvm/script/ir_builder/relax/parser_protocol.py
+++ b/python/tvm/relax/script/ir_builder/parser_protocol.py
@@ -32,7 +32,7 @@ from typing import Any, NoReturn
import tvm_ffi as _ffi
-import tvm.script.ir_builder.relax as _builder
+import tvm.relax.script.ir_builder as _builder
from tvm import ir as _ir
from tvm import relax as _relax
from tvm import tirx as _tir
diff --git a/python/tvm/relax/testing/attention.py
b/python/tvm/relax/testing/attention.py
index 3ea7d50228..33deda07d6 100644
--- a/python/tvm/relax/testing/attention.py
+++ b/python/tvm/relax/testing/attention.py
@@ -18,10 +18,10 @@
"""Relax script for attention module."""
import tvm
+from tvm.relax.script import ir_builder as relax_builder
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import relax as relax_builder
def get_relax_attention_module(
diff --git a/python/tvm/relax/testing/matmul.py
b/python/tvm/relax/testing/matmul.py
index 63986d84ca..2d2ab2d109 100644
--- a/python/tvm/relax/testing/matmul.py
+++ b/python/tvm/relax/testing/matmul.py
@@ -18,9 +18,9 @@
"""Utilities to construct matmul workloads."""
import tvm
+from tvm.relax.script import ir_builder as relax_builder
from tvm.script import relax as R
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import relax as relax_builder
def get_relax_matmul_module(
diff --git a/python/tvm/relax/transform/legalize_ops/grad.py
b/python/tvm/relax/transform/legalize_ops/grad.py
index 96e32ac2c4..a808e52cde 100644
--- a/python/tvm/relax/transform/legalize_ops/grad.py
+++ b/python/tvm/relax/transform/legalize_ops/grad.py
@@ -22,8 +22,8 @@ import logging
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.script.ir_builder.tirx.utils import buffer_proxy
+from tvm.tirx.script import ir_builder as T
+from tvm.tirx.script.ir_builder.utils import buffer_proxy
from ...block_builder import BlockBuilder
from ...expr import Expr
diff --git a/python/tvm/s_tir/tensor_intrin/arm_cpu.py
b/python/tvm/s_tir/tensor_intrin/arm_cpu.py
index b5fbe3a0ac..34472417f5 100644
--- a/python/tvm/s_tir/tensor_intrin/arm_cpu.py
+++ b/python/tvm/s_tir/tensor_intrin/arm_cpu.py
@@ -21,8 +21,8 @@
from tvm import tirx
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder.tirx import prim_func as build_prim_func
from tvm.target.codegen import llvm_version_major
+from tvm.tirx.script.ir_builder import prim_func as build_prim_func
from .. import TensorIntrin
from .dot_product_common import (
diff --git a/python/tvm/script/__init__.py b/python/tvm/script/__init__.py
index 6c846b2079..a4d9da860c 100644
--- a/python/tvm/script/__init__.py
+++ b/python/tvm/script/__init__.py
@@ -29,15 +29,15 @@ from collections.abc import Callable
from typing import Any
_DIALECT_REGISTRY: dict[str, str] = {}
+_DIALECT_BUILDER_REGISTRY: dict[str, str] = {}
# Language variants register whole-module checks at import, independently of
source decorators.
# Callbacks own concrete eligibility; this list retains no source or
construction state.
_MODULE_VALIDATORS: list[Callable[[Any], None]] = []
-# An empty suffix names the public language variant namespace itself.
-# Builder redirection applies only when no real package owns that name.
+# Public parser and builder compatibility paths consume dialect-owned
registrations.
_REDIRECTED_SUBPACKAGES = {
- "tvm.script.parser": "",
- "tvm.script.ir_builder": "builder",
+ "tvm.script.parser": _DIALECT_REGISTRY,
+ "tvm.script.ir_builder": _DIALECT_BUILDER_REGISTRY,
}
@@ -73,40 +73,44 @@ def register_module_validator(validator: Callable[[Any],
None], *, prepend: bool
_MODULE_VALIDATORS.append(validator)
-def register_dialect(name: str, module_path: str) -> None:
- """Register a language variant's script package path.
+def register_dialect(name: str, module_path: str, *, builder_path: str | None
= None) -> None:
+ """Register a dialect's script namespace and canonical builder package.
- 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 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``.
+ Registration is lazy: it records paths without importing either package.
+ ``tvm.script.<name>`` and ``tvm.script.parser.<name>`` expose
``module_path``;
+ ``tvm.script.ir_builder.<name>`` exposes ``builder_path``. Deep imports
under
+ these compatibility paths resolve to the same canonical modules, preserving
+ their identity and avoiding repeated initialization. Existing physical
+ out-of-tree packages under ``tvm.script.ir_builder`` retain normal package
+ lookup precedence.
- This function is idempotent — re-registering the same name with the same
- path is harmless.
-
- Each in-tree language variant calls this from its own ``__init__.py``::
+ Each dialect registers from its own package initialization::
import tvm.script
- tvm.script.register_dialect("tirx", "tvm.tirx.script")
+ tvm.script.register_dialect(
+ "example", "example.script",
builder_path="example.script.ir_builder"
+ )
- Out-of-tree language variants do the same in their own package init without
- editing any in-tree file.
+ Out-of-tree dialects register the same way without editing shared TVMScript
+ files. Re-registering the same name and paths is idempotent.
Parameters
----------
name : str
- The short name exposed under ``tvm.script.<name>`` (e.g. ``"tirx"``).
+ Short name exposed under ``tvm.script.<name>``.
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. Out-of-tree language variants provide imperative
- APIs in a ``builder`` submodule; in-tree builders live in
- ``tvm.script.ir_builder.<name>``.
+ Full dotted path of the dialect's script package, which owns its public
+ construction namespace.
+ builder_path : str, optional
+ Full dotted path of the dialect-owned imperative builder package,
+ normally ``module_path + ".ir_builder"``. When omitted, use
+ ``module_path + ".builder"`` to preserve existing out-of-tree
+ registrations. No imports or aliases are created at registration time.
"""
_DIALECT_REGISTRY[name] = module_path
+ _DIALECT_BUILDER_REGISTRY[name] = (
+ builder_path if builder_path is not None else module_path + ".builder"
+ )
def _redirect_target(fullname: str) -> str | None:
@@ -122,13 +126,13 @@ def _redirect_target(fullname: str) -> str | None:
target = _DIALECT_REGISTRY[head]
return f"{target}.{tail}" if tail else target
# Registered language variants may redirect parser and builder
subpackages.
- for prefix, sub in _REDIRECTED_SUBPACKAGES.items():
+ for prefix, registry in _REDIRECTED_SUBPACKAGES.items():
if fullname == prefix or not fullname.startswith(prefix + "."):
continue
rest = fullname[len(prefix) + 1 :]
head, _, tail = rest.partition(".")
- if head in _DIALECT_REGISTRY:
- target = _DIALECT_REGISTRY[head] + (f".{sub}" if sub else "")
+ if head in registry:
+ target = registry[head]
return f"{target}.{tail}" if tail else target
return None
@@ -156,8 +160,7 @@ 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.
+ # Preserve real external builder packages before the registered
fallback.
prefix = "tvm.script.ir_builder."
if fullname.startswith(prefix):
package_name = prefix + fullname[len(prefix) :].partition(".")[0]
diff --git a/python/tvm/script/ir_builder/__init__.py
b/python/tvm/script/ir_builder/__init__.py
index 5438e851e8..3906665490 100644
--- a/python/tvm/script/ir_builder/__init__.py
+++ b/python/tvm/script/ir_builder/__init__.py
@@ -75,7 +75,7 @@ __all__ = [
def __getattr__(name: str) -> _Any:
- # Real packages resolve directly; the finder supplies registered external
builders.
+ # The finder resolves each registered dialect's canonical builder package.
from tvm.script import _DIALECT_REGISTRY # pylint:
disable=import-outside-toplevel
if name in _DIALECT_REGISTRY:
diff --git a/python/tvm/script/ir_builder/parser_protocol.py
b/python/tvm/script/ir_builder/parser_protocol.py
index 4f97019d4f..06000c78f0 100644
--- a/python/tvm/script/ir_builder/parser_protocol.py
+++ b/python/tvm/script/ir_builder/parser_protocol.py
@@ -897,8 +897,8 @@ def function_(*, decl: bool = False, span: _Span = None,
**options: Any) -> IRBu
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.
+ body. See :func:`tvm.tirx.script.ir_builder.function_` and
+ :func:`tvm.relax.script.ir_builder.function_` for their concrete
options.
Returns
-------
diff --git a/python/tvm/script/parser/entry.py
b/python/tvm/script/parser/entry.py
index 39312c8d27..4060000453 100644
--- a/python/tvm/script/parser/entry.py
+++ b/python/tvm/script/parser/entry.py
@@ -251,15 +251,15 @@ def make_decorator(
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`.
+ :func:`tvm.tirx.script.ir_builder.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`.
+ :func:`tvm.tirx.script.ir_builder.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_`.
+ :func:`tvm.relax.script.ir_builder.function_`.
**options
Keyword options are forwarded to the selected language variant's
:func:`~tvm.script.ir_builder.parser_protocol.function_` hook,
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 38158bff41..a406ba8224 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -20,7 +20,7 @@
import tvm.script
-tvm.script.register_dialect("tirx", "tvm.tirx.script")
+tvm.script.register_dialect("tirx", "tvm.tirx.script",
builder_path="tvm.tirx.script.ir_builder")
from tvm.ir import Expr
@@ -117,12 +117,12 @@ if not _RUNTIME_ONLY_TIRX:
import tvm.script
-tvm.script.register_dialect("tirx", "tvm.tirx.script")
+tvm.script.register_dialect("tirx", "tvm.tirx.script",
builder_path="tvm.tirx.script.ir_builder")
def _check_script_module(module: "tvm.ir.IRModule") -> None:
# Delay builder imports until validation, keeping dialect/runtime
bootstrap safe.
- from tvm.script.ir_builder.tirx.parser_protocol import
_check_module_well_formed
+ from tvm.tirx.script.ir_builder.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 091ae9ea41..47d639947f 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.script.ir_builder.tirx.decl_buffer(
+ return tvm.tirx.script.ir_builder.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.script.ir_builder.tirx.decl_buffer(
+ return tvm.tirx.script.ir_builder.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 2ab9680e1d..9b32937e53 100644
--- a/python/tvm/tirx/script/__init__.py
+++ b/python/tvm/tirx/script/__init__.py
@@ -31,11 +31,11 @@ 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 ir_builder as builder
from . import tile
from .jit import make_jit
@@ -68,6 +68,8 @@ def _initialize() -> None:
def __getattr__(name: str) -> _Any:
+ if name == "ir_builder":
+ return _importlib.import_module(__name__ + ".ir_builder")
if name == "tile":
return _importlib.import_module(__name__ + ".tile")
if name.startswith("_") and name != "__all__":
@@ -75,7 +77,7 @@ def __getattr__(name: str) -> _Any:
_initialize()
if name in globals():
return globals()[name]
- return getattr(_importlib.import_module("tvm.script.ir_builder.tirx"),
name)
+ return getattr(_importlib.import_module("tvm.tirx.script.ir_builder"),
name)
from .jit import OptionalAnnotation as Optional
diff --git a/python/tvm/script/ir_builder/tirx/__init__.py
b/python/tvm/tirx/script/ir_builder/__init__.py
similarity index 100%
rename from python/tvm/script/ir_builder/tirx/__init__.py
rename to python/tvm/tirx/script/ir_builder/__init__.py
diff --git a/python/tvm/script/ir_builder/tirx/_ffi_api.py
b/python/tvm/tirx/script/ir_builder/_ffi_api.py
similarity index 100%
rename from python/tvm/script/ir_builder/tirx/_ffi_api.py
rename to python/tvm/tirx/script/ir_builder/_ffi_api.py
diff --git a/python/tvm/script/ir_builder/tirx/external_kernel.py
b/python/tvm/tirx/script/ir_builder/external_kernel.py
similarity index 100%
rename from python/tvm/script/ir_builder/tirx/external_kernel.py
rename to python/tvm/tirx/script/ir_builder/external_kernel.py
diff --git a/python/tvm/script/ir_builder/tirx/frame.py
b/python/tvm/tirx/script/ir_builder/frame.py
similarity index 100%
rename from python/tvm/script/ir_builder/tirx/frame.py
rename to python/tvm/tirx/script/ir_builder/frame.py
diff --git a/python/tvm/script/ir_builder/tirx/ir.py
b/python/tvm/tirx/script/ir_builder/ir.py
similarity index 99%
rename from python/tvm/script/ir_builder/tirx/ir.py
rename to python/tvm/tirx/script/ir_builder/ir.py
index 63a5e39606..1d35f2dea8 100644
--- a/python/tvm/script/ir_builder/tirx/ir.py
+++ b/python/tvm/tirx/script/ir_builder/ir.py
@@ -2312,7 +2312,7 @@ def launch_thread(
.. code-block:: python
- from tvm.script.ir_builder import tirx as T
+ from tvm.tirx.script import ir_builder as T
brow = T.env_thread("blockIdx.y")
T.launch_thread(brow, 1)
@@ -3112,7 +3112,7 @@ def register_script_namespace(name: str, namespace:
object) -> object:
import sys # pylint: disable=import-outside-toplevel
for module_name in [
- "tvm.script.ir_builder.tirx",
+ "tvm.tirx.script.ir_builder",
"tvm.tirx.script",
"tvm.script.tirx",
]:
diff --git a/python/tvm/script/ir_builder/tirx/parser_protocol.py
b/python/tvm/tirx/script/ir_builder/parser_protocol.py
similarity index 99%
rename from python/tvm/script/ir_builder/tirx/parser_protocol.py
rename to python/tvm/tirx/script/ir_builder/parser_protocol.py
index 32fa88528e..16ba16145b 100644
--- a/python/tvm/script/ir_builder/tirx/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -29,7 +29,7 @@ from collections.abc import Sequence
from functools import partial as _partial
from typing import Any
-import tvm.script.ir_builder.tirx as _builder
+import tvm.tirx.script.ir_builder as _builder
from tvm import ir as _ir
from tvm import tirx as _tir
from tvm.ir.prim import _ffi_api as _prim_ffi
diff --git a/python/tvm/script/ir_builder/tirx/tirx.py
b/python/tvm/tirx/script/ir_builder/tirx.py
similarity index 100%
rename from python/tvm/script/ir_builder/tirx/tirx.py
rename to python/tvm/tirx/script/ir_builder/tirx.py
diff --git a/python/tvm/script/ir_builder/tirx/triton.py
b/python/tvm/tirx/script/ir_builder/triton.py
similarity index 100%
rename from python/tvm/script/ir_builder/tirx/triton.py
rename to python/tvm/tirx/script/ir_builder/triton.py
diff --git a/python/tvm/script/ir_builder/tirx/utils.py
b/python/tvm/tirx/script/ir_builder/utils.py
similarity index 99%
rename from python/tvm/script/ir_builder/tirx/utils.py
rename to python/tvm/tirx/script/ir_builder/utils.py
index 7f72c46008..5edf425b95 100644
--- a/python/tvm/script/ir_builder/tirx/utils.py
+++ b/python/tvm/tirx/script/ir_builder/utils.py
@@ -213,7 +213,7 @@ def buffer_proxy(buf: Buffer) -> _BufferProxy:
--------
.. code-block:: python
- from tvm.script.ir_builder.tirx.utils import buffer_proxy
+ from tvm.tirx.script.ir_builder.utils import buffer_proxy
buf = tvm.tirx.decl_buffer([2, 3], "float32")
ptr = buffer_proxy(buf)
diff --git a/python/tvm/tirx/script/jit.py b/python/tvm/tirx/script/jit.py
index b93a5b9bb1..c636d8d92d 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.script.ir_builder import tirx as builder
+ from tvm.tirx.script import ir_builder 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 a9d4921f03..bad6774f99 100644
--- a/python/tvm/tirx/script/tile.py
+++ b/python/tvm/tirx/script/tile.py
@@ -19,8 +19,8 @@
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 tvm.tirx.script.ir_builder import tirx as _builder
def _get_arg(args, kwargs, index, name):
diff --git a/python/tvm/topi/gpu/scan.py b/python/tvm/topi/gpu/scan.py
index e576e7d3e8..5e69a613d2 100644
--- a/python/tvm/topi/gpu/scan.py
+++ b/python/tvm/topi/gpu/scan.py
@@ -24,7 +24,7 @@ import tvm
from tvm import te
from tvm.contrib.thrust import can_use_rocthrust, can_use_thrust
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script import ir_builder as T
from ..math import cast, ceil_log2
from ..transform import expand_dims, reshape, squeeze, transpose
diff --git a/python/tvm/topi/gpu/scatter_elements.py
b/python/tvm/topi/gpu/scatter_elements.py
index 5049670e35..d778a4f080 100644
--- a/python/tvm/topi/gpu/scatter_elements.py
+++ b/python/tvm/topi/gpu/scatter_elements.py
@@ -20,7 +20,7 @@
import tvm
from tvm import te, tirx
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script import ir_builder as T
from .. import utils
from ..math import cast
diff --git a/python/tvm/topi/gpu/scatter_nd.py
b/python/tvm/topi/gpu/scatter_nd.py
index 6f90477fc5..ed0015174c 100644
--- a/python/tvm/topi/gpu/scatter_nd.py
+++ b/python/tvm/topi/gpu/scatter_nd.py
@@ -21,7 +21,7 @@
import tvm
from tvm import te, tirx # hide redefinition of min and max
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script import ir_builder as T
from ..math import cast
from ..scatter import _verify_scatter_nd_inputs
diff --git a/python/tvm/topi/gpu/sort.py b/python/tvm/topi/gpu/sort.py
index 9b9871e6c3..5d38addf92 100644
--- a/python/tvm/topi/gpu/sort.py
+++ b/python/tvm/topi/gpu/sort.py
@@ -21,7 +21,7 @@
import tvm
from tvm import te
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script import ir_builder as T
from ..math import cast, ceil_log2
from ..searchsorted import binary_search
diff --git a/python/tvm/topi/index_put.py b/python/tvm/topi/index_put.py
index 08ba1fbecc..f2ee79a905 100644
--- a/python/tvm/topi/index_put.py
+++ b/python/tvm/topi/index_put.py
@@ -19,7 +19,7 @@
from tvm import te, tirx
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script import ir_builder as T
from . import utils
diff --git a/python/tvm/topi/scan.py b/python/tvm/topi/scan.py
index 591eb9b9cd..57bb003777 100644
--- a/python/tvm/topi/scan.py
+++ b/python/tvm/topi/scan.py
@@ -22,7 +22,7 @@ from collections.abc import Callable
import tvm
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script import ir_builder as T
from ..te import extern
from ..tirx import decl_buffer
diff --git a/python/tvm/topi/scatter.py b/python/tvm/topi/scatter.py
index 05c3dba215..1bfef89b5c 100644
--- a/python/tvm/topi/scatter.py
+++ b/python/tvm/topi/scatter.py
@@ -20,8 +20,8 @@
from tvm import DataTypeCode, ir, te, tirx # hide redefinition of min and max
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
from tvm.sym.analyzer import Analyzer
+from tvm.tirx.script import ir_builder as T
def _verify_scatter_nd_inputs(data, indices, updates):
diff --git a/python/tvm/topi/scatter_elements.py
b/python/tvm/topi/scatter_elements.py
index f1b28fed07..ebe37bfb01 100644
--- a/python/tvm/topi/scatter_elements.py
+++ b/python/tvm/topi/scatter_elements.py
@@ -18,7 +18,7 @@
from tvm import te, tirx
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script import ir_builder as T
from . import utils
from .math import cast
diff --git a/python/tvm/topi/searchsorted.py b/python/tvm/topi/searchsorted.py
index 36db4676f9..7adf407870 100644
--- a/python/tvm/topi/searchsorted.py
+++ b/python/tvm/topi/searchsorted.py
@@ -18,7 +18,7 @@
"""searchsorted operator"""
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script import ir_builder as T
from . import te, utils
from .math import cast
diff --git a/python/tvm/topi/signal.py b/python/tvm/topi/signal.py
index c3c38854b7..41c488d263 100644
--- a/python/tvm/topi/signal.py
+++ b/python/tvm/topi/signal.py
@@ -21,7 +21,7 @@ from math import pi
from tvm import ir, te, tirx
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script import ir_builder as T
def stft(
diff --git a/python/tvm/topi/sparse_reshape.py
b/python/tvm/topi/sparse_reshape.py
index 2017ced6cb..4676fbd6da 100644
--- a/python/tvm/topi/sparse_reshape.py
+++ b/python/tvm/topi/sparse_reshape.py
@@ -18,9 +18,9 @@
"""Sparse_Reshape operator"""
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
from tvm.te import div, extern, floordiv, floormod
from tvm.tirx import Cast, decl_buffer
+from tvm.tirx.script import ir_builder as T
def sparse_reshape(
diff --git a/python/tvm/topi/unique.py b/python/tvm/topi/unique.py
index 6135f6c8f6..d07aca847e 100644
--- a/python/tvm/topi/unique.py
+++ b/python/tvm/topi/unique.py
@@ -19,7 +19,7 @@
from tvm import te, tirx
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script import ir_builder as T
def _calc_adjacent_diff_ir(data, output, binop=tirx.Sub):
diff --git a/python/tvm/topi/vision/nms.py b/python/tvm/topi/vision/nms.py
index a82056f541..4953d0a099 100644
--- a/python/tvm/topi/vision/nms.py
+++ b/python/tvm/topi/vision/nms.py
@@ -20,8 +20,8 @@
import tvm
from tvm import te
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
from tvm.tirx import if_then_else
+from tvm.tirx.script import ir_builder as T
from .. import reduction
from ..math import cast
diff --git a/python/tvm/topi/vision/nms_util.py
b/python/tvm/topi/vision/nms_util.py
index a55bedb697..7a13279f14 100644
--- a/python/tvm/topi/vision/nms_util.py
+++ b/python/tvm/topi/vision/nms_util.py
@@ -21,7 +21,7 @@
import tvm
from tvm import te
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script import ir_builder as T
def _get_boundaries(output, box_idx):
diff --git a/tests/python/codegen/test_target_codegen_vulkan.py
b/tests/python/codegen/test_target_codegen_vulkan.py
index 9ac20f0236..597ccd084e 100644
--- a/tests/python/codegen/test_target_codegen_vulkan.py
+++ b/tests/python/codegen/test_target_codegen_vulkan.py
@@ -27,8 +27,8 @@ 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 tirx as T_builder
from tvm.testing import env
+from tvm.tirx.script import ir_builder as T_builder
dtype = tvm.testing.parameter("float32", "int32", "float16", "int8")
fuzz_seed = tvm.testing.parameter(range(25))
diff --git a/tests/python/relax/backend/adreno/mod_utils.py
b/tests/python/relax/backend/adreno/mod_utils.py
index ab08434c29..9bc518bc90 100644
--- a/tests/python/relax/backend/adreno/mod_utils.py
+++ b/tests/python/relax/backend/adreno/mod_utils.py
@@ -26,11 +26,11 @@ import tvm
import tvm.testing
from tvm import relax, rpc
from tvm.relax.backend.adreno import clml
+from tvm.relax.script import ir_builder as relax_builder
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import relax as relax_builder
def get_relax_conv2d_mod(
diff --git a/tests/python/relax/backend/adreno/test_clml_ops.py
b/tests/python/relax/backend/adreno/test_clml_ops.py
index 428f21abdf..df8d34adab 100644
--- a/tests/python/relax/backend/adreno/test_clml_ops.py
+++ b/tests/python/relax/backend/adreno/test_clml_ops.py
@@ -51,11 +51,11 @@ import tvm.testing
from tvm import relax, rpc
from tvm.relax.backend.adreno import clml
from tvm.relax.backend.adreno.clml import OpenCLMLOffLoad,
OpenCLMLOffLoadForLLM
+from tvm.relax.script import ir_builder as relax_builder
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import relax as relax_builder
CLML_VERSION = clml.clml_sdk_version()
TARGET_CLML_VERSION = int(os.environ.get("ADRENO_TARGET_CLML_VERSION", 4))
diff --git a/tests/python/relax/backend/adreno/test_texture_network.py
b/tests/python/relax/backend/adreno/test_texture_network.py
index 0b7bcfcd74..456415960b 100644
--- a/tests/python/relax/backend/adreno/test_texture_network.py
+++ b/tests/python/relax/backend/adreno/test_texture_network.py
@@ -31,12 +31,12 @@ import tvm
import tvm.testing
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx
+from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import relax as relax_builder
TARGETS = [tvm.target.Target("qcom/adreno-opencl-texture")]
diff --git a/tests/python/relax/test_codegen_cublas.py
b/tests/python/relax/test_codegen_cublas.py
index f273ebf8ef..2e5f431a2f 100644
--- a/tests/python/relax/test_codegen_cublas.py
+++ b/tests/python/relax/test_codegen_cublas.py
@@ -27,10 +27,10 @@ pytest.importorskip("scipy") # tvm.topi.testing imports
scipy
import tvm.topi.testing
from tvm import relax
from tvm.relax.backend.cuda.cublas import partition_for_cublas
+from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.testing import get_relax_matmul_module
from tvm.script import relax as R
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import relax as relax_builder
try:
import ml_dtypes
diff --git a/tests/python/relax/test_codegen_cudnn.py
b/tests/python/relax/test_codegen_cudnn.py
index b934cf9967..950d9cb7a4 100644
--- a/tests/python/relax/test_codegen_cudnn.py
+++ b/tests/python/relax/test_codegen_cudnn.py
@@ -28,10 +28,10 @@ import tvm.topi.testing
from tvm import relax
from tvm.contrib.pickle_memoize import memoize
from tvm.relax.backend.cuda.cudnn import partition_for_cudnn
+from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.testing import get_relax_stacked_attention_module
from tvm.script import relax as R
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import relax as relax_builder
@pytest.fixture(autouse=True)
diff --git a/tests/python/relax/test_codegen_cutlass.py
b/tests/python/relax/test_codegen_cutlass.py
index b26c15fc0c..ff4c764974 100644
--- a/tests/python/relax/test_codegen_cutlass.py
+++ b/tests/python/relax/test_codegen_cutlass.py
@@ -29,6 +29,7 @@ from tvm import relax
from tvm.contrib.cutlass.build import is_shape_valid_for_cutlass_matmul
from tvm.contrib.pickle_memoize import memoize
from tvm.relax.backend.cuda.cutlass import partition_for_cutlass
+from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.testing import (
get_relax_attention_module,
get_relax_matmul_module,
@@ -38,7 +39,6 @@ from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import relax as relax_builder
@pytest.fixture(autouse=True)
@@ -891,9 +891,9 @@ def attention_rewrite_size(request):
def get_relax_attention_rewrite_module(
q_shape, k_shape, v_shape, out_shape, dtype, bias_shape=None, scale=None
):
+ from tvm.relax.script import ir_builder as relax_builder
from tvm.script.ir_builder import IRBuilder
- from tvm.script.ir_builder import relax as relax_builder
- from tvm.script.ir_builder import tirx as T
+ from tvm.tirx.script import ir_builder as T
with IRBuilder() as builder:
with relax_builder.function():
diff --git
a/tests/python/relax/test_transform_attach_attr_layout_free_buffers.py
b/tests/python/relax/test_transform_attach_attr_layout_free_buffers.py
index c21b9a6475..65c428c7a2 100644
--- a/tests/python/relax/test_transform_attach_attr_layout_free_buffers.py
+++ b/tests/python/relax/test_transform_attach_attr_layout_free_buffers.py
@@ -19,12 +19,12 @@ import numpy as np
import tvm.testing
from tvm import relax, tirx
+from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.transform import CombineParallelMatmul
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import relax as relax_builder
def test_param():
diff --git a/tests/python/relax/test_transform_combine_parallel_matmul.py
b/tests/python/relax/test_transform_combine_parallel_matmul.py
index 57be71ae88..0269d7b070 100644
--- a/tests/python/relax/test_transform_combine_parallel_matmul.py
+++ b/tests/python/relax/test_transform_combine_parallel_matmul.py
@@ -19,11 +19,11 @@ import pytest
import tvm.testing
from tvm import relax, tirx
+from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.transform import CombineParallelMatmul
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import relax as relax_builder
def get_parallel_matmul(
diff --git a/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py
b/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py
index be2dcace3b..dbdd653fef 100644
--- a/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py
+++ b/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py
@@ -24,10 +24,10 @@ import tvm.testing
from tvm import relax
from tvm.ir.module import IRModule
from tvm.relax.expr_functor import PyExprVisitor, visitor
+from tvm.relax.script import ir_builder as relax_builder
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import relax as relax_builder
def get_conv2d_batchnorm_sample():
diff --git a/tests/python/relax/test_tvmscript_ir_builder.py
b/tests/python/relax/test_tvmscript_ir_builder.py
index 995e260f47..8f06282ca0 100644
--- a/tests/python/relax/test_tvmscript_ir_builder.py
+++ b/tests/python/relax/test_tvmscript_ir_builder.py
@@ -18,7 +18,7 @@
import tvm
import tvm.testing
from tvm import relax, tirx, topi
-from tvm.script.ir_builder import relax as R
+from tvm.relax.script import ir_builder as R
from tvm.script.ir_builder.base import IRBuilder
@@ -165,7 +165,7 @@ def test_dataflow_block():
def test_regression_py_print():
# Test that the py_print directs to python builtin print
- from tvm.script.ir_builder.relax.ir import py_print # pylint:
disable=import-outside-toplevel
+ from tvm.relax.script.ir_builder.ir import py_print # pylint:
disable=import-outside-toplevel
assert py_print == print
diff --git a/tests/python/tirx/test_layout.py b/tests/python/tirx/test_layout.py
index d7f20f4a1c..b562535760 100644
--- a/tests/python/tirx/test_layout.py
+++ b/tests/python/tirx/test_layout.py
@@ -26,7 +26,6 @@ from tvm.ir import assert_structural_equal
from tvm.ir.type import PointerType, PrimType
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as Tx_builder
from tvm.sym import Analyzer
from tvm.tirx import Var
from tvm.tirx.cuda.tile_primitive.tma_utils import (
@@ -54,6 +53,7 @@ from tvm.tirx.layout import (
wgid,
wid_in_wg,
)
+from tvm.tirx.script import ir_builder as Tx_builder
def test_axis():
diff --git a/tests/python/tirx/test_op.py b/tests/python/tirx/test_op.py
index bc728d7863..b9bfead5f9 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.script.ir_builder.tirx.tirx import gemm_async
+ from tvm.tirx.script.ir_builder.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 6f23681a3c..1811a59a40 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.script.ir_builder.tirx import ir as builder_ir
+ from tvm.tirx.script.ir_builder 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 ir_builder as builder
+ from tvm.tirx.script.ir_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 d507cad6bc..b42be2e233 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.script.ir_builder.tirx.buffer_store
+ original = tvm.tirx.script.ir_builder.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.script.ir_builder.tirx.buffer_store", side_effect=bomb):
+ with patch("tvm.tirx.script.ir_builder.buffer_store", side_effect=bomb):
with pytest.raises(ValueError, match="boom"):
from_source(src)
diff --git a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py
b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py
index 66631d4529..567398e0a8 100644
--- a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py
+++ b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py
@@ -16,7 +16,7 @@
# under the License.
# pylint: disable=invalid-name, missing-docstring
# ruff: noqa: F401, F841
-"""Unittests for tvm.script.ir_builder.tirx"""
+"""Unittests for tvm.tirx.script.ir_builder"""
import numpy as np
import pytest
@@ -27,7 +27,7 @@ import tvm.testing
from tvm import s_tir, tirx
from tvm.ir.base import SourceName, Span, assert_structural_equal
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script import ir_builder as T
def test_ir_builder_tir_primfunc_base():
diff --git a/tests/python/tvmscript/test_tvmscript_parser_evaluator.py
b/tests/python/tvmscript/test_tvmscript_parser_evaluator.py
index 684b191e4d..7c600e5641 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 ir_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 a18bdb6a7b..1b0b9b34b0 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_ir.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_ir.py
@@ -21,7 +21,7 @@ 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 tirx as T
+from tvm.tirx.script import ir_builder as T
def _assert_print(obj, expected):
diff --git a/tests/python/tvmscript/test_tvmscript_printer_tir.py
b/tests/python/tvmscript/test_tvmscript_printer_tir.py
index 61e6920a73..9c096fc968 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_tir.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_tir.py
@@ -25,7 +25,7 @@ import tvm.testing
from tvm import ir, s_tir, tirx
from tvm.ir import Range
from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script import ir_builder as T
def _assert_print(obj, expected):
