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 4086b5ff0e [REFACTOR][TIRx] Simplify Python APIs and source docstrings
(#20777)
4086b5ff0e is described below
commit 4086b5ff0ef0bbfc6463f7a9d7ba581a9f3d663c
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Oct 17 14:50:05 2026 -0400
[REFACTOR][TIRx] Simplify Python APIs and source docstrings (#20777)
---
docs/reference/api/python/script/script.rst | 2 +-
docs/tirx/api/cuda.rst | 2 +-
python/tvm/s_tir/dlight/analysis/gemv.py | 2 +-
python/tvm/s_tir/dlight/gpu/low_batch_gemv.py | 2 +-
python/tvm/s_tir/dlight/gpu/reduction.py | 2 +-
.../test_s_tir_analysis_verify_block_scope.py | 2 +-
.../s_tir/script/test_s_tir_script_printer.py | 2 +-
.../test_s_tir_transform_mma_tensor_layout.py | 2 +-
tests/python/tirx/base/test_tirx_constructor.py | 4 +-
.../python/tirx/script/test_tirx_script_printer.py | 2 +-
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 +-
31 files changed, 154 insertions(+), 165 deletions(-)
diff --git a/docs/reference/api/python/script/script.rst
b/docs/reference/api/python/script/script.rst
index 612b634c7a..8fd2ae0545 100644
--- a/docs/reference/api/python/script/script.rst
+++ b/docs/reference/api/python/script/script.rst
@@ -57,7 +57,7 @@ The canonical ``MemoryScope`` class is documented in
:mod:`tvm_tirx.runtime`.
.. automodule:: tvm_tirx.script.tirx
:members:
:exclude-members: AlreadyEmitted, DataTypeImm, FuncType, GenericConst,
IRBuilder,
- IRModuleFrame, MemoryScope, MissingType, Op, PrimType, StringImm,
StringType, TensorRegion,
+ IRModuleFrame, MemoryScope, MissingType, Op, PrimType, StringImm,
StringType, TensorLoad, TensorRegion,
TensorRegionType, Type, Range, meta_var, Var, Call, CommReducer, Reduce,
FloatImm, IntImm, Cast, Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod, LShift,
RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot, Min, Max, EQ, NE, LT,
LE, GT, GE, And, Or, Not, Select, Ramp, Broadcast, Shuffle, CallEffectKind,
IterVar, ComposeLayout, DtypeConstructor, ThreadScope, Iter, Layout,
LetAnnotation, LocalVectorAnnotation, TileLayout, CallFFIKernelAttr,
TensorMapEncodeTiledAttr, function, jit, inli [...]
.. autofunction:: tvm_tirx.script.tirx.function
diff --git a/docs/tirx/api/cuda.rst b/docs/tirx/api/cuda.rst
index 701263a222..045c0e7bd8 100644
--- a/docs/tirx/api/cuda.rst
+++ b/docs/tirx/api/cuda.rst
@@ -141,6 +141,6 @@ IKET profiling
CUDA-specific transforms
------------------------
-.. automodule:: tvm_tirx.backend.cuda.transforms
+.. automodule:: tvm_tirx.backend.cuda.transform
:members:
:no-index:
diff --git a/python/tvm/s_tir/dlight/analysis/gemv.py
b/python/tvm/s_tir/dlight/analysis/gemv.py
index 82da265d90..9c8477eeb6 100644
--- a/python/tvm/s_tir/dlight/analysis/gemv.py
+++ b/python/tvm/s_tir/dlight/analysis/gemv.py
@@ -56,7 +56,7 @@ def get_reduction_expr(block: s_tir.SBlock) -> tirx.Expr |
None:
return None
if not tvm_ffi.structural_equal(
tensor_store.value.a,
- tirx.TensorLoad(tensor_store.dest, tensor_store.indices),
+ tvm_tirx.ir.TensorLoad(tensor_store.dest, tensor_store.indices),
map_free_vars=True,
):
return None
diff --git a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
index 5107f66b76..51dac9eaae 100644
--- a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
@@ -51,7 +51,7 @@ def _get_reduction_expr(block: s_tir.SBlock) -> tirx.Expr |
None:
return None
if not tvm_ffi.structural_equal(
tensor_store.value.a,
- tirx.TensorLoad(tensor_store.dest, tensor_store.indices),
+ tvm_tirx.ir.TensorLoad(tensor_store.dest, tensor_store.indices),
map_free_vars=True,
):
return None
diff --git a/python/tvm/s_tir/dlight/gpu/reduction.py
b/python/tvm/s_tir/dlight/gpu/reduction.py
index 341ca9433d..7b56adf004 100644
--- a/python/tvm/s_tir/dlight/gpu/reduction.py
+++ b/python/tvm/s_tir/dlight/gpu/reduction.py
@@ -47,7 +47,7 @@ def _get_reduction_expr(block: s_tir.SBlock) -> tirx.Expr |
None:
return None
if not tvm_ffi.structural_equal(
tensor_store.value.a,
- tirx.TensorLoad(tensor_store.dest, tensor_store.indices),
+ tvm_tirx.ir.TensorLoad(tensor_store.dest, tensor_store.indices),
map_free_vars=True,
):
return None
diff --git
a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
index 19ef7b1a9d..dda027f75e 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
@@ -135,7 +135,7 @@ def test_error_undeclared_tensor_in_schedulable_tir():
iter_vars=[tvm.s_tir.IterVar(tvm_tirx.ir.Range(0, n), bi, 0)], # 0 =
kDataPar
reads=[TensorRegion(A, [tvm_tirx.ir.Range(bi, bi + 1)],
ty=TensorRegionType())],
writes=[TensorRegion(B, [tvm_tirx.ir.Range(bi, bi + 1)],
ty=TensorRegionType())],
- body=tvm_tirx.ir.TensorStore(B, [bi], tvm_tirx.tirx.TensorLoad(A,
[bi])),
+ body=tvm_tirx.ir.TensorStore(B, [bi], tvm_tirx.ir.TensorLoad(A, [bi])),
name_hint="write_B",
)
block_realize = tvm.s_tir.SBlockRealize(
diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py
b/tests/python/s_tir/script/test_s_tir_script_printer.py
index 958a590211..6dce514f71 100644
--- a/tests/python/s_tir/script/test_s_tir_script_printer.py
+++ b/tests/python/s_tir/script/test_s_tir_script_printer.py
@@ -102,7 +102,7 @@ def test_function_symbolic_alloc_tensor_roundtrip():
ty_args=[buf.ty],
),
),
- tvm_tirx.ir.Evaluate(tirx.TensorLoad(buf, [0])),
+ tvm_tirx.ir.Evaluate(tvm_tirx.ir.TensorLoad(buf, [0])),
]
),
).with_attr({"s_tir": True, "script.namespace": "s_tir"})
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
b/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
index a4d7f1d09e..a81351920e 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
@@ -30,7 +30,7 @@ from tvm_tirx.script import tirx as T
def test_explicit_matrix_ab_access_is_rejected(scope, shape, access_kind):
tensor = tirx.decl_tensor(shape, "float32", scope=scope)
if access_kind == "load":
- body = tvm_tirx.ir.Evaluate(tirx.TensorLoad(tensor, [0, 0]))
+ body = tvm_tirx.ir.Evaluate(tvm_tirx.ir.TensorLoad(tensor, [0, 0]))
else:
body = tvm_tirx.ir.TensorStore(tensor, [0, 0], 0.0)
block = s_tir.SBlock([], [], [], "root", body, alloc_tensors=[tensor])
diff --git a/tests/python/tirx/base/test_tirx_constructor.py
b/tests/python/tirx/base/test_tirx_constructor.py
index a4fd3e38c9..5aaa6059f5 100644
--- a/tests/python/tirx/base/test_tirx_constructor.py
+++ b/tests/python/tirx/base/test_tirx_constructor.py
@@ -110,7 +110,7 @@ def test_expr_constructor():
tensor_var = tvm_tirx.tirx.Var("buf",
tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("float32")))
tensor = tvm_tirx.tirx.decl_tensor([16], "float32", data=tensor_var)
- x = tvm_tirx.tirx.TensorLoad(tensor, [1])
+ x = tvm_tirx.ir.TensorLoad(tensor, [1])
assert isinstance(x, tvm_tirx.ir.TensorLoad)
assert x.ty == tvm_tirx.ir.PrimType("float32")
assert x.source == tensor
@@ -287,7 +287,7 @@ def test_operator_base_categories_have_primitive_type():
),
tvm_tirx.tirx.Cast("float32", var),
tvm_tirx.tirx.Select(var < 1, var, 1),
- tvm_tirx.tirx.TensorLoad(tensor, [0]),
+ tvm_tirx.ir.TensorLoad(tensor, [0]),
tvm_tirx.ir.prim.vector_expr.Ramp(0, 1, 4),
tvm_tirx.ir.prim.vector_expr.Broadcast(var, 4),
tvm_tirx.ir.prim.vector_expr.Shuffle([tvm_tirx.ir.prim.vector_expr.Broadcast(var,
2)], [0]),
diff --git a/tests/python/tirx/script/test_tirx_script_printer.py
b/tests/python/tirx/script/test_tirx_script_printer.py
index 7c8ade324f..c631a944da 100644
--- a/tests/python/tirx/script/test_tirx_script_printer.py
+++ b/tests/python/tirx/script/test_tirx_script_printer.py
@@ -185,7 +185,7 @@ src[64:128, 64:128]
def test_tensor_load():
a = tirx.decl_tensor((128, 128), "float16", name="A")
- obj = tirx.TensorLoad(a, [128, 128])
+ obj = tvm_tirx.ir.TensorLoad(a, [128, 128])
_assert_print(
obj,
"""A = T.Var("A", T.Tensor((128, 128), T.float16))
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")
