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 0defab0b3f86bef90c957e176faac0415db8d7d1 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. --- 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 ++++++++--------- 20 files changed, 117 insertions(+), 103 deletions(-) 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:
