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 f7c5b680f3 [REFACTOR][Python] Align namespace helpers with their
implementation owners (#20599)
f7c5b680f3 is described below
commit f7c5b680f33207d25dc473ac9f4588fbec1210ce
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Oct 8 14:36:01 2026 -0400
[REFACTOR][Python] Align namespace helpers with their implementation owners
(#20599)
Align primitive constants, script display, device conversion and scalar
extraction with their IR, Script, Target and Sym owners. Keep reducer
construction in the S-TIR operator module and expose it through its
script builder; remove redundant public Python hoisting configuration
wrappers and flags while retaining native pass options. Let the TIRx
source-kernel adapter assemble its CUDA options using existing
lower-layer path APIs.
---
docs/reference/api/python/script/printer.rst | 1 +
python/tvm/__init__.py | 3 +-
python/tvm/backend/cuda/op.py | 3 +-
python/tvm/backend/cuda/ptx/engine.py | 3 +-
python/tvm/ir/__init__.py | 2 +
.../{runtime/object_generic.py => ir/_constant.py} | 10 ++--
python/tvm/ir/_overload_prim_expr.py | 3 +-
python/tvm/ir/base.py | 24 +++------
python/tvm/ir/expr.py | 5 +-
python/tvm/ir/module.py | 11 ++--
python/tvm/ir/prim/__init__.py | 2 +-
python/tvm/ir/type.py | 4 +-
python/tvm/relax/expr.py | 2 +-
python/tvm/runtime/__init__.py | 8 +--
python/tvm/runtime/_ffi_node_api.py | 44 ----------------
python/tvm/runtime/_tensor.py | 24 ---------
python/tvm/runtime/object_generic.py | 47 -----------------
python/tvm/s_tir/op.py | 54 +++++++++++++++++++
python/tvm/s_tir/script/__init__.py | 3 ++
python/tvm/s_tir/script/ir_builder/ir.py | 5 ++
python/tvm/s_tir/stmt.py | 4 +-
python/tvm/s_tir/transform/__init__.py | 1 -
python/tvm/script/__init__.py | 4 ++
python/tvm/script/printer/__init__.py | 2 +
python/tvm/script/printer/doc_printer.py | 2 +-
.../printer/scriptable.py} | 61 +++++++++-------------
.../{s_tir/transform/__init__.py => sym/utils.py} | 34 ++++++++++--
python/tvm/target/__init__.py | 2 +-
python/tvm/target/target.py | 23 +++++++-
python/tvm/te/reduction.py | 4 +-
python/tvm/te/tensor.py | 4 +-
python/tvm/tirx/__init__.py | 2 +-
python/tvm/tirx/expr.py | 2 +-
python/tvm/tirx/function.py | 4 +-
python/tvm/tirx/op/__init__.py | 2 +-
.../tvm/tirx/script/ir_builder/external_kernel.py | 23 ++++----
python/tvm/tirx/script/ir_builder/op.py | 37 -------------
python/tvm/tirx/script/ir_builder/triton.py | 2 +-
python/tvm/tirx/stmt.py | 15 +++++-
python/tvm/tirx/tile_primitive.py | 4 +-
python/tvm/tirx/transform/transform.py | 60 ---------------------
python/tvm/topi/utils.py | 24 +--------
src/ir/prim/op.cc | 2 +-
src/relax/ir/dependent_type.cc | 2 +-
src/script/printer/config.cc | 9 ++--
src/te/script/printer/expr.cc | 4 +-
.../relax/script/test_relax_script_printer.py | 2 +-
.../test_s_tir_transform_hoist_expression.py | 52 +++++++++++++++++-
tests/python/script/test_script_printer_entry.py | 2 +-
tests/python/tirx-base/test_tir_index_map.py | 3 +-
tests/python/tirx/codegen/test_ptx_addr.py | 3 +-
.../python/tirx/script/test_tirx_script_printer.py | 8 +--
...test_tirx_script_printer_invisible_path_info.py | 2 +-
53 files changed, 289 insertions(+), 374 deletions(-)
diff --git a/docs/reference/api/python/script/printer.rst
b/docs/reference/api/python/script/printer.rst
index cc16c55c3b..098333de07 100644
--- a/docs/reference/api/python/script/printer.rst
+++ b/docs/reference/api/python/script/printer.rst
@@ -23,3 +23,4 @@ tvm.script.printer
.. automodule:: tvm.script.printer
:members:
:imported-members:
+ :exclude-members: Scriptable
diff --git a/python/tvm/__init__.py b/python/tvm/__init__.py
index 388e2180f3..02bc9e56cd 100644
--- a/python/tvm/__init__.py
+++ b/python/tvm/__init__.py
@@ -31,7 +31,7 @@ from .base import _RUNTIME_ONLY
# tvm.runtime
from .runtime import Object
-from .runtime._tensor import device, device_from_target, cpu, cuda, opencl,
vulkan, metal
+from .runtime._tensor import device, cpu, cuda, opencl, vulkan, metal
from .runtime._tensor import vpi, rocm, ext_dev, hexagon
from .runtime import DataType, DataTypeCode
@@ -57,6 +57,7 @@ from . import backend
# tvm.target
from . import target
+from .target import device_from_target
# tvm.te
from . import te
diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py
index c60f788576..4090f623ca 100644
--- a/python/tvm/backend/cuda/op.py
+++ b/python/tvm/backend/cuda/op.py
@@ -24,10 +24,9 @@ from enum import Enum
import tvm_ffi
from tvm import tirx
-from tvm.ir import Attrs, Call, Op, StringImm
+from tvm.ir import Attrs, Call, Op, StringImm, const
from tvm.ir.op import _init_op_api, _make_op_api
from tvm.ir.type import PointerType, PrimType
-from tvm.runtime import const
from tvm.tirx.op import bitwise_and, call_intrin, tvm_access_ptr
from tvm.tirx.operator.intrinsics._common import (
CP_ASYNC_BULK_CACHE_HINT as _CP_ASYNC_BULK_CACHE_HINT,
diff --git a/python/tvm/backend/cuda/ptx/engine.py
b/python/tvm/backend/cuda/ptx/engine.py
index 473b5bd20b..f9f15859f6 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -42,10 +42,9 @@ from tvm_ffi import get_global_func, register_global_func,
structural_equal
from tvm.backend.cuda.codegen.registry import register_codegen
from tvm.backend.cuda.codegen.utils import parse_str
from tvm.backend.cuda.op import cuda_cvta_generic_to_shared, cuda_func_call
-from tvm.ir import Call, Op, StringImm, TensorLoad
+from tvm.ir import Call, Op, StringImm, TensorLoad, const
from tvm.ir.op import register_op_attr
from tvm.ir.type import PointerType, PrimType
-from tvm.runtime import const
from tvm.tirx.expr import CallEffectKind, IntImm
from tvm.tirx.op import call_intrin, reinterpret
diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py
index 345b294468..6aa8a4ae7a 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -19,10 +19,12 @@
"""Common data structures across all IR variants."""
from . import instrument, transform
+from ._constant import const
from .attrs import Attrs, DictAttrs, make_node
from .base import (
EnvFunc,
Node,
+ Scriptable,
SourceName,
Span,
SequentialSpan,
diff --git a/python/tvm/runtime/object_generic.py b/python/tvm/ir/_constant.py
similarity index 89%
copy from python/tvm/runtime/object_generic.py
copy to python/tvm/ir/_constant.py
index c8cbdf596d..b07413b9a3 100644
--- a/python/tvm/runtime/object_generic.py
+++ b/python/tvm/ir/_constant.py
@@ -14,13 +14,9 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: F401
-"""Common implementation of object generic related logic"""
+"""Construction of primitive IR constants."""
-# pylint: disable=unused-import, invalid-name
-from tvm_ffi import ObjectConvertible
-
-from . import _ffi_node_api
+from . import _ffi_api
def _scalar_type_inference(value):
@@ -65,4 +61,4 @@ def const(value, dtype=None, span=None):
"""
if dtype is None:
dtype = _scalar_type_inference(value)
- return _ffi_node_api._const(value, dtype, span)
+ return _ffi_api._const(value, dtype, span)
diff --git a/python/tvm/ir/_overload_prim_expr.py
b/python/tvm/ir/_overload_prim_expr.py
index 26eaf1332a..12e187df63 100644
--- a/python/tvm/ir/_overload_prim_expr.py
+++ b/python/tvm/ir/_overload_prim_expr.py
@@ -16,7 +16,8 @@
# under the License.
"""Primitive-expression overloads for shared IR expressions."""
-from ..runtime import DataTypeCode, ObjectConvertible, const
+from ..runtime import DataTypeCode, ObjectConvertible
+from ._constant import const
from .base import Span
from .expr import Expr, TensorRegion, is_prim_expr
from .prim import _ffi_api as _prim_ffi_api
diff --git a/python/tvm/ir/base.py b/python/tvm/ir/base.py
index f572b02e37..3cf24e4961 100644
--- a/python/tvm/ir/base.py
+++ b/python/tvm/ir/base.py
@@ -26,19 +26,12 @@ from ..libinfo import __version__
from . import _ffi_api, json_compact
-class Node(Object):
- """Base class of all IR Nodes."""
+class Scriptable:
+ """Marker for display methods installed by the TVMScript owner."""
- def __repr__(self) -> str:
- from tvm.runtime.script_printer import PrinterConfig, _script
- try:
- return _script(
- self,
- PrinterConfig(extra_config={"ir.comment_imports": True}),
- )
- except Exception:
- return super().__repr__()
+class Node(Object):
+ """Base class of all IR Nodes."""
@register_object("ir.SourceMap")
@@ -190,16 +183,13 @@ def assert_structural_equal(lhs, rhs,
map_free_vars=False):
"""
first_mismatch = tvm_ffi.get_first_structural_mismatch(lhs, rhs,
map_free_vars)
if first_mismatch is not None:
- from tvm.runtime.script_printer import ( # pylint:
disable=import-outside-toplevel
- Scriptable,
- )
-
lhs_path, rhs_path = first_mismatch
# Diagnostics use the same display policy as Object.script(), including
# dialect selection and commented imports. The unbound method also
# accepts IR objects that do not inherit the convenience mixin.
- lhs_script = Scriptable.script(lhs, path_to_underline=[lhs_path])
- rhs_script = Scriptable.script(rhs, path_to_underline=[rhs_path])
+ script = getattr(Scriptable, "script", None)
+ lhs_script = script(lhs, path_to_underline=[lhs_path]) if script else
repr(lhs)
+ rhs_script = script(rhs, path_to_underline=[rhs_path]) if script else
repr(rhs)
raise ValueError(
f"StructuralEqual check failed, caused by lhs at {lhs_path}:\n"
f"{lhs_script}\n"
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 3d8f368c2d..e2da21fb75 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -23,9 +23,10 @@ import tvm_ffi
import tvm
-from ..runtime import Object, Scriptable, const
+from ..runtime import Object
from . import _ffi_api, _tensor_expr_overload
-from .base import Node, Span
+from ._constant import const
+from .base import Node, Scriptable, Span
def _convert_subscript_index(index):
diff --git a/python/tvm/ir/module.py b/python/tvm/ir/module.py
index fd5ebc59aa..9fec84e03d 100644
--- a/python/tvm/ir/module.py
+++ b/python/tvm/ir/module.py
@@ -18,13 +18,12 @@
import tvm_ffi
-import tvm
-from tvm.runtime import Object, Scriptable
+from tvm.runtime import Object
from . import _ffi_api
from . import expr as _expr
-from .attrs import DictAttrs
-from .base import Node
+from .attrs import DictAttrs, make_node
+from .base import Node, Scriptable
from .function import BaseFunc
@@ -55,7 +54,7 @@ class IRModule(Node, Scriptable):
attrs = None if not attrs else attrs
if attrs is not None:
- attrs = tvm.ir.make_node("ir.DictAttrs", **attrs)
+ attrs = make_node("ir.DictAttrs", **attrs)
if global_infos is None:
global_infos = {}
self.__init_handle_by_constructor__(
@@ -278,7 +277,7 @@ class IRModule(Node, Scriptable):
mod : IRModule
A new copy of the IRModule with the attribute
"""
- if isinstance(attr_map, tvm.ir.DictAttrs):
+ if isinstance(attr_map, DictAttrs):
attr_map = attr_map._dict()
return _ffi_api.Module_WithAttrs(self, attr_map)
diff --git a/python/tvm/ir/prim/__init__.py b/python/tvm/ir/prim/__init__.py
index 9cadc68f81..0a6687bb35 100644
--- a/python/tvm/ir/prim/__init__.py
+++ b/python/tvm/ir/prim/__init__.py
@@ -17,7 +17,7 @@
# pylint: disable=unused-import
"""Primitive expressions and construction helpers shared by TVM IR dialects."""
-from ...runtime import const
+from .._constant import const
from ..expr import Expr
from . import _ffi_api
from .expr import (
diff --git a/python/tvm/ir/type.py b/python/tvm/ir/type.py
index dfaf28623e..2d67c0ade2 100644
--- a/python/tvm/ir/type.py
+++ b/python/tvm/ir/type.py
@@ -18,10 +18,8 @@
import tvm_ffi
-from tvm.runtime import Scriptable
-
from . import _ffi_api
-from .base import Node
+from .base import Node, Scriptable
@tvm_ffi.register_object("ir.Type")
diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py
index 245053ef59..0795d06a49 100644
--- a/python/tvm/relax/expr.py
+++ b/python/tvm/relax/expr.py
@@ -30,10 +30,10 @@ import tvm.ir
import tvm.relax
import tvm.runtime
from tvm import DataType
+from tvm.ir import Scriptable
from ..ir import BaseFunc, Node, Span
from ..ir.expr import _CallableExprWithOp
-from ..runtime import Scriptable
from . import _ffi_api
# It is a workaround for mypy:
https://github.com/python/mypy/issues/7866#issuecomment-549454370
diff --git a/python/tvm/runtime/__init__.py b/python/tvm/runtime/__init__.py
index 56b904a519..7b950012e4 100644
--- a/python/tvm/runtime/__init__.py
+++ b/python/tvm/runtime/__init__.py
@@ -20,12 +20,7 @@
from tvm_ffi import convert, Object
from tvm_ffi._dtype import dtype as DataType, DataTypeCode
-# Import _ffi_node_api for its side effect of installing AsRepr as
-# tvm_ffi.core.__object_repr__.
-from . import _ffi_node_api
-
# class exposures
-from .script_printer import Scriptable
from .object_generic import ObjectConvertible
from .device import Device
from ._tensor import Tensor, tensor, empty
@@ -33,10 +28,9 @@ from .module import Module
from .executable import Executable
# function exposures
-from ._tensor import device, device_from_target, cpu, cuda, opencl, vulkan,
metal
+from ._tensor import device, cpu, cuda, opencl, vulkan, metal
from ._tensor import vpi, rocm, ext_dev, from_dlpack
from .module import load_module, enabled, system_lib, load_static_library,
num_threads
-from .object_generic import const
from .params import (
save_param_dict,
load_param_dict,
diff --git a/python/tvm/runtime/_ffi_node_api.py
b/python/tvm/runtime/_ffi_node_api.py
deleted file mode 100644
index 1c87b989b6..0000000000
--- a/python/tvm/runtime/_ffi_node_api.py
+++ /dev/null
@@ -1,44 +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.
-
-# pylint: disable=invalid-name, unused-argument
-"""FFI for tvm.node"""
-
-import tvm_ffi
-import tvm_ffi.core
-
-
-# The implementations below are default ones when the corresponding
-# functions are not available in the runtime only mode.
-# They will be overriden via tvm_ffi.init_ffi_api to the ones registered
-def AsRepr(obj):
- return type(obj).__name__ + "(" + obj.__ctypes_handle__().value + ")"
-
-
-def SaveJSON(obj):
- raise RuntimeError("Do not support object serialization in runtime only
mode")
-
-
-def LoadJSON(json_str):
- raise RuntimeError("Do not support object serialization in runtime only
mode")
-
-
-# Exports functions registered in node namespace.
-tvm_ffi.init_ffi_api("node", __name__)
-
-# Override the default repr function for tvm_ffi.core.Object.
-tvm_ffi.core.__object_repr__ = AsRepr
diff --git a/python/tvm/runtime/_tensor.py b/python/tvm/runtime/_tensor.py
index b53d88250e..58c5e60969 100644
--- a/python/tvm/runtime/_tensor.py
+++ b/python/tvm/runtime/_tensor.py
@@ -32,7 +32,6 @@ except ImportError:
import tvm_ffi
from tvm_ffi import DLDeviceType, device
-import tvm
from tvm.runtime import Device
from . import _ffi_api
@@ -359,29 +358,6 @@ def tensor(arr, device=None, mem_scope=None):
return empty(arr.shape, arr.dtype, device, mem_scope).copyfrom(arr)
-def device_from_target(target, index=None):
- """Construct a runtime device from a compilation target.
-
- Parameters
- ----------
- target : str or dict or tvm.target.Target
- The compilation target whose device type should be used.
-
- index : int, optional
- The integer device index.
-
- Returns
- -------
- dev : Device
- The created device.
- """
- from tvm.target import Target # pylint: disable=import-outside-toplevel
-
- if not isinstance(target, Target):
- target = Target(target)
- return device(target.get_target_device_type(), index)
-
-
def cpu(dev_id=0):
"""Construct a CPU device
diff --git a/python/tvm/runtime/object_generic.py
b/python/tvm/runtime/object_generic.py
index c8cbdf596d..14f54dc82d 100644
--- a/python/tvm/runtime/object_generic.py
+++ b/python/tvm/runtime/object_generic.py
@@ -19,50 +19,3 @@
# pylint: disable=unused-import, invalid-name
from tvm_ffi import ObjectConvertible
-
-from . import _ffi_node_api
-
-
-def _scalar_type_inference(value):
- if hasattr(value, "dtype"):
- return str(value.dtype)
- elif isinstance(value, bool):
- return "bool"
- elif isinstance(value, float):
- # We intentionally prefer convert the float to float32 since it's more
common in DL.
- if -3.40282347e38 <= value <= 3.40282347e38:
- return "float32"
- else:
- return "float64"
- elif isinstance(value, int):
- # We intentionally prefer convert the python int to int32 since it's
more common in DL.
- if -2147483648 <= value <= 2147483647:
- return "int32"
- else:
- return "int64"
- else:
- raise NotImplementedError(f"Cannot automatically inference the type.
value={value}")
-
-
-def const(value, dtype=None, span=None):
- """construct a constant
-
- Parameters
- ----------
- value : number
- The content of the constant number.
-
- dtype : str or None, optional
- The data type.
-
- span : Optional[Span]
- The location of the constant value in the source.
-
- Returns
- -------
- const_val: tvm.Expr
- The result expression.
- """
- if dtype is None:
- dtype = _scalar_type_inference(value)
- return _ffi_node_api._const(value, dtype, span)
diff --git a/python/tvm/s_tir/op.py b/python/tvm/s_tir/op.py
new file mode 100644
index 0000000000..2760558d65
--- /dev/null
+++ b/python/tvm/s_tir/op.py
@@ -0,0 +1,54 @@
+# 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.
+"""S-TIR operator construction helpers."""
+
+import inspect
+from collections.abc import Callable
+
+from tvm.te import CommReducer
+from tvm.tirx import Expr, Var
+
+
+def comm_reducer(combiner: Callable, identity: list[Expr]) -> CommReducer:
+ """
+ Create a CommReducer from lambda inputs/outputs and the identities
+
+ Parameters
+ ----------
+ combiner : Callable
+ A binary function which takes two Expr as input to return a Expr.
+
+ identity : List[Expr]
+ A list of types of output Expr.
+
+ Returns
+ -------
+ res : CommReducer
+ The CommReducer.
+ """
+ params = inspect.signature(combiner).parameters
+ num_args = len(params)
+ args = []
+ for name, i in zip(params.keys(), identity + identity):
+ if isinstance(i, int):
+ args.append(Var(name, "int32"))
+ else:
+ args.append(Var(name, i.ty))
+ res = combiner(*args)
+ if not isinstance(res, tuple):
+ res = (res,)
+ return CommReducer(args[: num_args // 2], args[num_args // 2 :], res,
identity)
diff --git a/python/tvm/s_tir/script/__init__.py
b/python/tvm/s_tir/script/__init__.py
index c7d43fdf76..dd0bc97707 100644
--- a/python/tvm/s_tir/script/__init__.py
+++ b/python/tvm/s_tir/script/__init__.py
@@ -41,6 +41,9 @@ def _initialize():
if name not in ("builder", "ir_builder", "jit")
)
for name in (
+ "CommReducer",
+ "Reduce",
+ "comm_reducer",
"async_copy_scope",
"async_commit",
"async_wait",
diff --git a/python/tvm/s_tir/script/ir_builder/ir.py
b/python/tvm/s_tir/script/ir_builder/ir.py
index f3e701d209..8c3a39d3a0 100644
--- a/python/tvm/s_tir/script/ir_builder/ir.py
+++ b/python/tvm/s_tir/script/ir_builder/ir.py
@@ -25,8 +25,10 @@ from tvm import ir
from tvm import tirx as tir
from tvm.ir import TensorLoad, TensorRegion, is_prim_expr
from tvm.s_tir.iter_var import IterVar
+from tvm.s_tir.op import comm_reducer
from tvm.script.ir_builder.base import MISSING
from tvm.script.parser.protocol_registry import register_mutable_decl as
_register_mutable_decl
+from tvm.te import CommReducer, Reduce
from tvm.tirx import Expr, IntImm, Var
from tvm.tirx.layout import Layout
from tvm.tirx.script.ir_builder.ir import _get_layout, _record_meta_resource
@@ -476,12 +478,15 @@ class axis: # pylint: disable=invalid-name
__all__ = [
+ "CommReducer",
"IterVar",
+ "Reduce",
"async_commit",
"async_copy_scope",
"async_wait",
"axis",
"block_name_suffix_context",
+ "comm_reducer",
"iter_var",
"manual_sync",
"match_buffer",
diff --git a/python/tvm/s_tir/stmt.py b/python/tvm/s_tir/stmt.py
index 9d61e76988..8ff0317d5f 100644
--- a/python/tvm/s_tir/stmt.py
+++ b/python/tvm/s_tir/stmt.py
@@ -20,8 +20,8 @@ from collections.abc import Mapping, Sequence
import tvm_ffi
-from tvm.ir import Expr, Span, TensorRegion, Var
-from tvm.runtime import Object, Scriptable, const
+from tvm.ir import Expr, Scriptable, Span, TensorRegion, Var, const
+from tvm.runtime import Object
from tvm.s_tir import IterVar
from tvm.tirx.stmt import SeqStmt, Stmt
diff --git a/python/tvm/s_tir/transform/__init__.py
b/python/tvm/s_tir/transform/__init__.py
index 19bc7281f4..1a195b3fcb 100644
--- a/python/tvm/s_tir/transform/__init__.py
+++ b/python/tvm/s_tir/transform/__init__.py
@@ -19,4 +19,3 @@
# pylint: disable=wildcard-import, invalid-name
from .transform import *
-from ...tirx.transform.transform import HoistedConditionals, HoistedLetBindings
diff --git a/python/tvm/script/__init__.py b/python/tvm/script/__init__.py
index a4d9da860c..cfab7ab503 100644
--- a/python/tvm/script/__init__.py
+++ b/python/tvm/script/__init__.py
@@ -248,3 +248,7 @@ def __getattr__(name: str) -> Any:
from .parser import register_namespace as _register_namespace
_register_namespace("script", sys.modules[__name__])
+
+
+# Install the Script-owned display methods after the shared IR types exist.
+from . import printer
diff --git a/python/tvm/script/printer/__init__.py
b/python/tvm/script/printer/__init__.py
index 8d2f73bb2b..77af75e0f0 100644
--- a/python/tvm/script/printer/__init__.py
+++ b/python/tvm/script/printer/__init__.py
@@ -19,3 +19,5 @@ TVMScript Unified Printer
This package provides a set of APIs to print supported TVM IR into TVMScript
in a roundtrippable way.
"""
+
+from .scriptable import PrinterConfig, Scriptable
diff --git a/python/tvm/script/printer/doc_printer.py
b/python/tvm/script/printer/doc_printer.py
index 42216a2c18..76e18d26f1 100644
--- a/python/tvm/script/printer/doc_printer.py
+++ b/python/tvm/script/printer/doc_printer.py
@@ -18,7 +18,7 @@
from tvm_ffi.access_path import AccessPath
-from tvm.runtime.script_printer import PrinterConfig
+from tvm.script.printer.scriptable import PrinterConfig
from . import _ffi_api
from .doc import Doc
diff --git a/python/tvm/runtime/script_printer.py
b/python/tvm/script/printer/scriptable.py
similarity index 88%
rename from python/tvm/runtime/script_printer.py
rename to python/tvm/script/printer/scriptable.py
index d50fe142dd..0199ca4c8d 100644
--- a/python/tvm/runtime/script_printer.py
+++ b/python/tvm/script/printer/scriptable.py
@@ -19,12 +19,13 @@
import os
from collections.abc import Sequence
-from tvm_ffi import get_global_func, register_object
+from tvm_ffi import register_object
from tvm_ffi.access_path import AccessPath
+from tvm.ir.base import Node, Scriptable
from tvm.runtime import Object
-from . import _ffi_node_api
+from . import _ffi_api
@register_object("script.PrinterConfig")
@@ -102,21 +103,20 @@ class PrinterConfig(Object):
if extra_config is not None:
cfg["extra_config"] = extra_config
self.__init_handle_by_constructor__(
- _ffi_node_api.PrinterConfig,
+ _ffi_api.PrinterConfig,
cfg, # type: ignore # pylint: disable=no-member
)
def _script(obj: Object, config: PrinterConfig) -> str:
- return _ffi_node_api.TVMScriptPrinterScript(obj, config) # type: ignore #
pylint: disable=no-member
+ return _ffi_api.Script(obj, config) # type: ignore # pylint:
disable=no-member
def _relax_script(obj: Object, config: PrinterConfig) -> str:
- func = get_global_func("script.printer.ReprPrintRelax")
- return func(obj, config)
+ return _ffi_api.ReprPrintRelax(obj, config)
-class Scriptable:
+class _ScriptableMethods:
"""A base class that enables the script() and show() method."""
def script(
@@ -186,10 +186,6 @@ class Scriptable:
The TVM Script of the given TVM IR
"""
- # Auto-switch to tirx (`T`/`tirx`) flavor only when explicitly
- # printing a Function / IRModule that has no s_tir-tagged content.
- # Free objects (Buffer, buffer-backed TensorRegion, ...) keep the
default `T`/`tir`
- # flavor -- they have no enclosing function to indicate tirx vs s_tir.
merged_extra: dict = {}
if extra_config is not None:
merged_extra.update(extra_config)
@@ -197,32 +193,6 @@ class Scriptable:
# retain executable imports for standalone parser round-trips.
merged_extra.setdefault("ir.comment_imports", True)
- # Only auto-switch if the caller has not already set a tirx.prefix
override.
- if "tirx.prefix" not in merged_extra:
- from tvm.ir import IRModule # pylint:
disable=import-outside-toplevel
- from tvm.tirx import Function # pylint:
disable=import-outside-toplevel
-
- switch_to_tirx = False
- if isinstance(self, Function):
- attrs = getattr(self, "attrs", None)
- if attrs is None or not attrs.get("s_tir", False):
- switch_to_tirx = True
- elif isinstance(self, IRModule):
- any_prim = False
- any_s_tir = False
- for _, base_func in self.functions.items():
- if isinstance(base_func, Function):
- any_prim = True
- if getattr(base_func, "attrs", None) and
base_func.attrs.get(
- "s_tir", False
- ):
- any_s_tir = True
- break
- if any_prim and not any_s_tir:
- switch_to_tirx = True
- if switch_to_tirx:
- merged_extra["tirx.prefix"] = "T"
-
return _script(
self,
PrinterConfig(
@@ -399,3 +369,20 @@ class Scriptable:
style=style,
black_format=black_format,
)
+
+
+def _repr(self) -> str:
+ try:
+ return _script(
+ self,
+ PrinterConfig(extra_config={"ir.comment_imports": True}),
+ )
+ except Exception:
+ return Object.__repr__(self)
+
+
+# The printer supplies its explicit display capability to existing IR wrappers.
+Scriptable.script = _ScriptableMethods.script
+Scriptable.show = _ScriptableMethods.show
+Scriptable._relax_script = _ScriptableMethods._relax_script
+Node.__repr__ = _repr
diff --git a/python/tvm/s_tir/transform/__init__.py b/python/tvm/sym/utils.py
similarity index 53%
copy from python/tvm/s_tir/transform/__init__.py
copy to python/tvm/sym/utils.py
index 19bc7281f4..02d4fbc833 100644
--- a/python/tvm/s_tir/transform/__init__.py
+++ b/python/tvm/sym/utils.py
@@ -1,4 +1,3 @@
-# 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
@@ -15,8 +14,33 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Namespace of all S-TIR transformations"""
-# pylint: disable=wildcard-import, invalid-name
+"""Utilities for extracting values from symbolic expressions."""
-from .transform import *
-from ...tirx.transform.transform import HoistedConditionals, HoistedLetBindings
+from numbers import Integral
+
+from tvm.ir.prim import IntImm
+
+from .analyzer import Analyzer
+
+
+def get_const_int(expr):
+ """Verifies expr is integer and get the constant value.
+
+ Parameters
+ ----------
+ expr : tvm.Expr or int
+ The input expression.
+
+ Returns
+ -------
+ out_value : int
+ The output.
+ """
+ if isinstance(expr, Integral):
+ return expr
+ if not isinstance(expr, IntImm):
+ ana = Analyzer()
+ expr = ana.simplify(expr)
+ if not isinstance(expr, IntImm):
+ raise ValueError("Expect value to be constant int")
+ return int(expr.value)
diff --git a/python/tvm/target/__init__.py b/python/tvm/target/__init__.py
index 20e26bccf3..efe4ddacc2 100644
--- a/python/tvm/target/__init__.py
+++ b/python/tvm/target/__init__.py
@@ -31,7 +31,7 @@ Use :py:func:`tvm.target.list_tags` to list all available
target tags,
and :py:func:`tvm.target.register_tag` to register new tags.
"""
-from .target import Target, TargetKind
+from .target import Target, TargetKind, device_from_target
from .tag import list_tags, register_tag
from . import codegen
from . import tag_registry # registers tags on import
diff --git a/python/tvm/target/target.py b/python/tvm/target/target.py
index 0ecfd1d34c..80bf739064 100644
--- a/python/tvm/target/target.py
+++ b/python/tvm/target/target.py
@@ -20,7 +20,7 @@ import tvm_ffi
from tvm_ffi import Map
from tvm_ffi.core import String
-from tvm.runtime import Device, Object, convert
+from tvm.runtime import Device, Object, convert, device
from . import _ffi_api
@@ -227,3 +227,24 @@ class Target(Object):
if target is None:
raise ValueError("Target is not set in env or passed as argument.")
return target
+
+
+def device_from_target(target, index=None):
+ """Construct a runtime device from a compilation target.
+
+ Parameters
+ ----------
+ target : str or dict or tvm.target.Target
+ The compilation target whose device type should be used.
+
+ index : int, optional
+ The integer device index.
+
+ Returns
+ -------
+ dev : Device
+ The created device.
+ """
+ if not isinstance(target, Target):
+ target = Target(target)
+ return device(target.get_target_device_type(), index)
diff --git a/python/tvm/te/reduction.py b/python/tvm/te/reduction.py
index 12a47c670d..687336b639 100644
--- a/python/tvm/te/reduction.py
+++ b/python/tvm/te/reduction.py
@@ -25,11 +25,11 @@ from tvm_ffi import Array
import tvm
from tvm import ir
from tvm import tirx as tir
-from tvm.ir import Expr, Var
+from tvm.ir import Expr, Scriptable, Var, const
from tvm.ir.base import Span
from tvm.ir.prim import _ffi_api as _prim_ffi_api
from tvm.ir.prim import max_value, min_value
-from tvm.runtime import Object, Scriptable, const
+from tvm.runtime import Object
from tvm.tirx.op import _primexpr_dtype
from . import _ffi_api
diff --git a/python/tvm/te/tensor.py b/python/tvm/te/tensor.py
index 34c42a3e07..cbd9e90e8e 100644
--- a/python/tvm/te/tensor.py
+++ b/python/tvm/te/tensor.py
@@ -19,8 +19,8 @@
# pylint: disable=invalid-name
import tvm_ffi
-from tvm.ir import OpaqueExpr, is_prim_expr
-from tvm.runtime import Object, ObjectConvertible, const
+from tvm.ir import OpaqueExpr, const, is_prim_expr
+from tvm.runtime import Object, ObjectConvertible
from tvm.tirx import expr as _expr
from tvm.tirx.stmt import SeqStmt
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index f69f0e9886..6014781758 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -24,7 +24,7 @@ tvm.script.register_dialect("tirx", "tvm.tirx.script",
builder_path="tvm.tirx.sc
from tvm.ir import Expr
-from tvm.runtime import const
+from tvm.ir import const
from .buffer import (
BufferAccessKind,
diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py
index f1c9b28616..efb24e4250 100644
--- a/python/tvm/tirx/expr.py
+++ b/python/tvm/tirx/expr.py
@@ -30,6 +30,7 @@ For example, you can use addexp.a to get the left operand of
an Add node.
from tvm import ir
from tvm.ir import Expr
+from tvm.ir import const as const
from tvm.ir._overload_prim_expr import ( # noqa: F401
EqualOp,
ExprOp,
@@ -78,7 +79,6 @@ from tvm.ir.prim.expr import ( # noqa: F401
Sub,
)
from tvm.runtime import ObjectConvertible
-from tvm.runtime import const as const
from . import _ffi_api
diff --git a/python/tvm/tirx/function.py b/python/tvm/tirx/function.py
index d5db3c9a3b..5761a85fcc 100644
--- a/python/tvm/tirx/function.py
+++ b/python/tvm/tirx/function.py
@@ -26,8 +26,8 @@ import tvm_ffi
import tvm
import tvm.runtime
-from tvm.ir import BaseFunc, Range, Var
-from tvm.runtime import Object, Scriptable
+from tvm.ir import BaseFunc, Range, Scriptable, Var
+from tvm.runtime import Object
from ..runtime._tensor import Tensor
from . import _ffi_api
diff --git a/python/tvm/tirx/op/__init__.py b/python/tvm/tirx/op/__init__.py
index f31082fdb3..264247fabf 100644
--- a/python/tvm/tirx/op/__init__.py
+++ b/python/tvm/tirx/op/__init__.py
@@ -37,6 +37,7 @@ from tvm.ir import (
TensorLoad,
TensorRegion,
Var,
+ const,
)
from tvm.ir.base import Span
from tvm.ir.prim import clz as clz
@@ -65,7 +66,6 @@ from tvm.ir.prim.op import shift_right as shift_right
from tvm.ir.prim.op import truncdiv as truncdiv
from tvm.ir.prim.op import truncmod as truncmod
from tvm.ir.prim.op import vscale as vscale
-from tvm.runtime import const
from .. import _ffi_api
from ..buffer import buffer_data, is_tensor_var
diff --git a/python/tvm/tirx/script/ir_builder/external_kernel.py
b/python/tvm/tirx/script/ir_builder/external_kernel.py
index fe07b346dc..3b4abbe7a9 100644
--- a/python/tvm/tirx/script/ir_builder/external_kernel.py
+++ b/python/tvm/tirx/script/ir_builder/external_kernel.py
@@ -27,9 +27,9 @@ from typing import Any
import tvm_ffi
from tvm import __version__ as tvm_version
-from tvm import tirx
-from tvm.ir import Expr, PointerType, is_prim_expr
-from tvm.runtime import Module, const
+from tvm import libinfo, tirx
+from tvm.ir import Expr, PointerType, const, is_prim_expr
+from tvm.runtime import Module
from tvm.support import nvcc
@@ -121,10 +121,6 @@ class SourceKernel(BaseKernel): # pylint:
disable=too-few-public-methods
**kwargs: dict[str, Any],
) -> tuple[str, Module, list[Any]]:
"""Compile the kernel to a device module."""
- from tvm.relax.frontend.nn import ( # pylint:
disable=import-outside-toplevel
- SourceModule,
- )
-
kernel_name = kwargs["kernel_name"]
assert len(grid) == 2, (
"grid should be two list of integers, representing the dimension
of "
@@ -147,8 +143,17 @@ class SourceKernel(BaseKernel): # pylint:
disable=too-few-public-methods
kernel_arg_types.append(str(arg.ty.dtype))
runtime_args = runtime_args + list(grid[0]) + list(grid[1])
- # Reuse compilation path from SourceModule
- compile_options = SourceModule.get_compile_options("cu")
+ include_paths = [
+ Path(libinfo.find_include_path()),
+ Path(tvm_ffi.libinfo.find_include_path()),
+ Path(tvm_ffi.libinfo.find_dlpack_include_path()),
+ ]
+ compile_options = []
+ for include_path in dict.fromkeys(include_paths):
+ assert include_path.exists(), f"Not found: {include_path!s}"
+ assert include_path.is_dir(), f"Not a directory: {include_path!s}"
+ compile_options += ["-I", str(include_path)]
+ compile_options += ["-c", "-O3", "-std=c++17", "-Xcompiler=-fPIC"]
source_code = self.source_code
try:
source_path = Path(source_code)
diff --git a/python/tvm/tirx/script/ir_builder/op.py
b/python/tvm/tirx/script/ir_builder/op.py
index 423aa1849c..3116446031 100644
--- a/python/tvm/tirx/script/ir_builder/op.py
+++ b/python/tvm/tirx/script/ir_builder/op.py
@@ -19,8 +19,6 @@
from __future__ import annotations
import functools
-import inspect
-from collections.abc import Callable
from typing import Any
import tvm_ffi as _ffi
@@ -40,7 +38,6 @@ from tvm.script.ir_builder.frame import IRModuleFrame
# pylint: disable=unused-import
from tvm.target.codegen import llvm_lookup_intrinsic_id
-from tvm.te import CommReducer, Reduce
from tvm.tirx import Expr, is_tensor_var
from tvm.tirx import op as _tir_op
from tvm.tirx.exec_scope import Var
@@ -154,37 +151,6 @@ def max(a: Expr, b: Expr) -> Expr: # pylint:
disable=redefined-builtin
return _ffi_api.max(a, b) # type: ignore[attr-defined] # pylint:
disable=no-member
-def comm_reducer(combiner: Callable, identity: list[Expr]) -> CommReducer:
- """
- Create a CommReducer from lambda inputs/outputs and the identities
-
- Parameters
- ----------
- combiner : Callable
- A binary function which takes two Expr as input to return a Expr.
-
- identity : List[Expr]
- A list of types of output Expr.
-
- Returns
- -------
- res : CommReducer
- The CommReducer.
- """
- params = inspect.signature(combiner).parameters
- num_args = len(params)
- args = []
- for name, i in zip(params.keys(), identity + identity):
- if isinstance(i, int):
- args.append(Var(name, "int32"))
- else:
- args.append(Var(name, i.ty))
- res = combiner(*args)
- if not isinstance(res, tuple):
- res = (res,)
- return CommReducer(args[: num_args // 2], args[num_args // 2 :], res,
identity)
-
-
def _llvm_result_type(func):
"""Keep symbolic intrinsic-name conversion when spelling the result as
ty."""
@@ -806,7 +772,6 @@ __all__ = [
"CallEffectKind",
"CallFFIKernelAttr",
"Cast",
- "CommReducer",
"Div",
"FloorDiv",
"FloorMod",
@@ -820,7 +785,6 @@ __all__ = [
"Or",
"RShift",
"Ramp",
- "Reduce",
"Select",
"Shuffle",
"Sub",
@@ -858,7 +822,6 @@ __all__ = [
"ceil",
"ceildiv",
"clz",
- "comm_reducer",
"cooperative_tensor_fill",
"cooperative_tensor_load",
"cooperative_tensor_multiply_accumulate",
diff --git a/python/tvm/tirx/script/ir_builder/triton.py
b/python/tvm/tirx/script/ir_builder/triton.py
index f10c8d27dc..f18e2dca0b 100644
--- a/python/tvm/tirx/script/ir_builder/triton.py
+++ b/python/tvm/tirx/script/ir_builder/triton.py
@@ -26,7 +26,7 @@ from triton.runtime.jit import type_canonicalisation_dict
from tvm import tirx
from tvm.ir import PointerType, PrimType, is_prim_expr
from tvm.runtime import Module
-from tvm.topi.utils import get_const_int
+from tvm.sym.utils import get_const_int
from .external_kernel import BaseKernel
diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py
index 9026d76411..1cdef7cbec 100644
--- a/python/tvm/tirx/stmt.py
+++ b/python/tvm/tirx/stmt.py
@@ -33,8 +33,19 @@ from typing import Any
import tvm_ffi
-from tvm.ir import DictAttrs, Expr, Op, Range, Span, StringImm, TensorRegion,
Var, make_node
-from tvm.runtime import Object, Scriptable
+from tvm.ir import (
+ DictAttrs,
+ Expr,
+ Op,
+ Range,
+ Scriptable,
+ Span,
+ StringImm,
+ TensorRegion,
+ Var,
+ make_node,
+)
+from tvm.runtime import Object
from . import _ffi_api
from .exec_scope import ScopeIdDef
diff --git a/python/tvm/tirx/tile_primitive.py
b/python/tvm/tirx/tile_primitive.py
index 13685dd6da..ed026a4e4f 100644
--- a/python/tvm/tirx/tile_primitive.py
+++ b/python/tvm/tirx/tile_primitive.py
@@ -25,8 +25,8 @@ from typing import Any, ClassVar
import tvm_ffi
from tvm_ffi import register_object
-from tvm.ir import Expr, Op, Range, Var
-from tvm.runtime import Object, Scriptable
+from tvm.ir import Expr, Op, Range, Scriptable, Var
+from tvm.runtime import Object
from tvm.target import Target
from . import _ffi_api
diff --git a/python/tvm/tirx/transform/transform.py
b/python/tvm/tirx/transform/transform.py
index ff7e536a1c..91ae6ec336 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -17,7 +17,6 @@
"""Wrapping existing transformations."""
# pylint: disable=invalid-name, unsupported-binary-operation
-import enum
from collections.abc import Callable
import tvm_ffi as _ffi
@@ -399,65 +398,6 @@ def VerifyMemory():
return _ffi_api.VerifyMemory() # type: ignore
-@_ffi.register_object("s_tir.transform.HoistIfThenElseConfig")
-class HoistIfThenElseConfig(_ffi.Object):
- """Config for hoist if then else pass"""
-
-
-class HoistedConditionals(enum.Flag):
- """Flags for use in HoistExpressionConfig.conditional_types
-
- Each bitflag represents a type of expression that should be
- hoisted to the outermost loop possible.
- """
-
- Never = 0
- """ No hoisting of conditionals """
-
- IfElseStmt = 1
- """ If set, look for hoist candidates in IfElseStmt """
-
- IfElseExpr = 2
- """ If set, look for hoist candidates in tirx.if_then_else """
-
- BooleanExpression = 4
- """ If set, look for hoist candidates in all boolean expressions """
-
- UsingBlockVar = 8
- """ If set, allow hoisting of conditionals that use a block variable (e.g.
threadIdx.x) """
-
- All = IfElseStmt | IfElseExpr | BooleanExpression | UsingBlockVar
- """ Enable all hoisting of conditionals"""
-
-
-class HoistedLetBindings(enum.Flag):
- """Flags for use in HoistExpressionConfig.let_binding_types
-
- Each bitflag represents a type of let binding expression that should be
- hoisted to the outermost loop possible.
- """
-
- Never = 0
- """ No hoisting of let bindings """
-
- RequiredByConditional = 1
- """ Bindings that are used by a hoisted conditional """
-
- Bind = 2
- """ Bindings occurring in Bind nodes """
-
- LetExpr = 4
- """ Bindings occurring in Let expressions """
-
- All = RequiredByConditional | Bind | LetExpr
- """ Enable all hoisting of let bindings """
-
-
-@_ffi.register_object("s_tir.transform.HoistExpressionConfig")
-class HoistExpressionConfig(_ffi.Object):
- """Config for hoist expression pass"""
-
-
def FlattenBuffer():
"""Flatten the multi-dimensional TensorLoad and TensorStore to single
dimensional
TensorLoad/TensorStore for the TIR not contains opaque block.
diff --git a/python/tvm/topi/utils.py b/python/tvm/topi/utils.py
index b1bbe29178..6f1151355c 100644
--- a/python/tvm/topi/utils.py
+++ b/python/tvm/topi/utils.py
@@ -25,6 +25,7 @@ import numpy as np
import tvm
from tvm import te
from tvm.s_tir import sbijective_layout, slayout
+from tvm.sym.utils import get_const_int
from . import cpp, tag
@@ -103,29 +104,6 @@ def prod(x):
return res
-def get_const_int(expr):
- """Verifies expr is integer and get the constant value.
-
- Parameters
- ----------
- expr : tvm.Expr or int
- The input expression.
-
- Returns
- -------
- out_value : int
- The output.
- """
- if isinstance(expr, Integral):
- return expr
- if not isinstance(expr, tvm.tirx.IntImm):
- ana = tvm.sym.Analyzer()
- expr = ana.simplify(expr)
- if not isinstance(expr, tvm.tirx.IntImm):
- raise ValueError("Expect value to be constant int")
- return int(expr.value)
-
-
def get_const_float(expr):
"""Verifies expr is a floating point and get the constant value.
diff --git a/src/ir/prim/op.cc b/src/ir/prim/op.cc
index 9236e6f641..55d7fc462a 100644
--- a/src/ir/prim/op.cc
+++ b/src/ir/prim/op.cc
@@ -835,7 +835,7 @@ PrimExpr prim::clz(PrimExpr x, Span span) {
TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef()
- .def_packed("node._const",
+ .def_packed("ir._const",
[](ffi::PackedArgs args, ffi::Any* ret) {
if (auto opt = args[0].try_cast<ffi::BigInt>();
opt.has_value()) {
*ret = prim::MakeConst(args[1].cast<PrimType>(), *opt,
args[2].cast<Span>());
diff --git a/src/relax/ir/dependent_type.cc b/src/relax/ir/dependent_type.cc
index 231b3c9f4f..a1aec4c20f 100644
--- a/src/relax/ir/dependent_type.cc
+++ b/src/relax/ir/dependent_type.cc
@@ -378,7 +378,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def("relax.UpdateType", [](Expr expr, Type ty) { UpdateType(expr, ty);
})
- .def("ir.ExprType", [](Expr expr) { return GetType(expr); });
+ .def("relax.ExprType", [](Expr expr) { return GetType(expr); });
}
} // namespace relax
diff --git a/src/script/printer/config.cc b/src/script/printer/config.cc
index 3e384a2628..51558a1316 100644
--- a/src/script/printer/config.cc
+++ b/src/script/printer/config.cc
@@ -182,9 +182,12 @@ ffi::Array<ffi::String>
PrinterConfigNode::GetBuiltinKeywords() {
TVM_FFI_STATIC_INIT_BLOCK() {
PrinterConfigNode::RegisterReflection();
- ffi::reflection::GlobalDef().def(
- "node.PrinterConfig",
- [](ffi::Map<ffi::String, Any> config_dict) { return
PrinterConfig(config_dict); });
+ auto make_config = [](ffi::Map<ffi::String, Any> config_dict) {
+ return PrinterConfig(config_dict);
+ };
+ ffi::reflection::GlobalDef()
+ .def("script.printer.PrinterConfig", make_config)
+ .def("node.PrinterConfig", make_config);
}
} // namespace tvm
diff --git a/src/te/script/printer/expr.cc b/src/te/script/printer/expr.cc
index da29b6da85..ca6c171804 100644
--- a/src/te/script/printer/expr.cc
+++ b/src/te/script/printer/expr.cc
@@ -41,7 +41,7 @@ ffi::Optional<ExprDoc>
CommReducerDocTranslate(DocTranslatorObj* d, ffi::AnyView
for (const PrimExpr& value : reducer->identity_element)
identity.push_back(d->Translate(value).value());
ExprDoc result_doc = result.size() == 1 ? result[0] :
ExprDoc(TupleDoc(result));
- return NamespaceDoc("tirx")
+ return NamespaceDoc("s_tir")
->Attr("comm_reducer")
->Call({LambdaDoc(args, result_doc), ListDoc(identity)});
}
@@ -76,7 +76,7 @@ ffi::Optional<ExprDoc> ReduceDocTranslate(DocTranslatorObj*
d, ffi::AnyView inpu
d->RecordOrigin(source, node->source);
d->RecordOrigin(axis, node->axis);
d->RecordOrigin(init, node->init);
- return NamespaceDoc("tirx")->Attr("Reduce")->Call(
+ return NamespaceDoc("s_tir")->Attr("Reduce")->Call(
{combiner, source, axis, condition, value_index, init});
}
diff --git a/tests/python/relax/script/test_relax_script_printer.py
b/tests/python/relax/script/test_relax_script_printer.py
index 86e8aff794..76dcba5285 100644
--- a/tests/python/relax/script/test_relax_script_printer.py
+++ b/tests/python/relax/script/test_relax_script_printer.py
@@ -28,11 +28,11 @@ from tvm import IRModule, relax, tirx
from tvm.ir import Range
from tvm.relax import TensorType
from tvm.relax.distributed import DeviceMesh, DTensorType, Placement
-from tvm.runtime.script_printer import PrinterConfig, _script
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
+from tvm.script.printer.scriptable import PrinterConfig, _script
def _assert_print(obj, expected):
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_hoist_expression.py
b/tests/python/s_tir/transform/test_s_tir_transform_hoist_expression.py
index fd502f55c9..4de0d34327 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_hoist_expression.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_hoist_expression.py
@@ -15,14 +15,64 @@
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: F401
+import enum
+
import tvm
import tvm.testing
from tvm import s_tir
-from tvm.s_tir.transform import HoistedConditionals, HoistedLetBindings
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
+class HoistedConditionals(enum.Flag):
+ """Test flags for the native hoisted_conditionals pass option
+
+ Each bitflag represents a type of expression that should be
+ hoisted to the outermost loop possible.
+ """
+
+ Never = 0
+ """ No hoisting of conditionals """
+
+ IfElseStmt = 1
+ """ If set, look for hoist candidates in IfElseStmt """
+
+ IfElseExpr = 2
+ """ If set, look for hoist candidates in tirx.if_then_else """
+
+ BooleanExpression = 4
+ """ If set, look for hoist candidates in all boolean expressions """
+
+ UsingBlockVar = 8
+ """ If set, allow hoisting of conditionals that use a block variable (e.g.
threadIdx.x) """
+
+ All = IfElseStmt | IfElseExpr | BooleanExpression | UsingBlockVar
+ """ Enable all hoisting of conditionals"""
+
+
+class HoistedLetBindings(enum.Flag):
+ """Test flags for the native hoisted_let_bindings pass option
+
+ Each bitflag represents a type of let binding expression that should be
+ hoisted to the outermost loop possible.
+ """
+
+ Never = 0
+ """ No hoisting of let bindings """
+
+ RequiredByConditional = 1
+ """ Bindings that are used by a hoisted conditional """
+
+ Bind = 2
+ """ Bindings occurring in Bind nodes """
+
+ LetExpr = 4
+ """ Bindings occurring in Let expressions """
+
+ All = RequiredByConditional | Bind | LetExpr
+ """ Enable all hoisting of let bindings """
+
+
def _run_transform(before, hoisted_conditionals, hoisted_let_bindings):
"""Run HoistExpression transform with the given configuration."""
before_mod = tvm.IRModule.from_expr(before)
diff --git a/tests/python/script/test_script_printer_entry.py
b/tests/python/script/test_script_printer_entry.py
index 6a90fa3c3e..e7f9cd51af 100644
--- a/tests/python/script/test_script_printer_entry.py
+++ b/tests/python/script/test_script_printer_entry.py
@@ -23,7 +23,7 @@ import tvm_ffi
from tvm_ffi.access_path import AccessPath
import tvm
-from tvm.runtime.script_printer import PrinterConfig, _script
+from tvm.script.printer.scriptable import PrinterConfig, _script
@pytest.mark.parametrize(
diff --git a/tests/python/tirx-base/test_tir_index_map.py
b/tests/python/tirx-base/test_tir_index_map.py
index 13f07c1849..9a1fd90536 100644
--- a/tests/python/tirx-base/test_tir_index_map.py
+++ b/tests/python/tirx-base/test_tir_index_map.py
@@ -21,8 +21,7 @@ import tvm_ffi
import tvm
import tvm.testing
-from tvm.ir import assert_structural_equal
-from tvm.runtime import const
+from tvm.ir import assert_structural_equal, const
from tvm.script import tirx as T
from tvm.tirx import IndexMap, IntImm, Var, floordiv, floormod
diff --git a/tests/python/tirx/codegen/test_ptx_addr.py
b/tests/python/tirx/codegen/test_ptx_addr.py
index a594073d79..8017eb32a3 100644
--- a/tests/python/tirx/codegen/test_ptx_addr.py
+++ b/tests/python/tirx/codegen/test_ptx_addr.py
@@ -20,8 +20,7 @@ import pytest
import tvm_ffi
import tvm
-from tvm.ir import Call, Op
-from tvm.runtime import const
+from tvm.ir import Call, Op, const
from tvm.script import tirx as T
from tvm.tirx.expr import Broadcast, CallEffectKind
diff --git a/tests/python/tirx/script/test_tirx_script_printer.py
b/tests/python/tirx/script/test_tirx_script_printer.py
index 465fa6e169..3be8e04769 100644
--- a/tests/python/tirx/script/test_tirx_script_printer.py
+++ b/tests/python/tirx/script/test_tirx_script_printer.py
@@ -26,11 +26,11 @@ import tvm.testing
from tvm import ir, s_tir, tirx
from tvm import tirx as tir
from tvm.ir import Range, assert_structural_equal
-from tvm.runtime.script_printer import _script
from tvm.script import ir as I
from tvm.script import ir_builder as IB
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
+from tvm.script.printer.scriptable import _script
from tvm.script.tirx import tile as Tx
from tvm.tirx.cuda import op as cuda_op
from tvm.tirx.script import ir_builder as TB
@@ -587,11 +587,13 @@ T.atan(T.float32(1.0))
def test_comm_reducer():
- obj = TB.comm_reducer(lambda x, y: x + y, identity=[TB.float32(0)])
+ from tvm.s_tir.script import ir_builder as SB
+
+ obj = SB.comm_reducer(lambda x, y: x + y, identity=[TB.float32(0)])
_assert_print(
obj,
"""
-T.comm_reducer(lambda x, y: x + y, [T.float32(0.0)])
+Ts.comm_reducer(lambda x, y: x + y, [T.float32(0.0)])
""",
)
diff --git
a/tests/python/tirx/script/test_tirx_script_printer_invisible_path_info.py
b/tests/python/tirx/script/test_tirx_script_printer_invisible_path_info.py
index ad9407d253..ab3a1bf16f 100644
--- a/tests/python/tirx/script/test_tirx_script_printer_invisible_path_info.py
+++ b/tests/python/tirx/script/test_tirx_script_printer_invisible_path_info.py
@@ -21,7 +21,7 @@ from tvm_ffi.access_path import AccessPath
import tvm
import tvm.testing
-from tvm.runtime.script_printer import PrinterConfig, _script
+from tvm.script.printer.scriptable import PrinterConfig, _script
def test_render_invisible_path_info_defaults_to_true():
