This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 9aa2b70cbe07babb3098a7f70d1f65f8e8aaec5b Author: Tianqi Chen <[email protected]> AuthorDate: Sat Oct 17 14:50:05 2026 -0400 [REFACTOR][TIRx] Simplify Python APIs and source docstrings (#20777) --- tirx/CMakeLists.txt | 2 +- tirx/python/tvm_tirx/backend/cuda/__init__.py | 4 +- .../tvm_tirx/backend/cuda/lang/tile_scheduler.py | 150 +++++++++++---------- .../cuda/{transforms => transform}/__init__.py | 8 +- tirx/python/tvm_tirx/script/parser/__init__.py | 20 ++- tirx/python/tvm_tirx/tirx/__init__.py | 2 +- tirx/python/tvm_tirx/tirx/compilation_pipeline.py | 10 +- tirx/python/tvm_tirx/tirx/expr.py | 22 +-- tirx/python/tvm_tirx/tirx/op/__init__.py | 3 +- tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py | 5 +- .../tvm_tirx/tirx/script/ir_builder/triton.py | 7 +- tirx/python/tvm_tirx/tirx/transform/transform.py | 12 -- .../bind_backend_config.cc | 6 +- .../cuda/{transforms => transform}/lower_iket.cc | 4 +- tirx/tests/python/sym/test_deduce_bound.py | 2 +- tirx/tests/python/sym/test_rewrite_simplify.py | 2 +- .../tirx/analysis/test_check_well_defined.py | 2 +- .../python/tirx/analysis/test_undefined_vars.py | 6 +- tirx/tests/python/tirx/base/test_nodes.py | 16 +-- tirx/tests/python/tirx/iket/test_profiler.py | 8 +- tirx/tests/python/tirx/test_cuda_backend_config.py | 6 +- 21 files changed, 143 insertions(+), 154 deletions(-) diff --git a/tirx/CMakeLists.txt b/tirx/CMakeLists.txt index e32a7f92aa..c2c31594f8 100644 --- a/tirx/CMakeLists.txt +++ b/tirx/CMakeLists.txt @@ -158,7 +158,7 @@ file(GLOB CODEGEN_SRCS CONFIGURE_DEPENDS src/codegen/llvm/canonicalizer/*.cc src/backend/cuda/codegen/*.cc src/backend/cuda/op/*.cc - src/backend/cuda/transforms/*.cc + src/backend/cuda/transform/*.cc src/backend/cuda/tile/*.cc src/backend/cuda/layout/*.cc src/backend/hexagon/codegen/*.cc diff --git a/tirx/python/tvm_tirx/backend/cuda/__init__.py b/tirx/python/tvm_tirx/backend/cuda/__init__.py index ee50466435..a26155860d 100644 --- a/tirx/python/tvm_tirx/backend/cuda/__init__.py +++ b/tirx/python/tvm_tirx/backend/cuda/__init__.py @@ -40,7 +40,7 @@ _LAZY_SUBMODULES = { "script", "target_tags", "tile", - "transforms", + "transform", } @@ -140,5 +140,5 @@ __all__ = [ "target_tags", "thread_scope", "tile", - "transforms", + "transform", ] diff --git a/tirx/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py b/tirx/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py index dc83f44393..24eba47b80 100644 --- a/tirx/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py +++ b/tirx/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py @@ -54,32 +54,35 @@ class BaseTileScheduler: class ClusterPersistentScheduler2D(BaseTileScheduler): - """ - Tile scheduler for cluster-based persistent kernels. + """Tile scheduler for cluster-based persistent kernels. Distributes a 2D tile grid across persistent clusters using group-major ordering for L2 cache locality. Each cluster starts at its cluster_id and strides by num_clusters to process tiles. Tile Ordering (group-major for L2 locality): + - Tiles are grouped into "L2 groups" of `l2_group_size` rows - Within a group, tiles are visited in column-major order within the group - Groups are processed in row-major order - Example with 4x4 tiles, l2_group_size=2: + Example with 4x4 tiles, l2_group_size=2:: + Group 0 (rows 0-1): 0 2 4 6 1 3 5 7 Group 1 (rows 2-3): 8 10 12 14 9 11 13 15 Serpentine Mode (serpentine=True): + - Uses CUTLASS-style 2D block swizzle with serpentine traversal - Grid is divided into swizzle_size x swizzle_size blocks - Within each block, tiles are visited in row-major order - Blocks are traversed in serpentine order (even block-rows forward, odd backward) - This provides better L2 locality by reusing both A and B tiles - Example with 4x4 tiles, swizzle_size=2, serpentine=True: + Example with 4x4 tiles, swizzle_size=2, serpentine=True:: + Block layout: Block(0,0) Block(0,1) Block(1,0) Block(1,1) @@ -127,77 +130,83 @@ class ClusterPersistentScheduler2D(BaseTileScheduler): Usage ----- - ```python - scheduler = ClusterPersistentScheduler2D( - "sched", num_m_tiles=M_TILES, num_n_tiles=N_TILES, - num_clusters=NUM_CLUSTERS, l2_group_size=8 - ) - scheduler.init(cluster_id) # cluster_id = cta_idx // CLUSTER_SIZE - - while scheduler.valid(): - m = T.meta_var(scheduler.m_idx) # current M tile - n = T.meta_var(scheduler.n_idx) # current N tile - # ... process tile (m, n) ... - scheduler.next_tile() - ``` + + .. code-block:: python + + scheduler = ClusterPersistentScheduler2D( + "sched", num_m_tiles=M_TILES, num_n_tiles=N_TILES, + num_clusters=NUM_CLUSTERS, l2_group_size=8 + ) + scheduler.init(cluster_id) # cluster_id = cta_idx // CLUSTER_SIZE + + while scheduler.valid(): + m = T.meta_var(scheduler.m_idx) # current M tile + n = T.meta_var(scheduler.n_idx) # current N tile + # ... process tile (m, n) ... + scheduler.next_tile() + Examples -------- Example 1: Basic persistent kernel - ``` - num_m_tiles=4, num_n_tiles=4, num_clusters=3, l2_group_size=2 - cluster_m=1, cluster_n=1 (default, no tile subdivision) - - Group-major tile numbering (l2_group_size=2): - n=0 n=1 n=2 n=3 - m=0 0 2 4 6 ┐ L2 group 0 - m=1 1 3 5 7 ┘ - m=2 8 10 12 14 ┐ L2 group 1 - m=3 9 11 13 15 ┘ - - Work distribution (cluster starts at cluster_id, strides by num_clusters=3): - cluster 0: work_idx 0,3,6,9,12,15 -> tiles 0,3,6,9,12,15 - cluster 1: work_idx 1,4,7,10,13 -> tiles 1,4,7,10,13 - cluster 2: work_idx 2,5,8,11,14 -> tiles 2,5,8,11,14 - - Tile grid (which cluster handles each tile): - n=0 n=1 n=2 n=3 - m=0 C0 C2 C1 C0 ┐ L2 group 0 - m=1 C1 C0 C2 C1 ┘ - m=2 C2 C1 C0 C2 ┐ L2 group 1 - m=3 C0 C2 C1 C0 ┘ - - Tile sequence per cluster (in execution order): - cluster 0: (0,0)->(1,1)->(0,3)->(2,0)->(2,3)->(3,3) - cluster 1: (1,0)->(0,2)->(1,3)->(2,1)->(3,2) - cluster 2: (0,1)->(1,2)->(2,0)->(3,1)->(2,3) - ``` + + .. code-block:: text + + num_m_tiles=4, num_n_tiles=4, num_clusters=3, l2_group_size=2 + cluster_m=1, cluster_n=1 (default, no tile subdivision) + + Group-major tile numbering (l2_group_size=2): + n=0 n=1 n=2 n=3 + m=0 0 2 4 6 ┐ L2 group 0 + m=1 1 3 5 7 ┘ + m=2 8 10 12 14 ┐ L2 group 1 + m=3 9 11 13 15 ┘ + + Work distribution (cluster starts at cluster_id, strides by num_clusters=3): + cluster 0: work_idx 0,3,6,9,12,15 -> tiles 0,3,6,9,12,15 + cluster 1: work_idx 1,4,7,10,13 -> tiles 1,4,7,10,13 + cluster 2: work_idx 2,5,8,11,14 -> tiles 2,5,8,11,14 + + Tile grid (which cluster handles each tile): + n=0 n=1 n=2 n=3 + m=0 C0 C2 C1 C0 ┐ L2 group 0 + m=1 C1 C0 C2 C1 ┘ + m=2 C2 C1 C0 C2 ┐ L2 group 1 + m=3 C0 C2 C1 C0 ┘ + + Tile sequence per cluster (in execution order): + cluster 0: (0,0)->(1,1)->(0,3)->(2,0)->(2,3)->(3,3) + cluster 1: (1,0)->(0,2)->(1,3)->(2,1)->(3,2) + cluster 2: (0,1)->(1,2)->(2,0)->(3,1)->(2,3) + Example 2: 2SM GEMM (typical B200 config) - ``` - M=1024, N=512, CTA_M=128, MMA_N=128, CLUSTER_M=2, CLUSTER_N=1 - => M_TILES=8, N_TILES=4 - => CLUSTER_M_TILES=4, CLUSTER_N_TILES=4 (scheduler at cluster granularity) - - Scheduler params: - num_m_tiles=4, num_n_tiles=4, num_clusters=74, l2_group_size=8 - cluster_m=1, cluster_n=1 - - Key: Scheduler outputs CLUSTER-level tiles. - All CTAs in same cluster get SAME (m_idx, n_idx) from scheduler. - CTAs differentiate via cluster_rank (computed OUTSIDE scheduler): - cluster_rank = cta_idx % CLUSTER_SIZE - cb_m = cluster_rank % CLUSTER_M # 0 or 1 for 2SM - cb_n = cluster_rank // CLUSTER_M # 0 for 2SM - - Final CTA tile: - cta_m = m_idx * CLUSTER_M + cb_m - cta_n = n_idx * CLUSTER_N + cb_n - - Example: cluster 5 gets scheduler tile (1,2) - CTA rank=0 (cb_m=0): actual tile (2,2) - CTA rank=1 (cb_m=1): actual tile (3,2) - ``` + + .. code-block:: text + + M=1024, N=512, CTA_M=128, MMA_N=128, CLUSTER_M=2, CLUSTER_N=1 + => M_TILES=8, N_TILES=4 + => CLUSTER_M_TILES=4, CLUSTER_N_TILES=4 (scheduler at cluster granularity) + + Scheduler params: + num_m_tiles=4, num_n_tiles=4, num_clusters=74, l2_group_size=8 + cluster_m=1, cluster_n=1 + + Key: Scheduler outputs CLUSTER-level tiles. + All CTAs in same cluster get SAME (m_idx, n_idx) from scheduler. + CTAs differentiate via cluster_rank (computed OUTSIDE scheduler): + cluster_rank = cta_idx % CLUSTER_SIZE + cb_m = cluster_rank % CLUSTER_M # 0 or 1 for 2SM + cb_n = cluster_rank // CLUSTER_M # 0 for 2SM + + Final CTA tile: + cta_m = m_idx * CLUSTER_M + cb_m + cta_n = n_idx * CLUSTER_N + cb_n + + Example: cluster 5 gets scheduler tile (1,2) + CTA rank=0 (cb_m=0): actual tile (2,2) + CTA rank=1 (cb_m=1): actual tile (3,2) + """ def __init__( @@ -469,7 +478,8 @@ class FlashAttentionLinearScheduler(BaseTileScheduler): """Linear 3D scheduler for flash attention (batch, head, m_block). Used for non-causal attention with simple linear decomposition. - Maps linear_idx -> (batch_idx, head_idx, m_block_idx) using: + Maps linear_idx -> (batch_idx, head_idx, m_block_idx) using:: + batch = linear_idx // (num_heads * num_m_blocks) head = (linear_idx % (num_heads * num_m_blocks)) // num_m_blocks m_block = linear_idx % num_m_blocks diff --git a/tirx/python/tvm_tirx/backend/cuda/transforms/__init__.py b/tirx/python/tvm_tirx/backend/cuda/transform/__init__.py similarity index 85% rename from tirx/python/tvm_tirx/backend/cuda/transforms/__init__.py rename to tirx/python/tvm_tirx/backend/cuda/transform/__init__.py index 6107a1250a..dd946a328b 100644 --- a/tirx/python/tvm_tirx/backend/cuda/transforms/__init__.py +++ b/tirx/python/tvm_tirx/backend/cuda/transform/__init__.py @@ -35,22 +35,22 @@ def LowerIket(): fpass : tvm_tirx.ir.transform.Pass The result pass. """ - return get_global_func("tirx.backend.cuda.transforms.LowerIket")() + return get_global_func("tirx.backend.cuda.transform.LowerIket")() def BindBackendConfig(): """Resolve per-entry compilation settings and targets before tile dispatch.""" - return get_global_func("tirx.backend.cuda.transforms.BindBackendConfig")("") + return get_global_func("tirx.backend.cuda.transform.BindBackendConfig")("") def SpecializeDeviceHelpers(): """Clone private device helpers into each independently compiled CUDA group.""" - return get_global_func("tirx.backend.cuda.transforms.SpecializeDeviceHelpers")() + return get_global_func("tirx.backend.cuda.transform.SpecializeDeviceHelpers")() def SpecializeEntryHelpers(): """Specialize shared helper calls before target-sensitive tile dispatch.""" - return get_global_func("tirx.backend.cuda.transforms.SpecializeEntryHelpers")() + return get_global_func("tirx.backend.cuda.transform.SpecializeEntryHelpers")() __all__ = ["BindBackendConfig", "LowerIket", "SpecializeDeviceHelpers", "SpecializeEntryHelpers"] diff --git a/tirx/python/tvm_tirx/script/parser/__init__.py b/tirx/python/tvm_tirx/script/parser/__init__.py index 3e84fd9ab4..508490d3b7 100644 --- a/tirx/python/tvm_tirx/script/parser/__init__.py +++ b/tirx/python/tvm_tirx/script/parser/__init__.py @@ -14,7 +14,25 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Canonical TVMScript parser with language variant namespace initialization.""" +"""Translate TVMScript syntax into IR-builder calls. + +Notes +----- +A registered language namespace supplies hooks for function construction, +value binding, statement emission, and control flow. Use ``register_namespace`` +and ``make_decorator`` to register a namespace and create its decorator. +The hook contracts are defined by ``script.ir_builder.parser_protocol``. + +Examples +-------- +.. code-block:: python + + from tvm_tirx.script import tirx as T + + @T.function + def example(a: T.int32): + T.evaluate(a + 1) +""" from __future__ import annotations diff --git a/tirx/python/tvm_tirx/tirx/__init__.py b/tirx/python/tvm_tirx/tirx/__init__.py index 508184eb50..ca2fc9604c 100644 --- a/tirx/python/tvm_tirx/tirx/__init__.py +++ b/tirx/python/tvm_tirx/tirx/__init__.py @@ -40,7 +40,7 @@ from .expr import Var, Cast from .expr import Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod from .expr import LShift, RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot from .expr import Min, Max, EQ, NE, LT, LE, GT, GE, And, Or, Not -from .expr import Select, TensorLoad +from .expr import Select from .expr import Let diff --git a/tirx/python/tvm_tirx/tirx/compilation_pipeline.py b/tirx/python/tvm_tirx/tirx/compilation_pipeline.py index f7f921f029..0a1c689bec 100644 --- a/tirx/python/tvm_tirx/tirx/compilation_pipeline.py +++ b/tirx/python/tvm_tirx/tirx/compilation_pipeline.py @@ -22,7 +22,7 @@ import inspect import tvm_tirx from tvm_tirx import tirx -from tvm_tirx.backend.cuda import transforms as cuda_transforms +from tvm_tirx.backend.cuda import transform as cuda_transform def default_tir_pipeline(*, target=None): @@ -70,8 +70,8 @@ def _bind_target(mod, target): return tvm_tirx.ir.transform.Sequential( [ tirx.transform.BindTarget(target), - cuda_transforms.BindBackendConfig(), - cuda_transforms.SpecializeEntryHelpers(), + cuda_transform.BindBackendConfig(), + cuda_transform.SpecializeEntryHelpers(), ] )(mod) @@ -125,8 +125,8 @@ def finalize_tir_pipeline(): [ tirx.transform.AnnotateEntryFunc(), tirx.transform.SplitHostDevice(), - cuda_transforms.LowerIket(), - cuda_transforms.SpecializeDeviceHelpers(), + cuda_transform.LowerIket(), + cuda_transform.SpecializeDeviceHelpers(), tirx.transform.BF16ComputeLegalize(), tirx.transform.FP8ComputeLegalize(), tirx.transform.MakePackedAPI(), diff --git a/tirx/python/tvm_tirx/tirx/expr.py b/tirx/python/tvm_tirx/tirx/expr.py index e87b678a55..2106659e80 100644 --- a/tirx/python/tvm_tirx/tirx/expr.py +++ b/tirx/python/tvm_tirx/tirx/expr.py @@ -29,7 +29,7 @@ For example, you can use addexp.a to get the left operand of an Add node. """ from tvm_tirx import ir -from tvm_tirx.ir import Expr +from tvm_tirx.ir import Expr as Expr from tvm_tirx.ir import const as const from tvm_tirx.ir._prim_expr_overload import ( # noqa: F401 EqualOp, @@ -37,7 +37,6 @@ from tvm_tirx.ir._prim_expr_overload import ( # noqa: F401 NotEqualOp, div_ambiguity_error, ) -from tvm_tirx.ir.location import UNKNOWN_LOC, Location from tvm_tirx.ir.prim import convert as convert # Retain historical imports as aliases of the canonical shared definitions. @@ -75,22 +74,3 @@ from tvm_tirx.ir.prim.expr import ( # noqa: F401 ) Var = ir.Var - - -def TensorLoad(tensor: Var, indices: list[Expr], loc: Location = UNKNOWN_LOC) -> ir.TensorLoad: - """Construct a validated tensor load. - - Parameters - ---------- - tensor : Var - The tensor to be loaded. - - indices : List[Expr] - The tensor indices to load values from. - - loc : Location, optional - The location of this expression in the source code. - - """ - - return ir.TensorLoad(tensor, indices, loc) diff --git a/tirx/python/tvm_tirx/tirx/op/__init__.py b/tirx/python/tvm_tirx/tirx/op/__init__.py index 15f6e662b7..1ffc80c7a0 100644 --- a/tirx/python/tvm_tirx/tirx/op/__init__.py +++ b/tirx/python/tvm_tirx/tirx/op/__init__.py @@ -112,7 +112,6 @@ from tvm_tirx.ir.prim.op import vscale as vscale from .. import _ffi_api from ..expr import ExprOp -from ..expr import TensorLoad as _make_tensor_load from ..tensor import is_tensor_var, tensor_data_ptr from ..type import TensorMapType @@ -790,7 +789,7 @@ def address_of(obj: Var | TensorLoad, loc: Location = UNKNOWN_LOC, *, ty=None) - """ if is_tensor_var(obj): n_dim = len(obj.ty.shape) - tensor_load = _make_tensor_load(obj, [0] * n_dim) + tensor_load = TensorLoad(obj, [0] * n_dim) return Call("tirx.address_of", [tensor_load], loc=loc, ty=ty) elif isinstance(obj, Var): if _is_tensormap_var(obj): diff --git a/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py index a79838854e..7bce0c7e00 100644 --- a/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py +++ b/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py @@ -29,9 +29,7 @@ import tvm_tirx from tvm_tirx import ir from tvm_tirx import ir as _ir from tvm_tirx import tirx as tir - -# import tirx.expr for direct ir construction to pass structural_equal comparison -from tvm_tirx.ir import FloatImm, IntImm, Range, TensorRegion, Type +from tvm_tirx.ir import FloatImm, IntImm, Range, TensorLoad, TensorRegion, Type from tvm_tirx.ir.location import UNKNOWN_LOC, Location from tvm_tirx.ir.type import _init_type_api, _normalize_type from tvm_tirx.runtime import MemoryScope @@ -47,7 +45,6 @@ from tvm_tirx.target import Target # pylint: disable=unused-import from tvm_tirx.tirx import Expr, IndexMap, Var, is_tensor_var, tile -from tvm_tirx.tirx.expr import TensorLoad from tvm_tirx.tirx.layout import Axis, ComposeLayout, Iter, Layout, R, S, TileLayout from tvm_tirx.tirx.thread_scope import ThreadScope diff --git a/tirx/python/tvm_tirx/tirx/script/ir_builder/triton.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/triton.py index d7dbfe6ae7..ba7582bfa1 100644 --- a/tirx/python/tvm_tirx/tirx/script/ir_builder/triton.py +++ b/tirx/python/tvm_tirx/tirx/script/ir_builder/triton.py @@ -42,9 +42,10 @@ class TritonKernel(BaseKernel): This class bridges the Triton kernel with TVM runtime. The compilation includes the following steps: - - Deduce the kernel signature and generate the Triton kernel - - Embed the compiled kernel into the current IRModule as an external module - - Generate a call to the Triton kernel following its calling convention via call_packed. + + - Deduce the kernel signature and generate the Triton kernel + - Embed the compiled kernel into the current IRModule as an external module + - Generate a call to the Triton kernel following its calling convention via call_packed. """ def __init__(self, func): diff --git a/tirx/python/tvm_tirx/tirx/transform/transform.py b/tirx/python/tvm_tirx/tirx/transform/transform.py index eb4b1e4589..4d4a7de295 100644 --- a/tirx/python/tvm_tirx/tirx/transform/transform.py +++ b/tirx/python/tvm_tirx/tirx/transform/transform.py @@ -19,8 +19,6 @@ from collections.abc import Callable -import tvm_ffi as _ffi - from . import _ffi_api from .function_pass import function_pass @@ -104,11 +102,6 @@ def PointerValueTypeRewrite(): return _ffi_api.PointerValueTypeRewrite() # type: ignore -@_ffi.register_object("tirx.transform.UnrollLoopConfig") -class UnrollLoopConfig(_ffi.Object): - """Config for unroll loop pass""" - - def UnrollLoop(): """Unroll the constant loop marked by unroll. @@ -122,11 +115,6 @@ def UnrollLoop(): return _ffi_api.UnrollLoop() # type: ignore -@_ffi.register_object("tirx.transform.RemoveNoOpConfig") -class RemoveNoOpConfig(_ffi.Object): - """Config for remove no op pass""" - - def RemoveNoOp(): """Remove No Op from the Stmt. diff --git a/tirx/src/backend/cuda/transforms/bind_backend_config.cc b/tirx/src/backend/cuda/transform/bind_backend_config.cc similarity index 97% rename from tirx/src/backend/cuda/transforms/bind_backend_config.cc rename to tirx/src/backend/cuda/transform/bind_backend_config.cc index 94fa1b94b4..acfacef5a5 100644 --- a/tirx/src/backend/cuda/transforms/bind_backend_config.cc +++ b/tirx/src/backend/cuda/transform/bind_backend_config.cc @@ -273,9 +273,9 @@ Pass SpecializeDeviceHelpers() { TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::GlobalDef() - .def("tirx.backend.cuda.transforms.BindBackendConfig", BindBackendConfig) - .def("tirx.backend.cuda.transforms.SpecializeDeviceHelpers", SpecializeDeviceHelpers) - .def("tirx.backend.cuda.transforms.SpecializeEntryHelpers", SpecializeEntryHelpers); + .def("tirx.backend.cuda.transform.BindBackendConfig", BindBackendConfig) + .def("tirx.backend.cuda.transform.SpecializeDeviceHelpers", SpecializeDeviceHelpers) + .def("tirx.backend.cuda.transform.SpecializeEntryHelpers", SpecializeEntryHelpers); } } // namespace transform } // namespace tirx diff --git a/tirx/src/backend/cuda/transforms/lower_iket.cc b/tirx/src/backend/cuda/transform/lower_iket.cc similarity index 99% rename from tirx/src/backend/cuda/transforms/lower_iket.cc rename to tirx/src/backend/cuda/transform/lower_iket.cc index fab8df2ba0..08a62bcee7 100644 --- a/tirx/src/backend/cuda/transforms/lower_iket.cc +++ b/tirx/src/backend/cuda/transform/lower_iket.cc @@ -18,7 +18,7 @@ */ /*! - * \file backend/cuda/transforms/lower_iket.cc + * \file backend/cuda/transform/lower_iket.cc * \brief Lower frontend-only TIRx IKET annotations to CUDA tracing helpers. */ @@ -1324,7 +1324,7 @@ Pass LowerIket() { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("tirx.backend.cuda.transforms.LowerIket", LowerIket); + refl::GlobalDef().def("tirx.backend.cuda.transform.LowerIket", LowerIket); } } // namespace transform diff --git a/tirx/tests/python/sym/test_deduce_bound.py b/tirx/tests/python/sym/test_deduce_bound.py index 6211902a16..ec41000dc9 100644 --- a/tirx/tests/python/sym/test_deduce_bound.py +++ b/tirx/tests/python/sym/test_deduce_bound.py @@ -236,7 +236,7 @@ def test_deduce_non_support(): test_non_support(tvm_tirx.tirx.EQ(a, 16)) test_non_support(tvm_tirx.tirx.NE(a, 16)) test_non_support(tvm_tirx.tirx.log(a)) - test_non_support(tvm_tirx.tirx.TensorLoad(decl_tensor([16], "int32"), [a])) + test_non_support(tvm_tirx.ir.TensorLoad(decl_tensor([16], "int32"), [a])) def test_deduce_floordiv(): diff --git a/tirx/tests/python/sym/test_rewrite_simplify.py b/tirx/tests/python/sym/test_rewrite_simplify.py index 002a3cb74a..4a6e1b80d4 100644 --- a/tirx/tests/python/sym/test_rewrite_simplify.py +++ b/tirx/tests/python/sym/test_rewrite_simplify.py @@ -1597,7 +1597,7 @@ class TestDivZero(BaseCompare): class TestSubBufferload(BaseCompare): buf = tvm_tirx.tirx.decl_tensor([1], dtype="float32") - load = tvm_tirx.tirx.TensorLoad(buf, [0]) + load = tvm_tirx.ir.TensorLoad(buf, [0]) test_case = tvm_tirx.testing.parameter( TestCase(load - load, 0.0), diff --git a/tirx/tests/python/tirx/analysis/test_check_well_defined.py b/tirx/tests/python/tirx/analysis/test_check_well_defined.py index 5e35f9ff2e..016fcc4bce 100644 --- a/tirx/tests/python/tirx/analysis/test_check_well_defined.py +++ b/tirx/tests/python/tirx/analysis/test_check_well_defined.py @@ -354,7 +354,7 @@ def test_tensor_load_asserted_type_matches_source_and_indices(): def test_tensor_load_malformed_indices_return_false_without_asserting(): tensor = tvm_tirx.tirx.decl_tensor((4, 4), "float32") vector_index = tvm_tirx.ir.prim.vector_expr.Ramp(0, 1, 4) - load = tvm_tirx.tirx.TensorLoad(tensor, [0, vector_index]) + load = tvm_tirx.ir.TensorLoad(tensor, [0, vector_index]) func = tvm_tirx.tirx.Function([tensor], tvm_tirx.ir.Evaluate(load)) graph = json.loads(tvm_tirx.ir.save_json(func)) diff --git a/tirx/tests/python/tirx/analysis/test_undefined_vars.py b/tirx/tests/python/tirx/analysis/test_undefined_vars.py index ba12177bf9..d3658116f4 100644 --- a/tirx/tests/python/tirx/analysis/test_undefined_vars.py +++ b/tirx/tests/python/tirx/analysis/test_undefined_vars.py @@ -34,7 +34,7 @@ def test_decl_tensor_data_is_use(): data_ptr = tirx.Var("buf_data", PtrType(PrimType("float32"))) buf = tirx.decl_tensor((n,), "float32", "buf", data=data_ptr) - body = tvm_tirx.ir.Evaluate(tirx.TensorLoad(buf, [0])) + body = tvm_tirx.ir.Evaluate(tvm_tirx.ir.TensorLoad(buf, [0])) decl = tvm_tirx.ir.Bind( buf, tvm_tirx.ir.Call( @@ -64,7 +64,7 @@ def test_decl_tensor_elem_offset_is_use(): elem_off = tirx.Var("buf_elem_offset", "int32") buf = tirx.decl_tensor((n,), "float32", "buf", data=data_ptr, elem_offset=elem_off) - body = tvm_tirx.ir.Evaluate(tirx.TensorLoad(buf, [0])) + body = tvm_tirx.ir.Evaluate(tvm_tirx.ir.TensorLoad(buf, [0])) decl = tvm_tirx.ir.Bind( buf, tvm_tirx.ir.Call( @@ -92,7 +92,7 @@ def test_alloc_tensor_data_is_def(): n = tirx.Var("n", "int32") buf = tirx.decl_tensor((n,), "float32", "buf") - body = tvm_tirx.ir.Evaluate(tirx.TensorLoad(buf, [0])) + body = tvm_tirx.ir.Evaluate(tvm_tirx.ir.TensorLoad(buf, [0])) alloc = tvm_tirx.ir.Bind( buf, tvm_tirx.ir.Call( diff --git a/tirx/tests/python/tirx/base/test_nodes.py b/tirx/tests/python/tirx/base/test_nodes.py index a0840dde38..b8ec2b40e9 100644 --- a/tirx/tests/python/tirx/base/test_nodes.py +++ b/tirx/tests/python/tirx/base/test_nodes.py @@ -342,16 +342,13 @@ def test_scoped_storage_vars(): def test_tensor_load_store(): b = tvm_tirx.tirx.decl_tensor((10,), "float32") - x = tvm_tirx.tirx.TensorLoad(b, [0]) + x = tvm_tirx.ir.TensorLoad(b, [0]) assert isinstance(x, tvm_tirx.ir.TensorLoad) - assert callable(tvm_tirx.tirx.TensorLoad) assert x.ty.dtype == "float32" assert x.source == b assert not hasattr(x, "tensor") loc = tvm_tirx.ir.SourceLoc(tvm_tirx.ir.SourceName("tensor_load"), 1, 2, 1, 3) core_load = tvm_tirx.ir.TensorLoad(b, [0], loc) - dialect_load = tvm_tirx.tirx.TensorLoad(b, [0], loc) - tvm_tirx.ir.assert_structural_equal(core_load, dialect_load) assert core_load.loc.same_as(loc) restored = tvm_tirx.ir.load_json(tvm_tirx.ir.save_json(core_load)) assert isinstance(restored.ty, tvm_tirx.ir.PrimType) @@ -366,15 +363,14 @@ def test_tensor_load_store(): ) [email protected]("constructor", [tvm_tirx.ir.TensorLoad, tvm_tirx.tirx.TensorLoad]) -def test_tensor_load_validation(constructor): +def test_tensor_load_validation(): tensor = tvm_tirx.tirx.decl_tensor((8, 8), "float32") with pytest.raises(RuntimeError, match="dimensional"): - constructor(tensor, [0]) + tvm_tirx.ir.TensorLoad(tensor, [0]) with pytest.raises(RuntimeError, match="Only the last index"): - constructor(tensor, [tvm_tirx.ir.prim.vector_expr.Ramp(0, 1, 4), 0]) + tvm_tirx.ir.TensorLoad(tensor, [tvm_tirx.ir.prim.vector_expr.Ramp(0, 1, 4), 0]) with pytest.raises(TypeError, match="does not support TensorLoad"): - constructor(tvm_tirx.tirx.Var("scalar", "float32"), [0]) + tvm_tirx.ir.TensorLoad(tvm_tirx.tirx.Var("scalar", "float32"), [0]) @pytest.mark.parametrize( @@ -464,7 +460,7 @@ def test_broadcast_to_scalable_vec(): def test_tensor_load_scalable_vec(): buf = tvm_tirx.tirx.decl_tensor((24,), "float32") index = tvm_tirx.ir.prim.vector_expr.Ramp(1, 1, 8 * tvm_tirx.tirx.vscale()) - load = tvm_tirx.tirx.TensorLoad(buf, [index]) + load = tvm_tirx.ir.TensorLoad(buf, [index]) assert isinstance(load, tvm_tirx.ir.TensorLoad) assert load.ty.dtype == "float32xvscalex8" diff --git a/tirx/tests/python/tirx/iket/test_profiler.py b/tirx/tests/python/tirx/iket/test_profiler.py index c4a549a3c7..c63a44afcc 100644 --- a/tirx/tests/python/tirx/iket/test_profiler.py +++ b/tirx/tests/python/tirx/iket/test_profiler.py @@ -32,7 +32,7 @@ from types import SimpleNamespace import pytest import tvm_tirx -from tvm_tirx.backend.cuda import transforms as cuda_transforms +from tvm_tirx.backend.cuda import transform as cuda_transform from tvm_tirx.codegen import Executable from tvm_tirx.script import tirx as T from tvm_tirx.testing import env @@ -379,7 +379,7 @@ def test_public_interface_is_official_only(): assert not hasattr(profiler, "capture") assert not hasattr(profiler, "export") assert not hasattr(tvm_tirx.tirx.transform, "LowerIket") - assert callable(cuda_transforms.LowerIket) + assert callable(cuda_transform.LowerIket) script = serial_a.script() assert 'T.cuda.iket_mark("a")' in script @@ -414,7 +414,7 @@ def test_annotation_ops_are_cuda_owned(name): def test_regular_lowering_strips_annotations_and_tokens(): - stripped = cuda_transforms.LowerIket()(tvm_tirx.IRModule({"main": token_loop})) + stripped = cuda_transform.LowerIket()(tvm_tirx.IRModule({"main": token_loop})) script = stripped.script() assert "cuda.iket" not in script assert "sentinel" not in script @@ -505,7 +505,7 @@ def test_tirx_pipelines_immediately_lower_iket(module_name, factory_name): pipeline(tvm_tirx.IRModule()) split = visited.index(("after", tvm_tirx.tirx.transform.SplitHostDevice().info.name)) - assert visited[split + 1] == ("before", cuda_transforms.LowerIket().info.name) + assert visited[split + 1] == ("before", cuda_transform.LowerIket().info.name) @pytest.mark.parametrize( diff --git a/tirx/tests/python/tirx/test_cuda_backend_config.py b/tirx/tests/python/tirx/test_cuda_backend_config.py index 0963e9311f..559af97eac 100644 --- a/tirx/tests/python/tirx/test_cuda_backend_config.py +++ b/tirx/tests/python/tirx/test_cuda_backend_config.py @@ -36,7 +36,7 @@ from tvm_tirx.backend.config import ( ) from tvm_tirx.backend.cuda import BackendConfig from tvm_tirx.backend.cuda.backend_config import compiler_options, resolve_backend_config -from tvm_tirx.backend.cuda.transforms import BindBackendConfig +from tvm_tirx.backend.cuda.transform import BindBackendConfig from tvm_tirx.script import tirx as T from tvm_tirx.testing import env @@ -407,7 +407,7 @@ np.testing.assert_array_equal(b.numpy(), x * np.float32(0.5)) def test_device_helpers_are_cloned_per_configuration(): _require_compiler("nvrtc") - from tvm_tirx.backend.cuda.transforms import SpecializeDeviceHelpers + from tvm_tirx.backend.cuda.transform import SpecializeDeviceHelpers from tvm_tirx.script import ir as I @I.ir_module @@ -609,7 +609,7 @@ import tvm_tirx import tvm_tirx.testing from tvm_tirx.backend.cuda import BackendConfig from tvm_tirx.backend.config import prepare_target -from tvm_tirx.backend.cuda.transforms import BindBackendConfig +from tvm_tirx.backend.cuda.transform import BindBackendConfig from tvm_tirx.script import tirx as T target = tvm_tirx.target.Target("cuda", host="llvm")
