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 c4328e3b97 [REFACTOR][TIRx] Return statements from tile dispatch
(#20767)
c4328e3b97 is described below
commit c4328e3b97b4a78da548768d2d0abe41dda632cc
Author: Bohan Hou <[email protected]>
AuthorDate: Fri Oct 16 18:47:17 2026 -0400
[REFACTOR][TIRx] Return statements from tile dispatch (#20767)
Tile dispatch currently returns a `tirx.Function` even though
`TileDispatch` only consumes its body. Change `FTileOpDispatch` and the
Python lowerer contract to return an `ir.Stmt` directly, then
recursively lower that replacement statement at the original
`Evaluate(Call)` site.
CUDA patterns now construct a statement sequence directly. Trainium
lowerers return their implementation bodies. Allocation/initialization
callbacks, shared state, and residual tile-call checks retain their
existing behavior. The dispatcher rejects non-statement results,
including the former `Function` return type.
Update the API documentation and C++/Python coverage for direct
statement results, recursive expansion of multi-statement
implementations, and invalid lowerer results. LowerTIRx
execution-context test probes now also return a native statement,
preserving their existing scope and structural-equality assertions.
Validation:
- Built TIRx and its C++ test executable in a fresh CPU-only build; both
`ThreadScope.*` tests passed, including typed statement dispatch.
- The non-GPU tile suite passed 153 tests, including all 14 dispatcher
tests and Trainium lowering coverage; existing hardware skips remained
in effect.
- The complete LowerTIRx test module passed 47 tests, with one existing
hardware-dependent skip.
- All five existing DSMEM dispatch cases passed their CPU-side
statement-type, instruction-count, and rejection checks.
- Pre-commit passed on all changed files.
- Python syntax compilation and/or pytest collection passed for all 16
modified Python files.
- C++ formatting comparison and `git diff --check` passed.
- GPU execution was not tested locally: CUDA pattern tests were skipped
by their existing hardware requirements. Upstream CI covers the broader
platform matrix.
---
docs/tirx/api/tile_dispatch.rst | 4 +-
docs/tirx/arch/tile_dispatch.rst | 3 +-
tirx/include/tvm/tirx/tile/dispatch.h | 8 ++--
.../tvm_tirx/backend/cuda/tile/patterns/common.py | 7 ++-
.../tvm_tirx/backend/trn/tile/binary/default.py | 7 ++-
.../backend/trn/tile/compose_op/binary_chain.py | 7 ++-
.../backend/trn/tile/compose_op/binary_reduce.py | 7 ++-
.../backend/trn/tile/compose_op/unary_reduce.py | 9 ++--
.../tvm_tirx/backend/trn/tile/copy/default.py | 6 +--
.../tvm_tirx/backend/trn/tile/gemm/default.py | 9 ++--
.../tvm_tirx/backend/trn/tile/reduction/utils.py | 10 ++---
.../tvm_tirx/backend/trn/tile/select/default.py | 7 ++-
.../tvm_tirx/backend/trn/tile/unary/default.py | 11 +++--
.../tvm_tirx/backend/trn/tile/unary/utils.py | 10 ++---
.../backend/trn/tile/unary/with_bias_scale.py | 11 +++--
tirx/python/tvm_tirx/tirx/tile/dispatch.py | 5 ++-
tirx/src/tirx/tile/dispatch.cc | 6 +--
tirx/src/tirx/tile/lower.cc | 5 +--
tirx/tests/cpp/tirx/thread_scope_test.cc | 2 +-
.../python/tirx/tile/cuda/test_cp_async_bulk.py | 2 +-
tirx/tests/python/tirx/tile/test_dispatcher.py | 51 ++++++++++++++++------
.../tests/python/tirx/transform/test_lower_tirx.py | 40 ++++++++---------
22 files changed, 121 insertions(+), 106 deletions(-)
diff --git a/docs/tirx/api/tile_dispatch.rst b/docs/tirx/api/tile_dispatch.rst
index 7feed70fe4..93da48bc9d 100644
--- a/docs/tirx/api/tile_dispatch.rst
+++ b/docs/tirx/api/tile_dispatch.rst
@@ -32,12 +32,12 @@ scope uses a reflected attribute. See
:doc:`../cuda_tile_patterns` for the
CUDA pattern extension workflow.
Lowerers receive a semantic view and ``DispatchContext`` and return a
-``tvm.tirx.Function`` body. They can prepare pointers/descriptors and expand
+``tvm.ir.Stmt`` directly. They can prepare pointers/descriptors and expand
layouts for the selected core instruction. They must report unsupported
contracts rather than select a different instruction family.
``FTileOpDispatch`` has the signature
-``(tvm.ir.Call, DispatchContext) -> tvm.tirx.Function``. Its presence
identifies
+``(tvm.ir.Call, DispatchContext) -> tvm.ir.Stmt``. Its presence identifies
a tile Op; no separate category registration is required. ``Instruction``
registers an adapter that decodes the semantic view and invokes its lowerer.
Custom implementations can register the callback directly using
diff --git a/docs/tirx/arch/tile_dispatch.rst b/docs/tirx/arch/tile_dispatch.rst
index 8d2fa2522b..2e17e168ae 100644
--- a/docs/tirx/arch/tile_dispatch.rst
+++ b/docs/tirx/arch/tile_dispatch.rst
@@ -46,7 +46,8 @@ The pass recognizes ``Evaluate(Call)`` by the presence of the
operator's
resolves the scope against the active thread set,
constructs a ``DispatchContext``, and
invokes the instruction's registered lowerer. It replaces the Evaluate with
-the returned function body and lowers any nested tensor calls.
+the returned statement and lowers any nested tensor calls. A lowerer may return
+a single statement or a ``SeqStmt`` containing multiple statements.
This preserves launch parameters, inter/intra scope maps, value ranges,
descriptor caches, host initialization, and allocation/initialization
diff --git a/tirx/include/tvm/tirx/tile/dispatch.h
b/tirx/include/tvm/tirx/tile/dispatch.h
index d03ff3270b..0e4983227f 100644
--- a/tirx/include/tvm/tirx/tile/dispatch.h
+++ b/tirx/include/tvm/tirx/tile/dispatch.h
@@ -25,7 +25,7 @@
#include <tvm/ffi/function.h>
#include <tvm/ir/op.h>
-#include <tvm/tirx/function.h>
+#include <tvm/tirx/stmt.h>
#include <tvm/tirx/tile/context.h>
namespace tvm::tirx::tile {
@@ -38,7 +38,7 @@ namespace op_attr {
* dispatcher can be implemented in C++ or registered from Python.
*/
inline constexpr ffi::reflection::AttrKey<
- ffi::TypedFunction<Function(const Call&, const DispatchContext&)>>
+ ffi::TypedFunction<Stmt(const Call&, const DispatchContext&)>>
kTileOpDispatch{"FTileOpDispatch"};
} // namespace op_attr
@@ -49,9 +49,9 @@ using FTileOpDispatch =
decltype(op_attr::kTileOpDispatch)::ValueType;
* \brief Invoke the Call's FTileOpDispatch attribute.
* \param call The tile instruction to lower.
* \param context The operator's execution context and accumulated callbacks.
- * \return The instruction implementation, whose body is expanded by
TileDispatch.
+ * \return The instruction implementation statement, expanded by TileDispatch.
*/
-TVM_TIRX_DLL Function Dispatch(const Call& call, const DispatchContext&
context);
+TVM_TIRX_DLL Stmt Dispatch(const Call& call, const DispatchContext& context);
} // namespace tvm::tirx::tile
#endif // TVM_TIRX_TILE_DISPATCH_H_
diff --git a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py
b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py
index 0382449415..272bb95092 100644
--- a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py
+++ b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py
@@ -35,9 +35,9 @@ from tvm_tirx.ir import (
ForKind,
IntImm,
Op,
- PrimType,
Range,
SeqStmt,
+ Stmt,
TensorRegion,
Var,
)
@@ -45,7 +45,6 @@ from tvm_tirx.ir.prim import Let
from tvm_tirx.runtime import DataType, MemoryScope
from tvm_tirx.script import tirx as T
from tvm_tirx.sym import Analyzer
-from tvm_tirx.tirx import Function
from tvm_tirx.tirx.layout import TileLayout
@@ -211,7 +210,7 @@ class Context:
def region(self, name, region, atom_shape, scopes, *, repeat_axes=None,
bounds=True):
return Region(self, name, region, atom_shape, scopes, repeat_axes,
bounds)
- def finish(self, statements):
+ def finish(self, statements) -> Stmt:
body = SeqStmt(
[
Evaluate(s) if not isinstance(s, Bind | SeqStmt) and
hasattr(s, "ty") else s
@@ -220,7 +219,7 @@ class Context:
)
for var, extent in reversed(list(zip(self.repeat or (),
self.repeat_shape or ()))):
body = For(var, 0, extent, ForKind.UNROLLED, body)
- return Function([], [*self.bindings, body], ret_type=PrimType("void"))
+ return SeqStmt([*self.bindings, body])
class Region:
diff --git a/tirx/python/tvm_tirx/backend/trn/tile/binary/default.py
b/tirx/python/tvm_tirx/backend/trn/tile/binary/default.py
index 8fc9014fe4..080434e625 100644
--- a/tirx/python/tvm_tirx/backend/trn/tile/binary/default.py
+++ b/tirx/python/tvm_tirx/backend/trn/tile/binary/default.py
@@ -17,9 +17,8 @@
"""Implementation of binary operator dispatches."""
-from tvm_tirx.ir import FloatImm
+from tvm_tirx.ir import FloatImm, Stmt
from tvm_tirx.script import tirx as T
-from tvm_tirx.tirx import Function
from tvm_tirx.tirx.tile import DispatchContext, fail
from tvm_tirx.tirx.tile.common import MapOpType
from tvm_tirx.tirx.tile.instruction import TensorCall
@@ -29,7 +28,7 @@ from ..instruction_generator import InstructionGenerator
from .utils import InstType, binary_map_ops, try_find_inst_nary
-def binary_trn(op: TensorCall, binary_op: MapOpType, sctx: DispatchContext) ->
Function | None:
+def binary_trn(op: TensorCall, binary_op: MapOpType, sctx: DispatchContext) ->
Stmt:
"""Generate a binary operation schedule for Trainium."""
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
fail("requires Trainium target and thread scope")
@@ -106,4 +105,4 @@ def binary_trn(op: TensorCall, binary_op: MapOpType, sctx:
DispatchContext) -> F
)
)
- return impl
+ return impl.body
diff --git a/tirx/python/tvm_tirx/backend/trn/tile/compose_op/binary_chain.py
b/tirx/python/tvm_tirx/backend/trn/tile/compose_op/binary_chain.py
index 455f4ba121..2fbdcd34a4 100644
--- a/tirx/python/tvm_tirx/backend/trn/tile/compose_op/binary_chain.py
+++ b/tirx/python/tvm_tirx/backend/trn/tile/compose_op/binary_chain.py
@@ -17,9 +17,8 @@
"""Implementation of BinaryChain dispatch."""
-from tvm_tirx.ir import TensorRegion
+from tvm_tirx.ir import Stmt, TensorRegion
from tvm_tirx.script import tirx as T
-from tvm_tirx.tirx import Function
from tvm_tirx.tirx.tile import DispatchContext
from tvm_tirx.tirx.tile.instruction import TensorCall
from tvm_tirx.tirx.tile.views import BinaryChain
@@ -30,7 +29,7 @@ from ..instruction_generator import InstructionGenerator
from .utils import opcode_table
-def binary_chain_trn(op: TensorCall, sctx: DispatchContext) -> Function | None:
+def binary_chain_trn(op: TensorCall, sctx: DispatchContext) -> Stmt:
"""Generate a TRN schedule for binary chain operations."""
op = TensorCall.decode(op)
assert isinstance(op, BinaryChain), f"invalid operator downcast: {op}"
@@ -114,4 +113,4 @@ def binary_chain_trn(op: TensorCall, sctx: DispatchContext)
-> Function | None:
T.evaluate(func(dst[tuple(dst_indices)], *srcs,
opcode0, opcode1, reverse[0], reverse[1])) # noqa: E501
# fmt: on
- return impl
+ return impl.body
diff --git a/tirx/python/tvm_tirx/backend/trn/tile/compose_op/binary_reduce.py
b/tirx/python/tvm_tirx/backend/trn/tile/compose_op/binary_reduce.py
index 70bfa96505..6797929eb1 100644
--- a/tirx/python/tvm_tirx/backend/trn/tile/compose_op/binary_reduce.py
+++ b/tirx/python/tvm_tirx/backend/trn/tile/compose_op/binary_reduce.py
@@ -17,9 +17,8 @@
"""Implementation of BinaryReduce dispatch."""
-from tvm_tirx.ir import TensorRegion
+from tvm_tirx.ir import Stmt, TensorRegion
from tvm_tirx.script import tirx as T
-from tvm_tirx.tirx import Function
from tvm_tirx.tirx.tile import DispatchContext
from tvm_tirx.tirx.tile.instruction import TensorCall
from tvm_tirx.tirx.tile.views import BinaryReduce
@@ -32,7 +31,7 @@ from ..reduction.utils import generate_intermediate_tensor
from .utils import opcode_table
-def binary_reduce_trn(op: TensorCall, sctx: DispatchContext) -> Function |
None:
+def binary_reduce_trn(op: TensorCall, sctx: DispatchContext) -> Stmt:
"""Generate a TRN schedule for binary reduction operations."""
op = TensorCall.decode(op)
assert isinstance(op, BinaryReduce), f"invalid operator downcast: {op}"
@@ -151,4 +150,4 @@ def binary_reduce_trn(op: TensorCall, sctx:
DispatchContext) -> Function | None:
T.nki.tensorreduce(dst2[tuple(dst_2_indices)],
intermediate_tensor[p_loop, f_loop], reduce_opcode, False, -1) # noqa: E501
# fmt: on
- return impl
+ return impl.body
diff --git a/tirx/python/tvm_tirx/backend/trn/tile/compose_op/unary_reduce.py
b/tirx/python/tvm_tirx/backend/trn/tile/compose_op/unary_reduce.py
index b571a4d581..cd8c0c8acf 100644
--- a/tirx/python/tvm_tirx/backend/trn/tile/compose_op/unary_reduce.py
+++ b/tirx/python/tvm_tirx/backend/trn/tile/compose_op/unary_reduce.py
@@ -17,9 +17,8 @@
"""Implementation of UnaryReduce dispatch."""
-from tvm_tirx.ir import TensorRegion
+from tvm_tirx.ir import Stmt, TensorRegion
from tvm_tirx.script import tirx as T
-from tvm_tirx.tirx import Function
from tvm_tirx.tirx.tile import DispatchContext
from tvm_tirx.tirx.tile.instruction import TensorCall
from tvm_tirx.tirx.tile.views import UnaryReduce
@@ -33,7 +32,7 @@ from ..unary.utils import get_const_bias_tensor,
try_find_inst_unary
from .utils import opcode_table
-def unary_reduce_trn(op: TensorCall, sctx: DispatchContext) -> Function | None:
+def unary_reduce_trn(op: TensorCall, sctx: DispatchContext) -> Stmt:
"""Generate a TRN schedule for unary reduction operations."""
op = TensorCall.decode(op)
assert isinstance(op, UnaryReduce), f"invalid operator downcast: {op}"
@@ -129,7 +128,7 @@ def unary_reduce_trn(op: TensorCall, sctx: DispatchContext)
-> Function | None:
mod = tvm_tirx.IRModule({"main": impl})
mod = tvm_tirx.tirx.transform.StmtSimplify()(mod)
- return mod["main"]
+ return mod["main"].body
else:
# fmt: off
# This fragment captures tensors and indices from its insertion scope.
@@ -159,4 +158,4 @@ def unary_reduce_trn(op: TensorCall, sctx: DispatchContext)
-> Function | None:
T.evaluate(T.nki.tensorreduce(dst2[tuple(dst_2_indices)],
intermediate_tensor[p_loop, f_loop], reduce_opcode, False, -1)) # noqa: E501
# fmt: on
- return impl
+ return impl.body
diff --git a/tirx/python/tvm_tirx/backend/trn/tile/copy/default.py
b/tirx/python/tvm_tirx/backend/trn/tile/copy/default.py
index 58597b549f..b64e9ab500 100644
--- a/tirx/python/tvm_tirx/backend/trn/tile/copy/default.py
+++ b/tirx/python/tvm_tirx/backend/trn/tile/copy/default.py
@@ -20,9 +20,9 @@
from functools import cache
from tvm_tirx.backend.trn.layout import is_trainium_layout
+from tvm_tirx.ir import Stmt
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.script import tirx as T
-from tvm_tirx.tirx import Function
from tvm_tirx.tirx.tile import DispatchContext, fail
from tvm_tirx.tirx.tile.instruction import TensorCall
@@ -43,7 +43,7 @@ def _scope_trn_sbuf():
return MemoryScope("trn.sbuf")
-def copy_trn(op: TensorCall, sctx: DispatchContext) -> Function | None:
+def copy_trn(op: TensorCall, sctx: DispatchContext) -> Stmt:
"""Schedule copy operation between global and shared memory on CUDA."""
# Basic validation checks
if sctx.scope_kind != "thread":
@@ -148,4 +148,4 @@ def copy_trn(op: TensorCall, sctx: DispatchContext) ->
Function | None:
dst_indices =
T.meta_var(inst_gen.generate_indices(dst_region))
func(dst[tuple(dst_indices)],
src[tuple(src_indices)])
# fmt: on
- return impl
+ return impl.body
diff --git a/tirx/python/tvm_tirx/backend/trn/tile/gemm/default.py
b/tirx/python/tvm_tirx/backend/trn/tile/gemm/default.py
index 582aef30cb..d395c2b183 100644
--- a/tirx/python/tvm_tirx/backend/trn/tile/gemm/default.py
+++ b/tirx/python/tvm_tirx/backend/trn/tile/gemm/default.py
@@ -22,11 +22,10 @@ import operator
from functools import cache
from tvm_tirx.backend.trn.layout import is_trainium_layout
-from tvm_tirx.ir import TensorRegion, assert_structural_equal
+from tvm_tirx.ir import Stmt, TensorRegion, assert_structural_equal
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.script import tirx as T
from tvm_tirx.sym.analyzer import Analyzer
-from tvm_tirx.tirx import Function
from tvm_tirx.tirx.tile import DispatchContext, fail
from tvm_tirx.tirx.tile.instruction import TensorCall
@@ -117,7 +116,7 @@ def get_pf_dim_from_tensor_region(
return p_dim, f_dim
-def matmul_trn(op: TensorCall, sctx: DispatchContext) -> Function | None:
+def matmul_trn(op: TensorCall, sctx: DispatchContext) -> Stmt:
"""Schedule GEMM operation on Trainium."""
# Basic validation checks
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
@@ -254,7 +253,7 @@ def matmul_trn(op: TensorCall, sctx: DispatchContext) ->
Function | None:
def impl_C_psum():
for lhs_b_loop, rhs_b_loop, reduction_b_loop in
T.grid(lhs_b_extent, rhs_b_extent, reduction_b_extent): # noqa: E501
matmul_inst_macro(lhs_b_loop, rhs_b_loop, reduction_b_loop, C,
True, None)
- return impl_C_psum
+ return impl_C_psum.body
# todo: generalize the process of generating composite matmul + another_op
pattern
# by generating TIR op and reusing existing dispatch rule
@@ -292,4 +291,4 @@ def matmul_trn(op: TensorCall, sctx: DispatchContext) ->
Function | None:
acc_indices =
T.meta_var(inst_gen.generate_indices(C_tensor_region))
T.evaluate(T.nki.tensor_copy(C[acc_indices],
acc_psum[(lhs_b_loop * rhs_b_extent + rhs_b_loop) % max_psum_slots, lhs_f_loop,
rhs_f_loop])) # noqa: E501
# fmt: on
- return impl_C_sbuf
+ return impl_C_sbuf.body
diff --git a/tirx/python/tvm_tirx/backend/trn/tile/reduction/utils.py
b/tirx/python/tvm_tirx/backend/trn/tile/reduction/utils.py
index d6ec5d1cbb..62530eb030 100644
--- a/tirx/python/tvm_tirx/backend/trn/tile/reduction/utils.py
+++ b/tirx/python/tvm_tirx/backend/trn/tile/reduction/utils.py
@@ -20,9 +20,9 @@
from functools import cache
from tvm_tirx.backend.trn.layout import is_trainium_layout
+from tvm_tirx.ir import Stmt
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.script import tirx as T
-from tvm_tirx.tirx import Function
from tvm_tirx.tirx.tile import DispatchContext, fail
from tvm_tirx.tirx.tile.common import ReduceOpType
from tvm_tirx.tirx.tile.instruction import TensorCall
@@ -74,7 +74,7 @@ def generate_intermediate_tensor(
def reduction_trn(
op: TensorCall, reduce_op: ReduceOpType, sctx: DispatchContext, negate:
bool = False
-) -> Function | None:
+) -> Stmt:
"""Schedule reduction operation on Trainium.
Args:
@@ -84,7 +84,7 @@ def reduction_trn(
negate: Whether to negate the result.
Returns:
- Optional[Function]: The scheduled function, or None if not applicable.
+ Stmt: The implementation statement.
"""
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
fail("requires Trainium target and thread scope")
@@ -156,7 +156,7 @@ def reduction_trn(
src_indices =
T.meta_var(inst_gen.generate_indices(src_tensor_region)) # noqa: E501
dst_indices =
T.meta_var(inst_gen.generate_indices(dst_tensor_region)) # noqa: E501
T.evaluate(T.nki.tensorreduce(dst[tuple(dst_indices)], src[tuple(src_indices)],
opcode, negate, -1)) # noqa: E501
- return impl
+ return impl.body
# Two-stage reduction implementation
else:
# This fragment captures tensors and indices from its insertion scope.
@@ -179,5 +179,5 @@ def reduction_trn(
if inst_gen.make_guard(src_tensor_region):
dst_indices =
T.meta_var(inst_gen.generate_indices(dst_tensor_region)) # noqa: E501
T.evaluate(T.nki.tensorreduce(dst[dst_indices], intermediate_tensor[p_loop,
f_loop], opcode, negate, -1)) # noqa: E501
- return two_stage_reduction
+ return two_stage_reduction.body
# fmt: on
diff --git a/tirx/python/tvm_tirx/backend/trn/tile/select/default.py
b/tirx/python/tvm_tirx/backend/trn/tile/select/default.py
index e3e738d925..2b9e9ec5dc 100644
--- a/tirx/python/tvm_tirx/backend/trn/tile/select/default.py
+++ b/tirx/python/tvm_tirx/backend/trn/tile/select/default.py
@@ -20,10 +20,9 @@
from functools import cache
from tvm_tirx.backend.trn.layout import is_trainium_layout
-from tvm_tirx.ir import FloatImm, TensorRegion
+from tvm_tirx.ir import FloatImm, Stmt, TensorRegion
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.script import tirx as T
-from tvm_tirx.tirx import Function
from tvm_tirx.tirx.tile import DispatchContext, fail
from tvm_tirx.tirx.tile.instruction import TensorCall
from tvm_tirx.tirx.tile.views import Select
@@ -38,7 +37,7 @@ def _scope_trn_sbuf():
return MemoryScope("trn.sbuf")
-def select_trn(op: TensorCall, sctx: DispatchContext) -> Function | None:
+def select_trn(op: TensorCall, sctx: DispatchContext) -> Stmt:
"""Generate schedule for select operation on Trainium."""
if sctx.scope_kind != "thread":
fail("requires thread scope for TRN select")
@@ -127,4 +126,4 @@ def select_trn(op: TensorCall, sctx: DispatchContext) ->
Function | None:
T.evaluate(T.nki.affine_select(dst_tensor[tuple(dst_indices)],
analyzer.simplify(op.predicate.apply(inst_gen.generate_axes(dst))),
true_value_tensor[tuple(true_value_indices)], false_value)) # noqa: E501
# fmt: on
- return impl
+ return impl.body
diff --git a/tirx/python/tvm_tirx/backend/trn/tile/unary/default.py
b/tirx/python/tvm_tirx/backend/trn/tile/unary/default.py
index f921915e6f..961035b404 100644
--- a/tirx/python/tvm_tirx/backend/trn/tile/unary/default.py
+++ b/tirx/python/tvm_tirx/backend/trn/tile/unary/default.py
@@ -17,8 +17,7 @@
"""Implementation of default unary operator dispatches."""
-from tvm_tirx.ir import FloatImm
-from tvm_tirx.tirx import Function
+from tvm_tirx.ir import FloatImm, Stmt
from tvm_tirx.tirx.tile import DispatchContext, fail
from tvm_tirx.tirx.tile.common import MapOpType
from tvm_tirx.tirx.tile.instruction import TensorCall
@@ -27,13 +26,13 @@ from ..common import init_analyzer
from ..instruction_generator import InstructionGenerator
from .utils import (
const_input_ops,
- generate_unary_func,
+ generate_unary_stmt,
non_activation_unary_map_ops,
try_find_inst_unary,
)
-def unary_trn(op: TensorCall, unary_op: MapOpType, sctx: DispatchContext) ->
Function | None:
+def unary_trn(op: TensorCall, unary_op: MapOpType, sctx: DispatchContext) ->
Stmt:
"""Schedule unary operation on Trainium."""
# Check execution environment
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
@@ -62,8 +61,8 @@ def unary_trn(op: TensorCall, unary_op: MapOpType, sctx:
DispatchContext) -> Fun
inst_repr = try_find_inst_unary(dst_tensor_region, src_tensor_region,
analyzer, inst_gen)
else:
inst_repr = try_find_inst_unary(dst_tensor_region, dst_tensor_region,
analyzer, inst_gen)
- # Generate and return the implementation function
- return generate_unary_func(
+ # Generate and return the implementation statement
+ return generate_unary_stmt(
dst_tensor_region,
_src,
inst_gen,
diff --git a/tirx/python/tvm_tirx/backend/trn/tile/unary/utils.py
b/tirx/python/tvm_tirx/backend/trn/tile/unary/utils.py
index dfd904691e..befba4d67d 100644
--- a/tirx/python/tvm_tirx/backend/trn/tile/unary/utils.py
+++ b/tirx/python/tvm_tirx/backend/trn/tile/unary/utils.py
@@ -20,7 +20,7 @@
from functools import cache
from tvm_tirx.backend.trn.layout import is_trainium_layout
-from tvm_tirx.ir import FloatImm, TensorRegion
+from tvm_tirx.ir import FloatImm, Stmt, TensorRegion
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.script import tirx as T
from tvm_tirx.sym.analyzer import Analyzer
@@ -136,7 +136,7 @@ def get_const_bias_tensor(bias, shape, dtype, workspace,
sctx):
return bias_tensor
-def generate_unary_func(
+def generate_unary_stmt(
dst_tensor_region,
_src,
inst_gen: InstructionGenerator,
@@ -148,8 +148,8 @@ def generate_unary_func(
workspace,
config,
sctx,
-):
- """Generate a function that implements a unary operation."""
+) -> Stmt:
+ """Generate a statement that implements a unary operation."""
# Prepare parameters
p_size = dst_tensor_region.source.ty.effective_layout().size("P")
@@ -203,4 +203,4 @@ def generate_unary_func(
T.evaluate(T.nki.activation(dst[tuple(dst_indices)], src[tuple(src_indices)],
opcode, scale=scale, bias=bias_tensor[p_loop, f_loop])) # noqa: E501
# fmt: on
- return impl
+ return impl.body
diff --git a/tirx/python/tvm_tirx/backend/trn/tile/unary/with_bias_scale.py
b/tirx/python/tvm_tirx/backend/trn/tile/unary/with_bias_scale.py
index c4f53d40c5..6eb6869ff9 100644
--- a/tirx/python/tvm_tirx/backend/trn/tile/unary/with_bias_scale.py
+++ b/tirx/python/tvm_tirx/backend/trn/tile/unary/with_bias_scale.py
@@ -17,8 +17,7 @@
"""Implementation of unary with bias and scale operator dispatches."""
-from tvm_tirx.ir import TensorRegion
-from tvm_tirx.tirx import Function
+from tvm_tirx.ir import Stmt, TensorRegion
from tvm_tirx.tirx.tile import DispatchContext, fail
from tvm_tirx.tirx.tile.common import MapOpType
from tvm_tirx.tirx.tile.instruction import TensorCall
@@ -26,12 +25,12 @@ from tvm_tirx.tirx.tile.instruction import TensorCall
from ..binary import try_find_inst_nary
from ..common import init_analyzer
from ..instruction_generator import InstructionGenerator
-from .utils import activation_map_ops, generate_unary_func, try_find_inst_unary
+from .utils import activation_map_ops, generate_unary_stmt, try_find_inst_unary
def unary_with_bias_scale_trn(
op: TensorCall, unary_op: MapOpType = MapOpType.SQRT, sctx:
DispatchContext = None
-) -> Function | None:
+) -> Stmt:
"""Schedule unary operation with bias and scale on Trainium."""
# Check execution environment
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
@@ -61,8 +60,8 @@ def unary_with_bias_scale_trn(
# Handle scalar bias
inst_repr = try_find_inst_unary(dst_tensor_region, src_tensor_region,
analyzer, inst_gen)
- # Generate and return the implementation function
- return generate_unary_func(
+ # Generate and return the implementation statement
+ return generate_unary_stmt(
dst_tensor_region,
src_tensor_region,
inst_gen,
diff --git a/tirx/python/tvm_tirx/tirx/tile/dispatch.py
b/tirx/python/tvm_tirx/tirx/tile/dispatch.py
index d5cb237e05..8da1f65e87 100644
--- a/tirx/python/tvm_tirx/tirx/tile/dispatch.py
+++ b/tirx/python/tvm_tirx/tirx/tile/dispatch.py
@@ -230,7 +230,8 @@ def fail(reason):
raise DispatchFail(reason)
-def run_dispatch(call, sctx):
+def run_dispatch(call, sctx) -> Stmt:
+ """Lower a tensor instruction to its replacement statement."""
from tvm_tirx.tirx.tile.instruction import TensorCall
view = TensorCall.decode(call)
@@ -243,6 +244,8 @@ def run_dispatch(call, sctx):
result = view.spec.lower(view, sctx)
if result is None:
raise DispatchFail("lowerer returned no implementation")
+ if not isinstance(result, Stmt):
+ raise TypeError(f"lowerer must return a Stmt, got
{type(result).__name__}")
return result
except Exception as error:
raise RuntimeError(
diff --git a/tirx/src/tirx/tile/dispatch.cc b/tirx/src/tirx/tile/dispatch.cc
index 4cf45ea60f..5dd1a0ddfa 100644
--- a/tirx/src/tirx/tile/dispatch.cc
+++ b/tirx/src/tirx/tile/dispatch.cc
@@ -20,15 +20,15 @@
namespace tvm::tirx::tile {
-Function Dispatch(const Call& call, const DispatchContext& context) {
+Stmt Dispatch(const Call& call, const DispatchContext& context) {
Op op = call->op.as_or_throw<Op>();
TVM_FFI_CHECK(Op::HasAttrMap(op_attr::kTileOpDispatch), ValueError)
<< "No FTileOpDispatch registered for " << op->name;
static const auto& dispatchers = Op::GetAttrMap(op_attr::kTileOpDispatch);
TVM_FFI_CHECK(dispatchers.count(op), ValueError)
<< "No FTileOpDispatch registered for " << op->name;
- Function implementation = dispatchers[op](call, context);
- TVM_FFI_ICHECK(implementation.defined()) << "TIRx dispatcher did not return
a Function";
+ Stmt implementation = dispatchers[op](call, context);
+ TVM_FFI_ICHECK(implementation.defined()) << "TIRx dispatcher did not return
a Stmt";
return implementation;
}
diff --git a/tirx/src/tirx/tile/lower.cc b/tirx/src/tirx/tile/lower.cc
index 3e133d1525..bf44889803 100644
--- a/tirx/src/tirx/tile/lower.cc
+++ b/tirx/src/tirx/tile/lower.cc
@@ -501,7 +501,7 @@ class TileDispatcher : public StmtExprMutator {
tile::DispatchContext sctx(target_, scope, launch_params_, var_range_map_,
/*alloc_only=*/false, /*callbacks=*/{},
shared_state_, inter_map,
intra_map, scope_kind);
- Function res = tile::Dispatch(ffi::GetRef<Call>(op), sctx);
+ Stmt res = tile::Dispatch(ffi::GetRef<Call>(op), sctx);
// Implementation found, handle callbacks
if (auto bufs = sctx->callbacks.Get(tirx::tile::callback::kPrivateAlloc)) {
auto buf_list = bufs.value().as<Array<Bind>>().value();
@@ -524,8 +524,7 @@ class TileDispatcher : public StmtExprMutator {
}
// Propagate shared_state changes back (Map uses COW semantics)
shared_state_ = sctx->shared_state;
- TVM_FFI_CHECK(res->body.has_value(), ValueError) << "A tile op
implementation must have a body";
- return Mutate(res->body.value(),
inplace_mode).ValueOrUnchanged(res->body.value());
+ return Mutate(res, inplace_mode).ValueOrUnchanged(res);
}
// --- Scope-id resolution at kernel scope ----------------------------------
diff --git a/tirx/tests/cpp/tirx/thread_scope_test.cc
b/tirx/tests/cpp/tirx/thread_scope_test.cc
index c6530ec9e2..c3adc3959a 100644
--- a/tirx/tests/cpp/tirx/thread_scope_test.cc
+++ b/tirx/tests/cpp/tirx/thread_scope_test.cc
@@ -89,7 +89,7 @@ TEST(ThreadScope, TypedCallAndDispatchWithoutPython) {
.cast<ThreadScope>()
.same_as(scope));
++*invoked;
- return Function({}, SeqStmt(Evaluate(37)));
+ return Evaluate(37);
});
Call call(PrimType::Void(), Op::Get("thread_scope_test.instruction"), {},
attrs);
Function func({}, SeqStmt(Evaluate(call)), VoidType(), DictAttrs({{"target",
target}}));
diff --git a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
index 1265c371ec..654eccaf9c 100644
--- a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
+++ b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
@@ -78,7 +78,7 @@ def _count_s2c_ops(impl):
visitor.default_visit(op)
tvm_ffi.structural_visit(
- impl.body,
+ impl,
[(tvm_tirx.ir.For, visit_for), (tvm_tirx.ir.Evaluate, visit_evaluate)],
)
return total
diff --git a/tirx/tests/python/tirx/tile/test_dispatcher.py
b/tirx/tests/python/tirx/tile/test_dispatcher.py
index eeb9000aec..bdf162ff03 100644
--- a/tirx/tests/python/tirx/tile/test_dispatcher.py
+++ b/tirx/tests/python/tirx/tile/test_dispatcher.py
@@ -20,7 +20,17 @@ from types import SimpleNamespace
import pytest
import tvm_tirx
-from tvm_tirx.ir import Call, Evaluate, Op, PrimType, Return, const,
make_node, register_op_attr
+from tvm_tirx.ir import (
+ Call,
+ Evaluate,
+ Op,
+ PrimType,
+ Return,
+ SeqStmt,
+ const,
+ make_node,
+ register_op_attr,
+)
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx import Function, decl_tensor
from tvm_tirx.tirx.tile import DispatchContext
@@ -118,7 +128,7 @@ def
test_op_attribute_dispatches_without_python_instruction_spec():
assert context.thread_scope.backend == "cuda"
assert context.scope_kind == "thread"
seen.append(actual_call)
- return Function([], Evaluate(23))
+ return Evaluate(23)
lowered = _lower(Evaluate(call))
assert len(seen) == 1
@@ -128,7 +138,7 @@ def
test_op_attribute_dispatches_without_python_instruction_spec():
register_op_attr(
call.op.name,
"FTileOpDispatch",
- lambda actual_call, context: Function([], Evaluate(41)),
+ lambda actual_call, context: Evaluate(41),
override=True,
)
assert _lower(Evaluate(call)).body.seq[0].value == 41
@@ -138,12 +148,12 @@ def
test_dispatch_recursively_expands_tile_implementations():
outer = _custom_call("outer")
inner = _custom_call("inner")
register_op_attr(
- outer.op.name, "FTileOpDispatch", lambda call, context: Function([],
Evaluate(inner))
- )
- register_op_attr(
- inner.op.name, "FTileOpDispatch", lambda call, context: Function([],
Evaluate(31))
+ outer.op.name,
+ "FTileOpDispatch",
+ lambda call, context: SeqStmt([Evaluate(inner), Evaluate(7)]),
)
- assert _lower(Evaluate(outer)).body.seq[0].value == 31
+ register_op_attr(inner.op.name, "FTileOpDispatch", lambda call, context:
Evaluate(31))
+ assert [stmt.value for stmt in _lower(Evaluate(outer)).body.seq] == [31, 7]
def test_op_without_dispatch_attribute_is_preserved():
@@ -154,16 +164,29 @@ def test_op_without_dispatch_attribute_is_preserved():
def test_residual_tile_call_is_rejected():
call = _custom_call("residual")
- register_op_attr(
- call.op.name, "FTileOpDispatch", lambda call, context: Function([],
Evaluate(0))
- )
+ register_op_attr(call.op.name, "FTileOpDispatch", lambda call, context:
Evaluate(0))
# Dispatch expands statement-boundary Calls; other positions must be
rejected.
with pytest.raises(RuntimeError, match="Failed to lower the TIRx program"):
_lower(Return(call))
-def test_dispatch_requires_function_result():
+def test_dispatch_requires_stmt_result():
call = _custom_call("invalid_result")
- register_op_attr(call.op.name, "FTileOpDispatch", lambda call, context:
Evaluate(0))
- with pytest.raises(TypeError, match="tirx.Function"):
+ register_op_attr(
+ call.op.name, "FTileOpDispatch", lambda call, context: Function([],
Evaluate(0))
+ )
+ with pytest.raises(TypeError, match="ir.Stmt"):
_lower(Evaluate(call))
+
+
[email protected]("result", [Function([], Evaluate(0)), const(0), None])
+def test_selected_lowerer_requires_stmt_result(monkeypatch, result):
+ tensor = decl_tensor((16,), "float32", scope="local", layout="default")
+ call = T.cuda.tile.apply(T.ptx["sqrt.rn.f32"], tensor, tensor)
+ monkeypatch.setattr(SPECS[call.op.name], "lower", lambda view, context:
result)
+ context = SimpleNamespace(target=tvm_tirx.target.Target("cuda"),
scope_kind="thread")
+ message = (
+ "lowerer returned no implementation" if result is None else "lowerer
must return a Stmt"
+ )
+ with pytest.raises(RuntimeError, match=message):
+ run_dispatch(call, context)
diff --git a/tirx/tests/python/tirx/transform/test_lower_tirx.py
b/tirx/tests/python/tirx/transform/test_lower_tirx.py
index d27649f222..c81cd2c662 100644
--- a/tirx/tests/python/tirx/transform/test_lower_tirx.py
+++ b/tirx/tests/python/tirx/transform/test_lower_tirx.py
@@ -81,9 +81,7 @@ def _int_range(side, axis):
return tuple(int(x) for x in side[axis])
[email protected](private=True)
-def _probe_impl():
- T.evaluate(0)
+_PROBE_STMT = tvm_tirx.ir.Evaluate(0)
L_LANE = T.TileLayout(T.S[32 : 1 @ laneid])
@@ -1018,7 +1016,7 @@ def
test_lower_exec_context_infers_plain_predicate_for_dispatch():
def _probe(op_call, sctx):
seen.append({"scope_kind": sctx.scope_kind, "inter": sctx.inter,
"intra": sctx.intra})
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1068,7 +1066,7 @@ def
test_lower_exec_context_tracks_let_aliases_and_hoists_warp_shuffle():
def _probe(op_call, sctx):
seen.append(sctx.inter)
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1123,7 +1121,7 @@ def
test_lower_exec_context_respects_preferred_cluster_bounds():
def _probe(op_call, sctx):
seen.append(sctx.inter)
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1191,7 +1189,7 @@ def
test_lower_exec_context_infers_warpgroup_range_predicate_for_dispatch():
def _probe(op_call, sctx):
seen.append({"scope_kind": sctx.scope_kind, "inter": sctx.inter,
"intra": sctx.intra})
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1250,7 +1248,7 @@ def
test_lower_exec_context_tracks_cta_thread_range_predicate_for_dispatch():
def _probe(op_call, sctx):
seen.append({"scope_kind": sctx.scope_kind, "inter": sctx.inter,
"intra": sctx.intra})
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1295,7 +1293,7 @@ def
test_lower_exec_context_tracks_cta_thread_single_warp_range_predicate():
def _probe(op_call, sctx):
seen.append({"scope_kind": sctx.scope_kind, "inter": sctx.inter,
"intra": sctx.intra})
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1340,7 +1338,7 @@ def
test_lower_exec_context_tracks_warpgroup_thread_range_predicate():
def _probe(op_call, sctx):
seen.append({"scope_kind": sctx.scope_kind, "inter": sctx.inter,
"intra": sctx.intra})
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1392,7 +1390,7 @@ def
test_lower_exec_context_tracks_dependent_conjunctive_predicate():
def _probe(op_call, sctx):
seen.append({"scope_kind": sctx.scope_kind, "inter": sctx.inter,
"intra": sctx.intra})
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1567,7 +1565,7 @@ def
test_lower_exec_context_selector_filter_for_elect_sync():
def _probe(op_call, sctx):
seen.append(sctx.inter["laneid"][1].script(extra_config={"tirx.prefix": "T"}))
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1655,7 +1653,7 @@ def
test_lower_exec_context_scope_guard_mixes_structural_and_selector():
def _probe(op_call, sctx):
seen.append({"inter": sctx.inter, "intra": sctx.intra})
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1707,7 +1705,7 @@ def
test_lower_exec_context_tracks_factorized_cta_predicate():
def _probe(op_call, sctx):
seen.append(sctx.inter)
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1751,13 +1749,13 @@ def
test_lower_exec_context_keeps_kernel_cta_predicate_out_of_cluster_active_set
def _probe_kernel(op_call, sctx):
seen["kernel"] = sctx.inter
- return _probe_impl
+ return _PROBE_STMT
@_register_probe(cluster_variant)
def _probe_cluster(op_call, sctx):
seen["cluster"] = sctx.inter
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1805,7 +1803,7 @@ def
test_lower_exec_context_tracks_cta_axis_modulo_predicate():
def _probe(op_call, sctx):
seen.append(sctx.inter)
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1848,7 +1846,7 @@ def
test_lower_exec_context_tracks_cta_id_in_pair_predicate():
def _probe(op_call, sctx):
seen.append(sctx.inter)
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1902,13 +1900,13 @@ def
test_lower_exec_context_tracks_two_cta_pair_predicates():
def _probe_zero(op_call, sctx):
seen["zero"] = sctx.inter
- return _probe_impl
+ return _PROBE_STMT
@_register_probe(one_variant)
def _probe_one(op_call, sctx):
seen["one"] = sctx.inter
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
@@ -1962,7 +1960,7 @@ def
test_lower_exec_context_tracks_cta_id_in_pair_after_axis_predicate():
def _probe(op_call, sctx):
seen.append(sctx.inter)
- return _probe_impl
+ return _PROBE_STMT
@I.ir_module
class Before:
