You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/09 19:28:36 UTC
(tvm) branch main updated: refactor(op)!: replace tile statements with tensor instruction calls (#20632)
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 c18a389e22 refactor(op)!: replace tile statements with tensor instruction calls (#20632)
c18a389e22 is described below
commit c18a389e22ad1945958be0de5028f4230f5cdfbd
Author: Bohan Hou <sp...@gmail.com>
AuthorDate: Fri Oct 9 15:28:28 2026 -0400
refactor(op)!: replace tile statements with tensor instruction calls (#20632)
---
docs/tirx/api/ptx.rst | 2 +-
docs/tirx/api/script.rst | 2 +-
docs/tirx/api/tile.rst | 166 +-
docs/tirx/api/tile_dispatch.rst | 63 +-
docs/tirx/arch/backends.rst | 4 +-
docs/tirx/arch/lowering_pipeline.rst | 5 +-
docs/tirx/arch/tile_dispatch.rst | 124 +-
docs/tirx/layout.rst | 2 +-
docs/tirx/overview.rst | 10 +-
docs/tirx/programming/index.rst | 2 +-
docs/tirx/tile_primitives.rst | 287 ++-
docs/tirx/tile_primitives/copy.rst | 56 +-
docs/tirx/tile_primitives/copy/fallback.rst | 171 --
docs/tirx/tile_primitives/copy/gmem_smem.rst | 211 ---
docs/tirx/tile_primitives/copy/ldstmatrix.rst | 22 +-
docs/tirx/tile_primitives/copy/reg.rst | 28 +-
docs/tirx/tile_primitives/copy_async.rst | 58 +-
docs/tirx/tile_primitives/copy_async/dsmem.rst | 46 +-
docs/tirx/tile_primitives/copy_async/ldgsts.rst | 35 +-
.../tirx/tile_primitives/copy_async/tcgen05_cp.rst | 34 +-
.../tile_primitives/copy_async/tcgen05_ldst.rst | 55 +-
docs/tirx/tile_primitives/copy_async/tma.rst | 45 +-
docs/tirx/tile_primitives/elementwise.rst | 51 +-
docs/tirx/tile_primitives/elementwise/reg.rst | 27 +-
docs/tirx/tile_primitives/elementwise/smem.rst | 30 +-
docs/tirx/tile_primitives/gemm.rst | 31 +-
docs/tirx/tile_primitives/gemm_async.rst | 35 +-
docs/tirx/tile_primitives/permute_layout.rst | 179 --
docs/tirx/tile_primitives/reduction.rst | 41 +-
docs/tirx/tile_primitives/reduction/local.rst | 157 --
docs/tirx/tile_primitives/reduction/shared.rst | 143 --
.../tile_primitives/reduction/sm100_packed.rst | 157 --
include/tvm/tirx/exec_scope.h | 2 +-
include/tvm/tirx/op/tile.h | 393 -----
include/tvm/tirx/script/ir_builder/ir.h | 2 -
include/tvm/tirx/stmt.h | 67 -
include/tvm/tirx/stmt_functor.h | 6 -
python/tvm/backend/cuda/script.py | 4 +-
python/tvm/backend/cuda/tensor_instructions.py | 407 +++++
python/tvm/backend/cuda/tile_primitive/__init__.py | 8 -
python/tvm/backend/cuda/tile_primitive/common.py | 10 +-
.../backend/cuda/tile_primitive/copy/__init__.py | 12 +-
.../backend/cuda/tile_primitive/copy/_common.py | 2 +-
.../backend/cuda/tile_primitive/copy/fallback.py | 137 --
.../cuda/tile_primitive/copy/ld_stmatrix.py | 21 +-
.../tvm/backend/cuda/tile_primitive/copy/utils.py | 12 +-
.../backend/cuda/tile_primitive/copy/vec_auto.py | 53 -
.../cuda/tile_primitive/copy/vec_auto_gmem_smem.py | 184 --
.../cuda/tile_primitive/copy/vec_auto_reg.py | 22 +-
.../backend/cuda/tile_primitive/copy/vec_forced.py | 57 +-
.../cuda/tile_primitive/copy_async/__init__.py | 7 +-
.../cuda/tile_primitive/copy_async/dsmem.py | 67 +-
.../cuda/tile_primitive/copy_async/ldgsts.py | 35 +-
.../cuda/tile_primitive/copy_async/tcgen05_cp.py | 45 +-
.../cuda/tile_primitive/copy_async/tcgen05_ldst.py | 35 +-
.../backend/cuda/tile_primitive/copy_async/tma.py | 124 +-
.../cuda/tile_primitive/elementwise/__init__.py | 9 -
.../tile_primitive/elementwise/ops/__init__.py | 6 +-
.../cuda/tile_primitive/elementwise/ops/binary.py | 14 +-
.../cuda/tile_primitive/elementwise/ops/cast.py | 4 +-
.../cuda/tile_primitive/elementwise/ops/fma.py | 4 +-
.../cuda/tile_primitive/elementwise/ops/unary.py | 4 +-
.../backend/cuda/tile_primitive/elementwise/reg.py | 7 +-
.../cuda/tile_primitive/elementwise/register.py | 61 -
.../cuda/tile_primitive/elementwise/smem.py | 7 +-
.../cuda/tile_primitive/exec_scope_utils.py | 6 +-
.../backend/cuda/tile_primitive/gemm/__init__.py | 7 +-
.../cuda/tile_primitive/gemm/mma_m16n8k_.py | 25 +-
.../cuda/tile_primitive/gemm_async/tcgen05.py | 79 +-
.../tvm/backend/cuda/tile_primitive/gemm_utils.py | 4 +-
.../cuda/tile_primitive/permute_layout/__init__.py | 18 -
.../permute_layout/warp_xor_swizzle.py | 436 -----
.../cuda/tile_primitive/reduction/__init__.py | 20 -
.../backend/cuda/tile_primitive/reduction/local.py | 487 ------
.../cuda/tile_primitive/reduction/shared.py | 295 ----
.../cuda/tile_primitive/reduction/sm100_packed.py | 255 ---
.../backend/cuda/tile_primitive/reduction/utils.py | 262 ---
python/tvm/backend/trn/__init__.py | 7 +-
python/tvm/backend/trn/tensor_instructions.py | 307 ++++
python/tvm/backend/trn/tile_primitive/__init__.py | 9 -
.../backend/trn/tile_primitive/binary/default.py | 29 +-
.../trn/tile_primitive/compose_op/__init__.py | 1 -
.../trn/tile_primitive/compose_op/binary_chain.py | 38 +-
.../trn/tile_primitive/compose_op/binary_reduce.py | 33 +-
.../trn/tile_primitive/compose_op/reduce_negate.py | 51 -
.../trn/tile_primitive/compose_op/unary_reduce.py | 51 +-
.../backend/trn/tile_primitive/compose_op/utils.py | 27 +-
.../tvm/backend/trn/tile_primitive/copy/default.py | 195 +--
.../tvm/backend/trn/tile_primitive/gemm/default.py | 35 +-
.../backend/trn/tile_primitive/private_alloc.py | 83 +-
.../trn/tile_primitive/reduction/__init__.py | 2 -
.../trn/tile_primitive/reduction/default.py | 33 -
.../backend/trn/tile_primitive/reduction/utils.py | 10 +-
.../backend/trn/tile_primitive/select/default.py | 36 +-
.../backend/trn/tile_primitive/unary/default.py | 20 +-
.../trn/tile_primitive/unary/with_bias_scale.py | 23 +-
.../backend/trn/transform/private_buffer_alloc.py | 26 +-
python/tvm/script/__init__.py | 4 +-
python/tvm/tirx/__init__.py | 1 -
python/tvm/tirx/op/tile.py | 136 +-
.../tvm/tirx/operator/tile_primitive/__init__.py | 15 +-
.../tvm/tirx/operator/tile_primitive/dispatcher.py | 319 +---
.../tvm/tirx/operator/tile_primitive/registry.py | 52 +-
python/tvm/tirx/script/__init__.py | 7 -
python/tvm/tirx/script/ir_builder/__init__.py | 26 -
python/tvm/tirx/script/ir_builder/tirx.py | 1841 --------------------
python/tvm/tirx/script/tile.py | 131 --
python/tvm/tirx/stmt.py | 182 +-
python/tvm/tirx/tensor_instruction.py | 378 ++++
python/tvm/tirx/tile_dispatch.py | 2 +-
python/tvm/tirx/transform/common.py | 21 +-
src/backend/trn/transform/lower_trainium_layout.cc | 16 -
src/script/printer/script_printer.cc | 1 -
src/tirx/analysis/verify_well_formed.h | 19 -
src/tirx/ir/stmt_functor.cc | 31 -
src/tirx/ir/tir_visitor_with_path.cc | 20 -
src/tirx/ir/tir_visitor_with_path.h | 1 -
src/tirx/ir/tirx_stmt.cc | 166 --
src/tirx/op/tile.cc | 676 +++----
src/tirx/script/ir_builder/ir.cc | 3 -
src/tirx/script/printer/stmt.cc | 129 +-
src/tirx/transform/lower_tirx_cleanup.cc | 17 -
src/tirx/transform/tile_primitive_dispatch.cc | 41 +-
.../python/tirx/codegen/test_codegen_blackwell.py | 102 +-
tests/python/tirx/codegen/test_codegen_cuda.py | 2 +-
tests/python/tirx/codegen/test_ptx_ld_st_ops.py | 3 +-
.../tile_primitive/cuda/copy/test_fallback.py | 239 ---
.../tile_primitive/cuda/copy/test_gmem_smem.py | 370 +++-
.../tile_primitive/cuda/copy/test_ld_stmatrix.py | 17 +-
.../operator/tile_primitive/cuda/copy/test_reg.py | 115 +-
.../cuda/copy/test_vec_forced_cache.py | 25 +-
.../tile_primitive/cuda/copy_async/test_dsmem.py | 97 +-
.../tile_primitive/cuda/copy_async/test_ldgsts.py | 41 +-
.../cuda/copy_async/test_tcgen05_cp.py | 220 ++-
.../cuda/copy_async/test_tcgen05_ldst.py | 114 +-
.../tile_primitive/cuda/copy_async/test_tma.py | 74 +-
.../tile_primitive/cuda/elementwise/test_binary.py | 621 ++++++-
.../tile_primitive/cuda/elementwise/test_fma.py | 51 +-
.../tile_primitive/cuda/elementwise/test_unary.py | 656 ++++++-
.../cuda/gemm/test_gemm_mma_m16n8k_.py | 72 +-
.../cuda/gemm_async/test_gemm_async.py | 641 ++++---
.../cuda/permute_layout/test_permute_layout.py | 482 -----
.../cuda/reduction/test_reduction.py | 1149 ------------
.../operator/tile_primitive/test_dispatcher.py | 178 +-
.../operator/tile_primitive/trn/test_binary_trn.py | 58 +-
.../tile_primitive/trn/test_compose_op_trn.py | 53 +-
.../operator/tile_primitive/trn/test_copy_trn.py | 331 +---
.../operator/tile_primitive/trn/test_gemm_trn.py | 45 +-
.../tile_primitive/trn/test_private_alloc_trn.py | 87 +-
.../tile_primitive/trn/test_reduction_trn.py | 18 +-
.../operator/tile_primitive/trn/test_select_trn.py | 9 +-
.../operator/tile_primitive/trn/test_unary_trn.py | 25 +-
.../tirx/script/test_tirx_script_basic_usage.py | 33 +-
.../python/tirx/script/test_tirx_script_buffers.py | 7 +-
tests/python/tirx/script/test_tirx_script_gpu.py | 151 +-
.../python/tirx/script/test_tirx_script_printer.py | 5 +-
.../script/test_tirx_script_source_locations.py | 17 +-
tests/python/tirx/test_hint.py | 67 -
tests/python/tirx/test_jit.py | 15 +-
tests/python/tirx/test_op.py | 235 +--
tests/python/tirx/test_op_namespace_cleanup.py | 70 +-
tests/python/tirx/test_verifier.py | 5 +-
.../tirx/transform/test_transform_lower_tirx.py | 129 +-
.../transform/test_transform_naive_allocator.py | 25 +-
164 files changed, 5446 insertions(+), 12460 deletions(-)
diff --git a/docs/tirx/api/ptx.rst b/docs/tirx/api/ptx.rst
index de63107acc..9bc3f129f2 100644
--- a/docs/tirx/api/ptx.rst
+++ b/docs/tirx/api/ptx.rst
@@ -31,7 +31,7 @@ modifier combination, operand count, and operand types while constructing one
``tirx.ptx.*`` IR call. A new instruction is added by extending the table and,
when it needs a new representation, its code-generation support.
-This layer is below the layout-aware ``Tx.tile.*`` primitives and does not run
+This layer is below the layout-aware ``Tx.cuda.tile.*`` / ``Tx.trn.tile.*`` primitives and does not run
tile-primitive dispatch.
Supported instruction families
diff --git a/docs/tirx/api/script.rst b/docs/tirx/api/script.rst
index cb2b46c13d..781773c017 100644
--- a/docs/tirx/api/script.rst
+++ b/docs/tirx/api/script.rst
@@ -67,4 +67,4 @@ thread-placement and tensor alias policies remain in the TIRx/S-TIR builders.
C = tmem_pool.alloc_tcgen05_mma_D(
(64, 128), "float32", M=128, cta_group=2)
frag = Tx.alloc_tcgen05_ldst_frag("32x32b", (64, 128), "float32")
- Tx.tile.wg.copy_async(frag[:, :], C[:, :])
+ Tx.cuda.tile.tcgen05.ld(frag[:, :], C[:, :], scope="warpgroup")
diff --git a/docs/tirx/api/tile.rst b/docs/tirx/api/tile.rst
index 2abf16986b..e72824cc8c 100644
--- a/docs/tirx/api/tile.rst
+++ b/docs/tirx/api/tile.rst
@@ -15,144 +15,114 @@
specific language governing permissions and limitations
under the License.
-Tile Primitive Authoring API
-============================
+Tensor Instruction Authoring API
+================================
-Kernel code calls tile primitives through ``Tx.tile``::
+Use ``T.cuda.tile`` or ``T.trn.tile`` after ``from tvm.script import tirx as T``.
+The signatures below list runtime operands first and static keyword qualifiers
+after ``*``. Tensor variables are expanded to full regions. An omitted unary
+source defaults to its destination. Named workspace arguments are optional.
- from tvm.script import tirx as Tx
+See :doc:`../tile_primitives` for memory, scope, descriptor, and synchronization
+contracts. These calls produce ordinary IR Calls; the raw scalar instruction
+namespaces remain ``T.ptx`` and ``T.nki``.
- Tx.tile.copy(dst, src)
- Tx.tile.wg.gemm_async(accum, lhs, rhs)
+CUDA
+----
-The scope namespaces bind an operation to an execution scope.
-``Tx.tile.copy(...)`` uses thread scope, while ``Tx.tile.warp.copy(...)``,
-``Tx.tile.wg.copy(...)``, and ``Tx.tile.cta.copy(...)`` request wider
-cooperation. The dialect also exposes those scope objects at its root: for
-example, ``Tx.cta.copy(...)``
-and ``Tx.tile.cta.copy(...)`` construct the same call. This documentation uses
-the explicit ``Tx.tile`` form consistently. See the
-:doc:`programming guide <../tile_primitives>` for the model, primitive catalog,
-and dispatch configuration.
+.. py:function:: T.cuda.tile.ld(dst, src, *, scope='thread', vec_bits=0, cache=None, l1_evict=None, l2_evict=None, prefetch_size=None)
-Scope namespaces
-----------------
+.. py:function:: T.cuda.tile.st(dst, src, *, scope='thread', vec_bits=0, cache=None, l1_evict=None, l2_evict=None, prefetch_size=None)
-.. autodata:: tvm.tirx.script.tile.thread
- :no-index:
+.. py:function:: T.cuda.tile.ldmatrix(dst, src, *, scope='thread')
-.. autodata:: tvm.tirx.script.tile.warp
- :no-index:
+.. py:function:: T.cuda.tile.stmatrix(dst, src, *, scope='thread')
-.. autodata:: tvm.tirx.script.tile.wg
- :no-index:
+.. py:function:: T.cuda.tile.mov(dst, src, *, scope='thread')
-.. autodata:: tvm.tirx.script.tile.warpgroup
- :no-index:
+.. py:function:: T.cuda.tile.cp_async(dst, src, predicate=-1, *, scope='thread', direct=False, fill_mode='', prefetch_size=-1)
-.. autodata:: tvm.tirx.script.tile.cta
- :no-index:
+.. py:function:: T.cuda.tile.cp_async_bulk(dst, src, mbar, remote_cta_id, *, scope='thread')
-.. autodata:: tvm.tirx.script.tile.cluster
- :no-index:
+.. py:function:: T.cuda.tile.cp_async_bulk_tensor_load(dst, src, mbar, cta_mask=0, mbarrier_addr=False, gather4=None, src_selector=None, cache_policy=None, *, scope='thread', descriptor_mode='auto', cta_group=1, cache_hint='', tma_dtype=None, oob=None, prefetch_tensormap=False, tensormap_l2_promotion=None, reduce_op=None)
-Operations
-----------
+.. py:function:: T.cuda.tile.cp_async_bulk_tensor_store(dst, src, cache_policy=None, *, scope='thread', descriptor_mode='auto', cta_group=1, cache_hint='', tma_dtype=None, oob=None, prefetch_tensormap=False, tensormap_l2_promotion=None, reduce_op=None)
-``ScopedOp`` objects are callable but are not ordinary Python functions, so
-they are listed explicitly. Autodoc follows each wrapper to its underlying
-builder function, whose signature includes an internal ``scope`` parameter.
-Kernel authors do not pass that parameter directly: choose a scope through
-``Tx.tile.thread`` / ``warp`` / ``wg`` / ``cta`` / ``cluster``. Bound scope
-objects expose a generic ``(*args, **kwargs)`` Python signature and forward to
-the same underlying operation.
+.. py:function:: T.cuda.tile.cp_reduce_async_bulk_tensor(dst, src, cache_policy=None, *, scope='thread', descriptor_mode='auto', cta_group=1, cache_hint='', tma_dtype=None, oob=None, prefetch_tensormap=False, tensormap_l2_promotion=None, reduce_op=None)
-.. autofunction:: tvm.tirx.script.tile.add
- :no-index:
+.. py:function:: T.cuda.tile.tcgen05.cp(dst, src, *, scope='thread', shape=None, multicast=None, cta_group=1, decompress=None)
-.. autofunction:: tvm.tirx.script.tile.binary_chain
- :no-index:
+.. py:function:: T.cuda.tile.tcgen05.ld(dst, src, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.binary_reduce
- :no-index:
+.. py:function:: T.cuda.tile.tcgen05.st(dst, src, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.cast
- :no-index:
+.. py:function:: T.cuda.tile.mma_sync(D, A, B, C, transpose_A=False, transpose_B=False, alpha=1.0, beta=0.0, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.copy
- :no-index:
+.. py:function:: T.cuda.tile.tcgen05.mma(C, A, B, transA=False, transB=False, accum=False, descI=None, *, scope='thread', cta_group=1, mma_m=None, mma_n=None, smem_desc='hoist', is_AB_tf32=False, weight_stationary=None)
-.. autofunction:: tvm.tirx.script.tile.copy_async
- :no-index:
+.. py:function:: T.cuda.tile.tcgen05.mma_block_scale(C, A, B, SFA, SFB, transA=False, transB=False, accum=False, descI=None, *, scope='thread', cta_group=1, mma_m=None, mma_n=None, smem_desc='hoist', is_AB_tf32=False, weight_stationary=None)
-.. autofunction:: tvm.tirx.script.tile.exp
- :no-index:
+.. py:function:: T.cuda.tile.add(dst, lhs, rhs, *, scope='thread', rounding_mode=None)
-.. autofunction:: tvm.tirx.script.tile.exp2
- :no-index:
+.. py:function:: T.cuda.tile.sub(dst, lhs, rhs, *, scope='thread', rounding_mode=None)
-.. autofunction:: tvm.tirx.script.tile.log2
- :no-index:
+.. py:function:: T.cuda.tile.mul(dst, lhs, rhs, *, scope='thread', rounding_mode=None)
-.. autofunction:: tvm.tirx.script.tile.fdiv
- :no-index:
+.. py:function:: T.cuda.tile.div(dst, lhs, rhs, *, scope='thread', rounding_mode=None)
-.. autofunction:: tvm.tirx.script.tile.fill
- :no-index:
+.. py:function:: T.cuda.tile.max(dst, lhs, rhs, *, scope='thread', rounding_mode=None)
-.. autofunction:: tvm.tirx.script.tile.fma
- :no-index:
+.. py:function:: T.cuda.tile.cvt(dst, src=None, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.gemm
- :no-index:
+.. py:function:: T.cuda.tile.sqrt(dst, src=None, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.gemm_async
- :no-index:
+.. py:function:: T.cuda.tile.ex2(dst, src=None, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.max
- :no-index:
+.. py:function:: T.cuda.tile.lg2(dst, src=None, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.maximum
- :no-index:
+.. py:function:: T.cuda.tile.compose.exp(dst, src=None, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.memset
- :no-index:
+.. py:function:: T.cuda.tile.compose.silu(dst, src=None, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.min
- :no-index:
+.. py:function:: T.cuda.tile.fma(dst, src, scale, bias, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.minimum
- :no-index:
+.. py:function:: T.cuda.tile.compose.exp_with_scale_bias(dst, src, scale, bias, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.mul
- :no-index:
+.. py:function:: T.cuda.tile.compose.exp2_with_scale_bias(dst, src, scale, bias, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.permute_layout
- :no-index:
+.. py:function:: T.cuda.tile.compose.log2_with_scale_bias(dst, src, scale, bias, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.reciprocal
- :no-index:
+.. py:function:: T.cuda.tile.compose.sqrt_with_scale_bias(dst, src, scale, bias, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.reduce_negate
- :no-index:
+TRN
+---
-.. autofunction:: tvm.tirx.script.tile.select
- :no-index:
+.. py:function:: T.trn.tile.load(dst, src, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.silu
- :no-index:
+.. py:function:: T.trn.tile.store(dst, src, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.sqrt
- :no-index:
+.. py:function:: T.trn.tile.tensor_copy(dst, src, *, scope='thread', max_inst_size=512)
-.. autofunction:: tvm.tirx.script.tile.sub
- :no-index:
+.. py:function:: T.trn.tile.matmul(D, A, B, C=None, transpose_A=False, transpose_B=False, alpha=1.0, beta=0.0, acc_psum=None, *, scope='thread')
-.. autofunction:: tvm.tirx.script.tile.sum
- :no-index:
+.. py:function:: T.trn.tile.reciprocal(dst, src, *, scope='thread', max_inst_size=512)
-.. autofunction:: tvm.tirx.script.tile.unary_reduce
- :no-index:
+.. py:function:: T.trn.tile.memset(dst, src, *, scope='thread', max_inst_size=512)
-.. autofunction:: tvm.tirx.script.tile.zero
- :no-index:
+.. py:function:: T.trn.tile.activation(dst, src, scale=1.0, bias=0.0, const_bias=None, *, scope='thread', opcode='exp', max_inst_size=512)
+
+.. py:function:: T.trn.tile.tensortensor(dst, lhs, rhs, *, scope='thread', opcode='add', max_inst_size=512)
+
+.. py:function:: T.trn.tile.tensorscalar(dst, lhs, rhs, *, scope='thread', opcode='add', max_inst_size=512)
+
+.. py:function:: T.trn.tile.tensorreduce(dst, src, partial_reduce=None, *, scope='thread', reduce_op='sum', axes=[-1], negate=False, max_inst_size=None)
+
+.. py:function:: T.trn.tile.scalar_tensor_scalar(dst, data, operand0, operand1, *, scope='thread', op0='mul', op1='add', reverse1=False, max_inst_size=512)
+
+.. py:function:: T.trn.tile.scalar_tensor_tensor(dst, data, operand0, operand1, *, scope='thread', op0='mul', op1='add', reverse1=False, max_inst_size=512)
+
+.. py:function:: T.trn.tile.tensorscalar_reduce(dst, reduced, lhs, rhs, partial_reduce=None, *, scope='thread', opcode='add', reduce_op='sum', axes=[-1], max_inst_size=None)
+
+.. py:function:: T.trn.tile.activation_reduce(dst, reduced, src, scale=1.0, bias=0.0, const_bias=None, partial_reduce=None, *, scope='thread', opcode='exp', reduce_op='sum', axes=[-1], max_inst_size=None)
+
+.. py:function:: T.trn.tile.affine_select(dst, true_value, false_value, pred, *, scope='thread')
diff --git a/docs/tirx/api/tile_dispatch.rst b/docs/tirx/api/tile_dispatch.rst
index 81dfc49b04..51099d48b2 100644
--- a/docs/tirx/api/tile_dispatch.rst
+++ b/docs/tirx/api/tile_dispatch.rst
@@ -15,43 +15,26 @@
specific language governing permissions and limitations
under the License.
-Tile Dispatch Extension API
-===========================
-
-Raw node constructors
----------------------
-
-``tvm.tirx.script.ir_builder.tirx`` constructs the same
-``TileOpCall`` nodes as the validated ``Tx.tile`` facade. This surface
-is intended for parser, builder, and extension authors; kernel code should
-normally use the :doc:`tile authoring API <tile>`. Its public builder spelling
-is ``tvm.tirx.script.ir_builder.tile``::
-
- from tvm.tirx.script import ir_builder as Tx_builder
-
- Tx_builder.tile.copy(dst, src)
-
-Each callable tile operation has a same-name raw constructor. The authoring
-page lists the public facade without duplicating the full operation list here.
-
-.. py:class:: ScopedOp(fn)
-
- Internal callable wrapper that supplies a default thread scope or binds a
- scope selected through a namespace.
-
-.. py:class:: ScopeNamespace(scope, label)
-
- Internal resolver that binds the named execution scope to a tile operation.
-
-.. _tirx-api-dispatch-extension-points:
-
-Dispatch registration
----------------------
-
-.. automodule:: tvm.tirx.operator.tile_primitive
- :members:
- :imported-members:
- :no-index:
-
-See :doc:`../arch/tile_dispatch` for the selection algorithm and the registered
-CUDA and Trainium variants.
+Tensor Instruction Registration
+===============================
+
+Backend authors define an ``Instruction`` with its canonical operator name,
+ordered ``Operand`` list, reflected static attribute schema, and one lowerer.
+The shared implementation is ``tvm.tirx.tensor_instruction``; CUDA and
+Trainium's ``tensor_instructions.py`` files contain the concrete contracts.
+
+Registration installs the fixed argument count, void return type, opaque
+side effects, category, printer name, and native Call validator on the Op.
+The validator also runs for directly constructed Calls and printer entry.
+Runtime expressions must be operands. New static qualifiers require a typed
+attribute field rather than an arbitrary dictionary.
+
+Lowerers receive a semantic view and ``DispatchContext`` and return a
+``tvm.tirx.Function`` body. 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.
+
+Named workspace operands have fixed optional slots. Declare a workspace policy
+only for instructions that own such scratch storage; do not use it to hide
+caller-owned algorithms. See :doc:`../arch/tile_dispatch` for callbacks and
+the statement replacement boundary.
diff --git a/docs/tirx/arch/backends.rst b/docs/tirx/arch/backends.rst
index 64fa7ad47f..d989e9f697 100644
--- a/docs/tirx/arch/backends.rst
+++ b/docs/tirx/arch/backends.rst
@@ -54,7 +54,7 @@ CUDA ownership
* - ``op``
- Define IR builders used by the CUDA script namespaces.
* - ``tile_primitive``
- - Register CUDA implementations of common ``Tx.tile`` operations.
+ - Register named ``Tx.cuda.tile`` instruction contracts and lowerers.
* - ``codegen`` and ``cpp``
- Register source-generation callbacks and CUDA C++ helpers.
* - ``transforms``
@@ -73,7 +73,7 @@ Trainium ownership
------------------
``tvm.backend.trn`` follows the same boundary. ``script`` and ``op`` construct
-the ``Tx.nki`` surface, ``tile_primitive`` registers target dispatches,
+the ``Tx.nki`` surface, ``tensor_instructions`` registers ``Tx.trn.tile`` instruction contracts,
``layout`` and ``transform`` lower Trainium-specific memory mappings, and
``pipeline`` assembles the Trainium pass sequence. ``target_tags`` registers
the named AWS Trainium targets.
diff --git a/docs/tirx/arch/lowering_pipeline.rst b/docs/tirx/arch/lowering_pipeline.rst
index 16567ef77e..e4a7844560 100644
--- a/docs/tirx/arch/lowering_pipeline.rst
+++ b/docs/tirx/arch/lowering_pipeline.rst
@@ -128,9 +128,8 @@ Inside LowerTIRx
LowerTIRx = Sequential([ TilePrimitiveDispatch, LowerTIRxCleanup ])
-- **``TilePrimitiveDispatch``** replaces every ``TileOpCall`` (``copy``,
- ``gemm``, ``reduction``, …) with the body emitted by its selected backend
- dispatch — the variant-selection and codegen described in
+- **``TilePrimitiveDispatch``** replaces every tensor ``Evaluate(Call)`` with the body emitted by its
+ instruction lowerer, including delayed mathematical composites as described in
:doc:`tile_dispatch`. In the same pass it removes the ``device_entry``
marker, resolves standalone scope-id definitions to ``Bind`` statements,
and wraps the device body in the corresponding thread-extent attributes.
diff --git a/docs/tirx/arch/tile_dispatch.rst b/docs/tirx/arch/tile_dispatch.rst
index 18175ed588..a62738bbca 100644
--- a/docs/tirx/arch/tile_dispatch.rst
+++ b/docs/tirx/arch/tile_dispatch.rst
@@ -15,87 +15,43 @@
specific language governing permissions and limitations
under the License.
-Tile Primitive Dispatch
-=======================
-
-This chapter documents how unresolved ``TileOpCall`` nodes are selected
-and lowered. Kernel authors should start with the
-:doc:`Tile Primitives programming guide <../tile_primitives>`; extension
-authors can find the callable registration interfaces in
-:doc:`../api/tile_dispatch`.
-
-Implementation surfaces
------------------------
-
-The authoritative operation list is the C++ registry
-(``src/tirx/op/tirx.cc``, with operations named ``tirx.tile.<name>``). IR
-wrapper classes live in
-``python/tvm/tirx/op/tile.py``. Raw
-``TileOpCall`` constructors live in
-``python/tvm/tirx/script/ir_builder/tirx.py``, while the validated authoring facade
-is in ``python/tvm/tirx/script/tile.py``. Both Python construction surfaces
-produce the same IR node type.
-
-Dispatch pipeline
------------------
-
-Dispatch runs in the ``tirx.TilePrimitiveDispatch`` pass, the first phase of
-``LowerTIRx()``, before layout and execution-scope cleanup. The C++ mutator
-``TilePrimitiveDispatcher`` walks the IR and, for each call:
-
-#. resolves the ``(inter, intra)`` execution split for the call's scope from the
- active set tracked through control flow (``if wg_id == ...``, ``warp_id``,
- and ``Tx.cuda.elect_sync()``);
-#. builds a ``DispatchContext`` carrying the target, scope, launch parameters,
- value ranges, and encoded ``inter``/``intra`` maps plus ``scope_kind``;
-#. invokes the global FFI hook ``tirx.f_op_dispatcher`` with the call and
- context, which returns a ``Function``;
-#. splices that ``Function`` body in place of the call and drains side-effect
- callbacks for private allocations and device or host initialization.
-
-If a ``TileOpCall`` survives lowering, the verifier reports a fatal
-error.
-
-Variant selection
------------------
-
-The Python dispatcher holds a table keyed by ``(Op, target_kind)``. Backends
-register each case with ``register_dispatch``, including its variant name,
-priority, predicates, and implementation. ``run_dispatch(op_call, sctx)``:
-
-#. looks up the primitive and target pair;
-#. filters to ``op_call.dispatch`` when the caller explicitly requests a
- variant;
-#. sorts candidates by descending priority and then by variant name;
-#. evaluates each candidate's predicates and runs the first implementation that
- accepts the call; a predicate exception is recorded as a rejection reason;
-#. continues searching when an implementation raises ``DispatchFail`` or another
- exception; and
-#. reports every rejection reason when no candidate accepts the call. If an
- implementation raised an unexpected exception, the final ``RuntimeError`` is
- chained from the last such exception so its traceback is retained.
-
-Dispatch is therefore target-specific, priority-ordered, and predicate-guarded,
-with an optional ``dispatch=`` override. Common predicates validate matching
-operand shapes and layouts or require a complete active thread group, preventing
-a partial warp or warpgroup from being lowered by an implementation that needs
-full participation.
-
-Variants by primitive
----------------------
-
-The following chapters document the registered variants, their selection
-conditions, the IR they emit, and the conditions under which they decline.
-Their URLs remain under ``tile_primitives`` for compatibility with existing
-links.
-
-.. toctree::
- :maxdepth: 1
-
- ../tile_primitives/copy
- ../tile_primitives/copy_async
- ../tile_primitives/gemm
- ../tile_primitives/gemm_async
- ../tile_primitives/elementwise
- ../tile_primitives/reduction
- ../tile_primitives/permute_layout
+Tensor Instruction Lowering
+===========================
+
+Backend tensor operations are ordinary ``tvm.ir.Call`` nodes with opaque
+side effects and void return types. CUDA and Trainium register their contracts
+in ``python/tvm/backend/{cuda,trn}/tensor_instructions.py`` using the shared
+``python/tvm/tirx/tensor_instruction.py`` machinery. C++ reflected attribute
+types and the native Call validator are in ``src/tirx/op/tile.cc``.
+
+``tirx.TilePrimitiveDispatch`` remains the first phase of ``LowerTIRx``.
+The pass recognizes ``Evaluate(Call)`` by the operator's ``TIRxOpCategory``:
+``tile_primitive`` or ``tile_composite``. It validates the call, 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.
+
+This preserves launch parameters, inter/intra scope maps, value ranges,
+descriptor caches, host initialization, and allocation/initialization
+callbacks. Lowering occurs at a statement boundary because one instruction
+may require address preparation, layout expansion, and multiple occurrences
+of the same core instruction. A residual-call verifier rejects either
+category after expansion.
+
+There is one lowerer per instruction identity. Unsupported layouts and scopes
+produce an error that includes the instruction, target, scope, and original
+cause. There is no cross-instruction priority search or forced variant bag.
+
+The transient Python ``TensorCall`` view supplies semantic operand names to
+existing backend emission code. It is not a registered IR node and is never
+serialized. Its decoded options are private lowering data. Every expression
+in the original IR resides in ``Call.args``; shared traversal does not inspect
+attributes for hidden expressions.
+
+The TVMScript printer validates tensor Calls and prints canonical backend
+names, tensor regions, nested tuple operands, and static qualifiers. Parsing
+reconstructs the fixed signature. Private Trainium workspace allocation
+rewrites Call argument slots before instruction lowering.
+
+See :doc:`../tile_primitives` for instruction contracts and caller-owned
+algorithms, and :doc:`../api/tile_dispatch` for registration.
diff --git a/docs/tirx/layout.rst b/docs/tirx/layout.rst
index 38a3446841..d86788e3db 100644
--- a/docs/tirx/layout.rst
+++ b/docs/tirx/layout.rst
@@ -351,7 +351,7 @@ fragment API:
.. code-block:: python
frag = Tx.alloc_tcgen05_ldst_frag("32x32b", (64, N), "float32")
- Tx.tile.wg.copy_async(frag[:, :], paired_accumulator[:, :])
+ Tx.cuda.tile.tcgen05.ld(frag[:, :], paired_accumulator[:, :], scope="warpgroup")
Tx.ptx.tcgen05.wait__ld.sync.aligned()
The logical ``(64, N)`` fragment is one physical ``.32x32b`` transfer over
diff --git a/docs/tirx/overview.rst b/docs/tirx/overview.rst
index e3deb79d1e..e6f4293415 100644
--- a/docs/tirx/overview.rst
+++ b/docs/tirx/overview.rst
@@ -102,8 +102,8 @@ program, but enter the compiler at different levels:
- ``Tx.*``
- Creates TIRx statements, expressions, buffers, and control flow.
* - Tile primitives
- - ``Tx.tile.*``
- - Creates a ``TileOpCall`` that is replaced by a target-specific
+ - ``Tx.cuda.tile.*`` / ``Tx.trn.tile.*``
+ - Creates an ordinary void ``Call`` that is replaced by its backend instruction
implementation during tile dispatch.
* - Backend operations
- ``Tx.cuda.*`` and ``Tx.ptx.*``
@@ -117,7 +117,7 @@ convert the program to the separate ``tvm.tir`` object model.
.. code-block:: text
Tx.* ──────────────────────────────▶ TIRx statements and expressions ──────┐
- Tx.tile.* ─▶ TileOpCall ─▶ target dispatch ─────────────────────────┤
+ backend.tile.* ─▶ Evaluate(Call) ─▶ instruction lowering ─────────────────────────┤
Tx.cuda.* / Tx.ptx.* ─────────────▶ backend calls ─────────────────────────┤
▼
layout and scope lowering ─▶ codegen
@@ -135,8 +135,8 @@ issuing lane within each active warp. To obtain one issuer for a wider scope,
first select one warp and then use ``elect_sync`` inside it.
The primitive namespace is also part of the scope. For example,
-``Tx.tile.wg.*`` denotes warpgroup-level primitives, while an unqualified
-``Tx.tile.*`` call defaults to thread-level invocation.
+``scope="warpgroup"`` selects warpgroup cooperation; ``scope="thread"``
+is the default on backend tensor instructions.
Tensor layout
~~~~~~~~~~~~~~
diff --git a/docs/tirx/programming/index.rst b/docs/tirx/programming/index.rst
index c0b556b198..63d21cc95b 100644
--- a/docs/tirx/programming/index.rst
+++ b/docs/tirx/programming/index.rst
@@ -28,7 +28,7 @@ The examples use a single alias for the TIRx TVMScript dialect::
Core language constructs are written as ``Tx.*``. Backend-specific helpers
and direct CUDA instructions are available through ``Tx.cuda.*`` and
-``Tx.ptx.*``. Reusable tile operations are written as ``Tx.tile.*``.
+``Tx.ptx.*``. Reusable tile operations are written as ``Tx.cuda.tile.*`` / ``Tx.trn.tile.*``.
.. toctree::
:maxdepth: 1
diff --git a/docs/tirx/tile_primitives.rst b/docs/tirx/tile_primitives.rst
index 0831f55ece..4dd6a64ea5 100644
--- a/docs/tirx/tile_primitives.rst
+++ b/docs/tirx/tile_primitives.rst
@@ -15,187 +15,126 @@
specific language governing permissions and limitations
under the License.
-Tile Primitives
-===============
+Tensor Instructions
+===================
-.. note::
+Tensor instructions expose a backend instruction over tensors with layouts::
- This page documents the tile-primitive surface and dispatch as it exists in
- the source today; signatures and variants may change.
+ from tvm.script import tirx as T
-Tile primitives are the dispatchable, hardware-level operations a TIRx kernel
-issues — data movement (``copy``, ``copy_async``), matrix multiply (``gemm``,
-``gemm_async``), reductions, elementwise math, and a few fused/compose forms.
-A primitive call is recorded as an **unresolved** ``TileOpCall`` IR node;
-the compiler later *dispatches* it — selecting a concrete lowering from the
-primitive, the execution scope, the operand layouts, the target, and an optional
-explicit hint — and replaces it with native IR (loops, address arithmetic,
-synchronization, and backend intrinsics).
+ T.cuda.tile.ld(registers, global_tile, scope="warpgroup")
+ T.cuda.tile.st(shared_tile, registers, scope="warpgroup")
+ T.cuda.tile.tcgen05.mma(accumulator, shared_a, shared_b)
-Calling convention
-------------------
-
-The examples use one TIRx dialect alias and reach tile primitives through its
-``tile`` namespace::
+They construct ordinary, opaque, void ``tvm.ir.Call`` expressions. A statement
+in TVMScript becomes one ``Evaluate(Call(...))``. ``T.ptx`` and ``T.nki``
+continue to expose raw instructions over scalar registers and addresses.
- from tvm.script import tirx as Tx
+Each tensor instruction has a fixed operand list and typed static attributes.
+Tensors, regions, pointers, predicates, coordinates, runtime cache policies,
+accumulation controls, and optional workspace tensors are **Call arguments**.
+They participate in normal free-variable analysis, substitution, serialization,
+and traversal. Static qualifiers such as instruction shape, cache hint, and
+execution scope live in reflected attribute types.
-``Tx`` is an ordinary Python module alias, not an injected language keyword.
-Under ``Tx.tile``, the next namespace prefix selects the **cooperation
-scope**:
+``scope="thread"`` is the default. Use ``"warp"``, ``"warpgroup"``, ``"cta"``,
+or ``"cluster"`` only where the selected instruction supports that scope.
+Layouts and the active thread set must satisfy the instruction's contract.
+There is no generic ``copy``/``gemm`` operation or priority dispatch, and no
+``config``, ``workspace``, or ``dispatch`` keyword bag.
-- ``Tx.tile.<name>(...)`` — unqualified, runs at **thread** scope.
-- ``Tx.tile.warp.<name>`` / ``Tx.tile.wg.<name>`` (alias
- ``Tx.tile.warpgroup``) / ``Tx.tile.cta.<name>`` /
- ``Tx.tile.cluster.<name>`` / ``Tx.tile.thread.<name>`` — bind a wider scope.
+CUDA instructions
+-----------------
-Most primitive constructors also carry ``workspace: dict[str, Var] | None``,
-``dispatch: str | None`` (force a named lowering variant), and ``**kwargs``
-collected into a ``config`` dict that tunes the chosen lowering. ``ScopedOp``
-fills the underlying ``scope`` argument from the namespace prefix; select a
-scope with ``Tx.tile.warp`` / ``Tx.tile.wg`` / ``Tx.tile.cta`` rather than
-passing it to the callable directly. Operands are tensor variables / ``TensorRegion``
-values, each carrying a :doc:`TileLayout <layout>` that dispatch reads.
+* ``ld`` and ``st`` transfer between local registers and global/shared memory.
+ Their register layout determines the participating threads. ``vec_bits``
+ selects an explicit 16-, 32-, 64-, 128-, or 256-bit thread transfer.
+ Fixed-width global loads accept ``cache``, ``l1_evict``, ``l2_evict``, and
+ ``prefetch_size`` qualifiers.
+* ``mov`` copies registers or fills a destination with a scalar. ``ldmatrix``
+ and ``stmatrix`` explicitly request matrix memory instructions.
+* ``cp_async(dst, src, predicate=-1)`` requests global-to-shared asynchronous
+ copying. ``direct=True`` requests one thread-level instruction.
+ ``fill_mode="zero"`` uses the instruction's source-size operand for zero
+ filling; it must not become a predicate that skips the instruction.
+* ``cp_async_bulk(dst, src, mbar, remote_cta_id)`` requests DSMEM copying.
+ ``dst`` and ``mbar`` refer to the executing CTA's local shared addresses;
+ lowering maps both addresses to the remote CTA before issuing the copy.
+* ``cp_async_bulk_tensor_load(dst, src, mbar, cta_mask=0,
+ mbarrier_addr=False, gather4=None, src_selector=None, cache_policy=None)``
+ requests a TMA load. ``cp_async_bulk_tensor_store(dst, src,
+ cache_policy=None)`` and ``cp_reduce_async_bulk_tensor(dst, src,
+ cache_policy=None, reduce_op=...)`` request TMA stores and reductions.
+* TMA ``descriptor_mode="auto"`` derives a legal descriptor and issue loops
+ from layouts. ``descriptor_mode="explicit"`` follows the declared global
+ tensor descriptor and issues once. Gather4 row coordinates and source
+ selectors are explicit-mode operands. ``cache_hint`` is a static string;
+ ``cache_policy`` is a separate runtime operand. They are mutually exclusive.
+* ``tcgen05.cp`` copies shared memory to TMEM. ``tcgen05.ld/st`` transfer
+ between TMEM and registers at ``scope="warpgroup"``; each warp issues its
+ own instruction. The caller supplies waits, fences, and barriers.
+* ``mma_sync`` performs register MMA. ``tcgen05.mma`` and
+ ``tcgen05.mma_block_scale`` perform asynchronous TMEM accumulation. The
+ latter requires both scale-factor tensor operands. Matrix descriptor and
+ pointer preparation remain inside lowering.
+* ``add/sub/mul/div/max``, ``cvt``, ``fma``, ``sqrt``, ``ex2``, and ``lg2``
+ expose tensor arithmetic. Reciprocal is ``div(dst, 1, src)``. Packed
+ arithmetic remains available when the operand layouts permit it.
+
+The mathematical compositions ``compose.silu``, ``compose.exp``,
+``compose.exp_with_scale_bias``, ``compose.exp2_with_scale_bias``,
+``compose.log2_with_scale_bias``, and ``compose.sqrt_with_scale_bias`` are
+also ordinary Calls. They expand during lowering and have the distinct
+``tile_composite`` category.
+
+Trainium instructions
+---------------------
+
+``T.trn.tile`` exposes ``load``, ``store``, ``tensor_copy``, ``matmul``,
+``activation``, ``reciprocal``, ``memset``, ``tensortensor``, ``tensorscalar``,
+``tensorreduce``, and ``affine_select``. The selected binary instruction must
+match the operand layout: broadcasting along a free dimension can require
+``tensorscalar`` even when the scalar is supplied by another tensor.
+
+``scalar_tensor_scalar``, ``scalar_tensor_tensor``, ``tensorscalar_reduce``,
+and ``activation_reduce`` represent native fused instructions. Static
+``opcode``, ``op0``, ``op1``, ``reduce_op``, and ``axes`` qualifiers specify
+their arithmetic. Scale, bias, and all tensor operands remain in ``args``.
+
+Named optional ``acc_psum``, ``const_bias``, and ``partial_reduce`` operands
+allow caller-provided workspace. The private allocation pass still allocates
+and reuses omitted workspace for the instructions that need it. In IR, an
+absent operand occupies its fixed slot as an empty ``Tuple``; TVMScript prints
+it as ``None``.
+
+Algorithms owned by the caller
+------------------------------
+
+CUDA reductions, layout permutation, synchronous global/shared copying, and
+scalar fallback copying are no longer tensor primitives or compositions.
+Write the loops, temporary registers, and synchronization at the call site.
+For an in-place shared-memory permutation, finish all reads into registers,
+synchronize the warp, write using the destination layout, and synchronize
+again before reusing the storage.
+
+Trainium ``tensor_copy`` rejects partition-axis transposition. A caller using
+matrix multiplication for transposition allocates and initializes an identity
+tensor, allocates PSUM, calls ``matmul(..., transpose_A=True)``, and optionally
+calls ``tensor_copy`` to move the result back to SBUF. Copy lowering does not
+allocate an identity or select a matrix multiplication algorithm.
+
+See :doc:`api/tile` for the constructor signatures and
+:doc:`arch/tile_dispatch` for the lowering boundary.
+
+Instruction guides
+------------------
-Primitive catalog
------------------
+.. toctree::
+ :maxdepth: 1
-The C++ registry currently defines 31 operation names. This programming guide
-groups them by purpose; the :doc:`API reference <api/tile>` lists their current
-Python callables and explains the internal ``scope`` parameter exposed by
-introspection.
-
-.. list-table::
- :header-rows: 1
- :widths: 22 34 44
-
- * - Group
- - Operations
- - Purpose
- * - Data movement
- - ``copy``, ``copy_async``, ``permute_layout``
- - synchronous and asynchronous transfer, or rearrangement between layouts
- * - Matrix multiply
- - ``gemm``, ``gemm_async``
- - synchronous register MMA or asynchronous/backend-specific MMA, including
- block scaling where supported
- * - Initialization
- - ``fill``, ``memset``, ``zero``
- - initialize a tile or region
- * - Unary and cast
- - ``cast``, ``sqrt``, ``exp``, ``exp2``, ``log2``, ``reciprocal``, ``silu``
- - per-element conversion or unary math
- * - Binary and ternary
- - ``add``, ``sub``, ``mul``, ``fdiv``, ``maximum``, ``minimum``, ``fma``,
- ``select``
- - per-element arithmetic and selection
- * - Reductions
- - ``sum``, ``max``, ``min``
- - reduce selected axes, optionally accumulating into the destination
- * - Fused and composed
- - ``binary_reduce``, ``unary_reduce``, ``binary_chain``, ``reduce_negate``
- - combine several primitive operations for backends that dispatch them as
- one unit
-
-Dispatch config
----------------
-
-A call is materialized as a ``TileOpCall`` node whose fields carry
-everything dispatch needs (``python/tvm/tirx/tile_dispatch.py``):
-
-.. list-table::
- :header-rows: 1
- :widths: 14 22 64
-
- * - Field
- - Type
- - Meaning
- * - ``op``
- - ``tvm.Op``
- - primitive identity, e.g. ``tirx.tile.copy_async``
- * - ``args``
- - ``Array``
- - operands (regions / scalars), in the order shown above
- * - ``workspace``
- - ``Map[str, Var]``
- - pre-allocated scratch buffers
- * - ``config``
- - ``Map[str, Any]``
- - open-ended tuning bag (table below)
- * - ``dispatch``
- - ``Optional[str]``
- - forced variant name; ``None`` = auto-select
- * - ``scope``
- - ``ExecScope``
- - cooperation scope (default ``thread``)
-
-``config`` has **no central schema**. Each dispatch implementation defines and
-interprets the keys it consumes. Some implementations ignore unrelated keys,
-while others (notably the TMA variants) reject unknown keys. Only ``dispatch``
-is interpreted generically by the dispatcher. The current target
-implementations consume the following keys:
-
-.. list-table::
- :header-rows: 1
- :widths: 28 34 38
-
- * - Primitive / variant
- - Keys consumed
- - Notes
- * - ``copy``: ``vec_16b`` / ``vec_32b`` / ``vec_64b`` /
- ``vec_128b`` / ``vec_256b``
- - ``cache``, ``l1_evict``, ``l2_evict``, ``prefetch_size``
- - Explicit, thread-scope fixed-width copies. Cache controls require a
- global-memory source.
- * - ``copy_async``: ``ldgsts``
- - ``prefetch_size``, ``predicate``, ``fill_mode``, ``direct``
- - ``direct=True`` requires thread scope and an exact 4-, 8-, or 16-byte
- region.
- * - ``copy_async``: ``tma_auto`` / ``tma_explicit``
- - ``cache_hint``, ``cta_group``, ``cta_mask``, ``mbar``,
- ``mbarrier_addr``, ``oob``, ``prefetch_tensormap``,
- ``tensormap_l2_promotion``, ``tma_dtype``, ``use_tma_reduce``
- - ``tma_explicit`` additionally accepts ``gather4`` and ``src_selector``.
- Direction-specific constraints are documented in
- :doc:`tile_primitives/copy_async/tma`.
- * - ``copy_async``: ``dsmem``
- - ``remote_cta_id``, ``mbar``
- - Both are required: the destination CTA id and its completion barrier.
- * - ``copy_async``: ``smem->tmem``
- - ``shape``, ``multicast``, ``cta_group``
- - ``decompress`` is detected but currently rejected as unsupported.
- * - ``gemm_async``: ``tcgen05``
- - ``cta_group``, ``mma_m``, ``mma_n``, ``descI``, ``is_AB_tf32``,
- ``weight_stationary``, ``smem_desc``
- - ``mma_m`` and ``mma_n`` are an all-or-nothing pair. ``descI`` is
- accepted only by the block-scaled path.
- * - ``sum`` / ``max`` / ``min``: ``local``
- - ``thread_reduce``
- - Enables the per-thread shuffle-reduction mode where supported.
- * - binary elementwise: ``reg`` / ``smem``
- - ``rounding_mode``
- - Passed to packed floating-point forms that expose a rounding mode.
- * - Trainium tile implementations
- - ``max_inst_size``
- - Instruction-size limit used by copy, elementwise, select, reduction,
- and composed-op implementations.
-
-Vector widths selected by ``vec_auto`` and ``ldgsts`` are derived internally
-from dtype, alignment, layout, and execution scope; ``vec_len`` is not a user
-configuration key.
-
-Three dispatch inputs are **implicit**, not config keys: the **execution scope**
-(set by the namespace, then refined against the active thread set tracked through
-control flow into ``inter``/``intra`` maps and a ``scope_kind``), the **operand
-layouts** (each tensor variable's ``layout``), and the **target** (the dispatch table is
-keyed by its kind, e.g. ``"cuda"``).
-
-See also
---------
-
-- :doc:`layout` — the ``TileLayout`` model dispatch reads from operands.
-- :doc:`api/tile` — current ``Tx.tile.*`` callables and introspected signatures.
-- :doc:`arch/tile_dispatch` — dispatch selection, extension points, and the
- target-specific variants for each primitive.
-- :doc:`overview` — execution scope, tensor layout, and tile primitive dispatch
- as the three core constructs.
+ tile_primitives/copy
+ tile_primitives/copy_async
+ tile_primitives/elementwise
+ tile_primitives/gemm
+ tile_primitives/gemm_async
+ tile_primitives/reduction
diff --git a/docs/tirx/tile_primitives/copy.rst b/docs/tirx/tile_primitives/copy.rst
index f95b6bfad4..000d2cec2b 100644
--- a/docs/tirx/tile_primitives/copy.rst
+++ b/docs/tirx/tile_primitives/copy.rst
@@ -15,56 +15,22 @@
specific language governing permissions and limitations
under the License.
-copy
-====
+CUDA Tensor Loads, Stores, and Moves
+====================================
-Synchronous element copy ``src → dst`` between global, shared, and register
-(``local``) memory. CUDA currently registers eight variants: five explicit
-fixed-width variants, ``ldstmatrix``, ``vec_auto``, and ``fallback``. The
-``vec_auto`` implementation contains separate global/shared and register paths;
-``gmem_smem`` and ``reg`` below name those implementation paths, not selectable
-dispatch variants.
+``T.cuda.tile.ld(dst_registers, src_memory, scope=...)`` and
+``T.cuda.tile.st(dst_memory, src_registers, scope=...)`` expand the register
+layout into memory instructions. ``vec_bits`` requests a fixed-width
+thread-level instruction. ``ldmatrix`` and ``stmatrix`` explicitly select
+matrix load/store instructions. ``mov`` copies registers or fills a tensor.
-.. list-table::
- :header-rows: 1
- :widths: 28 16 12 44
-
- * - Variant
- - Pair
- - Prio
- - Lowering
- * - ``vec_16b`` / ``vec_32b`` / ``vec_64b`` / ``vec_128b`` /
- ``vec_256b``
- - global ↔ shared/local, or shared ↔ local
- - 20
- - explicit thread-scope transfer of exactly the named width; optional
- global-load cache controls
- * - ``vec_auto``: :doc:`copy/gmem_smem` path
- - global ↔ shared
- - 10
- - synthesized ``[outer, threads, vec]`` partition with direct-PTX vector
- loads and stores
- * - ``vec_auto``: :doc:`copy/reg` path
- - register ↔ shared/global
- - 10
- - partition induced by the register layout's thread axes
- * - :doc:`copy/ldstmatrix`
- - register ↔ shared
- - 10
- - warp-collective ``ldmatrix`` / ``stmatrix`` (m8n8 fragments)
- * - :doc:`copy/fallback`
- - global / shared / local
- - 0
- - scalar single-thread copy (last resort)
-
-The detailed pages cover the two ``vec_auto`` paths, ``ldstmatrix``, and the
-fallback, including accepted input, lowering algorithm, emitted TIRx IR, and
-generated CUDA:
+Synchronous global/shared copying requires explicit local storage and separate
+load and store operations, or caller-written scalar loops. There is no scalar
+fallback dispatcher. See :doc:`../tile_primitives` for scope and cache
+qualifiers.
.. toctree::
:maxdepth: 1
- copy/gmem_smem
copy/reg
copy/ldstmatrix
- copy/fallback
diff --git a/docs/tirx/tile_primitives/copy/fallback.rst b/docs/tirx/tile_primitives/copy/fallback.rst
deleted file mode 100644
index 584f6ebe06..0000000000
--- a/docs/tirx/tile_primitives/copy/fallback.rst
+++ /dev/null
@@ -1,171 +0,0 @@
-.. Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
-.. http://www.apache.org/licenses/LICENSE-2.0
-
-.. Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
-
-copy → fallback
-===============
-
-The ``fallback`` variant is the **priority-0 last resort**: a scalar, single-thread
-copy that runs only when ``ldstmatrix`` and both ``vec_auto`` paths
-(:doc:`gmem_smem` and :doc:`reg`) have declined. Explicit ``vec_*`` variants
-are considered only when requested through ``dispatch=``. The fallback works for
-valid copies among directly addressable global, shared, and local memory at
-thread, warp, warpgroup, or CTA scope. It is intentionally slow, so it emits a
-``UserWarning`` when chosen. Source:
-``python/tvm/backend/cuda/tile_primitive/copy/fallback.py``.
-
-What it accepts
----------------
-
-The registered predicate checks only ``_is_valid_copy`` (layouts present, equal
-dtype, and equal non-unit extents); it has no separate execution-scope or
-memory-pair predicate. The emitter has an unguarded thread path and first-thread
-paths for warp, warpgroup, and CTA scope. Because its body uses ordinary buffer
-loads and stores, use it with directly addressable global, shared, or local
-buffers. There is no divisibility restriction. It is registered at
-``priority=0`` so it is the last candidate the dispatcher tries:
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - priority
- - ``0`` — only reached after the applicable priority-10 variants decline
- * - target / scope
- - ``cuda``; the emitter implements ``thread`` / ``warp`` / ``warpgroup`` /
- ``cta``. The registration itself has no scope-kind filter
- * - memory pair
- - no registration-time filter; the scalar body is usable for directly
- addressable global / shared / local buffers. It cannot ordinarily load
- from or store to tensor memory
- * - dtype / shape
- - ``_is_valid_copy`` only — equal dtype and equal non-unit extents
-
-A typical reason the fast variants all decline: the region's element count does not
-divide the thread count, so the ``[outer, threads, vec]`` split has no solution
-(``vec_auto``'s global/shared path declines), and neither side is a supported
-register-layout pair (its register path and ``ldstmatrix`` decline). Example: a
-warp (32 threads) copying a ``4×6 = 24``-element tile (``24 ∤ 32``).
-
-Demonstration program
-----------------------
-
-A warp copies a ``4×6`` ``float32`` tile global → shared and back. ``24`` is not
-divisible by ``32``, so this falls through to ``fallback`` (from
-``test_fallback.py``):
-
-.. code-block:: python
-
- shape, dtype = (4, 6), "float32" # 24 elements, 32 threads -> 24 ∤ 32
- s_layout = TileLayout(S[shape])
- full = (slice(0, 4), slice(0, 6))
-
-
- @Tx.function
- def kernel(A: Tx.Tensor(shape, dtype), B: Tx.Tensor(shape, dtype)):
-
- Tx.device_entry()
- Tx.cta_id([1])
- Tx.lane_id([32])
- Tx.thread_id([32])
- A_smem = Tx.alloc_tensor(shape, dtype, scope="shared", layout=s_layout)
- Tx.tile.warp.copy(A_smem[full], A[full]) # fallback
- Tx.cuda.cta_sync()
- Tx.tile.warp.copy(B[full], A_smem[full]) # fallback
-
-Algorithm
----------
-
-**1. Elect one thread.** For a multi-thread scope the copy is done entirely by the
-**first thread of the scope** — its ``laneid`` base, plus the warp offset within a
-warpgroup / cta (radix-32 composition of the per-axis offsets); every other thread
-skips it. For ``thread`` scope there is only one thread, so no guard is emitted:
-
-.. code-block:: python
-
- if scope_kind == "thread":
- def impl():
- _copy_body(dst, src) # the single thread copies everything
- else:
- first_tid = int(sctx.intra["laneid"][1]) # first thread of the scope:
- if scope_kind == "warpgroup":
- first_tid += 32 * int(sctx.intra["wid_in_wg"][1]) # + warp offset within the wg
- elif scope_kind == "cta":
- first_tid += 32 * int(sctx.intra["warpid"][1]) # + warp offset within the cta
- def impl():
- tid = _axis_decl(tid_axis_name, sctx)
- if tid == first_tid: # only the scope's first thread copies
- _copy_body(dst, src)
-
-**2. Scalar nested loop over the region.** ``_copy_body`` iterates the non-unit
-extents with ``Tx.grid`` and copies one element per step — no vectorization, no
-partition:
-
-.. code-block:: python
-
- with Tx.grid(*copy_extents) as lvs: # copy_extents = non-unit dst extents
- dst[_dst_coord(lvs)] = src[_src_coord(lvs)]
-
-Generated TIRx IR
------------------
-
-``LowerTIRx`` produces the guarded scalar grid (global → shared shown):
-
-.. code-block:: python
-
- if tid == 0: # first_tid (lane 0 for this warp scope)
- for v_3, v_4 in Tx.grid(4, 6):
- A_smem[v_3, v_4] = A[v_3, v_4]
-
-Generated CUDA
---------------
-
-.. code-block:: c++
-
- __shared__ alignas(64) float A_smem_ptr[24];
- if (((int)threadIdx.x) == 0) { // lane 0 does the whole copy
- for (int v_3 = 0; v_3 < 4; ++v_3)
- for (int v_4 = 0; v_4 < 6; ++v_4)
- A_smem_ptr[(v_3 * 6) + v_4] = A_ptr[(v_3 * 6) + v_4];
- }
-
-Lane 0 copies all 24 elements one at a time; the other 31 lanes do nothing. (At
-lowering the dispatch also prints ``UserWarning: copy/fallback (scalar
-single-thread) picked … all faster variants rejected``.)
-
-How inputs change the algorithm
--------------------------------
-
-The **scope** decides who runs the loop:
-
-.. list-table::
- :header-rows: 1
- :widths: 30 70
-
- * - scope
- - who copies
- * - ``thread``
- - the single thread, no ``if`` guard
- * - ``warp`` / ``warpgroup`` / ``cta``
- - only ``first_tid`` (the scope's first thread = lane base + warp offset); all
- other threads skip
-
-The **shape** sets the ``Tx.grid`` bounds (the non-unit extents); the loop body is
-always one scalar element copy, regardless of dtype. The algorithm remains scalar
-for every dtype; element width and the target memory spaces can still affect the
-cost of each load and store. This variant exists for correctness, not speed.
diff --git a/docs/tirx/tile_primitives/copy/gmem_smem.rst b/docs/tirx/tile_primitives/copy/gmem_smem.rst
deleted file mode 100644
index b9837d963d..0000000000
--- a/docs/tirx/tile_primitives/copy/gmem_smem.rst
+++ /dev/null
@@ -1,211 +0,0 @@
-.. Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
-.. http://www.apache.org/licenses/LICENSE-2.0
-
-.. Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
-
-copy → vec_auto global/shared path
-==================================
-
-The ``gmem_smem`` implementation path inside the registered ``vec_auto`` variant
-lowers a synchronous ``copy`` between **global and shared** memory (either
-direction) when **neither side is a register**. It is not a separate dispatch
-name; automatic selection or ``dispatch="vec_auto"`` reaches this path. Because
-neither operand carries a thread partition, the implementation synthesizes one
-from the execution scope: it splits the region into ``[outer, threads, vec]`` and
-emits a serial loop of vectorized loads/stores. Source:
-``python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py``.
-
-What it accepts
----------------
-
-The predicate ``_is_gmem_smem`` gates this ``vec_auto`` path:
-
-.. code-block:: python
-
- def _is_gmem_smem(op_call, sctx):
- if not sctx.is_target("cuda"):
- return False, "non-cuda target"
- if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
- return False, f"unsupported exec_scope {sctx.scope_kind}"
- for check in (
- lambda: _all_threads_active(sctx), # full scope, no narrowing
- lambda: _is_valid_copy(op_call, sctx), # layouts, equal dtype/extents
- lambda: _scope_allowed(op_call, sctx, allowed_pairs=_GMEM_SMEM_PAIRS),
- lambda: _divides_thread_cnt(op_call, sctx),
- ):
- ok, msg = check()
- if not ok:
- return False, msg
- return True, None
-
-So the accepted input is:
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target
- - ``cuda``
- * - scope
- - ``thread`` / ``warp`` / ``warpgroup`` / ``cta``, and **all threads active**
- (``_all_threads_active`` — ``laneid`` spans 32, etc., none narrowed by an
- enclosing ``if``)
- * - memory pair
- - ``(global, shared*)`` or ``(shared*, global)`` — ``_GMEM_SMEM_PAIRS``;
- neither side is ``local``
- * - dtype / shape
- - both operands have a layout, equal dtype, equal non-unit extents
- (``_is_valid_copy`` → ``validate_copy_op``)
- * - divisibility
- - the region's element count is divisible by the thread count
- (``_divides_thread_cnt``) — otherwise the ``[outer, threads, vec]`` split has
- no integer solution and the variant declines
-
-Demonstration program
-----------------------
-
-A warp (32 threads) copies a ``32×32`` ``float32`` tile global → shared and back
-(the round trip from ``test_gmem_smem.py``):
-
-.. code-block:: python
-
- from tvm.script import tirx as Tx
- from tvm.tirx.layout import S, TileLayout
-
- shape, dtype = (32, 32), "float32"
- s_layout = TileLayout(S[shape])
- fs = (slice(0, 32), slice(0, 32))
-
-
- @Tx.function
- def kernel(A: Tx.Tensor(shape, dtype), B: Tx.Tensor(shape, dtype)):
-
- Tx.device_entry()
- Tx.cta_id([1])
- Tx.lane_id([32])
- Tx.thread_id([32])
- A_smem = Tx.alloc_tensor(shape, dtype, scope="shared", layout=s_layout)
- Tx.tile.warp.copy(A_smem[fs], A[fs]) # global -> shared (this dispatch)
- Tx.cuda.cta_sync()
- Tx.tile.warp.copy(B[fs], A_smem[fs]) # shared -> global (this dispatch)
-
-Algorithm
----------
-
-**1. Synthesize the partition.** With 32 threads and ``32×32 = 1024`` elements,
-the dispatch builds a 3-D iteration ``[outer, threads, vec]`` via
-``align_layouts_gs``: it slices both layouts to the region, makes the **global**
-side drive the canonical (stride-descending) order, then carves a contiguous
-``vec`` tail and a ``threads`` chunk off it and re-groups the shared side to match.
-
-**2. Choose the vector width, widest first.** It tries element counts for
-``{128, 64, 32, 16, 8}`` bits and accepts the widest where (a) the contiguous tail
-divides it and (b) every **non-vec iter stride** (on both sides, the thread iter
-included) and both base offsets is a multiple of it, so the per-thread, per-round
-vector pointer is naturally aligned. (Only the innermost ``vec`` iter is excluded
-from that check.) For ``float32`` that is ``vec = 4`` (``4 × 4 B = 16 B = 128 bit``),
-giving ``outer = 1024 / (32 × 4) = 8``.
-
-**3. Emit a serial loop** (``vec_auto_gmem_smem.py``) — deliberately an ordinary
-``range`` loop, *not* ``Tx.unroll``. This leaves any final unrolling decision to
-ptxas:
-
-.. code-block:: python
-
- for f in range(total_outer):
- g_lin = g_p.apply(f, tid, v0, shape=apply_shape)["m"]
- s_off = s_apply_layout.apply(f, tid, v0, shape=apply_shape)["m"]
- s_ptr = _ptr_off(s_buf.ptr_to(s_zero), s_off)
- g_ptr = _ptr_off(g_buf.ptr_to(g_zero), g_lin)
- if g_is_src:
- Tx.ptx[ld_g](*[tmp[i] for i in range(lanes)], g_ptr)
- Tx.ptx[st_s](s_ptr, *[tmp[i] for i in range(lanes)])
- else:
- Tx.ptx[ld_s](*[tmp[i] for i in range(lanes)], s_ptr)
- Tx.ptx[st_g](g_ptr, *[tmp[i] for i in range(lanes)])
-
-Each ``(f, tid, 0)`` coordinate is flattened by ``layout.apply`` against
-``[outer, threads, vec]``, so the emit never needs to know how the partition split
-the iters. ``ld_g``, ``st_s``, ``ld_s``, and ``st_g`` are registered direct-PTX
-forms selected for the memory direction and vector width. A 128-bit transfer uses
-four ``uint32`` registers and the ``v4.u32`` forms.
-
-Generated TIRx IR
------------------
-
-Running ``LowerTIRx`` on the program above turns each ``Tx.tile.warp.copy`` into the
-synthesized loop (global → shared shown, trimmed):
-
-.. code-block:: python
-
- tid: Tx.let = threadIdx_x % 32
- A_smem = Tx.alloc_shared((1024,))
- tmp = Tx.alloc_local((4,), "uint32")
- for f in range(8): # outer = 8
- s_lin = f * 128 + tid * 4 # 32 threads × vec 4 = 128 / round
- g_lin = f * 128 + tid * 4
- s_ptr = pointer_offset(A_smem, s_lin)
- g_ptr = pointer_offset(A_1, g_lin) # A_1 = A.view(1024)
- Tx.ptx.ld.global_.v4.u32(tmp[0], tmp[1], tmp[2], tmp[3], g_ptr)
- Tx.ptx.st.shared.v4.u32(s_ptr, tmp[0], tmp[1], tmp[2], tmp[3])
-
-Generated PTX instructions
---------------------------
-
-The CUDA code generator emits one vector load and one vector store per round::
-
- ld.global.v4.u32 {r0, r1, r2, r3}, [g_ptr];
- st.shared.v4.u32 [s_ptr], {r0, r1, r2, r3};
-
-The shared-to-global direction uses ``ld.shared.v4.u32`` followed by
-``st.global.v4.u32``.
-
-Thread ``tid`` handles elements ``[f·128 + tid·4 .. +4)`` each round; across 8
-rounds and 32 lanes that covers all 1024 elements, each as one 128-bit transfer.
-
-How inputs change the algorithm
--------------------------------
-
-The element **dtype** sets the vector width (widest 128-bit transfer that stays
-aligned), which sets the round count. For the same ``32×32`` tile and 32 threads:
-
-.. list-table::
- :header-rows: 1
- :widths: 20 20 20 40
-
- * - dtype
- - ``vec``
- - transfer width
- - ``outer = 1024 / (32 · vec)``
- * - ``float32``
- - 4
- - 16 B (``v4.u32``)
- - 8
- * - ``float16``
- - 8
- - 16 B (``v4.u32``)
- - 4
- * - ``uint8``
- - 16
- - 16 B (``v4.u32``)
- - 2
-
-The **scope** sets which axis names the thread id (``warp`` → ``laneid``,
-``cta`` → ``tx``, …) and the thread count, hence the partition. A **swizzled**
-shared layout caps ``vec`` to one swizzle chunk and routes ``s_off`` through the
-swizzle (a recognized swizzle becomes a few register adds per round; otherwise
-``swizzle.apply`` per round).
diff --git a/docs/tirx/tile_primitives/copy/ldstmatrix.rst b/docs/tirx/tile_primitives/copy/ldstmatrix.rst
index 9aafabea2b..1e265bd16b 100644
--- a/docs/tirx/tile_primitives/copy/ldstmatrix.rst
+++ b/docs/tirx/tile_primitives/copy/ldstmatrix.rst
@@ -15,16 +15,16 @@
specific language governing permissions and limitations
under the License.
-copy → ldstmatrix
-=================
+ldmatrix / stmatrix
+===================
-The ``ldstmatrix`` variant lowers a ``copy`` between **register and shared** memory
+``Tx.cuda.tile.ldmatrix`` and ``Tx.cuda.tile.stmatrix`` transfer between register and shared memory
to the warp-collective PTX ``ldmatrix`` / ``stmatrix`` instructions: one
instruction moves ``num`` 8×8 16-bit matrix tiles between shared memory and the
warp's registers, with the hardware performing the lane↔element shuffle that an
MMA fragment needs. It only applies when the register and shared **layouts match
-the m8n8 fragment geometry**; otherwise the copy falls back to the
-:doc:`reg` path in ``vec_auto``. Source:
+the m8n8 fragment geometry**. An incompatible layout is rejected; use
+:doc:`reg` for ordinary memory instructions. Source:
``python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py``.
What it accepts
@@ -49,10 +49,8 @@ The predicate is lean — scope, a valid copy, and a register↔shared pair:
return False, msg
return True, None
-The **real** gate is the layout fit, applied during emit. Both this variant and
-``vec_auto`` are priority 10 and accept ``local ↔ shared``; ``ldstmatrix`` is
-tried first and **declines** (via ``fail(...)``) if the layouts are not ldmatrix
-fragments, leaving the :doc:`reg` path in ``vec_auto`` to handle the copy:
+The layout fit is validated during emission. The selected matrix instruction
+must match an m8n8 fragment, or lowering reports an error:
.. code-block:: python
@@ -95,7 +93,7 @@ fragments, leaving the :doc:`reg` path in ``vec_auto`` to handle the copy:
the contiguous eight-element shared-memory row addressed by the instruction
Demonstration program
-----------------------
+---------------------
A warp loads ``num = 2`` row-major matrix tiles (``M, N = 8, 16`` fp16) shared →
register, from ``test_ld_stmatrix.py`` (register layout = the m8n8 fragment,
@@ -123,7 +121,7 @@ register, from ``test_ld_stmatrix.py`` (register layout = the m8n8 fragment,
# ... stage A into A_smem (row = tid//4, cp = tid%4) ...
Tx.cuda.cta_sync()
R = Tx.alloc_tensor((8, 4, num, 2), "float16", scope="local", layout=r_layout)
- Tx.tile.warp.copy(R[full], A_smem[full]) # shared -> register (ldmatrix)
+ Tx.cuda.tile.ldmatrix(R[full], A_smem[full], scope="warp") # shared -> register (ldmatrix)
# ... write R back out to B ...
Algorithm
@@ -176,7 +174,7 @@ words:
else:
Tx.ptx[chain](smem_ptr, *words) # stmatrix takes the address first
-(This is the one copy variant that **does** use ``Tx.unroll`` — ``m_outer`` is tiny.)
+This lowerer unrolls the small ``m_outer`` loop.
Generated TIRx IR
-----------------
diff --git a/docs/tirx/tile_primitives/copy/reg.rst b/docs/tirx/tile_primitives/copy/reg.rst
index d2e4974b68..8cc37193cf 100644
--- a/docs/tirx/tile_primitives/copy/reg.rst
+++ b/docs/tirx/tile_primitives/copy/reg.rst
@@ -15,16 +15,14 @@
specific language governing permissions and limitations
under the License.
-copy → vec_auto register path
-==============================
-
-The register implementation path inside the registered ``vec_auto`` variant
-lowers a synchronous ``copy`` where **exactly one side is a register** (``local``)
-buffer and the other is ``shared*`` or ``global``. It is not a separate
-``reg`` dispatch name; automatic selection or ``dispatch="vec_auto"`` reaches
-this path. Unlike :doc:`gmem_smem`, the partition is **not synthesized** — it is
-*induced* by the register operand's layout: that layout's thread-axis iters
-already say which thread owns which logical coordinate, so the dispatch drops
+ld / st: register transfers
+===========================
+
+``Tx.cuda.tile.ld`` loads from shared or global memory into local registers;
+``Tx.cuda.tile.st`` stores registers to memory. With ``vec_bits=0``, the
+register layout determines the vector width. A positive ``vec_bits`` fixes
+the width. The partition is induced by the register operand's layout: that layout's thread-axis iters
+already say which thread owns which logical coordinate, so the lowerer drops
those axes, leaves each thread its private bundle of elements, and copies them
in a vectorized serial loop. Source:
``python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py``.
@@ -73,7 +71,7 @@ What it accepts
its region
Demonstration program
-----------------------
+---------------------
A warp round-trips a ``32×8`` ``float32`` tile shared → register → shared, with the
register layout ``S[(32,8):(1@laneid, 1)]`` — **lane ``i`` owns row ``i``** (8
@@ -101,16 +99,16 @@ contiguous elements). From ``test_reg.py``:
A_smem[tid, kk] = Tx.cast(tid * 100 + kk + 1, dtype)
Tx.cuda.cta_sync()
R = Tx.alloc_tensor(shape, dtype, scope="local", layout=r_layout)
- Tx.tile.warp.copy(R[fs], A_smem[fs]) # shared -> register (this dispatch)
+ Tx.cuda.tile.ld(R[fs], A_smem[fs], scope="warp") # shared -> register (this dispatch)
# ... clear A_smem, cta_sync ...
- Tx.tile.warp.copy(A_smem[fs], R[fs]) # register -> shared (this dispatch)
+ Tx.cuda.tile.st(A_smem[fs], R[fs], scope="warp") # register -> shared (this dispatch)
# ... cta_sync; B[tid, kk] = A_smem[tid, kk] ...
Algorithm
---------
**1. Inherit the partition from R.** The register layout's thread axis (``laneid``)
-states that lane ``i`` owns row ``i``; the dispatch aligns the other (shared) side
+states that lane ``i`` owns row ``i``; the lowerer aligns the other (shared) side
to that order, then **drops the thread iters** — what remains is each thread's
private memory bundle: here ``8`` contiguous elements per lane.
@@ -125,7 +123,7 @@ address). For ``8`` contiguous ``float32`` that is ``vec = 4``, so ``outer = 2``
**3. Per-thread base offset + serial loop.** The shared-side base offset is built
from thread-axis placeholders (substituted with the real ``Tx.lane_id()`` etc.),
and the register side is a flat per-thread ``local`` buffer. The emit is a serial
-loop (not ``Tx.unroll`` — same flooding rationale as :doc:`gmem_smem`):
+loop to limit generated code size:
.. code-block:: python
diff --git a/docs/tirx/tile_primitives/copy_async.rst b/docs/tirx/tile_primitives/copy_async.rst
index 6414d96db2..aa4b64b705 100644
--- a/docs/tirx/tile_primitives/copy_async.rst
+++ b/docs/tirx/tile_primitives/copy_async.rst
@@ -15,45 +15,25 @@
specific language governing permissions and limitations
under the License.
-copy_async
-==========
-
-Asynchronous copy. Every variant emits the transfer's *issue* instruction; the
-caller supplies the matching completion protocol. ``ldgsts`` uses ``cp.async``
-commit/wait, TMA loads and distributed-shared-memory copies signal an mbarrier,
-TMA stores use the bulk async-group commit/wait operations, and tensor-memory
-paths use their ``tcgen05`` commit/wait operations. Selection is by the
-source/destination memory pair and scope. CUDA currently registers six variant
-names; the TMA page below covers two of them.
-
-.. list-table::
- :header-rows: 1
- :widths: 26 16 14 44
-
- * - Variant
- - Pair
- - Prio
- - Issue instruction
- * - :doc:`copy_async/ldgsts`
- - global → shared
- - 20
- - ``cp.async`` (LDGSTS), per-thread vectorized
- * - :doc:`copy_async/tma` (``tma_auto`` / ``tma_explicit``)
- - global ↔ shared
- - 10
- - ``cp.async.bulk.tensor`` (TMA, descriptor-driven, single-thread)
- * - :doc:`copy_async/dsmem`
- - shared → shared (cross-CTA)
- - 10
- - ``cp.async.bulk`` shared::cluster (``mapa`` remote address)
- * - :doc:`copy_async/tcgen05_cp` (``smem->tmem``)
- - shared → tmem
- - 10
- - one of the registered ``tcgen05.cp`` shapes (matrix-descriptor driven)
- * - :doc:`copy_async/tcgen05_ldst` (``tmem<->local``)
- - tmem ↔ register
- - 10
- - ``tcgen05.ld`` / ``tcgen05.st`` (warpgroup, atom-matched)
+Asynchronous tensor transfers
+=============================
+
+The instruction name fixes the transfer mechanism. Every interface issues the
+transfer; its caller supplies the corresponding completion protocol.
+
+* :doc:`copy_async/ldgsts`: ``Tx.cuda.tile.cp_async`` copies global to shared.
+ The caller commits and waits with the ``cp.async`` group operations.
+* :doc:`copy_async/tma`: ``cp_async_bulk_tensor_load``,
+ ``cp_async_bulk_tensor_store``, and ``cp_reduce_async_bulk_tensor`` use TMA.
+ Loads signal an mbarrier; stores and reductions use bulk async groups.
+* :doc:`copy_async/dsmem`: ``cp_async_bulk`` copies between CTAs' shared memory
+ and signals the destination CTA's mbarrier.
+* :doc:`copy_async/tcgen05_cp`: ``tcgen05.cp`` copies shared memory to TMEM.
+* :doc:`copy_async/tcgen05_ldst`: ``tcgen05.ld`` and ``tcgen05.st`` transfer
+ between TMEM and registers at warpgroup scope.
+
+All names above are in ``Tx.cuda.tile``. Tensor-memory transfers require the
+matching ``tcgen05`` completion operations.
.. toctree::
:maxdepth: 1
diff --git a/docs/tirx/tile_primitives/copy_async/dsmem.rst b/docs/tirx/tile_primitives/copy_async/dsmem.rst
index 4db7e1cb84..e733feffce 100644
--- a/docs/tirx/tile_primitives/copy_async/dsmem.rst
+++ b/docs/tirx/tile_primitives/copy_async/dsmem.rst
@@ -15,11 +15,11 @@
specific language governing permissions and limitations
under the License.
-copy_async → dsmem
-==================
+cp_async_bulk: DSMEM
+====================
-The ``dsmem`` variant lowers a ``copy_async`` whose **source and destination are
-both shared** memory but in **different CTAs of a cluster** (distributed shared
+``Tx.cuda.tile.cp_async_bulk`` transfers a tile whose source and destination are
+both shared memory but in different CTAs of a cluster (distributed shared
memory). A call already placed in a single-thread execution scope maps the
destination CTA's shared
address into its own address space (PTX ``mapa``) and issues a bulk copy
@@ -32,27 +32,14 @@ What it accepts
Three predicates: a valid copy, a single-thread scope, and a shared → shared pair:
-.. code-block:: python
-
- # register_dispatch(..., priority=10, when=[
- predicate("validate_copy_op", ...),
- predicate("single_thread", lambda op, sctx: (single_thread(op, sctx), "expected single thread")),
- predicate("is_shared_to_shared", lambda op, sctx: (_is_shared_to_shared(op), "not shared-to-shared")),
- # ])
-
- def _is_shared_to_shared(op_call):
- src_scope = op_call.src.buffer.scope()
- dst_scope = op_call.dst.buffer.scope()
- return src_scope.startswith("shared") and dst_scope.startswith("shared")
-
.. list-table::
:header-rows: 1
:widths: 22 78
* - Property
- Requirement
- * - target / priority
- - ``cuda``; priority ``10``
+ * - target
+ - ``cuda``
* - scope
- the call must already be in a **single-thread** scope; the caller normally
selects that thread with control flow
@@ -61,13 +48,13 @@ Three predicates: a valid copy, a single-thread scope, and a shared → shared p
via ``remote_cta_id``
* - chunk size
- the contiguous chunk must be **≥ 16 bytes and a multiple of 16**
- (``cp.async.bulk`` requirement) — else the dispatch declines (``fail``)
+ (``cp.async.bulk`` requirement) — else lowering fails
* - environment
- a **cluster launch** (so a remote CTA's shared memory exists), plus a caller
mbarrier on the destination CTA
Demonstration program
-----------------------
+---------------------
A 2-CTA cluster: CTA 0 stages a ``128×64`` ``float16`` tile global → its shared,
then bulk-copies it into **CTA 1's** shared via ``dsmem``; CTA 1 waits on the
@@ -116,25 +103,28 @@ mbarrier and writes the result out (from ``test_dsmem.py``):
Tx.cuda.cluster_sync()
if tid == 0:
if cbx == 0: # source CTA
- Tx.tile.copy(src_smem[r], A[r]) # global -> local shared
+ for row, col in Tx.grid(128, 64):
+ value = A[row, col]
+ src_smem[row, col] = value
Tx.ptx.fence.proxy.async_.shared__cta()
- Tx.tile.copy_async(
+ Tx.cuda.tile.cp_async_bulk(
dst_smem[r],
src_smem[r],
- dispatch="dsmem",
mbar=mbar.ptr_to([0]),
remote_cta_id=Tx.int32(1),
) # -> CTA 1
else: # destination CTA
Tx.ptx.mbarrier.arrive.expect_tx.shared.b64(mbar.ptr_to([0]), Tx.uint32(copy_bytes))
mbar.wait(0, 0)
- Tx.tile.copy(B[r], dst_smem[r]) # remote shared -> global
+ for row, col in Tx.grid(128, 64):
+ value = dst_smem[row, col]
+ B[row, col] = value
Tx.cuda.cluster_sync()
Algorithm
---------
-**1. Find the contiguous chunk.** The dispatch slices and groups both layouts to the
+**1. Find the contiguous chunk.** The lowerer slices and groups both layouts to the
copy region, walks inward to the longest matching contiguous stride-1 shard chain,
and multiplies those extents into ``chunk_elements``; ``chunk_bytes`` must be ≥ 16
and a multiple of 16 (a ``cp.async.bulk`` constraint), else it declines:
@@ -174,7 +164,7 @@ chunk's offsets each step:
Tx.cast(chunk_bytes, "uint32"), Tx.cast(mapped[0], "uint32"))
The ``complete_tx::bytes`` form makes the hardware decrement ``remote_mbar`` by
-``chunk_bytes`` on completion; the dispatch emits no wait — the caller arms the
+``chunk_bytes`` on completion; the lowerer emits no wait — the caller arms the
mbarrier (``arrive.expect_tx``) and waits.
Generated TIRx IR
@@ -221,4 +211,4 @@ How inputs change the algorithm
- the ``mapa`` rank — which cluster CTA receives the data
* - incompatible layouts
- e.g. row-major source vs column-major destination → no matching contiguous
- chain → the dispatch declines (``fail``)
+ chain → lowering fails
diff --git a/docs/tirx/tile_primitives/copy_async/ldgsts.rst b/docs/tirx/tile_primitives/copy_async/ldgsts.rst
index 0fa5b50e71..de30b232a2 100644
--- a/docs/tirx/tile_primitives/copy_async/ldgsts.rst
+++ b/docs/tirx/tile_primitives/copy_async/ldgsts.rst
@@ -15,15 +15,14 @@
specific language governing permissions and limitations
under the License.
-copy_async → ldgsts
-===================
+cp_async
+========
-The ``ldgsts`` variant lowers ``copy_async`` for a **global → shared** transfer to
+``Tx.cuda.tile.cp_async`` lowers a global-to-shared transfer to
the PTX ``cp.async`` (LDGSTS) instruction: each thread issues an *asynchronous*
vectorized copy that the hardware completes in the background, so the warp can keep
-computing while the load is in flight. It reuses the exact ``[outer, threads, vec]``
-partition of the synchronous :doc:`../copy/gmem_smem` ``vec_auto`` path; the
-differences are all in *what* is emitted and *when* it completes. Source:
+computing while the load is in flight. Its layout planner constructs an
+``[outer, threads, vec]`` partition. Source:
``python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py``.
What it accepts
@@ -65,13 +64,11 @@ What it accepts
* - vector width
- ``cp.async`` only accepts cp_size ∈ {4, 8, 16} bytes, so the candidate set is
restricted to ``_LDGSTS_VEC_BITS`` = ``{128, 64, 32}`` bits
- * - priority
- - ``20`` — selected for ``copy_async`` global→shared (vs the bulk/TMA variants)
Configuration
-------------
-``ldgsts`` reads four optional config keys. ``prefetch_size`` is ``-1`` by
+``cp_async`` takes a predicate operand and three static qualifiers. ``prefetch_size`` is ``-1`` by
default (no qualifier), or ``64``, ``128``, or ``256`` for the corresponding
PTX ``.L2::<N>B`` qualifier. ``predicate`` defaults to ``-1`` (unpredicated);
with ``fill_mode="zero"``, a false predicate uses the ``src-size`` form to
@@ -84,7 +81,7 @@ integer and constructs the PTX spelling, whose table lookup accepts the
registered qualifiers.
Demonstration program
-----------------------
+---------------------
A CTA (128 threads) asynchronously loads a ``128×32`` ``float16`` tile global →
shared, then commits and waits before reading it back (from ``test_ldgsts.py``):
@@ -105,16 +102,18 @@ shared, then commits and waits before reading it back (from ``test_ldgsts.py``):
Tx.lane_id([32])
tid = Tx.thread_id([128])
A_smem = Tx.alloc_tensor(shape, dtype, scope="shared", layout=s_layout)
- Tx.tile.cta.copy_async(A_smem[full], A[full], dispatch="ldgsts") # async global -> shared
+ Tx.cuda.tile.cp_async(A_smem[full], A[full], scope="cta") # async global -> shared
Tx.ptx.cp.async_.commit_group() # caller commits ...
Tx.ptx.cp.async_.wait_group(0) # ... and waits
Tx.cuda.cta_sync()
- Tx.tile.cta.copy(B[full], A_smem[full])
+ for k in Tx.serial(32):
+ value = A_smem[tid, k]
+ B[tid, k] = value
Algorithm
---------
-**1. Same partition as** :doc:`../copy/gmem_smem`. ``align_layouts_gs`` builds the
+**1. Cooperative partition.** ``align_layouts_gs`` builds the
``[outer, threads, vec]`` split with the global side driving the canonical order —
but the vector candidates are clamped to ``{128, 64, 32}`` bits so the byte size is
a legal ``cp.async`` cp_size. For ``128×32 = 4096`` ``float16`` over 128 threads the
@@ -134,7 +133,7 @@ widest legal width is ``vec = 8`` (``8 × 2 B = 16 B``), giving ``outer = 4``.
Completion is the caller's responsibility
(``Tx.ptx.cp.async_.commit_group()`` then
-``Tx.ptx.cp.async_.wait_group(0)``); the dispatch only issues the in-flight
+``Tx.ptx.cp.async_.wait_group(0)``); the lowerer only issues the in-flight
loads.
Generated TIRx IR
@@ -163,8 +162,8 @@ caller's ``wait_group``.
How inputs change the algorithm
-------------------------------
-The dtype/alignment set ``vec`` (hence ``cp_size`` and ``outer``), but unlike the
-synchronous variant the width is capped at 16 B (cp.async maximum):
+The dtype and alignment set ``vec`` (hence ``cp_size`` and ``outer``).
+The width is capped at 16 B, the cp.async maximum:
.. list-table::
:header-rows: 1
@@ -192,5 +191,5 @@ synchronous variant the width is capped at 16 B (cp.async maximum):
- (×4)
If the region can't satisfy even a 4-byte (32-bit) cp_size, ``align_layouts_gs``
-finds no candidate and the variant declines. The **direction is fixed**: a
-shared → global ``copy_async`` is never ``ldgsts`` (hardware has no store form).
+finds no candidate and lowering reports an error. The **direction is fixed**: a
+shared-to-global transfer cannot use ``cp_async`` (hardware has no store form).
diff --git a/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst b/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst
index bdcf118430..5c1c94fe20 100644
--- a/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst
+++ b/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst
@@ -15,23 +15,23 @@
specific language governing permissions and limitations
under the License.
-copy_async → smem->tmem (tcgen05.cp)
-====================================
+tcgen05.cp
+==========
-The ``smem->tmem`` variant lowers a ``copy_async`` from **shared memory to tensor
-memory** (Blackwell ``tmem``) through the ``tcgen05_cp`` planner, which covers every
+``Tx.cuda.tile.tcgen05.cp`` transfers from shared memory to tensor
+memory (Blackwell ``tmem``) through the ``tcgen05_cp`` planner, which covers every
``tcgen05.cp`` shape. A shared **matrix descriptor** names the source tile; all
descriptor fields (ldo/sdo/swizzle) and the cp issue sequence are derived from
-the two buffer layouts. The dispatch issues only the copy; the caller signals
+the two buffer layouts. The lowerer issues only the copy; the caller signals
completion with ``tcgen05.commit``. Source:
``python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py``.
Shape selection
---------------
-- ``shape=`` config: forces that PTX shape (with ``multicast=`` where the shape
+- ``shape=`` qualifier: forces that PTX shape (with ``multicast=`` where the shape
has more than one legal qualifier).
-- No ``shape`` config: the planner tries each candidate **widest atom first**
+- No ``shape`` qualifier: the planner tries each candidate **widest atom first**
and takes the first whose plan validates against the layouts:
``128x256b`` → ``4x256b`` → ``128x128b`` → ``64x128b.warpx2::02_13`` →
``64x128b.warpx2::01_23`` → ``32x128b.warpx4``. All candidates but the
@@ -80,22 +80,14 @@ What it accepts
Two predicates — the memory-scope envelope and a single-thread exec scope; all
shape/layout validation happens in the planner with readable errors:
-.. code-block:: python
-
- # register_dispatch(..., variant="smem->tmem", priority=10, when=[
- predicate("validate_smem_tmem_copy", _validate_smem_tmem_copy),
- predicate("exec_scope", _single_thread_exec), # exec_scope == "thread"
- # ])
-
.. list-table::
:header-rows: 1
:widths: 22 78
* - Property
- Requirement
- * - target / priority
- - ``cuda`` target with ``tcgen05`` support (tested with ``sm_100a``);
- priority ``10``
+ * - target
+ - ``cuda`` target with ``tcgen05`` support (tested with ``sm_100a``)
* - scope
- **single thread** issues the copy
* - CTA group
@@ -113,10 +105,10 @@ shape/layout validation happens in the planner with readable errors:
the buffer's swizzle (if any)
Demonstration program
-----------------------
+---------------------
A warpgroup allocates 16 tmem columns, fills a ``32×16`` ``uint8`` shared tile,
-and copies it into tmem — no shape config, the planner infers
+and copies it into tmem — no shape qualifier, the planner infers
``32x128b.warpx4`` from the layouts (from ``test_tcgen05_cp.py``; readback /
dealloc tail elided):
@@ -135,7 +127,7 @@ dealloc tail elided):
tmem = Tx.decl_tensor([32, 16], "uint8", scope="tmem", allocated_addr=tmem_addr[0],
layout=TileLayout(S[(32, 16) : (1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane]))
if tid_in_wg == 0:
- Tx.tile.copy_async(tmem[0:32, 0:16], A_smem[0:32, 0:16], cta_group=1) # smem -> tmem
+ Tx.cuda.tile.tcgen05.cp(tmem[0:32, 0:16], A_smem[0:32, 0:16], cta_group=1) # smem -> tmem
# caller signals
Tx.ptx.tcgen05.commit.cta_group__1.mbarrier__arrive__one.shared__cluster.b64(
cp_mbar.ptr_to([0]))
@@ -174,7 +166,7 @@ bits, plus the lane half-word for lane-tiled atoms):
Tx.cast(t_addr[0] + t_addr_off + t_off, "uint32"),
smem_desc_add_16B_offset(desc_buf[0], init_off_16B + s_off))
-The dispatch emits **no** ``tcgen05.commit`` / ``wait`` — the caller commits
+The lowerer emits **no** ``tcgen05.commit`` / ``wait`` — the caller commits
against an mbarrier (as in the demo).
Generated CUDA
diff --git a/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst b/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst
index 6d62d56150..0bbe390d9d 100644
--- a/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst
+++ b/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst
@@ -15,35 +15,23 @@
specific language governing permissions and limitations
under the License.
-copy_async → tmem<->local (tcgen05.ld/st)
-=========================================
+tcgen05.ld / tcgen05.st
+=======================
-The ``tmem<->local`` variant lowers a ``copy_async`` between **tensor memory and
-registers** through the ``tcgen05_ldst`` implementation (Blackwell
+``Tx.cuda.tile.tcgen05.ld`` and ``Tx.cuda.tile.tcgen05.st`` transfer between
+tensor memory and registers through the ``tcgen05_ldst`` implementation (Blackwell
``tcgen05.ld`` / ``tcgen05.st``). It is warpgroup-collective:
the four warps cooperatively move a tensor-memory tile to/from their per-thread
-registers. One registration handles both directions — ``tmem → local`` lowers to
-``tcgen05.ld``, ``local → tmem`` to ``tcgen05.st`` — and the dispatch picks the
+registers. ``ld`` requires tensor-memory-to-local operands, while ``st``
+requires local-to-tensor-memory operands. The lowerer picks the
widest instruction shape the register layout matches. As with the other async
-variants, completion (``tcgen05.wait.ld`` / ``tcgen05.wait.st``) is the caller's. Source:
+instructions, completion (``tcgen05.wait.ld`` / ``tcgen05.wait.st``) is the caller's. Source:
``python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py``.
What it accepts
---------------
-A single registration (``variant="tmem<->local"``); direction is inferred at
-lowering:
-
-.. code-block:: python
-
- @register_dispatch("copy_async", "cuda", variant="tmem<->local", priority=10, when=[
- predicate("validate_copy_op", _is_valid_copy),
- predicate("exec_scope", exec_scope_ok, expected_scopes=["warpgroup"]),
- predicate("storage_scope", _scope_allowed,
- allowed_pairs=[("tmem", "local"), ("local", "tmem")]),
- ])
- # direction inferred in copy_tmem_local_impl:
- # src tmem + dst local -> "tmem2local" (ld); else "local2tmem" (st)
+The named instruction fixes the transfer direction. Both require warpgroup scope.
.. list-table::
:header-rows: 1
@@ -51,9 +39,8 @@ lowering:
* - Property
- Requirement
- * - target / priority
- - ``cuda`` target with ``tcgen05`` support (tested with ``sm_100a``);
- priority ``10``
+ * - target
+ - ``cuda`` target with ``tcgen05`` support (tested with ``sm_100a``)
* - scope
- **warpgroup** (``exec_scope_ok(expected_scopes=["warpgroup"])``) — the four
warps act together
@@ -70,7 +57,7 @@ lowering:
the lower or upper 16-lane sub-slab of each warp partition
Demonstration program
-----------------------
+---------------------
A warpgroup round-trips a ``128×8`` ``float16`` tile registers → tmem → registers
(the GPU smoke test ``test_copy_tmem2reg_async``; ``WIDTH = 8`` for ``width_32b=4``,
@@ -112,10 +99,10 @@ fp16):
A_local = A_reg.view(128, WIDTH, layout=local_view)
B_local = B_reg.view(128, WIDTH, layout=local_view)
# ... load A into A_reg, zero B_reg, cta_sync ...
- Tx.tile.wg.copy_async(tmem[:, :], A_local[:, :])
+ Tx.cuda.tile.tcgen05.st(tmem[:, :], A_local[:, :], scope="warpgroup")
Tx.ptx.tcgen05.wait__st.sync.aligned() # store (local -> tmem)
Tx.cuda.cta_sync()
- Tx.tile.wg.copy_async(B_local[:, :], tmem[:, :])
+ Tx.cuda.tile.tcgen05.ld(B_local[:, :], tmem[:, :], scope="warpgroup")
Tx.ptx.tcgen05.wait__ld.sync.aligned() # load (tmem -> local)
# ... write B_reg out; tcgen05.dealloc ...
@@ -125,7 +112,7 @@ Algorithm
**1. Infer direction.** ``tmem → local`` is a load (``tcgen05.ld``); ``local → tmem``
is a store (``tcgen05.st``).
-**2. Pick the instruction shape.** The dispatch matches the register layout against
+**2. Pick the instruction shape.** The lowerer matches the register layout against
``tcgen05_atom_layout`` for ``.16x64b`` / ``.16x128b`` / ``.16x256b``
(``_match_tcgen05_atom_layout``); the matched shape sets the column factor (2/4/8
fp32 columns) and the ``num`` count. If nothing matches, the implementation
@@ -154,12 +141,12 @@ Layout B is routed before the ordinary atom matching. Its logical
``(64, N)`` fragment is emitted as one physical ``.32x32b.x{N/2}``
instruction over all 128 lanes.
-The dispatch emits **no** wait. When synchronization is required, the caller
+The lowerer emits **no** wait. When synchronization is required, the caller
issues ``Tx.ptx.tcgen05.wait__ld.sync.aligned()`` or
``Tx.ptx.tcgen05.wait__st.sync.aligned()`` (as in the demo).
Selecting the upper F sub-slab
--------------------------------
+------------------------------
``sub_slab`` is part of the tensor-memory layout rather than an option on
``copy_async``. This keeps physical TMEM occupation explicit and lets two
@@ -184,9 +171,9 @@ Selecting the upper F sub-slab
layout=tmem_datapath_layout("F", 64, cols, sub_slab=1),
)
- Tx.tile.wg.copy_async(lower_frag, lower)
+ Tx.cuda.tile.tcgen05.ld(lower_frag, lower, scope="warpgroup")
Tx.ptx.tcgen05.wait__ld.sync.aligned()
- Tx.tile.wg.copy_async(upper_frag, upper)
+ Tx.cuda.tile.tcgen05.ld(upper_frag, upper, scope="warpgroup")
Tx.ptx.tcgen05.wait__ld.sync.aligned()
The lower view emits ``row=0`` and the upper view emits ``row=16`` for
@@ -210,11 +197,11 @@ Use the public allocation and fragment APIs together:
(64, N), "float32", M=128, cta_group=2)
frag = Tx.alloc_tcgen05_ldst_frag("32x32b", (64, N), "float32")
- Tx.tile.wg.copy_async(frag[:, :], accumulator[:, :])
+ Tx.cuda.tile.tcgen05.ld(frag[:, :], accumulator[:, :], scope="warpgroup")
Tx.ptx.tcgen05.wait__ld.sync.aligned()
# The inverse direction emits tcgen05.st with the same physical image.
- Tx.tile.wg.copy_async(accumulator[:, :], frag[:, :])
+ Tx.cuda.tile.tcgen05.st(accumulator[:, :], frag[:, :], scope="warpgroup")
Tx.ptx.tcgen05.wait__st.sync.aligned()
This is a single ``tcgen05.{ld,st}.32x32b.x{N/2}`` issue. ``N`` must be
@@ -260,7 +247,7 @@ How inputs change the algorithm
* - register layout
- matches a ``.16x64b`` / ``.16x128b`` / ``.16x256b`` atom → that shape; no
match → ``.32x32b`` only for the compatible 128-row/Layout D form used by
- this demo; otherwise dispatch fails
+ this demo; otherwise lowering fails
* - column width / dtype
- sets ``num`` (the ``.xN`` count) and the registers per thread
(``elem_per_32b = 32 / dtype_bits``)
diff --git a/docs/tirx/tile_primitives/copy_async/tma.rst b/docs/tirx/tile_primitives/copy_async/tma.rst
index de6aa836fe..93d3f7ad2e 100644
--- a/docs/tirx/tile_primitives/copy_async/tma.rst
+++ b/docs/tirx/tile_primitives/copy_async/tma.rst
@@ -15,12 +15,13 @@
specific language governing permissions and limitations
under the License.
-copy_async → tma_auto / tma_explicit
-=====================================
+TMA tensor instructions
+=======================
-The ``tma_auto`` and ``tma_explicit`` variants lower ``copy_async`` between
-global and shared memory to CUDA Tensor Memory Accelerator instructions. Both
-variants:
+``Tx.cuda.tile.cp_async_bulk_tensor_load`` and
+``Tx.cuda.tile.cp_async_bulk_tensor_store`` select the transfer direction.
+``cp_reduce_async_bulk_tensor`` selects a shared-to-global reduction.
+The ``descriptor_mode="auto"`` and ``descriptor_mode="explicit"`` planners:
* require the call to be in a single-thread execution scope (the caller
normally selects the issuing thread);
@@ -31,7 +32,7 @@ variants:
* require the source and destination regions to contain the same total number
of bytes. Their ranks and per-dimension shapes need not match.
-The variants differ only in how the TensorMap and issue count are planned.
+The descriptor modes differ only in how the TensorMap and issue count are planned.
``tma_auto``
------------
@@ -42,10 +43,10 @@ can be proven statically:
.. code-block:: python
- Tx.tile.copy_async(
+ Tx.cuda.tile.cp_async_bulk_tensor_load(
A_smem[:, :],
A[tile_m : tile_m + 64, tile_k : tile_k + 64],
- dispatch="tma_auto",
+ descriptor_mode="auto",
mbar=mbar.ptr_to([0]),
)
@@ -89,16 +90,16 @@ are only known at runtime.
* Buffer data plus ``elem_offset`` becomes the TensorMap base.
It never regroups, compresses, promotes, shrinks, or splits a copy. One
-``Tx.tile.copy_async`` call emits exactly one TMA instruction, so a caller must
+explicit descriptor call emits exactly one TMA instruction, so a caller must
explicitly tile a wider transfer:
.. code-block:: python
for atom in Tx.unroll(8):
- Tx.tile.copy_async(
+ Tx.cuda.tile.cp_async_bulk_tensor_store(
O[:, atom * 64 : (atom + 1) * 64],
O_smem[:, atom * 64 : (atom + 1) * 64],
- dispatch="tma_explicit",
+ descriptor_mode="explicit",
)
The sliced shared layout must canonicalize to a trivial box after its pointer
@@ -116,10 +117,10 @@ coordinates in ``{column, row0, row1, row2, row3}`` order:
.. code-block:: python
- Tx.tile.copy_async(
+ Tx.cuda.tile.cp_async_bulk_tensor_load(
K_smem[0:4, :],
K[0:1, :],
- dispatch="tma_explicit",
+ descriptor_mode="explicit",
mbar=mbar.ptr_to([0]),
gather4=[row0, row1, row2, row3],
)
@@ -136,10 +137,10 @@ the main operand's region and gather coordinates:
.. code-block:: python
- Tx.tile.copy_async(
+ Tx.cuda.tile.cp_async_bulk_tensor_load(
K_smem[0:4, :],
K_main[0:1, :],
- dispatch="tma_explicit",
+ descriptor_mode="explicit",
mbar=mbar.ptr_to([0]),
gather4=[row0, row1, row2, row3],
src_selector=[
@@ -148,6 +149,7 @@ the main operand's region and gather coordinates:
],
)
+Candidates must cover the entire tensor view; use a subview to change its base.
Conditions use first-true priority and the main Buffer is the default. Every
candidate gets its own validated and encoded TensorMap. Candidates may have
different bases, global shapes, and strides, but must have the same descriptor
@@ -178,9 +180,10 @@ Common configuration
- ``cta_group`` is ``1`` or ``2``. Group 2 requires SM100+ and a
precomputed shared mbarrier address. ``cta_mask`` is a uint16 integer or
TIR expression and is valid only for global-to-shared multicast.
- * - ``cache_hint``
- - ``evict_normal``, ``evict_first``, ``evict_last``, or a runtime
- ``uint64`` cache-policy expression.
+ * - ``cache_hint`` / ``cache_policy``
+ - The static hint is ``evict_normal``, ``evict_first``, or ``evict_last``.
+ A runtime ``uint64`` policy is a separate ``cache_policy`` operand.
+ Supplying both is rejected.
* - ``prefetch_tensormap``
- Deduplicated device-side prefetch of the main descriptor. Requires the
``warp_id_in_cta`` launch parameter.
@@ -190,8 +193,8 @@ Common configuration
``L2::256B``.
* - ``tma_dtype``
- ``tf32`` or ``tfloat32`` for a float32 descriptor conversion.
- * - ``use_tma_reduce``
- - Shared-to-global reduction operation: ``add``, ``min``, ``max``,
+ * - ``reduce_op``
+ - Operation for ``cp_reduce_async_bulk_tensor``: ``add``, ``min``, ``max``,
``inc``, ``dec``, ``and``, ``or``, or ``xor``.
* - ``oob``
- ``tma_explicit`` global-to-shared only. ``None`` is the default zero-fill
@@ -233,7 +236,7 @@ choose among alternative boxes or issue loops.
Completion
----------
-The dispatch emits the TMA instruction but no completion operation.
+The lowerer emits the TMA instruction but no completion operation.
Global-to-shared callers initialize the barrier, call
``arrive.expect_tx`` with the transferred byte count, and wait for its phase.
Shared-to-global callers use the bulk-group commit/wait operations required by
diff --git a/docs/tirx/tile_primitives/elementwise.rst b/docs/tirx/tile_primitives/elementwise.rst
index 4696af3b9f..2e98b7559c 100644
--- a/docs/tirx/tile_primitives/elementwise.rst
+++ b/docs/tirx/tile_primitives/elementwise.rst
@@ -15,37 +15,26 @@
specific language governing permissions and limitations
under the License.
-elementwise
-===========
-
-Covers ``cast``, ``fill``, the unary ops (``zero``, ``reciprocal``, ``sqrt``,
-``exp``, ``exp2``, ``log2``, ``silu``), the binary ops (``add``, ``sub``,
-``mul``, ``fdiv``, ``maximum``), and ``fma``. Every op registers **two** variants — ``reg`` and
-``smem`` — both at priority 10; the buffer-operand storage scope (all-local vs
-all-shared) is the mutually-exclusive discriminator. Scalar inputs do not have a
-storage scope. Each op is described by an
-``OpSpec`` (a ``parse`` that builds the destination + source list, optional dtype
-checks, and the scalar expression applied per element).
-
-This is the current CUDA elementwise set. Other constructors in the global
-``Tx.tile`` catalog, including ``minimum``, ``memset``, and ``select``, do not
-currently have these CUDA ``reg`` / ``smem`` variants; a call needs a variant
-registered by its selected target backend.
-
-.. list-table::
- :header-rows: 1
- :widths: 24 18 58
-
- * - Variant
- - Operands
- - Lowering
- * - :doc:`elementwise/reg`
- - all buffer operands local
- - partition induced by the local-buffer layout; op applied per-thread
- * - :doc:`elementwise/smem`
- - all buffer operands shared
- - synthesized ``[outer, threads, vec]`` partition; op applied per (vectorized)
- element
+Elementwise tensor instructions
+===============================
+
+``Tx.cuda.tile`` provides ``mov``, ``cvt``, ``add``, ``sub``, ``mul``, ``div``,
+``max``, ``fma``, ``sqrt``, ``ex2``, and ``lg2``. The instruction's tensor layout
+and execution scope determine how its scalar operation is repeated across the tile.
+``mov(dst, 0)`` initializes a tile, and ``div(dst, 1.0, src)`` computes a reciprocal.
+
+The mathematical compositions ``exp``, ``silu``, and the four
+``*_with_scale_bias`` functions live in ``Tx.cuda.tile.compose`` and remain opaque
+Calls until lowering. Their opcode category is ``tile_composite``.
+
+All tensor operands of an elementwise instruction must be local or all shared.
+The corresponding lowerer expands the same operation over the selected storage:
+
+* :doc:`elementwise/reg` uses the register layout's per-thread ownership.
+* :doc:`elementwise/smem` constructs a cooperative partition in shared memory.
+
+CUDA reduction algorithms and layout permutation are caller-owned; they have no
+replacement in ``compose``.
.. toctree::
:maxdepth: 1
diff --git a/docs/tirx/tile_primitives/elementwise/reg.rst b/docs/tirx/tile_primitives/elementwise/reg.rst
index 3267778b06..bceba0434e 100644
--- a/docs/tirx/tile_primitives/elementwise/reg.rst
+++ b/docs/tirx/tile_primitives/elementwise/reg.rst
@@ -18,13 +18,13 @@
elementwise → reg
=================
-The ``reg`` variant lowers an elementwise op (``sqrt``, ``exp``, ``add``,
+The register lowerer expands an elementwise instruction (``sqrt``, ``exp``, ``add``,
``fma``, …) when **all buffer operands use ``local`` scope**. Scalar inputs are
also accepted where the operation's authoring API permits them. Like the copy
-:doc:`../copy/reg` ``vec_auto`` path the partition is *induced* by the operands'
+:doc:`../copy/reg` instruction the partition is *induced* by the operands'
local-buffer layout — the thread axes are dropped, leaving each thread its
private bundle — and the op is applied to every element in that bundle. The
-variant name describes this local-buffer path; final register allocation
+lowering path uses local operands; final register allocation
remains a CUDA compiler decision. Source:
``python/tvm/backend/cuda/tile_primitive/elementwise/reg.py``.
@@ -54,9 +54,8 @@ What it accepts
* - Property
- Requirement
- * - target / scope / priority
- - ``cuda``; ``thread`` / ``warp`` / ``warpgroup`` / ``cta`` (all active);
- priority ``10``
+ * - target / scope
+ - ``cuda``; ``thread`` / ``warp`` / ``warpgroup`` / ``cta`` (all active)
* - operands
- **every buffer operand** in ``local``; scalar sources are allowed by
``fill``, the binary ops, and ``fma``
@@ -72,7 +71,7 @@ What it accepts
the destination region
Demonstration program
-----------------------
+---------------------
A warp takes the elementwise ``sqrt`` of a ``32×8`` ``float32`` local tile
(local layout ``S[(32,8):(1@laneid,1)]`` — lane ``i`` owns row ``i``):
@@ -93,14 +92,18 @@ A warp takes the elementwise ``sqrt`` of a ``32×8`` ``float32`` local tile
Tx.lane_id([32])
tid = Tx.thread_id([32])
A_smem = Tx.alloc_tensor((32, 8), "float32", scope="shared", layout=TileLayout(S[(32, 8)]))
- Tx.tile.warp.copy(A_smem[fs], A[fs])
+ for k in Tx.serial(8):
+ value = A[tid, k]
+ A_smem[tid, k] = value
Tx.cuda.cta_sync()
R = Tx.alloc_tensor((32, 8), "float32", scope="local", layout=r_layout)
- Tx.tile.warp.copy(R[fs], A_smem[fs])
- Tx.tile.warp.sqrt(R[fs], R[fs]) # elementwise reg dispatch
- Tx.tile.warp.copy(A_smem[fs], R[fs])
+ Tx.cuda.tile.ld(R[fs], A_smem[fs], scope="warp")
+ Tx.cuda.tile.sqrt(R[fs], R[fs], scope="warp") # elementwise reg dispatch
+ Tx.cuda.tile.st(A_smem[fs], R[fs], scope="warp")
Tx.cuda.cta_sync()
- Tx.tile.warp.copy(B[fs], A_smem[fs])
+ for k in Tx.serial(8):
+ value = A_smem[tid, k]
+ B[tid, k] = value
Algorithm
---------
diff --git a/docs/tirx/tile_primitives/elementwise/smem.rst b/docs/tirx/tile_primitives/elementwise/smem.rst
index 683c0687af..62d79bbc2a 100644
--- a/docs/tirx/tile_primitives/elementwise/smem.rst
+++ b/docs/tirx/tile_primitives/elementwise/smem.rst
@@ -18,10 +18,9 @@
elementwise → smem
==================
-The ``smem`` variant lowers an elementwise op (``sqrt``, ``exp``, ``add``,
+The shared-memory lowerer expands an elementwise instruction (``sqrt``, ``exp``, ``add``,
``fma``, …) when **all buffer operands are in shared memory**. Scalar inputs are
-also accepted where the operation's authoring API permits them. Like the copy
-:doc:`../copy/gmem_smem` ``vec_auto`` path it *synthesizes* a
+also accepted where the instruction permits them. It synthesizes an
``[outer, threads, vec]``
partition from the execution scope, then applies the op to each (vectorized)
element. Source:
@@ -51,9 +50,8 @@ What it accepts
* - Property
- Requirement
- * - target / scope / priority
- - ``cuda``; ``thread`` / ``warp`` / ``warpgroup`` / ``cta`` (all active);
- priority ``10``
+ * - target / scope
+ - ``cuda``; ``thread`` / ``warp`` / ``warpgroup`` / ``cta`` (all active)
* - operands
- **every buffer operand** (including the output) in ``shared*``; scalar
sources are allowed by ``fill``, the binary ops, and ``fma``
@@ -67,7 +65,7 @@ What it accepts
thread bound the scheduling chunk width
Demonstration program
-----------------------
+---------------------
A CTA takes the elementwise ``sqrt`` of a ``32×32`` ``float32`` shared tile
(adapted from ``test_unary.py`` — here a 256-thread CTA, so the partition is one
@@ -86,11 +84,19 @@ round):
Tx.cta_id([1])
Tx.warp_id([8])
Tx.lane_id([32])
- Tx.thread_id([256])
+ tid = Tx.thread_id([256])
A_smem = Tx.alloc_tensor((32, 32), "float32", scope="shared", layout=s_layout)
- Tx.tile.cta.copy(A_smem[full], A[full])
- Tx.tile.cta.sqrt(A_smem[full], A_smem[full]) # elementwise smem dispatch
- Tx.tile.cta.copy(A[full], A_smem[full])
+ for k in Tx.serial(4):
+ index = tid + k * 256
+ value = A[index // 32, index % 32]
+ A_smem[index // 32, index % 32] = value
+ Tx.cuda.cta_sync()
+ Tx.cuda.tile.sqrt(A_smem[full], A_smem[full], scope="cta") # elementwise smem dispatch
+ Tx.cuda.cta_sync()
+ for k in Tx.serial(4):
+ index = tid + k * 256
+ value = A_smem[index // 32, index % 32]
+ A[index // 32, index % 32] = value
Algorithm
---------
@@ -99,7 +105,7 @@ Algorithm
(inputs, output, the op); the predicate confirms every buffer operand is shared.
**2. Synthesize the partition** from the scope's **thread count** (as
-:doc:`../copy/gmem_smem` does): split the region into ``[outer, threads, vec]``.
+the cooperative transfer planner does): split the region into ``[outer, threads, vec]``.
The candidate width must divide the elements per thread and every operand's
logical innermost region extent. ``_max_layout_vec`` does not inspect physical
layout strides when choosing this width. For the dense identity layout in this
diff --git a/docs/tirx/tile_primitives/gemm.rst b/docs/tirx/tile_primitives/gemm.rst
index 2a91af4a9e..ebc9630932 100644
--- a/docs/tirx/tile_primitives/gemm.rst
+++ b/docs/tirx/tile_primitives/gemm.rst
@@ -15,17 +15,17 @@
specific language governing permissions and limitations
under the License.
-gemm
-====
+mma_sync
+========
-``gemm`` computes ``D = alpha·A@B + beta·C`` as a fully-unrolled nest of
+``Tx.cuda.tile.mma_sync`` computes ``D = alpha·A@B + beta·C`` as a fully-unrolled nest of
warp-collective ``mma.sync.aligned.m16n8k{16,8}`` instructions. A warp call is
the usual form. A full warpgroup or CTA scope is also accepted; each contained
warp executes the fragment program, and the operand thread-axis tiling determines
whether those warps own distinct output tiles or repeat the same tile. A and B
fragments and the C/D accumulators **all live in registers** — the caller stages
A and B into register fragments first (typically via :doc:`copy/ldstmatrix`).
-The dispatch tiles M/N/K into ``m16n8k`` atoms and emits one ``mma`` per output
+The lowerer tiles M/N/K into ``m16n8k`` atoms and emits one ``mma`` per output
tile, accumulating over K in place. Source:
``python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py``. (For the
Blackwell async tensor-core path see :doc:`gemm_async`.)
@@ -33,25 +33,14 @@ Blackwell async tensor-core path see :doc:`gemm_async`.)
What it accepts
---------------
-.. code-block:: python
-
- # register_dispatch("gemm", "cuda", priority=10, when=[
- predicate("full_active_lanes", _full_active_lanes), # complete warp(s), un-narrowed
- predicate("no_replica", _no_replica), # no broadcast axes on D/A/B/C
- # ])
- # in the impl:
- for buf, name in ((D, "D"), (A, "A"), (B, "B"), (C, "C")):
- if buf.scope() != "local":
- fail(f"gemm mma requires {name} in register (local) scope, got {buf.scope()}")
-
.. list-table::
:header-rows: 1
:widths: 22 78
* - Property
- Requirement
- * - target / scope / priority
- - ``cuda``; priority ``10``. The predicate has no scope-kind allowlist: it
+ * - target / scope
+ - ``cuda``. The predicate has no scope-kind allowlist: it
requires every axis present in ``sctx.intra`` to be a complete,
zero-offset ``laneid`` / ``wid_in_wg`` / ``warpid`` axis. This admits the
normal warp / warpgroup / CTA call sites and rejects unrecognized axes
@@ -73,7 +62,7 @@ What it accepts
- ``alpha == 1.0``; ``beta ∈ {0.0, 1.0}`` (0 → ``D = A@B``; 1 → ``D = A@B + C``)
Demonstration program
-----------------------
+---------------------
A single warp computes ``D[16,8] = A[16,16] @ B[16,8]`` in ``float16`` (f32
accumulate) — one ``m16n8k16`` atom (from ``test_gemm_mma_m16n8k_.py``):
@@ -111,8 +100,8 @@ accumulate) — one ``m16n8k16`` atom (from ``test_gemm_mma_m16n8k_.py``):
for s in Tx.unroll(4):
kp, kHi = s % 2, s // 2
B_reg[s] = B_g[2 * (lane % 4) + kp + 8 * kHi, lane // 4]
- Tx.tile.warp.gemm(
- D_f, A_f, B_f, D_f, transpose_A=False, transpose_B=False, alpha=1.0, beta=0.0
+ Tx.cuda.tile.mma_sync(
+ D_f, A_f, B_f, D_f, transpose_A=False, transpose_B=False, alpha=1.0, beta=0.0, scope="warp"
)
D_reg = D_f.local(4) # write the 4 result regs out
for s in Tx.unroll(4):
@@ -122,7 +111,7 @@ accumulate) — one ``m16n8k16`` atom (from ``test_gemm_mma_m16n8k_.py``):
Algorithm
---------
-**1. Tile and fragment-group.** The dispatch slices each operand's layout to its
+**1. Tile and fragment-group.** The lowerer slices each operand's layout to its
region and, for each candidate instruction (``m16n8k16`` then ``m16n8k8``), tries to
group the operand sub-layouts (``D_M, D_N, A_M, A_K, B_K, B_N, C_*``) into the fixed
m16n8k frame, anchoring A/C on D's M, B/C on D's N, and B on A's K. The first
diff --git a/docs/tirx/tile_primitives/gemm_async.rst b/docs/tirx/tile_primitives/gemm_async.rst
index b601fb988d..da6dc2aceb 100644
--- a/docs/tirx/tile_primitives/gemm_async.rst
+++ b/docs/tirx/tile_primitives/gemm_async.rst
@@ -15,10 +15,10 @@
specific language governing permissions and limitations
under the License.
-gemm_async
-==========
+tcgen05.mma / tcgen05.mma_block_scale
+=====================================
-``gemm_async`` lowers a matrix multiply to the **Blackwell asynchronous
+``Tx.cuda.tile.tcgen05.mma`` lowers a matrix multiply to the **Blackwell asynchronous
tensor-core** instruction ``tcgen05.mma``. B, and normally A, live in **shared
memory** and are named by 64-bit *matrix descriptors*; A can instead use the
tensor-memory operand path. The accumulator lives in **tensor memory**, and one
@@ -26,7 +26,7 @@ thread launches the MMA, which runs asynchronously. A
single-thread call site has already selected that issuer; a warp-scoped call uses
``elect_sync`` internally. The caller signals completion with ``tcgen05.commit``
against an mbarrier. It also
-supports **block-scaled** low precision (fp8 / fp4 with per-block scale factors
+provides ``Tx.cuda.tile.tcgen05.mma_block_scale`` for block-scaled low precision (fp8 / fp4 with per-block scale factors
``SFA`` / ``SFB`` in tensor memory). Source:
``python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py``. (For the
synchronous warp-register path see :doc:`gemm`.)
@@ -36,23 +36,15 @@ What it accepts
A single predicate — single-thread or warp scope:
-.. code-block:: python
-
- # register_dispatch("gemm_async", "cuda", priority=10, when=[
- predicate("single_thread_or_warp",
- lambda op, sctx: (single_thread(op, sctx) or sctx.is_warp,
- f"unsupported exec_scope {sctx.exec_scope}"))
- # ])
-
.. list-table::
:header-rows: 1
:widths: 22 78
* - Property
- Requirement
- * - target / scope / priority
+ * - target / scope
- ``cuda`` target with ``tcgen05`` support (this implementation is tested
- with ``sm_100a``); **single thread or warp**; priority ``10``
+ with ``sm_100a``); **single thread or warp**
* - operands
- B is in **shared**; A is in **shared** or, for the TMEM-A path, **tmem**;
the accumulator C/D is in **tmem** (``float32``)
@@ -84,14 +76,14 @@ A single predicate — single-thread or warp scope:
``M * cta_group``; ``mma_n`` must divide N exactly; the resulting
physical tile is also checked by the tcgen05 hardware-shape validator
* - instruction descriptor
- - dense MMA always encodes its descriptor in the dispatcher and rejects
+ - dense MMA always encodes its descriptor during lowering and rejects
``descI``. Block-scaled MMA may accept a pre-encoded uint32 ``descI``
* - layout forms
- swizzled shared layouts, no-swizzle packed shared layouts, regular tmem
accumulators, and FlashMLA-style packed ``N/2`` tmem accumulator layouts
Demonstration program
-----------------------
+---------------------
One selected thread in a warpgroup multiplies a ``128×64`` × ``64×128``
``float16`` tile (f32 accumulate) into a tmem accumulator, after TMA-loading A/B
@@ -115,7 +107,7 @@ into shared (from
layout=TileLayout(S[(128, 512) : (1 @ TLane, 1 @ TCol)]))
# ... TMA-load A_smem, B_smem from global, wait ...
if tid_in_wg == 0:
- Tx.tile.gemm_async(tmem[0:128, 256:384], A_smem[1:2, :, :], B_smem[2:3, :, :], dispatch="tcgen05")
+ Tx.cuda.tile.tcgen05.mma(tmem[0:128, 256:384], A_smem[1:2, :, :], B_smem[2:3, :, :])
# caller signals completion
Tx.ptx.tcgen05.commit.cta_group__1.mbarrier__arrive__one.shared__cluster.b64(
mma_mbar.ptr_to([0]))
@@ -165,7 +157,7 @@ destination's existing value.
For **block-scaled** fp8/fp4 the chain gains ``.block_scale.scale_vec::<n>X``
with two extra tmem addresses — ``SFA`` / ``SFB`` — and the scale-factor dtypes;
the instruction descriptor is encoded at runtime unless the caller supplies
-``descI``. As with the other async ops, the dispatch emits **no** completion —
+``descI``. As with the other async ops, the lowerer emits **no** completion —
the caller's ``tcgen05.commit`` + mbarrier wait close it.
For row-0 schedules, the lowering folds ``Tx.cuda.get_tmem_addr(base, 0, col)`` to
@@ -173,7 +165,7 @@ For row-0 schedules, the lowering folds ``Tx.cuda.get_tmem_addr(base, 0, col)``
FlashMLA kernels while preserving the helper call for nonzero row offsets.
On the dense path, ``weight_stationary=True`` with ``cta_group=1`` selects the
``tcgen05.mma.ws`` ABI. The PTX table has no ``.ws.cta_group::2`` form. The
-dispatcher also infers this mode from the packed M=64 Layout-E accumulator and
+lowerer also infers this mode from the packed M=64 Layout-E accumulator and
rejects layout/flag combinations that would place tensor-memory rows
incorrectly. Block-scaled MMA uses its own instruction chain and does not append
``.ws``.
@@ -195,16 +187,15 @@ Allocate and read a Layout B result as follows:
accumulator = tmem_pool.alloc_tcgen05_mma_D(
(64, N), "float32", M=128, cta_group=2)
- Tx.tile.gemm_async(
+ Tx.cuda.tile.tcgen05.mma(
accumulator[:, :],
A_smem[:, :],
B_smem[:, :],
- dispatch="tcgen05",
cta_group=2,
)
frag = Tx.alloc_tcgen05_ldst_frag("32x32b", (64, N), "float32")
- Tx.tile.wg.copy_async(frag[:, :], accumulator[:, :])
+ Tx.cuda.tile.tcgen05.ld(frag[:, :], accumulator[:, :], scope="warpgroup")
Tx.ptx.tcgen05.wait__ld.sync.aligned()
The fragment is a logical ``(64, N)`` view of one physical
diff --git a/docs/tirx/tile_primitives/permute_layout.rst b/docs/tirx/tile_primitives/permute_layout.rst
deleted file mode 100644
index 68864178c0..0000000000
--- a/docs/tirx/tile_primitives/permute_layout.rst
+++ /dev/null
@@ -1,179 +0,0 @@
-.. Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
-.. http://www.apache.org/licenses/LICENSE-2.0
-
-.. Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
-
-permute_layout
-==============
-
-``permute_layout`` rearranges a warp's data from a source ``TileLayout`` to a
-destination one — typically an in-place transpose. The single CUDA variant
-(``warp_xor_swizzle``) stages each lane's elements through a ``local`` buffer and
-writes them back under the destination layout. It chooses a per-lane **XOR
-swizzle** using the shared-memory bank model; for shared operands, the selected
-iteration order makes both read and write phases bank-conflict-free. A
-``warp_sync`` separates the two phases so the op is safe even when source and
-destination alias. Source:
-``python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py``.
-
-What it accepts
----------------
-
-The implementation first runs its ``_why_reject`` validator:
-
-.. code-block:: python
-
- if sctx.scope_kind != "warp": return "scope is not 'warp'"
- if "threadIdx.y" in launch or "threadIdx.z" in launch: return "multi-dim threadIdx"
- if src_buf.dtype != dst_buf.dtype: return "dtype mismatch"
- if src_ext_i != dst_ext_i: return "extent mismatch"
- if dtype_bytes not in (1, 2, 4, 8, 16): return "unsupported dtype byte width"
- if not isinstance(src_buf.layout, TileLayout): return "src not a plain TileLayout"
- if not isinstance(dst_buf.layout, TileLayout): return "dst not a plain TileLayout"
- # + layouts must slice, regroup, and define bijections on the slice
- # + V % 32 == 0 and P = V/32 is a power of two in [1, 32]
- # + _choose_xor_k must find a valid k (else fail)
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target / scope / priority
- - ``cuda``; **warp** scope only; priority ``20``
- * - operands
- - equal dtype, equal (compile-time) extents; both plain ``TileLayout`` (no
- swizzle wrapper); dtype byte width ∈ {1, 2, 4, 8, 16}. The dispatcher has
- no storage-scope predicate; shared 32/64-bit operands use direct PTX
- ``ld.shared`` / ``st.shared``, while other accepted cases use ordinary
- buffer loads and stores
- * - launch / volume
- - one-dimensional ``threadIdx``; slice volume ``V`` is divisible by 32 and
- ``P = V/32`` is a power of two in ``[1, 32]``
- * - layout mapping
- - after slicing and regrouping, source and destination describe the same
- iteration extents and each is a bijection on the slice
- * - bank-freedom
- - ``_choose_xor_k`` must find an XOR-bit count ``k ∈ [0, log2(P)]`` that makes
- **both** phases bank-conflict-free, else the dispatch declines (``fail``)
-
-Demonstration program
-----------------------
-
-A warp transposes the inner ``4×32`` block of a scale-factor tile — source layout
-strides ``(…, 32, 1)``, destination ``(…, 1, 4)`` — for two pipeline stages (the
-canonical SF-transpose, from ``test_permute_layout.py``):
-
-.. code-block:: python
-
- pipe, blk, dtype = 2, 128, "float32"
- high = 1
- shape = (pipe, high, 4, 32)
- pre = TileLayout(S[shape : (blk, 128, 32, 1)]) # source
- post = TileLayout(S[shape : (blk, 128, 1, 4)]) # destination (4↔32 transposed)
-
-
- @Tx.function
- def f(A_buf: Tx.Tensor(shape, dtype, layout=pre), B_buf: Tx.Tensor(shape, dtype, layout=post)):
-
- Tx.device_entry()
- Tx.cta_id([1])
- Tx.thread_id([32])
- for s in Tx.serial(0, pipe):
- Tx.tile.warp.permute_layout(B_buf[s, 0:1, 0:4, 0:32], A_buf[s, 0:1, 0:4, 0:32])
-
-Algorithm
----------
-
-**1. Align the two layouts.** Both layouts are sliced to the region and
-canonicalized; if their shards differ in structure (a linear layout collapses to 1-D
-under canon, a transposed one keeps its multi-dim shape) the source is regrouped to
-the destination's shape. From the destination shard come the iteration ``extent``
-and the per-side strides ``src_str`` / ``dst_str``. ``P`` = elements per lane =
-``prod(extent) / 32`` (here ``4``).
-
-**2. Choose the XOR swizzle.** ``_choose_xor_k`` simulates the shared-memory bank
-pattern at shard granularity for ``k = 0, 1, … log2(P)`` and picks the smallest
-``k`` whose ``shift`` / ``mask`` make *both* phases conflict-free (here ``shift = 3``,
-``mask = 3``).
-
-**3. Emit two local-staged phases.** Each lane reads its ``P`` elements through
-the source layout into a local temporary (the swizzle permutes which slot holds
-which iteration), a ``warp_sync`` follows, then the elements are written back through the
-destination layout:
-
-.. code-block:: python
-
- regs = Tx.alloc_tensor((P,), dtype, scope="local")
- for r in Tx.unroll(0, P): # read via src layout
- j = r ^ ((lane_id >> shift) & mask)
- idx = decompose(lane_id + j * 32, extent)
- regs[r] = src_buf[project(idx, src_st)]
- Tx.cuda.warp_sync()
- for r in Tx.unroll(0, P): # write via dst layout
- j = r ^ ((lane_id >> shift) & mask)
- idx = decompose(lane_id + j * 32, extent)
- dst_buf[project(idx, dst_st)] = regs[r]
- Tx.cuda.warp_sync()
-
-Generated TIRx IR
------------------
-
-.. code-block:: python
-
- regs[r] = A_buf[s*128 + (r ^ ((tx >> 3) & 3)) % 4 * 32 + tx] # phase 1 (src order)
- Tx.cuda.warp_sync()
- B_buf[s*128 + tx * 4 + (r ^ ((tx >> 3) & 3)) % 4] = regs[r] # phase 2 (dst order)
- Tx.cuda.warp_sync()
-
-Generated CUDA
---------------
-
-.. code-block:: c++
-
- alignas(64) float regs_ptr[4];
- regs_ptr[0] = A_buf_ptr[(s*128) + (((0 ^ ((threadIdx.x >> 3) & 3)) & 3) * 32) + threadIdx.x];
- regs_ptr[1] = A_buf_ptr[(s*128) + (((1 ^ ((threadIdx.x >> 3) & 3)) & 3) * 32) + threadIdx.x];
- regs_ptr[2] = A_buf_ptr[(s*128) + (((2 ^ ((threadIdx.x >> 3) & 3)) & 3) * 32) + threadIdx.x];
- regs_ptr[3] = A_buf_ptr[(s*128) + (((3 ^ ((threadIdx.x >> 3) & 3)) & 3) * 32) + threadIdx.x];
- __syncwarp();
- // ... 4 transposed writes into B_buf_ptr, then __syncwarp();
-
-Each lane owns column ``threadIdx.x`` and stages its 4 rows through ``regs``; the
-``(threadIdx.x >> 3)`` XOR rotates the register order per lane-group of 8 so the
-write phase hits distinct banks. Verified on ``sm_100a`` — the ``4×32`` block is
-transposed for every pipeline stage.
-
-How inputs change the algorithm
--------------------------------
-
-.. list-table::
- :header-rows: 1
- :widths: 28 72
-
- * - input
- - effect
- * - layout strides (the permutation)
- - define ``extent`` / ``src_str`` / ``dst_str`` and hence ``P`` and the
- per-element index math (the transpose pattern)
- * - dtype byte width
- - feeds the bank simulation in ``_choose_xor_k``; no dtype width alone
- guarantees a match, and any layout whose two phases cannot both be made
- bank-free causes the dispatch to fail
- * - chosen ``k``
- - sets ``shift`` / ``mask`` of the XOR swizzle (``k = 0`` ⇒ no swizzle)
- * - ``P`` (= elements/lane)
- - the number of staged local elements and unrolled iterations per phase
diff --git a/docs/tirx/tile_primitives/reduction.rst b/docs/tirx/tile_primitives/reduction.rst
index b752a673d3..5877d944f5 100644
--- a/docs/tirx/tile_primitives/reduction.rst
+++ b/docs/tirx/tile_primitives/reduction.rst
@@ -15,37 +15,16 @@
specific language governing permissions and limitations
under the License.
-reduction
-=========
+Reduction Instructions
+======================
-Covers ``sum``, ``max``, ``min`` (reduce over ``axes``). All three operations
-register ``local`` and ``shared`` variants at priority 10, discriminated by
-operand storage scope. On CUDA SM100 or newer, ``sum`` additionally registers
-``packed_add_sum`` and ``max`` / ``min`` register ``3input_maxmin`` at priority
-20. The implementation file and detailed page group those two names as the
-SM100 packed paths; ``sm100_packed`` is not itself a dispatch name.
+CUDA reduction algorithms belong to the caller. Use scalar or packed
+arithmetic, warp shuffle instructions, shared storage, and explicit
+synchronization as appropriate for the kernel's decomposition.
-.. list-table::
- :header-rows: 1
- :widths: 26 14 60
+Trainium exposes the native ``T.trn.tile.tensorreduce`` instruction and the
+fused ``tensorscalar_reduce`` and ``activation_reduce`` instructions. Their
+axes and opcode are static qualifiers; tensors and optional reduction
+workspace are ordinary Call arguments.
- * - Variant
- - Prio
- - Lowering
- * - :doc:`reduction/local`
- - 10
- - local-buffer src/dst; sequential thread reduction (+ optional warp shuffle)
- * - :doc:`reduction/shared`
- - 10
- - shared src/dst; adaptive group-size ``__shfl_xor`` tree
- * - :doc:`reduction/sm100_packed` (``packed_add_sum`` / ``3input_maxmin``)
- - 20
- - CUDA SM100+ thread-scope fp32 ≥8: packed ``add.f32x2`` /
- ``max3``/``min3``
-
-.. toctree::
- :maxdepth: 1
-
- reduction/local
- reduction/shared
- reduction/sm100_packed
+See :doc:`../tile_primitives` for the complete contract.
diff --git a/docs/tirx/tile_primitives/reduction/local.rst b/docs/tirx/tile_primitives/reduction/local.rst
deleted file mode 100644
index 031b247893..0000000000
--- a/docs/tirx/tile_primitives/reduction/local.rst
+++ /dev/null
@@ -1,157 +0,0 @@
-.. Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
-.. http://www.apache.org/licenses/LICENSE-2.0
-
-.. Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
-
-reduction → local
-=================
-
-The ``local`` variant lowers a reduction (``sum`` / ``max`` / ``min``) when **both
-source and destination are ``local`` buffers**. At thread scope it is a
-plain sequential reduction over each thread's own elements. Warp scope has two
-layout-driven paths: a specialized ``laneid`` shard-to-replica reduction, or a
-general local-axis reduction that can optionally add cross-lane shuffle steps.
-Source:
-``python/tvm/backend/cuda/tile_primitive/reduction/local.py``.
-
-What it accepts
----------------
-
-.. code-block:: python
-
- @register_dispatch(op_name, "cuda", variant="local", priority=10, when=[
- predicate("storage_scope", _match_reduction_storage_scope, expected_scope=["local"]),
- predicate("local_valid", validate_reduction_local),
- ])
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target / priority
- - ``cuda``; priority ``10``
- * - operand scope
- - src **and** dst in ``local``, equal dtype
- * - exec scope
- - ``thread`` (sequential and thread-local); ``warp`` / ``warpgroup``
- require valid non-swizzled ``TileLayout`` values. At warp scope a
- ``laneid`` shard→replica pattern automatically selects the specialized
- shuffle path; otherwise ``thread_reduce=True`` optionally adds shuffles
- to the general path. Warpgroup scope rejects ``thread_reduce=True``
- * - shape
- - at thread scope the predicate does not inspect axes or compare source and
- destination sizes; axis analysis and loop construction happen during
- emit. Wide-scope view reductions require matching spatial layout
- dimensions; reduced dimensions have dst local extent 1
-
-Demonstration program
-----------------------
-
-A single thread reduces a 4-element ``float32`` local vector to a scalar
-(thread-wise path, from ``test_reduction.py``):
-
-.. code-block:: python
-
- @Tx.function
- def test_func(
- A: Tx.Tensor([4], "float32", layout=TileLayout(S[4,])),
- B: Tx.Tensor([1], "float32", layout=TileLayout(S[1,])),
- ):
-
- Tx.device_entry()
- Tx.cta_id([1])
- Tx.thread_id([1])
- A_local = Tx.alloc_tensor([4], "float32", scope="local")
- B_local = Tx.alloc_tensor([1], "float32", scope="local")
- for i in Tx.serial(4):
- A_local[i] = A[i]
- Tx.tile.sum(B_local, A_local, accum=False) # reduction local dispatch
- B[0] = B_local[0]
-
-(4 < 8 elements, so this stays on ``local`` rather than the
-:doc:`sm100_packed` ``packed_add_sum`` / ``3input_maxmin`` fast paths.)
-
-Algorithm
----------
-
-**Thread-wise** (``_emit_reduction_local_thread_wise``): a spatial loop over the
-output positions, each initialized to the op's identity (unless ``accum``), then a
-reduction loop accumulating the source — no cross-thread communication:
-
-.. code-block:: python
-
- for spa in Tx.serial(spatial_len):
- if not accum: dst[spa] = identity
- for red in Tx.serial(reduction_len):
- dst[spa] = op(dst[spa], src[spa, red])
-
-**Specialized shard→replica shuffle** (``_gen_warp_shuffle_reduce``): when the
-source has a full-span ``laneid`` shard and the destination has a power-of-two
-``laneid`` replica, the implementation copies each lane's corresponding local
-values and applies ``Tx.cuda.warp_reduce`` across the replica width. This path is
-selected automatically, independently of ``thread_reduce``; it does not first
-run the general local-axis loop. This specialized path does not branch on the
-``accum`` argument, so it overwrites the destination with the shuffle result.
-
-**General warp/warpgroup view path** (``_emit_reduction_local_view``): reduces
-the source's local reduction axes into each destination position. At warp scope,
-``thread_reduce=True`` additionally emits explicit
-``gpu_warp_shuffle_xor`` steps using ``__activemask()``. Warpgroup scope supports
-the local part only.
-
-Generated TIRx IR
------------------
-
-For the 4-element thread reduction:
-
-.. code-block:: python
-
- for spa in Tx.serial(1):
- dst[...] = Tx.float32(0)
- for red in Tx.serial(4):
- dst[...] = dst[...] + src[...] # op = sum
-
-Generated CUDA
---------------
-
-.. code-block:: c++
-
- for (int red = 0; red < 4; ++red)
- B_local_ptr[0] = B_local_ptr[0] + A_local_ptr[red];
-
-(Verified on ``sm_100a`` — ``B == sum(A)``.)
-
-How inputs change the algorithm
--------------------------------
-
-.. list-table::
- :header-rows: 1
- :widths: 28 72
-
- * - input
- - effect
- * - op
- - ``sum`` → ``+``, ``max`` → ``max``, ``min`` → ``min`` (and the identity)
- * - exec scope
- - ``thread`` → sequential; matching warp shard→replica layout → specialized
- ``warp_reduce``; otherwise warp/warpgroup use the general local view path,
- with optional warp shuffles only when ``thread_reduce=True``
- * - axes / shape
- - set the spatial vs reduction loop extents
- * - accum
- - thread-wise and general view paths reuse the old dst value as described
- above; the specialized shard→replica path ignores this flag
diff --git a/docs/tirx/tile_primitives/reduction/shared.rst b/docs/tirx/tile_primitives/reduction/shared.rst
deleted file mode 100644
index c15893642b..0000000000
--- a/docs/tirx/tile_primitives/reduction/shared.rst
+++ /dev/null
@@ -1,143 +0,0 @@
-.. Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
-.. http://www.apache.org/licenses/LICENSE-2.0
-
-.. Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
-
-reduction → shared
-==================
-
-The ``shared`` variant lowers a reduction (``sum`` / ``max`` / ``min``) when
-**both source and destination are shared** memory. At CTA / warpgroup / warp scope
-it partitions the threads into groups — one group per output position — has each
-thread gather a chunk of the reduction axis, then folds the group with an adaptive
-``__shfl_xor`` tree. Source:
-``python/tvm/backend/cuda/tile_primitive/reduction/shared.py``.
-
-What it accepts
----------------
-
-.. code-block:: python
-
- @register_dispatch(op_name, "cuda", variant="shared", priority=10, when=[
- predicate("storage_scope", _match_reduction_storage_scope, expected_scope=["shared*"]),
- predicate("shared_valid", validate_reduction_shared),
- ])
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target / priority
- - ``cuda``; priority ``10``
- * - operand scope
- - src **and** dst in ``shared*``, equal dtype
- * - exec scope
- - ``cta`` / ``warpgroup`` / ``warp`` (shuffle tree) or ``thread`` (sequential)
- * - thread binding
- - ``threadIdx.x`` present and **1-D** (no ``threadIdx.y`` / ``z``). The
- validator does not test divisibility by the selected shuffle-group size
- * - shape
- - ``dst`` size equals the source's spatial extent (product of the non-reduced
- dims)
-
-Demonstration program
-----------------------
-
-A 32-thread CTA reduces each row of a ``4×8`` ``float32`` shared tile (reduce
-axis ``-1``) to a ``4``-vector (from ``test_reduction.py``):
-
-.. code-block:: python
-
- @Tx.function
- def test_reduction(
- A: Tx.Tensor((4, 8), "float32", layout=TileLayout(S[4, 8])),
- B: Tx.Tensor((4,), "float32", layout=TileLayout(S[4,])),
- ):
-
- Tx.device_entry()
- Tx.cta_id([1])
- Tx.thread_id([32])
- A_smem = Tx.alloc_tensor((4, 8), "float32", scope="shared", layout=TileLayout(S[(4, 8)]))
- B_smem = Tx.alloc_tensor((4,), "float32", scope="shared", layout=TileLayout(S[(4,)]))
- Tx.tile.cta.copy(A_smem, A)
- Tx.cuda.cta_sync()
- Tx.tile.cta.sum(B_smem, A_smem, axes=(-1,), accum=False) # reduction shared dispatch
- Tx.cuda.cta_sync()
- Tx.tile.cta.copy(B, B_smem)
-
-Algorithm
----------
-
-**1. Choose the group size.** ``group_size = min(next_power_of_2(reduction_len),
-32, thread_cnt)`` — here ``reduction_len = 8`` ⇒ ``group_size = 8``. Each group of 8
-lanes cooperatively reduces one row; the CTA processes the 4 rows in parallel.
-
-**2. Gather + shuffle tree.** Each lane loads its slice of the reduction axis into a
-register, then ``log2(group_size)`` ``shfl_xor`` steps (masks ``1, 2, 4``) fold the
-group; lane 0 of each group writes the result, followed by a barrier:
-
-.. code-block:: python
-
- mask = Tx.gpu_warp_activemask()
- for i in range(n_shuffles): # n_shuffles = log2(group_size)
- thread_data[0] = op(thread_data[0],
- Tx.gpu_warp_shuffle_xor(mask, thread_data[0], 1 << i, group_size, 32))
-
-(``warp`` uses ``warp_sync``; ``warpgroup`` ``warpgroup_sync(8)``; ``cta``
-``cta_sync``. Thread scope is instead the sequential loop of :doc:`local`.)
-
-Generated TIRx IR
------------------
-
-.. code-block:: python
-
- thread_data[0] = thread_data[0] + Tx.gpu_warp_shuffle_xor(Tx.gpu_warp_activemask(), thread_data[0], 1, 8, 32)
- thread_data[0] = thread_data[0] + Tx.gpu_warp_shuffle_xor(Tx.gpu_warp_activemask(), thread_data[0], 2, 8, 32)
- thread_data[0] = thread_data[0] + Tx.gpu_warp_shuffle_xor(Tx.gpu_warp_activemask(), thread_data[0], 4, 8, 32)
-
-Generated CUDA
---------------
-
-.. code-block:: c++
-
- thread_data_ptr[0] = thread_data_ptr[0] + __shfl_xor_sync(__activemask(), thread_data_ptr[0], 1, 8);
- thread_data_ptr[0] = thread_data_ptr[0] + __shfl_xor_sync(__activemask(), thread_data_ptr[0], 2, 8);
- thread_data_ptr[0] = thread_data_ptr[0] + __shfl_xor_sync(__activemask(), thread_data_ptr[0], 4, 8);
-
-(Verified on ``sm_100a`` — each ``B[r] == sum(A[r, :])``. The shuffle width ``8``
-is the group size, not the full warp.)
-
-How inputs change the algorithm
--------------------------------
-
-.. list-table::
- :header-rows: 1
- :widths: 28 72
-
- * - input
- - effect
- * - op
- - ``sum`` → ``+`` shuffle tree; ``max`` / ``min`` → the corresponding combine
- * - reduction length / thread count
- - set ``group_size = min(next_pow2(reduction_len), 32, thread_cnt)`` and hence
- the number of shuffle steps; ``spatial_par`` uses integer division
- ``thread_cnt // group_size``
- * - exec scope
- - ``cta`` / ``warpgroup`` / ``warp`` → shuffle tree (different sync); ``thread``
- → sequential loop
- * - accum
- - ``True`` combines the reduced value with the old dst
diff --git a/docs/tirx/tile_primitives/reduction/sm100_packed.rst b/docs/tirx/tile_primitives/reduction/sm100_packed.rst
deleted file mode 100644
index f5047fa7e8..0000000000
--- a/docs/tirx/tile_primitives/reduction/sm100_packed.rst
+++ /dev/null
@@ -1,157 +0,0 @@
-.. Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
-.. http://www.apache.org/licenses/LICENSE-2.0
-
-.. Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
-
-reduction → SM100 packed paths
-==============================
-
-The registered ``packed_add_sum`` (for ``sum``) and ``3input_maxmin`` (for
-``max`` / ``min``) variants are **CUDA SM100+ fast paths** at priority **20**,
-so they pre-empt :doc:`local`. Both accept a thread-scope operation from a 1-D
-``float32`` vector of at least 8 elements to a scalar. The registration does not
-inspect ``reduce_axes``; the implementation folds the full source region. They
-use the SM100 packed math instructions — ``add.f32x2`` for ``sum``,
-``max3.f32`` / ``min3.f32`` for ``max`` / ``min`` — to fold two (or three)
-values per instruction. Source:
-``python/tvm/backend/cuda/tile_primitive/reduction/sm100_packed.py``.
-
-What it accepts
----------------
-
-All of the following must hold (else the dispatch declines and :doc:`local` runs):
-
-.. code-block:: python
-
- @register_dispatch(op_name, "cuda", variant=variant_name, priority=20,
- when=[
- predicate("exec_scope", exec_scope_ok, expected_scopes=["thread"]),
- predicate("local_scope", _local_scope_match), # src & dst local
- predicate("dst_len", _dst_len_ok, expected_len=1), # reduce to scalar
- predicate("src_ndim", _src_ndim_ok, expected_ndim=1),
- predicate("dtype", _dtype_ok, expected_dtype="float32"),
- predicate("sm_version", sm_version_ok, min_version=100),
- predicate("reduction_len", _reduction_len_ok, min_len=8),
- ])
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target / priority
- - ``cuda`` ``sm_100+``; priority ``20`` (beats ``local``)
- * - exec scope
- - ``thread`` only
- * - operands
- - src & dst ``local``, both ``float32``; dst length ``1``; src **1-D** with
- ``≥ 8`` elements
- * - axes
- - not checked by this variant's predicates; the implementation folds the
- full 1-D source region
-
-Demonstration program
-----------------------
-
-A single thread sums a 32-element ``float32`` local vector on ``sm_100a`` (from
-``test_reduction.py``):
-
-.. code-block:: python
-
- @Tx.function
- def test_func(
- A: Tx.Tensor([32], "float32", layout=TileLayout(S[32,])),
- B: Tx.Tensor([1], "float32", layout=TileLayout(S[1,])),
- ):
-
- Tx.device_entry()
- Tx.cta_id([1])
- Tx.thread_id([1])
- A_local = Tx.alloc_tensor([32], "float32", scope="local")
- B_local = Tx.alloc_tensor([1], "float32", scope="local")
- for i in Tx.serial(32):
- A_local[i] = A[i]
- Tx.tile.sum(B_local, A_local, accum=False) # -> packed_add_sum
- B[0] = B_local[0]
-
-
- target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
-
-Algorithm
----------
-
-**sum (packed_add_sum).** Keep an 8-wide ``float32`` accumulator. Load the first 8
-elements; for each further chunk of 8, pairwise-add it in with four ``add.f32x2``
-(each adds two ``float2`` lanes at once); handle the remainder scalar; then collapse
-the accumulator ``8 → 4 → 2 → 1`` with three more ``add.f32x2``:
-
-.. code-block:: python
-
- # final tree (8 -> 4 -> 2 -> 1); mov.b64 packs/unpacks the float2 lanes
- Tx.ptx.mov.b64(acc, local_sum[0], local_sum[1])
- Tx.ptx.mov.b64(rhs, local_sum[2], local_sum[3])
- Tx.ptx.add.rn.ftz.f32x2(acc, acc, rhs)
- Tx.ptx.mov.b64(local_sum[0], local_sum[1], acc)
- # ... same for local_sum[4:8], then fold the two halves together ...
- dst[...] = local_sum[0] + local_sum[1]
-
-**max / min (3input_maxmin).** A 4-wide accumulator folded three-at-a-time with the
-``max3.f32`` / ``min3.f32`` instructions.
-
-Generated TIRx IR
------------------
-
-.. code-block:: python
-
- Tx.ptx.mov.b64(acc, local_sum[0], local_sum[1])
- Tx.ptx.mov.b64(rhs, local_sum[2], local_sum[3])
- Tx.ptx.add.rn.ftz.f32x2(acc, acc, rhs) # ... the 8->4->2->1 tree
- Tx.ptx.mov.b64(local_sum[0], local_sum[1], acc)
-
-Generated CUDA
---------------
-
-.. code-block:: c++
-
- // packed pairwise add: two float lanes per instruction
- "add.rn.ftz.f32x2 %0, %1, %2;"
- // call: tvm_builtin_ptx_add_rn_ftz_f32x2(acc, acc, rhs);
- // with acc / rhs packed by tvm_builtin_ptx_mov_pack_b32x2_b64_u64_f32(...)
- // and unpacked by tvm_builtin_ptx_mov_unpack_b32x2_b64_f32_u64(...)
-
-(Verified on ``sm_100a`` — ``B == sum(A)`` for a 32-element vector.)
-
-How inputs change the algorithm
--------------------------------
-
-.. list-table::
- :header-rows: 1
- :widths: 28 72
-
- * - input
- - effect
- * - op
- - ``sum`` → the ``add.f32x2`` packed tree; ``max`` / ``min`` → the
- ``max3.f32`` / ``min3.f32`` 3-input fold
- * - reduction length
- - the chunk-of-8 (sum) / chunk handling and the scalar remainder loop; must be
- ``≥ 8``
- * - accum
- - ``True`` folds the old dst value into the first accumulator slot
- * - anything outside the gate
- - non-fp32, 2-D src, dst length > 1, pre-SM100, or ``< 8`` elements →
- the dispatch declines and :doc:`local` handles it when that variant's
- requirements hold
diff --git a/include/tvm/tirx/exec_scope.h b/include/tvm/tirx/exec_scope.h
index b83324f6e1..9fd2fc1785 100644
--- a/include/tvm/tirx/exec_scope.h
+++ b/include/tvm/tirx/exec_scope.h
@@ -39,7 +39,7 @@ namespace tirx {
* \brief The target execution scope kind of a tile primitive call.
*
* Identifies the granularity at which an op executes (the per-call
- * ``scope`` on a ``TileOpCall``, e.g. ``Tx.warp.copy(...)``).
+ * ``scope`` on tensor Call attributes, e.g. ``T.cuda.tile.ld(..., scope="warp")``).
* Ordered from coarsest to finest; smaller integer = wider scope, so
* ``ScopeKindHigher`` is a plain ``<``.
*/
diff --git a/include/tvm/tirx/op/tile.h b/include/tvm/tirx/op/tile.h
deleted file mode 100644
index 416e70ba57..0000000000
--- a/include/tvm/tirx/op/tile.h
+++ /dev/null
@@ -1,393 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-/*!
- * \file tvm/tirx/op/tile.h
- * \brief TIRX tile primitive statements, operators, and reified lambda expressions.
- */
-#ifndef TVM_TIRX_OP_TILE_H_
-#define TVM_TIRX_OP_TILE_H_
-
-#include <tvm/ffi/object.h>
-#include <tvm/ir/op.h>
-
-namespace tvm {
-namespace tirx {
-
-/*!
- * \brief The type of the function that sanitizes the arguments of a TIRX operator.
- * \param op The operator.
- * \param args The arguments.
- */
-using FArgSanitizer = ffi::TypedFunction<void(tvm::Op, ffi::Array<ffi::ObjectRef>)>;
-
-/*! \brief Tile primitive operator handles. */
-namespace tile {
-
-/*!
- * \brief See pesudo code below:
- *
- * Tx.cast(TensorRegion dst, TensorRegion src)
- */
-TVM_DLL const Op& cast_op();
-
-/*!
- * \brief See pesudo code below:
- *
- * Tx.copy(TensorRegion dst, TensorRegion src)
- */
-TVM_DLL const Op& copy_op();
-
-/*!
- * \brief See pesudo code below:
- *
- * Tx.Async.copy(TensorRegion dst, TensorRegion src)
- */
-TVM_DLL const Op& copy_async_op();
-
-/*!
- * \brief See pesudo code below:
- *
- * Tx.fill(TensorRegion dst, PrimExpr value)
- */
-TVM_DLL const Op& fill_op();
-
-/*!
- * \brief See pesudo code below:
- *
- * Tx.gemm(TensorVar A, TensorVar B, TensorVar C, TensorVar D, PrimExpr alpha, PrimExpr beta)
- */
-TVM_DLL const Op& gemm_op();
-
-/*!
- * \brief See pesudo code below:
- *
- * Tx.gemm_async(TensorRegion C, TensorRegion A, TensorRegion B, bool transA, bool transB,
- * bool accum)
- */
-TVM_DLL const Op& gemm_async_op();
-
-/*!
- * \brief Get the tile zero operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- */
-TVM_DLL const Op& zero_op();
-
-/*!
- * \brief Get the tile sqrt operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- */
-TVM_DLL const Op& sqrt_op();
-/*!
- * \brief Get the tile sqrt with scale bias operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- * - args[2]: scale, The scale factor.
- * - args[3]: bias, The bias.
- */
-TVM_DLL const Op& sqrt_with_scale_bias_op();
-
-/*!
- * \brief Get the tile exp operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- */
-TVM_DLL const Op& exp_op();
-/*!
- * \brief Get the tile exp with scale bias operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- * - args[2]: scale, The scale factor.
- * - args[3]: bias, The bias.
- */
-TVM_DLL const Op& exp_with_scale_bias_op();
-
-/*!
- * \brief Get the tile exp2 operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- */
-TVM_DLL const Op& exp2_op();
-/*!
- * \brief Get the tile exp2 with scale bias operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- * - args[2]: scale, The scale factor.
- * - args[3]: bias, The bias.
- */
-TVM_DLL const Op& exp2_with_scale_bias_op();
-
-/*!
- * \brief Get the tile log2 operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- */
-TVM_DLL const Op& log2_op();
-/*!
- * \brief Get the tile log2 with scale bias operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- * - args[2]: scale, The scale factor.
- * - args[3]: bias, The bias.
- */
-TVM_DLL const Op& log2_with_scale_bias_op();
-
-/*!
- * \brief Get the tile add operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src1, The first source.
- * - args[2]: src2, The second source.
- */
-TVM_DLL const Op& add_op();
-
-/*!
- * \brief Get the tile sub operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src1, The first source.
- * - args[2]: src2, The second source.
- */
-TVM_DLL const Op& sub_op();
-
-/*!
- * \brief Get the tile mul operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src1, The first source.
- * - args[2]: src2, The second source.
- */
-TVM_DLL const Op& mul_op();
-
-/*!
- * \brief Get the tile fdiv operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src1, The first source.
- * - args[2]: src2, The second source.
- */
-TVM_DLL const Op& fdiv_op();
-
-/*!
- * \brief Get the tile minimum operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src1, The first source.
- * - args[2]: src2, The second source.
- */
-TVM_DLL const Op& minimum_op();
-
-/*!
- * \brief Get the tile maximum operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src1, The first source.
- * - args[2]: src2, The second source.
- */
-TVM_DLL const Op& maximum_op();
-
-/*!
- * \brief Get the tile reciprocal operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- */
-TVM_DLL const Op& reciprocal_op();
-
-/*!
- * \brief Get the tile sum operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- * - args[2]: axes, The axes.
- * - args[3]: accum, The accumulator.
- */
-TVM_DLL const Op& sum_op();
-
-/*!
- * \brief Get the tile max operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- * - args[2]: axes, The axes.
- * - args[3]: accum, The accumulator.
- */
-TVM_DLL const Op& max_op();
-
-/*!
- * \brief Get the tile min operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- * - args[2]: axes, The axes.
- * - args[3]: accum, The accumulator.
- */
-TVM_DLL const Op& min_op();
-
-/*!
- * \brief Get the tile memset operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: value, The value to use.
- */
-TVM_DLL const Op& memset_op();
-
-/*!
- * \brief Get the tile reduce negate operation.
- *
- * Arguments, in order:
- * - args[0]: output, The output.
- * - args[1]: input, The input.
- * - args[2]: reduce_op, The reduction operation.
- * - args[3]: reduce_axes, The reduction axes.
- * - args[4]: accum, The accumulator.
- */
-TVM_DLL const Op& reduce_negate_op();
-
-/*!
- * \brief Get the tile binary reduce operation.
- *
- * Arguments, in order:
- * - args[0]: binary_output, The binary operation output.
- * - args[1]: reduce_output, The reduction output.
- * - args[2]: binary_input1, The first binary input.
- * - args[3]: binary_input2, The second binary input.
- * - args[4]: binary_op, The binary operation.
- * - args[5]: reduce_op, The reduction operation.
- * - args[6]: reduce_axes, The reduction axes.
- */
-TVM_DLL const Op& binary_reduce_op();
-
-/*!
- * \brief Get the tile unary reduce operation.
- *
- * Arguments, in order:
- * - args[0]: unary_output, The unary operation output.
- * - args[1]: reduce_output, The reduction output.
- * - args[2]: unary_input, The unary operation input.
- * - args[3]: unary_op, The unary operation.
- * - args[4]: reduce_op, The reduction operation.
- * - args[5]: reduce_axes, The reduction axes.
- */
-TVM_DLL const Op& unary_reduce_op();
-/*!
- * \brief Get the tile unary reduce with scale bias operation.
- *
- * Arguments, in order:
- * - args[0]: unary_output, The unary operation output.
- * - args[1]: reduce_output, The reduction output.
- * - args[2]: unary_input, The unary operation input.
- * - args[3]: unary_op, The unary operation.
- * - args[4]: reduce_op, The reduction operation.
- * - args[5]: scale, The scale factor.
- * - args[6]: bias, The bias.
- * - args[7]: reduce_axes, The reduction axes.
- */
-TVM_DLL const Op& unary_reduce_with_scale_bias_op();
-
-/*!
- * \brief Get the tile binary chain operation.
- *
- * Arguments, in order:
- * - args[0]: output, The output.
- * - args[1]: data, The input data.
- * - args[2]: operand0, The first operand.
- * - args[3]: operand1, The second operand.
- * - args[4]: op0, The first operation.
- * - args[5]: op1, The second operation.
- * - args[6]: reverse1, Whether to reverse the second operation.
- */
-TVM_DLL const Op& binary_chain_op();
-
-/*!
- * \brief Get the tile select operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: true_value, The value when the condition is true.
- * - args[2]: false_value, The value when the condition is false.
- * - args[3]: pred, The predicate.
- */
-TVM_DLL const Op& select_op();
-
-/*!
- * \brief Get the tile fma operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- * - args[2]: scale, The scale factor.
- * - args[3]: bias, The bias.
- */
-TVM_DLL const Op& fma_op();
-
-/*!
- * \brief Get the tile silu operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- */
-TVM_DLL const Op& silu_op();
-
-/*!
- * \brief Get the tile permute layout operation.
- *
- * Arguments, in order:
- * - args[0]: dst, The destination.
- * - args[1]: src, The source.
- */
-TVM_DLL const Op& permute_layout_op();
-
-} // namespace tile
-
-} // namespace tirx
-} // namespace tvm
-
-#endif // TVM_TIRX_OP_TILE_H_
diff --git a/include/tvm/tirx/script/ir_builder/ir.h b/include/tvm/tirx/script/ir_builder/ir.h
index dabe216d18..500839ef23 100644
--- a/include/tvm/tirx/script/ir_builder/ir.h
+++ b/include/tvm/tirx/script/ir_builder/ir.h
@@ -119,8 +119,6 @@ void FuncAttrs(ffi::Map<ffi::String, ffi::Any> attrs);
*/
Type FuncRet(Type ret_type);
-void TileOpCall(tvm::tirx::TileOpCall op_call);
-
/*!
* \brief Define a scope id. Pass `extents=std::nullopt` to defer the extent; it is
* inferred at LowerTIRx from the sibling ScopeIdDef closure.
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 8bdc98e868..44905ed0f3 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -112,73 +112,6 @@ inline bool IsNoOp(const Stmt& stmt) {
return false;
}
-/*!
- * \brief TIRX TileOpCall stmt.
- */
-class TileOpCallNode : public StmtNode {
- public:
- explicit TileOpCallNode(ffi::UnsafeInit tag) : op(tag) {}
-
- TileOpCallNode(tvm::Op op, ffi::Array<Expr> args, ffi::Map<ffi::String, TensorVar> workspace,
- ffi::Map<ffi::String, Expr> config, ffi::Optional<ffi::String> dispatch,
- ExecScope scope)
- : op(std::move(op)),
- args(std::move(args)),
- workspace(std::move(workspace)),
- config(std::move(config)),
- dispatch(std::move(dispatch)),
- scope(std::move(scope)) {}
-
- // tvm::Op which corresponds to the TIRX operator.
- tvm::Op op;
-
- // Arguments to the operator.
- ffi::Array<Expr> args;
-
- // Workspace (pre-allocated buffers) for the operator.
- ffi::Map<ffi::String, TensorVar> workspace;
-
- // Config for the operator/scheduler.
- ffi::Map<ffi::String, Expr> config;
-
- // Optional dispatch variant name registered via @register_dispatch.
- ffi::Optional<ffi::String> dispatch{std::nullopt};
-
- // Cooperation scope of this call. Default thread (an unscoped call).
- ExecScope scope = ExecScope(ScopeKind::kThread);
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<TileOpCallNode>()
- .def_ro("op", &TileOpCallNode::op)
- .def_ro("args", &TileOpCallNode::args)
- .def_ro("workspace", &TileOpCallNode::workspace)
- .def_ro("config", &TileOpCallNode::config)
- .def_ro("dispatch", &TileOpCallNode::dispatch)
- .def_ro("scope", &TileOpCallNode::scope);
- }
-
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TileOpCall", TileOpCallNode, StmtNode);
-};
-
-/*!
- * \brief Managed reference to TileOpCallNode
- * \sa TileOpCallNode
- */
-class TileOpCall : public Stmt {
- public:
- TVM_DLL TileOpCall(tvm::Op op, ffi::Array<Expr> args,
- ffi::Map<ffi::String, TensorVar> workspace = {},
- ffi::Map<ffi::String, Expr> config = {},
- ffi::Optional<ffi::String> dispatch = std::nullopt,
- ExecScope scope = ExecScope(ScopeKind::kThread));
-
- explicit TileOpCall(ffi::ObjectPtr<TileOpCallNode> node) : Stmt(std::move(node)) {}
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TileOpCall, Stmt, TileOpCallNode);
- TVM_DEFINE_OBJECT_REF_COW_METHOD(TileOpCallNode);
-};
-
} // namespace tirx
} // namespace tvm
#endif // TVM_TIRX_STMT_H_
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index 87488c6c45..c9b9024cc2 100644
--- a/include/tvm/tirx/stmt_functor.h
+++ b/include/tvm/tirx/stmt_functor.h
@@ -111,9 +111,6 @@ class StmtFunctor<R(const Stmt&, Args...)> {
virtual R Dispatch_(const ScopeIdDefStmtNode* node, Args... args) {
return DispatchDefault_(node, std::forward<Args>(args)...);
}
- virtual R Dispatch_(const tirx::TileOpCallNode* node, Args... args) {
- return DispatchDefault_(node, std::forward<Args>(args)...);
- }
/*! \brief Default behavior for statement hooks not overridden by a subclass. */
virtual R DispatchDefault_(const ffi::Object* node, Args...) {
TVM_FFI_THROW(InternalError) << "Do not have a default for " << node->GetTypeKey();
@@ -144,7 +141,6 @@ class StmtFunctor<R(const Stmt&, Args...)> {
SetDispatch<TSelf, SeqStmtNode>(vtable);
SetDispatch<TSelf, EvaluateNode>(vtable);
SetDispatch<TSelf, ScopeIdDefStmtNode>(vtable);
- SetDispatch<TSelf, tirx::TileOpCallNode>(vtable);
}
/*! \brief Register an additional node hook implemented by Self. */
template <typename Self, typename Node>
@@ -201,7 +197,6 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor {
virtual ffi::Optional<VisitInterrupt> Visit_(const SeqStmtNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const EvaluateNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const ScopeIdDefStmtNode* op);
- virtual ffi::Optional<VisitInterrupt> Visit_(const TileOpCallNode* op);
// Preserve TIRx operand traversal where it differs from the shared defaults.
ffi::Optional<VisitInterrupt> Visit_(const VarNode* op) override;
@@ -285,7 +280,6 @@ class TVM_DLL StmtExprMutator : public tvm::ExprMutator {
virtual UnchangedOr<Stmt> Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const EvaluateNode* op, InplaceMode inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const ScopeIdDefStmtNode* op, InplaceMode inplace_mode);
- virtual UnchangedOr<Stmt> Mutate_(const TileOpCallNode* op, InplaceMode inplace_mode);
protected:
explicit StmtExprMutator(const VTable* vtable) : tvm::ExprMutator(vtable) {}
diff --git a/python/tvm/backend/cuda/script.py b/python/tvm/backend/cuda/script.py
index bf86e49d3b..26f05b09b5 100644
--- a/python/tvm/backend/cuda/script.py
+++ b/python/tvm/backend/cuda/script.py
@@ -21,6 +21,8 @@ from __future__ import annotations
from tvm import ir as _ir
from tvm.backend.cuda import op as _cuda_op
+from .tensor_instructions import make_namespace as _make_tensor_namespace
+
# pylint: disable=protected-access
@@ -138,7 +140,7 @@ __shfl_down_sync = _cuda_op.__shfl_down_sync
__shfl_xor_sync = _cuda_op.__shfl_xor_sync
__activemask = _cuda_op.__activemask
-
+tile = _make_tensor_namespace()
_ir.op._init_op_api("tirx.cuda", __name__)
diff --git a/python/tvm/backend/cuda/tensor_instructions.py b/python/tvm/backend/cuda/tensor_instructions.py
new file mode 100644
index 0000000000..e0a8b3a2d6
--- /dev/null
+++ b/python/tvm/backend/cuda/tensor_instructions.py
@@ -0,0 +1,407 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""CUDA tensor instruction interfaces and their fixed lowering contracts."""
+
+from tvm.ir import TensorRegion
+from tvm.tirx.tensor_instruction import Instruction, Operand, namespace
+
+
+def _check(result):
+ if isinstance(result, tuple):
+ ok, reason = result
+ else:
+ ok, reason = result, "incompatible operands or execution scope"
+ if not ok:
+ raise ValueError(reason)
+
+
+def _copy_project(values, attrs):
+ return (
+ "copy",
+ [values["dst"], values["src"]],
+ {k: v for k, v in values.items() if k not in {"dst", "src"}},
+ )
+
+
+def _async_project(values, attrs):
+ _, args, extra = _copy_project(values, attrs)
+ return "copy_async", args, extra
+
+
+def _memory(view, sctx):
+ from .tile_primitive.copy import ld_stmatrix, vec_auto_reg, vec_forced
+
+ dst, src = view.args
+ name = view.op.name.rsplit(".", 1)[-1]
+ dst_scope, src_scope = dst.source.scope(), src.source.scope()
+ if name in {"ld", "ldmatrix"}:
+ if dst_scope != "local" or src_scope not in {"global", "shared", "shared.dyn"}:
+ raise ValueError(f"{name} requires memory -> local operands")
+ elif name in {"st", "stmatrix"}:
+ if src_scope != "local" or dst_scope not in {"global", "shared", "shared.dyn"}:
+ raise ValueError(f"{name} requires local -> memory operands")
+ elif dst_scope != "local" or src_scope != "local":
+ raise ValueError("register mov requires local -> local operands")
+ if name in {"ldmatrix", "stmatrix"}:
+ _check(ld_stmatrix._is_ldstmatrix(view, sctx))
+ return ld_stmatrix._emit(view, sctx)
+ bits = int(view.call.attrs.vec_bits) if hasattr(view.call.attrs, "vec_bits") else 0
+ if bits:
+ if bits not in (16, 32, 64, 128, 256):
+ raise ValueError("vec_bits must be 16, 32, 64, 128 or 256")
+ _check(
+ vec_forced._is_forced_vec_copy(view, sctx, variant=f"vec_{bits}b", num_bytes=bits // 8)
+ )
+ return vec_forced._emit_forced_vec_copy(view, sctx, bits // 8)
+ _check(vec_auto_reg._is_reg_copy(view, sctx))
+ return vec_auto_reg._emit_reg(view, sctx)
+
+
+def _elementwise(view, sctx):
+ from .tile_primitive.elementwise.ops import ALL_OPS
+ from .tile_primitive.elementwise.reg import emit_reg, is_reg_ewise
+ from .tile_primitive.elementwise.smem import emit_smem, is_smem_ewise
+
+ spec = ALL_OPS[view.kind]
+ if view.args[0].source.scope() == "local":
+ _check(is_reg_ewise(spec)(view, sctx))
+ return emit_reg(view, spec, sctx)
+ _check(is_smem_ewise(spec)(view, sctx))
+ return emit_smem(view, spec, sctx)
+
+
+def _mov_project(values, attrs):
+ kind = "copy" if isinstance(values["src"], TensorRegion) else "fill"
+ return kind, [values["dst"], values["src"]], {}
+
+
+def _mov(view, sctx):
+ if view.kind == "copy":
+ if any(region.source.scope() != "local" for region in view.args):
+ raise ValueError("register mov requires local -> local operands")
+ view.kind = "cast"
+ return _elementwise(view, sctx)
+
+
+def _ldgsts(view, sctx):
+ from .tile_primitive.copy_async.ldgsts import _emit_ldgsts, _is_ldgsts
+
+ _check(_is_ldgsts(view, sctx))
+ return _emit_ldgsts(view, sctx)
+
+
+def _dsmem(view, sctx):
+ from .tile_primitive.common import validate_copy_op
+ from .tile_primitive.copy_async.dsmem import copy_dsmem_impl
+
+ _check(validate_copy_op(view, sctx))
+ if not sctx.is_thread or any(
+ not region.source.scope().startswith("shared") for region in view.args
+ ):
+ raise ValueError("cp_async_bulk requires shared -> shared operands at thread scope")
+ return copy_dsmem_impl(view, sctx)
+
+
+def _tma(view, sctx):
+ from .tile_primitive.copy_async.tma import (
+ _validate_tma_copy_op,
+ copy_tma_auto_impl,
+ copy_tma_explicit_impl,
+ )
+
+ if not sctx.is_thread:
+ raise ValueError("TMA tensor instructions require thread scope")
+ name = view.op.name.rsplit(".", 1)[-1]
+ dst, src = view.args
+ load = name == "cp_async_bulk_tensor_load"
+ if (src.source.scope() == "global") != load or (dst.source.scope() == "global") == load:
+ raise ValueError(f"{name}: incompatible transfer direction")
+ view.options.pop("descriptor_mode", None)
+ reduction = view.options.pop("reduce_op", None)
+ if name == "cp_reduce_async_bulk_tensor":
+ if reduction is None:
+ raise ValueError("cp_reduce_async_bulk_tensor requires reduce_op")
+ view.options["use_tma_reduce"] = reduction
+ elif reduction is not None:
+ raise ValueError("reduce_op requires cp_reduce_async_bulk_tensor")
+ policy = view.options.pop("cache_policy", None)
+ if policy is not None:
+ if view.call.attrs.cache_hint:
+ raise ValueError("cache_hint and cache_policy are mutually exclusive")
+ view.options["cache_hint"] = policy
+ _check(_validate_tma_copy_op(view, sctx))
+ mode = view.call.attrs.descriptor_mode
+ if mode == "auto":
+ return copy_tma_auto_impl(view, sctx)
+ if mode == "explicit":
+ return copy_tma_explicit_impl(view, sctx)
+ raise ValueError("descriptor_mode must be auto or explicit")
+
+
+def _tc_cp(view, sctx):
+ from .tile_primitive.copy_async.tcgen05_cp import _validate_smem_tmem_copy, copy_smem_tmem_impl
+
+ if not sctx.is_thread:
+ raise ValueError("tcgen05.cp requires thread scope")
+ _check(_validate_smem_tmem_copy(view, sctx))
+ return copy_smem_tmem_impl(view, sctx)
+
+
+def _tc_ldst(view, sctx):
+ from .tile_primitive.copy.utils import _is_valid_copy
+ from .tile_primitive.copy_async.tcgen05_ldst import copy_tmem_local_impl
+
+ _check(_is_valid_copy(view, sctx))
+ dst, src = view.args
+ load = view.op.name.endswith(".ld")
+ if (dst.source.scope(), src.source.scope()) != (
+ ("local", "tmem") if load else ("tmem", "local")
+ ):
+ raise ValueError("tcgen05.ld/st requires the named TMEM transfer direction")
+ if not sctx.is_warpgroup:
+ raise ValueError("tcgen05.ld/st requires warpgroup scope")
+ return copy_tmem_local_impl(view, sctx)
+
+
+def _mma(view, sctx):
+ from .tile_primitive.gemm.mma_m16n8k_ import (
+ _full_active_lanes,
+ _no_replica,
+ gemm_cuda_mma_dispatch,
+ )
+
+ _check(_full_active_lanes(view, sctx))
+ _check(_no_replica(view, sctx))
+ return gemm_cuda_mma_dispatch(view, sctx)
+
+
+def _tc_mma(view, sctx):
+ from .tile_primitive.gemm_async.tcgen05 import gemm_async_tcgen05_impl
+
+ if not (sctx.is_thread or sctx.is_warp):
+ raise ValueError("tcgen05.mma requires thread or warp scope")
+ return gemm_async_tcgen05_impl(view, sctx)
+
+
+def make_namespace():
+ def spec(name, kind, operands, schema="ScopeAttrs", options=None, project=None, lower=None):
+ return Instruction(
+ "tirx.cuda.tile." + name,
+ kind,
+ tuple(operands),
+ schema,
+ options or {},
+ project=project,
+ lower=lower,
+ )
+
+ pair = [Operand("dst", "region"), Operand("src", "region")]
+ memory_options = dict(vec_bits=0, cache=None, l1_evict=None, l2_evict=None, prefetch_size=None)
+ result = [
+ spec(n, "copy", pair, "MemoryAttrs", memory_options, _copy_project, _memory)
+ for n in ("ld", "st")
+ ]
+ result += [
+ spec(n, "copy", pair, project=_copy_project, lower=_memory)
+ for n in ("ldmatrix", "stmatrix")
+ ]
+ result += [
+ spec(
+ "mov",
+ "fill",
+ [Operand("dst", "region"), Operand("src")],
+ project=_mov_project,
+ lower=_mov,
+ )
+ ]
+ result += [
+ spec(
+ "cp_async",
+ "copy_async",
+ [*pair, Operand("predicate", "scalar", -1)],
+ "AsyncAttrs",
+ dict(direct=False, fill_mode="", prefetch_size=-1),
+ _async_project,
+ _ldgsts,
+ )
+ ]
+ result += [
+ spec(
+ "cp_async_bulk",
+ "copy_async",
+ [*pair, Operand("mbar", "address"), Operand("remote_cta_id", "scalar")],
+ project=_async_project,
+ lower=_dsmem,
+ )
+ ]
+ tma_options = dict(
+ descriptor_mode="auto",
+ cta_group=1,
+ cache_hint="",
+ tma_dtype=None,
+ oob=None,
+ prefetch_tensormap=False,
+ tensormap_l2_promotion=None,
+ reduce_op=None,
+ )
+ for name in (
+ "cp_async_bulk_tensor_load",
+ "cp_async_bulk_tensor_store",
+ "cp_reduce_async_bulk_tensor",
+ ):
+ operands = (
+ pair
+ + (
+ [
+ Operand("mbar", "address"),
+ Operand("cta_mask", "scalar", 0),
+ Operand("mbarrier_addr", "scalar", False),
+ Operand("gather4", "coordinates", None),
+ Operand("src_selector", "selectors", None),
+ ]
+ if name.endswith("_load")
+ else []
+ )
+ + [Operand("cache_policy", "scalar", None)]
+ )
+ result.append(
+ spec(name, "copy_async", operands, "TMAAttrs", tma_options, _async_project, _tma)
+ )
+ result += [
+ spec(
+ "tcgen05.cp",
+ "copy_async",
+ pair,
+ "TcCopyAttrs",
+ dict(shape=None, multicast=None, cta_group=1, decompress=None),
+ _async_project,
+ _tc_cp,
+ )
+ ]
+ result += [
+ spec("tcgen05." + n, "copy_async", pair, project=_async_project, lower=_tc_ldst)
+ for n in ("ld", "st")
+ ]
+ result += [
+ spec(
+ "mma_sync",
+ "gemm",
+ [
+ Operand("D", "region"),
+ Operand("A", "region"),
+ Operand("B", "region"),
+ Operand("C", "region"),
+ Operand("transpose_A", "scalar", False),
+ Operand("transpose_B", "scalar", False),
+ Operand("alpha", "scalar", 1.0),
+ Operand("beta", "scalar", 0.0),
+ ],
+ lower=_mma,
+ )
+ ]
+ mma_options = dict(
+ cta_group=1,
+ mma_m=None,
+ mma_n=None,
+ smem_desc="hoist",
+ is_AB_tf32=False,
+ weight_stationary=None,
+ )
+ for scaled in (False, True):
+ operands = [Operand("C", "region"), Operand("A", "region"), Operand("B", "region")]
+ if scaled:
+ operands += [Operand("SFA", "region"), Operand("SFB", "region")]
+ operands += [
+ Operand("transA", "scalar", False),
+ Operand("transB", "scalar", False),
+ Operand("accum", "scalar", False),
+ ]
+ # Descriptor I is a runtime operand, separate from static instruction qualifiers.
+ n_base = len(operands)
+ operands += [Operand("descI", "scalar", None)]
+
+ def project(values, attrs, names=tuple(a.name for a in operands[:n_base])):
+ return "gemm_async", [values[n] for n in names], {"descI": values["descI"]}
+
+ result += [
+ spec(
+ "tcgen05.mma_block_scale" if scaled else "tcgen05.mma",
+ "gemm_async",
+ operands,
+ "TcMmaAttrs",
+ mma_options,
+ project,
+ _tc_mma,
+ )
+ ]
+ for name, kind in {
+ "add": "add",
+ "sub": "sub",
+ "mul": "mul",
+ "div": "fdiv",
+ "max": "maximum",
+ }.items():
+ result += [
+ spec(
+ name,
+ kind,
+ [Operand("dst", "region"), Operand("lhs"), Operand("rhs")],
+ "MathAttrs",
+ dict(rounding_mode=None),
+ lower=_elementwise,
+ )
+ ]
+ for name, kind in {
+ "cvt": "cast",
+ "sqrt": "sqrt",
+ "ex2": "exp2",
+ "lg2": "log2",
+ "compose.exp": "exp",
+ "compose.silu": "silu",
+ }.items():
+ result += [
+ spec(
+ name,
+ kind,
+ [Operand("dst", "region"), Operand("src", "region", "@dst")],
+ lower=_elementwise,
+ )
+ ]
+ result += [
+ spec(
+ "fma",
+ "fma",
+ [Operand("dst", "region"), Operand("src"), Operand("scale"), Operand("bias")],
+ lower=_elementwise,
+ )
+ ]
+ for op in ("exp", "exp2", "log2", "sqrt"):
+ name = op + "_with_scale_bias"
+ result += [
+ spec(
+ "compose." + name,
+ name,
+ [
+ Operand("dst", "region"),
+ Operand("src", "region"),
+ Operand("scale", "scalar"),
+ Operand("bias"),
+ ],
+ lower=_elementwise,
+ )
+ ]
+ return namespace("cuda", result)
diff --git a/python/tvm/backend/cuda/tile_primitive/__init__.py b/python/tvm/backend/cuda/tile_primitive/__init__.py
index c4f37b0253..13a83393a9 100644
--- a/python/tvm/backend/cuda/tile_primitive/__init__.py
+++ b/python/tvm/backend/cuda/tile_primitive/__init__.py
@@ -14,11 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-
-from .copy import *
-from .copy_async import *
-from .elementwise import *
-from .gemm import *
-from .gemm_async import *
-from .permute_layout import *
-from .reduction import *
diff --git a/python/tvm/backend/cuda/tile_primitive/common.py b/python/tvm/backend/cuda/tile_primitive/common.py
index 140c9ca401..0f1a46dfa1 100644
--- a/python/tvm/backend/cuda/tile_primitive/common.py
+++ b/python/tvm/backend/cuda/tile_primitive/common.py
@@ -28,7 +28,7 @@ from tvm.script import tirx as T
from tvm.sym.analyzer import Analyzer
from tvm.tirx import Function, Var
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
def next_power_of_2(x: int) -> int:
@@ -82,7 +82,7 @@ class CopyInstType(Enum):
def validate_copy_op(
- op_call: TileOpCall,
+ op_call: TensorCall,
sctx: DispatchContext, # pylint: disable=unused-argument
) -> bool:
"""Sanity check for copy op"""
@@ -152,7 +152,7 @@ def get_vec_len(
def copy_vec_load_impl(
- op_call: TileOpCall, sctx: DispatchContext, inst_type: CopyInstType
+ op_call: TensorCall, sctx: DispatchContext, inst_type: CopyInstType
) -> Function | None:
"""Schedule copy operation between global and local/shared memory on CUDA across a CTA/thread.
The implementation tries to vectorize the copy operation and parallelize over
@@ -181,7 +181,7 @@ def copy_vec_load_impl(
fail(f"unsupported exec_scope {sctx.scope_kind}")
elem_size = DataType(src.dtype).bits # in bits
- vec_len = op_call.config.get("vec_len", None)
+ vec_len = op_call.options.get("vec_len", None)
if vec_len is None:
vec_len = get_vec_len(
dst_buffer_region,
@@ -272,7 +272,7 @@ def get_thread_cnt(sctx: DispatchContext) -> int | None:
def sm_version_ok(
- op: TileOpCall, sctx: DispatchContext, min_version: int
+ op: TensorCall, sctx: DispatchContext, min_version: int
) -> tuple[bool, str | None]:
"""Check if SM version >= min_version. Usable as a dispatch predicate."""
target_arch = sctx.target.arch if hasattr(sctx.target, "arch") else ""
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/__init__.py b/python/tvm/backend/cuda/tile_primitive/copy/__init__.py
index 6ba001eaf3..56059157be 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/__init__.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/__init__.py
@@ -15,14 +15,4 @@
# specific language governing permissions and limitations
# under the License.
-from .fallback import *
-from .ld_stmatrix import *
-from .utils import (
- _is_valid_copy,
- _scope_allowed,
- _single_thread_exec,
-)
-from .vec_auto import *
-from .vec_auto_gmem_smem import *
-from .vec_auto_reg import *
-from .vec_forced import *
+from .utils import _is_valid_copy, _scope_allowed, _single_thread_exec
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/_common.py b/python/tvm/backend/cuda/tile_primitive/copy/_common.py
index 42dae4debc..3f6f2bd7a9 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/_common.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/_common.py
@@ -27,7 +27,7 @@ to call, allowed vec widths). All the layout/partition logic lives here.
from tvm import sym
from tvm.tirx.layout import ComposeLayout, Iter, TileLayout
-from tvm.tirx.operator.tile_primitive.registry import DispatchContext
+from tvm.tirx.tile_dispatch import DispatchContext
from ..layout_utils import strip_swizzle_to_tile
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/fallback.py b/python/tvm/backend/cuda/tile_primitive/copy/fallback.py
deleted file mode 100644
index fa1bafe597..0000000000
--- a/python/tvm/backend/cuda/tile_primitive/copy/fallback.py
+++ /dev/null
@@ -1,137 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""Scalar single-thread copy fallback (priority=0)."""
-
-import warnings
-
-import tvm
-from tvm.script import tirx as T
-from tvm.tirx import Function, Var
-from tvm.tirx.layout import ComposeLayout, TileLayout
-from tvm.tirx.operator.tile_primitive.dispatcher import (
- predicate,
- register_dispatch,
-)
-from tvm.tirx.operator.tile_primitive.registry import DispatchContext
-from tvm.tirx.stmt import TileOpCall
-
-from ._common import _TID_AXIS_FOR_SCOPE
-from .utils import _is_valid_copy
-from .vec_auto_reg import _axis_decl
-
-
-def _region_st_extent(buffer_region):
- region = buffer_region.region
- return [r.min for r in region], [r.extent for r in region]
-
-
-def _is_scalar_copy(op_call: TileOpCall, sctx: DispatchContext):
- ok, reason = _is_valid_copy(op_call, sctx)
- if not ok:
- return ok, reason
- op_call = TileOpCall.downcast(op_call)
- for region in (op_call.src, op_call.dst):
- layout = region.source.layout
- if isinstance(layout, ComposeLayout):
- layout = layout.tile_layout
- if isinstance(layout, TileLayout) and any(
- it.axis.is_thread() for it in (*layout.shard, *layout.replica)
- ):
- return False, "single-thread fallback cannot access a tensor with thread axes"
- return True, None
-
-
-def _emit_fallback(op_call: TileOpCall, sctx: DispatchContext) -> Function:
- op_call = TileOpCall.downcast(op_call)
- src: Var = op_call.src.source
- dst: Var = op_call.dst.source
- src_st, src_extent = _region_st_extent(op_call.src)
- dst_st, dst_extent = _region_st_extent(op_call.dst)
-
- warnings.warn(
- f"copy/fallback (scalar single-thread) picked for {src.scope()} -> "
- f"{dst.scope()} at scope_kind={sctx.scope_kind}; all faster variants "
- f"rejected.",
- stacklevel=2,
- )
-
- def _copy_body(dst_buf, src_buf):
- dst_indices = [i for i in range(len(dst_buf.shape)) if dst_extent[i] != 1]
- src_indices = [i for i in range(len(src_buf.shape)) if src_extent[i] != 1]
- assert len(dst_indices) == len(src_indices)
- copy_extents = [dst_extent[i] for i in dst_indices]
-
- def _dst_coord(lvs):
- if isinstance(lvs, tvm.ir.Var):
- lvs = [lvs]
- coord = list(dst_st)
- for k, lv in enumerate(lvs):
- coord[dst_indices[k]] += lv
- return coord
-
- def _src_coord(lvs):
- if isinstance(lvs, tvm.ir.Var):
- lvs = [lvs]
- coord = list(src_st)
- for k, lv in enumerate(lvs):
- coord[src_indices[k]] += lv
- return coord
-
- if not copy_extents:
- T.tensor_store(dst_buf, dst_st, src_buf[tuple(src_st)])
- return
-
- with T.grid(*copy_extents) as lvs:
- T.tensor_store(dst_buf, _dst_coord(lvs), src_buf[tuple(_src_coord(lvs))])
-
- scope_kind = sctx.scope_kind
-
- if scope_kind == "thread":
-
- @T.function(check_well_formed=False)
- def impl():
- _copy_body(dst, src)
-
- return impl
-
- tid_axis_name = _TID_AXIS_FOR_SCOPE[scope_kind]
- # first-active tid = composition of per-axis offsets (radix-32, since a warp is 32 lanes)
- first_tid = int(sctx.intra["laneid"][1])
- if scope_kind == "warpgroup":
- first_tid += 32 * int(sctx.intra["wid_in_wg"][1])
- elif scope_kind == "cta":
- first_tid += 32 * int(sctx.intra["warpid"][1])
-
- @T.function(check_well_formed=False)
- def impl():
- tid = _axis_decl(tid_axis_name, sctx)
- if tid == first_tid:
- _copy_body(dst, src)
-
- return impl
-
-
-@register_dispatch(
- "copy",
- "cuda",
- variant="fallback",
- priority=0,
- when=[predicate("scalar_copy_applicable", _is_scalar_copy)],
-)
-def copy_schedule_fallback(op_call: TileOpCall, sctx: DispatchContext) -> Function:
- return _emit_fallback(op_call, sctx)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py b/python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py
index 9038b7f798..69be117587 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py
@@ -27,9 +27,9 @@ from math import prod
from tvm.script import tirx as T
from tvm.tirx import Function
from tvm.tirx.layout import ComposeLayout, S, TileLayout
-from tvm.tirx.operator.tile_primitive.dispatcher import fail, predicate, register_dispatch
-from tvm.tirx.operator.tile_primitive.registry import DispatchContext
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.operator.tile_primitive.dispatcher import fail
+from tvm.tirx.tensor_instruction import TensorCall
+from tvm.tirx.tile_dispatch import DispatchContext
from ._common import ( # noqa: F401 (_carve_tail reserved for future variants)
_carve_tail,
@@ -56,7 +56,7 @@ def _compute_r_perm(r):
return [i for i, _ in sorted(enumerate(r.shard), key=key)]
-def _is_ldstmatrix(op_call: TileOpCall, sctx: DispatchContext) -> tuple[bool, str | None]:
+def _is_ldstmatrix(op_call: TensorCall, sctx: DispatchContext) -> tuple[bool, str | None]:
if not sctx.is_target("cuda"):
return False, "non-cuda target"
if sctx.scope_kind not in ("warp", "warpgroup", "cta"):
@@ -72,8 +72,8 @@ def _is_ldstmatrix(op_call: TileOpCall, sctx: DispatchContext) -> tuple[bool, st
return True, None
-def _emit(op_call: TileOpCall, sctx: DispatchContext) -> Function:
- op_call = TileOpCall.downcast(op_call)
+def _emit(op_call: TensorCall, sctx: DispatchContext) -> Function:
+ op_call = TensorCall.decode(op_call)
# Step 1: identify reg / smem sides and pull their tensor shape + layout.
src_br = op_call.src
@@ -385,14 +385,7 @@ def _emit(op_call: TileOpCall, sctx: DispatchContext) -> Function:
return impl
-@register_dispatch(
- "copy",
- "cuda",
- variant="ldstmatrix",
- priority=10,
- when=[predicate("ldstmatrix_applicable", _is_ldstmatrix)],
-)
-def copy_schedule_ldstmatrix(op_call: TileOpCall, sctx: DispatchContext) -> Function:
+def copy_schedule_ldstmatrix(op_call: TensorCall, sctx: DispatchContext) -> Function:
return _emit(op_call, sctx)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/utils.py b/python/tvm/backend/cuda/tile_primitive/copy/utils.py
index c8c0435c46..d684bd9a1c 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/utils.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/utils.py
@@ -18,13 +18,13 @@
from collections.abc import Iterable
-from tvm.tirx.operator.tile_primitive.registry import DispatchContext
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
+from tvm.tirx.tile_dispatch import DispatchContext
from ..common import match_scope, validate_copy_op
-def _single_thread_exec(op_call: TileOpCall, sctx: DispatchContext):
+def _single_thread_exec(op_call: TensorCall, sctx: DispatchContext):
"""Predicate: exec scope must be single-thread."""
exec_scope = sctx.scope_kind
ok = exec_scope == "thread"
@@ -42,11 +42,11 @@ DEFAULT_ALLOWED_PAIRS: tuple[tuple[str, str], ...] = (
def _scope_allowed(
- op_call: TileOpCall,
+ op_call: TensorCall,
sctx: DispatchContext,
allowed_pairs: Iterable[tuple[str, str]] = DEFAULT_ALLOWED_PAIRS,
):
- op_call = TileOpCall.downcast(op_call)
+ op_call = TensorCall.decode(op_call)
dst_buffer_region, src_buffer_region = (op_call.dst, op_call.src)
src_scope = src_buffer_region.source.scope()
dst_scope = dst_buffer_region.source.scope()
@@ -65,5 +65,5 @@ def _scope_allowed(
return (True, None)
-def _is_valid_copy(op_call: TileOpCall, sctx: DispatchContext):
+def _is_valid_copy(op_call: TensorCall, sctx: DispatchContext):
return (validate_copy_op(op_call, sctx), "validate_copy_op failed")
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto.py b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto.py
deleted file mode 100644
index f58f36ccae..0000000000
--- a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto.py
+++ /dev/null
@@ -1,53 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""Auto-vectorizing synchronous copy dispatch."""
-
-from tvm.tirx import Function
-from tvm.tirx.operator.tile_primitive.dispatcher import fail, predicate, register_dispatch
-from tvm.tirx.operator.tile_primitive.registry import DispatchContext
-from tvm.tirx.stmt import TileOpCall
-
-from .vec_auto_gmem_smem import _emit_gmem_smem, _is_gmem_smem
-from .vec_auto_reg import _emit_reg, _is_reg_copy
-
-
-def _is_vec_auto_copy(op_call: TileOpCall, sctx: DispatchContext):
- g_ok, g_reason = _is_gmem_smem(op_call, sctx)
- if g_ok:
- return True, None
- r_ok, r_reason = _is_reg_copy(op_call, sctx)
- if r_ok:
- return True, None
- return False, f"gmem_smem: {g_reason}; reg: {r_reason}"
-
-
-@register_dispatch(
- "copy",
- "cuda",
- variant="vec_auto",
- priority=10,
- when=[predicate("vec_auto_applicable", _is_vec_auto_copy)],
-)
-def copy_schedule_vec_auto(op_call: TileOpCall, sctx: DispatchContext) -> Function:
- g_ok, g_reason = _is_gmem_smem(op_call, sctx)
- if g_ok:
- return _emit_gmem_smem(op_call, sctx)
- r_ok, r_reason = _is_reg_copy(op_call, sctx)
- if r_ok:
- return _emit_reg(op_call, sctx)
- fail(f"gmem_smem: {g_reason}; reg: {r_reason}")
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py
deleted file mode 100644
index 8b6928de31..0000000000
--- a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py
+++ /dev/null
@@ -1,184 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""Copy dispatch for ``global ↔ shared`` (no register side).
-
-There's no per-thread register side to inherit a partition from — both sides
-are cross-thread storage. The partition is synthesized from the surrounding
-scope context (warp / warpgroup / cta / thread): ``thread_cnt`` is derived
-from ``sctx.intra`` and each thread takes ``n_elements / thread_cnt``
-consecutive fused-index slots. Layout / partition algorithm lives in
-``_common.py`` and is shared with ``ldgsts.py``.
-"""
-
-from tvm.runtime import DataType
-from tvm.script import tirx as T
-from tvm.tirx import Function, Var
-from tvm.tirx.expr import IntImm as _IntImm
-from tvm.tirx.operator.tile_primitive.registry import DispatchContext
-from tvm.tirx.stmt import TileOpCall
-
-from ..layout_utils import recompose_swizzle
-from ._common import (
- _TID_AXIS_FOR_SCOPE,
- _thread_cnt,
- align_layouts_gs,
- copy_ptx_form,
-)
-from .utils import _is_valid_copy, _scope_allowed
-from .vec_auto_reg import _all_threads_active, _axis_decl, _ptr_off
-
-_GMEM_SMEM_PAIRS = [
- ("global", "shared*"),
- ("shared*", "global"),
-]
-
-
-def _divides_thread_cnt(op_call: TileOpCall, sctx: DispatchContext) -> tuple[bool, str | None]:
- """Reject copies whose region element count does not divide ``thread_cnt``.
-
- Without this guard the emit's ``[outer, T, vec]`` partition has no
- integer solution: either every thread gets fractional work, or
- ``thread_cnt=0`` (degenerate scope) hits a modulo-by-zero. Both cases
- indicate a poorly-shaped copy (e.g. 1024-thread CTA writing a 64-elem
- tail) that this dispatch refuses to paper over with a slow scalar emit.
- """
- op_call = TileOpCall.downcast(op_call)
- thread_cnt = _thread_cnt(sctx)
- if thread_cnt <= 0:
- return False, f"degenerate thread_cnt={thread_cnt} (scope has empty intra)"
- g_br = op_call.src if op_call.src.source.scope() == "global" else op_call.dst
- n_elements = 1
- for r in g_br.region:
- ext = r.extent
- try:
- n_elements *= int(ext)
- except (TypeError, ValueError):
- return False, f"non-constant region extent {ext}"
- if n_elements % thread_cnt != 0:
- return False, (f"region size {n_elements} not divisible by thread_cnt={thread_cnt}")
- return True, None
-
-
-def _is_gmem_smem(op_call: TileOpCall, sctx: DispatchContext) -> tuple[bool, str | None]:
- if not sctx.is_target("cuda"):
- return False, "non-cuda target"
- if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
- return False, f"unsupported exec_scope {sctx.scope_kind}"
- for check in (
- lambda: _all_threads_active(sctx),
- lambda: _is_valid_copy(op_call, sctx),
- lambda: _scope_allowed(op_call, sctx, allowed_pairs=_GMEM_SMEM_PAIRS),
- lambda: _divides_thread_cnt(op_call, sctx),
- ):
- ok, msg = check()
- if not ok:
- return False, msg
- return True, None
-
-
-def _emit_gmem_smem(op_call: TileOpCall, sctx: DispatchContext) -> Function:
- op_call = TileOpCall.downcast(op_call)
- src: Var = op_call.src.source
- dst: Var = op_call.dst.source
- if src.scope() == "global":
- g_buf, g_br, s_buf, s_br = src, op_call.src, dst, op_call.dst
- g_is_src = True
- else:
- g_buf, g_br, s_buf, s_br = dst, op_call.dst, src, op_call.src
- g_is_src = False
-
- g_region = [(r.min, r.min + r.extent) for r in g_br.region]
- s_region = [(r.min, r.min + r.extent) for r in s_br.region]
-
- elem_bits = DataType(src.dtype).bits
- thread_cnt = _thread_cnt(sctx)
-
- with sctx.target:
- g_p, s_p, vec_len = align_layouts_gs(
- g_buf.layout,
- g_buf.shape,
- g_region,
- s_buf.layout,
- s_buf.shape,
- s_region,
- elem_bits,
- thread_cnt,
- )
- s_apply_layout = recompose_swizzle(s_buf.layout, s_p)
-
- # vec_len=1 is the scalar fallback — uses the same unified
- # [outer x thread x vec] coord scheme below.
-
- vec_bits = vec_len * elem_bits
- num_bytes = vec_bits // 8
- tail, lanes, reg_dtype = copy_ptx_form(num_bytes)
- # Chains are built here: a Python string bound inside the traced body is
- # not something the parser can carry.
- ld_g, st_s = f"ld.global.{tail}", f"st.shared.{tail}"
- ld_s, st_g = f"ld.shared.{tail}", f"st.global.{tail}"
-
- # Express the per-thread per-round address as a 3D coord ``(f, tid, 0)`` vs
- # ``[total_outer, thread_cnt, vec_len]``; ``layout.apply`` flattens the rest.
- n_elements = 1
- for it in s_p.shard:
- n_elements *= int(it.extent)
- assert n_elements % (thread_cnt * vec_len) == 0, (
- f"partition produced {n_elements} elements but thread_cnt({thread_cnt}) * "
- f"vec_len({vec_len}) = {thread_cnt * vec_len} doesn't divide it"
- )
- total_outer = n_elements // (thread_cnt * vec_len)
- apply_shape = [
- _IntImm("int32", total_outer),
- _IntImm("int32", thread_cnt),
- _IntImm("int32", vec_len),
- ]
-
- s_zero = [0] * len(s_buf.shape)
- g_zero = [0] * len(g_buf.shape)
-
- tid_axis_name = _TID_AXIS_FOR_SCOPE[sctx.scope_kind] if thread_cnt > 1 else None
-
- def _decl_tid():
- if tid_axis_name is not None:
- return _axis_decl(tid_axis_name, sctx)
- return _IntImm("int32", 0)
-
- v0 = _IntImm("int32", 0)
-
- # fmt: off
- @T.function(check_well_formed=False)
- def impl():
- tid = _decl_tid()
- # The scratch only shuttles bits, so it is allocated in the PTX
- # container type rather than the element type.
- tmp = T.alloc_local((lanes,), reg_dtype)
- # Pass typed ptr_to(...) directly to _ptr_off (caching → byte math,
- # misaligned vec ops); keep a serial loop, T.unroll floods the kernel.
- for f in range(total_outer):
- g_lin = g_p.apply(f, tid, v0, shape=apply_shape)["m"]
- s_off = s_apply_layout.apply(f, tid, v0, shape=apply_shape)["m"]
- s_ptr = _ptr_off(s_buf.ptr_to(s_zero), s_off)
- g_ptr = _ptr_off(g_buf.ptr_to(g_zero), g_lin)
- if T.constexpr(g_is_src):
- T.ptx[ld_g](*[tmp[i] for i in range(lanes)], g_ptr)
- T.ptx[st_s](s_ptr, *[tmp[i] for i in range(lanes)])
- else:
- T.ptx[ld_s](*[tmp[i] for i in range(lanes)], s_ptr)
- T.ptx[st_g](g_ptr, *[tmp[i] for i in range(lanes)])
- # fmt: on
- return impl
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
index 063a0479b5..a411e1a502 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
@@ -37,8 +37,8 @@ from tvm.tirx import Function, Var
from tvm.tirx import Var as _TirVar
from tvm.tirx.expr import IntImm as _IntImm
from tvm.tirx.layout import ComposeLayout, Iter, TileLayout
-from tvm.tirx.operator.tile_primitive.registry import DispatchContext
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
+from tvm.tirx.tile_dispatch import DispatchContext
from ..layout_utils import recompose_swizzle, strip_swizzle_to_tile
from ._common import _alignment_ok, copy_ptx_form, copy_ptx_ld_chain
@@ -90,8 +90,8 @@ def _all_threads_active(sctx: DispatchContext) -> tuple[bool, str | None]:
return True, None
-def _r_side_layout_valid(op_call: TileOpCall, sctx: DispatchContext) -> tuple[bool, str | None]:
- op_call = TileOpCall.downcast(op_call)
+def _r_side_layout_valid(op_call: TensorCall, sctx: DispatchContext) -> tuple[bool, str | None]:
+ op_call = TensorCall.decode(op_call)
src: Var = op_call.src.source
dst: Var = op_call.dst.source
r_buf = src if src.scope() == "local" else dst
@@ -131,9 +131,9 @@ def _r_side_layout_valid(op_call: TileOpCall, sctx: DispatchContext) -> tuple[bo
return True, None
-def _s_side_slice_ok(op_call: TileOpCall) -> tuple[bool, str | None]:
+def _s_side_slice_ok(op_call: TensorCall) -> tuple[bool, str | None]:
"""S is the non-local side (shared* or global). Slice must succeed."""
- op_call = TileOpCall.downcast(op_call)
+ op_call = TensorCall.decode(op_call)
src_br = op_call.src
dst_br = op_call.dst
s_br = dst_br if src_br.source.scope() == "local" else src_br
@@ -147,7 +147,7 @@ def _s_side_slice_ok(op_call: TileOpCall) -> tuple[bool, str | None]:
return True, None
-def _is_reg_copy(op_call: TileOpCall, sctx: DispatchContext) -> tuple[bool, str | None]:
+def _is_reg_copy(op_call: TensorCall, sctx: DispatchContext) -> tuple[bool, str | None]:
if not sctx.is_target("cuda"):
return False, "non-cuda target"
if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
@@ -264,8 +264,8 @@ def _split_atoms_for_vec(atoms, vec_len):
raise ValueError(f"tail too short for vec_len {vec_len}")
-def _align_layouts(op_call: TileOpCall, sctx: DispatchContext):
- op_call = TileOpCall.downcast(op_call)
+def _align_layouts(op_call: TensorCall, sctx: DispatchContext):
+ op_call = TensorCall.decode(op_call)
src_br = op_call.src
dst_br = op_call.dst
if src_br.source.scope() == "local":
@@ -521,8 +521,8 @@ def _outer_const_offsets(outer_atoms, flat_idx: int) -> tuple[int, int]:
return ds, dr
-def _emit_reg(op_call: TileOpCall, sctx: DispatchContext) -> Function:
- op_call = TileOpCall.downcast(op_call)
+def _emit_reg(op_call: TensorCall, sctx: DispatchContext) -> Function:
+ op_call = TensorCall.decode(op_call)
src: Var = op_call.src.source
dst: Var = op_call.dst.source
if src.scope() == "local":
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py b/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
index 2df6d79230..5d3491a415 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
@@ -17,7 +17,7 @@
"""Explicit fixed-width vector copy dispatches.
-Cache-semantics config (all variants, read from ``op_call.config``):
+Cache-semantics config (all variants, read from ``op_call.options``):
- ``cache``: ``None`` (default, plain ``T.ptx.ld``) or ``"nc"`` (load via
``T.ptx.ld.global_.nc``, i.e. PTX ``ld.global.nc``). Requires a global src.
@@ -36,10 +36,9 @@ from tvm.ir import TensorRegion
from tvm.runtime import DataType
from tvm.script import tirx as T
from tvm.sym.analyzer import Analyzer
-from tvm.tirx import Function, Var
-from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
-from tvm.tirx.operator.tile_primitive.registry import DispatchContext
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx import Var
+from tvm.tirx.tensor_instruction import TensorCall
+from tvm.tirx.tile_dispatch import DispatchContext
from ._common import copy_ptx_form, copy_ptx_ld_chain
from .utils import _scope_allowed
@@ -71,33 +70,31 @@ def _ptx_space(scope: str) -> str:
_LD_CACHE_HINT_KEYS = ("l1_evict", "l2_evict", "prefetch_size")
-def _ld_cache_config(op_call: TileOpCall) -> tuple[str | None, dict[str, str]]:
- """Read the cache-semantics config from ``op_call.config``.
+def _ld_cache_config(op_call: TensorCall) -> tuple[str | None, dict[str, str]]:
+ """Read the cache-semantics config from ``op_call.options``.
Returns ``(cache, hints)``: ``cache`` is ``None`` or ``"nc"``, and
``hints`` maps the ``T.ptx.ld``/``T.ptx.ld.global_.nc`` L1/L2 hint kwargs
(``l1_evict``, ``l2_evict``, ``prefetch_size``) to their string values.
"""
- cache = op_call.config.get("cache", None)
+ cache = op_call.options.get("cache", None)
if cache is not None:
cache = cache.value
hints: dict[str, str] = {}
for key in _LD_CACHE_HINT_KEYS:
- value = op_call.config.get(key, None)
+ value = op_call.options.get(key, None)
if value is not None and value.value:
hints[key] = value.value
return cache, hints
def _is_forced_vec_copy(
- op_call: TileOpCall,
+ op_call: TensorCall,
sctx: DispatchContext,
*,
variant: str,
num_bytes: int,
):
- if getattr(op_call, "dispatch", None) != variant:
- return False, f"requires explicit dispatch={variant!r}"
if sctx.scope_kind != "thread":
return False, f"expected thread exec_scope, got {sctx.scope_kind}"
@@ -105,7 +102,7 @@ def _is_forced_vec_copy(
if not scope_ok:
return False, scope_reason
- op_call = TileOpCall.downcast(op_call)
+ op_call = TensorCall.decode(op_call)
src: Var = op_call.src.source
dst: Var = op_call.dst.source
if src.dtype != dst.dtype:
@@ -135,8 +132,8 @@ def _is_forced_vec_copy(
return True, None
-def _emit_forced_vec_copy(op_call: TileOpCall, _sctx: DispatchContext, num_bytes: int):
- op_call = TileOpCall.downcast(op_call)
+def _emit_forced_vec_copy(op_call: TensorCall, _sctx: DispatchContext, num_bytes: int):
+ op_call = TensorCall.decode(op_call)
src: Var = op_call.src.source
dst: Var = op_call.dst.source
src_scope = src.scope()
@@ -201,33 +198,3 @@ def _emit_forced_vec_copy(op_call: TileOpCall, _sctx: DispatchContext, num_bytes
T.ptx[st_chain](dst_ptr, *[tmp[i] for i in range(lanes)])
# fmt: on
return impl
-
-
-def _register_forced_vec_copy(variant: str, num_bytes: int) -> None:
- @register_dispatch(
- "copy",
- "cuda",
- variant=variant,
- priority=20,
- when=[
- predicate(
- f"{variant}_applicable",
- _is_forced_vec_copy,
- variant=variant,
- num_bytes=num_bytes,
- )
- ],
- )
- def _copy_schedule_forced_vec(
- op_call: TileOpCall,
- sctx: DispatchContext,
- _num_bytes=num_bytes,
- ) -> Function:
- return _emit_forced_vec_copy(op_call, sctx, _num_bytes)
-
-
-_register_forced_vec_copy("vec_256b", 32)
-_register_forced_vec_copy("vec_128b", 16)
-_register_forced_vec_copy("vec_64b", 8)
-_register_forced_vec_copy("vec_32b", 4)
-_register_forced_vec_copy("vec_16b", 2)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/__init__.py b/python/tvm/backend/cuda/tile_primitive/copy_async/__init__.py
index 15ced951fc..37cf03d8c0 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/__init__.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/__init__.py
@@ -15,12 +15,7 @@
# specific language governing permissions and limitations
# under the License.
-"""Implementation of copy_async operator dispatches for CUDA targets.
-
-Registered op: copy_async (4 variants).
-See the @register_dispatch blocks in each submodule for detailed documentation
-with before/after IR examples.
-"""
+"""CUDA asynchronous tensor transfer instruction lowerers."""
from .dsmem import *
from .ldgsts import *
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py b/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
index d763c1eada..5bac05b376 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
@@ -23,28 +23,21 @@ import operator
import tvm
from tvm.script import tirx as T
from tvm.tirx import Function, Var
-from tvm.tirx.operator.tile_primitive import (
- DispatchContext,
- fail,
- predicate,
- register_dispatch,
-)
-from tvm.tirx.stmt import TileOpCall
-
-from ..common import validate_copy_op
-from ..exec_scope_utils import single_thread
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.tensor_instruction import TensorCall
+
from .utils import find_contiguous_region, to_tile_layout
-def _is_shared_to_shared(op_call: TileOpCall) -> bool:
+def _is_shared_to_shared(op_call: TensorCall) -> bool:
"""Check if both src and dst are in shared memory."""
- op_call = TileOpCall.downcast(op_call)
+ op_call = TensorCall.decode(op_call)
src_scope = op_call.src.source.scope()
dst_scope = op_call.dst.source.scope()
return src_scope.startswith("shared") and dst_scope.startswith("shared")
-def copy_dsmem_impl(op_call: TileOpCall, sctx: DispatchContext) -> Function:
+def copy_dsmem_impl(op_call: TensorCall, sctx: DispatchContext) -> Function:
"""Implement shared-to-shared cross-CTA copy using cp.async.bulk.
Uses cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes
@@ -55,13 +48,13 @@ def copy_dsmem_impl(op_call: TileOpCall, sctx: DispatchContext) -> Function:
analysis of both src and dst buffers. Non-contiguous dimensions are iterated
over, emitting one cp.async.bulk instruction per contiguous chunk.
"""
- op_call = TileOpCall.downcast(op_call)
+ op_call = TensorCall.decode(op_call)
# Extract config
- remote_cta_id = op_call.config.get("remote_cta_id", None)
+ remote_cta_id = op_call.options.get("remote_cta_id", None)
if remote_cta_id is None:
fail("remote_cta_id not set in config")
- mbar = op_call.config.get("mbar", None)
+ mbar = op_call.options.get("mbar", None)
if mbar is None:
fail("mbar not set in config")
@@ -192,45 +185,3 @@ def copy_dsmem_impl(op_call: TileOpCall, sctx: DispatchContext) -> Function:
# fmt: on
return impl
-
-
-# === Variant: copy_async/dsmem (priority=10) ===
-#
-# When: valid async copy at single-thread scope where both src and dst are in
-# shared memory. Used for intra-cluster DSMEM copies (shared::cta -> shared::cluster).
-#
-# Before (TileOpCall):
-# Tx.copy_async(
-# dst_smem[0:128, 0:64],
-# src_smem[0:128, 0:64],
-# config={"mbar": mbar, "remote_cta_id": cta_id}
-# )
-#
-# After (emits cp.async.bulk.shared::cluster.shared::cta):
-# cluster_dst = mapa(dst_smem.ptr, cta_id)
-# cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes
-# [cluster_dst], [src_smem.ptr], size, [mbar]
-@register_dispatch(
- "copy_async",
- "cuda",
- variant="dsmem",
- priority=10,
- when=[
- predicate(
- "validate_copy_op", lambda op, sctx: (validate_copy_op(op, sctx), "not a valid copy op")
- ),
- predicate(
- "single_thread",
- lambda op, sctx: (
- single_thread(op, sctx),
- f"unsupported exec_scope {sctx.exec_scope}, expected single thread",
- ),
- ),
- predicate(
- "is_shared_to_shared",
- lambda op, sctx: (_is_shared_to_shared(op), "not a shared-to-shared copy"),
- ),
- ],
-)
-def copy_async_dispatch_dsmem(op: TileOpCall, sctx: DispatchContext) -> Function:
- return copy_dsmem_impl(op, sctx)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py b/python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py
index da4792f943..6ed39ff39c 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py
@@ -45,12 +45,8 @@ from tvm.runtime import DataType
from tvm.script import tirx as T
from tvm.tirx import Function, Var
from tvm.tirx.expr import IntImm as _IntImm
-from tvm.tirx.operator.tile_primitive.dispatcher import (
- predicate,
- register_dispatch,
-)
-from tvm.tirx.operator.tile_primitive.registry import DispatchContext
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
+from tvm.tirx.tile_dispatch import DispatchContext
from ..copy._common import (
_TID_AXIS_FOR_SCOPE,
@@ -101,12 +97,12 @@ def _config_bool(value) -> bool:
def _divides_thread_cnt_ldgsts(
- op_call: TileOpCall, sctx: DispatchContext
+ op_call: TensorCall, sctx: DispatchContext
) -> tuple[bool, str | None]:
"""Mirror of ``gmem_smem._divides_thread_cnt``: reject copies whose
region element count doesn't divide ``thread_cnt`` (and reject
``thread_cnt=0`` scopes outright). See that docstring for rationale."""
- op_call = TileOpCall.downcast(op_call)
+ op_call = TensorCall.decode(op_call)
thread_cnt = _thread_cnt(sctx)
if thread_cnt <= 0:
return False, f"degenerate thread_cnt={thread_cnt} (scope has empty intra)"
@@ -123,7 +119,7 @@ def _divides_thread_cnt_ldgsts(
return True, None
-def _is_ldgsts(op_call: TileOpCall, sctx: DispatchContext) -> tuple[bool, str | None]:
+def _is_ldgsts(op_call: TensorCall, sctx: DispatchContext) -> tuple[bool, str | None]:
if not sctx.is_target("cuda"):
return False, "non-cuda target"
if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
@@ -140,8 +136,8 @@ def _is_ldgsts(op_call: TileOpCall, sctx: DispatchContext) -> tuple[bool, str |
return True, None
-def _emit_ldgsts(op_call: TileOpCall, sctx: DispatchContext) -> Function:
- op_call = TileOpCall.downcast(op_call)
+def _emit_ldgsts(op_call: TensorCall, sctx: DispatchContext) -> Function:
+ op_call = TensorCall.decode(op_call)
src: Var = op_call.src.source
dst: Var = op_call.dst.source
# Predicate above guarantees src is global, dst is shared.
@@ -149,11 +145,11 @@ def _emit_ldgsts(op_call: TileOpCall, sctx: DispatchContext) -> Function:
s_buf, s_br = dst, op_call.dst
elem_bits = DataType(src.dtype).bits
- prefetch_size = op_call.config.get("prefetch_size", -1)
- predicate_expr = op_call.config.get("predicate", -1)
- fill_mode = op_call.config.get("fill_mode", "")
+ prefetch_size = op_call.options.get("prefetch_size", -1)
+ predicate_expr = op_call.options.get("predicate", -1)
+ fill_mode = op_call.options.get("fill_mode", "")
- if _config_bool(op_call.config.get("direct", False)):
+ if _config_bool(op_call.options.get("direct", False)):
if sctx.scope_kind != "thread":
raise ValueError("ldgsts direct=True is only valid in thread scope")
n_elements = 1
@@ -261,12 +257,5 @@ def _emit_ldgsts(op_call: TileOpCall, sctx: DispatchContext) -> Function:
return impl
-@register_dispatch(
- "copy_async",
- "cuda",
- variant="ldgsts",
- priority=20,
- when=[predicate("ldgsts_applicable", _is_ldgsts)],
-)
-def copy_schedule_ldgsts(op_call: TileOpCall, sctx: DispatchContext) -> Function:
+def copy_schedule_ldgsts(op_call: TensorCall, sctx: DispatchContext) -> Function:
return _emit_ldgsts(op_call, sctx)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
index 18c36b5e33..264208f742 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
@@ -116,10 +116,8 @@ from tvm.sym import Analyzer
from tvm.tirx import Function, Var
from tvm.tirx.layout import ComposeLayout, TCol, TileLayout, TLane
from tvm.tirx.layout import m as m_axis
-from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
-from tvm.tirx.stmt import TileOpCall
-
-from ..copy import _single_thread_exec
+from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.tensor_instruction import TensorCall
# Allowed .multicast qualifiers per shape (PTX ISA 8.8 §9.7.16.9.2): "" = none,
# 64x128b needs an explicit warpx2 pick.
@@ -193,10 +191,10 @@ def _cp_lane_replica_pattern(shape: str, multicast: str):
raise ValueError(f"unknown tcgen05.cp multicast {multicast!r}")
-def _resolve_cp_shape(op_call: TileOpCall):
+def _resolve_cp_shape(op_call: TensorCall):
"""Resolve (shape, multicast) from an explicit ``shape=`` config."""
- shape = op_call.config["shape"].value
- multicast = op_call.config.get("multicast")
+ shape = op_call.options["shape"].value
+ multicast = op_call.options.get("multicast")
if isinstance(multicast, tvm.ir.StringImm):
multicast = multicast.value
allowed = _CP_SHAPE_MULTICASTS.get(shape)
@@ -314,7 +312,7 @@ def _align_middles(t_middle, s_middle):
# -----------------------------------------------------------------------------
# Plan (state object)
# -----------------------------------------------------------------------------
-def _build_plan(op_call: TileOpCall):
+def _build_plan(op_call: TensorCall):
"""Run A..H and return a dispatch plan.
Plan fields:
@@ -329,12 +327,12 @@ def _build_plan(op_call: TileOpCall):
- t_addr_off (Expr, taddr offset of the first cp: 32-bit col
offset plus the region row offset in the lane half-word)
"""
- op_call = TileOpCall.downcast(op_call)
- if op_call.config.get("shape") is not None:
+ op_call = TensorCall.decode(op_call)
+ if op_call.options.get("shape") is not None:
shape, multicast = _resolve_cp_shape(op_call)
return _plan_for_shape(op_call, shape, multicast)
# No shape config: infer from the buffer layouts.
- multicast_cfg = op_call.config.get("multicast")
+ multicast_cfg = op_call.options.get("multicast")
if isinstance(multicast_cfg, tvm.ir.StringImm):
multicast_cfg = multicast_cfg.value
errors = []
@@ -350,7 +348,7 @@ def _build_plan(op_call: TileOpCall):
)
-def _plan_for_shape(op_call: TileOpCall, shape: str, multicast: str):
+def _plan_for_shape(op_call: TensorCall, shape: str, multicast: str):
"""Run A..I for one (shape, multicast); raises ValueError on any mismatch."""
dst_region, src_region = op_call.args[:2]
s_buf: Var = src_region.source
@@ -719,7 +717,7 @@ def _desc_set_addr(desc_val, addr_ptr):
return T.bitwise_or(T.bitwise_and(desc_val, T.bitwise_not(T.uint64(0x3FFF))), start_addr)
-def _validate_smem_tmem_copy(op_call: TileOpCall, sctx: DispatchContext):
+def _validate_smem_tmem_copy(op_call: TensorCall, sctx: DispatchContext):
"""Memory-scope envelope only; shape resolution/inference and the detailed
layout validation raise readable ValueErrors in ``_build_plan``."""
dst_region, src_region = op_call.args[:2]
@@ -742,8 +740,8 @@ def _validate_smem_tmem_copy(op_call: TileOpCall, sctx: DispatchContext):
# is responsible for issuing ``tcgen05.commit`` against a barrier if they
# need synchronization.
# -----------------------------------------------------------------------------
-def copy_smem_tmem_impl(op_call: TileOpCall, sctx: DispatchContext) -> Function | None:
- decompress = op_call.config.get("decompress")
+def copy_smem_tmem_impl(op_call: TensorCall, sctx: DispatchContext) -> Function | None:
+ decompress = op_call.options.get("decompress")
if isinstance(decompress, StringImm):
decompress = decompress.value
if decompress:
@@ -767,7 +765,7 @@ def copy_smem_tmem_impl(op_call: TileOpCall, sctx: DispatchContext) -> Function
# keep the legacy LDO=0 encoding. 256b atoms carry the derived field.
LDO_field = plan["LDO_field"]
- cta_group = op_call.config.get("cta_group", 1)
+ cta_group = op_call.options.get("cta_group", 1)
desc_buf = _get_or_create_desc(sctx, s_buf, LDO_field, SDO_field, sw)
t_addr = t_buf.allocated_addr
@@ -816,18 +814,3 @@ def copy_smem_tmem_impl(op_call: TileOpCall, sctx: DispatchContext) -> Function
# fmt: on
return impl
-
-
-# === Variant: copy_async/smem->tmem (priority=10) ===
-@register_dispatch(
- "copy_async",
- "cuda",
- variant="smem->tmem",
- priority=10,
- when=[
- predicate("validate_smem_tmem_copy", _validate_smem_tmem_copy),
- predicate("exec_scope", _single_thread_exec),
- ],
-)
-def copy_async_schedule_smem_tmem(op_call: TileOpCall, sctx: DispatchContext) -> Function:
- return copy_smem_tmem_impl(op_call, sctx)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py
index 53eb699cde..43529c6f81 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py
@@ -35,12 +35,10 @@ from tvm.tirx.layout import (
tcgen05_atom_layout,
tmem_datapath_layout,
)
-from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.tensor_instruction import TensorCall
from ..common import get_st_extent
-from ..copy import _is_valid_copy, _scope_allowed
-from ..exec_scope_utils import exec_scope_ok
# Per-warp fp32-column factor for each instr_shape (mirrors
# ``_TCGEN05_COL_FACTOR_FP32`` in ``tvm.tirx.layout``; .16x64b → 2,
@@ -222,8 +220,8 @@ def _tmem_window(tmem_buf, tmem_region, atom_kind, frag_rows, analyzer):
return width, window.offset.get(TCol, 0), lane_off // 16
-def copy_tmem_local_impl(op_call: TileOpCall, sctx: DispatchContext) -> Function | None:
- op_call = TileOpCall.downcast(op_call)
+def copy_tmem_local_impl(op_call: TensorCall, sctx: DispatchContext) -> Function | None:
+ op_call = TensorCall.decode(op_call)
dst_buffer_region, src_buffer_region = op_call.dst, op_call.src
dst: Var = dst_buffer_region.source
src: Var = src_buffer_region.source
@@ -567,28 +565,3 @@ def _emit_datapath_b_path(
emit(tmem_buf.allocated_addr[0], 0, 0, [local_32b[i] for i in range(n_half)])
# fmt: on
return impl
-
-
-# === Variant: copy_async/tmem<->local (priority=10) ===
-#
-# When: one buffer is in tmem (tensor memory, Blackwell SM100+) and the other
-# is in local scope, at warpgroup exec scope.
-#
-# Emits: T.ptx.tcgen05.ld / T.ptx.tcgen05.st (async). The caller is
-# responsible for issuing the matching ``T.ptx.tcgen05.wait__ld`` /
-# ``T.ptx.tcgen05.wait__st`` when synchronization is required.
-@register_dispatch(
- "copy_async",
- "cuda",
- variant="tmem<->local",
- priority=10,
- when=[
- predicate("validate_copy_op", _is_valid_copy),
- predicate("exec_scope", exec_scope_ok, expected_scopes=["warpgroup"]),
- predicate(
- "storage_scope", _scope_allowed, allowed_pairs=[("tmem", "local"), ("local", "tmem")]
- ),
- ],
-)
-def copy_async_schedule_tmem_local_async(op_call: TileOpCall, sctx: DispatchContext) -> Function:
- return copy_tmem_local_impl(op_call, sctx)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
index 7edde08c3c..ff4be6d364 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
@@ -37,16 +37,10 @@ from tvm.script import tirx as T
from tvm.sym import Analyzer
from tvm.tirx import Function, IntImm, Var, is_tensor_var
from tvm.tirx.layout import Layout, TileLayout
-from tvm.tirx.operator.tile_primitive import (
- DispatchContext,
- fail,
- predicate,
- register_dispatch,
-)
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.tensor_instruction import TensorCall
from ...op import _is_static_unicast_cta_mask, _resolve_cache_policy
-from ..exec_scope_utils import single_thread
from ..layout_utils import strip_swizzle_to_tile
from ..tma_utils import SwizzleMode, get_swizzle_mode_from_layout
@@ -244,7 +238,7 @@ def _require_proven(predicate, analyzer: Analyzer, stage: str, detail: str) -> N
def _auto_fail(stage: str, detail: str):
fail(
- f'tma_auto stage={stage}: {detail}; use dispatch="tma_explicit" '
+ f'tma_auto stage={stage}: {detail}; use descriptor_mode="explicit" '
"when the mapping or hardware legality is only known at runtime"
)
@@ -362,6 +356,8 @@ def _normalize_l2_promotion(value) -> int:
def _normalize_oob(value) -> int:
+ if isinstance(value, int | IntImm) and int(value) in (0, 1):
+ return int(value)
if isinstance(value, tvm.ir.StringImm):
value = value.value
if value is None or value == "zero":
@@ -1541,26 +1537,37 @@ def _repair_auto_candidate(plan: TMAPlan):
def _runtime_config(op_call, sctx, direction: str, *, explicit: bool):
+ op_call = TensorCall.decode(op_call)
+ # Fixed named Call operands are decoded once into the planner's semantic view.
+ op_call.options.pop("descriptor_mode", None)
+ reduction = op_call.options.pop("reduce_op", None)
+ if reduction is not None:
+ op_call.options["use_tma_reduce"] = reduction
+ policy = op_call.options.pop("cache_policy", None)
+ if policy is not None:
+ if op_call.call.attrs.cache_hint:
+ fail("cache_hint and cache_policy are mutually exclusive")
+ op_call.options["cache_hint"] = policy
allowed = _EXPLICIT_CONFIG if explicit else _COMMON_CONFIG
- unknown = sorted(set(op_call.config) - allowed)
+ unknown = sorted(set(op_call.options) - allowed)
if unknown:
fail(
f"dispatch={'tma_explicit' if explicit else 'tma_auto'} does not support "
f"config key(s) {unknown}"
)
- if not explicit and op_call.config.get("oob") is not None:
- fail('tma_auto does not support non-default oob; use dispatch="tma_explicit"')
- if direction != "g2s" and op_call.config.get("oob") is not None:
+ if not explicit and op_call.options.get("oob") is not None:
+ fail('tma_auto does not support non-default oob; use descriptor_mode="explicit"')
+ if direction != "g2s" and op_call.options.get("oob") is not None:
fail("TensorMap oob is only valid for explicit global-to-shared copies")
- cta_group = op_call.config.get("cta_group", 1)
+ cta_group = op_call.options.get("cta_group", 1)
if isinstance(cta_group, IntImm):
cta_group = int(cta_group)
if cta_group not in (1, 2):
fail(f"cta_group must be 1 or 2, got {cta_group}")
- cta_mask = op_call.config.get("cta_mask", 0)
+ cta_mask = op_call.options.get("cta_mask", 0)
if isinstance(cta_mask, IntImm):
cta_mask_value = int(cta_mask)
if not 0 <= cta_mask_value <= 0xFFFF:
@@ -1570,8 +1577,8 @@ def _runtime_config(op_call, sctx, direction: str, *, explicit: bool):
fail(f"cta_mask must fit uint16, got {cta_mask}")
elif not isinstance(cta_mask, tvm.tirx.Expr):
fail("cta_mask must be an integer or TIR expression")
- mbar = op_call.config.get("mbar")
- mbarrier_addr = op_call.config.get("mbarrier_addr", False)
+ mbar = op_call.options.get("mbar")
+ mbarrier_addr = op_call.options.get("mbarrier_addr", False)
if isinstance(mbarrier_addr, IntImm):
mbarrier_addr = bool(int(mbarrier_addr))
if not isinstance(mbarrier_addr, bool | tvm.tirx.Expr):
@@ -1590,7 +1597,7 @@ def _runtime_config(op_call, sctx, direction: str, *, explicit: bool):
):
fail("cta_mask is only valid for global-to-shared TMA")
- use_tma_reduce = op_call.config.get("use_tma_reduce")
+ use_tma_reduce = op_call.options.get("use_tma_reduce")
if isinstance(use_tma_reduce, tvm.ir.StringImm):
use_tma_reduce = use_tma_reduce.value
if use_tma_reduce is not None:
@@ -1599,7 +1606,7 @@ def _runtime_config(op_call, sctx, direction: str, *, explicit: bool):
if use_tma_reduce not in ("add", "min", "max", "inc", "dec", "and", "or", "xor"):
fail(f"unsupported TMA reduce operation {use_tma_reduce!r}")
- cache_hint, cache_policy = _normalize_cache_hint(op_call.config.get("cache_hint", ""))
+ cache_hint, cache_policy = _normalize_cache_hint(op_call.options.get("cache_hint", ""))
return {
"cta_group": cta_group,
"cta_mask": cta_mask,
@@ -1617,16 +1624,18 @@ def _runtime_config(op_call, sctx, direction: str, *, explicit: bool):
"use_tma_reduce": use_tma_reduce,
"cache_hint": cache_hint,
"cache_policy": cache_policy,
- "l2_promotion": _normalize_l2_promotion(op_call.config.get("tensormap_l2_promotion")),
- "oob_fill": _normalize_oob(op_call.config.get("oob")),
- "prefetch": bool(op_call.config.get("prefetch_tensormap", False)),
- "tma_dtype": (op_call.config["tma_dtype"].value if "tma_dtype" in op_call.config else None),
+ "l2_promotion": _normalize_l2_promotion(op_call.options.get("tensormap_l2_promotion")),
+ "oob_fill": _normalize_oob(op_call.options.get("oob")),
+ "prefetch": bool(op_call.options.get("prefetch_tensormap", False)),
+ "tma_dtype": (
+ op_call.options["tma_dtype"].value if "tma_dtype" in op_call.options else None
+ ),
"target_arch": sctx.target.arch,
}
def _copy_direction(op_call):
- op_call = TileOpCall.downcast(op_call)
+ op_call = TensorCall.decode(op_call)
dst_region, src_region = op_call.dst, op_call.src
src_scope = src_region.source.scope()
dst_scope = dst_region.source.scope()
@@ -1637,7 +1646,8 @@ def _copy_direction(op_call):
fail(f"TMA requires global<->shared operands, got src={src_scope}, dst={dst_scope}")
-def _build_auto_plan(op_call: TileOpCall, sctx: DispatchContext) -> TMAPlan:
+def _build_auto_plan(op_call: TensorCall, sctx: DispatchContext) -> TMAPlan:
+ op_call = TensorCall.decode(op_call)
direction, shared_region, global_region = _copy_direction(op_call)
s_buf = shared_region.source
g_buf = global_region.source
@@ -1647,8 +1657,8 @@ def _build_auto_plan(op_call: TileOpCall, sctx: DispatchContext) -> TMAPlan:
f"shared dtype={s_buf.dtype} and global dtype={g_buf.dtype} differ",
)
runtime = _runtime_config(op_call, sctx, direction, explicit=False)
- if "gather4" in op_call.config or "src_selector" in op_call.config:
- fail('gather4 and src_selector are only supported by dispatch="tma_explicit"')
+ if "gather4" in op_call.options or "src_selector" in op_call.options:
+ fail('gather4 and src_selector are only supported by descriptor_mode="explicit"')
s_starts, s_extents = _copy_region_parts(shared_region.region)
g_starts, g_extents = _copy_region_parts(global_region.region)
@@ -1916,6 +1926,8 @@ def _normalize_src_selector(value):
condition, buffer = item
if not isinstance(condition, tvm.tirx.Expr):
fail(f"tma_explicit src_selector[{idx}] condition must be a TIR expression")
+ if isinstance(buffer, tvm.ir.TensorRegion):
+ buffer = buffer.source
if not is_tensor_var(buffer):
fail(
f"tma_explicit src_selector[{idx}] candidate must be a global "
@@ -1959,13 +1971,14 @@ def _selector_compatibility(main: TensorMapSpec, candidate: TensorMapSpec, index
)
-def _build_explicit_plan(op_call: TileOpCall, sctx: DispatchContext):
+def _build_explicit_plan(op_call: TensorCall, sctx: DispatchContext):
+ op_call = TensorCall.decode(op_call)
direction, shared_region, global_region = _copy_direction(op_call)
s_buf = shared_region.source
g_buf = global_region.source
runtime = _runtime_config(op_call, sctx, direction, explicit=True)
- gather4 = _normalize_gather4(op_call.config.get("gather4"))
- selectors = _normalize_src_selector(op_call.config.get("src_selector"))
+ gather4 = _normalize_gather4(op_call.options.get("gather4"))
+ selectors = _normalize_src_selector(op_call.options.get("src_selector"))
if (gather4 or selectors) and direction != "g2s":
fail("tma_explicit gather4 and src_selector are only valid for global-to-shared")
if gather4 and len(g_buf.shape) != 2:
@@ -2220,19 +2233,21 @@ def _emit_plan(
return impl
-def copy_tma_auto_impl(op_call: TileOpCall, sctx: DispatchContext) -> Function:
- """Lower one ``tma_auto`` call."""
+def copy_tma_auto_impl(op_call: TensorCall, sctx: DispatchContext) -> Function:
+ """Lower one automatic TensorMap instruction call."""
+ op_call = TensorCall.decode(op_call)
plan = _build_auto_plan(op_call, sctx)
tensor_map, key = _get_or_encode_descriptor(plan.spec, sctx)
impl = _emit_plan(plan, tensor_map, None, False, sctx)
- if bool(op_call.config.get("prefetch_tensormap", False)):
+ if bool(op_call.options.get("prefetch_tensormap", False)):
_prefetch_main_descriptor(tensor_map, key, sctx)
return impl
-def copy_tma_explicit_impl(op_call: TileOpCall, sctx: DispatchContext) -> Function:
+def copy_tma_explicit_impl(op_call: TensorCall, sctx: DispatchContext) -> Function:
"""Lower one direct TensorMap and exactly one TMA instruction."""
+ op_call = TensorCall.decode(op_call)
plan, candidates = _build_explicit_plan(op_call, sctx)
main_map, main_key = _get_or_encode_descriptor(plan.spec, sctx)
@@ -2244,12 +2259,12 @@ def copy_tma_explicit_impl(op_call: TileOpCall, sctx: DispatchContext) -> Functi
main_map, candidate_maps
)
impl = _emit_plan(plan, selected_map, selector_bind, selected_is_address, sctx)
- if bool(op_call.config.get("prefetch_tensormap", False)):
+ if bool(op_call.options.get("prefetch_tensormap", False)):
_prefetch_main_descriptor(main_map, main_key, sctx)
return impl
-def _validate_tma_copy_op(op_call: TileOpCall, _sctx: DispatchContext) -> bool:
+def _validate_tma_copy_op(op_call: TensorCall, _sctx: DispatchContext) -> bool:
dst_region, src_region = op_call.args[:2]
src = src_region.source
dst = dst_region.source
@@ -2259,40 +2274,3 @@ def _validate_tma_copy_op(op_call: TileOpCall, _sctx: DispatchContext) -> bool:
return (src_scope == "global" and dst_scope.startswith("shared")) or (
src_scope.startswith("shared") and dst_scope == "global"
)
-
-
-_COMMON_PREDICATES = [
- predicate(
- "validate_tma_copy_op",
- lambda op, sctx: (_validate_tma_copy_op(op, sctx), "not a global<->shared TMA copy"),
- ),
- predicate(
- "single_thread",
- lambda op, sctx: (
- single_thread(op, sctx),
- f"unsupported exec_scope {sctx.exec_scope}, expected single thread",
- ),
- ),
-]
-
-
-@register_dispatch(
- "copy_async",
- "cuda",
- variant="tma_auto",
- priority=10,
- when=_COMMON_PREDICATES,
-)
-def copy_async_dispatch_tma_auto(op: TileOpCall, sctx: DispatchContext) -> Function:
- return copy_tma_auto_impl(op, sctx)
-
-
-@register_dispatch(
- "copy_async",
- "cuda",
- variant="tma_explicit",
- priority=10,
- when=_COMMON_PREDICATES,
-)
-def copy_async_dispatch_tma_explicit(op: TileOpCall, sctx: DispatchContext) -> Function:
- return copy_tma_explicit_impl(op, sctx)
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/__init__.py b/python/tvm/backend/cuda/tile_primitive/elementwise/__init__.py
index 4760792d96..a5d5fe6076 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/__init__.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/__init__.py
@@ -26,12 +26,3 @@ Each op in ``ops.ALL_OPS`` is registered under both variants. Per-op packed
PTX/CUDA intrinsics live in ``vec_emit/`` (``binary_f32x2`` / ``cast_vec2``
/ ``fma_f32x2``) and are attached to the relevant ``OpSpec.vec_impls``.
"""
-
-from .register import *
-
-# Suppress submodule-attribute leakage. Without an explicit ``__all__`` here,
-# ``from tvm.backend.cuda.tile_primitive.elementwise import *`` (run by
-# tile_primitive/__init__.py) re-exports the implicit submodule attributes
-# (``ops``, ``reg``, ``smem``, ``vec_emit``) into the public dispatcher namespace.
-# Tile operator classes are owned by tvm.tirx.op.tile.
-__all__: list[str] = []
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/__init__.py b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/__init__.py
index a0af69421d..8c5e36cef6 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/__init__.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/__init__.py
@@ -32,7 +32,7 @@ from typing import Any
from tvm.ir import TensorRegion
from tvm.ir.expr import Expr
-from tvm.tirx import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
@dataclass
@@ -78,7 +78,7 @@ class VecImpl:
"""
vec_len: int # elements per packed call
- applies: Callable[[TileOpCall, Any, Plan], tuple[bool, str | None]]
+ applies: Callable[[TensorCall, Any, Plan], tuple[bool, str | None]]
# emit(dst_ptr, src_ptrs, extras) -> Stmt
# dst_ptr: typed ptr to ``vec_len`` consecutive dst elements
# src_ptrs[i]: typed ptr to ``vec_len`` consecutive src[i] elements,
@@ -95,7 +95,7 @@ class OpSpec:
name: str
# parse(op_call) -> (Plan, msg|None); msg explains why parse failed.
- parse: Callable[[TileOpCall], tuple[Plan | None, str | None]]
+ parse: Callable[[TensorCall], tuple[Plan | None, str | None]]
# Scalar compute used by the fallback emit path (wrapped in Tx.vectorized).
# compute_scalar(src_vals_at_one_idx, extras, dst_dtype) -> Expr
compute_scalar: Callable[[list, dict, str], Any]
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py
index 7bd741ae08..c08e0bcc71 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py
@@ -35,7 +35,7 @@ from typing import Any
from tvm.ir import TensorRegion
from tvm.script import tirx as Tx
-from tvm.tirx import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
from ..vec_emit.binary_f32x2 import BINARY_F32X2_IMPLS
from . import OpSpec, Plan, SrcSpec
@@ -46,7 +46,7 @@ _COMMUTATIVE = frozenset({"add", "mul", "maximum"})
def _parse_binary_for(op_name: str):
"""Build a ``parse(op_call) -> (Plan, msg)`` for a specific binary op."""
- def parse(op: TileOpCall) -> tuple[Plan | None, str | None]:
+ def parse(op: TensorCall) -> tuple[Plan | None, str | None]:
_dst: TensorRegion = op.args[0]
_src1 = op.args[1]
_src2 = op.args[2]
@@ -57,7 +57,7 @@ def _parse_binary_for(op_name: str):
return None, "both inputs are constants"
# Move constant to rhs (commute if allowed; else reject).
- if s1_scalar:
+ if s1_scalar and op_name != "fdiv":
if op_name not in _COMMUTATIVE:
return None, f"non-commutative op {op_name} cannot have constant lhs"
_src1, _src2 = _src2, _src1
@@ -65,7 +65,7 @@ def _parse_binary_for(op_name: str):
# If rhs is a smaller buffer (broadcast), swap if commutative so the
# bigger one is in src1 — keeps src1 == dst convention.
- if not s2_scalar:
+ if not s1_scalar and not s2_scalar:
s1_n = functools.reduce(operator.mul, [r.extent for r in _src1.region], 1)
s2_n = functools.reduce(operator.mul, [r.extent for r in _src2.region], 1)
if s1_n < s2_n:
@@ -73,14 +73,16 @@ def _parse_binary_for(op_name: str):
return None, f"non-commutative op {op_name} cannot swap to broadcast"
_src1, _src2 = _src2, _src1
- srcs: list[SrcSpec] = [SrcSpec(buf_region=_src1)]
+ srcs: list[SrcSpec] = [
+ SrcSpec(scalar=_src1) if s1_scalar and op_name == "fdiv" else SrcSpec(buf_region=_src1)
+ ]
if s2_scalar:
srcs.append(SrcSpec(scalar=_src2))
else:
srcs.append(SrcSpec(buf_region=_src2))
extras: dict[str, Any] = {}
- rm = op.config.get("rounding_mode", None)
+ rm = op.options.get("rounding_mode", None)
if rm is not None:
extras["rounding_mode"] = rm.value
return Plan(dst=_dst, srcs=srcs, extras=extras), None
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/cast.py b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/cast.py
index 1a7d5edb61..0a3ff4fee2 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/cast.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/cast.py
@@ -22,13 +22,13 @@ CUDA intrinsics like ``__float22half2_rn``."""
from __future__ import annotations
from tvm.ir import TensorRegion
-from tvm.tirx import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
from ..vec_emit.cast_vec2 import CAST_VEC2_IMPL
from . import OpSpec, Plan, SrcSpec
-def _parse_cast(op: TileOpCall) -> tuple[Plan | None, str | None]:
+def _parse_cast(op: TensorCall) -> tuple[Plan | None, str | None]:
_dst: TensorRegion = op.args[0]
_src = op.args[1]
if not isinstance(_src, TensorRegion):
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/fma.py b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/fma.py
index bf4bb671a4..f4e705c6c8 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/fma.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/fma.py
@@ -24,13 +24,13 @@ Attaches ``fma_f32x2`` VecImpl for sm_100+ f32; falls back to scalar
from __future__ import annotations
from tvm.ir import TensorRegion
-from tvm.tirx import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
from ..vec_emit.fma_f32x2 import FMA_F32X2_IMPL
from . import OpSpec, Plan, SrcSpec
-def _parse_fma(op: TileOpCall) -> tuple[Plan | None, str | None]:
+def _parse_fma(op: TensorCall) -> tuple[Plan | None, str | None]:
_dst: TensorRegion = op.args[0]
args = op.args[1:4]
srcs: list[SrcSpec] = []
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py
index dac46abad0..a78d5cb4a8 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py
@@ -27,13 +27,13 @@ from typing import Any
from tvm.ir import TensorRegion, is_prim_expr
from tvm.script import tirx as T
-from tvm.tirx import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
from .._common import scalar_dtype
from . import OpSpec, Plan, SrcSpec
-def _parse_unary(op: TileOpCall) -> tuple[Plan | None, str | None]:
+def _parse_unary(op: TensorCall) -> tuple[Plan | None, str | None]:
"""T.<unary>(dst, src[, scale, bias]) → Plan."""
_dst: TensorRegion = op.args[0]
_src = op.args[1]
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py b/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
index ac1a9d619d..155b3acbd4 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
@@ -36,10 +36,11 @@ import operator
from tvm.script import tirx as T
from tvm.sym import Analyzer
-from tvm.tirx import Function, TileOpCall
+from tvm.tirx import Function
from tvm.tirx.layout import TileLayout
from tvm.tirx.operator.tile_primitive import DispatchContext
from tvm.tirx.operator.tile_primitive.dispatcher import fail
+from tvm.tirx.tensor_instruction import TensorCall
from ..common import get_st_extent
from ..copy._common import _carve_tail, _verify_s_tail_contig
@@ -175,7 +176,7 @@ def _check_layout_operands_agree(plan, sctx) -> tuple[bool, str | None]:
def is_reg_ewise(spec):
"""Predicate factory: dispatch accepted iff all operands in ``local`` scope."""
- def check(op_call: TileOpCall, sctx: DispatchContext) -> tuple[bool, str | None]:
+ def check(op_call: TensorCall, sctx: DispatchContext) -> tuple[bool, str | None]:
if not sctx.is_target("cuda"):
return False, "non-cuda target"
if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
@@ -230,7 +231,7 @@ def _prod(it) -> int:
# -----------------------------------------------------------------------------
# Main entry
# -----------------------------------------------------------------------------
-def emit_reg(op_call: TileOpCall, spec, sctx: DispatchContext) -> Function:
+def emit_reg(op_call: TensorCall, spec, sctx: DispatchContext) -> Function:
plan, msg = spec.parse(op_call)
if msg is not None or plan is None:
fail(msg or "parse failed")
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/register.py b/python/tvm/backend/cuda/tile_primitive/elementwise/register.py
deleted file mode 100644
index 4e34cb8490..0000000000
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/register.py
+++ /dev/null
@@ -1,61 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""Register each op in ``ALL_OPS`` for both dispatch variants (``reg``, ``smem``).
-
-Mirrors copy PR-640's two-variant model: scope-pair drives the dispatch
-selection, the underlying algorithm (induced vs synthesized) follows.
-"""
-
-from tvm.tirx import Function, TileOpCall
-from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
-
-from .ops import ALL_OPS
-from .reg import emit_reg, is_reg_ewise
-from .smem import emit_smem, is_smem_ewise
-
-
-def _register_reg(spec) -> None:
- @register_dispatch(
- spec.name,
- "cuda",
- variant="reg",
- priority=10,
- when=[predicate(f"{spec.name}_reg", is_reg_ewise(spec))],
- )
- def _dispatch(op: TileOpCall, sctx: DispatchContext, _spec=spec) -> Function:
- return emit_reg(op, _spec, sctx)
-
-
-def _register_smem(spec) -> None:
- @register_dispatch(
- spec.name,
- "cuda",
- variant="smem",
- priority=10,
- when=[predicate(f"{spec.name}_smem", is_smem_ewise(spec))],
- )
- def _dispatch(op: TileOpCall, sctx: DispatchContext, _spec=spec) -> Function:
- return emit_smem(op, _spec, sctx)
-
-
-for _spec in ALL_OPS.values():
- _register_reg(_spec)
- _register_smem(_spec)
-
-
-__all__: list[str] = []
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/smem.py b/python/tvm/backend/cuda/tile_primitive/elementwise/smem.py
index 1fd7726c9c..6918d5a52e 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/smem.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/smem.py
@@ -31,9 +31,10 @@ codegen time. Packed-vec emit requires the innermost dim to have stride 1
from __future__ import annotations
from tvm.script import tirx as T
-from tvm.tirx import Function, TileOpCall
+from tvm.tirx import Function
from tvm.tirx.operator.tile_primitive import DispatchContext
from tvm.tirx.operator.tile_primitive.dispatcher import fail
+from tvm.tirx.tensor_instruction import TensorCall
from ..common import get_indices, get_st_extent, get_thread_cnt
from ._common import (
@@ -60,7 +61,7 @@ from .vec_emit import _emit_vec
def is_smem_ewise(spec):
"""Predicate factory: dispatch accepted iff all operands in ``shared*``."""
- def check(op_call: TileOpCall, sctx: DispatchContext) -> tuple[bool, str | None]:
+ def check(op_call: TensorCall, sctx: DispatchContext) -> tuple[bool, str | None]:
if not sctx.is_target("cuda"):
return False, "non-cuda target"
if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
@@ -129,7 +130,7 @@ def _max_layout_vec(plan, total: int, thread_cnt: int) -> int:
# -----------------------------------------------------------------------------
# Main entry
# -----------------------------------------------------------------------------
-def emit_smem(op_call: TileOpCall, spec, sctx: DispatchContext) -> Function:
+def emit_smem(op_call: TensorCall, spec, sctx: DispatchContext) -> Function:
plan, msg = spec.parse(op_call)
if msg is not None or plan is None:
fail(msg or "parse failed")
diff --git a/python/tvm/backend/cuda/tile_primitive/exec_scope_utils.py b/python/tvm/backend/cuda/tile_primitive/exec_scope_utils.py
index 6167881027..ae47ec3980 100644
--- a/python/tvm/backend/cuda/tile_primitive/exec_scope_utils.py
+++ b/python/tvm/backend/cuda/tile_primitive/exec_scope_utils.py
@@ -21,7 +21,7 @@ from collections.abc import Callable
from tvm.script import tirx as T
from tvm.tirx import Function
from tvm.tirx.operator.tile_primitive import DispatchContext
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
def macro_or_function(macro: Callable, need_macro: bool = False) -> Callable:
@@ -89,14 +89,14 @@ def thread_selector(sctx: DispatchContext, inner_impl, macro: bool = False) -> C
raise ValueError(f"thread_selector: unsupported exec_scope {name!r}")
-def single_thread(op_call: TileOpCall, sctx: DispatchContext) -> bool:
+def single_thread(op_call: TensorCall, sctx: DispatchContext) -> bool:
"""Predicate for dispatchers that require a single-thread execution scope."""
del op_call
return sctx.is_thread
def exec_scope_ok(
- op_call: TileOpCall, sctx: DispatchContext, expected_scopes: list[str]
+ op_call: TensorCall, sctx: DispatchContext, expected_scopes: list[str]
) -> tuple[bool, str | None]:
"""Predicate helper: check that ``sctx.scope_kind`` is in *expected_scopes*."""
del op_call
diff --git a/python/tvm/backend/cuda/tile_primitive/gemm/__init__.py b/python/tvm/backend/cuda/tile_primitive/gemm/__init__.py
index f6caa8d58f..6d25095b04 100644
--- a/python/tvm/backend/cuda/tile_primitive/gemm/__init__.py
+++ b/python/tvm/backend/cuda/tile_primitive/gemm/__init__.py
@@ -15,11 +15,6 @@
# specific language governing permissions and limitations
# under the License.
-"""CUDA synchronous ``gemm`` lowerings (warp-level ``mma.sync`` tensor core).
-
-Importing this package registers every synchronous CUDA ``gemm`` dispatch
-candidate as a side effect (each submodule calls ``register_dispatch`` at
-import time). It is the synchronous counterpart to ``gemm_async``.
-"""
+"""CUDA warp-level mma.sync tensor instruction lowering."""
from .mma_m16n8k_ import *
diff --git a/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py b/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py
index f9cff4fa97..b27d514a1c 100644
--- a/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py
+++ b/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py
@@ -23,13 +23,8 @@ from tvm.script import tirx as T
from tvm.sym.analyzer import Analyzer
from tvm.tirx import Function
from tvm.tirx.layout import TileLayout
-from tvm.tirx.operator.tile_primitive import (
- DispatchContext,
- fail,
- predicate,
- register_dispatch,
-)
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.tensor_instruction import TensorCall
@dataclass(frozen=True)
@@ -220,7 +215,7 @@ def _same_iters(a, b):
)
-def _full_active_lanes(op: TileOpCall, sctx: DispatchContext):
+def _full_active_lanes(op: TensorCall, sctx: DispatchContext):
"""The active thread set (sctx.intra) must be complete and un-narrowed.
mma.sync.aligned is collective over every active thread; an enclosing if
@@ -249,7 +244,7 @@ def _full_active_lanes(op: TileOpCall, sctx: DispatchContext):
return True
-def _no_replica(op: TileOpCall, sctx: DispatchContext):
+def _no_replica(op: TensorCall, sctx: DispatchContext):
"""All operand layouts must have no replica (no broadcast/duplicated axes)."""
for region, name in zip(op.args[:4], ("D", "A", "B", "C")):
if region.source.layout.replica:
@@ -257,17 +252,7 @@ def _no_replica(op: TileOpCall, sctx: DispatchContext):
return True
-@register_dispatch(
- "gemm",
- "cuda",
- variant="mma.m16n8k*",
- priority=10,
- when=[
- predicate("full_active_lanes", _full_active_lanes),
- predicate("no_replica", _no_replica),
- ],
-)
-def gemm_cuda_mma_dispatch(op: TileOpCall, sctx: DispatchContext) -> Function:
+def gemm_cuda_mma_dispatch(op: TensorCall, sctx: DispatchContext) -> Function:
"""``gemm`` -> warp-level ``mma.sync`` of the m16n8k* family.
This is the ``"mma.m16n8k*"`` variant. It targets the m16n8k* tensor-core
diff --git a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
index f798828467..f33c484ebc 100644
--- a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
+++ b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
@@ -15,12 +15,7 @@
# specific language governing permissions and limitations
# under the License.
-"""Implementation of gemm_async operator dispatch for CUDA targets.
-
-Registered op: gemm_async (1 variant: "tcgen05").
-See the @register_dispatch block below for detailed documentation with
-before/after IR examples.
-"""
+"""CUDA tcgen05.mma and tcgen05.mma.block_scale tensor instruction lowering."""
import functools
import operator
@@ -53,8 +48,8 @@ from tvm.tirx.layout import (
tmem_datapath_layout,
tmem_mma_operand_layout,
)
-from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.tensor_instruction import TensorCall
from ...cpp.descriptors import (
_check_tcgen05_mma_matrix_shape,
@@ -64,7 +59,6 @@ from ...cpp.descriptors import (
encode_instr_descriptor_dense_uint32,
)
from ..common import get_st_extent, smem_desc_add_16B_offset
-from ..exec_scope_utils import single_thread
from ..layout_utils import strip_swizzle_to_tile
from ..tma_utils import (
SwizzleMode,
@@ -361,7 +355,7 @@ def _layout_matches_datapath_f(tmem_buf) -> bool:
return False
-def gemm_async_tcgen05_impl(op_call: TileOpCall, sctx: DispatchContext) -> Function:
+def gemm_async_tcgen05_impl(op_call: TensorCall, sctx: DispatchContext) -> Function:
"""Schedule an asynchronous GEMM operation using tcgen05.mma (Blackwell Tensor Core).
Computes C = A @ B (with optional transpose on A/B and accumulation).
@@ -371,7 +365,7 @@ def gemm_async_tcgen05_impl(op_call: TileOpCall, sctx: DispatchContext) -> Funct
so that only one thread in the warp issues the MMA instruction.
Args:
- op_call: The TileOpCall containing:
+ op_call: The TensorCall containing:
Regular (6 args):
- args[0:3]: C, A, B buffer regions
- args[3:6]: transA, transB, accum flags
@@ -398,7 +392,7 @@ def gemm_async_tcgen05_impl(op_call: TileOpCall, sctx: DispatchContext) -> Funct
AssertionError: If shape/layout constraints are not satisfied.
"""
warp_scope = sctx.is_warp
- op_call = TileOpCall.downcast(op_call)
+ op_call = TensorCall.decode(op_call)
is_block_scaled = op_call.is_block_scaled
C_buffer_region: TensorRegion = op_call.output
@@ -433,10 +427,10 @@ def gemm_async_tcgen05_impl(op_call: TileOpCall, sctx: DispatchContext) -> Funct
assert C_type == "float32", f"tcgen05 schedule expected C_type=float32, got {C_type}"
# fp32/bf16 storage may still use tf32 MMA semantics via is_AB_tf32.
- is_AB_tf32 = op_call.config.get("is_AB_tf32", False)
+ is_AB_tf32 = op_call.options.get("is_AB_tf32", False)
# Emit the PTX ``tcgen05.mma.ws`` weight-stationary form only for kernels
# that explicitly require that tcgen05 ABI.
- weight_stationary = bool(op_call.config.get("weight_stationary", False))
+ weight_stationary = bool(op_call.options.get("weight_stationary", False))
A_sem = "tf32" if is_AB_tf32 else A_type
B_sem = "tf32" if is_AB_tf32 else B_type
@@ -513,13 +507,13 @@ def gemm_async_tcgen05_impl(op_call: TileOpCall, sctx: DispatchContext) -> Funct
transA = bool(transA) if isinstance(transA, IntImm) else transA
transB = bool(transB) if isinstance(transB, IntImm) else transB
- cta_group = op_call.config.get("cta_group", 1)
+ cta_group = op_call.options.get("cta_group", 1)
if isinstance(cta_group, IntImm):
cta_group = cta_group.value
assert cta_group in [1, 2], f"tcgen05 schedule expected cta_group=1 or 2, got {cta_group}"
# descI (pre-encoded uint32 instruction descriptor): rejected on the dense
# path (dispatcher encodes it); block-scaled callers may still pass it in.
- descI = op_call.config.get("descI", None)
+ descI = op_call.options.get("descI", None)
if descI is not None and not is_block_scaled:
raise ValueError(
"descI was removed: the dispatcher encodes the instruction descriptor itself"
@@ -890,7 +884,7 @@ def gemm_async_tcgen05_impl(op_call: TileOpCall, sctx: DispatchContext) -> Funct
MMA_K = 16
MMA_N_MIN = 8 if cta_group == 1 else 16 # Minimum N dimension
- explicit_mma_tile = _get_explicit_mma_tile(op_call.config)
+ explicit_mma_tile = _get_explicit_mma_tile(op_call.options)
if explicit_mma_tile is None:
M_mma, N_mma = _choose_mma_tile(M, N, cta_group, MMA_N_MIN)
else:
@@ -965,7 +959,7 @@ def gemm_async_tcgen05_impl(op_call: TileOpCall, sctx: DispatchContext) -> Funct
# Packed Layout-E C (M, 2, N//2) is uniquely the cta_group::1 M=64 .ws datapath
# (PTX §9.7.16.10.5), so .ws is inferred; weight_stationary=False on it is rejected.
if packed_n2 and not is_2x2:
- explicit_ws = op_call.config.get("weight_stationary")
+ explicit_ws = op_call.options.get("weight_stationary")
if (
isinstance(explicit_ws, IntImm)
and str(explicit_ws.ty.dtype) == "bool"
@@ -1227,7 +1221,7 @@ def gemm_async_tcgen05_impl(op_call: TileOpCall, sctx: DispatchContext) -> Funct
# smem_desc modes: "hoist" (default, encode once after alloc), "local_hoist"
# (encode at call site, reuse via add_16B_offset), "encode"/"recompute" (per MMA).
- smem_desc_mode = op_call.config.get("smem_desc", "hoist")
+ smem_desc_mode = op_call.options.get("smem_desc", "hoist")
if isinstance(smem_desc_mode, tvm.ir.StringImm):
smem_desc_mode = smem_desc_mode.value
local_hoist = smem_desc_mode == "local_hoist"
@@ -1502,50 +1496,3 @@ def gemm_async_tcgen05_impl(op_call: TileOpCall, sctx: DispatchContext) -> Funct
# fmt: on
return impl
-
-
-# === Variant: gemm_async/tcgen05 (priority=10) ===
-#
-# When: gemm_async op at single-thread exec scope on Blackwell (SM100+).
-# Requires A in smem (with TMA-compatible swizzle layout) or tmem, B in smem, accum in tmem.
-#
-# Before (TileOpCall — regular MMA):
-# Tx.gemm_async(C_tmem[0:64, 0:256], A_smem[0:64, 0:64], B_smem[0:256, 0:64])
-# # A: shared float16, B: shared float16, C: tmem float32
-#
-# After (encodes instruction descriptor + calls tcgen05.mma):
-# descI_local: uint32
-# T.cuda.tcgen05.encode_instr_descriptor(
-# &descI_local, C_type="f32", A_type="f16", B_type="f16",
-# M=64, N=256, MMA_K=64, transA=False, transB=True, cta_group=1)
-# T.ptx[mma_chain](..., descA_buf[0], descB_buf[0], descI_local, ...)
-#
-# Before (TileOpCall — block-scaled fp8 MMA):
-# Tx.gemm_async(C_tmem, A_smem, B_smem,
-# scale_A=SFA_tmem, scale_B=SFB_tmem)
-# # A/B: shared float8_e4m3, SFA/SFB: tmem float8_e8m0fnu
-#
-# After (adds scale factor descriptors):
-# T.ptx[mma_chain](..., descA, descB, descI,
-# scale_A=sfA_desc, scale_B=sfB_desc)
-#
-# Scale factor layout (sf_tmem_layout) must match tcgen05 hardware requirements:
-# rows = M or N, sf_mma_k = ceil(MMA_K / sf_block_size), specific TileLayout
-# structure with direct_sum atom tiling.
-@register_dispatch(
- "gemm_async",
- "cuda",
- variant="tcgen05",
- priority=10,
- when=[
- predicate(
- "single_thread_or_warp",
- lambda op, sctx: (
- single_thread(op, sctx) or sctx.is_warp,
- f"unsupported exec_scope {sctx.exec_scope}, expected single thread or warp scope",
- ),
- )
- ],
-)
-def gemm_async_dispatch_tcgen05(op_call: TileOpCall, sctx: DispatchContext) -> Function:
- return gemm_async_tcgen05_impl(op_call, sctx)
diff --git a/python/tvm/backend/cuda/tile_primitive/gemm_utils.py b/python/tvm/backend/cuda/tile_primitive/gemm_utils.py
index c6a1275638..269e593dbb 100644
--- a/python/tvm/backend/cuda/tile_primitive/gemm_utils.py
+++ b/python/tvm/backend/cuda/tile_primitive/gemm_utils.py
@@ -20,10 +20,10 @@
from tvm.sym.analyzer import Analyzer
from tvm.tirx import Var
from tvm.tirx.operator.tile_primitive import DispatchContext
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
-def validate_gemm_op(op_call: TileOpCall, sctx: DispatchContext) -> bool:
+def validate_gemm_op(op_call: TensorCall, sctx: DispatchContext) -> bool:
"""Sanity check for gemm op"""
C_buffer_region, A_buffer_region, B_buffer_region = op_call.args[:3]
C: Var = C_buffer_region.source
diff --git a/python/tvm/backend/cuda/tile_primitive/permute_layout/__init__.py b/python/tvm/backend/cuda/tile_primitive/permute_layout/__init__.py
deleted file mode 100644
index e406e9c3fd..0000000000
--- a/python/tvm/backend/cuda/tile_primitive/permute_layout/__init__.py
+++ /dev/null
@@ -1,18 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-from .warp_xor_swizzle import *
diff --git a/python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py b/python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py
deleted file mode 100644
index 1f8bff6da3..0000000000
--- a/python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py
+++ /dev/null
@@ -1,436 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""CUDA permute_layout dispatch: warp register-staged in-place transpose with
-optional per-lane XOR-swizzle to avoid SMEM bank conflicts on the write phase.
-
-The dispatcher reasons about the **layout's shard**, not the buffer's
-declared shape (the two can differ — a buffer with ``shape=(PIPE, M, K)``
-may carry a layout whose shard has more dims internally, with grouping
-mapping shard segments onto buffer dims). Concretely:
-
- src_sliced = src.layout.slice(src.shape, region).canonicalize()
- dst_sliced = dst.layout.slice(dst.shape, region).canonicalize()
- # If the two sliced shards have different structures (which is common —
- # a linear layout collapses to 1D under canon while a transposed one
- # keeps its multi-dim structure), regroup src to dst's shape.
- if src_sliced.shard != dst_sliced.shard:
- src_sliced, _ = src_sliced.group(dst.shard.extents)
- extent = [int(it.extent) for it in dst_sliced.shard] # iteration shape
- src_str = [int(it.stride) for it in src_sliced.shard]
- dst_str = [int(it.stride) for it in dst_sliced.shard]
-
-The algorithm:
-
- regs[P]
- for r in 0..P:
- j = r XOR ((lane >> SHIFT) & MASK)
- i = lane + j * 32 # flat logical index
- idx = decompose(i, extent) # iter multi-dim index
- regs[r] = src[project(idx, src.shape, slice_starts)]
- warp_sync()
- for r in 0..P:
- j = r XOR ((lane >> SHIFT) & MASK)
- i = lane + j * 32
- idx = decompose(i, extent)
- dst[project(idx, dst.shape, slice_starts)] = regs[r]
- warp_sync()
-
-where ``project`` mixed-radix-folds the iter shard dims back onto the
-buffer's iterated slice dims (so the emit's index matches buf.shape rank,
-which TIR's TensorLoad/Store requires).
-
-SHIFT and MASK are chosen by simulating the bank pattern at the **shard
-granularity** (where strides are affine), trying k = 0, 1, …, log2(P)
-and picking the smallest k that makes both phases bank-conflict-free.
-
-Correctness rests on:
-
-* For each lane, ``r ↦ r XOR const`` is a bijection on ``[0, P)``.
-* Therefore (lane, r) ↔ flat over [0, V).
-* Both layouts are verified bijections on the slice (every logical
- position has a unique byte offset under that layout).
-* The mixed-radix projection from iter shard idx to buf coord is exactly
- what TIR's TensorLoad does internally when buf.shape rank < shard rank
- — so iter shard's strides and the buffer-indexed byte offset agree.
-"""
-
-from __future__ import annotations
-
-import math
-
-from tvm.ir import TensorRegion
-from tvm.runtime import DataType
-from tvm.script import tirx as T
-from tvm.tirx import Function, IntImm, is_tensor_var
-from tvm.tirx.layout import TileLayout, _flatten_coord
-from tvm.tirx.operator.tile_primitive import DispatchContext, fail, register_dispatch
-from tvm.tirx.stmt import TileOpCall
-
-from ..common import get_indices, get_st_extent
-
-# ---------- helpers ----------------------------------------------------------
-
-
-def _as_buffer_and_region(arg):
- """Normalize a Buffer or TensorRegion to (buffer, start_list, extent_list)."""
- if is_tensor_var(arg):
- buf = arg
- extent = list(buf.ty.shape)
- st = [0] * len(extent)
- elif isinstance(arg, TensorRegion):
- if not is_tensor_var(arg.source):
- raise TypeError("permute_layout requires a TensorRegion with a TensorVar source")
- buf = arg.source
- st, extent = get_st_extent(arg)
- else:
- raise TypeError(f"unexpected permute_layout arg type: {type(arg)}")
- return buf, list(st), list(extent)
-
-
-def _as_int(x):
- """Return int(x) if x is int-like, else None."""
- if isinstance(x, int):
- return x
- if isinstance(x, IntImm):
- return int(x.value)
- if hasattr(x, "value") and isinstance(x.value, int):
- return int(x.value)
- try:
- return int(x)
- except (TypeError, ValueError):
- return None
-
-
-def _layout_shard_int(layout):
- """Return (extents, strides) as int lists from a TileLayout's shard, or (None, None)."""
- if not isinstance(layout, TileLayout):
- return None, None
- extents, strides = [], []
- for it in layout.shard:
- e = _as_int(it.extent)
- s = _as_int(it.stride)
- if e is None or s is None:
- return None, None
- extents.append(e)
- strides.append(s)
- return extents, strides
-
-
-def _decompose_row_major(i, extent):
- out, rem = [], i
- for e in reversed(extent):
- out.append(rem % e)
- rem //= e
- return list(reversed(out))
-
-
-def _eval_offset(idx, strides):
- return sum(i * s for i, s in zip(idx, strides))
-
-
-def _check_bijection(extent, strides):
- """Iteration extents + strides define a bijection on [0, V)?"""
- V = math.prod(extent)
- seen = set()
- for i in range(V):
- off = _eval_offset(_decompose_row_major(i, extent), strides)
- if off in seen:
- return False
- seen.add(off)
- return len(seen) == V
-
-
-def _bank_free(extent, strides, dtype_bytes, P, k):
- """For every register slot r ∈ [0, P), do the 32 lanes hit 32 distinct banks?"""
- T, BANKS, BANK_W = 32, 32, 4
- shift = 5 - k
- mask = (1 << k) - 1
- for r in range(P):
- seen = set()
- for lane in range(T):
- j = r ^ ((lane >> shift) & mask)
- flat = lane + j * T
- idx = _decompose_row_major(flat, extent)
- off_bytes = _eval_offset(idx, strides) * dtype_bytes
- bank = (off_bytes // BANK_W) % BANKS
- if bank in seen:
- return False
- seen.add(bank)
- return True
-
-
-def _choose_xor_k(extent, src_strides, dst_strides, dtype_bytes, P):
- max_k = int(math.log2(P)) if P > 0 else 0
- for k in range(max_k + 1):
- if _bank_free(extent, src_strides, dtype_bytes, P, k) and _bank_free(
- extent, dst_strides, dtype_bytes, P, k
- ):
- return k
- return None
-
-
-# ---------- validator + dispatch impl ---------------------------------------
-
-
-def _gather(op_call):
- op_call = TileOpCall.downcast(op_call)
- dst_arg, src_arg = op_call.args[0], op_call.args[1]
- src_buf, src_st, src_ext = _as_buffer_and_region(src_arg)
- dst_buf, dst_st, dst_ext = _as_buffer_and_region(dst_arg)
- return src_buf, src_st, src_ext, dst_buf, dst_st, dst_ext
-
-
-def _why_reject(op_call, sctx):
- if not sctx.is_warp:
- return f"scope {sctx.scope_kind!r} is not 'warp'"
- if "threadIdx.y" in sctx.launch_params or "threadIdx.z" in sctx.launch_params:
- return "multi-dim threadIdx is not supported"
-
- src_buf, src_st, src_ext, dst_buf, dst_st, dst_ext = _gather(op_call)
-
- if src_buf.dtype != dst_buf.dtype:
- return f"dtype mismatch: dst={dst_buf.dtype} vs src={src_buf.dtype}"
-
- src_ext_i = [_as_int(e) for e in src_ext]
- dst_ext_i = [_as_int(e) for e in dst_ext]
- if None in src_ext_i or None in dst_ext_i:
- return "extents must be compile-time integers"
- if src_ext_i != dst_ext_i:
- return f"slice shape mismatch: src={src_ext_i} vs dst={dst_ext_i}"
-
- dtype_bytes = DataType(src_buf.dtype).bits // 8
- if dtype_bytes not in (1, 2, 4, 8, 16):
- return f"unsupported dtype byte width: {dtype_bytes}"
-
- if not isinstance(src_buf.layout, TileLayout):
- return "src buffer's layout is not a plain TileLayout"
- if not isinstance(dst_buf.layout, TileLayout):
- return "dst buffer's layout is not a plain TileLayout"
-
- # Slice + canonicalize both layouts. The result's shard describes the
- # iteration domain; runtime starts (like ``ks``) are folded into the
- # layout's offset, separate from the shard's affine part.
- src_region = [(s, s + e) for s, e in zip(src_st, src_ext)]
- dst_region = [(s, s + e) for s, e in zip(dst_st, dst_ext)]
- src_sliced = src_buf.layout.slice(list(src_buf.shape), src_region)
- dst_sliced = dst_buf.layout.slice(list(dst_buf.shape), dst_region)
- if src_sliced is None or dst_sliced is None:
- return "layout.slice failed"
- src_sliced = src_sliced.canonicalize()
- dst_sliced = dst_sliced.canonicalize()
-
- # Iteration shape: regroup dst onto the iterated buf dims; the result's
- # shard may stay finer than iter_buf_extents (one buf dim ↔ several shard
- # dims via seps), which is fine. Then regroup src to match dst's shard
- # extents exactly so both phases share the same iteration index space.
- iter_buf_extents = [e for e in src_ext_i if e != 1]
- try:
- dst_grouped, dst_seps = dst_sliced.group(iter_buf_extents)
- src_grouped, _ = src_sliced.group([int(it.extent) for it in dst_grouped.shard])
- except Exception as e:
- return f"layout.group failed: {e}"
-
- dst_ext_, dst_str_ = _layout_shard_int(dst_grouped)
- src_ext_, src_str_ = _layout_shard_int(src_grouped)
- if dst_ext_ is None or src_ext_ is None:
- return "regrouped layout shard contains non-integer extent/stride"
- if src_ext_ != dst_ext_:
- return f"src shard {src_ext_} doesn't match dst shard {dst_ext_} after regrouping"
-
- extent = dst_ext_
- V = math.prod(extent)
- T = 32
- if V == 0 or V % T != 0:
- return f"volume {V} not divisible by warp size {T}"
- P = V // T
- if P == 0 or (P & (P - 1)) != 0 or P > T:
- return f"per-thread count {P} must be power of 2 in [1, {T}]"
-
- if not _check_bijection(extent, src_str_):
- return "src layout (regrouped) is not a bijection on the slice"
- if not _check_bijection(extent, dst_str_):
- return "dst layout is not a bijection on the slice"
- return None
-
-
-def _impl(op_call, sctx):
- src_buf, src_st, src_ext, dst_buf, dst_st, dst_ext = _gather(op_call)
- src_ext_i = [_as_int(e) for e in src_ext]
-
- src_region = [(s, s + e) for s, e in zip(src_st, src_ext)]
- dst_region = [(s, s + e) for s, e in zip(dst_st, dst_ext)]
- src_sliced = src_buf.layout.slice(list(src_buf.shape), src_region).canonicalize()
- dst_sliced = dst_buf.layout.slice(list(dst_buf.shape), dst_region).canonicalize()
-
- iter_buf_extents = [e for e in src_ext_i if e != 1]
- dst_grouped, dst_seps = dst_sliced.group(iter_buf_extents)
- src_grouped, _ = src_sliced.group([int(it.extent) for it in dst_grouped.shard])
-
- extent, dst_str_ = _layout_shard_int(dst_grouped)
- _, src_str_ = _layout_shard_int(src_grouped)
- V = math.prod(extent)
- P = V // 32
- dtype_bytes = DataType(src_buf.dtype).bits // 8
-
- k_opt = _choose_xor_k(extent, src_str_, dst_str_, dtype_bytes, P)
- if k_opt is None:
- fail(f"no XOR-bits k ∈ [0, log2(P)={int(math.log2(P))}] makes both phases bank-free")
-
- shift = 5 - k_opt
- mask = (1 << k_opt) - 1
-
- iter_buf_dims = [i for i, e in enumerate(src_ext_i) if e != 1]
- seps = list(dst_seps)
-
- def _project(iter_idx, st_list):
- buf_idx = list(st_list)
- for bi in range(len(seps) - 1):
- lo, hi = seps[bi], seps[bi + 1]
- flat = _flatten_coord(iter_idx[lo:hi], extent[lo:hi])
- buf_idx[iter_buf_dims[bi]] = st_list[iter_buf_dims[bi]] + flat
- return tuple(buf_idx)
-
- tid_x = sctx.launch_params["threadIdx.x"][0]
- dtype = src_buf.dtype
-
- # Shared 32/64b: base ptr + stride offset avoids buf[] flatten IMAD path.
- direct = (
- dtype_bytes in (4, 8)
- and "shared" in str(src_buf.scope())
- and "shared" in str(dst_buf.scope())
- )
- ptx_t = f"b{dtype_bytes * 8}"
- # Built out here: a Python string bound inside the traced body is not
- # something the parser can carry.
- ld_chain, st_chain = f"ld.shared.{ptx_t}", f"st.shared.{ptx_t}"
- # ld/st only move bits, so use an unsigned container of the matching width:
- # ``ptx.ld(..., bN)`` rejects a float return dtype, and the permute is a pure
- # byte shuffle, so a float32/float64 tile loads/stores correctly as uint.
- bits_dtype = f"uint{dtype_bytes * 8}"
-
- def _iter_off(iter_idx, strides):
- return sum(iter_idx[d] * strides[d] for d in range(len(strides)))
-
- # fmt: off
- if direct:
- @T.function(check_well_formed=False)
- def impl():
- warp_size = T.meta_var(32)
- lane_id = T.meta_var(tid_x % warp_size)
- regs = T.alloc_tensor((P,), bits_dtype, scope="local")
- base_src = T.meta_var(src_buf.ptr_to(list(src_st)))
- base_dst = T.meta_var(dst_buf.ptr_to(list(dst_st)))
- # Phase 1: read via L_src
- for r in T.unroll(0, P):
- j = T.meta_var(r ^ ((lane_id >> shift) & mask))
- flat = T.meta_var(lane_id + j * warp_size)
- iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent))
- off = T.meta_var(_iter_off(iter_idx, src_str_))
- ptr = T.meta_var(T.ptr_byte_offset(base_src, off * dtype_bytes, dtype))
- T.ptx[ld_chain](regs[r], ptr)
- T.cuda.warp_sync()
- # Phase 2: write via L_dst
- for r in T.unroll(0, P):
- j = T.meta_var(r ^ ((lane_id >> shift) & mask))
- flat = T.meta_var(lane_id + j * warp_size)
- iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent))
- off = T.meta_var(_iter_off(iter_idx, dst_str_))
- ptr = T.meta_var(T.ptr_byte_offset(base_dst, off * dtype_bytes, dtype))
- T.ptx[st_chain](ptr, regs[r])
- T.cuda.warp_sync()
- else:
- @T.function(check_well_formed=False)
- def impl():
- warp_size = T.meta_var(32)
- lane_id = T.meta_var(tid_x % warp_size)
- regs = T.alloc_tensor((P,), dtype, scope="local")
- # Phase 1: read via L_src
- for r in T.unroll(0, P):
- j = T.meta_var(r ^ ((lane_id >> shift) & mask))
- flat = T.meta_var(lane_id + j * warp_size)
- iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent))
- src_idx = T.meta_var(_project(iter_idx, src_st))
- regs[r] = src_buf[tuple(src_idx)]
- T.cuda.warp_sync()
- # Phase 2: write via L_dst
- for r in T.unroll(0, P):
- j = T.meta_var(r ^ ((lane_id >> shift) & mask))
- flat = T.meta_var(lane_id + j * warp_size)
- iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent))
- dst_idx = T.meta_var(_project(iter_idx, dst_st))
- dst_buf[tuple(dst_idx)] = regs[r]
- T.cuda.warp_sync()
- # fmt: on
- return impl
-
-
-# === Variant: permute_layout/warp_xor_swizzle (priority=20) ============
-#
-# When: warp scope; matching dst/src dtype + slice shape; both buffers carry
-# a plain TileLayout; after slice + canonicalize (and regrouping src to dst's
-# structure if needed), the iteration extents form a power-of-2 ≤32 elements
-# per lane; both layouts are bijections on the slice; and there exists an
-# XOR-bits ``k`` that makes both phases bank-conflict-free.
-#
-# Buffer ``shape`` rank does NOT need to equal layout ``shard`` rank — the
-# dispatcher uses the layout shard for iteration (after slice+canon) and
-# projects back onto ``buf.shape`` via mixed-radix grouping for the emit.
-#
-# Before (TileOpCall):
-# with T.warp():
-# # SFA_smem: u32 (PIPE, BLK_SFA//32, 32), layout shard 4D
-# # (PIPE, BLK_SFA//128, 4, 32) strides (BLK_SFA, 128, 32, 1)
-# # SFA_post: same shape; layout shard 4D, strides (BLK_SFA, 128, 1, 4)
-# Tx.permute_layout(SFA_post[ks, :, :], SFA_smem[ks, :, :])
-#
-# After (BLK_SFA=128, P=4, k=2, shift=3):
-# lane_id = threadIdx.x % 32
-# regs = T.alloc_tensor((4,), "uint32", scope="local")
-# for r in T.unroll(4):
-# j = r ^ ((lane_id >> 3) & 0x3)
-# flat = lane_id + j * 32
-# (g, l) = decompose(flat, extent=[4, 32])
-# regs[r] = src[ks, g, l]
-# T.cuda.warp_sync()
-# for r in T.unroll(4):
-# j = r ^ ((lane_id >> 3) & 0x3)
-# flat = lane_id + j * 32
-# (g, l) = decompose(flat, extent=[4, 32])
-# dst[ks, g, l] = regs[r]
-# T.cuda.warp_sync()
-@register_dispatch(
- "permute_layout",
- "cuda",
- variant="warp_xor_swizzle",
- priority=20,
-)
-def permute_layout_dispatch(op: TileOpCall, sctx: DispatchContext) -> Function:
- reason = _why_reject(op, sctx)
- if reason is not None:
- fail(reason)
- return _impl(op, sctx)
-
-
-__all__ = [
- "_bank_free",
- "_check_bijection",
- "_choose_xor_k",
- "_decompose_row_major",
- "_eval_offset",
- "permute_layout_dispatch",
-]
diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/__init__.py b/python/tvm/backend/cuda/tile_primitive/reduction/__init__.py
deleted file mode 100644
index 8b7ad67057..0000000000
--- a/python/tvm/backend/cuda/tile_primitive/reduction/__init__.py
+++ /dev/null
@@ -1,20 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-from .local import *
-from .shared import *
-from .sm100_packed import *
diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/local.py b/python/tvm/backend/cuda/tile_primitive/reduction/local.py
deleted file mode 100644
index d82e7eafc1..0000000000
--- a/python/tvm/backend/cuda/tile_primitive/reduction/local.py
+++ /dev/null
@@ -1,487 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""CUDA reduction operator dispatch: local-memory variant.
-
-Registered ops: sum, max, min.
-
-When: dst and src are both local-scope buffers with matching dtype, on CUDA.
-
-(A) Thread scope -- sequential per-element reduction
- (_emit_reduction_local_thread_wise):
-
-Before:
- Tx.sum(B_local[0:2, 0:3], A_local[0:2, 0:3, 0:4], [-1], False)
-
-After (scheduled Function, spatial_len=6, reduction_len=4):
- for spa in range(6):
- B_local[spa] = T.float32(0.0) # init (skipped if accum)
- for red in range(4):
- B_local[spa] = B_local[spa] + A_local[spa * 4 + red]
-
-(B) Warp/Warpgroup scope -- layout-driven reduction
- (_emit_reduction_local_view):
- Requires TileLayout with valid thread-partition. Decomposes layout to
- identify thread-local elements, then optionally shuffles partial sums.
-
- thread_reduce=False: local-only, no shuffle (warp and warpgroup).
- thread_reduce=True: local reduction + cross-thread shfl_xor steps (warp only).
- accum=True + shuffle: saves old dst before reduce+shuffle, combines after (warp only).
-
-Before:
- Tx.warp.sum(red_view[0:16, 0:4], acc_view[0:16, 0:128], [-1], False,
- thread_reduce=True)
-
-After (scheduled Function, local_total=2, local_red=32, 2 shuffle steps):
- src_local = acc_view.view(64)
- dst_local = red_view.view(2)
- for spa in range(2):
- dst_local[spa] = T.float32(0.0)
- for red in range(32):
- dst_local[spa] = dst_local[spa] + src_local[...]
- dst_local[spa] = dst_local[spa] + shfl_xor(..., 1, 32, 32)
- dst_local[spa] = dst_local[spa] + shfl_xor(..., 2, 32, 32)
-"""
-
-import functools
-import operator
-from typing import Any
-
-from tvm.ir import TensorRegion
-from tvm.script import tirx as T
-from tvm.sym.analyzer import Analyzer
-from tvm.tirx import Function
-from tvm.tirx.layout import TileLayout, laneid
-from tvm.tirx.operator.tile_primitive import DispatchContext, fail
-from tvm.tirx.operator.tile_primitive.common import ReduceOpType
-from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
-from tvm.tirx.stmt import TileOpCall
-
-from ..common import get_indices, get_st_extent
-from ..layout_utils import get_local_region, get_sublayout_from_region
-from .utils import (
- _REDUCE_OP_TO_STR,
- _analyze_axes,
- _analyze_layout_dims,
- _build_local_dim_map,
- _compute_shuffle_masks,
- _match_reduction_storage_scope,
- _reduction_args,
- _validate_reduction_layout,
- reduce_default_value_table,
- reduce_op_table,
-)
-
-
-def _analyze_shuffle_reduce(src_layout, dst_layout):
- """Analyze src/dst layouts for laneid shard->replica reduce pattern.
-
- Returns (reduce_width, local_elems) if the pattern matches, or None.
- - reduce_width: number of lanes participating in each group's reduction
- - local_elems: per-thread element count (product of non-laneid shard extents)
- """
- if src_layout.is_swizzle() or dst_layout.is_swizzle():
- return None
-
- src_canon = src_layout.canonicalize()
- dst_canon = dst_layout.canonicalize()
-
- # Extract laneid iters from shard and replica
- src_laneid_shard = [it for it in src_canon.shard if it.axis == laneid]
- dst_laneid_replica = [it for it in dst_canon.replica if it.axis == laneid]
-
- # src shard must contain laneid (data distributed across lanes)
- if not src_laneid_shard:
- return None
- # dst replica must contain laneid (result broadcast to lanes)
- if not dst_laneid_replica:
- return None
-
- # laneid span must be 32 (full warp)
- src_laneid_span = 1 + sum(abs(int(it.stride)) * (int(it.extent) - 1) for it in src_laneid_shard)
- if src_laneid_span != 32:
- return None
-
- reduce_width = functools.reduce(operator.mul, [int(it.extent) for it in dst_laneid_replica], 1)
- if reduce_width <= 0 or reduce_width > 32 or (reduce_width & (reduce_width - 1)) != 0:
- return None # must be power of 2
-
- # local_elems = product of non-laneid shard extents in src
- src_non_laneid = [it for it in src_canon.shard if it.axis != laneid]
- local_elems = functools.reduce(operator.mul, [int(it.extent) for it in src_non_laneid], 1)
-
- return reduce_width, local_elems
-
-
-def _gen_warp_shuffle_reduce(src, dst, reduce_width, local_elems, accum, op_type, init_value):
- """Generate warp shuffle reduce codegen for laneid shard->replica pattern.
-
- Unified for both full warp (reduce_width=32) and partial warp (e.g. reduce_width=8).
- """
- is_same_buffer = src.same_as(dst)
- op_str = _REDUCE_OP_TO_STR[op_type]
-
- # fmt: off
- @T.function(check_well_formed=False)
- def impl():
- src_local = src.local(local_elems, layout=src.layout.storage())
- dst_local = dst.local(local_elems, layout=dst.layout.storage())
- for k in T.serial(local_elems):
- if T.constexpr(not is_same_buffer):
- dst_local[k] = src_local[k]
- dst_local[k] = T.cuda.warp_reduce(dst_local[k], op_str, reduce_width)
- # fmt: on
-
- return impl
-
-
-def validate_reduction_local(op: TileOpCall, sctx: DispatchContext) -> tuple[bool, str | None]:
- """Validate reduction in local memory."""
- op = TileOpCall.downcast(op)
- dst_br, src_br = op.output, op.input
- dst, src = dst_br.source, src_br.source
-
- if not (src.scope() == "local" and dst.scope() == "local" and sctx.is_target("cuda")):
- return False, "expected local scope and CUDA target"
- if src.dtype != dst.dtype:
- return False, f"dtype mismatch: src={src.dtype} dst={dst.dtype}"
-
- if sctx.is_thread:
- return True, None # thread-wise reduction
- elif sctx.scope_kind in ["warp", "warpgroup"]:
- if not sctx.is_warp and op.config.get("thread_reduce", False):
- return (
- False,
- "thread_reduce=True is only supported in warp scope; "
- "warpgroup local reduction is thread-local only",
- )
- # VIEW: need layouts and layout analysis
- if not (src.layout and dst.layout):
- return False, "layouts required for view-based local reduction"
- if not (isinstance(src.layout, TileLayout) and isinstance(dst.layout, TileLayout)):
- return False, "TileLayout required for view-based local reduction"
- if src.layout.is_swizzle() or dst.layout.is_swizzle():
- return False, "swizzle layout unsupported for local reduction"
-
- analyzer = Analyzer()
-
- # Validate get_local_region succeeds for both
- src_st, src_extent = get_st_extent(src_br)
- dst_st, dst_extent = get_st_extent(dst_br)
-
- if sctx.is_warp:
- # Check for laneid shard->replica shuffle reduce pattern first.
- # This pattern has laneid in dst replica (broadcast), which the
- # general validation below would reject.
- shuffle_info = _analyze_shuffle_reduce(src.layout, dst.layout)
- if shuffle_info is not None:
- return True, None
-
- for layout, buf, st, ext, name in [
- (src.layout, src, src_st, src_extent, "src"),
- (dst.layout, dst, dst_st, dst_extent, "dst"),
- ]:
- for it in layout.shard:
- if it.axis.is_thread() and analyzer.can_prove_equal(it.stride, 0):
- return False, f"thread dim with zero stride in {name}"
- replica = getattr(layout, "replica", None) or []
- if any(it.axis.is_thread() for it in replica):
- return False, f"thread axis in replica for {name}"
- if get_local_region(layout, list(buf.shape), st, ext) is None:
- return False, f"get_local_region failed for {name}"
-
- # Validate layout compatibility
- # Spatial dims match, reduce dims in dst have local_extent==1
- reduce_axes = tuple(int(a) for a in op.reduce_axes)
- src_ndim = len(src_br.region)
- try:
- reduce_dims, _ = _analyze_axes(src_ndim, reduce_axes)
- except AssertionError as e:
- return False, str(e)
- src_sliced = get_sublayout_from_region(src.layout, src.shape, src_st, src_extent)
- dst_sliced = get_sublayout_from_region(dst.layout, dst.shape, dst_st, dst_extent)
- ok, msg = _validate_reduction_layout(
- src_sliced, dst_sliced, list(src_extent), list(dst_extent), reduce_dims
- )
- return ok, msg
- else:
- return False, f"unsupported exec_scope {sctx.scope_kind} for local reduction"
-
-
-def _emit_reduction_local_thread_wise(
- dst_br: TensorRegion,
- src_br: TensorRegion,
- accum: bool,
- reduce_op: ReduceOpType,
- reduce_dims: list[int],
- spatial_dims: list[int],
-) -> Function:
- dst, src = dst_br.source, src_br.source
- dtype = src.dtype
- src_st, src_extent = get_st_extent(src_br)
- dst_st, dst_extent = get_st_extent(dst_br)
- src_ndim = len(src_extent)
- spa_extents = [src_extent[d] for d in spatial_dims]
- red_extents = [src_extent[d] for d in reduce_dims]
- spatial_len = functools.reduce(operator.mul, spa_extents, 1)
- reduction_len = functools.reduce(operator.mul, red_extents, 1)
-
- op_func = reduce_op_table.get(reduce_op)
- assert op_func is not None
- init_value = reduce_default_value_table(dtype).get(reduce_op)
-
- def get_src_indices(spa_fused, red_fused):
- spa_indices = []
- rem = spa_fused
- for e in reversed(spa_extents):
- spa_indices.append(rem % e)
- rem //= e
- spa_indices.reverse()
-
- red_indices = []
- rem = red_fused
- for e in reversed(red_extents):
- red_indices.append(rem % e)
- rem //= e
- red_indices.reverse()
-
- full = [None] * src_ndim
- for i, d in enumerate(spatial_dims):
- full[d] = spa_indices[i] + src_st[d]
- for i, d in enumerate(reduce_dims):
- full[d] = red_indices[i] + src_st[d]
- return full
-
- # fmt: off
- @T.function(check_well_formed=False)
- def impl():
- for spa in T.serial(spatial_len):
- dst_idx = T.meta_var(get_indices(spa, dst_st, dst_extent))
- if T.constexpr(not accum):
- dst[tuple(dst_idx)] = init_value
- for red in T.serial(reduction_len):
- src_idx = T.meta_var(get_src_indices(spa, red))
- dst[tuple(dst_idx)] = op_func(dst[tuple(dst_idx)], src[tuple(src_idx)])
- # fmt: on
-
- return impl
-
-
-def _emit_reduction_local_view(
- dst_br: TensorRegion,
- src_br: TensorRegion,
- accum: bool,
- reduce_op: ReduceOpType,
- config: dict[str, Any],
- reduce_dims: set[int],
- spatial_dims: list[int],
- src_local_info,
- dst_local_info,
- shuffle_masks: list[int],
-) -> Function:
- dst, src = dst_br.source, src_br.source
- dtype = src.dtype
-
- op_func = reduce_op_table.get(reduce_op)
- assert op_func is not None
- init_value = reduce_default_value_table(dtype).get(reduce_op)
-
- src_local_shape, src_local_st, src_local_ext = src_local_info
- dst_local_shape, dst_local_st, dst_local_ext = dst_local_info
-
- # Build maps from original dim index to position in get_local_region output
- src_dim_map = _build_local_dim_map(src.layout, list(src.shape))
- dst_dim_map = _build_local_dim_map(dst.layout, list(dst.shape))
-
- # Only include reduction dims that have local parts in src
- src_ndim = len(src_br.region)
- reduce_local_dims = [d for d in reduce_dims if src_dim_map[d] is not None]
- reduction_local_ext = [src_local_ext[src_dim_map[d]] for d in reduce_local_dims]
- reduction_local_st = [src_local_st[src_dim_map[d]] for d in reduce_local_dims]
-
- reduction_local_total = functools.reduce(operator.mul, reduction_local_ext, 1)
- dst_local_total = functools.reduce(operator.mul, dst_local_ext, 1)
-
- def _get_src_local_index(dst_fused, red_fused):
- """Compute src local multi-dim index from dst fused index and reduction fused index."""
- dst_indices = get_indices(dst_fused, dst_local_st, dst_local_ext)
- red_indices = get_indices(red_fused, reduction_local_st, reduction_local_ext)
-
- # Interleave into src local indices (skipping pure-thread dims)
- src_local = []
- ri = 0
- for d in range(src_ndim):
- if src_dim_map[d] is None:
- continue # pure-thread in src, not in src.local()
- if d in reduce_dims:
- src_local.append(red_indices[ri])
- ri += 1
- else:
- # Spatial dim: use corresponding dst local position
- src_local.append(dst_indices[dst_dim_map[d]])
-
- return src_local
-
- # is_same_buffer = src.same_as(dst)
- shuffle = bool(config.get("thread_reduce", False))
- in_place = dst.same_as(src)
-
- def shuffle_data(mask, dst_local, dst_idx):
- @T.inline
- def inner_shuffle(v, shuffle_mask):
- dst_local[tuple(dst_idx)] = op_func(
- v, T.gpu_warp_shuffle_xor(mask, v, shuffle_mask, 32, 32)
- )
-
- for i in range(len(shuffle_masks)):
- inner_shuffle(dst_local[tuple(dst_idx)], shuffle_masks[i])
-
- need_save_accum = accum and shuffle
-
- # fmt: off
- if need_save_accum:
- @T.function(check_well_formed=False)
- def impl():
- # These dimensions are storage-iterator coordinates, so request
- # the mediated layout explicitly. Bare local() is physical-order.
- src_local = src.local(*src_local_shape, layout=src.layout.storage())
- dst_local = dst.local(*dst_local_shape, layout=dst.layout.storage())
- old_val = T.alloc_tensor([1], dtype, scope="local")
-
- for spa in T.serial(dst_local_total):
- dst_idx = T.meta_var(get_indices(spa, dst_local_st, dst_local_ext))
- old_val[0] = dst_local[tuple(dst_idx)]
- if T.constexpr(not in_place):
- dst_local[tuple(dst_idx)] = init_value
- for red in T.serial(reduction_local_total):
- src_idx = T.meta_var(_get_src_local_index(spa, red))
- dst_local[tuple(dst_idx)] = op_func(dst_local[tuple(dst_idx)], src_local[tuple(src_idx)]) # noqa: E501
- if T.constexpr(shuffle):
- mask = T.gpu_warp_activemask()
- shuffle_data(mask, dst_local, dst_idx)
- dst_local[tuple(dst_idx)] = op_func(dst_local[tuple(dst_idx)], old_val[0])
- else:
- @T.function(check_well_formed=False)
- def impl():
- # These dimensions are storage-iterator coordinates, so request
- # the mediated layout explicitly. Bare local() is physical-order.
- src_local = src.local(*src_local_shape, layout=src.layout.storage())
- dst_local = dst.local(*dst_local_shape, layout=dst.layout.storage())
-
- for spa in T.serial(dst_local_total):
- dst_idx = T.meta_var(get_indices(spa, dst_local_st, dst_local_ext))
- if T.constexpr(not in_place):
- if T.constexpr(not accum):
- dst_local[tuple(dst_idx)] = init_value
- for red in T.serial(reduction_local_total):
- src_idx = T.meta_var(_get_src_local_index(spa, red))
- dst_local[tuple(dst_idx)] = op_func(dst_local[tuple(dst_idx)], src_local[tuple(src_idx)]) # noqa: E501
- if T.constexpr(shuffle):
- mask = T.gpu_warp_activemask()
- shuffle_data(mask, dst_local, dst_idx)
- # fmt: on
-
- return impl
-
-
-def reduction_local_impl(
- op: TileOpCall, op_type: ReduceOpType, sctx: DispatchContext
-) -> Function | None:
- dst_br, src_br, reduce_axes, accum, config = _reduction_args(op)
- src_ndim = len(src_br.region)
- reduce_dims, spatial_dims = _analyze_axes(src_ndim, reduce_axes)
-
- if sctx.is_thread:
- return _emit_reduction_local_thread_wise(
- dst_br, src_br, accum, op_type, reduce_dims, spatial_dims
- )
- elif sctx.scope_kind in ["warp", "warpgroup"]:
- src = src_br.source
- dst = dst_br.source
-
- if sctx.is_warp:
- # --- Try laneid shard->replica shuffle reduce ---
- shuffle_info = _analyze_shuffle_reduce(src.layout, dst.layout)
- if shuffle_info is not None:
- reduce_width, local_elems = shuffle_info
- if op_type not in _REDUCE_OP_TO_STR:
- fail(f"unsupported reduce op: {op_type}")
- dtype = src.dtype
- init_value = reduce_default_value_table(dtype).get(op_type)
- return _gen_warp_shuffle_reduce(
- src, dst, reduce_width, local_elems, accum, op_type, init_value
- )
- elif config.get("thread_reduce", False):
- fail(
- "thread_reduce=True is only supported in warp scope; "
- "warpgroup local reduction is thread-local only"
- )
-
- # --- Existing WGMMA layout path below ---
- src_st, src_extent = get_st_extent(src_br)
- dst_st, dst_extent = get_st_extent(dst_br)
-
- src_local_info = get_local_region(src.layout, list(src.shape), src_st, src_extent)
- dst_local_info = get_local_region(dst.layout, list(dst.shape), dst_st, dst_extent)
- assert src_local_info is not None and dst_local_info is not None
-
- src_dim_info = _analyze_layout_dims(src.layout, list(src.shape))
- shuffle_masks = (
- _compute_shuffle_masks(src_dim_info, reduce_dims)
- if config.get("thread_reduce", False)
- else []
- )
-
- return _emit_reduction_local_view(
- dst_br,
- src_br,
- accum,
- op_type,
- config,
- reduce_dims,
- spatial_dims,
- src_local_info,
- dst_local_info,
- shuffle_masks,
- )
- else:
- fail(f"unsupported exec_scope {sctx.scope_kind} for reduction_local_impl")
-
-
-# ---------------------------------------------------------------------------
-# Registration: local memory reduction (priority=10)
-# ---------------------------------------------------------------------------
-
-for op_name, op_type in [
- ("sum", ReduceOpType.SUM),
- ("max", ReduceOpType.MAX),
- ("min", ReduceOpType.MIN),
-]:
-
- @register_dispatch(
- op_name,
- "cuda",
- variant="local",
- priority=10,
- when=[
- predicate("storage_scope", _match_reduction_storage_scope, expected_scope=["local"]),
- predicate("local_valid", validate_reduction_local),
- ],
- )
- def _local_dispatch(op: TileOpCall, sctx: DispatchContext, _op_type=op_type) -> Function:
- op = TileOpCall.downcast(op)
- return reduction_local_impl(op, _op_type, sctx)
diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/shared.py b/python/tvm/backend/cuda/tile_primitive/reduction/shared.py
deleted file mode 100644
index e61e551a19..0000000000
--- a/python/tvm/backend/cuda/tile_primitive/reduction/shared.py
+++ /dev/null
@@ -1,295 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""CUDA reduction operator dispatch: shared-memory variant.
-
-Registered ops: sum, max, min.
-
-When: dst and src are both shared-memory buffers, exec scope is one of
-{cta, warpgroup, warp, thread}, threadIdx.x bound, reduce axes valid.
-
-(A) CTA/warpgroup/warp scope -- adaptive-group shuffle tree
- (_emit_reduction_shared_cta):
- group_size = min(next_power_of_2(reduction_len), 32).
- Each group of threads reduces one spatial position via shfl_xor.
-
-Before:
- Tx.cta.sum(B_smem[0:4], A_smem[0:4, 0:8], [-1], False)
-
-After (scheduled Function, group_size=8, spatial_par=4):
- thread_data[0] = T.float32(0.0)
- thread_data[0] = thread_data[0] + A_smem[tid_in_scope] # gather
- # log2(8) = 3 shuffle-xor steps with width=8
- thread_data[0] = thread_data[0] + shfl_xor(thread_data[0], 1, 8, 32)
- thread_data[0] = thread_data[0] + shfl_xor(thread_data[0], 2, 8, 32)
- thread_data[0] = thread_data[0] + shfl_xor(thread_data[0], 4, 8, 32)
- if tid_in_scope % 8 == 0:
- B_smem[tid_in_scope // 8] = thread_data[0]
-
-(B) Thread scope -- sequential loop (_emit_reduction_shared_thread):
-
-Before:
- if tid == 65:
- Tx.sum(B_smem[0:4], A_smem[0:4, 0:8], [-1], False)
-
-After (scheduled Function):
- for spa in range(4):
- B_smem[spa] = T.float32(0.0) # init (skipped if accum)
- for red in range(8):
- B_smem[spa] = B_smem[spa] + A_smem[spa * 8 + red]
-"""
-
-import functools
-import math
-import operator
-
-from tvm.ir import TensorRegion
-from tvm.script import tirx as T
-from tvm.sym.analyzer import Analyzer
-from tvm.tirx import Function
-from tvm.tirx.operator.tile_primitive import DispatchContext, fail
-from tvm.tirx.operator.tile_primitive.common import ReduceOpType
-from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
-from tvm.tirx.stmt import TileOpCall
-
-from ..common import get_indices, get_st_extent, next_power_of_2
-from .utils import (
- _analyze_axes,
- _match_reduction_storage_scope,
- _reduction_args,
- build_src_indices,
- reduce_default_value_table,
- reduce_op_table,
-)
-
-
-def validate_reduction_shared(op: TileOpCall, sctx: DispatchContext) -> tuple[bool, str | None]:
- """Validate reduction in shared memory."""
- if sctx.scope_kind not in ["cta", "warpgroup", "warp", "thread"]:
- return False, f"unsupported exec_scope {sctx.scope_kind} for shared reduction"
-
- op = TileOpCall.downcast(op)
- dst, src = op.output.source, op.input.source
- if not (src.scope().startswith("shared") and dst.scope().startswith("shared")):
- return False, "expected shared scope for both src and dst"
- if src.dtype != dst.dtype:
- return False, f"dtype mismatch: src={src.dtype} dst={dst.dtype}"
-
- if "threadIdx.x" not in sctx.launch_params:
- return False, "threadIdx.x not in launch_params"
- if "threadIdx.y" in sctx.launch_params or "threadIdx.z" in sctx.launch_params:
- return False, "multi-dimensional thread binding not supported for shared reduction"
-
- reduce_axes = tuple(int(a) for a in op.reduce_axes)
- src_region = op.input.region
- dst_region = op.output.region
- src_ndim = len(src_region)
- try:
- reduce_dims, spatial_dims = _analyze_axes(src_ndim, reduce_axes)
- except AssertionError as e:
- return False, str(e)
-
- # Validate dst shape matches spatial dims of src
- src_extent = [r.extent for r in src_region]
- dst_extent = [r.extent for r in dst_region]
- expected_dst_len = functools.reduce(operator.mul, [src_extent[d] for d in spatial_dims], 1)
- actual_dst_len = functools.reduce(operator.mul, dst_extent, 1)
- analyzer = Analyzer()
- if not analyzer.can_prove_equal(expected_dst_len, actual_dst_len):
- return (False, f"dst size {actual_dst_len} != expected spatial size {expected_dst_len}")
-
- return True, None
-
-
-def _emit_reduction_shared_cta(
- dst_br: TensorRegion,
- src_br: TensorRegion,
- accum: bool,
- reduce_op: ReduceOpType,
- sctx: DispatchContext,
- reduce_dims: list[int],
- spatial_dims: list[int],
-) -> Function:
- exec_scope_name = sctx.scope_kind
-
- def get_thread_cnt():
- if exec_scope_name == "cta":
- return sctx.launch_params["threadIdx.x"][1]
- elif exec_scope_name == "warpgroup":
- return 128
- elif exec_scope_name == "warp":
- return 32
- elif exec_scope_name == "thread":
- return 1
-
- thread_cnt = get_thread_cnt()
- dst, src = dst_br.source, src_br.source
- src_st, src_extent = get_st_extent(src_br)
- dst_st, dst_extent = get_st_extent(dst_br)
- dtype = src.dtype
-
- # Compute spatial/reduction from the explicit axes
- spatial_len = functools.reduce(operator.mul, [src_extent[d] for d in spatial_dims], 1)
- reduction_len = functools.reduce(operator.mul, [src_extent[d] for d in reduce_dims], 1)
-
- op_func = reduce_op_table.get(reduce_op)
- assert op_func is not None
- init_value = reduce_default_value_table(dtype).get(reduce_op)
-
- # Adaptive group size: nearest power-of-2 for reduction length, capped at warp size and thread count. # noqa: E501
- group_size = min(next_power_of_2(int(reduction_len)), 32, int(thread_cnt))
- group_size = max(group_size, 1) # ensure at least 1
- n_shuffles = int(math.log2(group_size)) if group_size > 1 else 0
- spatial_par = int(thread_cnt) // group_size
-
- def get_tid_in_scope():
- tx_var = sctx.launch_params["threadIdx.x"][0]
- if exec_scope_name == "cta":
- return tx_var
- elif exec_scope_name in ("warp", "warpgroup"):
- return tx_var % thread_cnt
- elif exec_scope_name == "thread":
- return 0
-
- def shuffle_data(thread_data):
- @T.inline
- def inner_shuffle(mask, v, shuffle_mask):
- v[0] = op_func(v[0], T.gpu_warp_shuffle_xor(mask, v[0], shuffle_mask, group_size, 32))
-
- if n_shuffles > 0:
- mask = T.gpu_warp_activemask()
- for i in range(n_shuffles):
- inner_shuffle(mask, thread_data, 1 << i)
-
- @T.inline
- def sync():
- if T.constexpr(exec_scope_name == "cta"):
- T.cuda.cta_sync()
- elif T.constexpr(exec_scope_name == "warpgroup"):
- T.cuda.warpgroup_sync(8) # TODO: fix this hardcoded value
- elif T.constexpr(exec_scope_name == "warp"):
- T.cuda.warp_sync()
- elif T.constexpr(exec_scope_name == "thread"):
- pass
-
- # fmt: off
- @T.function(check_well_formed=False)
- def impl():
- tid_in_scope = get_tid_in_scope()
- thread_data = T.alloc_tensor([1], dtype=dtype, scope="local")
- group_id = T.meta_var(T.floordiv(tid_in_scope, group_size))
- lane_in_grp = T.meta_var(tid_in_scope % group_size)
- for step in T.serial(T.ceildiv(spatial_len, spatial_par)):
- spa_fused = T.meta_var(step * spatial_par + group_id)
- if spa_fused < spatial_len:
- thread_data[0] = init_value
- for t in T.serial(T.ceildiv(reduction_len, group_size)):
- red_fused = T.meta_var(t * group_size + lane_in_grp)
- if red_fused < reduction_len:
- src_indices = T.meta_var(build_src_indices(spa_fused, red_fused, spatial_dims, reduce_dims, src_extent, src_st)) # noqa: E501
- thread_data[0] = op_func(thread_data[0], src[tuple(src_indices)])
- shuffle_data(thread_data)
- if lane_in_grp == 0:
- dst_indices = T.meta_var(get_indices(spa_fused, dst_st, dst_extent))
- dst[tuple(dst_indices)] = T.if_then_else(T.bool(accum), op_func(dst[tuple(dst_indices)], thread_data[0]), thread_data[0]) # noqa: E501
-
- sync()
- # fmt: on
-
- return impl
-
-
-def _emit_reduction_shared_thread(
- dst_br: TensorRegion,
- src_br: TensorRegion,
- accum: bool,
- reduce_op: ReduceOpType,
- sctx: DispatchContext,
- reduce_dims: list[int],
- spatial_dims: list[int],
-) -> Function:
- dst, src = dst_br.source, src_br.source
- src_st, src_extent = get_st_extent(src_br)
- dst_st, dst_extent = get_st_extent(dst_br)
- dtype = src.dtype
-
- # Compute spatial/reduction from the explicit axes
- spatial_len = functools.reduce(operator.mul, [src_extent[d] for d in spatial_dims], 1)
- reduction_len = functools.reduce(operator.mul, [src_extent[d] for d in reduce_dims], 1)
-
- op_func = reduce_op_table.get(reduce_op)
- assert op_func is not None
- init_value = reduce_default_value_table(dtype).get(reduce_op)
-
- @T.function(check_well_formed=False)
- def impl():
- for spa_fused in T.serial(spatial_len):
- dst_indices = T.meta_var(get_indices(spa_fused, dst_st, dst_extent))
- if T.constexpr(not accum):
- dst[tuple(dst_indices)] = init_value
- for red_fused in T.serial(reduction_len):
- src_indices = T.meta_var(
- build_src_indices(
- spa_fused, red_fused, spatial_dims, reduce_dims, src_extent, src_st
- )
- )
- dst[tuple(dst_indices)] = op_func(dst[tuple(dst_indices)], src[tuple(src_indices)])
-
- return impl
-
-
-def reduction_shared_impl(
- op: TileOpCall, op_type: ReduceOpType, sctx: DispatchContext
-) -> Function | None:
- dst_br, src_br, reduce_axes, accum, config = _reduction_args(op)
- src_ndim = len(src_br.region)
- reduce_dims, spatial_dims = _analyze_axes(src_ndim, reduce_axes)
- if sctx.scope_kind in ["cta", "warpgroup", "warp"]:
- return _emit_reduction_shared_cta(
- dst_br, src_br, accum, op_type, sctx, reduce_dims, spatial_dims
- )
- elif sctx.is_thread:
- return _emit_reduction_shared_thread(
- dst_br, src_br, accum, op_type, sctx, reduce_dims, spatial_dims
- )
- else:
- fail(f"unsupported exec_scope {sctx.scope_kind} for reduction_shared_impl")
-
-
-# ---------------------------------------------------------------------------
-# Registration: shared memory reduction (priority=10)
-# ---------------------------------------------------------------------------
-
-for op_name, op_type in [
- ("sum", ReduceOpType.SUM),
- ("max", ReduceOpType.MAX),
- ("min", ReduceOpType.MIN),
-]:
-
- @register_dispatch(
- op_name,
- "cuda",
- variant="shared",
- priority=10,
- when=[
- predicate("storage_scope", _match_reduction_storage_scope, expected_scope=["shared*"]),
- predicate("shared_valid", validate_reduction_shared),
- ],
- )
- def _shared_dispatch(op: TileOpCall, sctx: DispatchContext, _op_type=op_type) -> Function:
- op = TileOpCall.downcast(op)
- return reduction_shared_impl(op, _op_type, sctx)
diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/sm100_packed.py b/python/tvm/backend/cuda/tile_primitive/reduction/sm100_packed.py
deleted file mode 100644
index 7740b3bccf..0000000000
--- a/python/tvm/backend/cuda/tile_primitive/reduction/sm100_packed.py
+++ /dev/null
@@ -1,255 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""CUDA reduction operator dispatch: SM100+ packed optimized variant.
-
-Registered ops: sum, max, min.
-
-When: thread scope, all local buffers, float32, 1D src with len >= 8,
-SM100+ (uses packed PTX instructions not available on older GPUs).
-
-Before (TileOpCall -- sum example):
- Tx.sum(dst_local[0:1], src_local[0:32]) # float32, reduce 32 -> 1 (thread scope)
-
-After -- packed_add_sum (uses add.f32x2 to reduce pairs):
- # Iteratively reduce: 32 -> 16 -> 8 -> 4 -> 2 -> 1
- # Each step: add.f32x2 combines adjacent pairs
- for i in T.serial(16):
- T.cuda.func_call("add_f32x2", &buf[i*2], &buf[i*2], &buf[i*2+2])
- # ... repeat halving until scalar result
- dst_local[0] = buf[0]
-
-After -- 3input_maxmin (uses 3-input PTX max/min):
- # Tree reduction with 3-input instructions:
- # max(a, b, c) in one PTX instruction
- for i in T.serial(n // 3):
- T.cuda.func_call("max3_f32", &buf[i*3], &buf[i*3+1], &buf[i*3+2])
-
-With accum=True: accumulator folded into first element/pair of the reduction.
-"""
-
-import functools
-import operator
-
-from tvm.ir import TensorRegion
-from tvm.script import tirx as T
-from tvm.tirx import Function
-from tvm.tirx.operator.tile_primitive import DispatchContext
-from tvm.tirx.operator.tile_primitive.common import ReduceOpType
-from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
-from tvm.tirx.stmt import TileOpCall
-
-from ..common import sm_version_ok
-from ..exec_scope_utils import exec_scope_ok
-from .utils import (
- _dst_len_ok,
- _dtype_ok,
- _local_scope_match,
- _reduction_len_ok,
- _src_ndim_ok,
- reduce_op_table,
-)
-
-
-def _emit_reduction_local_thread_packed_add_sum(
- dst_buffer_region: TensorRegion,
- src_buffer_region: TensorRegion,
- accum: bool,
- reduce_op: ReduceOpType,
- sctx: DispatchContext,
-) -> Function:
- dst, src = dst_buffer_region.source, src_buffer_region.source
- src_region, dst_region = src_buffer_region.region, dst_buffer_region.region
- dtype = src.dtype
-
- src_extent = [r.extent for r in src_region]
- [r.extent for r in dst_region]
- src_st = [r.min for r in src_region]
- dst_st = [r.min for r in dst_region]
-
- reduction_len = functools.reduce(operator.mul, src_extent, 1)
-
- src_base = src_st[0]
- num_full_chunks = reduction_len // 8
- remainder = reduction_len % 8
- remainder_base = num_full_chunks * 8
-
- # fmt: off
- @T.function(check_well_formed=False)
- def impl():
- local_sum = T.alloc_tensor([8], dtype, scope="local")
- # add.f32x2's operands are .b64 register pairs, so each packed add is
- # mov (pack) -> add -> mov (unpack). nvcc emitted exactly these movs for
- # the old make_float2/float2_x glue too; writing them keeps every
- # instruction visible and ptxas folds them away either way.
- acc = T.local_scalar("uint64")
- rhs = T.local_scalar("uint64")
- # First pass: copy first 8 elements (with optional accumulator)
- for i in T.unroll(8):
- if accum and i == 0:
- local_sum[i] = src[src_base + i] + dst[tuple(dst_st)]
- else:
- local_sum[i] = src[src_base + i]
-
- # Process remaining full chunks of 8
- for outer in T.serial(num_full_chunks - 1):
- for j in T.unroll(4):
- T.ptx.mov.b64(acc, local_sum[2 * j], local_sum[2 * j + 1])
- T.ptx.mov.b64(
- rhs,
- src[src_base + 8 * (outer + 1) + 2 * j],
- src[src_base + 8 * (outer + 1) + 2 * j + 1],
- )
- T.ptx.add.rn.ftz.f32x2(acc, acc, rhs)
- T.ptx.mov.b64(local_sum[2 * j], local_sum[2 * j + 1], acc)
-
- # Handle remainder elements (0 to 7)
- for i in T.serial(remainder):
- local_sum[0] = local_sum[0] + src[src_base + remainder_base + i]
-
- # Final packed add sum: 8 -> 4 -> 2 -> 1
- T.ptx.mov.b64(acc, local_sum[0], local_sum[1])
- T.ptx.mov.b64(rhs, local_sum[2], local_sum[3])
- T.ptx.add.rn.ftz.f32x2(acc, acc, rhs)
- T.ptx.mov.b64(local_sum[0], local_sum[1], acc)
-
- T.ptx.mov.b64(acc, local_sum[4], local_sum[5])
- T.ptx.mov.b64(rhs, local_sum[6], local_sum[7])
- T.ptx.add.rn.ftz.f32x2(acc, acc, rhs)
- T.ptx.mov.b64(local_sum[4], local_sum[5], acc)
-
- T.ptx.mov.b64(acc, local_sum[0], local_sum[1])
- T.ptx.mov.b64(rhs, local_sum[4], local_sum[5])
- T.ptx.add.rn.ftz.f32x2(acc, acc, rhs)
- T.ptx.mov.b64(local_sum[0], local_sum[1], acc)
-
- dst[tuple(dst_st)] = local_sum[0] + local_sum[1]
- # fmt: on
-
- return impl
-
-
-def _emit_reduction_local_thread_3input_maxmin(
- dst_buffer_region: TensorRegion,
- src_buffer_region: TensorRegion,
- accum: bool,
- reduce_op: ReduceOpType,
- sctx: DispatchContext,
-) -> Function:
- dst, src = dst_buffer_region.source, src_buffer_region.source
- src_region, dst_region = src_buffer_region.region, dst_buffer_region.region
- dtype = src.dtype
-
- src_extent = [r.extent for r in src_region]
- src_st = [r.min for r in src_region]
- dst_st = [r.min for r in dst_region]
-
- reduction_len = functools.reduce(operator.mul, src_extent, 1)
-
- op_func = reduce_op_table[reduce_op]
- # The three-source `max{.ftz}{.NaN}{.abs}.f32 d, a, b, c` line (ISA
- # 9.7.3.12); the shared `max` mnemonic picks it by operand count.
- reduce3_name = "max" if reduce_op == ReduceOpType.MAX else "min"
-
- src_base = src_st[0]
- num_full_chunks = reduction_len // 8
- remainder = reduction_len % 8
- remainder_base = num_full_chunks * 8
-
- # fmt: off
- @T.function(check_well_formed=False)
- def impl():
- temp = T.alloc_tensor([4], dtype, scope="local")
- # First pass: process first 8 elements into 4 temps
- for i in T.unroll(4):
- if accum and i == 0:
- T.ptx[f"{reduce3_name}.f32"](temp[i], src[src_base + 2 * i], src[src_base + 2 * i + 1], dst[tuple(dst_st)]) # noqa: E501
- else:
- temp[i] = op_func(src[src_base + 2 * i], src[src_base + 2 * i + 1])
-
- # Process remaining full chunks of 8
- for outer in T.serial(num_full_chunks - 1):
- for i in T.unroll(4):
- T.ptx[f"{reduce3_name}.f32"](
- temp[i],
- temp[i],
- src[src_base + 8 * (outer + 1) + 2 * i],
- src[src_base + 8 * (outer + 1) + 2 * i + 1],
- )
-
- # Process remainder elements (0 to 7 elements)
- for i in T.serial(remainder):
- temp[0] = op_func(temp[0], src[src_base + remainder_base + i])
-
- # Final merge: combine 4 temps into result
- dst[tuple(dst_st)] = op_func(temp[0], temp[1])
- T.ptx[f"{reduce3_name}.f32"](dst[tuple(dst_st)], dst[tuple(dst_st)], temp[2], temp[3])
- # fmt: on
-
- return impl
-
-
-def _sm100_packed_add_sum_impl(op: TileOpCall, op_type: ReduceOpType, sctx: DispatchContext):
- op = TileOpCall.downcast(op)
- return _emit_reduction_local_thread_packed_add_sum(op.output, op.input, op.accum, op_type, sctx)
-
-
-def _sm100_3input_maxmin_impl(op: TileOpCall, op_type: ReduceOpType, sctx: DispatchContext):
- op = TileOpCall.downcast(op)
- return _emit_reduction_local_thread_3input_maxmin(op.output, op.input, op.accum, op_type, sctx)
-
-
-_optimized_local_reduction_predicates = [
- predicate("exec_scope", exec_scope_ok, expected_scopes=["thread"]),
- predicate("local_scope", _local_scope_match),
- predicate("dst_len", _dst_len_ok, expected_len=1),
- predicate("src_ndim", _src_ndim_ok, expected_ndim=1),
- predicate("dtype", _dtype_ok, expected_dtype="float32"),
- predicate("sm_version", sm_version_ok, min_version=100),
- predicate("reduction_len", _reduction_len_ok, min_len=8),
-]
-
-_optimized_impl_table = {
- ReduceOpType.SUM: ("packed_add_sum", _sm100_packed_add_sum_impl),
- ReduceOpType.MAX: ("3input_maxmin", _sm100_3input_maxmin_impl),
- ReduceOpType.MIN: ("3input_maxmin", _sm100_3input_maxmin_impl),
-}
-
-
-# ---------------------------------------------------------------------------
-# Registration: SM100+ optimized local reduction (priority=20)
-# ---------------------------------------------------------------------------
-
-for op_name, op_type in [
- ("sum", ReduceOpType.SUM),
- ("max", ReduceOpType.MAX),
- ("min", ReduceOpType.MIN),
-]:
- variant_name, optimized_impl = _optimized_impl_table[op_type]
-
- @register_dispatch(
- op_name,
- "cuda",
- variant=variant_name,
- priority=20,
- when=_optimized_local_reduction_predicates,
- )
- def _optimized_dispatch(
- op: TileOpCall, sctx: DispatchContext, _impl=optimized_impl, _op_type=op_type
- ) -> Function:
- op = TileOpCall.downcast(op)
- return _impl(op, _op_type, sctx)
diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/utils.py b/python/tvm/backend/cuda/tile_primitive/reduction/utils.py
deleted file mode 100644
index 274c46ddd3..0000000000
--- a/python/tvm/backend/cuda/tile_primitive/reduction/utils.py
+++ /dev/null
@@ -1,262 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""Shared helpers for reduction operator dispatches on CUDA targets."""
-
-import functools
-import math
-import operator
-
-from tvm.ir import TensorRegion
-from tvm.script import tirx as T
-from tvm.sym.analyzer import Analyzer
-from tvm.tirx.operator.tile_primitive import DispatchContext
-from tvm.tirx.operator.tile_primitive.common import ReduceOpType
-from tvm.tirx.stmt import TileOpCall
-
-from ..common import match_scope
-
-reduce_op_table = {
- ReduceOpType.SUM: lambda a, b: a + b,
- ReduceOpType.MAX: T.max,
- ReduceOpType.MIN: T.min,
-}
-
-
-def reduce_default_value_table(dtype):
- return {
- ReduceOpType.SUM: 0.0,
- ReduceOpType.MAX: T.min_value(dtype),
- ReduceOpType.MIN: T.max_value(dtype),
- }
-
-
-def _reduction_args(
- op: TileOpCall,
-) -> tuple[TensorRegion, TensorRegion, tuple[int, ...], bool, dict]:
- """Parse ReduceOp -> (dst, src, reduce_axes, accum, config)."""
- op = TileOpCall.downcast(op)
- dst = op.output
- src = op.input
- reduce_axes = tuple(int(a) for a in op.reduce_axes)
- accum = op.accum
- config = op.config
- return dst, src, reduce_axes, accum, config
-
-
-def _match_reduction_storage_scope(
- op: TileOpCall, sctx: DispatchContext, expected_scope: list[str]
-) -> tuple[bool, str | None]:
- """Check that dst and src scopes match one of the expected patterns."""
- op = TileOpCall.downcast(op)
- dst_scope = op.output.source.scope()
- src_scope = op.input.source.scope()
-
- ok = any(match_scope(dst_scope, p) and match_scope(src_scope, p) for p in expected_scope)
- msg = f"storage scope mismatch: dst {dst_scope}, src {src_scope}; expected {expected_scope}"
- return (ok, None if ok else msg)
-
-
-def _analyze_axes(src_ndim: int, reduce_axes: tuple[int, ...]) -> tuple[list[int], list[int]]:
- """Normalize negative axes -> (reduce_dim_set, spatial_dim_list)."""
- reduce_dims = set()
- for ax in reduce_axes:
- a = ax if ax >= 0 else ax + src_ndim
- assert 0 <= a < src_ndim, f"reduce axis {ax} out of range for ndim={src_ndim}"
- reduce_dims.add(a)
- spatial_dims = [d for d in range(src_ndim) if d not in reduce_dims]
- return sorted(reduce_dims), spatial_dims
-
-
-def _analyze_layout_dims(layout, shape):
- """layout.group(shape) -> decompose each dim into thread/local iters.
-
- Returns list of per-dim (thread_extent, local_extent, thread_strides):
- thread_extent = product of thread iter extents in this dim
- local_extent = product of local iter extents in this dim
- thread_strides = list of (stride, extent) for thread iters in this dim
- """
- grouped, seps = layout.group(list(shape))
- result = []
- for d in range(len(shape)):
- shard_range = list(range(seps[d], seps[d + 1]))
- thread_extent = 1
- local_extent = 1
- thread_strides = []
- for s_idx in shard_range:
- it = grouped.shard[s_idx]
- if it.axis.is_thread():
- thread_extent *= it.extent
- thread_strides.append((it.stride, it.extent))
- else:
- local_extent *= it.extent
- result.append((thread_extent, local_extent, thread_strides))
- return result
-
-
-def _compute_shuffle_masks(dim_info, reduce_dims: set[int]) -> list[int]:
- """From reduction dims' thread iter (stride, extent) pairs, compute XOR masks.
-
- For each thread iter in a reduction dim:
- masks += [stride * 2^i for i in range(log2(extent))]
- Sorted ascending.
- """
- masks = []
- for d in reduce_dims:
- _, _, thread_strides = dim_info[d]
- for stride, extent in thread_strides:
- ext_int = int(extent) if hasattr(extent, "__int__") else extent
- n_bits = int(math.log2(ext_int))
- for i in range(n_bits):
- stride_int = int(stride) if hasattr(stride, "__int__") else stride
- masks.append(stride_int * (1 << i))
- masks.sort()
- return masks
-
-
-def _build_local_dim_map(layout, buffer_shape):
- """Map original dim index to position in get_local_region output (None if pure-thread)."""
- grouped, seps = layout.group(list(buffer_shape))
- dim_map = {}
- local_pos = 0
- for d in range(len(buffer_shape)):
- shard_range = list(range(seps[d], seps[d + 1]))
- has_local = any(not grouped.shard[s].axis.is_thread() for s in shard_range)
- if has_local:
- dim_map[d] = local_pos
- local_pos += 1
- else:
- dim_map[d] = None
- return dim_map
-
-
-def _validate_reduction_layout(
- src_layout, dst_layout, src_shape, dst_shape, reduce_dims: list[int]
-) -> tuple[bool, str | None]:
- """Validate that spatial dims of src/dst have matching thread+local structure,
- and that reduction dims in dst have local_extent == 1.
- """
- src_dim_info = _analyze_layout_dims(src_layout, src_shape)
- dst_dim_info = _analyze_layout_dims(dst_layout, dst_shape)
- analyzer = Analyzer()
-
- # Spatial dims: src/dst must match in both thread and local extents.
- # Reduce dims: src/dst thread extent must match, and dst local extent must be 1.
-
- # get expected simplified dst layout
- expected_dst_dim = []
- for src_idx in range(len(src_shape)):
- if analyzer.can_prove_equal(src_dim_info[src_idx][0], 1) and analyzer.can_prove_equal(
- src_dim_info[src_idx][1], 1
- ):
- continue # skip if extent=1
- if src_idx in reduce_dims: # reduce dims
- if not analyzer.can_prove_equal(src_dim_info[src_idx][0], 1):
- expected_dst_dim.append((src_dim_info[src_idx][0], 1))
- else: # spatial dims
- expected_dst_dim.append((src_dim_info[src_idx][0], src_dim_info[src_idx][1]))
-
- # check dst layout
- check_idx = 0
- for dst_idx in range(len(dst_shape)):
- if analyzer.can_prove_equal(dst_dim_info[dst_idx][0], 1) and analyzer.can_prove_equal(
- dst_dim_info[dst_idx][1], 1
- ):
- continue
- if not (
- analyzer.can_prove_equal(dst_dim_info[dst_idx][0], expected_dst_dim[check_idx][0])
- and analyzer.can_prove_equal(dst_dim_info[dst_idx][1], expected_dst_dim[check_idx][1])
- ):
- return False, "mismatch dst/src layout for reduction"
- check_idx += 1
- if check_idx != len(expected_dst_dim):
- return False, "mismatch dst/src layout for reduction"
- return True, None
-
-
-def build_src_indices(spa_fused, red_fused, spatial_dims, reduce_dims, src_extent, src_st):
- """Combine spatial and reduction indices into full src index tuple."""
-
- # Build index helpers that work with the explicit axis split
- def get_spatial_or_reduction_src_indices(spa_or_red_fused, is_spatial):
- dims = spatial_dims if is_spatial else reduce_dims
- spa_extents = [src_extent[d] for d in dims]
- indices = []
- rem = spa_or_red_fused
- for e in reversed(spa_extents):
- indices.append(rem % e)
- rem //= e
- indices.reverse()
- return [idx + src_st[d] for idx, d in zip(indices, dims)]
-
- spa_vals = get_spatial_or_reduction_src_indices(spa_fused, is_spatial=True)
- red_vals = get_spatial_or_reduction_src_indices(red_fused, is_spatial=False)
- full = [None] * len(src_extent)
- for i, d in enumerate(spatial_dims):
- full[d] = spa_vals[i]
- for i, d in enumerate(reduce_dims):
- full[d] = red_vals[i]
- return full
-
-
-_REDUCE_OP_TO_STR = {ReduceOpType.SUM: "sum", ReduceOpType.MAX: "max", ReduceOpType.MIN: "min"}
-
-
-def _dtype_ok(op: TileOpCall, sctx: DispatchContext, expected_dtype: str):
- op = TileOpCall.downcast(op)
- dtype = op.input.source.dtype
- ok = dtype == expected_dtype
- return (ok, None if ok else f"dtype {dtype} != {expected_dtype}")
-
-
-def _reduction_len_ok(op: TileOpCall, sctx: DispatchContext, min_len: int):
- op = TileOpCall.downcast(op)
- src_extent = [r.extent for r in op.input.region]
- reduction_len = functools.reduce(operator.mul, src_extent, 1)
- ok = reduction_len >= min_len
- return (ok, None if ok else f"reduction_len {reduction_len} < {min_len}")
-
-
-def _dst_len_ok(op: TileOpCall, sctx: DispatchContext, expected_len: int):
- op = TileOpCall.downcast(op)
- dst_extent = [r.extent for r in op.output.region]
- dst_len = functools.reduce(operator.mul, dst_extent, 1)
- ok = dst_len == expected_len
- return (ok, None if ok else f"dst_len {dst_len} != {expected_len}")
-
-
-def _src_ndim_ok(op: TileOpCall, sctx: DispatchContext, expected_ndim: int):
- op = TileOpCall.downcast(op)
- src_extent = [r.extent for r in op.input.region]
- ok = len(src_extent) == expected_ndim
- return (ok, None if ok else f"src ndim {len(src_extent)} != {expected_ndim}")
-
-
-def _local_scope_match(op: TileOpCall, sctx: DispatchContext):
- op = TileOpCall.downcast(op)
- src, dst = op.input.source, op.output.source
- ok = all(
- [
- src.scope() == "local",
- dst.scope() == "local",
- src.dtype == dst.dtype,
- sctx.is_target("cuda"),
- ]
- )
- if not ok:
- return (False, "src/dst must be local scope with matching dtype on CUDA")
- return (True, None)
diff --git a/python/tvm/backend/trn/__init__.py b/python/tvm/backend/trn/__init__.py
index e9f1a80bc5..f568a0ce8e 100644
--- a/python/tvm/backend/trn/__init__.py
+++ b/python/tvm/backend/trn/__init__.py
@@ -55,7 +55,12 @@ def script_namespace():
def script_namespaces(**_):
"""Return Trainium-owned TVMScript namespaces."""
- return {"nki": script_namespace()}
+ from types import SimpleNamespace
+
+ from .tensor_instructions import make_namespace
+ from .tile_primitive import private_alloc # register workspace policies
+
+ return {"nki": script_namespace(), "trn": SimpleNamespace(tile=make_namespace())}
def __getattr__(name: str):
diff --git a/python/tvm/backend/trn/tensor_instructions.py b/python/tvm/backend/trn/tensor_instructions.py
new file mode 100644
index 0000000000..cfc66f5467
--- /dev/null
+++ b/python/tvm/backend/trn/tensor_instructions.py
@@ -0,0 +1,307 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Trainium tensor instructions with explicit NKI instruction identity."""
+
+from tvm.ir import LambdaExpr, PrimType, Tuple, const
+from tvm.tirx.tensor_instruction import Instruction, Operand, namespace
+
+
+def _copy(view, sctx):
+ from .tile_primitive.copy.default import copy_trn
+
+ dst, src = view.args
+ name = view.op.name.rsplit(".", 1)[-1]
+ expected = (
+ "load"
+ if src.source.scope() == "global"
+ else "store"
+ if dst.source.scope() == "global"
+ else "tensor_copy"
+ )
+ if name != expected:
+ raise ValueError(f"{name} does not accept this memory direction; use {expected}")
+ return copy_trn(view, sctx)
+
+
+def _matmul(view, sctx):
+ from .tile_primitive.gemm.default import matmul_trn
+
+ return matmul_trn(view, sctx)
+
+
+def _unary(view, sctx):
+ from tvm.tirx.operator.tile_primitive.common import MapOpType
+
+ from .tile_primitive.unary.default import unary_trn
+ from .tile_primitive.unary.with_bias_scale import unary_with_bias_scale_trn
+
+ if view.kind.endswith("_with_scale_bias"):
+ opcode = view.call.attrs.opcode
+ return unary_with_bias_scale_trn(view, MapOpType[opcode.upper()], sctx)
+ return unary_trn(view, MapOpType.FILL if view.kind == "memset" else MapOpType.RECIPROCAL, sctx)
+
+
+def _binary(view, sctx):
+ from tvm.tirx.operator.tile_primitive.common import MapOpType
+
+ from .tile_primitive.binary.default import binary_trn
+
+ return binary_trn(view, MapOpType[view.call.attrs.opcode.upper()], sctx)
+
+
+def _reduce(view, sctx):
+ from tvm.tirx.operator.tile_primitive.common import ReduceOpType
+
+ from .tile_primitive.reduction.utils import reduction_trn
+
+ return reduction_trn(
+ view, ReduceOpType[view.call.attrs.reduce_op.upper()], sctx, negate=view.call.attrs.negate
+ )
+
+
+def _chain(view, sctx):
+ from .tile_primitive.compose_op.binary_chain import binary_chain_trn
+
+ return binary_chain_trn(view, sctx)
+
+
+def _binary_reduce(view, sctx):
+ from .tile_primitive.compose_op.binary_reduce import binary_reduce_trn
+
+ return binary_reduce_trn(view, sctx)
+
+
+def _activation_reduce(view, sctx):
+ from .tile_primitive.compose_op.unary_reduce import unary_reduce_trn
+
+ return unary_reduce_trn(view, sctx)
+
+
+def _select(view, sctx):
+ from .tile_primitive.select.default import select_trn
+
+ return select_trn(view, sctx)
+
+
+def _axes(attrs):
+ return Tuple([const(i) for i in attrs.axes])
+
+
+def make_namespace():
+ def spec(name, kind, operands, options=None, workspaces=(), project=None, lower=None):
+ return Instruction(
+ "tirx.trn.tile." + name,
+ kind,
+ tuple(operands) + tuple(Operand(w, "tensor", None) for w in workspaces),
+ "TrnAttrs",
+ options or {},
+ workspaces,
+ project,
+ lower,
+ )
+
+ pair = [Operand("dst", "region"), Operand("src", "region")]
+ result = [
+ spec(
+ name, "copy", pair, {"max_inst_size": 512} if name == "tensor_copy" else {}, lower=_copy
+ )
+ for name in ("load", "store", "tensor_copy")
+ ]
+ result += [
+ spec(
+ "matmul",
+ "gemm",
+ [
+ Operand("D", "region"),
+ Operand("A", "region"),
+ Operand("B", "region"),
+ Operand("C", "region", "@dst"),
+ Operand("transpose_A", "scalar", False),
+ Operand("transpose_B", "scalar", False),
+ Operand("alpha", "scalar", 1.0),
+ Operand("beta", "scalar", 0.0),
+ ],
+ workspaces=("acc_psum",),
+ lower=_matmul,
+ )
+ ]
+ for name in ("reciprocal", "memset"):
+ result += [
+ spec(
+ name,
+ name,
+ [Operand("dst", "region"), Operand("src")],
+ {"max_inst_size": 512},
+ lower=_unary,
+ )
+ ]
+
+ def activation_project(v, a):
+ return a.opcode + "_with_scale_bias", [v["dst"], v["src"], v["scale"], v["bias"]], {}
+
+ result += [
+ spec(
+ "activation",
+ "exp_with_scale_bias",
+ [*pair, Operand("scale", "scalar", 1.0), Operand("bias", "expr", 0.0)],
+ {"opcode": "exp", "max_inst_size": 512},
+ ("const_bias",),
+ activation_project,
+ _unary,
+ )
+ ]
+
+ def binary_project(v, a):
+ return (
+ {"max": "maximum", "min": "minimum"}.get(a.opcode, a.opcode),
+ [v["dst"], v["lhs"], v["rhs"]],
+ {},
+ )
+
+ for name in ("tensortensor", "tensorscalar"):
+ result += [
+ spec(
+ name,
+ "add",
+ [Operand("dst", "region"), Operand("lhs"), Operand("rhs")],
+ {"opcode": "add", "max_inst_size": 512},
+ project=binary_project,
+ lower=_binary,
+ )
+ ]
+
+ def reduce_project(v, a):
+ return a.reduce_op, [v["dst"], v["src"], _axes(a), const(False)], {}
+
+ result += [
+ spec(
+ "tensorreduce",
+ "sum",
+ pair,
+ {"reduce_op": "sum", "axes": [-1], "negate": False, "max_inst_size": None},
+ ("partial_reduce",),
+ reduce_project,
+ _reduce,
+ )
+ ]
+
+ def chain_project(v, a):
+ return (
+ "binary_chain",
+ [v["dst"], v["data"], v["operand0"], v["operand1"], a.op0, a.op1, a.reverse1],
+ {},
+ )
+
+ for name in ("scalar_tensor_scalar", "scalar_tensor_tensor"):
+ result += [
+ spec(
+ name,
+ "binary_chain",
+ [
+ Operand("dst", "region"),
+ Operand("data"),
+ Operand("operand0"),
+ Operand("operand1"),
+ ],
+ {"op0": "mul", "op1": "add", "reverse1": False, "max_inst_size": 512},
+ project=chain_project,
+ lower=_chain,
+ )
+ ]
+
+ def br_project(v, a):
+ return (
+ "binary_reduce",
+ [v["dst"], v["reduced"], v["lhs"], v["rhs"], a.opcode, a.reduce_op, _axes(a)],
+ {},
+ )
+
+ result += [
+ spec(
+ "tensorscalar_reduce",
+ "binary_reduce",
+ [
+ Operand("dst", "region"),
+ Operand("reduced", "region"),
+ Operand("lhs"),
+ Operand("rhs"),
+ ],
+ {"opcode": "add", "reduce_op": "sum", "axes": [-1], "max_inst_size": None},
+ ("partial_reduce",),
+ br_project,
+ _binary_reduce,
+ )
+ ]
+
+ def ar_project(v, a):
+ return (
+ "unary_reduce_with_scale_bias",
+ [
+ v["dst"],
+ v["reduced"],
+ v["src"],
+ a.opcode,
+ a.reduce_op,
+ v["scale"],
+ v["bias"],
+ _axes(a),
+ ],
+ {},
+ )
+
+ result += [
+ spec(
+ "activation_reduce",
+ "unary_reduce_with_scale_bias",
+ [
+ Operand("dst", "region"),
+ Operand("reduced", "region"),
+ Operand("src", "region"),
+ Operand("scale", "scalar", 1.0),
+ Operand("bias", "expr", 0.0),
+ ],
+ {"opcode": "exp", "reduce_op": "sum", "axes": [-1], "max_inst_size": None},
+ ("const_bias", "partial_reduce"),
+ ar_project,
+ _activation_reduce,
+ )
+ ]
+
+ def select_project(v, a):
+ return "select", [v["dst"], v["true_value"], v["false_value"], v["pred"]], {}
+
+ select_spec = spec(
+ "affine_select",
+ "select",
+ [Operand("dst", "region"), Operand("true_value"), Operand("false_value"), Operand("pred")],
+ project=select_project,
+ lower=_select,
+ )
+ result += [select_spec]
+ ns = namespace("trn", result)
+ original = ns.affine_select
+
+ def affine_select(dst, true_value, false_value, pred, **kwargs):
+ from tvm.tirx.tensor_instruction import _region
+
+ dst = _region(dst)
+ if not isinstance(pred, LambdaExpr):
+ pred = LambdaExpr([PrimType("int32")] * len(dst.region), pred)
+ return original(dst, true_value, false_value, pred, **kwargs)
+
+ ns.affine_select = affine_select
+ return ns
diff --git a/python/tvm/backend/trn/tile_primitive/__init__.py b/python/tvm/backend/trn/tile_primitive/__init__.py
index 6334a8b19b..13a83393a9 100644
--- a/python/tvm/backend/trn/tile_primitive/__init__.py
+++ b/python/tvm/backend/trn/tile_primitive/__init__.py
@@ -14,12 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-
-from .binary import *
-from .compose_op import *
-from .copy import *
-from .gemm import *
-from .private_alloc import *
-from .reduction import *
-from .select import *
-from .unary import *
diff --git a/python/tvm/backend/trn/tile_primitive/binary/default.py b/python/tvm/backend/trn/tile_primitive/binary/default.py
index 5aae2335f3..487e258347 100644
--- a/python/tvm/backend/trn/tile_primitive/binary/default.py
+++ b/python/tvm/backend/trn/tile_primitive/binary/default.py
@@ -21,14 +21,14 @@ from tvm.script import tirx as T
from tvm.tirx import FloatImm, Function
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
from tvm.tirx.operator.tile_primitive.common import MapOpType
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
from ..common import init_analyzer, nki_dim
from ..instruction_generator import InstructionGenerator
from .utils import InstType, binary_map_ops, try_find_inst_nary
-def binary_trn(op: TileOpCall, binary_op: MapOpType, sctx: DispatchContext) -> Function | None:
+def binary_trn(op: TensorCall, binary_op: MapOpType, sctx: DispatchContext) -> Function | None:
"""Generate a binary operation schedule for Trainium."""
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
fail("requires Trainium target and thread exec_scope")
@@ -55,7 +55,7 @@ def binary_trn(op: TileOpCall, binary_op: MapOpType, sctx: DispatchContext) -> F
b_var = T.Var("B", "int32")
f_var = T.Var("F", "int32")
p_size = dst.ty.layout.size("P")
- inst_size_limit = op.config.get("max_inst_size", 512)
+ inst_size_limit = op.options.get("max_inst_size", 512)
inst_repr.bound_inst_size(inst_size_limit, analyzer)
inst_gen.bind_inst_iter(_dst, p_var, p_size, 1, False)
inst_gen.bind_inst_iter(_dst, f_var, inst_repr.size, inst_repr.stride, True)
@@ -64,7 +64,10 @@ def binary_trn(op: TileOpCall, binary_op: MapOpType, sctx: DispatchContext) -> F
opcode = binary_map_ops[binary_op]
# Select appropriate NKI function based on instruction type
- _func = T.nki.tensortensor if inst_types[0] == InstType.TENSOR_TENSOR else T.nki.tensorscalar
+ instruction = "tensortensor" if inst_types[0] == InstType.TENSOR_TENSOR else "tensorscalar"
+ if not op.op.name.endswith("." + instruction):
+ raise ValueError(f"{op.op.name} operands require {instruction}")
+ _func = getattr(T.nki, instruction)
def func(*args):
return _func(*args, reverse[0]) if inst_types[0] == InstType.TENSOR_SCALAR else _func(*args)
@@ -103,21 +106,3 @@ def binary_trn(op: TileOpCall, binary_op: MapOpType, sctx: DispatchContext) -> F
)
return impl
-
-
-# ---------------------------------------------------------------------------
-# Registration: bind each binary op name to its TRN schedule candidates.
-# ---------------------------------------------------------------------------
-from tvm.tirx.operator.tile_primitive import register_dispatch # noqa: E402
-
-for _op_name, _op_type in {
- "add": MapOpType.ADD,
- "sub": MapOpType.SUB,
- "mul": MapOpType.MUL,
- "maximum": MapOpType.MAX,
- "minimum": MapOpType.MIN,
-}.items():
-
- @register_dispatch(_op_name, "trn", variant="binary", priority=0)
- def _binary_dispatch(op, sctx, _ty=_op_type):
- return binary_trn(op, _ty, sctx)
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/__init__.py b/python/tvm/backend/trn/tile_primitive/compose_op/__init__.py
index ebca132fa2..58c9a2c405 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/__init__.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/__init__.py
@@ -17,5 +17,4 @@
from .binary_chain import *
from .binary_reduce import *
-from .reduce_negate import *
from .unary_reduce import *
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py b/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
index 08a7108254..cff07e752e 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
@@ -19,9 +19,10 @@
from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import Function, TileOpCall
+from tvm.tirx import Function
from tvm.tirx.op.tile import BinaryChain
-from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
+from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.tensor_instruction import TensorCall
from ..binary.utils import InstType, try_find_inst_nary
from ..common import init_analyzer, nki_dim
@@ -29,9 +30,9 @@ from ..instruction_generator import InstructionGenerator
from .utils import opcode_table
-def binary_chain_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
+def binary_chain_trn(op: TensorCall, sctx: DispatchContext) -> Function | None:
"""Generate a TRN schedule for binary chain operations."""
- op = TileOpCall.downcast(op)
+ op = TensorCall.decode(op)
assert isinstance(op, BinaryChain), f"invalid operator downcast: {op}"
# Extract operation components
@@ -61,7 +62,7 @@ def binary_chain_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
b_var = T.Var("B", "int32")
f_var = T.Var("F", "int32")
p_size = output.source.ty.layout.size("P")
- inst_size_limit = op.config.get("max_inst_size", 512)
+ inst_size_limit = op.options.get("max_inst_size", 512)
inst_repr.bound_inst_size(inst_size_limit, analyzer)
inst_gen.bind_inst_iter(output, p_var, p_size, 1, False)
inst_gen.bind_inst_iter(output, f_var, inst_repr.size, inst_repr.stride, True)
@@ -78,6 +79,14 @@ def binary_chain_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
else T.nki.scalar_tensor_tensor
)
+ instruction = (
+ "scalar_tensor_scalar"
+ if inst_types[1] == InstType.TENSOR_SCALAR
+ else "scalar_tensor_tensor"
+ )
+ if not op.op.name.endswith("." + instruction):
+ raise ValueError(f"{op.op.name} operands require {instruction}")
+
# Helper function to get source indices
def get_srcs(inst_gen):
return [
@@ -106,22 +115,3 @@ def binary_chain_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
# fmt: on
return impl
-
-
-@register_dispatch(
- "binary_chain",
- "trn",
- variant="default",
- priority=10,
- when=[
- predicate(
- "exec_scope",
- lambda op, sctx: (
- sctx.scope_kind == "thread",
- f"unsupported exec_scope {sctx.scope_kind}",
- ),
- )
- ],
-)
-def binary_chain_trn_dispatch(op: TileOpCall, sctx: DispatchContext) -> Function:
- return binary_chain_trn(op, sctx)
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py b/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
index 3e693c94fe..e953323c27 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
@@ -19,9 +19,10 @@
from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import Function, TileOpCall
+from tvm.tirx import Function
from tvm.tirx.op.tile import BinaryReduce
-from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
+from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.tensor_instruction import TensorCall
from ..binary.utils import InstType, try_find_inst_nary
from ..common import init_analyzer, nki_dim
@@ -31,9 +32,9 @@ from ..reduction.utils import generate_intermediate_buffer
from .utils import opcode_table
-def binary_reduce_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
+def binary_reduce_trn(op: TensorCall, sctx: DispatchContext) -> Function | None:
"""Generate a TRN schedule for binary reduction operations."""
- op = TileOpCall.downcast(op)
+ op = TensorCall.decode(op)
assert isinstance(op, BinaryReduce), f"invalid operator downcast: {op}"
# Extract operation components
@@ -63,7 +64,7 @@ def binary_reduce_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
)
# Apply instruction size limits
- inst_size_limit = op.config.get("max_inst_size", None)
+ inst_size_limit = op.options.get("max_inst_size", None)
inst_repr.bound_inst_size(inst_size_limit, analyzer)
# Generate axes and validate
@@ -87,7 +88,7 @@ def binary_reduce_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
spatial_b_extent = inst_gen.fill_in_block_dim(binary_output, spatial_b_var)
if reduction_b_extent != 1:
intermediate_buffer = generate_intermediate_buffer(
- reduce_output, reduction_b_extent, op.workspace, sctx
+ reduce_output, reduction_b_extent, op.workspaces, sctx
)
# Handle source 2 (either buffer region or constant)
@@ -151,23 +152,3 @@ def binary_reduce_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
# fmt: on
return impl
-
-
-# Rich dispatcher variants for TRN compose ops
-@register_dispatch(
- "binary_reduce",
- "trn",
- variant="default",
- priority=10,
- when=[
- predicate(
- "exec_scope",
- lambda op, sctx: (
- sctx.scope_kind == "thread",
- f"unsupported exec_scope {sctx.scope_kind}",
- ),
- )
- ],
-)
-def binary_reduce_trn_dispatch(op: TileOpCall, sctx: DispatchContext) -> Function:
- return binary_reduce_trn(op, sctx)
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/reduce_negate.py b/python/tvm/backend/trn/tile_primitive/compose_op/reduce_negate.py
deleted file mode 100644
index 1ae97d79d0..0000000000
--- a/python/tvm/backend/trn/tile_primitive/compose_op/reduce_negate.py
+++ /dev/null
@@ -1,51 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""Implementation of ReduceNegate dispatch."""
-
-from tvm.tirx import Function, TileOpCall
-from tvm.tirx.op.tile import ReduceNegate
-from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
-
-from ..reduction.utils import reduction_trn
-from .utils import optype_table
-
-
-def reduce_negate_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
- """Generate a TRN schedule for reduce negate operations."""
- op = TileOpCall.downcast(op)
- assert isinstance(op, ReduceNegate), f"invalid operator downcast: {op}"
- return reduction_trn(op, optype_table[op.reduce_op], sctx, negate=True)
-
-
-@register_dispatch(
- "reduce_negate",
- "trn",
- variant="default",
- priority=10,
- when=[
- predicate(
- "exec_scope",
- lambda op, sctx: (
- sctx.scope_kind == "thread",
- f"unsupported exec_scope {sctx.scope_kind}",
- ),
- )
- ],
-)
-def reduce_negate_trn_dispatch(op: TileOpCall, sctx: DispatchContext) -> Function:
- return reduce_negate_trn(op, sctx)
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py b/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
index 529c91ca01..24a612583c 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
@@ -19,9 +19,10 @@
from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import Function, TileOpCall
+from tvm.tirx import Function
from tvm.tirx.op.tile import UnaryReduce
-from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
+from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.tensor_instruction import TensorCall
from ..binary.utils import try_find_inst_nary
from ..common import init_analyzer, nki_dim
@@ -32,15 +33,15 @@ from ..unary.utils import get_const_bias_tensor, try_find_inst_unary
from .utils import opcode_table
-def unary_reduce_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
+def unary_reduce_trn(op: TensorCall, sctx: DispatchContext) -> Function | None:
"""Generate a TRN schedule for unary reduction operations."""
- op = TileOpCall.downcast(op)
+ op = TensorCall.decode(op)
assert isinstance(op, UnaryReduce), f"invalid operator downcast: {op}"
# Extract operation components
unary_output, reduce_output = op.dsts
unary_input = op.unary_input
- if op.op.name == "tirx.tile.unary_reduce_with_scale_bias":
+ if op.kind == "unary_reduce_with_scale_bias":
scale, bias = op.scale, op.bias
else:
scale, bias = 1.0, 0.0
@@ -71,7 +72,7 @@ def unary_reduce_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
)
# Apply instruction size limits
- inst_size_limit = op.config.get("max_inst_size", None)
+ inst_size_limit = op.options.get("max_inst_size", None)
inst_repr.bound_inst_size(inst_size_limit, analyzer)
p_var = T.Var("P", "int32")
@@ -85,7 +86,7 @@ def unary_reduce_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
spatial_b_extent = inst_gen.fill_in_block_dim(unary_output, spatial_b_var)
if reduction_b_extent != 1:
intermediate_buffer = generate_intermediate_buffer(
- reduce_output, reduction_b_extent, op.workspace, sctx
+ reduce_output, reduction_b_extent, op.workspaces, sctx
)
# Extract buffers and opcodes
src, dst1, dst2 = unary_input.source, unary_output.source, reduce_output.source
@@ -97,7 +98,7 @@ def unary_reduce_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
bias.source
if isinstance(bias, TensorRegion)
else get_const_bias_tensor(
- bias, (p_size, inst_repr.size), dst1.ty.dtype, op.workspace, sctx
+ bias, (p_size, inst_repr.size), dst1.ty.dtype, op.workspaces, sctx
)
)
@@ -159,37 +160,3 @@ def unary_reduce_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
# fmt: on
return impl
-
-
-@register_dispatch(
- "unary_reduce_with_scale_bias",
- "trn",
- variant="default",
- priority=10,
- when=[
- predicate(
- "exec_scope",
- lambda op, sctx: (
- sctx.scope_kind == "thread",
- f"unsupported exec_scope {sctx.scope_kind}",
- ),
- )
- ],
-)
-@register_dispatch(
- "unary_reduce",
- "trn",
- variant="default",
- priority=10,
- when=[
- predicate(
- "exec_scope",
- lambda op, sctx: (
- sctx.scope_kind == "thread",
- f"unsupported exec_scope {sctx.scope_kind}",
- ),
- )
- ],
-)
-def unary_reduce_trn_dispatch(op: TileOpCall, sctx: DispatchContext) -> Function:
- return unary_reduce_trn(op, sctx)
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/utils.py b/python/tvm/backend/trn/tile_primitive/compose_op/utils.py
index 1509482be8..f7ebbfd01f 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/utils.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/utils.py
@@ -17,25 +17,24 @@
"""Shared helpers for compose operator dispatches."""
-from tvm.ir import Op
from tvm.tirx.operator.tile_primitive.common import ReduceOpType
# Operation code mappings
opcode_table = {
- Op.get("tirx.tile.add"): "add",
- Op.get("tirx.tile.sub"): "sub",
- Op.get("tirx.tile.mul"): "mul",
- Op.get("tirx.tile.maximum"): "max",
- Op.get("tirx.tile.minimum"): "min",
- Op.get("tirx.tile.sqrt"): "sqrt",
- Op.get("tirx.tile.sum"): "add",
- Op.get("tirx.tile.max"): "max",
- Op.get("tirx.tile.min"): "min",
- Op.get("tirx.tile.exp"): "exp",
+ "add": "add",
+ "sub": "sub",
+ "mul": "mul",
+ "maximum": "max",
+ "minimum": "min",
+ "sqrt": "sqrt",
+ "sum": "add",
+ "max": "max",
+ "min": "min",
+ "exp": "exp",
}
optype_table = {
- Op.get("tirx.tile.sum"): ReduceOpType.SUM,
- Op.get("tirx.tile.max"): ReduceOpType.MAX,
- Op.get("tirx.tile.min"): ReduceOpType.MIN,
+ "sum": ReduceOpType.SUM,
+ "max": ReduceOpType.MAX,
+ "min": ReduceOpType.MIN,
}
diff --git a/python/tvm/backend/trn/tile_primitive/copy/default.py b/python/tvm/backend/trn/tile_primitive/copy/default.py
index 7857898c2b..fa5b2749e0 100644
--- a/python/tvm/backend/trn/tile_primitive/copy/default.py
+++ b/python/tvm/backend/trn/tile_primitive/copy/default.py
@@ -20,175 +20,15 @@
from tvm.backend.trn.layout import is_trainium_layout
from tvm.script import tirx as T
from tvm.tirx import Function
-from tvm.tirx.operator.tile_primitive import (
- DispatchContext,
- fail,
- predicate,
- register_dispatch,
-)
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.tensor_instruction import TensorCall
from ..common import init_analyzer, nki_dim
from ..dim_utils import get_ewise_dim_map
from ..instruction_generator import InstructionGenerator
-from ..workspace_utils import check_workspace_buffer, largest_psum_per_bank, max_psum_banks
-def transpose_schedule(
- op: TileOpCall, inst_gen: InstructionGenerator, sctx: DispatchContext
-) -> Function | None:
- dst_region, src_region = op.args
- assert src_region.source.scope() != "trn.psum", "Transpose on psum buffer is not supported"
-
- inst_repr_dst, inst_repr_src = inst_gen.find_max_inst_size_transpose(dst_region, src_region)
-
- lhs_f = T.Var("lhs_F", "int32")
- lhs_p = T.Var("lhs_P", "int32")
- dst_f = T.Var("dst_F", "int32")
- b_var = T.Var("B", "int32")
- extend_b = T.Var("extend_B", "int32")
- p_size = src_region.source.ty.layout.size("P")
- lhs_f_size = dst_region.source.ty.layout.size("P")
- rhs_f_size = p_size
- inst_gen.bind_inst_iter(
- src_region, lhs_f, inst_repr_src.size, inst_repr_src.stride, is_free_dim=True
- )
- inst_gen.bind_inst_iter(
- dst_region,
- dst_f,
- inst_repr_dst.size,
- inst_repr_dst.stride,
- is_free_dim=True,
- no_propagate=True,
- )
- inst_gen.bind_inst_iter(src_region, lhs_p, p_size, 1, is_free_dim=False, no_propagate=True)
- if dst_region.source.scope() == "trn.sbuf":
- max_extend_num = (
- inst_gen.find_max_inst_size_from_one_region(
- dst_region, min_stride=inst_repr_dst.stride
- ).size
- // rhs_f_size
- )
- max_elem_in_a_bank = largest_psum_per_bank // rhs_f_size
- if max_extend_num < max_elem_in_a_bank:
- extend_len = max_extend_num
- elif max_extend_num % max_elem_in_a_bank == 0:
- extend_len = max_elem_in_a_bank
- else:
- extend_len = 1
- inst_gen.bind_inst_iter(
- dst_region,
- extend_b,
- extend_len,
- inst_repr_dst.stride * inst_repr_dst.size,
- is_free_dim=True,
- )
- b_extent = inst_gen.fill_in_block_dim(dst_region, b_var)
-
- if "identity" not in op.workspace:
- assert sctx.alloc_only, (
- "Identity tensor must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first." # noqa: E501
- )
- identity_tensor = T.Var(
- "identity", T.Tensor((p_size, rhs_f_size), src_region.source.ty.dtype, scope="trn.sbuf")
- )
- sctx.add_alloc_buffer(identity_tensor)
-
- # This fragment captures buffers and indices from its insertion scope.
- @T.function(check_well_formed=False)
- def identity_init():
- with T.nki.tensorized_instruction():
- for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
- for rhs_f_loop in T.serial(0, rhs_f_size, annotations={nki_dim: "F"}):
- T.evaluate(T.nki.identity(identity_tensor[p_loop, rhs_f_loop], p_size))
- T.kernel_replace_point()
-
- sctx.add_init_stmt(identity_init.body)
- else:
- identity_tensor = op.workspace["identity"]
- check_workspace_buffer(identity_tensor, (p_size, rhs_f_size), "trn.sbuf")
-
- dst_buffer = dst_region.source
- src_buffer = src_region.source
- if dst_buffer.scope() == "trn.psum":
- # This fragment captures buffers and indices from its insertion scope.
- @T.function(check_well_formed=False)
- def transpose_psum_output():
- for b_loop in T.serial(0, b_extent):
- with T.nki.tensorized_instruction():
- for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
- for lhs_f_loop in T.serial(0, lhs_f_size, annotations={nki_dim: "lhs_F"}):
- for rhs_f_loop in T.serial(
- 0, rhs_f_size, annotations={nki_dim: "rhs_F"}
- ):
- inst_gen.set_bind_map(
- dst_region,
- {b_var: b_loop, lhs_f: lhs_f_loop, dst_f: rhs_f_loop},
- )
- inst_gen.set_bind_map(
- src_region, {b_var: b_loop, lhs_f: lhs_f_loop, lhs_p: p_loop}
- )
- src_indices = T.meta_var(inst_gen.generate_indices(src_region))
- dst_indices = T.meta_var(inst_gen.generate_indices(dst_region))
- if inst_gen.make_guard(src_region) and inst_gen.make_guard(
- dst_region
- ):
- T.evaluate(
- T.nki.matmul(
- dst_buffer[tuple(dst_indices)],
- src_buffer[tuple(src_indices)],
- identity_tensor[p_loop, rhs_f_loop],
- )
- )
-
- return transpose_psum_output
-
- if "acc_psum" not in op.workspace:
- assert sctx.alloc_only, (
- "Accumulation psum buffer must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first." # noqa: E501
- )
- acc_psum = T.Var(
- "acc_psum",
- T.Tensor(
- (max_psum_banks, p_size, largest_psum_per_bank),
- "float32",
- scope="trn.psum",
- allocated_addr=(0, 0),
- ),
- )
- sctx.add_alloc_buffer(acc_psum)
- max_psum_slots = max_psum_banks
- else:
- acc_psum = op.workspace["acc_psum"]
- check_workspace_buffer(acc_psum, (p_size, largest_psum_per_bank), "trn.psum")
- max_psum_slots = acc_psum.ty.shape[0]
-
- # fmt: off
- # This fragment captures buffers and indices from its insertion scope.
- @T.function(check_well_formed=False)
- def transpose_sbuf_output():
- for b_loop in T.serial(0, b_extent):
- for extend_b_loop in T.serial(0, extend_len):
- with T.nki.tensorized_instruction():
- for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
- for lhs_f_loop in T.serial(0, lhs_f_size, annotations={nki_dim: "lhs_F"}):
- for rhs_f_loop in T.serial(0, rhs_f_size, annotations={nki_dim: "rhs_F"}): # noqa: E501
- inst_gen.set_bind_map(src_region, {b_var: b_loop, lhs_f: lhs_f_loop, lhs_p: p_loop, extend_b: extend_b_loop}) # noqa: E501
- src_indices = T.meta_var(inst_gen.generate_indices(src_region))
- if inst_gen.make_guard(src_region):
- T.evaluate(T.nki.matmul(acc_psum[b_loop % max_psum_slots, lhs_f_loop,extend_b_loop * rhs_f_size + rhs_f_loop], src_buffer[tuple(src_indices)], identity_tensor[p_loop, rhs_f_loop])) # noqa: E501
- with T.nki.tensorized_instruction():
- for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
- for f_loop in T.serial(0, rhs_f_size * extend_len, annotations={nki_dim: "F"}):
- inst_gen.set_bind_map(dst_region, {b_var: b_loop, lhs_f: p_loop, dst_f: f_loop % rhs_f_size, extend_b: f_loop // rhs_f_size}) # noqa: E501
- dst_indices = T.meta_var(inst_gen.generate_indices(dst_region))
- if inst_gen.make_guard(dst_region):
- T.evaluate(T.nki.tensor_copy(dst_buffer[tuple(dst_indices)], acc_psum[b_loop % max_psum_slots, p_loop, f_loop])) # noqa: E501
- # fmt: on
- return transpose_sbuf_output
-
-
-def copy_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
+def copy_trn(op: TensorCall, sctx: DispatchContext) -> Function | None:
"""Schedule copy operation between global and shared memory on CUDA."""
# Basic validation checks
if sctx.scope_kind != "thread":
@@ -233,7 +73,10 @@ def copy_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
inst_gen.link_buffer_regions(src_region, dst_region, dim_map)
if not inst_gen.check_partition_dim_match(src_region, dst_region):
- return transpose_schedule(op, inst_gen, sctx)
+ raise ValueError(
+ "tensor_copy cannot transpose the partition dimension; "
+ "prepare identity and use matmul explicitly"
+ )
if is_trainium_layout(src.ty.layout):
inst = inst_gen.find_max_inst_size_from_one_region(src_region)
@@ -252,10 +95,10 @@ def copy_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
func = T.nki.tensor_copy
if func == T.nki.tensor_copy:
- inst_size_limit = op.config.get("max_inst_size", 512)
+ inst_size_limit = op.options.get("max_inst_size", 512)
inst.bound_inst_size(inst_size_limit, analyzer)
else:
- assert "max_inst_size" not in op.config, "max_inst_size is not supported for load/store"
+ assert "max_inst_size" not in op.options, "max_inst_size is not supported for load/store"
p_var = T.Var("P", "int32")
f_var = T.Var("F", "int32")
@@ -285,23 +128,3 @@ def copy_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
func(dst[tuple(dst_indices)], src[tuple(src_indices)])
# fmt: on
return impl
-
-
-# Rich dispatcher variant for TRN copy
-@register_dispatch(
- "copy",
- "trn",
- variant="default",
- priority=10,
- when=[
- predicate(
- "exec_scope",
- lambda op, sctx: (
- sctx.scope_kind == "thread",
- f"unsupported exec_scope {sctx.scope_kind}",
- ),
- )
- ],
-)
-def copy_trn_dispatch(op: TileOpCall, sctx: DispatchContext) -> Function:
- return copy_trn(op, sctx)
diff --git a/python/tvm/backend/trn/tile_primitive/gemm/default.py b/python/tvm/backend/trn/tile_primitive/gemm/default.py
index 2a1e690e77..92ac93b694 100644
--- a/python/tvm/backend/trn/tile_primitive/gemm/default.py
+++ b/python/tvm/backend/trn/tile_primitive/gemm/default.py
@@ -25,13 +25,8 @@ from tvm.ir import TensorRegion, assert_structural_equal
from tvm.script import tirx as T
from tvm.sym.analyzer import Analyzer
from tvm.tirx import Function
-from tvm.tirx.operator.tile_primitive import (
- DispatchContext,
- fail,
- predicate,
- register_dispatch,
-)
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.tensor_instruction import TensorCall
from ..common import init_analyzer
from ..dim_utils import normalize_and_group
@@ -110,7 +105,7 @@ def get_pf_dim_from_buffer_region(
return p_dim, f_dim
-def matmul_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
+def matmul_trn(op: TensorCall, sctx: DispatchContext) -> Function | None:
"""Schedule GEMM operation on Trainium."""
# Basic validation checks
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
@@ -254,7 +249,7 @@ def matmul_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
# and a matmul fusion pass can help infer the pattern
acc_psum_shape = (max_psum_banks, p_size, largest_psum_per_bank)
- if "acc_psum" not in op.workspace:
+ if "acc_psum" not in op.workspaces:
assert sctx.alloc_only, "Accumulation psum buffer must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first." # noqa: E501
acc_psum = T.Var(
"acc_psum",
@@ -263,7 +258,7 @@ def matmul_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
sctx.add_alloc_buffer(acc_psum)
max_psum_slots = max_psum_banks
else:
- acc_psum = op.workspace["acc_psum"]
+ acc_psum = op.workspaces["acc_psum"]
check_workspace_buffer(acc_psum, (p_size, largest_psum_per_bank), "trn.psum")
max_psum_slots = acc_psum.ty.shape[0]
@@ -282,23 +277,3 @@ def matmul_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
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
-
-
-# Rich dispatcher variant for TRN gemm
-@register_dispatch(
- "gemm",
- "trn",
- variant="default",
- priority=10,
- when=[
- predicate(
- "exec_scope",
- lambda op, sctx: (
- sctx.scope_kind == "thread",
- f"unsupported exec_scope {sctx.scope_kind}",
- ),
- )
- ],
-)
-def gemm_trn_dispatch(op: TileOpCall, sctx: DispatchContext) -> Function:
- return matmul_trn(op, sctx)
diff --git a/python/tvm/backend/trn/tile_primitive/private_alloc.py b/python/tvm/backend/trn/tile_primitive/private_alloc.py
index fb8b318696..0975b82a17 100644
--- a/python/tvm/backend/trn/tile_primitive/private_alloc.py
+++ b/python/tvm/backend/trn/tile_primitive/private_alloc.py
@@ -17,15 +17,12 @@
from typing import Any
-from tvm.backend.trn.tile_primitive.common import init_analyzer, nki_dim
-from tvm.backend.trn.tile_primitive.dim_utils import get_ewise_dim_map
-from tvm.backend.trn.tile_primitive.instruction_generator import InstructionGenerator
+from tvm.backend.trn.tile_primitive.common import nki_dim
from tvm.ir import Stmt
from tvm.script import tirx as T
-from tvm.tirx import FloatImm, IntImm, Var
+from tvm.tirx import DispatchContext, FloatImm, IntImm, Var
from tvm.tirx.op.tile import (
BinaryReduce,
- Copy,
Exp,
Gemm,
ReduceOp,
@@ -34,8 +31,7 @@ from tvm.tirx.op.tile import (
UnaryReduce,
)
from tvm.tirx.operator.tile_primitive.registry import f_op_dispatcher
-from tvm.tirx.stmt import TileOpCall
-from tvm.tirx.tile_dispatch import DispatchContext
+from tvm.tirx.tensor_instruction import TensorCall
def _scalar_dtype(scalar) -> str:
@@ -50,16 +46,16 @@ def _scalar_dtype(scalar) -> str:
def alloc_const_bias_trn(
- op: TileOpCall, buffer_dict: dict[Any, tuple[Var, Stmt | None]], sctx: DispatchContext
+ op: TensorCall, buffer_dict: dict[Any, tuple[Var, Stmt | None]], sctx: DispatchContext
) -> dict[str, Any]:
bias = getattr(op, "bias", FloatImm(op.dsts[0].source.ty.dtype, 0.0))
- if "const_bias" in op.workspace:
+ if "const_bias" in op.workspaces:
return {}
if not isinstance(bias, (FloatImm)):
return {}
bias_key = ("const_bias", _scalar_dtype(bias), float(bias.value).hex())
par_size = op.dsts[0].source.ty.layout.size("P")
- max_inst_size = op.config.get("max_inst_size", 512)
+ max_inst_size = op.options.get("max_inst_size", 512)
if isinstance(max_inst_size, int | IntImm) and int(max_inst_size) == -1:
raise ValueError("Constant bias workspace allocation requires a finite max_inst_size")
if bias_key in buffer_dict:
@@ -88,9 +84,9 @@ def alloc_const_bias_trn(
def alloc_partial_reduce_trn(
- op: TileOpCall, buffer_dict: dict[Any, tuple[Var, Stmt | None]], sctx: DispatchContext
+ op: TensorCall, buffer_dict: dict[Any, tuple[Var, Stmt | None]], sctx: DispatchContext
) -> dict[str, Any]:
- if "partial_reduce" in op.workspace:
+ if "partial_reduce" in op.workspaces:
return {}
f_op_dispatcher(op, sctx)
partial_reduce_buffer = None
@@ -107,41 +103,10 @@ def alloc_partial_reduce_trn(
return {"partial_reduce": partial_reduce_buffer}
-def alloc_identity_trn(
- op: TileOpCall, buffer_dict: dict[Any, tuple[Var, Stmt | None]], sctx: DispatchContext
-) -> dict[str, Any]:
- if "identity" in op.workspace:
- return {}
- par_size = op.srcs[0].source.ty.layout.size("P")
- if "identity" in buffer_dict:
- identity_buffer, identity_init_stmt = buffer_dict["identity"]
- old_shape = identity_buffer.ty.shape
- new_shape = [max(par_size, old_shape[0]), max(par_size, old_shape[1])]
- if new_shape[0] == old_shape[0] and new_shape[1] == old_shape[1]:
- return {"identity": "identity"}
- else:
- new_shape = (par_size, par_size)
- new_buffer = T.Var(
- "identity", T.Tensor(new_shape, dtype=op.srcs[0].source.ty.dtype, scope="trn.sbuf")
- )
-
- # This fragment captures buffers and indices from its insertion scope.
- @T.function(check_well_formed=False)
- def identity_init():
- with T.nki.tensorized_instruction():
- for p_loop in T.serial(0, par_size, annotations={nki_dim: "P"}):
- for rhs_f_loop in T.serial(0, par_size, annotations={nki_dim: "F"}):
- T.evaluate(T.nki.identity(new_buffer[p_loop, rhs_f_loop], par_size))
- T.kernel_replace_point()
-
- buffer_dict["identity"] = (new_buffer, identity_init.body)
- return {"identity": "identity"}
-
-
def alloc_acc_psum_trn(
- op: TileOpCall, buffer_dict: dict[Any, tuple[Var, Stmt | None]], sctx: DispatchContext
+ op: TensorCall, buffer_dict: dict[Any, tuple[Var, Stmt | None]], sctx: DispatchContext
) -> dict[str, Any]:
- if "acc_psum" in op.workspace or op.dsts[0].source.scope() == "trn.psum":
+ if "acc_psum" in op.workspaces or op.dsts[0].source.scope() == "trn.psum":
return {}
par_size = op.dsts[0].source.ty.layout.size("P")
acc_psum = T.Var(
@@ -152,32 +117,15 @@ def alloc_acc_psum_trn(
return {"acc_psum": acc_psum}
-def alloc_copy_trn(
- op: TileOpCall, buffer_dict: dict[Any, tuple[Var, Stmt | None]], sctx: DispatchContext
-) -> dict[str, Var]:
- src_region = op.srcs[0]
- dst_region = op.dsts[0]
- analyzer = init_analyzer(sctx)
- dim_map = get_ewise_dim_map(src_region, dst_region, analyzer)
- inst_gen = InstructionGenerator([src_region, dst_region], analyzer)
- inst_gen.link_buffer_regions(src_region, dst_region, dim_map)
- if inst_gen.check_partition_dim_match(src_region, dst_region):
- return {}
-
- identity_dict = alloc_identity_trn(op, buffer_dict, sctx)
- acc_psum_dict = alloc_acc_psum_trn(op, buffer_dict, sctx)
- return identity_dict | acc_psum_dict
-
-
def alloc_unary_reduce_trn(
- op: TileOpCall, buffer_dict: dict[Any, tuple[Var, Stmt | None]], sctx: DispatchContext
+ op: TensorCall, buffer_dict: dict[Any, tuple[Var, Stmt | None]], sctx: DispatchContext
) -> dict[str, Var]:
- if "max_inst_size" in op.config:
+ if "max_inst_size" in op.options:
partial_reduce_dict = alloc_partial_reduce_trn(op, buffer_dict, sctx)
const_bias_dict = alloc_const_bias_trn(op, buffer_dict, sctx)
return partial_reduce_dict | const_bias_dict
else:
- if "const_bias" in op.workspace and "partial_reduce" in op.workspace:
+ if "const_bias" in op.workspaces and "partial_reduce" in op.workspaces:
return {}
f_op_dispatcher(op, sctx)
partial_reduce_buffer = None
@@ -191,10 +139,10 @@ def alloc_unary_reduce_trn(
const_bias_buffer = buffer
# no reuse opportunity
workspace_dict = {}
- if partial_reduce_buffer is not None and "partial_reduce" not in op.workspace:
+ if partial_reduce_buffer is not None and "partial_reduce" not in op.workspaces:
buffer_dict[partial_reduce_buffer] = (partial_reduce_buffer, None)
workspace_dict["partial_reduce"] = partial_reduce_buffer
- if const_bias_buffer is not None and "const_bias" not in op.workspace:
+ if const_bias_buffer is not None and "const_bias" not in op.workspaces:
assert len(sctx.callbacks[DispatchContext.kDeviceInitStmt]) == 1, (
"const_bias should have init"
)
@@ -208,7 +156,6 @@ UnaryOpWithScaleBias.get_private_buffers_trn = alloc_const_bias_trn
Sqrt.get_private_buffers_trn = alloc_const_bias_trn
Exp.get_private_buffers_trn = alloc_const_bias_trn
ReduceOp.get_private_buffers_trn = alloc_partial_reduce_trn
-Copy.get_private_buffers_trn = alloc_copy_trn
Gemm.get_private_buffers_trn = alloc_acc_psum_trn
BinaryReduce.get_private_buffers_trn = alloc_partial_reduce_trn
UnaryReduce.get_private_buffers_trn = alloc_unary_reduce_trn
diff --git a/python/tvm/backend/trn/tile_primitive/reduction/__init__.py b/python/tvm/backend/trn/tile_primitive/reduction/__init__.py
index 358e449317..13a83393a9 100644
--- a/python/tvm/backend/trn/tile_primitive/reduction/__init__.py
+++ b/python/tvm/backend/trn/tile_primitive/reduction/__init__.py
@@ -14,5 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-
-from .default import *
diff --git a/python/tvm/backend/trn/tile_primitive/reduction/default.py b/python/tvm/backend/trn/tile_primitive/reduction/default.py
deleted file mode 100644
index 0a1d8b6cbb..0000000000
--- a/python/tvm/backend/trn/tile_primitive/reduction/default.py
+++ /dev/null
@@ -1,33 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""Reduction dispatch variant registrations."""
-
-from tvm.tirx.operator.tile_primitive import register_dispatch
-from tvm.tirx.operator.tile_primitive.common import ReduceOpType
-
-from .utils import reduction_trn
-
-for _op_name, _op_type in {
- "sum": ReduceOpType.SUM,
- "max": ReduceOpType.MAX,
- "min": ReduceOpType.MIN,
-}.items():
-
- @register_dispatch(_op_name, "trn", variant="reduction", priority=0)
- def _reduction_dispatch(op, sctx, _ty=_op_type):
- return reduction_trn(op, _ty, sctx)
diff --git a/python/tvm/backend/trn/tile_primitive/reduction/utils.py b/python/tvm/backend/trn/tile_primitive/reduction/utils.py
index 94ed3876c2..7ff2c85192 100644
--- a/python/tvm/backend/trn/tile_primitive/reduction/utils.py
+++ b/python/tvm/backend/trn/tile_primitive/reduction/utils.py
@@ -22,7 +22,7 @@ from tvm.script import tirx as T
from tvm.tirx import Function
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
from tvm.tirx.operator.tile_primitive.common import ReduceOpType
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
from ..common import init_analyzer, nki_dim
from ..dim_utils import get_reduction_dim_map
@@ -59,7 +59,7 @@ def generate_intermediate_buffer(
def reduction_trn(
- op: TileOpCall, reduce_op: ReduceOpType, sctx: DispatchContext, negate: bool = False
+ op: TensorCall, reduce_op: ReduceOpType, sctx: DispatchContext, negate: bool = False
) -> Function | None:
"""Schedule reduction operation on Trainium.
@@ -75,7 +75,7 @@ def reduction_trn(
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
fail("requires Trainium target and thread exec_scope")
- op = TileOpCall.downcast(op)
+ op = TensorCall.decode(op)
dst_buffer_region, src_buffer_region = op.output, op.input
axes, accum = op.reduce_axes, op.accum
assert not accum, "Accumulation is not supported for reduction on Trainium"
@@ -104,7 +104,7 @@ def reduction_trn(
inst_gen = InstructionGenerator([src_buffer_region, dst_buffer_region], analyzer)
inst_gen.link_buffer_regions(src_buffer_region, dst_buffer_region, dim_map)
inst_repr = inst_gen.find_max_inst_size_from_one_region(src_buffer_region, axes)
- inst_size_limit = op.config.get("max_inst_size", None)
+ inst_size_limit = op.options.get("max_inst_size", None)
inst_repr.bound_inst_size(inst_size_limit, analyzer)
assert analyzer.can_prove(inst_repr.size > 1), "Instruction size must be greater than 1"
@@ -124,7 +124,7 @@ def reduction_trn(
# Generate intermediate buffer if needed
if reduction_b_extent != 1:
intermediate_buffer = generate_intermediate_buffer(
- dst_buffer_region, reduction_b_extent, op.workspace, sctx
+ dst_buffer_region, reduction_b_extent, op.workspaces, sctx
)
# fmt: off
diff --git a/python/tvm/backend/trn/tile_primitive/select/default.py b/python/tvm/backend/trn/tile_primitive/select/default.py
index effeafa3d1..f2780a130b 100644
--- a/python/tvm/backend/trn/tile_primitive/select/default.py
+++ b/python/tvm/backend/trn/tile_primitive/select/default.py
@@ -20,26 +20,22 @@
from tvm.backend.trn.layout import is_trainium_layout
from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import FloatImm, Function, TileOpCall
+from tvm.tirx import FloatImm, Function
from tvm.tirx.op.tile import Select
-from tvm.tirx.operator.tile_primitive import (
- DispatchContext,
- fail,
- predicate,
- register_dispatch,
-)
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.tensor_instruction import TensorCall
from ..common import init_analyzer, nki_dim
from ..dim_utils import get_ewise_dim_map
from ..instruction_generator import InstructionGenerator
-def select_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
+def select_trn(op: TensorCall, sctx: DispatchContext) -> Function | None:
"""Generate schedule for select operation on Trainium."""
if sctx.scope_kind != "thread":
fail("requires thread exec_scope for TRN select")
- op = TileOpCall.downcast(op)
+ op = TensorCall.decode(op)
assert isinstance(op, Select), f"{op} is not a Select"
# Unpack operands
@@ -94,7 +90,7 @@ def select_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
inst_repr = inst_gen.find_max_inst_size_from_one_region(dst)
inst_repr = inst_gen.fit_inst_tile_to_region(inst_repr, true_value)
inst_repr = inst_gen.restrict_inst_to_one_dim(inst_repr)
- inst_repr.bound_inst_size(op.config.get("max_inst_size", 512), analyzer)
+ inst_repr.bound_inst_size(op.options.get("max_inst_size", 512), analyzer)
p_var = T.Var("p", "int32")
b_var = T.Var("b", "int32")
@@ -124,23 +120,3 @@ def select_trn(op: TileOpCall, sctx: DispatchContext) -> Function | None:
# fmt: on
return impl
-
-
-# Rich dispatcher variant for TRN select
-@register_dispatch(
- "select",
- "trn",
- variant="default",
- priority=10,
- when=[
- predicate(
- "exec_scope",
- lambda op, sctx: (
- sctx.scope_kind == "thread",
- f"unsupported exec_scope {sctx.scope_kind}",
- ),
- )
- ],
-)
-def select_trn_dispatch(op: TileOpCall, sctx: DispatchContext) -> Function:
- return select_trn(op, sctx)
diff --git a/python/tvm/backend/trn/tile_primitive/unary/default.py b/python/tvm/backend/trn/tile_primitive/unary/default.py
index d40b6f4653..59dfb94751 100644
--- a/python/tvm/backend/trn/tile_primitive/unary/default.py
+++ b/python/tvm/backend/trn/tile_primitive/unary/default.py
@@ -20,7 +20,7 @@
from tvm.tirx import FloatImm, Function
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
from tvm.tirx.operator.tile_primitive.common import MapOpType
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
from ..common import init_analyzer
from ..instruction_generator import InstructionGenerator
@@ -32,7 +32,7 @@ from .utils import (
)
-def unary_trn(op: TileOpCall, unary_op: MapOpType, sctx: DispatchContext) -> Function | None:
+def unary_trn(op: TensorCall, unary_op: MapOpType, sctx: DispatchContext) -> Function | None:
"""Schedule unary operation on Trainium."""
# Check execution environment
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
@@ -71,19 +71,7 @@ def unary_trn(op: TileOpCall, unary_op: MapOpType, sctx: DispatchContext) -> Fun
None, # No bias
None, # No scale
analyzer,
- op.workspace,
- op.config,
+ op.workspaces,
+ op.options,
sctx,
)
-
-
-# ---------------------------------------------------------------------------
-# Registration: bind each default unary op name to its TRN schedule candidates.
-# ---------------------------------------------------------------------------
-from tvm.tirx.operator.tile_primitive import register_dispatch # noqa: E402
-
-for _op_name, _op_type in {"reciprocal": MapOpType.RECIPROCAL, "memset": MapOpType.FILL}.items():
-
- @register_dispatch(_op_name, "trn", variant="unary", priority=0)
- def _unary_dispatch(op, sctx, _ty=_op_type):
- return unary_trn(op, _ty, sctx)
diff --git a/python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py b/python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py
index 797ef2d9ec..a523aca8cf 100644
--- a/python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py
+++ b/python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py
@@ -21,7 +21,7 @@ from tvm.ir import TensorRegion
from tvm.tirx import Function
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
from tvm.tirx.operator.tile_primitive.common import MapOpType
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.tensor_instruction import TensorCall
from ..binary import try_find_inst_nary
from ..common import init_analyzer
@@ -30,7 +30,7 @@ from .utils import activation_map_ops, generate_unary_func, try_find_inst_unary
def unary_with_bias_scale_trn(
- op: TileOpCall, unary_op: MapOpType = MapOpType.SQRT, sctx: DispatchContext = None
+ op: TensorCall, unary_op: MapOpType = MapOpType.SQRT, sctx: DispatchContext = None
) -> Function | None:
"""Schedule unary operation with bias and scale on Trainium."""
# Check execution environment
@@ -71,22 +71,7 @@ def unary_with_bias_scale_trn(
_bias,
scale,
analyzer,
- op.workspace,
- op.config,
+ op.workspaces,
+ op.options,
sctx,
)
-
-
-# ---------------------------------------------------------------------------
-# Registration: bind each unary_with_bias_scale op name to its TRN schedule candidates.
-# ---------------------------------------------------------------------------
-from tvm.tirx.operator.tile_primitive import register_dispatch # noqa: E402
-
-for _op_name, _op_type in {"sqrt": MapOpType.SQRT, "exp": MapOpType.EXP}.items():
-
- @register_dispatch(_op_name, "trn", variant="unary_with_bias_scale", priority=0)
- @register_dispatch(
- _op_name + "_with_scale_bias", "trn", variant="unary_with_bias_scale", priority=0
- )
- def _unary_bs_dispatch(op, sctx, _ty=_op_type):
- return unary_with_bias_scale_trn(op, _ty, sctx)
diff --git a/python/tvm/backend/trn/transform/private_buffer_alloc.py b/python/tvm/backend/trn/transform/private_buffer_alloc.py
index c0acafd314..0e253fd770 100644
--- a/python/tvm/backend/trn/transform/private_buffer_alloc.py
+++ b/python/tvm/backend/trn/transform/private_buffer_alloc.py
@@ -32,7 +32,7 @@ from tvm.ir import (
Var,
)
from tvm.target import Target
-from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.tensor_instruction import SPECS, TensorCall
from tvm.tirx.tile_dispatch import DispatchContext
from tvm.tirx.transform.common import seek_kernel_replace_point
from tvm.tirx.transform.function_pass import function_pass
@@ -60,8 +60,11 @@ def _collect_private_allocations(stmt: Stmt, target: Target):
def visit_for(op: For):
var_range_map[op.loop_var] = Range.from_min_extent(op.min, op.extent)
- def visit_op_call(op: TileOpCall):
- # Scope is a per-call field on the node; read it directly.
+ def visit_op_call(op: Call):
+ if getattr(op.op, "name", "") not in SPECS:
+ return
+ call = op
+ op = TensorCall.decode(call)
exec_scope = op.scope
scope_kind = op.scope.name
sctx = DispatchContext(
@@ -72,13 +75,12 @@ def _collect_private_allocations(stmt: Stmt, target: Target):
alloc_only=True,
scope_kind=scope_kind,
)
- op = TileOpCall.downcast(op)
- private_buf_refs[op] = op.get_private_buffers(buffer_dict, sctx)
+ private_buf_refs[call] = op.get_private_buffers(buffer_dict, sctx)
def visit(node):
tvm_ffi.structural_walk(
node,
- [(RegionStmt, visit_region), (For, visit_for), (TileOpCall, visit_op_call)],
+ [(RegionStmt, visit_region), (For, visit_for), (Call, visit_op_call)],
order="pre",
)
@@ -90,7 +92,7 @@ def _inject_private_allocations(
stmt: Stmt,
alloc_buffers: list[Var],
init_stmts: list[Stmt],
- added_workspace: dict[TileOpCall, dict[str, Var]],
+ added_workspace: dict[Call, dict[str, Var]],
) -> Stmt:
is_outer_block = True
@@ -125,16 +127,16 @@ def _inject_private_allocations(
)
return op
- def visit_op_call(op: TileOpCall):
+ def visit_op_call(op: Call):
if op not in added_workspace:
return op
- new_workspace = dict(op.workspace)
+ new_workspace = dict(TensorCall.decode(op).workspaces)
new_workspace.update(added_workspace[op])
- return TileOpCall.downcast(op).with_workspace(new_workspace)
+ return TensorCall.decode(op).with_workspaces(new_workspace)
return tvm_ffi.structural_map(
stmt,
- [(RegionStmt, visit_region), (TileOpCall, visit_op_call)],
+ [(RegionStmt, visit_region), (Call, visit_op_call)],
order="pre",
)
@@ -154,7 +156,7 @@ def private_alloc(stmt: Stmt, target: Target) -> Stmt:
@function_pass(opt_level=0, name="TrnPrivateBufferAlloc")
class TrnPrivateBufferAlloc:
- """Generate private buffer allocations for each TileOpCall"""
+ """Generate private buffer allocations for each tensor instruction Call"""
def transform_function(self, func, mod, ctx):
target = func.attrs.get("target", None)
diff --git a/python/tvm/script/__init__.py b/python/tvm/script/__init__.py
index cfab7ab503..22f41188fb 100644
--- a/python/tvm/script/__init__.py
+++ b/python/tvm/script/__init__.py
@@ -144,14 +144,14 @@ class _DialectRedirectFinder:
(``from tvm.script import tirx``). It cannot intercept deep
statement-form imports such as::
- import tvm.script.tirx.tile
+ import tvm.script.tirx.ir_builder
import tvm.script.ir_builder.relax.ir
This finder is installed on ``sys.meta_path`` to cover those cases.
When the import machinery asks for a module whose full name starts with
``tvm.script.<language variant>`` (or ``tvm.script.parser.<language variant>``, etc.) and
that language variant is in ``_DIALECT_REGISTRY``, :meth:`find_spec` imports the
- real target module (e.g. ``tvm.tirx.script.tile``) and returns an
+ real target module (e.g. ``tvm.tirx.script.ir_builder``) and returns an
alias spec whose loader hands back that module, so the import machinery
registers it in ``sys.modules`` under the public alias and all subsequent
imports and attribute walks resolve without going through the redirect
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 5c07d9b64a..cced97669a 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -47,7 +47,6 @@ from .expr import CallEffectKind, Let
from .stmt import BufferRegion
from .stmt import ScopeIdDefStmt
from .tile_dispatch import DispatchContext
-from .stmt import TileOpCall
from .function import Function, IndexMap, renew_def
diff --git a/python/tvm/tirx/op/tile.py b/python/tvm/tirx/op/tile.py
index da44ee373d..ffc1ae2d87 100644
--- a/python/tvm/tirx/op/tile.py
+++ b/python/tvm/tirx/op/tile.py
@@ -17,14 +17,8 @@
"""Implementation of TIR operator."""
-from tvm.ir import Op
from tvm.tirx import Expr
-from tvm.tirx.stmt import TileOpCall
-
-
-def get_tirx_op(op_name: str):
- assert isinstance(op_name, str)
- return Op.get("tirx.tile." + op_name)
+from tvm.tirx.tensor_instruction import TensorCall
class ArgProperty:
@@ -32,12 +26,12 @@ class ArgProperty:
self.index = index
def __get__(self, obj, objtype=None):
- assert obj is not None, "TileOpCall cannot be None"
+ assert obj is not None, "TensorCall cannot be None"
return obj.args[self.index]
### Base Operator Classes ###
-class UnaryOp(TileOpCall):
+class UnaryOp(TensorCall):
"""Base class for unary operators: unary(output, input).
Unary operators take a single input tensor and produce a single output tensor.
@@ -74,7 +68,7 @@ class UnaryOpWithScaleBias(UnaryOp):
return [self.input, self.scale, self.bias]
-class BinaryOp(TileOpCall):
+class BinaryOp(TensorCall):
"""Base class for binary operators: binary(output, input0, input1).
Binary operators take two input tensors and produce a single output tensor.
@@ -95,7 +89,7 @@ class BinaryOp(TileOpCall):
return [self.output]
-class ReduceOp(TileOpCall):
+class ReduceOp(TensorCall):
"""Base class for reduction operators: reduce(output, input, reduce_axes, accum).
Reduction operators reduce one or more dimensions of the input tensor.
@@ -121,53 +115,53 @@ class ReduceOp(TileOpCall):
class Zero(UnaryOp):
"""Zero out all elements in src and store to dst."""
- op = get_tirx_op("zero")
+ kind = "zero"
class Sqrt(UnaryOp):
"""Compute sqrt(src) and store to dst."""
- op = get_tirx_op("sqrt")
+ kind = "sqrt"
class SqrtWithScaleBias(UnaryOpWithScaleBias):
"""Compute sqrt(src * scale + bias) and store to dst."""
- op = get_tirx_op("sqrt_with_scale_bias")
+ kind = "sqrt_with_scale_bias"
class Fill(UnaryOp):
"""Fill dst with a scalar value."""
- op = get_tirx_op("fill")
+ kind = "fill"
scalar_input = True
class Add(BinaryOp):
"""Add src1 and src2 element-wise and store to dst."""
- op = get_tirx_op("add")
+ kind = "add"
class Sub(BinaryOp):
"""Subtract src2 from src1 element-wise and store to dst."""
- op = get_tirx_op("sub")
+ kind = "sub"
class Mul(BinaryOp):
"""Multiply src1 and src2 element-wise and store to dst."""
- op = get_tirx_op("mul")
+ kind = "mul"
class FDiv(BinaryOp):
"""Divide src1 by src2 element-wise using floating point division and store to dst."""
- op = get_tirx_op("fdiv")
+ kind = "fdiv"
-class FMA(TileOpCall):
+class FMA(TensorCall):
"""Fused multiply-add: output = input * scale + bias.
fma(output, input, scale, bias)
@@ -175,7 +169,7 @@ class FMA(TileOpCall):
scale and bias can each be either a TensorRegion or an Expr scalar.
"""
- op = get_tirx_op("fma")
+ kind = "fma"
output = ArgProperty(0)
input = ArgProperty(1)
@@ -196,10 +190,10 @@ class FMA(TileOpCall):
class Cast(UnaryOp):
"""Cast src to dst."""
- op = get_tirx_op("cast")
+ kind = "cast"
-class Copy(TileOpCall):
+class Copy(TensorCall):
"""Copy all elements from src to dst.
Args:
@@ -207,7 +201,7 @@ class Copy(TileOpCall):
src: Source buffer region
"""
- op = get_tirx_op("copy")
+ kind = "copy"
dst = ArgProperty(0)
src = ArgProperty(1)
@@ -223,7 +217,7 @@ class Copy(TileOpCall):
return [self.dst]
-class CopyAsync(TileOpCall):
+class CopyAsync(TensorCall):
"""Copy all elements from src to dst asynchronously.
Args:
@@ -231,7 +225,7 @@ class CopyAsync(TileOpCall):
src: Source buffer region
"""
- op = get_tirx_op("copy_async")
+ kind = "copy_async"
dst = ArgProperty(0)
src = ArgProperty(1)
@@ -247,7 +241,7 @@ class CopyAsync(TileOpCall):
return [self.dst]
-class Gemm(TileOpCall):
+class Gemm(TensorCall):
"""General matrix multiplication: D = A * B * alpha + C * beta.
Args:
@@ -261,7 +255,7 @@ class Gemm(TileOpCall):
beta: Scalar multiplier for C
"""
- op = get_tirx_op("gemm")
+ kind = "gemm"
output = ArgProperty(0)
lhs = ArgProperty(1)
rhs = ArgProperty(2)
@@ -282,7 +276,7 @@ class Gemm(TileOpCall):
return [self.output]
-class GemmAsync(TileOpCall):
+class GemmAsync(TensorCall):
"""General matrix multiplication asynchronously.
Supports two arg layouts:
@@ -290,7 +284,7 @@ class GemmAsync(TileOpCall):
- Block-scaled (8 args): C, A, B, SFA, SFB, transA, transB, accum
"""
- op = get_tirx_op("gemm_async")
+ kind = "gemm_async"
output = ArgProperty(0)
lhs = ArgProperty(1)
rhs = ArgProperty(2)
@@ -339,86 +333,86 @@ class GemmAsync(TileOpCall):
class Sum(ReduceOp):
"""Sum elements in src along specified axes and store in dst."""
... 15242 lines suppressed ...