You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by bo...@apache.org on 2026/10/13 19:05:17 UTC
(tvm) branch main updated: refactor(op-dispatch): replace adaptive CUDA tile dispatch with explicit patterns (#20687)
This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 08e48d8142 refactor(op-dispatch): replace adaptive CUDA tile dispatch with explicit patterns (#20687)
08e48d8142 is described below
commit 08e48d8142fb49c42d5fb2a30e934bc5bff8b4fa
Author: Bohan Hou <sp...@gmail.com>
AuthorDate: Tue Oct 13 15:05:09 2026 -0400
refactor(op-dispatch): replace adaptive CUDA tile dispatch with explicit patterns (#20687)
CUDA tile operations now take an explicit PTX instruction and canonical
operand views. Lowering checks their physical mapping and emits the
selected instruction for each repeat coordinate. Instruction selection,
layout reorganization, staging, packing, thread election, and
synchronization are controlled by the caller.
- Replace adaptive CUDA dispatch with
`T.cuda.tile.apply(T.ptx["exact.opcode"], ...)` and
`T.cuda.tile.ewise(scalar_function, ...)`.
- Define operand contracts for memory and register instructions, matrix
loads/stores, asynchronous copies, TMA, TMEM transfers, warp MMA, and
tcgen05 MMA, including native packing, decompression, descriptors, and
block scales.
- Handle gapped allocation spans, dtype views, swizzled origins,
symbolic bounds, and runtime TMA descriptor fields.
- Organize CUDA tests and instruction guides in flat directories by PTX
instruction family. Document operand mappings, explicit views,
synchronization, and runnable kernels; keep scalar compositions under
`ewise`.
Companion TIRx-lite adaptation:
https://github.com/mlc-ai/TIRx-kernels/pull/9.
---
docs/tirx/api/ptx.rst | 2 +-
docs/tirx/api/script.rst | 42 +-
docs/tirx/api/tile.rst | 74 +-
docs/tirx/api/tile_dispatch.rst | 6 +-
docs/tirx/arch/backends.rst | 2 +-
docs/tirx/arch/lowering_pipeline.rst | 6 +-
docs/tirx/arch/tile_dispatch.rst | 2 +-
docs/tirx/cuda_tile_patterns.rst | 514 +++
docs/tirx/get_started/index.rst | 2 +-
docs/tirx/layout.rst | 25 +-
docs/tirx/native_basics.rst | 4 +-
docs/tirx/native_basics/cuda/compiling.rst | 6 +-
docs/tirx/native_basics/cuda/tensors.rst | 4 +-
docs/tirx/native_basics/cuda/threads_sync.rst | 2 +-
docs/tirx/overview.rst | 64 +-
docs/tirx/programming/index.rst | 4 +-
docs/tirx/tile_ops.rst | 117 +
docs/tirx/tile_ops/abs_neg.rst | 56 +
docs/tirx/tile_ops/add_sub_mul.rst | 113 +
docs/tirx/tile_ops/cp_async.rst | 153 +
docs/tirx/tile_ops/cp_async_bulk.rst | 124 +
docs/tirx/tile_ops/cp_async_bulk_tensor.rst | 184 +
docs/tirx/tile_ops/cvt.rst | 75 +
docs/tirx/tile_ops/div_rcp_sqrt.rst | 77 +
docs/tirx/tile_ops/ewise.rst | 133 +
docs/tirx/tile_ops/ex2_lg2.rst | 56 +
docs/tirx/tile_ops/fma.rst | 58 +
docs/tirx/tile_ops/ld_st.rst | 128 +
docs/tirx/tile_ops/ldmatrix_stmatrix.rst | 147 +
docs/tirx/tile_ops/min_max.rst | 66 +
docs/tirx/tile_ops/mma_sync.rst | 159 +
docs/tirx/tile_ops/mov.rst | 67 +
.../{tile_primitives => tile_ops}/reduction.rst | 2 +-
docs/tirx/tile_ops/tcgen05_cp.rst | 171 +
docs/tirx/tile_ops/tcgen05_ld_st.rst | 174 +
docs/tirx/tile_ops/tcgen05_mma.rst | 230 ++
docs/tirx/tile_primitives.rst | 140 -
docs/tirx/tile_primitives/copy.rst | 36 -
docs/tirx/tile_primitives/copy/ldstmatrix.rst | 241 --
docs/tirx/tile_primitives/copy/reg.rst | 222 --
docs/tirx/tile_primitives/copy_async.rst | 45 -
docs/tirx/tile_primitives/copy_async/dsmem.rst | 222 --
docs/tirx/tile_primitives/copy_async/ldgsts.rst | 195 -
.../tirx/tile_primitives/copy_async/tcgen05_cp.rst | 222 --
.../tile_primitives/copy_async/tcgen05_ldst.rst | 263 --
docs/tirx/tile_primitives/copy_async/tma.rst | 243 --
docs/tirx/tile_primitives/elementwise.rst | 43 -
docs/tirx/tile_primitives/elementwise/reg.rst | 159 -
docs/tirx/tile_primitives/elementwise/smem.rst | 160 -
docs/tirx/tile_primitives/gemm.rst | 195 -
docs/tirx/tile_primitives/gemm_async.rst | 298 --
python/tvm/backend/cuda/cpp/descriptors.py | 4 +-
python/tvm/backend/cuda/lang/alloc_pool.py | 10 +-
python/tvm/backend/cuda/layout.py | 225 +-
python/tvm/backend/cuda/ptx/__init__.py | 11 +-
python/tvm/backend/cuda/ptx/engine.py | 45 +-
python/tvm/backend/cuda/tile/common.py | 278 --
python/tvm/backend/cuda/tile/copy/__init__.py | 18 -
python/tvm/backend/cuda/tile/copy/_common.py | 590 ---
python/tvm/backend/cuda/tile/copy/ld_stmatrix.py | 401 ---
python/tvm/backend/cuda/tile/copy/utils.py | 74 -
python/tvm/backend/cuda/tile/copy/vec_auto_reg.py | 633 ----
python/tvm/backend/cuda/tile/copy/vec_forced.py | 205 --
.../tvm/backend/cuda/tile/copy_async/__init__.py | 24 -
python/tvm/backend/cuda/tile/copy_async/dsmem.py | 197 -
python/tvm/backend/cuda/tile/copy_async/ldgsts.py | 269 --
.../tvm/backend/cuda/tile/copy_async/tcgen05_cp.py | 823 -----
.../backend/cuda/tile/copy_async/tcgen05_ldst.py | 572 ---
python/tvm/backend/cuda/tile/copy_async/tma.py | 2284 ------------
python/tvm/backend/cuda/tile/copy_async/utils.py | 76 -
.../tvm/backend/cuda/tile/elementwise/__init__.py | 28 -
.../tvm/backend/cuda/tile/elementwise/_common.py | 412 ---
.../backend/cuda/tile/elementwise/ops/__init__.py | 122 -
.../backend/cuda/tile/elementwise/ops/binary.py | 144 -
.../tvm/backend/cuda/tile/elementwise/ops/cast.py | 46 -
.../tvm/backend/cuda/tile/elementwise/ops/fma.py | 51 -
.../tvm/backend/cuda/tile/elementwise/ops/unary.py | 138 -
python/tvm/backend/cuda/tile/elementwise/reg.py | 447 ---
python/tvm/backend/cuda/tile/elementwise/smem.py | 271 --
.../cuda/tile/elementwise/vec_emit/__init__.py | 51 -
.../cuda/tile/elementwise/vec_emit/binary_f32x2.py | 105 -
.../cuda/tile/elementwise/vec_emit/cast_vec2.py | 92 -
.../cuda/tile/elementwise/vec_emit/fma_f32x2.py | 86 -
python/tvm/backend/cuda/tile/exec_scope_utils.py | 104 -
python/tvm/backend/cuda/tile/gemm/__init__.py | 20 -
python/tvm/backend/cuda/tile/gemm/mma_m16n8k_.py | 580 ---
.../tvm/backend/cuda/tile/gemm_async/__init__.py | 18 -
python/tvm/backend/cuda/tile/gemm_async/tcgen05.py | 1495 --------
python/tvm/backend/cuda/tile/gemm_utils.py | 67 -
python/tvm/backend/cuda/tile/instructions.py | 433 +--
python/tvm/backend/cuda/tile/layout_utils.py | 373 --
python/tvm/backend/cuda/tile/patterns/__init__.py | 98 +
.../tvm/backend/cuda/tile/patterns/async_copy.py | 90 +
python/tvm/backend/cuda/tile/patterns/common.py | 383 ++
.../tvm/backend/cuda/tile/patterns/descriptors.py | 82 +
python/tvm/backend/cuda/tile/patterns/ewise.py | 115 +
python/tvm/backend/cuda/tile/patterns/ldmatrix.py | 60 +
python/tvm/backend/cuda/tile/patterns/memory.py | 62 +
python/tvm/backend/cuda/tile/patterns/mma.py | 107 +
python/tvm/backend/cuda/tile/patterns/register.py | 86 +
.../tvm/backend/cuda/tile/patterns/tcgen05_mma.py | 267 ++
python/tvm/backend/cuda/tile/patterns/tma.py | 356 ++
python/tvm/backend/cuda/tile/patterns/tmem.py | 154 +
python/tvm/backend/cuda/tile/patterns/tmem_copy.py | 100 +
python/tvm/backend/cuda/tile/tma_utils.py | 151 -
python/tvm/backend/trn/transform/__init__.py | 2 +-
python/tvm/tirx/_tensor_view.py | 30 +-
python/tvm/tirx/layout.py | 54 +-
python/tvm/tirx/tile/instruction.py | 52 +-
python/tvm/tirx/transform/transform.py | 2 +-
src/backend/cuda/runtime/cuda_device_api.cc | 7 +-
src/backend/cuda/tile/exec_analysis.cc | 7 +-
src/sym/rewrite_simplify.cc | 33 +
src/sym/rewrite_simplify.h | 3 +
src/tirx/ir/layout/tile_core.cc | 16 +-
src/tirx/tile/lower.cc | 3 +-
src/tirx/tile/op.cc | 135 -
src/tirx/transform/lower_tirx_cleanup.cc | 25 +-
tests/python/sym/test_sym_rewrite_simplify.py | 15 +
tests/python/tirx/codegen/test_ptx_ld_st_ops.py | 14 +-
.../tirx/script/test_tirx_script_basic_usage.py | 42 +-
tests/python/tirx/script/test_tirx_script_gpu.py | 106 +-
.../python/tirx/script/test_tirx_script_printer.py | 4 +-
.../script/test_tirx_script_source_locations.py | 6 +-
.../python/tirx/script/test_tirx_script_tensors.py | 33 +-
tests/python/tirx/test_alloc_pool.py | 2 +-
tests/python/tirx/test_jit.py | 8 +-
tests/python/tirx/test_layout.py | 49 +-
tests/python/tirx/test_op.py | 91 +-
tests/python/tirx/test_op_namespace_cleanup.py | 20 +-
tests/python/tirx/test_tile_namespace.py | 13 +-
tests/python/tirx/test_verifier.py | 4 +-
tests/python/tirx/tile/cuda/conftest.py | 2 +-
tests/python/tirx/tile/cuda/copy/test_gmem_smem.py | 779 ----
tests/python/tirx/tile/cuda/copy_async/test_tma.py | 2144 -----------
.../tirx/tile/cuda/elementwise/test_binary.py | 1421 --------
.../tirx/tile/cuda/elementwise/test_unary.py | 2312 ------------
.../python/tirx/tile/cuda/test_add_sub_mul_div.py | 1207 +++++++
.../test_ldgsts.py => test_cp_async.py} | 36 +-
.../test_dsmem.py => test_cp_async_bulk.py} | 74 +-
.../tirx/tile/cuda/test_cp_async_bulk_tensor.py | 2538 +++++++++++++
tests/python/tirx/tile/cuda/test_cvt.py | 1045 ++++++
tests/python/tirx/tile/cuda/test_ewise.py | 662 ++++
.../tirx/tile/cuda/test_explicit_patterns.py | 396 +++
.../tirx/tile/cuda/{elementwise => }/test_fma.py | 243 +-
..._forced_cache.py => test_ld_st_cache_policy.py} | 61 +-
.../python/tirx/tile/cuda/test_ld_st_gmem_smem.py | 354 ++
.../{copy/test_reg.py => test_ld_st_registers.py} | 770 ++--
...st_ld_stmatrix.py => test_ldmatrix_stmatrix.py} | 175 +-
.../test_gemm_mma_m16n8k_.py => test_mma_sync.py} | 347 +-
tests/python/tirx/tile/cuda/test_mov.py | 157 +
tests/python/tirx/tile/cuda/test_pattern_groups.py | 209 ++
tests/python/tirx/tile/cuda/test_pattern_math.py | 111 +
tests/python/tirx/tile/cuda/test_pattern_mma.py | 742 ++++
.../tirx/tile/cuda/test_pattern_transfers.py | 586 +++
.../tile/cuda/{copy_async => }/test_tcgen05_cp.py | 765 ++--
.../test_tcgen05_ldst.py => test_tcgen05_ld_st.py} | 556 ++-
.../test_gemm_async.py => test_tcgen05_mma.py} | 3757 ++++++++------------
.../tirx/tile/cuda/test_tcgen05_mma_codegen.py | 650 ++++
.../tirx/tile/cuda/test_tcgen05_mma_ws_split_n.py | 143 +
tests/python/tirx/tile/test_dispatcher.py | 23 +-
.../tirx/transform/test_transform_lower_tirx.py | 71 +-
162 files changed, 17389 insertions(+), 25389 deletions(-)
diff --git a/docs/tirx/api/ptx.rst b/docs/tirx/api/ptx.rst
index 9bc3f129f2..567b9cc6ee 100644
--- a/docs/tirx/api/ptx.rst
+++ b/docs/tirx/api/ptx.rst
@@ -32,7 +32,7 @@ modifier combination, operand count, and operand types while constructing one
when it needs a new representation, its code-generation support.
This layer is below the layout-aware ``Tx.cuda.tile.*`` / ``Tx.trn.tile.*`` primitives and does not run
-tile-primitive dispatch.
+tile-op dispatch.
Supported instruction families
------------------------------
diff --git a/docs/tirx/api/script.rst b/docs/tirx/api/script.rst
index 781773c017..618cd35abe 100644
--- a/docs/tirx/api/script.rst
+++ b/docs/tirx/api/script.rst
@@ -24,7 +24,7 @@ TIRx kernels use ``tvm.script.tirx`` for the parser and core IR builders::
Tx.alloc_tensor(...)
-Tile primitives and backend-specific namespaces are documented separately in
+Tile ops and backend-specific namespaces are documented separately in
:doc:`tile`, :doc:`cuda`, and :doc:`ptx`. For the relationship between these
authoring layers and TIRx IR, see :ref:`tirx-programming-model`.
@@ -62,9 +62,41 @@ thread-placement and tensor alias policies remain in the TIRx/S-TIR builders.
Allocate a local register fragment whose layout matches a
``tcgen05.{ld,st}`` atom. ``instr_shape`` accepts ``"32x32b"``,
``"16x64b"``, ``"16x128b"``, or ``"16x256b"``. For example, a
- two-CTA Layout-B accumulator and its readback fragment can be allocated as::
+ two-CTA Layout-B accumulator uses this layout, defined outside the kernel::
- C = tmem_pool.alloc_tcgen05_mma_D(
- (64, 128), "float32", M=128, cta_group=2)
+ from tvm.backend.cuda.layout import tmem_mma_operand_layout
+
+ layout_b = tmem_mma_operand_layout(
+ "D", (64, 128), "float32", M=128, cta_group=2)
+
+ Inside the kernel, declare the accumulator over its allocated TMEM address,
+ allocate the register fragment, and prepare matching instruction views::
+
+ C = Tx.cuda.decl_tmem(tmem_address, ty_args=[
+ Tx.Tensor((64, 128), "float32", scope="tmem", layout=layout_b)
+ ])
frag = Tx.alloc_tcgen05_ldst_frag("32x32b", (64, 128), "float32")
- Tx.cuda.tile.tcgen05.ld(frag[:, :], C[:, :], scope="warpgroup")
+ C_words = C.rearrange("(wh r) (wl c) -> (wl wh) r c",
+ wh=2, r=32, wl=2)
+ frag_words = frag.rearrange("(wh r) (wl c) -> (wl wh) r c",
+ wh=2, r=32, wl=2)
+ Tx.cuda.tile.apply(
+ Tx.ptx["tcgen05.ld.sync.aligned.32x32b.x64.b32"],
+ frag_words, C_words, scope="warpgroup")
+ Tx.ptx.tcgen05.wait__ld.sync.aligned()
+
+ ``tmem_address`` must name a live allocation with at least 64 columns.
+ The readback requires a full active warpgroup.
+ Complete the producing MMA's commit/mbarrier wait and
+ the required ``tcgen05.fence::after_thread_sync`` before the load.
+
+ Layout B stores logical ``C[i,j]`` in TMEM lane
+ ``i + 64*(j//64)`` and column ``j%64``. The views above both have shape
+ ``(4,32,64)``: ``warp=2*wl+wh`` chooses the physical 32-lane slab,
+ ``r`` selects its row, and ``c`` selects its 64 columns. Each warp issues
+ one ``.32x32b.x64`` instruction and each lane receives 64 registers.
+ No temporary copy is needed. ``wait::ld`` makes the loaded registers
+ available to subsequent instructions; it does not wait for an earlier MMA.
+
+ See :doc:`../tile_ops/tcgen05_ld_st` for other native shapes,
+ pack/unpack, and a complete store/load example.
diff --git a/docs/tirx/api/tile.rst b/docs/tirx/api/tile.rst
index e72824cc8c..fd1d570f1d 100644
--- a/docs/tirx/api/tile.rst
+++ b/docs/tirx/api/tile.rst
@@ -20,79 +20,27 @@ Tensor Instruction Authoring API
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.
+after ``*``. Tensor variables are expanded to full regions. Trainium unary sources
+may default to the destination; its named workspace arguments are optional.
-See :doc:`../tile_primitives` for memory, scope, descriptor, and synchronization
+See :doc:`../tile_ops` for memory, scope, descriptor, and synchronization
contracts. These calls produce ordinary IR Calls; the raw scalar instruction
namespaces remain ``T.ptx`` and ``T.nki``.
CUDA
----
-.. py:function:: T.cuda.tile.ld(dst, src, *, scope='thread', vec_bits=0, cache=None, l1_evict=None, l2_evict=None, prefetch_size=None)
+.. py:function:: T.cuda.tile.apply(instruction, *operands, scope='thread', **parameters)
-.. py:function:: T.cuda.tile.st(dst, src, *, scope='thread', vec_bits=0, cache=None, l1_evict=None, l2_evict=None, prefetch_size=None)
+ Apply an exact ``T.ptx[opcode]`` to canonical ``group × repeat × atom``
+ tensor views. Keyword parameters are serialized Call arguments. See
+ :doc:`../cuda_tile_patterns` for each family's required fields.
-.. py:function:: T.cuda.tile.ldmatrix(dst, src, *, scope='thread')
+.. py:function:: T.cuda.tile.ewise(function, dst, *inputs, scope='thread')
-.. py:function:: T.cuda.tile.stmatrix(dst, src, *, scope='thread')
-
-.. py:function:: T.cuda.tile.mov(dst, src, *, scope='thread')
-
-.. py:function:: T.cuda.tile.cp_async(dst, src, predicate=-1, *, scope='thread', direct=False, fill_mode='', prefetch_size=-1)
-
-.. py:function:: T.cuda.tile.cp_async_bulk(dst, src, mbar, remote_cta_id, *, scope='thread')
-
-.. 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)
-
-.. 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)
-
-.. 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)
-
-.. py:function:: T.cuda.tile.tcgen05.cp(dst, src, *, scope='thread', shape=None, multicast=None, cta_group=1, decompress=None)
-
-.. py:function:: T.cuda.tile.tcgen05.ld(dst, src, *, scope='thread')
-
-.. py:function:: T.cuda.tile.tcgen05.st(dst, src, *, scope='thread')
-
-.. 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')
-
-.. 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)
-
-.. 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)
-
-.. py:function:: T.cuda.tile.add(dst, lhs, rhs, *, scope='thread', rounding_mode=None)
-
-.. py:function:: T.cuda.tile.sub(dst, lhs, rhs, *, scope='thread', rounding_mode=None)
-
-.. py:function:: T.cuda.tile.mul(dst, lhs, rhs, *, scope='thread', rounding_mode=None)
-
-.. py:function:: T.cuda.tile.div(dst, lhs, rhs, *, scope='thread', rounding_mode=None)
-
-.. py:function:: T.cuda.tile.max(dst, lhs, rhs, *, scope='thread', rounding_mode=None)
-
-.. py:function:: T.cuda.tile.cvt(dst, src=None, *, scope='thread')
-
-.. py:function:: T.cuda.tile.sqrt(dst, src=None, *, scope='thread')
-
-.. py:function:: T.cuda.tile.ex2(dst, src=None, *, scope='thread')
-
-.. py:function:: T.cuda.tile.lg2(dst, src=None, *, scope='thread')
-
-.. py:function:: T.cuda.tile.compose.exp(dst, src=None, *, scope='thread')
-
-.. py:function:: T.cuda.tile.compose.silu(dst, src=None, *, scope='thread')
-
-.. py:function:: T.cuda.tile.fma(dst, src, scale, bias, *, scope='thread')
-
-.. py:function:: T.cuda.tile.compose.exp_with_scale_bias(dst, src, scale, bias, *, scope='thread')
-
-.. py:function:: T.cuda.tile.compose.exp2_with_scale_bias(dst, src, scale, bias, *, scope='thread')
-
-.. py:function:: T.cuda.tile.compose.log2_with_scale_bias(dst, src, scale, bias, *, scope='thread')
-
-.. py:function:: T.cuda.tile.compose.sqrt_with_scale_bias(dst, src, scale, bias, *, scope='thread')
+ Inline a typed scalar TIRx function at every register coordinate. Tensor
+ inputs have the destination shape; scalar inputs broadcast. There is no
+ implicit memory staging or packed arithmetic selection.
TRN
---
diff --git a/docs/tirx/api/tile_dispatch.rst b/docs/tirx/api/tile_dispatch.rst
index 6e3b4d3ef2..c48ba7466c 100644
--- a/docs/tirx/api/tile_dispatch.rst
+++ b/docs/tirx/api/tile_dispatch.rst
@@ -26,8 +26,10 @@ Trainium's ``tile/instructions.py`` files contain the concrete contracts.
Registration installs the fixed argument count, void return type, opaque
side effects, ``FTileOpDispatch``, 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.
+Runtime expressions must be operands. CUDA encodes the opcode and named
+instruction/descriptor parameters as normal Call arguments; only execution
+scope uses a reflected attribute. See :doc:`../cuda_tile_patterns` for the
+CUDA pattern extension workflow.
Lowerers receive a semantic view and ``DispatchContext`` and return a
``tvm.tirx.Function`` body. They can prepare pointers/descriptors and expand
diff --git a/docs/tirx/arch/backends.rst b/docs/tirx/arch/backends.rst
index cf9d171ec8..106e906fe9 100644
--- a/docs/tirx/arch/backends.rst
+++ b/docs/tirx/arch/backends.rst
@@ -29,7 +29,7 @@ Loading and registration
backend can register the following target-owned behavior:
#. TVMScript namespaces such as ``Tx.cuda``, ``Tx.ptx``, and ``Tx.nki``;
-#. tile-primitive dispatch implementations for that target kind;
+#. tile-op dispatch implementations for that target kind;
#. target tags and, when the target owns one, compilation-pipeline entry
points; and
#. code-generation helpers that translate backend calls to target source.
diff --git a/docs/tirx/arch/lowering_pipeline.rst b/docs/tirx/arch/lowering_pipeline.rst
index 48b8e1ac5c..55aa5e9a51 100644
--- a/docs/tirx/arch/lowering_pipeline.rst
+++ b/docs/tirx/arch/lowering_pipeline.rst
@@ -20,7 +20,7 @@ TIRx lowering pipeline
``tvm.compile(mod, target, tir_pipeline="tirx")`` runs an authored TIRx module
through the **tirx pipeline** — an ordered sequence of IR passes that turns the
-high-level constructs you write (tile primitives, ``TileLayout``-typed tensors,
+high-level constructs you write (tile ops, ``TileLayout``-typed tensors,
CUDA index calls) into split **host** + **device** functions, which the CUDA
backend then renders to source. The pipeline is defined in
``python/tvm/tirx/compilation_pipeline.py`` (``tirx_pipeline``); this page walks the
@@ -144,7 +144,7 @@ Inside LowerTIRx
and removes tensor offsets that have been folded into the resulting views.
After ``LowerTIRx`` the module remains a ``tvm.tirx.Function``, but contains no
-tile primitives or ``TileLayout`` indirection, and CUDA indices read hardware coordinates. Later TIRx passes lower the remaining opaque constructs and
+tile ops or ``TileLayout`` indirection, and CUDA indices read hardware coordinates. Later TIRx passes lower the remaining opaque constructs and
the target code generators consume ``tirx::Function`` directly; there is no
conversion to the separate ``tvm.tir`` object model.
@@ -203,7 +203,7 @@ IR snippets across these docs were produced:
target = tvm.target.Target("cuda")
mod = TT.BindTarget(target.with_host("llvm"))(tvm.IRModule({"main": scale}))
- mod = TT.LowerTIRx()(mod) # tile primitives dispatched, layouts applied
+ mod = TT.LowerTIRx()(mod) # tile ops dispatched, layouts applied
print(mod.script()) # inspect the lowered TIRx IR
Or compile the whole module and read the generated CUDA:
diff --git a/docs/tirx/arch/tile_dispatch.rst b/docs/tirx/arch/tile_dispatch.rst
index fbde2a6eec..4e000d1e7f 100644
--- a/docs/tirx/arch/tile_dispatch.rst
+++ b/docs/tirx/arch/tile_dispatch.rst
@@ -73,5 +73,5 @@ 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
+See :doc:`../tile_ops` for instruction contracts and caller-owned
algorithms, and :doc:`../api/tile_dispatch` for registration.
diff --git a/docs/tirx/cuda_tile_patterns.rst b/docs/tirx/cuda_tile_patterns.rst
new file mode 100644
index 0000000000..b8b7588909
--- /dev/null
+++ b/docs/tirx/cuda_tile_patterns.rst
@@ -0,0 +1,514 @@
+.. 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.
+
+Explicit CUDA tile patterns
+================================================================================
+
+Choose an instruction, prepare zero-copy operand views, then apply it::
+
+ T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], registers, memory,
+ scope="thread")
+
+The instruction and modifiers determine the atom contract. The caller supplies
+regions with the canonical ``group × repeat × atom`` coordinates described
+below. Lowering checks the contract and emits one specified PTX
+instruction per repeat coordinate, in row-major repeat order. Lowering
+unrolls static repeats. It does not infer a different opcode, width, matrix
+count, transpose mode, or coordinate permutation.
+
+CUDA exposes two tile operations: ``apply`` for an explicitly selected PTX
+instruction and ``ewise`` for an explicitly supplied scalar TIRx function.
+The raw ``T.ptx`` namespace remains available for scalar/address operands.
+There is no automatic instruction selection or fallback to the retired CUDA
+dispatch algorithms. Trainium's API and lowering are unchanged.
+
+This page is the contract reference. The :doc:`tile_ops` instruction
+guides provide complete kernels, caller views, synchronization, and numerical
+checks, organized by PTX instruction family.
+
+Group, repeat, and atom axes
+--------------------------------------------------------------------------------
+
+A thread instruction at ``scope="thread"`` has no group prefix. At
+``scope="warp"`` it has a leading lane group axis; at ``scope="warpgroup"``
+its groups are ``(warp_in_group, lane)``; at ``scope="cta"`` they are
+``(warp_in_cta, lane)``. Extents come from the active execution scope, including
+its explicitly guarded subrange. Warp-collective instructions consume the
+lane axis inside their atom instead: their group prefix is empty at warp
+scope, one warp axis at warpgroup/CTA scope. All 32 lanes must be active.
+Cluster scope requires an explicit launch cluster and prepends its CTA axes
+in z/y/x order, omitting undeclared leading axes. A one-dimensional cluster
+therefore has ``(cta_x, warp, lane)`` groups for thread instructions.
+Local tensor ownership uses ``cbx/cby/cbz`` for these CTA coordinates.
+
+For example, 64 threads, three 128b loads per thread::
+
+ # A and R both have shape (2, 32, 3, 4).
+ registers = TileLayout(S[(2,32,3,4):(1 @ warpid,1 @ laneid,8,2)])
+ T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], R, A, scope="cta")
+
+The first two coordinates bind hardware IDs, the third generates three
+unrolled calls per thread, and the last selects four words of one instruction.
+The separated register slots are legal; the four memory words must be aligned
+and contiguous. A changed opcode or modifier is always written at the call site.
+
+Views such as ``sub``, ``permute``, ``view(dtype)``, and ``rearrange`` prepare
+these coordinates without moving data. A failed contract requires a different
+explicit view, caller-written preparation, or a different explicit instruction.
+It never invokes a permutation/factorization solver.
+
+128-bit global load
+--------------------------------------------------------------------------------
+
+``ld.global.v4.b32``, ``scope="thread"``:
+
+* Both regions have shape ``(*repeat, 4)`` and scalar 32-bit elements.
+* The source is global memory. For each repeat, its four words have addresses
+ ``base + 0, base + 1, base + 2, base + 3`` (in words); the base is 16B aligned.
+* The destination is four distinct registers of the executing thread. Their
+ storage offsets may be separated or permuted. No thread axes are accepted.
+* The operation copies raw bits; it performs no numeric conversion.
+
+The earlier example keeps its original float32 tensor layout::
+
+ original = TileLayout(S[(15, 2, 7, 2):(4, 1, 60, 2)])
+ A_std = A.rearrange("i b j d -> j i (d b)")
+ R_std = R.rearrange("i b j d -> j i (d b)")
+ T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], R_std, A_std,
+ scope="thread")
+
+``A_std`` has shape ``(7,15,4)`` and addresses ``60*j + 4*i + word``.
+The repeat axes are ``j`` and ``i``: this emits 7*15 = 105 vector loads.
+The destination may instead use original strides ``(8,2,120,4)``; each
+instruction then writes slots separated by two, without a remap copy.
+
+ldmatrix with interleaved rows and separated registers
+--------------------------------------------------------------------------------
+
+``ldmatrix.sync.aligned.m8n8.x2.shared.b16``, ``scope="warp"``:
+
+* Both regions have shape ``(*repeat, 2, 8, 4)``. The suffix is
+ ``(matrix, row, word)``.
+* Each word is a 32-bit container holding two adjacent 16-bit matrix values.
+ The pattern accepts float32/int32/uint32 containers as raw bits.
+* In shared memory, each ``(matrix,row)`` is four consecutive words with a
+ 16B-aligned start. Different rows can be separated; they need not form a
+ contiguous 8x8 allocation.
+* Destination element ``(matrix,row,word)`` belongs to lane ``4*row+word``.
+ Each lane receives two words; the register slots are selected by its layout.
+* All 32 lanes must execute the instruction. Shared initialization and any
+ preceding synchronization belong to the caller.
+
+Use the original shared shape ``(16,4,2)`` and strides ``(8,1,128)``::
+
+ S_std = SM.rearrange("(r m) w k -> k m r w", r=8, m=2)
+ registers = TileLayout(S[(2,2,8,4):(8,2,4 @ laneid,1 @ laneid)])
+ # R is declared over 12 caller-allocated words per thread, using registers.
+ T.cuda.tile.apply(T.ptx["ldmatrix.sync.aligned.m8n8.x2.shared.b16"],
+ R, S_std, scope="warp")
+
+``S_std[k,m,r,w] = SM[2*r+m,w,k]``. Its strides are ``(128,8,16,1)``:
+first matrix = even original rows; second matrix = odd original rows.
+For each ``k``, lanes 0..7 provide the eight matrix-0 row addresses, lanes
+8..15 provide the eight matrix-1 row addresses, and lanes 16..31 duplicate
+those valid addresses. These address providers are distinct from result owners.
+
+For lane ``l``, the complete output mapping is::
+
+ R_raw[0] = SM[2*(l//4), l%4, 0]
+ R_raw[2] = SM[2*(l//4)+1, l%4, 0]
+ R_raw[8] = SM[2*(l//4), l%4, 1]
+ R_raw[10] = SM[2*(l//4)+1, l%4, 1]
+
+There are two repeat coordinates, ``k=0,1``, so this is exactly two ``.x2``
+instructions. Slots 1,3,4,5,6,7,9,11 remain untouched. Lowering binds a
+physical view of the caller's registers and passes these slots directly as
+PTX output operands; it allocates no intermediate tile.
+
+Memory and register instructions
+--------------------------------------------------------------------------------
+
+``ld`` / ``st`` use ``(*group,*repeat,word)``, where ``word`` has the vector
+length selected in the opcode (one for scalar instructions). Both operands use
+the written PTX storage width; use an explicit dtype view for packed words.
+Global, shared, and local address spaces are explicit. Cache operators,
+eviction policies, prefetch widths, and vector widths stay in the PTX spelling.
+``pred`` and an opcode's optional ``cache_policy`` are ordinary runtime operands.
+A predicated load can request ``preserve_dst=True``.
+
+``ldmatrix`` / ``stmatrix`` support m8n8.b16 with x1/x2/x4 and optional
+``.trans``. The canonical suffix is ``(matrix,8,4)`` in b32 containers.
+For ``.trans``, the register image is the transpose of each shared b16 matrix;
+its lane ownership is still ``4*row+word``. The shared address provider and
+register owner are different maps.
+
+Register instructions (mov, add/sub/mul/div, max/min, fma, cvt, sqrt, ex2, lg2,
+rcp, abs, neg) end each tensor operand with the native register-group length.
+A scalar arithmetic operand therefore uses suffix ``(1,)``. A read-only scalar
+expression broadcasts over repeats. Packed ``f16x2`` / ``bf16x2`` / ``f32x2``
+use b32 / b32 / b64 storage views, respectively. Packing is explicit; lowering
+never chooses packed arithmetic because a layout happens to permit it.
+
+Scalar compositions
+--------------------------------------------------------------------------------
+
+Define exactly the scalar expression to apply::
+
+ @T.function(private=True)
+ def silu_affine(x: T.float32, scale: T.float32, bias: T.float32) -> T.float32:
+ y = x / (T.float32(1) + T.exp(-x))
+ return y * scale + bias
+
+ T.cuda.tile.ewise(silu_affine, R, R, T.float32(2), T.float32(1), scope="cta")
+
+The atom is one scalar element. Tensor operands share the group/repeat shape;
+scalar inputs broadcast. A tensor broadcast is an explicit zero-stride view.
+Types must match the function parameters and result. The scalar function may
+contain scalar bindings and returns, including conditional returns. Its typed
+expression is serialized as ``LambdaExpr`` and inlined at each register store.
+No shared/global staging, synchronization, or packed-vector selection occurs.
+Shared-memory compositions explicitly load into registers, run ``ewise``, and
+store back, with caller-owned synchronization.
+
+Asynchronous copy
+--------------------------------------------------------------------------------
+
+``cp.async.{ca,cg}.shared.global`` uses ``(*group,*repeat,element)``.
+``chunk_bytes=4/8/16`` fixes the last axis's byte count; .cg requires 16 bytes.
+The corresponding source and destination chunks are contiguous and naturally
+aligned. ``src_size`` or ``ignore_src`` uses the hardware's zero-fill operand;
+``pred`` gates execution and has different semantics. Commit and wait are
+caller instructions.
+
+``cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes``
+uses the same suffix with positive ``chunk_bytes`` divisible by 16. ``mbar``
+and ``remote_cta_id`` are required. Source, destination, and barrier arguments
+name local shared addresses; lowering maps the destination and barrier with
+``mapa`` before issuing the requested bulk copy. Election, barrier initialization,
+transaction accounting, and cluster synchronization belong to the caller.
+
+TMA
+--------------------------------------------------------------------------------
+
+Select the rank, direction, tile/gather4/scatter4 mode, reduction operator,
+multicast and cache qualifiers in the PTX opcode. The global operand is the
+rectangular tensor view described by the TensorMap. Without interleave its
+rank is exactly the instruction rank; shape, strides, starting region and
+dtype define the map. Alternatively, pass an ordinary address expression for
+a descriptor encoded by the caller. Its fields must agree with the pattern's
+box, format and swizzle; the dispatcher does not inspect opaque descriptors.
+
+Required keyword operands are ``box`` and ``coords``, both in outer-to-inner
+logical axis order. The shared operand is ``(*group,*repeat,*box_image)``.
+Every atom's unswizzled origin is 128-byte aligned, including unswizzled boxes.
+Swizzled allocation bases are aligned to the full swizzle period; a subregion
+uses its unswizzled origin so that the instruction applies its phase once.
+``coord_steps`` has one coordinate vector per group/repeat axis. Lowering adds
+these explicit steps to ``coords``, encodes the descriptor on the host, and
+issues one instruction at each repeat. It does not fold dimensions, promote
+dtypes, infer a box, prefetch a descriptor, select a source, or elect a thread.
+Use ordinary control flow, or select explicitly encoded descriptor addresses,
+to choose between sources. Prefetch is an explicit ``T.ptx.prefetch.tensormap``.
+
+Descriptor operands are ``format`` (storage dtype by default, or ``tf32``),
+``swizzle=0/1/2/3``, ``l2_promotion=0`` and ``oob_fill=0``. ``element_strides``
+has one entry in 1..8 per descriptor dimension; without interleave the inner
+entry is one. The non-interleaved shared image has extents
+``ceil(box[i]/element_strides[i])`` and compact storage before swizzling.
+Symbolic global dimensions and outer strides pass to runtime descriptor
+encoding; statically invalid values are rejected when lowering. A load requires
+``mbar``; a multicast also requires ``cta_mask``; a cache-hint form requires
+``cache_policy``. All are Call arguments and participate in substitution.
+
+Gather4/scatter4 use a 2D descriptor, ``box=(1,width)``, a shared atom of
+``(4,width)``, and four explicit row coordinates in ``gather4``. Dynamic row
+coordinates are allowed. No row-selection algorithm is hidden in lowering.
+
+``format="float4_e2m1fn"`` selects CUDA's ``16U4_ALIGN16B`` representation.
+Its inner global dimension is a multiple of 128, its inner box is 128, and
+the global base and outer byte strides are 32-byte aligned. Each sixteen
+nibbles occupy the low eight bytes of a sixteen-byte shared block. The other
+eight bytes are undefined padding. The mbarrier counts **payload bytes**,
+not padding. This representation supports loads with swizzle 0 or 3; gather4
+uses the same packed shared image.
+
+For ``interleave=1`` (16B) or ``2`` (32B), an opcode of rank 3..5 takes a
+global view ``(N,C_slices,*spatial,Ci)``. ``Ci`` is the number of scalar
+elements in one 16B/32B slice. The descriptor omits this last, implicit axis;
+``box``, ``coords`` and ``coord_steps`` have the opcode rank. The W stride
+is exactly one slice; Ci is contiguous. Other strides are explicit and can
+include padding. Mode 2 requires swizzle 1 and 32-byte global alignment.
+
+One interleaved tiled issue visits **one C slice**, selected by ``coords[1]``.
+The C box and traversal-stride fields do not iterate additional slices. Its
+shared image is ``(ceil(boxN/stepN),1,*ceil(boxSpatial/stepSpatial),Ci)``
+before swizzling. To visit several channel slices, add a repeat axis and a
+coordinate step in C. For example, rank-3 uint16 with interleave 2 uses
+``Ci=16``; ``box=(1,1,8)`` copies a ``(1,1,8,16)`` shared atom. A shared
+operand of shape ``(3,1,1,8,16)`` and ``coord_steps=((0,1,0),)`` issues three
+copies for three successive channel slices. Gather4/scatter4 do not interleave.
+
+TMEM transfers
+--------------------------------------------------------------------------------
+
+``tcgen05.ld/st.sync.aligned.<shape>.xN.b32`` use suffix ``(row,column)``
+in b32 storage on both sides. Each atom is a warp instruction. The four
+supported shapes and register images (lane ``l``, register ``r``) are:
+
+.. list-table::
+ :header-rows: 1
+
+ * - Shape
+ - Atom shape
+ - Register coordinate (row, column)
+ * - 32x32b
+ - (32,N)
+ - (l,r)
+ * - 16x64b
+ - (16,2N)
+ - (l//4+8*(l%2), l//2%2+2*r)
+ * - 16x128b
+ - (16,4N)
+ - (l//4+8*(r%2), l%4+4*(r//2))
+ * - 16x256b
+ - (16,8N)
+ - (l//4+8*(r//2%2), r%2+2*(l%4)+8*(r//4))
+
+TMEM coordinates are consecutive rows/columns in the executing warp's 32-lane
+slab. Explicit group axes distribute atoms across warps. Register storage can
+be strided.
+
+``.pack::16b`` / ``.unpack::16b`` keep the same b32 register image. Their TMEM
+operand instead has suffix ``(row,column,half)`` with ``half=2`` and a 16-bit
+dtype. The two halves of register word ``(row,column)`` correspond to the low
+16 bits of TMEM cells ``(row,2*column)`` and ``(row,2*column+1)``. In 16-bit
+layout units this means TCol strides four for ``column`` and two for ``half``.
+This is different from reinterpreting two adjacent 16-bit elements in one
+dense b32 cell. The hardware performs the packing; lowering adds no conversion.
+
+``.16x32bx2`` has register suffix ``(access,row,column)=(2,16,N)``. Lane ``l``
+owns ``(l//16,l%16,r)`` for register ``r``. Its two TMEM access regions are
+explicit operands, each with suffix ``(16,N)`` (or ``(16,N,2)`` for packing)::
+
+ T.cuda.tile.apply(
+ T.ptx["tcgen05.ld.sync.aligned.16x32bx2.x4.b32"],
+ registers, memory[:, :16, :4], memory[:, 16:32, 8:12],
+ imm_half_splitoff=(16 << 16) + 8, scope="warpgroup",
+ )
+
+For a store the operand order is ``memory_low, memory_high, registers``.
+Both regions share the same backing TMEM tensor. The second region's encoded
+base address must equal the first plus ``imm_half_splitoff``. There is no
+implicit sixteen-row offset for the second access. The immediate can select
+different rows, columns, or both, within the executing warp's TMEM partition.
+Each repeat emits exactly one instruction, including both half-accesses.
+Wait, fence, allocation and deallocation instructions remain at the call site.
+
+``tcgen05.cp`` uses ``(row,b32_word)`` according to its explicit shape. Multicast
+replicas must exactly match the written qualifier: warpx4 adds four replicas
+32 lanes apart, 02_13 adds two 64 lanes apart, 01_23 adds two 32 lanes apart.
+4x256b maps its four rows to TLanes 0/32/64/96; the other shapes follow their
+named contiguous slabs (01_23 uses slabs 0 and 2).
+
+Native decompression uses the same b32-word footprint and explicit qualifiers:
+``.b8x16.b4x16_p64`` or ``.b8x16.b6x16_p32``. Each 16-byte source block holds
+sixteen packed 4-bit or 6-bit values followed by eight or four bytes of padding,
+respectively. The destination holds sixteen bytes: the 4-bit values gain two
+low zero bits, while the 6-bit values gain two high zero bits. A 256-bit row
+contains two such blocks. These are raw packed-word views, not sixteen
+ordinary byte elements on the source side. Each repeat still emits one cp;
+there is no software unpacking or implicit dtype conversion.
+
+Shared matrix descriptors
+--------------------------------------------------------------------------------
+
+``tcgen05.cp`` takes ``ldo_bytes``, ``sdo_bytes`` and ``swizzle``.
+MMA takes ``a_descriptor=(ldo_bytes,sdo_bytes,swizzle)`` and the analogous B
+operand. These are explicit hardware descriptor fields, not tuning hints.
+They are mechanically encoded, with no layout search.
+
+For ``tcgen05.cp``, an 8-row group has row byte pitch ``16 << swizzle``.
+Consecutive row groups are separated by ``sdo_bytes``. Without swizzling,
+16-byte chunks within a row are separated by ``ldo_bytes``. With swizzling,
+words are contiguous before swizzling and hardware ignores LDO; changing
+that field does not make gapped words legal.
+
+For MMA, ``row`` means M/N in K-major order and K in MN-major order. The
+contiguous coordinate advances in units of the operand dtype. Its low bits
+occupy a ``16 << swizzle`` byte row for modes 0..3. The offsets for each subsequent 8-row
+group and each subsequent contiguous chunk are:
+
+.. list-table:: MMA descriptor strides in bytes
+ :header-rows: 1
+
+ * - Mode
+ - Next 8-row group
+ - Next contiguous chunk
+ * - K-major
+ - SDO
+ - LDO (unswizzled; swizzled instruction atoms fit within one row)
+ * - MN-major, swizzled
+ - SDO
+ - LDO
+ * - MN-major, unswizzled
+ - LDO
+ - SDO
+
+MMA mode ``swizzle=4`` is the distinct 128B span / 32B atomicity format
+for 32-bit MN-major operands. It has 128-byte rows and groups **four** K rows
+per SDO step. Consecutive 128-byte M/N chunks use LDO. Its byte permutation
+is ``offset ^ ((offset >> 2) & 96)`` with a 512-byte period, unlike mode 3's
+``offset ^ ((offset >> 3) & 112)``. The mode is explicit in the descriptor;
+it is never selected from a tensor layout. The same TF32 format in K-major
+order uses modes 0..3.
+
+This distinction follows PTX's canonical shared layouts; in particular,
+unswizzled MN-major exchanges the descriptor fields' roles. The selected
+swizzle permutes the resulting physical address. Patterns check this walk
+against the supplied view, including its start and element offset.
+Descriptor origins must be 16B aligned; a swizzled backing allocation must
+align to its full period (256/512/1024 bytes). The descriptor encodes the
+unswizzled start address; hardware applies the permutation to that start and
+the subsequent accesses. In particular, cp may start in the middle of a
+swizzle atom if its entire address walk matches the requested region.
+Descriptor base-offset and absolute-LDO modes are not inferred.
+Allocate a complete descriptor storage block even for a
+partial logical atom; in particular, a 4x256b copy still needs an 8-row shared
+backing block. The logical region may contain only the four copied rows.
+
+Warp MMA
+--------------------------------------------------------------------------------
+
+``mma.sync.aligned.m16n8k{8,16}.row.col.f32.{f16,bf16}.{f16,bf16}.f32``
+uses these canonical operands, after optional group and batch prefixes:
+
+* D/C: ``(M_repeat,N_repeat,16,8)`` in f32.
+* A: ``(M_repeat,K_repeat,16,K/2)`` in b32 words packing two K elements.
+* B: ``(N_repeat,K_repeat,8,K/2)`` in the same representation.
+
+Lane ``l`` owns A word ``(l//4+8*(r%2), l%4+4*(r//2))``, B word
+``(l//4,l%4+4*r)``, and D/C element ``(l//4+8*(r//2),2*(l%4)+r%2)``.
+Repeat order is batch, M, N, K. The first K step uses C (tensor or scalar
+broadcast); later steps use D. Alpha/beta preparation and logical transpose
+views are caller code.
+
+tcgen05 MMA
+--------------------------------------------------------------------------------
+
+Select dense, .ws, or block_scale and CTA group in the opcode. Specify
+``instruction_shape=(M,N,K)``; M is the CTA-group total. Per-CTA canonical
+suffixes after group/batch axes are:
+
+* D: ``(M_repeat,N_repeat,M/cta_group,N)``.
+* A: ``(M_repeat,K_repeat,M/cta_group,K)``.
+* B: ``(N_repeat,K_repeat,N/cta_group,K)``.
+* Block scale adds SFA/SFB with the corresponding A/B repeat axes and atom
+ ``(M/cta_group,scale_vec)`` / ``(N,scale_vec)``. SFB contains the
+ full instruction N range in each CTA, including two-CTA operations.
+
+A may be shared or TMEM; B is shared; D and scales are TMEM. D uses
+16-bit or 32-bit accumulation storage as required by the selected format.
+Each accumulator occupies one b32 TMEM cell; f16 occupies its low half
+(``TCol`` stride two in f16 element units). A-in-TMEM remains packed. ``a_major`` and
+``b_major`` explicitly select ``"k"`` (default) or ``"mn"``. Storage formats
+come from the tensor dtype unless ``a_format`` / ``b_format`` is supplied
+(for example, ``tf32`` over f32 storage). ``accum=False`` zeros the first K
+step's accumulator; later K steps accumulate. Negation, saturation, output-lane
+masks, .ws zero-column mask, and ``pred`` are explicit operands.
+
+The MMA TMEM datapath follows ``tmem_mma_operand_layout`` in
+``tvm.backend.cuda.layout`` for the chosen instruction shape and mode.
+Weight-stationary M=128 uses Layout D, like non-.ws M=128.
+Weight-stationary M=64 A-in-TMEM includes an explicit leading bank axis of
+size two. Scale rows use lane ``row%32``, byte-column ``4*(row//32)+sf`` and
+four 32-lane replicas. The scale ID is encoded from each scale tensor's own
+byte-column origin. Preparation, CTA election, fences, commit and wait remain
+caller-owned. Layouts never select .ws or a different M/N/K instruction shape.
+
+Splitting an M64 .ws N128 operation into two N64 instructions changes the
+bank boundary inside each instruction from 64 columns to 32. The caller must
+account for this: ``test_tcgen05_mma_ws_split_n.py`` preserves the original
+banked A and packed C, explicitly stages B in each instruction's column order,
+and permutes the C view. Simply cutting B and C into consecutive N64 logical
+slices does not preserve that operation.
+
+Choosing instructions
+--------------------------------------------------------------------------------
+
+Choose each instruction and its canonical view at the call site:
+
+.. list-table:: Instruction choices
+ :header-rows: 1
+ :widths: 25 75
+
+ * - Operation
+ - Instruction or composition
+ * - ld/st, register copy
+ - ``apply`` with a written ld/st width or mov opcode.
+ * - ldmatrix/stmatrix
+ - ``apply`` with matrix count and transpose qualifier in the opcode.
+ * - Scalar/packed arithmetic and casts
+ - ``apply`` with the exact arithmetic/cvt form and native register representation.
+ * - exp, silu, scale/bias combinations
+ - ``ewise`` with the caller's scalar TIRx function and explicit constants.
+ * - Arithmetic on shared memory
+ - Explicit shared load, register ``apply``/``ewise``, then shared store.
+ * - cp.async, DSMEM, TMA
+ - ``apply`` with explicit chunk/box/coordinates and runtime completion operands.
+ * - tcgen05.ld/st/cp
+ - ``apply`` with explicit shape, repeat count and multicast qualifiers.
+ * - mma/tcgen05.mma
+ - ``apply`` with explicit matrix mode, canonical M/N/K repeats and descriptors.
+
+Keep barriers, thread election, packing, scale-factor preparation and pipeline
+scheduling in caller code. Changing a view changes only logical coordinates;
+when the desired physical arrangement requires a copy, write that copy
+explicitly. Unsupported hardware modes remain available through scalar
+``T.ptx`` calls and can acquire their own documented pattern separately.
+
+Implementation and extension
+--------------------------------------------------------------------------------
+
+``python/tvm/backend/cuda/tile/patterns/`` has one module per fixed instruction
+family. ``common.py`` evaluates canonical coordinates, validates group/repeat
+shapes and physical storage, and emits repeat loops. ``descriptors.py`` only
+encodes/checks explicit hardware fields. There are no imports of retired
+CUDA dispatch implementations.
+
+To add a pattern:
+
+1. Define its canonical operand suffixes, ownership and physical address walk,
+ all semantic qualifiers, and required caller preparation in this document.
+2. Resolve the exact instruction through ``ptx.instruction_forms``. Reuse the
+ existing PTX table for legality and native operand emission.
+3. Use ``Context`` and ``Region`` to check the supplied views at those fixed
+ coordinates. Reject mismatches; do not search for another layout or opcode.
+4. Emit one selected instruction inside ``Context.finish``. Put algorithms,
+ packing, synchronization and thread election at the call site.
+5. Add positive/negative contracts, script/JSON roundtrips, exact PTX spelling
+ and issue-count checks, and GPU results reconstructed independently from
+ the documented hardware mapping.
+
+The common IR is ``Call(opcode, operands, parameters)`` with only execution
+scope in ``ScopeAttrs``. Runtime and static descriptor parameters both live in
+ordinary arguments. ``ewise`` stores its typed scalar body as ``LambdaExpr``.
+No process-local Python callable survives serialization.
+
+Region extents are positive static integers. Starts may be symbolic when their
+bounds and alignment can be proved. General physical-equivalence discovery is
+outside the contract. Gapped allocation covers the layout's physical span,
+including the address range introduced by a composed swizzle.
diff --git a/docs/tirx/get_started/index.rst b/docs/tirx/get_started/index.rst
index dc31e59cea..eb7351a845 100644
--- a/docs/tirx/get_started/index.rst
+++ b/docs/tirx/get_started/index.rst
@@ -29,4 +29,4 @@ launch, and generated-source inspection.
../native_basics/cuda/first_kernel
Continue with the :doc:`programming guide <../programming/index>` for the core
-language, CUDA intrinsics, layouts, and tile primitives.
+language, CUDA intrinsics, layouts, and tile ops.
diff --git a/docs/tirx/layout.rst b/docs/tirx/layout.rst
index 8142879926..3f28a1edcb 100644
--- a/docs/tirx/layout.rst
+++ b/docs/tirx/layout.rst
@@ -327,8 +327,8 @@ their 64 logical rows according to
Thus C/F with ``sub_slab=0`` and ``sub_slab=1`` can describe lower- and
upper-half aliases of the same 128-row Layout D allocation. The
-``tcgen05_ldst`` copy dispatch recognizes the layout and emits the matching
-``row=0`` or ``row=16`` instruction. Layout D already occupies both halves,
+explicit TMEM pattern validates the selected row slab and emits its
+``row=0`` or ``row=16`` address. Layout D already occupies both halves,
so a nonzero ``sub_slab`` is rejected.
Layout B is the per-CTA accumulator placement for an M=64
@@ -345,17 +345,14 @@ across the two 64-lane halves:
Thus the tile occupies all 128 lanes and ``N/2`` tensor-memory columns, the
same physical footprint as Layout D ``(128, N/2)``. ``N`` must be even and
-Layout B does not accept ``sub_slab``. Its register image uses the existing
-fragment API:
-
-.. code-block:: python
-
- frag = Tx.alloc_tcgen05_ldst_frag("32x32b", (64, N), "float32")
- 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
-all 128 lanes; each thread owns ``N/2`` contiguous fp32 registers.
+Layout B does not accept ``sub_slab``. Its logical
+``(64, N)`` register fragment has ``N/2`` values per thread. To load it,
+expose the physical slab as ``(4, 32, N/2)``: the leading axis selects a
+warp, and the remaining axes describe rows and words. Choose an explicit
+``tcgen05.ld.32x32b.xN`` width and split the word axis into repetitions of
+that width. The matching register and TMEM patterns are documented in
+:doc:`cuda_tile_patterns`; the caller issues ``tcgen05.wait::ld`` before
+using the loaded registers.
Beyond GPU registers
~~~~~~~~~~~~~~~~~~~~~~
@@ -486,5 +483,5 @@ Design rationale
itself, so an expression is unambiguous without external context. For instance
``1@tid`` (block-wide thread id) and ``1@tid_in_wg`` (thread id within a
warpgroup) are distinct rather than a generic ``t`` whose meaning depends on the
- definition site. Legality and feasibility checks are left to tile primitive
+ definition site. Legality and feasibility checks are left to tile op
dispatch.
diff --git a/docs/tirx/native_basics.rst b/docs/tirx/native_basics.rst
index 2c02a12c7c..293050823b 100644
--- a/docs/tirx/native_basics.rst
+++ b/docs/tirx/native_basics.rst
@@ -32,8 +32,8 @@ A native-level TIRx kernel reads like a structured device kernel: you place
threads yourself, allocate shared/per-thread local tensors, write loops and barriers, and
call device intrinsics directly. You explicitly choose the orchestration and
layouts; standard lowering still dispatches primitives, applies layouts, and
-vectorizes or unrolls marked loops. This is the foundation the tile primitives
-(:doc:`tile_primitives`) are built on; everything here is what those primitives
+vectorizes or unrolls marked loops. This is the foundation the tile ops
+(:doc:`tile_ops`) are built on; everything here is what those ops
ultimately lower to, so it is also where you go when a hardware feature does not
have a primitive yet.
diff --git a/docs/tirx/native_basics/cuda/compiling.rst b/docs/tirx/native_basics/cuda/compiling.rst
index 73c3ef18e9..af6b7f948c 100644
--- a/docs/tirx/native_basics/cuda/compiling.rst
+++ b/docs/tirx/native_basics/cuda/compiling.rst
@@ -211,11 +211,11 @@ Rung 2 in full — a 256-element block sum via a shared-memory tree reduction
exe(a, out) # out[0] == 32640.0
The full tile-level GEMM/attention ladder (sync → TMA → warp specialization →
-2-CTA cluster) is built on top of these and the dispatchable tile primitives in
-:doc:`../../tile_primitives`.
+2-CTA cluster) is built on top of these and the dispatchable tile ops in
+:doc:`../../tile_ops`.
Next steps
----------
- :doc:`../../layout` — how tensors map to physical resources (``TileLayout``).
-- :doc:`../../tile_primitives` — the dispatchable ops these native idioms lower to.
+- :doc:`../../tile_ops` — the dispatchable ops these native idioms lower to.
diff --git a/docs/tirx/native_basics/cuda/tensors.rst b/docs/tirx/native_basics/cuda/tensors.rst
index 2de8e146c0..ac02641b79 100644
--- a/docs/tirx/native_basics/cuda/tensors.rst
+++ b/docs/tirx/native_basics/cuda/tensors.rst
@@ -390,7 +390,7 @@ tensor as a view at a column offset, and one warp frees it at the end:
You add any column offsets to ``addr[0]`` and manage the ``tmem_layout`` (the
matching datapath A–G layout) yourself.
-See the :doc:`../../tile_primitives` walkthroughs for full examples.
+See the :doc:`../../tile_ops` walkthroughs for full examples.
Tensor variable APIs
--------------------
@@ -482,7 +482,7 @@ sees the 256-element tensor as ``64×4``; ``A.permute(1, 0)`` transposes the axe
At_ptr[(j * 4) + i] // permute: swapped strides
**Per-thread view — ``local``.** Decomposes a thread-axis ``local`` layout into the
-calling thread's storage bundle (used pervasively by the tile primitives).
+calling thread's storage bundle (used pervasively by the tile ops).
Both forms expose the raw physical storage span by default, including layout
gaps and offsets: ``R.local()`` infers a flat 1-D span, while
``R.local(d0, d1, ...)`` is a row-major reshape whose product must equal that
diff --git a/docs/tirx/native_basics/cuda/threads_sync.rst b/docs/tirx/native_basics/cuda/threads_sync.rst
index a889b92d3f..0dd9826383 100644
--- a/docs/tirx/native_basics/cuda/threads_sync.rst
+++ b/docs/tirx/native_basics/cuda/threads_sync.rst
@@ -18,7 +18,7 @@
CUDA C++/PTX intrinsics
=======================
-When no tile primitive covers what you need, two escape hatches reach the hardware
+When no tile op covers what you need, two escape hatches reach the hardware
directly: **call a backend intrinsic** (the ``Tx.cuda.*`` / ``Tx.ptx.*`` namespaces
from ``tvm.backend.cuda``), or **inline raw CUDA** source.
diff --git a/docs/tirx/overview.rst b/docs/tirx/overview.rst
index f4930a2579..1484b93e53 100644
--- a/docs/tirx/overview.rst
+++ b/docs/tirx/overview.rst
@@ -37,7 +37,7 @@ is built on top of the TVM compiler infrastructure.
A TIRx kernel keeps orchestration — pipeline state, roles, synchronization,
and backend intrinsics — in hardware-native source, while execution scope,
- tensor layout, and tile primitive dispatch expose the recurring tile
+ tensor layout, and tile op dispatch expose the recurring tile
structure to the compiler. Higher-level systems build on the same structure.
Design Philosophy
@@ -58,15 +58,15 @@ kernel that frequently require expert control — pipeline structure,
synchronization, role assignment, memory placement, and backend intrinsics — in
hardware-native source code. At the same time, it exposes the recurring tile-level
structure to the compiler through three constructs: *execution scope*, *tensor
-layout*, and *tile primitive dispatch*. Orchestration stays in hardware-native
+layout*, and *tile op dispatch*. Orchestration stays in hardware-native
source code, while the recurring tile-level structure becomes visible to the
compiler. Hardware-native control is powerful but costs engineering effort;
-exposing recurring operations as tile primitives relieves this, since authors
+exposing recurring operations as tile ops relieves this, since authors
reuse a dispatched implementation instead of re-writing the same operation for
each kernel and backend.
The result is a DSL that can **grow with the hardware**. A new feature can first
-be used directly as a native intrinsic, and later become a reusable primitive
+be used directly as a native intrinsic, and later become a reusable tile op
once the pattern stabilizes. This is the core design philosophy behind TIRx:
keep the foundation small and explicit, and let the backend library evolve as
new accelerator generations arrive.
@@ -78,7 +78,7 @@ The Programming Model
A TIRx program reads as a structured native kernel: loops, branches, tensors,
synchronization, pipeline state, backend intrinsics, and hardware roles are
-written directly. Tile primitives appear exactly where a repeated hardware-level
+written directly. Tile ops appear exactly where a repeated hardware-level
operation should become reusable and dispatchable.
Authoring layers and IR
@@ -101,7 +101,7 @@ program, but enter the compiler at different levels:
* - Core language
- ``Tx.*``
- Creates TIRx statements, expressions, tensors, and control flow.
- * - Tile primitives
+ * - Tile ops
- ``Tx.cuda.tile.*`` / ``Tx.trn.tile.*``
- Creates an ordinary void ``Call`` that is replaced by its backend instruction
implementation during tile dispatch.
@@ -128,13 +128,13 @@ Execution scope
~~~~~~~~~~~~~~~~~
Execution scope describes both the active participants and the logical scope of
-a primitive invocation. Control flow such as ``if wg_id == ...``,
+a tile op invocation. Control flow such as ``if wg_id == ...``,
``warp_id == ...``, or ``cbx == ...`` selects which hardware roles enter a
region, while predicates such as ``Tx.cuda.elect_sync()`` further select the
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,
+Tile ops declare their execution scope explicitly. For example,
``scope="warpgroup"`` selects warpgroup cooperation; ``scope="thread"``
is the default on backend tensor instructions.
@@ -145,10 +145,10 @@ Tensor layout, with a storage-first interface, describes how logical
tensors map to physical resources. A tile may live in global memory, shared
memory, per-thread local storage, tensor memory, or accelerator SRAM.
Users declare where each tile lives and how its elements are spread across
-lanes, warps, and per-thread storage elements; tile primitive dispatch reads
-those declarations to choose an implementation. A layout is a storage
-description, not a loop-transformation utility: users may construct a tile's
-layout, but never use layouts to transform loops.
+lanes, warps, and per-thread storage elements; tile op lowering uses
+those declarations to validate operand mappings and generate addresses. A layout
+is a storage description, not a loop-transformation utility: users may construct
+a tile's layout, but never use layouts to transform loops.
.. seealso::
@@ -156,31 +156,33 @@ layout, but never use layouts to transform loops.
design rationale — is described in :doc:`layout`, which also has an
interactive explorer.
-Tile primitive dispatch
-~~~~~~~~~~~~~~~~~~~~~~~~~
+Tile op dispatch
+~~~~~~~~~~~~~~~~
-Tile primitive dispatch selects an implementation according to the primitive,
-the current execution scope, the operand layouts, and the target backend. For
-example:
+Each backend registers lowering implementations for its tile ops. On CUDA,
+the caller selects an exact PTX instruction, including its qualifiers::
-- On CUDA, ``copy`` currently selects fixed-width or automatically vectorized
- loads/stores, ``ldmatrix`` / ``stmatrix``, or a scalar fallback. The
- asynchronous ``copy_async`` primitive has separate ``ldgsts``, TMA,
- distributed-shared-memory, and tensor-memory variants.
-- On CUDA, ``gemm`` currently selects an ``mma.m16n8k*`` implementation, while
- ``gemm_async`` selects ``tcgen05``. Trainium registers its own ``copy``,
- ``gemm``, reduction, and elementwise implementations.
+ Tx.cuda.tile.apply(Tx.ptx["ld.global.v4.b32"], registers, global_words)
-Once a variant is selected, dispatch generates the loops and addressing to
-apply that instruction across the whole tile.
+Each instruction defines a canonical operand pattern. The caller prepares
+views with the required atom coordinates and repeat axes; lowering checks the
+operand shapes, physical layouts, and execution scope, then emits that same
+instruction for each repeat coordinate. The caller controls vector width,
+matrix shape, transpose mode, temporary storage, and synchronization.
+``Tx.cuda.tile.ewise`` repeats a caller-defined scalar function over registers.
+
+Trainium exposes named instruction ops such as ``load``, ``store``,
+``tensor_copy``, ``matmul``, and ``tensorreduce``. Their lowering checks the
+instruction's operand and layout requirements and generates the corresponding
+NKI calls.
Putting it together
~~~~~~~~~~~~~~~~~~~~~
Across a full kernel, orchestration stays in ordinary source code, and the
-recurring hardware operations appear as tile primitives. For the building blocks —
-TMA loads, ``tcgen05`` async MMA, reductions, and the rest — see the per-dispatch
-walkthroughs in :doc:`tile_primitives`.
+recurring hardware operations appear as tile ops. For the building blocks —
+TMA loads, ``tcgen05`` async MMA, Trainium reductions, and the rest — see the
+instruction walkthroughs in :doc:`tile_ops`.
What TIRx Enables
-----------------
@@ -194,7 +196,7 @@ A stable extension boundary for future hardware
New hardware support is a staged process rather than a redesign of the DSL: a
feature is first exposed as a backend intrinsic, then promoted into a tile
-primitive once the usage pattern repeats. Future hardware grows the backend
+op once the usage pattern repeats. Future hardware grows the backend
library, not the core language.
Megakernels and composable tile tasks
@@ -239,6 +241,6 @@ Next Steps
- :doc:`install` — install TIRx and the kernel library.
- :doc:`native_basics/cuda/first_kernel` — write and compile a first CUDA kernel.
- :doc:`programming/index` — the complete programming guide.
-- :doc:`layout` and :doc:`tile_primitives` — reusable tensor and tile abstractions.
+- :doc:`layout` and :doc:`tile_ops` — reusable tensor and tile abstractions.
- :doc:`arch/index` — compiler internals (lowering pipeline, passes, codegen).
- :doc:`api/index` — authoring, IR, compiler, and extension API signatures.
diff --git a/docs/tirx/programming/index.rst b/docs/tirx/programming/index.rst
index 63d21cc95b..4bd47c14f4 100644
--- a/docs/tirx/programming/index.rst
+++ b/docs/tirx/programming/index.rst
@@ -20,7 +20,7 @@ Programming Guide
This guide explains how to write, compile, and inspect TIRx kernels. Start
with the CUDA guide for the core language and native programming model, then
-use tensor layouts and tile primitives to express reusable tile-level work.
+use tensor layouts and tile ops to express reusable tile-level work.
The examples use a single alias for the TIRx TVMScript dialect::
@@ -35,7 +35,7 @@ and direct CUDA instructions are available through ``Tx.cuda.*`` and
../native_basics
../layout
- ../tile_primitives
+ ../tile_ops
For callable reference entries, see the :doc:`API reference <../api/index>`.
For lowering, dispatch, and backend implementation details, see
diff --git a/docs/tirx/tile_ops.rst b/docs/tirx/tile_ops.rst
new file mode 100644
index 0000000000..ef9bfb3e36
--- /dev/null
+++ b/docs/tirx/tile_ops.rst
@@ -0,0 +1,117 @@
+.. 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.
+
+Tile Operations
+===============
+
+CUDA tile operations apply an explicit PTX instruction or a scalar function to
+canonical tensor views::
+
+ T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], registers, global_words)
+ T.cuda.tile.apply(T.ptx["st.shared.v4.b32"], shared_words, registers)
+
+See :doc:`cuda_tile_patterns` for every operand pattern, descriptor contract,
+execution-group axis, and caller view. ``T.cuda.tile.ewise`` applies a
+caller-defined scalar TIRx function to registers. No CUDA tile call selects a
+vector width, matrix shape, transpose mode, or staging algorithm from layouts.
+
+Both backends use ordinary opaque void Calls. Tensor regions, pointers,
+predicates, coordinates and descriptor parameters are normal IR arguments.
+The shared ``FTileOpDispatch`` boundary remains unchanged. ``T.ptx`` and
+``T.nki`` expose raw scalar/address instructions.
+
+.. toctree::
+ :hidden:
+
+ cuda_tile_patterns
+
+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 implemented directly by the caller.
+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.
+
+CUDA instruction guides
+-----------------------
+
+Pages follow PTX instruction families. Each page specifies the operand image,
+an explicit caller view, and the resulting instruction count. Asynchronous
+completion belongs to the selected instruction family, not to a generic
+``copy_async`` or ``gemm_async`` operation.
+
+.. toctree::
+ :maxdepth: 1
+
+ tile_ops/ld_st
+ tile_ops/ldmatrix_stmatrix
+ tile_ops/mov
+ tile_ops/add_sub_mul
+ tile_ops/div_rcp_sqrt
+ tile_ops/min_max
+ tile_ops/fma
+ tile_ops/cvt
+ tile_ops/ex2_lg2
+ tile_ops/abs_neg
+ tile_ops/cp_async
+ tile_ops/cp_async_bulk
+ tile_ops/cp_async_bulk_tensor
+ tile_ops/tcgen05_ld_st
+ tile_ops/tcgen05_cp
+ tile_ops/mma_sync
+ tile_ops/tcgen05_mma
+ tile_ops/ewise
+
+Trainium reductions
+-------------------
+
+.. toctree::
+ :maxdepth: 1
+
+ tile_ops/reduction
diff --git a/docs/tirx/tile_ops/abs_neg.rst b/docs/tirx/tile_ops/abs_neg.rst
new file mode 100644
index 0000000000..2490de1a76
--- /dev/null
+++ b/docs/tirx/tile_ops/abs_neg.rst
@@ -0,0 +1,56 @@
+.. 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.
+
+abs / neg
+=========
+
+``abs`` and ``neg`` operate independently on local register values. Their
+opcode selects the numeric type and applicable qualifiers. Operand order
+is ``dst,src``; each scalar tensor operand ends in ``(1,)`` after the group
+and repeat axes. An ordinary dtype view does not change the numeric type
+of the requested operation.
+
+Example: an explicit two-instruction scalar expression
+------------------------------------------------------
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+
+
+ @T.function
+ def negative_magnitude(A: T.Tensor((4, 1), "float32"), Out: T.Tensor((4, 1), "float32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=1))
+ r = T.alloc_local((4, 1), "float32", layout="default")
+ T.cuda.tile.apply(T.ptx["ld.global.f32"], r, A)
+ T.cuda.tile.apply(T.ptx["abs.f32"], r, r)
+ T.cuda.tile.apply(T.ptx["neg.f32"], r, r)
+ T.cuda.tile.apply(T.ptx["st.global.f32"], Out, r)
+
+For finite nonzero inputs this computes ``-abs(A)``. The tile expansion has
+four abs calls followed by four neg calls. This order comes from the two
+explicit operations, not from a hidden unary-op combination. Later compiler
+optimizations may simplify the expression. For a single caller-defined
+scalar expression with no exact-instruction requirement, use :doc:`ewise`.
+
+NaN, signed-zero, and integer boundary behavior belongs to the selected PTX
+type/form, not to a generic Python ``abs`` reference. Destination storage may
+be gapped but must be distinct across repeats. Register ownership cannot
+change between source and destination. Predication follows :doc:`mov`.
+
+Implementation: ``python/tvm/backend/cuda/tile/patterns/register.py``.
+No temporary memory, numeric cast, barrier, or packed form is inserted.
diff --git a/docs/tirx/tile_ops/add_sub_mul.rst b/docs/tirx/tile_ops/add_sub_mul.rst
new file mode 100644
index 0000000000..3894b583c6
--- /dev/null
+++ b/docs/tirx/tile_ops/add_sub_mul.rst
@@ -0,0 +1,113 @@
+.. 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.
+
+add / sub / mul
+===============
+
+The opcode chooses the numeric type, rounding, saturation and packing mode.
+The operand order is ``dst, lhs, rhs``. For a scalar instruction, each tensor
+ends in a singleton register-group axis; for example ``(8,1)`` at thread
+scope means eight scalar instructions. A read-only scalar expression may
+replace an input tensor and broadcasts over all eight repeats.
+
+Native packed arithmetic
+-------------------------
+
+Packed arithmetic still has one register **carrier** per operand:
+
+.. list-table::
+ :header-rows: 1
+
+ * - Example opcode
+ - Carrier dtype
+ - Values per carrier
+ * - ``add.rn.f32``
+ - float32
+ - One float32
+ * - ``add.rn.f16x2``
+ - uint32 (raw bits)
+ - Two float16
+ * - ``mul.rn.bf16x2``
+ - uint32 (raw bits)
+ - Two bfloat16
+ * - ``add.rn.f32x2``
+ - uint64 (raw bits)
+ - Two float32
+
+Not every qualifier is valid for every type or architecture. PTX form
+validation checks the spelling and the compiler checks the target. Packing
+is a caller decision, not the consequence of an even tensor extent.
+
+Example: eight packed half additions
+------------------------------------
+
+This kernel accepts packed input words and leaves the result packed:
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+
+
+ @T.function
+ def packed_add(
+ A: T.Tensor((8, 1), "uint32"), B: T.Tensor((8, 1), "uint32"), Out: T.Tensor((8, 1), "uint32")
+ ):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=1))
+ a = T.alloc_local((8, 1), "uint32", layout="default")
+ b = T.alloc_local((8, 1), "uint32", layout="default")
+ T.cuda.tile.apply(T.ptx["ld.global.b32"], a, A)
+ T.cuda.tile.apply(T.ptx["ld.global.b32"], b, B)
+ T.cuda.tile.apply(T.ptx["add.rn.f16x2"], a, a, b)
+ T.cuda.tile.apply(T.ptx["st.global.b32"], Out, a)
+
+.. code-block:: python
+
+ import numpy as np
+ import tvm
+
+ x = np.arange(16, dtype="float16").reshape(8, 2)
+ y = np.full((8, 2), 2, dtype="float16")
+ compiled = tvm.tirx.compile(packed_add, target="cuda", pipeline="tirx")
+ out = tvm.runtime.tensor(np.zeros((8, 1), dtype="uint32"), tvm.cuda())
+ compiled(
+ tvm.runtime.tensor(x.view("uint32"), tvm.cuda()),
+ tvm.runtime.tensor(y.view("uint32"), tvm.cuda()),
+ out,
+ )
+ np.testing.assert_array_equal(out.numpy().view("float16"), x + y)
+
+The arithmetic expansion is eight ``add.rn.f16x2`` instructions, processing
+sixteen half values. It does not insert a pack/unpack conversion. A scalar
+constant passed to a packed opcode must already represent the whole packed
+carrier; a numeric uint32 constant ``2`` is not the encoding of two half 2s.
+
+Subtraction and multiplication
+------------------------------
+
+``sub`` preserves the written operand order, including a scalar on the left.
+``mul`` uses the selected floating-point form or the integer ``lo``/``hi``/
+``wide`` form and its declared result type. Different result widths require
+the matching destination dtype, not a silent reinterpretation.
+
+Tensor operands share repeat extents and must have the correct thread owner.
+An input broadcast is a scalar or an explicit zero-stride view. Output
+storage cannot overlap across repeats. ``pred`` and ``preserve_dst`` follow
+the register instruction contract described in :doc:`mov`.
+
+For exp/SiLU plus scale/bias, define the scalar expression in :doc:`ewise`.
+Shared-memory inputs require explicit loads and stores around the arithmetic.
+Implementation: ``patterns/register.py`` under ``python/tvm/backend/cuda/tile``.
diff --git a/docs/tirx/tile_ops/cp_async.rst b/docs/tirx/tile_ops/cp_async.rst
new file mode 100644
index 0000000000..23ed4ac86b
--- /dev/null
+++ b/docs/tirx/tile_ops/cp_async.rst
@@ -0,0 +1,153 @@
+.. 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.
+
+cp.async
+========
+
+``cp.async`` copies a short contiguous chunk from global to shared memory
+without passing the payload through registers. Each issuing thread owns its
+async-copy group. Select the cache/prefetch qualifiers in the opcode and the
+copy width with ``chunk_bytes``::
+
+ T.cuda.tile.apply(
+ T.ptx["cp.async.cg.shared.global"], dst, src,
+ chunk_bytes=16, scope="cta",
+ )
+
+Operand contract
+----------------
+
+The operand order is ``dst_shared, src_global``. At thread scope each region
+has shape ``(*repeat,element)``. The final axis occupies exactly
+``chunk_bytes`` bytes and is contiguous and naturally aligned on both sides.
+For float16 with ``chunk_bytes=16``, that axis has eight elements. An explicit
+b32 view of the same bytes has four elements.
+
+.. list-table::
+ :header-rows: 1
+
+ * - Field
+ - Meaning
+ * - ``.ca`` / ``.cg``
+ - Select in the opcode. ``.ca`` supports 4/8/16 bytes; ``.cg`` requires 16.
+ * - ``chunk_bytes``
+ - Required constant 4, 8 or 16; no width is inferred from the layout.
+ * - ``.L2::64B`` / ``.L2::128B`` / ``.L2::256B``
+ - Optional prefetch qualifier in the opcode, not a larger copy size.
+ * - ``cache_policy``
+ - Runtime operand for an opcode containing ``.L2::cache_hint``.
+ * - ``scope``
+ - Thread by default. CTA scope adds explicit ``(warp,lane)`` group axes.
+
+The same byte count is copied when source and destination use different
+storage element widths, provided each side's atom contains that many whole
+contiguous bytes. This is a bit copy, not a conversion.
+
+Example: four 16-byte copies per thread
+---------------------------------------
+
+This complete kernel uses ``(128,32)`` float16 global and
+shared tensors. It explicitly assigns four eight-element chunks to each of
+128 threads:
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+
+
+ @T.function
+ def copy_async(A: T.Tensor((128, 32), "float16"), B: T.Tensor((128, 32), "float16")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=128))
+ tid = T.cuda.thread_idx("x")
+ shared = T.alloc_shared((128, 32), "float16", layout="default")
+ src = A.rearrange("(w t) (r v) -> w t r v", w=4, t=32, v=8)
+ dst = shared.rearrange("(w t) (r v) -> w t r v", w=4, t=32, v=8)
+ T.cuda.tile.apply(
+ T.ptx["cp.async.cg.shared.global"],
+ dst,
+ src,
+ chunk_bytes=16,
+ scope="cta",
+ )
+ T.ptx.cp.async_.commit_group()
+ T.ptx.cp.async_.wait_group(0)
+ T.cuda.cta_sync()
+ for j in range(32):
+ B[tid, j] = shared[tid, j]
+
+The views have shape ``(4,32,4,8)``. Coordinates ``w,t`` bind the warp and
+lane, ``r`` repeats the instruction four times, and ``v`` selects eight
+float16 elements. For thread ``tid=32*w+t``, issue ``r`` accesses
+``A[tid,8*r:8*r+8]`` and the matching shared row. There is no data rearrangement.
+
+Compile and check it with an Ampere-or-later CUDA device:
+
+.. code-block:: python
+
+ import numpy as np
+ import tvm
+
+ compiled = tvm.tirx.compile(copy_async, target="cuda", pipeline="tirx")
+ x = (np.arange(4096) % 97).astype("float16").reshape(128, 32)
+ a = tvm.runtime.tensor(x, tvm.cuda())
+ b = tvm.runtime.tensor(np.zeros_like(x), tvm.cuda())
+ compiled(a, b)
+ np.testing.assert_array_equal(b.numpy(), x)
+
+The transfer part emits four instances of this instruction per thread,
+followed by the caller's commit and wait. Pointer calculations and inline
+helper names are omitted here:
+
+.. code-block:: text
+
+ cp.async.cg.shared.global [shared_address], [global_address], 16;
+ // repeated for r = 0, 1, 2, 3
+ cp.async.commit_group;
+ cp.async.wait_group 0;
+
+Predication is different from zero fill
+---------------------------------------
+
+``pred=False`` suppresses the instruction; it does not write zeros to shared
+memory. To copy a partial chunk and zero the remaining bytes, supply
+``src_size`` in the range ``0..chunk_bytes``. To zero the entire chunk, use
+``src_size=0`` or ``ignore_src=True``. ``src_size`` and ``ignore_src`` are
+alternative operand forms and cannot both be supplied.
+
+For example, inside a kernel where ``tid`` selects one 16-byte row::
+
+ T.cuda.tile.apply(
+ T.ptx["cp.async.cg.shared.global.L2::128B"],
+ shared.sub[tid, :], A.sub[tid, :], chunk_bytes=16,
+ src_size=T.if_then_else(tid < valid_rows, T.uint32(16), T.uint32(0)),
+ )
+
+Here both tensors have uint8 rows of length 16. The final allocated row is
+still part of the declared source view: a zero-fill operand does not relax
+the tile pattern's bounds/shape proof.
+
+Completion and errors
+----------------------
+
+Each thread commits and waits for its own issues. Use a CTA barrier before
+other threads read the resulting shared tile. These calls are not inserted
+by ``apply``. A shared-to-global request, a noncontiguous chunk, an unaligned
+base, or a missing group axis fails its contract; lowering never substitutes
+a synchronous copy or a narrower instruction.
+
+Implementation: ``python/tvm/backend/cuda/tile/patterns/async_copy.py``.
+See :doc:`../cuda_tile_patterns` for the shared group/repeat rules.
diff --git a/docs/tirx/tile_ops/cp_async_bulk.rst b/docs/tirx/tile_ops/cp_async_bulk.rst
new file mode 100644
index 0000000000..e730de04a0
--- /dev/null
+++ b/docs/tirx/tile_ops/cp_async_bulk.rst
@@ -0,0 +1,124 @@
+.. 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.
+
+cp.async.bulk: shared to cluster shared
+=======================================
+
+This pattern issues
+``cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes``.
+It copies one contiguous shared-memory chunk to another CTA in the launch
+cluster and completes transaction bytes on a barrier in the destination CTA.
+For global tensor maps, use :doc:`cp_async_bulk_tensor` instead.
+
+The pointer contract
+---------------------
+
+``apply`` takes ``dst_shared, src_shared`` followed by explicit
+``chunk_bytes``, ``mbar`` and ``remote_cta_id``. Source and destination
+regions use coordinates ``(*group,*repeat,element)``. The final axis covers
+exactly ``chunk_bytes``, a positive multiple of 16; its bytes are consecutive
+and its origin is 16-byte aligned. Both shared/shared.dyn storage are accepted.
+
+All three addresses supplied by the caller are **local** shared addresses:
+
+* ``src`` names storage in the issuing CTA.
+* ``dst`` names the local counterpart of the destination allocation. Lowering
+ maps its address to ``remote_cta_id`` with ``mapa`` for each repeat.
+* ``mbar`` names the local counterpart of the completion barrier. Lowering
+ maps it to the same remote CTA once before the repeat loop.
+
+Do not pass an already mapped destination or barrier. Both CTAs must have
+corresponding live shared allocations at those offsets. ``remote_cta_id``
+is the rank within the cluster, not an absolute block index. ``pred`` gates
+the copy instruction and must agree with the caller's transaction accounting.
+
+Example: CTA 0 sends 256 bytes to CTA 1
+---------------------------------------
+
+The one-thread blocks here keep election and initialization visible. The
+same operation can be issued by an elected thread in a larger CTA.
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+
+
+ @T.function
+ def cluster_copy(A: T.Tensor((64,), "uint32"), Out: T.Tensor((64,), "uint32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=2, block=1, cluster=2))
+ rank = T.cuda.cluster_cta_id("x")
+ src = T.alloc_shared((64,), "uint32", layout="default", align=16)
+ dst = T.alloc_shared((64,), "uint32", layout="default", align=16)
+ mbar = T.alloc_shared((1,), "uint64", layout="default")
+ T.ptx.mbarrier.init.shared.b64(mbar.ptr_to([0]), T.uint32(1))
+ T.ptx.fence.mbarrier_init.release.cluster()
+ T.cuda.cluster_sync()
+ if rank == 0:
+ for i in range(64):
+ src[i] = A[i]
+ T.ptx.fence.proxy.async_.shared__cta()
+ T.cuda.tile.apply(
+ T.ptx["cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes"],
+ dst,
+ src,
+ chunk_bytes=256,
+ mbar=mbar.ptr_to([0]),
+ remote_cta_id=T.int32(1),
+ )
+ else:
+ T.ptx.mbarrier.arrive.expect_tx.shared.b64(mbar.ptr_to([0]), T.uint32(256))
+ T.cuda.mbarrier_wait(mbar.ptr_to([0]), 0)
+ for i in range(64):
+ Out[i] = dst[i]
+ T.cuda.cluster_sync()
+
+The sender emits one barrier-address mapa, one destination-address mapa,
+and one 256-byte bulk copy. The source address stays local. CTA 1 expects
+256 transaction bytes and waits for phase zero before consuming the data.
+The last cluster barrier keeps both CTAs' shared allocations alive until
+the communication is finished.
+
+Compile and verify on a device supporting thread-block clusters:
+
+.. code-block:: python
+
+ import numpy as np
+ import tvm
+
+ x = np.arange(64, dtype="uint32")
+ compiled = tvm.tirx.compile(cluster_copy, target="cuda", pipeline="tirx")
+ out = tvm.runtime.tensor(np.zeros_like(x), tvm.cuda())
+ compiled(tvm.runtime.tensor(x, tvm.cuda()), out)
+ np.testing.assert_array_equal(out.numpy(), x)
+
+Repeated chunks and synchronization
+-----------------------------------
+
+A float16 view ``(8,128):(256,1)`` with ``chunk_bytes=256`` represents eight
+separate contiguous chunks with gaps between them. It emits eight bulk
+copies. Set expected transaction bytes to ``8*256``. The source and
+destination may have different outer strides, provided each matching
+128-element chunk is contiguous and aligned on its own side.
+
+Do not flatten across a physical gap and request one larger chunk: it would
+describe different bytes and is rejected. There is no automatic split or
+fallback to a scalar loop. Cluster launch configuration, issuer election,
+barrier initialization, expected-byte accounting, proxy fences, phase
+progression and storage lifetime are caller responsibilities. Ordinary
+``cp.async.commit_group``/``wait_group`` is not this completion mechanism.
+
+Implementation: ``python/tvm/backend/cuda/tile/patterns/async_copy.py``.
diff --git a/docs/tirx/tile_ops/cp_async_bulk_tensor.rst b/docs/tirx/tile_ops/cp_async_bulk_tensor.rst
new file mode 100644
index 0000000000..ec10b94c45
--- /dev/null
+++ b/docs/tirx/tile_ops/cp_async_bulk_tensor.rst
@@ -0,0 +1,184 @@
+.. 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.
+
+cp.async.bulk.tensor / cp.reduce.async.bulk.tensor
+==================================================
+
+TMA transfers a descriptor-defined global box to or from a shared-memory
+image. Select direction, rank, tile/gather/scatter mode, multicast, cache
+qualifiers, and any reduction operator in the opcode. ``apply`` encodes
+the specified descriptor and issues that instruction at explicit coordinates;
+it does not plan a decomposition of an arbitrary tensor layout.
+
+Operands and descriptor fields
+------------------------------
+
+For a load the operands are ``shared, global``; for a store/reduction they
+are ``global, shared``. The global operand is either a rectangular global
+tensor region or an ordinary pointer to a caller-encoded TensorMap. In the
+latter case the caller is responsible for matching its fields to the tile
+contract; lowering does not inspect or select opaque descriptors.
+
+.. list-table::
+ :header-rows: 1
+
+ * - Keyword
+ - Contract
+ * - ``box``
+ - Required extents in outer-to-inner order; one per instruction dimension, each 1..256.
+ * - ``coords``
+ - Required runtime starting coordinates in the same order, relative to the descriptor view.
+ * - ``coord_steps``
+ - One rank-length increment vector per shared group/repeat axis; required when such axes exist.
+ * - ``format``
+ - Storage dtype by default, or an explicit matching format such as ``tf32`` over float32.
+ * - ``swizzle``
+ - 0/1/2/3: none/32B/64B/128B shared swizzle.
+ * - ``element_strides``
+ - Traversal steps in 1..8; default all ones. Without interleave, the innermost step is one.
+ * - ``interleave``
+ - 0/1/2: none/16B/32B. The global and shared images change as described below.
+ * - ``l2_promotion``, ``oob_fill``
+ - Explicit descriptor policies; both default to zero.
+ * - ``mbar`` / ``cta_mask`` / ``cache_policy``
+ - Runtime operands required by load / multicast / cache-hint forms, respectively.
+
+Without interleave, the global region has exactly the opcode rank, unit
+innermost stride, and suitably aligned outer byte strides. Its dimensions
+and outer strides can be runtime values and are passed to the host descriptor
+encoder. Known invalid fields are rejected rather than silently promoted,
+folded, transposed or reinterpreted as another dtype.
+
+The shared suffix is the box image: ``ceil(box[i]/element_strides[i])``
+per dimension, compact before swizzling. Every atom's unswizzled origin is
+128-byte aligned, even for swizzle zero; a swizzled backing base is aligned
+to the whole swizzle period. Swizzled subregions use their inverse-swizzled
+origins so the hardware applies the phase once.
+
+Example: two explicit 8x16 boxes
+--------------------------------
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+
+
+ @T.function
+ def tma_load(A: T.Tensor((16, 16), "uint32"), Out: T.Tensor((16, 16), "uint32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
+ lane = T.cuda.lane_id()
+ shared = T.alloc_shared((2, 8, 16), "uint32", layout="default", align=128)
+ mbar = T.alloc_shared((1,), "uint64", layout="default")
+ if lane == 0:
+ T.ptx.mbarrier.init.shared.b64(mbar.ptr_to([0]), T.uint32(1))
+ T.ptx.mbarrier.arrive.expect_tx.shared.b64(mbar.ptr_to([0]), T.uint32(1024))
+ T.cuda.cta_sync()
+ if lane == 0:
+ T.cuda.tile.apply(
+ T.ptx["cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx::bytes"],
+ shared,
+ A,
+ box=(8, 16),
+ coords=(0, 0),
+ coord_steps=((8, 0),),
+ mbar=mbar.ptr_to([0]),
+ )
+ T.cuda.mbarrier_wait(mbar.ptr_to([0]), 0)
+ T.cuda.cta_sync()
+ for g, r in T.grid(2, 8):
+ if lane < 16:
+ Out[g * 8 + r, lane] = shared[g, r, lane]
+
+The shared shape is ``(repeat=2,box_row=8,box_col=16)``. One elected lane
+issues two instructions at global coordinates ``(0,0)`` and ``(8,0)``.
+Each copies 512 bytes; one phase of the barrier accounts for 1024 bytes.
+Descriptor construction is host initialization associated with this call.
+It is not repeated on the device for each box.
+
+.. code-block:: python
+
+ import numpy as np
+ import tvm
+
+ x = np.arange(256, dtype="uint32").reshape(16, 16)
+ compiled = tvm.tirx.compile(tma_load, target="cuda", pipeline="tirx")
+ out = tvm.runtime.tensor(np.zeros_like(x), tvm.cuda())
+ compiled(tvm.runtime.tensor(x, tvm.cuda()), out)
+ np.testing.assert_array_equal(out.numpy(), x)
+
+Stores, reductions, and completion
+----------------------------------
+
+The corresponding store spelling is
+``cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group``. Reduction
+uses, for example,
+``cp.reduce.async.bulk.tensor.2d.global.shared::cta.add.bulk_group``.
+The coordinate/box contract is the same, with global destination first.
+Complete shared initialization and issue
+``T.ptx.fence.proxy.async_.shared__cta()`` before the transfer; the issuing
+thread then calls ``T.ptx.cp.async_.bulk.commit_group()`` and
+``T.ptx.cp.async_.bulk.wait_group(0)``. Initialize a reduction destination
+to the intended starting value. No mbarrier operand is used for these stores.
+
+Loads use their supplied mbarrier, with explicit initialization, expected
+transaction bytes, and phase wait. Multicast adds the explicit CTA mask and
+the corresponding cluster/barrier setup. The tile call elects no thread and
+inserts neither waits nor descriptor prefetch. A desired prefetch is an
+ordinary ``T.ptx.prefetch.tensormap`` call.
+
+Gather4 and scatter4
+--------------------
+
+These forms use a 2D descriptor with ``box=(1,width)`` and a shared image
+``(4,width)``. The four row coordinates are ordinary runtime operands in
+``gather4``; they need not be consecutive. For a prepared ``shared`` of
+shape ``(4,32)`` and a matching barrier, an issue can be written as::
+
+ T.cuda.tile.apply(
+ T.ptx["cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes"],
+ shared, A, box=(1,32), coords=(0,0),
+ gather4=(row,row+3,row+7,row+11), mbar=mbar.ptr_to([0]),
+ )
+
+For uint32 this is one 512-byte gather4 issue. Scatter4 uses global
+destination first and the matching store opcode/completion protocol.
+Gather/scatter coordinate selection is caller logic.
+
+Interleave and packed FP4
+-------------------------
+
+With ``interleave=1`` or ``2``, a rank-3..5 opcode takes global axes
+``(N,C_slices,*spatial,Ci)``. ``Ci`` holds one 16B/32B slice and is implicit
+in the descriptor rank. Its elements are contiguous and the innermost
+spatial stride is one slice. Interleave 2 requires swizzle 1 and 32-byte
+global alignment. One issue selects **one** C slice, even if a larger C box
+field is supplied; channel iteration requires an explicit repeat and C step.
+
+For example, uint16/interleave-2 with ``box=(1,1,8)`` has shared atom
+``(1,1,8,16)``. Add a leading repeat of three and
+``coord_steps=((0,1,0),)`` to issue three successive channel slices.
+Gather4/scatter4 do not use interleave. On B200, 16B interleave with
+swizzle 1/2/3 can trap, including when issued through raw PTX.
+
+``format="float4_e2m1fn"`` selects the load-only ``16U4_ALIGN16B`` format.
+The inner global dimension is a multiple of 128, the inner box is 128, and
+global origins/outer strides are 32-byte aligned. Each sixteen nibbles
+occupy the low eight bytes of a sixteen-byte shared block. The remaining
+eight bytes are undefined padding; barrier accounting counts payload bytes,
+not this padding. Swizzle 0/3 and gather4 use this explicit packed image.
+
+Implementation: ``python/tvm/backend/cuda/tile/patterns/tma.py``.
diff --git a/docs/tirx/tile_ops/cvt.rst b/docs/tirx/tile_ops/cvt.rst
new file mode 100644
index 0000000000..a143b783fa
--- /dev/null
+++ b/docs/tirx/tile_ops/cvt.rst
@@ -0,0 +1,75 @@
+.. 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.
+
+cvt
+===
+
+``cvt`` performs a numeric conversion selected by the opcode. A dtype view
+only reinterprets bits; it is used here to expose the native carrier layout,
+not to perform the conversion itself. Destination and source widths, operand
+group lengths, rounding, saturation and packing come from the PTX form.
+
+Scalar conversion uses tensor suffix ``(1,)`` on both sides. Packed
+``cvt.rn.f16x2.f32`` has three operands: one b32 destination, a float32 source
+for its **high** halfword, and a float32 source for its **low** halfword.
+All three have suffix ``(1,)``. The pair is not an inferred vectorization of
+a scalar cast.
+
+Example: eight float32 values to four packed half pairs
+-------------------------------------------------------
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+
+
+ @T.function
+ def narrow_pairs(A: T.Tensor((4, 2), "float32"), Out: T.Tensor((4, 1), "uint32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=1))
+ a = T.alloc_local((4, 2), "float32", layout="default")
+ d = T.alloc_local((4, 1), "uint32", layout="default")
+ T.cuda.tile.apply(T.ptx["ld.global.v2.f32"], a, A)
+ T.cuda.tile.apply(T.ptx["cvt.rn.f16x2.f32"], d, a[:, 1:2], a[:, 0:1])
+ T.cuda.tile.apply(T.ptx["st.global.b32"], Out, d)
+
+For each input row ``[x0,x1]``, the output word contains converted ``x0``
+in its low half and ``x1`` in its high half. Reversing the two source operands
+would exchange the output elements. The expansion emits four packed cvt
+instructions, one per row.
+
+.. code-block:: python
+
+ import numpy as np
+ import tvm
+
+ x = np.linspace(-2, 2, 8, dtype="float32").reshape(4, 2)
+ compiled = tvm.tirx.compile(narrow_pairs, target="cuda", pipeline="tirx")
+ out = tvm.runtime.tensor(np.zeros((4, 1), dtype="uint32"), tvm.cuda())
+ compiled(tvm.runtime.tensor(x, tvm.cuda()), out)
+ np.testing.assert_array_equal(out.numpy().view("float16"), x.astype("float16"))
+
+The inverse operation is two scalar ``cvt.f32.f16`` instructions per packed
+word; there is no corresponding ``cvt.f32x2.f16x2`` spelling in this API.
+Expose the packed source with a uint16 view, add a singleton register-group
+axis per half, and apply the scalar opcode. For a composite conversion
+expression, including explicit intermediate rounding, use :doc:`ewise`.
+
+Register ownership must agree on both sides even when their storage widths
+differ. A wrong carrier width, missing group axis, or overlapping destination
+is rejected. There is no implicit staging through shared memory.
+
+Implementation: ``python/tvm/backend/cuda/tile/patterns/register.py``.
diff --git a/docs/tirx/tile_ops/div_rcp_sqrt.rst b/docs/tirx/tile_ops/div_rcp_sqrt.rst
new file mode 100644
index 0000000000..d904d1294e
--- /dev/null
+++ b/docs/tirx/tile_ops/div_rcp_sqrt.rst
@@ -0,0 +1,77 @@
+.. 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.
+
+div / rcp / sqrt
+================
+
+Select accuracy, rounding and flush-to-zero behavior in the PTX opcode.
+These choices change arithmetic semantics; the tile layer does not replace
+a division with multiplication by an approximate reciprocal.
+
+.. list-table::
+ :header-rows: 1
+
+ * - Example
+ - Operand order
+ - Atom
+ * - ``div.rn.f32``
+ - ``dst, numerator, denominator``
+ - One float32 per operand
+ * - ``rcp.approx.ftz.f32``
+ - ``dst, src``
+ - One float32 per operand
+ * - ``sqrt.rn.f32``
+ - ``dst, src``
+ - One float32 per operand
+
+Tensor operands end in ``(1,)``, preceded by execution-group and repeat
+axes. Scalar read-only inputs broadcast over those repeats. A constant
+numerator is allowed; its position is not exchanged with the denominator.
+
+Example: explicit division followed by square root
+--------------------------------------------------
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+
+
+ @T.function
+ def root_ratio(A: T.Tensor((4, 1), "float32"), Out: T.Tensor((4, 1), "float32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=1))
+ r = T.alloc_local((4, 1), "float32", layout="default")
+ T.cuda.tile.apply(T.ptx["ld.global.f32"], r, A)
+ T.cuda.tile.apply(T.ptx["div.rn.f32"], r, r, T.float32(4))
+ T.cuda.tile.apply(T.ptx["sqrt.rn.f32"], r, r)
+ T.cuda.tile.apply(T.ptx["st.global.f32"], Out, r)
+
+For input ``[[4],[16],[36],[64]]``, the result is ``[[1],[2],[3],[4]]``.
+The arithmetic expansion contains four ``div.rn.f32`` and four
+``sqrt.rn.f32`` calls. Constant folding or instruction simplification later
+in the compiler may optimize them; no tile-layout heuristic chooses a form.
+
+To use an approximate reciprocal explicitly, replace the arithmetic with::
+
+ T.cuda.tile.apply(T.ptx["rcp.approx.ftz.f32"], r, r)
+
+This is a distinct numerical contract. A half-precision division composed
+from conversion and scalar arithmetic belongs in :doc:`ewise` unless a
+registered native form expresses exactly the requested operation.
+
+All tensor storage is local to the executing thread. Views may select gapped
+registers, but cannot move values between lanes. No memory staging or barrier
+is inserted. Implementation: ``python/tvm/backend/cuda/tile/patterns/register.py``.
diff --git a/docs/tirx/tile_ops/ewise.rst b/docs/tirx/tile_ops/ewise.rst
new file mode 100644
index 0000000000..36af0a9671
--- /dev/null
+++ b/docs/tirx/tile_ops/ewise.rst
@@ -0,0 +1,133 @@
+.. 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.
+
+ewise: caller-defined scalar compositions
+=========================================
+
+``T.cuda.tile.ewise(function, dst, *inputs, scope="thread")`` applies a
+typed scalar TIRx function independently to register elements. This is the
+one CUDA tile operation whose atom is a caller-defined expression rather
+than one PTX instruction. It can represent exp, SiLU, scale/bias, casts, or
+another scalar combination. Use :doc:`add_sub_mul`, :doc:`fma`, or
+:doc:`cvt` when the native instruction and its qualifiers are the contract.
+
+The scalar atom
+---------------
+
+Tensor operands have coordinates ``(*group,*repeat)``. There is no trailing
+PTX register-group axis: every coordinate after the group prefix selects
+one scalar invocation. Inputs and destination have matching extents and
+belong to the executing thread. Read-only zero-stride input views express
+tensor broadcasts; a destination must not overlap itself.
+
+Scalar arguments broadcast to every invocation. Their types must match the
+function parameters, and its result type must match the destination dtype.
+The function supports scalar bindings and returns, including an if/else
+with a scalar return on both paths. It does not express reductions, memory
+access algorithms, or inter-thread communication.
+
+Example: SiLU, scale, bias
+--------------------------
+
+This complete kernel uses four scalar register values per thread. Its
+memory instructions have a four-word suffix; ``ewise`` treats those same
+four coordinates as scalar repeats.
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+ from tvm.backend.cuda.axis import laneid, warpid
+ from tvm.tirx.layout import S, TileLayout
+
+ reg_layout = TileLayout(S[(2, 32, 4) : (1 @ warpid, 1 @ laneid, 1)])
+
+
+ @T.function(private=True)
+ def silu_affine(x: T.float32, scale: T.float32, bias: T.float32) -> T.float32:
+ y = x / (T.float32(1) + T.exp(-x))
+ return y * scale + bias
+
+
+ @T.function
+ def activation(A: T.Tensor((2, 32, 4), "float32"), Out: T.Tensor((2, 32, 4), "float32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
+ R = T.alloc_local((2, 32, 4), "float32", layout=reg_layout)
+ T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], R, A, scope="cta")
+ T.cuda.tile.ewise(silu_affine, R, R, T.float32(2), T.float32(1), scope="cta")
+ T.cuda.tile.apply(T.ptx["st.global.v4.b32"], Out, R, scope="cta")
+
+There is one vector load, four inlined scalar expressions, and one vector
+store per thread. The expression computes ``silu(x)*2+1``; it does not mean
+``silu(2*x+1)``. The ordinary scalar compiler lowers the expression; ``ewise``
+does not promise one PTX instruction, choose a packed form, or fuse an
+unspecified activation. For an exact exponential instruction use :doc:`ex2_lg2`.
+
+Shared-memory load/compute/store
+--------------------------------
+
+Both the destination and tensor inputs of ``ewise`` must be local registers.
+The following complete kernel makes the shared staging explicit. The final
+read deliberately uses another thread's row, so the second CTA barrier has
+a visible purpose.
+
+.. code-block:: python
+
+ @T.function(private=True)
+ def affine(x: T.float32) -> T.float32:
+ return x * T.float32(2) + T.float32(1)
+
+
+ @T.function
+ def shared_affine(A: T.Tensor((64, 4), "float32"), Out: T.Tensor((64, 4), "float32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
+ tid = T.cuda.thread_idx("x")
+ S = T.alloc_shared((64, 4), "float32", layout="default", align=16)
+ for j in range(4):
+ S[tid, j] = A[tid, j]
+ T.cuda.cta_sync()
+ sv = S.view(2, 32, 4)
+ R = T.alloc_local((2, 32, 4), "float32", layout=reg_layout)
+ T.cuda.tile.apply(T.ptx["ld.shared.v4.b32"], R, sv, scope="cta")
+ T.cuda.tile.ewise(affine, R, R, scope="cta")
+ T.cuda.tile.apply(T.ptx["st.shared.v4.b32"], sv, R, scope="cta")
+ T.cuda.cta_sync()
+ for j in range(4):
+ Out[tid, j] = S[(tid + 1) % 64, j]
+
+The two views refer to the existing shared and local allocations. The tile
+operation adds no temporary storage or barrier. If the shared store used a
+different ownership map that overwrote another thread's unread input, an
+additional barrier would be needed between all loads and all stores.
+
+Compile and verify both examples:
+
+.. code-block:: python
+
+ import numpy as np
+ import tvm
+
+ x = np.linspace(-2, 2, 256, dtype="float32")
+ for kernel, shape, ref in [
+ (activation, (2, 32, 4), (x / (1 + np.exp(-x)) * 2 + 1).reshape(2, 32, 4)),
+ (shared_affine, (64, 4), np.roll((x * 2 + 1).reshape(64, 4), -1, axis=0)),
+ ]:
+ compiled = tvm.tirx.compile(kernel, target="cuda", pipeline="tirx")
+ out = tvm.runtime.tensor(np.zeros(shape, dtype="float32"), tvm.cuda())
+ compiled(tvm.runtime.tensor(x.reshape(shape), tvm.cuda()), out)
+ np.testing.assert_allclose(out.numpy(), ref, rtol=2e-5, atol=2e-5)
+
+Implementation: ``python/tvm/backend/cuda/tile/patterns/ewise.py``.
diff --git a/docs/tirx/tile_ops/ex2_lg2.rst b/docs/tirx/tile_ops/ex2_lg2.rst
new file mode 100644
index 0000000000..cf339917be
--- /dev/null
+++ b/docs/tirx/tile_ops/ex2_lg2.rst
@@ -0,0 +1,56 @@
+.. 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.
+
+ex2 / lg2
+=========
+
+These PTX instructions calculate base-two exponential and logarithm. They
+are not natural exp/log. Select the explicit approximation and subnormal
+policy, for example ``ex2.approx.ftz.f32`` or ``lg2.approx.ftz.f32``.
+The tile layer forwards this spelling without selecting a replacement.
+
+Both tensor operands use local registers and end in ``(1,)``; each preceding
+repeat selects another scalar input/output. Output repeats cannot alias.
+A scalar read-only input broadcasts over all repeats.
+
+Example: four base-two exponentials
+-----------------------------------
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+
+
+ @T.function
+ def exp2_native(A: T.Tensor((4, 1), "float32"), Out: T.Tensor((4, 1), "float32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=1))
+ r = T.alloc_local((4, 1), "float32", layout="default")
+ T.cuda.tile.apply(T.ptx["ld.global.f32"], r, A)
+ T.cuda.tile.apply(T.ptx["ex2.approx.ftz.f32"], r, r)
+ T.cuda.tile.apply(T.ptx["st.global.f32"], Out, r)
+
+Inputs ``[-1,0,1,2]`` give the power-of-two reference ``[0.5,1,2,4]``.
+The arithmetic expansion contains four ex2 instructions. For general inputs,
+test with tolerances appropriate to the approximate opcode; ``.ftz`` is part
+of the chosen handling of subnormal values.
+
+Use ``lg2.approx.ftz.f32`` with the same operand pattern for base-two
+logarithms. If the operation is natural exp, SiLU, or an exp/scale/bias
+combination, write that scalar expression through :doc:`ewise`. The scalar
+compiler then owns its expansion; the call is not documented as one ex2.
+
+Implementation: ``python/tvm/backend/cuda/tile/patterns/register.py``.
diff --git a/docs/tirx/tile_ops/fma.rst b/docs/tirx/tile_ops/fma.rst
new file mode 100644
index 0000000000..2d4fd81ec6
--- /dev/null
+++ b/docs/tirx/tile_ops/fma.rst
@@ -0,0 +1,58 @@
+.. 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.
+
+fma
+===
+
+``T.cuda.tile.apply(T.ptx["fma.rn.f32"], dst, a, b, c)`` computes the
+selected fused multiply-add for each repeat. ``a*b+c`` is one native FMA
+atom, including its explicit rounding qualifier. It is not a request for
+two independently rounded multiply/add instructions.
+
+Scalar tensor operands have suffix ``(1,)``. Read-only scalar operands
+broadcast over repeats. Packed f16x2/bf16x2 forms use uint32 carriers;
+packed f32x2 uses uint64. A packed constant must encode both values.
+The opcode and storage representation stay explicit even if a contiguous
+layout would permit several other forms.
+
+Example: four fused operations
+------------------------------
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+
+
+ @T.function
+ def affine_fma(A: T.Tensor((4, 1), "float32"), Out: T.Tensor((4, 1), "float32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=1))
+ r = T.alloc_local((4, 1), "float32", layout="default")
+ T.cuda.tile.apply(T.ptx["ld.global.f32"], r, A)
+ T.cuda.tile.apply(T.ptx["fma.rn.f32"], r, r, T.float32(2), T.float32(1))
+ T.cuda.tile.apply(T.ptx["st.global.f32"], Out, r)
+
+The arithmetic expansion has four ``fma.rn.f32`` calls. For small integer
+inputs the reference is exactly ``2*A+1``. Tests intended to distinguish
+rounding modes must instead use inputs whose exact product-plus-addend lies
+between representable outputs; a rounded multiply followed by a rounded
+add is not an adequate FMA reference.
+
+Use ``.rm``, ``.rp`` or ``.rz`` only when that form's rounding is intended
+and supported by the selected type. Predication and destination preservation
+follow :doc:`mov`. Memory loads/stores and synchronization remain caller code.
+
+Implementation: ``python/tvm/backend/cuda/tile/patterns/register.py``.
diff --git a/docs/tirx/tile_ops/ld_st.rst b/docs/tirx/tile_ops/ld_st.rst
new file mode 100644
index 0000000000..6820bc544a
--- /dev/null
+++ b/docs/tirx/tile_ops/ld_st.rst
@@ -0,0 +1,128 @@
+.. 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.
+
+ld / st
+========
+
+Use ``T.cuda.tile.apply`` to repeat a chosen memory instruction. The operand
+order is always destination first: a load takes ``registers, memory`` and a
+store takes ``memory, registers``. For example::
+
+ T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], R, A)
+ T.cuda.tile.apply(T.ptx["st.shared.v4.b32"], S, R)
+
+Both calls transfer four 32-bit words per repeat. Global-to-shared copying
+through registers consists of these two explicit calls and caller-allocated
+register storage. To issue a direct asynchronous transfer, use :doc:`cp_async`.
+
+Contract and instruction selection
+----------------------------------
+
+Operands have shape ``(*group,*repeat,word)``. The final extent is the vector
+length in the opcode, or one for a scalar instruction. Both operands use
+scalar elements with the opcode's storage bit width. Each memory atom must
+be consecutive and aligned to the whole vector's byte size. Its start can
+vary across repeats. Local register slots can be separated or permuted, but
+each must belong to the issuing thread; loaded destinations cannot overlap.
+
+The supported address-space spellings are ``global``, ``shared``,
+``shared::cta``, and ``local``. The pointer side of ``ld.local``/``st.local``
+is addressable local memory, distinct from the register operands. Legal
+types, vector lengths, and qualifiers come from the selected PTX form.
+
+Cache operators such as ``.nc`` and ``.cg`` and prefetch/eviction qualifiers
+are part of the opcode. A form with a cache-policy operand takes
+``cache_policy=...``. ``pred`` gates execution. A predicated load can use
+``preserve_dst=True`` to keep an already initialized destination when the
+predicate is false. Lowering does not reduce the width when alignment fails.
+
+Example: a strided tensor and gapped registers
+---------------------------------------------------------
+
+The following kernel reads a float32 tensor whose logical axes are not in
+physical order. It keeps that layout and prepares a view with four physically
+adjacent words at the end. The destination deliberately has twice the storage
+stride, so each vector writes separated local slots.
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+ from tvm.tirx.layout import S, TileLayout
+
+ global_layout = TileLayout(S[(15, 2, 7, 2) : (4, 1, 60, 2)])
+ register_layout = TileLayout(S[(15, 2, 7, 2) : (8, 2, 120, 4)])
+
+
+ @T.function
+ def strided_load(
+ A: T.Tensor((15, 2, 7, 2), "float32", layout=global_layout, align=16),
+ Out: T.Tensor((420,), "float32"),
+ ):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=1))
+ R = T.alloc_tensor(
+ ty_args=[T.Tensor((15, 2, 7, 2), "float32", scope="local", layout=register_layout)]
+ )
+ av = A.rearrange("i b j d -> j i (d b)")
+ rv = R.rearrange("i b j d -> j i (d b)")
+ T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], rv, av)
+ for j, i, d, b in T.grid(7, 15, 2, 2):
+ Out[60 * j + 4 * i + 2 * d + b] = R[i, b, j, d]
+
+Both views have shape ``(7,15,4)``. For one repeat ``(j,i)``, source word
+``w`` lives at ``60*j+4*i+w`` and destination word ``w`` lives at local
+offset ``120*j+8*i+2*w``. One instruction binds its four outputs directly to
+these four slots. This emits **105 ld.global.v4.b32 instructions**, with no
+intermediate tile or register-remapping copy. The final scalar stores are
+ordinary caller code used to inspect the result.
+
+Compile and check the complete kernel:
+
+.. code-block:: python
+
+ import numpy as np
+ import tvm
+
+ compiled = tvm.tirx.compile(strided_load, target="cuda", pipeline="tirx")
+ x = np.arange(420, dtype="float32")
+ a = tvm.runtime.tensor(x.reshape(15, 2, 7, 2), tvm.cuda())
+ out = tvm.runtime.tensor(np.zeros_like(x), tvm.cuda())
+ compiled(a, out)
+ np.testing.assert_array_equal(out.numpy(), x)
+
+The raw load in the generated PTX has four output operands, regardless of
+the local offsets selected by the view:
+
+.. code-block:: text
+
+ ld.global.v4.b32 {r0, r1, r2, r3}, [address];
+
+Sharing the work across threads
+-------------------------------
+
+At CTA scope, prepend ``(warp,lane)``. With 128 threads,
+``(4,32,3,4)`` means three four-word instructions per thread. A register
+layout can assign the first two axes strides ``1 @ warpid`` and
+``1 @ laneid``, and the last two ordinary local strides ``4,1``. Memory
+views must expose the same group/repeat coordinates. These axes distribute
+the operation; they are not searched for in an arbitrary layout.
+
+Loads and stores do not insert a barrier. If another thread will consume
+shared stores, synchronize at that dependency. To change ownership or
+reorder shared data in place, finish the old reads before overwriting bytes.
+
+Implementation: ``python/tvm/backend/cuda/tile/patterns/memory.py``.
+See :doc:`../cuda_tile_patterns` for group, repeat, and bounds rules.
diff --git a/docs/tirx/tile_ops/ldmatrix_stmatrix.rst b/docs/tirx/tile_ops/ldmatrix_stmatrix.rst
new file mode 100644
index 0000000000..699ed1ec5d
--- /dev/null
+++ b/docs/tirx/tile_ops/ldmatrix_stmatrix.rst
@@ -0,0 +1,147 @@
+.. 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.
+
+ldmatrix / stmatrix
+===================
+
+The tile patterns support ``m8n8.b16`` with ``.x1``, ``.x2``, or ``.x4``
+and optional ``.trans``. Select the exact form in the opcode::
+
+ T.cuda.tile.apply(
+ T.ptx["ldmatrix.sync.aligned.m8n8.x2.shared.b16"],
+ registers, shared_words, scope="warp",
+ )
+
+``stmatrix`` reverses the operand order: ``shared_words, registers``.
+The implementation repeats the specified warp instruction, without picking
+a matrix count or transpose mode from the layout.
+
+Address providers and result owners
+-----------------------------------
+
+Both regions end in ``(matrix,row,word)=(count,8,4)`` and use scalar 32-bit
+containers. A container holds two adjacent 16-bit values. Use a dtype view
+to expose packed half/bfloat16 storage as b32 words; this is not a conversion.
+
+* Each shared row has four consecutive words and a 16-byte-aligned start.
+ Different rows may have gaps and may be interleaved with other matrices.
+* Lane ``8*matrix+row`` provides that row's shared address. For x1/x2,
+ the remaining address providers repeat valid row addresses.
+* Register word ``(matrix,row,word)`` belongs to lane ``4*row+word``.
+ Each lane receives ``count`` words. Their register offsets need not be
+ adjacent, and need not have the same order as another tensor declaration.
+
+Thus, the lane supplying a shared address generally differs from the lanes
+receiving the values from that row. The contract checks both maps. All 32
+lanes must participate, including those whose addresses are duplicates.
+
+Example: interleaved shared rows and register slots 0, 2, 8, 10
+---------------------------------------------------------------
+
+Here the original shared tensor is ``(16,4,2):(8,1,128)``. The first matrix
+uses even original rows, the second uses odd rows, and the final original
+axis selects the second instruction. The factory below includes shared
+initialization, synchronization, and readback for one warp.
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+ from tvm.backend.cuda.axis import laneid
+ from tvm.tirx.layout import S, TileLayout
+
+ LDMATRIX = "ldmatrix.sync.aligned.m8n8.x2.shared.b16"
+
+ def matrix_kernel():
+ shared = TileLayout(S[(16, 4, 2) : (8, 1, 128)])
+ registers = TileLayout(S[(2, 2, 8, 4) : (8, 2, 4 @ laneid, 1 @ laneid)])
+
+ @T.function
+ def kernel(A: T.Tensor((16, 4, 2), "uint32"), Out: T.Tensor((32, 12), "uint32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
+ lane = T.cuda.lane_id()
+ SM = T.alloc_tensor(
+ ty_args=[T.Tensor((16, 4, 2), "uint32", scope="shared", layout=shared, align=16)]
+ )
+ for q in range(4):
+ n = lane + 32 * q
+ SM[n // 8, (n // 2) % 4, n % 2] = A[n // 8, (n // 2) % 4, n % 2]
+ T.cuda.cta_sync()
+ raw = T.alloc_tensor(ty_args=[T.Tensor((12,), "uint32", scope="local", layout="default")])
+ for i in T.unroll(12):
+ raw[i] = T.uint32(4294966519) # bit pattern of int32(-777)
+ R = T.decl_tensor(
+ raw.ptr_to([0]),
+ ty_args=[T.Tensor((2, 2, 8, 4), "uint32", scope="local", layout=registers)],
+ )
+ S_std = SM.rearrange("(r m) w k -> k m r w", r=8, m=2)
+ T.cuda.tile.apply(T.ptx[LDMATRIX], R, S_std, scope="warp")
+ for i in T.unroll(12):
+ Out[lane, i] = raw[i]
+
+ return kernel
+
+``S_std[k,m,r,w]`` is ``SM[2*r+m,w,k]`` and has strides ``(128,8,16,1)``.
+Both operands have shape ``(2,2,8,4)``: one repeat axis of length two,
+then the x2 matrix atom. For lane ``l`` the loaded values are::
+
+ raw[0] = SM[2*(l//4), l%4, 0]
+ raw[2] = SM[2*(l//4)+1, l%4, 0]
+ raw[8] = SM[2*(l//4), l%4, 1]
+ raw[10] = SM[2*(l//4)+1, l%4, 1]
+
+All other local slots retain their initialized sentinel. The transfer emits
+**two x2 instructions per warp**. It does not emit x4 and does not copy from
+a contiguous temporary fragment into the separated output slots.
+
+Run the complete factory example on an Ampere-or-later device:
+
+.. code-block:: python
+
+ import numpy as np
+ import tvm
+
+ compiled = tvm.tirx.compile(matrix_kernel(), target="cuda", pipeline="tirx")
+ x = np.arange(128, dtype="uint32").reshape(16, 4, 2)
+ out = tvm.runtime.tensor(np.zeros((32, 12), dtype="uint32"), tvm.cuda())
+ compiled(tvm.runtime.tensor(x, tvm.cuda()), out)
+ expected = np.full((32, 12), 4294966519, dtype="uint32")
+ for lane in range(32):
+ for k in range(2):
+ for matrix in range(2):
+ expected[lane, 8 * k + 2 * matrix] = x[2 * (lane // 4) + matrix, lane % 4, k]
+ np.testing.assert_array_equal(out.numpy(), expected)
+
+Transpose and stores
+---------------------
+
+``.trans`` transposes each 8x8 matrix of **16-bit elements**, not a matrix
+of b32 words. Register image element ``(row,2*word+half)`` comes from shared
+element ``(2*word+half,row)``. The register word's lane owner remains
+``4*row+word``. Reorganize operands with this image in mind; toggling
+``.trans`` is a semantic change, not a layout optimization hint.
+
+``stmatrix`` uses the corresponding register image as input and each
+provider's shared row address as output. Its supported architecture is
+checked by the PTX compiler. Synchronize before a load consumes another
+thread's shared stores, and after a store before other threads consume or
+reuse that shared storage. No fence or barrier is inserted by the pattern.
+
+An insufficiently aligned row, wrong register owner, overlapping loaded
+registers, or a partially active warp is rejected. Gaps **between** rows or
+register slots are legal; gaps **within** a shared 16-byte row are not.
+
+Implementation: ``python/tvm/backend/cuda/tile/patterns/ldmatrix.py``.
diff --git a/docs/tirx/tile_ops/min_max.rst b/docs/tirx/tile_ops/min_max.rst
new file mode 100644
index 0000000000..efefd59ac4
--- /dev/null
+++ b/docs/tirx/tile_ops/min_max.rst
@@ -0,0 +1,66 @@
+.. 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.
+
+min / max
+=========
+
+``min`` and ``max`` apply the native instruction independently at every
+register repeat. They do not perform a warp, block, or tensor reduction.
+Choose NaN, sign, absolute-value, packed-type and other supported semantics
+through the instruction spelling.
+
+Each scalar tensor operand ends in ``(1,)``. Binary and three-input forms
+are distinguished by the explicitly supplied operands, not by a layout
+search. ``max.f32`` with ``dst,a,b,c`` uses the registered three-input form;
+``dst,a,b`` uses its binary form. The selected target must support that form.
+Packed f16x2/bf16x2 variants use b32 carriers as in :doc:`add_sub_mul`.
+
+Example: three-input maximum
+----------------------------
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+
+
+ @T.function
+ def maximum3(
+ A: T.Tensor((4, 1), "float32"), B: T.Tensor((4, 1), "float32"), Out: T.Tensor((4, 1), "float32")
+ ):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=1))
+ a = T.alloc_local((4, 1), "float32", layout="default")
+ b = T.alloc_local((4, 1), "float32", layout="default")
+ T.cuda.tile.apply(T.ptx["ld.global.f32"], a, A)
+ T.cuda.tile.apply(T.ptx["ld.global.f32"], b, B)
+ T.cuda.tile.apply(T.ptx["max.f32"], a, a, b, T.float32(0))
+ T.cuda.tile.apply(T.ptx["st.global.f32"], Out, a)
+
+The result for finite inputs is ``maximum(maximum(A,B),0)``. The tile
+expansion issues **four three-input max instructions**, not eight binary
+ones. NaN and signed-zero behavior is that of the selected PTX form and must
+not be inferred from this finite-input reference.
+
+Changing ``max.f32`` to ``min.f32`` changes the operation on the same
+coordinate pattern. Reducing many values still requires explicit reduction
+loops, shuffle operations and/or shared communication around these atoms.
+In particular, a repeat axis of length four produces four independent
+results; it does not accumulate into a single output.
+
+All inputs share group/repeat extents unless provided as scalar broadcasts.
+Output locations must be distinct across repeats. Predication follows
+:doc:`mov`. Implementation: ``python/tvm/backend/cuda/tile/patterns/register.py``;
+native form definitions are in ``python/tvm/backend/cuda/ptx/table.py``.
diff --git a/docs/tirx/tile_ops/mma_sync.rst b/docs/tirx/tile_ops/mma_sync.rst
new file mode 100644
index 0000000000..1baa449c59
--- /dev/null
+++ b/docs/tirx/tile_ops/mma_sync.rst
@@ -0,0 +1,159 @@
+.. 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.
+
+mma.sync
+========
+
+The warp tile pattern supports the explicit
+``mma.sync.aligned.m16n8k{8,16}.row.col.f32.{f16,bf16}.{f16,bf16}.f32``
+forms registered by the PTX dialect. Each issue multiplies fixed register
+fragments and accumulates into four float32 registers per lane. The tile
+layer does not select K=8 versus K=16, change input type, load shared memory,
+or transpose an operand implicitly.
+
+Operand coordinates and ownership
+---------------------------------
+
+The call takes ``D,A,B,C``. After optional execution-group and batch axes,
+the canonical shapes are:
+
+.. list-table::
+ :header-rows: 1
+
+ * - Operand
+ - Shape
+ - Storage
+ * - D / tensor C
+ - ``(M_repeat,N_repeat,16,8)``
+ - float32
+ * - A
+ - ``(M_repeat,K_repeat,16,k/2)``
+ - b32 words containing two adjacent f16/bf16 K values
+ * - B
+ - ``(N_repeat,K_repeat,8,k/2)``
+ - The same packed representation
+
+B's logical atom axes are ``(N,K/2)``, even though the mathematical product
+uses B transposed. For lane ``l``, native register ``r`` corresponds to:
+
+.. list-table::
+ :header-rows: 1
+
+ * - Operand
+ - Atom coordinate
+ - Registers per lane
+ * - A
+ - ``(l//4+8*(r%2), l%4+4*(r//2))``
+ - k/4
+ * - B
+ - ``(l//4, l%4+4*r)``
+ - k/8
+ * - D / C
+ - ``(l//4+8*(r//2), 2*(l%4)+r%2)``
+ - 4
+
+Local offsets can be strided; lane ownership must match these coordinates.
+All 32 lanes execute the warp instruction. Warpgroup/CTA scope prepends a
+warp group axis; that axis is not an additional matrix dimension.
+
+Example: a 16x8 result with two K16 issues
+------------------------------------------
+
+The complete kernel below loads the exact packed A/B words each lane needs.
+The host prepares these words from ordinary row-major half matrices.
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+ from tvm.backend.cuda.axis import laneid
+ from tvm.tirx.layout import S, TileLayout
+
+ la = TileLayout(S[(1, 2, 2, 8, 2, 4) : (8, 4, 1, 4 @ laneid, 2, 1 @ laneid)])
+ lb = TileLayout(S[(1, 2, 8, 2, 4) : (4, 2, 4 @ laneid, 1, 1 @ laneid)])
+ ld = TileLayout(S[(1, 1, 2, 8, 4, 2) : (4, 4, 2, 4 @ laneid, 1 @ laneid, 1)])
+
+
+ @T.function
+ def warp_mma(
+ A: T.Tensor((1, 2, 16, 8), "uint32"),
+ B: T.Tensor((1, 2, 8, 8), "uint32"),
+ Out: T.Tensor((16, 8), "float32"),
+ ):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
+ lane = T.cuda.lane_id()
+ ar = T.alloc_local((8,), "uint32", layout="default")
+ br = T.alloc_local((4,), "uint32", layout="default")
+ dr = T.alloc_local((4,), "float32", layout="default")
+ av = ar.view(1, 2, 16, 8, layout=la)
+ bv = br.view(1, 2, 8, 8, layout=lb)
+ dv = dr.view(1, 1, 16, 8, layout=ld)
+ for kk, r in T.grid(2, 4):
+ ar[kk * 4 + r] = A[0, kk, lane // 4 + 8 * (r % 2), lane % 4 + 4 * (r // 2)]
+ for kk, r in T.grid(2, 2):
+ br[kk * 2 + r] = B[0, kk, lane // 4, lane % 4 + 4 * r]
+ T.cuda.tile.apply(
+ T.ptx["mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32"],
+ dv,
+ av,
+ bv,
+ T.float32(0),
+ scope="warp",
+ )
+ for r in range(4):
+ Out[lane // 4 + 8 * (r // 2), 2 * (lane % 4) + r % 2] = dr[r]
+
+.. code-block:: python
+
+ import numpy as np
+ import tvm
+
+ rng = np.random.default_rng(0)
+ a = (rng.integers(-4, 5, (16, 32)) / 4).astype("float16")
+ b = (rng.integers(-4, 5, (8, 32)) / 4).astype("float16")
+ ap = np.ascontiguousarray(a.reshape(1, 16, 2, 16).transpose(0, 2, 1, 3)).view("uint32")
+ bp = np.ascontiguousarray(b.reshape(1, 8, 2, 16).transpose(0, 2, 1, 3)).view("uint32")
+ compiled = tvm.tirx.compile(warp_mma, target="cuda", pipeline="tirx")
+ out = tvm.runtime.tensor(np.zeros((16, 8), dtype="float32"), tvm.cuda())
+ compiled(tvm.runtime.tensor(ap, tvm.cuda()), tvm.runtime.tensor(bp, tvm.cuda()), out)
+ np.testing.assert_allclose(
+ out.numpy(), a.astype("float32") @ b.astype("float32").T, rtol=2e-4, atol=2e-4
+ )
+
+``M_repeat=N_repeat=1`` and ``K_repeat=2``, so each warp issues exactly two
+K16 instructions. The first uses scalar C=0; the second uses the preceding
+D registers. No separate zeroing of dr or add instruction is required.
+
+Accumulation, transpose and alpha/beta
+--------------------------------------
+
+Repeat order is batch, M, N, then K. C can be a matching register tensor or
+an explicit scalar broadcast. Only the first K issue reads the original C;
+later issues read D. This is the documented recurrence, rather than an
+arbitrary traversal inferred from layouts.
+
+To transpose a logical input, prepare a view whose canonical A/B coordinates
+refer to the intended values while preserving the required lane map. A view
+alone cannot exchange values between lanes; an incompatible physical owner
+requires explicit preparation. ``alpha``/``beta`` scaling likewise belongs
+before or after MMA at the call site, with the intended rounding order.
+
+The tile call inserts no shared loads, ``ldmatrix``, communication, or
+barrier. Use :doc:`ldmatrix_stmatrix` to prepare fragments when that native
+shared/register image matches your operands. Sparse and other MMA shapes
+require their own explicit pattern; this page does not claim those forms.
+
+Implementation: ``python/tvm/backend/cuda/tile/patterns/mma.py``.
diff --git a/docs/tirx/tile_ops/mov.rst b/docs/tirx/tile_ops/mov.rst
new file mode 100644
index 0000000000..51b0e02b77
--- /dev/null
+++ b/docs/tirx/tile_ops/mov.rst
@@ -0,0 +1,67 @@
+.. 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.
+
+mov
+===
+
+``mov`` copies register bits or assembles/disassembles a native register
+group. It performs no numeric conversion and cannot change which thread
+owns a value. Use :doc:`cvt` for conversion and :doc:`ld_st` for memory.
+
+The scalar form ends both tensor operands in ``(1,)``. All earlier axes
+after the execution-group prefix are repeats. A scalar source expression
+can broadcast, for example to fill a register tile. Vector-group mov forms
+instead use the group length declared by that PTX form; each input/output
+word remains an explicit operand slot.
+
+Example: filling and copying separated register slots
+-----------------------------------------------------
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+ from tvm.tirx.layout import S, TileLayout
+
+ gapped = TileLayout(S[(4, 1) : (2, 1)])
+
+
+ @T.function
+ def register_move(Out: T.Tensor((4, 1), "uint32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=1))
+ a = T.alloc_local((4, 1), "uint32", layout="default")
+ b = T.alloc_local((4, 1), "uint32", layout=gapped)
+ T.cuda.tile.apply(T.ptx["mov.b32"], a, T.uint32(17))
+ T.cuda.tile.apply(T.ptx["mov.b32"], b, a)
+ T.cuda.tile.apply(T.ptx["st.global.b32"], Out, b)
+
+The first call emits four moves of the same constant. The second emits four
+moves to offsets 0, 2, 4, and 6 of ``b``; its gaps require no extra remap
+buffer. All four output values are 17. Ordinary compiler optimizations may
+subsequently eliminate redundant moves; the tile expansion itself selects
+no alternative instruction or ownership map.
+
+``pred`` is an ordinary runtime predicate. For a predicated output that must
+keep its old value, initialize it and pass ``preserve_dst=True``. A suppressed
+move is not a zero fill. Destination repeats must not overlap each other.
+
+For packed arithmetic, a bit reinterpretation such as
+``halves.view("uint32")`` can expose the existing low/high halfword storage
+without a mov. Use an explicit group mov when assembling separately stored
+values. Group arity and carrier dtype must agree with the selected native
+form; the pattern does not guess a packing arrangement.
+
+Implementation: ``python/tvm/backend/cuda/tile/patterns/register.py``.
diff --git a/docs/tirx/tile_primitives/reduction.rst b/docs/tirx/tile_ops/reduction.rst
similarity index 95%
rename from docs/tirx/tile_primitives/reduction.rst
rename to docs/tirx/tile_ops/reduction.rst
index 5877d944f5..b15c89101e 100644
--- a/docs/tirx/tile_primitives/reduction.rst
+++ b/docs/tirx/tile_ops/reduction.rst
@@ -27,4 +27,4 @@ fused ``tensorscalar_reduce`` and ``activation_reduce`` instructions. Their
axes and opcode are static qualifiers; tensors and optional reduction
workspace are ordinary Call arguments.
-See :doc:`../tile_primitives` for the complete contract.
+See :doc:`../tile_ops` for the complete contract.
diff --git a/docs/tirx/tile_ops/tcgen05_cp.rst b/docs/tirx/tile_ops/tcgen05_cp.rst
new file mode 100644
index 0000000000..844c63a496
--- /dev/null
+++ b/docs/tirx/tile_ops/tcgen05_cp.rst
@@ -0,0 +1,171 @@
+.. 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.
+
+tcgen05.cp
+==========
+
+``tcgen05.cp`` copies a descriptor-defined shared image into TMEM. It is
+issued by a single thread, not collectively by the warps whose TMEM lanes
+receive the data. Select shape, CTA group, multicast and optional
+decompression directly in the opcode.
+
+TMEM shape and replicas
+------------------------
+
+Operands are ``dst_tmem, src_shared``, both with suffix ``(row,b32_word)``
+and 32-bit storage elements. An earlier axis is an explicit repeat.
+
+.. list-table::
+ :header-rows: 1
+
+ * - Shape / multicast
+ - Atom extents
+ - TMEM lanes
+ * - ``128x256b``
+ - ``(128,8)``
+ - Consecutive 128 lanes
+ * - ``128x128b``
+ - ``(128,4)``
+ - Consecutive 128 lanes
+ * - ``4x256b``
+ - ``(4,8)``
+ - 0, 32, 64, 96 relative to the origin
+ * - ``64x128b.warpx2::02_13``
+ - ``(64,4)``
+ - 64 consecutive lanes, replicated 64 lanes apart
+ * - ``64x128b.warpx2::01_23``
+ - ``(64,4)``
+ - Logical row r maps to ``64*(r//32)+r%32``, replicated 32 lanes apart
+ * - ``32x128b.warpx4``
+ - ``(32,4)``
+ - 32 consecutive lanes, four replicas 32 lanes apart
+
+The destination layout must state these replicas explicitly; the pattern
+does not infer multicast. Columns are consecutive b32 cells with a 128-bit
+aligned origin. The whole footprint, including replicas, must stay inside
+the 128-lane, 512-column TMEM address space. A CTA-group-2 opcode makes both
+CTAs copy from their own corresponding shared storage into their own TMEM.
+The launch, allocation and commit must use the matching CTA-group protocol.
+
+Shared descriptor address walk
+-------------------------------
+
+Supply ``ldo_bytes``, ``sdo_bytes`` and ``swizzle=0/1/2/3``. Before swizzling,
+row ``r`` starts at ``(r//8)*sdo_bytes + (r%8)*(16 << swizzle)`` bytes
+relative to the descriptor origin. Without swizzling, word ``w`` adds
+``(w//4)*ldo_bytes + (w%4)*4``. With swizzling, words are contiguous and add
+``4*w``; hardware ignores LDO in that case.
+
+Descriptor strides are nonnegative 16-byte multiples below 2^18; origins
+are 16-byte aligned. Swizzled bases require 256/512/1024-byte period
+alignment for modes 1/2/3. The complete physical walk is checked, including
+the origin's swizzle phase. A 4-row logical region still needs a backing
+storage block large enough for the descriptor's eight-row organization.
+
+Example: one 128x128b copy and readback
+---------------------------------------
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+ from tvm.backend.cuda.axis import TLane, TCol
+ from tvm.tirx.layout import S, TileLayout
+
+ tmem_layout = TileLayout(S[(128, 4) : (1 @ TLane, 1 @ TCol)])
+
+
+ @T.function
+ def copy_to_tmem(A: T.Tensor((128, 4), "uint32"), Out: T.Tensor((128, 4), "uint32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=128))
+ tid = T.cuda.thread_idx("x")
+ warp = T.cuda.warp_id()
+ shared = T.alloc_shared((128, 4), "uint32", layout="default", align=16)
+ address = T.alloc_shared((1,), "uint32", layout="default")
+ mbar = T.alloc_shared((1,), "uint64", layout="default")
+ if warp == 0:
+ T.ptx.tcgen05.alloc.cta_group__1.sync.aligned.shared__cta.b32(
+ T.address_of(address), T.uint32(32)
+ )
+ if tid == 0:
+ T.ptx.mbarrier.init.shared.b64(mbar.ptr_to([0]), T.uint32(1))
+ for j in range(4):
+ shared[tid, j] = A[tid, j]
+ T.cuda.cta_sync()
+ T.ptx.fence.proxy.async_.shared__cta()
+ mem = T.cuda.decl_tmem(
+ address[0], ty_args=[T.Tensor((128, 4), "uint32", scope="tmem", layout=tmem_layout)]
+ )
+ if tid == 0:
+ T.cuda.tile.apply(
+ T.ptx["tcgen05.cp.cta_group::1.128x128b"],
+ mem,
+ shared,
+ ldo_bytes=16,
+ sdo_bytes=128,
+ swizzle=0,
+ )
+ T.ptx.tcgen05.commit.cta_group__1.mbarrier__arrive__one.shared__cluster.b64(
+ mbar.ptr_to([0])
+ )
+ T.cuda.mbarrier_wait(mbar.ptr_to([0]), 0)
+ T.ptx.tcgen05.fence__after_thread_sync()
+ word = T.alloc_local((1,), "uint32", layout="default")
+ for j in range(4):
+ T.ptx["tcgen05.ld.sync.aligned.32x32b.x1.b32"](word[0], address[0] + T.uint32(j))
+ T.ptx.tcgen05.wait__ld.sync.aligned()
+ Out[tid, j] = word[0]
+ T.cuda.cta_sync()
+ if warp == 0:
+ T.ptx.tcgen05.relinquish_alloc_permit.cta_group__1.sync.aligned()
+ T.ptx.tcgen05.dealloc.cta_group__1.sync.aligned.b32(address[0], T.uint32(32))
+
+One cp instruction copies 2048 bytes. Its commit triggers one barrier
+arrival, rather than a TMA-style expected-byte transaction. The scalar x1
+readback issues are caller code and independently expose each TMEM word.
+
+.. code-block:: python
+
+ import numpy as np
+ import tvm
+
+ x = np.arange(512, dtype="uint32").reshape(128, 4)
+ compiled = tvm.tirx.compile(copy_to_tmem, target="cuda", pipeline="tirx")
+ out = tvm.runtime.tensor(np.zeros_like(x), tvm.cuda())
+ compiled(tvm.runtime.tensor(x, tvm.cuda()), out)
+ np.testing.assert_array_equal(out.numpy(), x)
+
+Native FP4/FP6 decompression
+----------------------------
+
+Append ``.b8x16.b4x16_p64`` or ``.b8x16.b6x16_p32`` to the opcode. The b32
+footprint above stays unchanged: each 16-byte shared block contains sixteen
+packed values followed by eight bytes (FP4) or four bytes (FP6) of padding.
+The destination contains sixteen expanded bytes. Four-bit values gain two
+low zero bits; six-bit values gain two high zero bits. A 256-bit row contains
+two such blocks. The pattern emits one native decompression instruction per
+repeat, with no software unpacking loop.
+
+For instance, the example's opcode can be explicitly replaced with
+``tcgen05.cp.cta_group::1.128x128b.b8x16.b4x16_p64`` **when the source data
+has been prepared in that packed format**. The reference then compares the
+expanded values, not the original words or padding. A dtype name alone
+neither prepares the payload nor chooses this qualifier.
+
+Allocation, shared proxy fencing, election, commit, wait and deallocation
+are caller code. ``pred`` gates the cp but does not adjust a surrounding
+completion protocol. Implementation: ``patterns/tmem_copy.py`` and
+``patterns/descriptors.py`` under ``python/tvm/backend/cuda/tile``.
diff --git a/docs/tirx/tile_ops/tcgen05_ld_st.rst b/docs/tirx/tile_ops/tcgen05_ld_st.rst
new file mode 100644
index 0000000000..e97a4977d6
--- /dev/null
+++ b/docs/tirx/tile_ops/tcgen05_ld_st.rst
@@ -0,0 +1,174 @@
+.. 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.
+
+tcgen05.ld / tcgen05.st
+=======================
+
+These warp instructions transfer a fixed image between TMEM and registers.
+The opcode specifies the row/bit shape, repetition qualifier ``.xN``, and
+optional pack/unpack mode. The tile repeats that exact instruction over
+caller-supplied views. A load takes ``registers,tmem``; a store takes
+``tmem,registers``. Both require all 32 lanes of each participating warp.
+
+Native images
+--------------
+
+For the ordinary b32 forms both operands end in ``(row,column)``. Let ``l``
+be the lane ID and ``r`` a register index within one instruction:
+
+.. list-table::
+ :header-rows: 1
+
+ * - Shape
+ - Atom extents for xN
+ - Registers per lane
+ - Register r contains coordinate
+ * - ``32x32b``
+ - ``(32,N)``
+ - N
+ - ``(l,r)``
+ * - ``16x64b``
+ - ``(16,2N)``
+ - N
+ - ``(l//4+8*(l%2), l//2%2+2*r)``
+ * - ``16x128b``
+ - ``(16,4N)``
+ - 2N
+ - ``(l//4+8*(r%2), l%4+4*(r//2))``
+ * - ``16x256b``
+ - ``(16,8N)``
+ - 4N
+ - ``(l//4+8*(r//2%2), r%2+2*(l%4)+8*(r//4))``
+
+TMEM rows and columns are consecutive physical lanes/cells inside each atom.
+Each warp can access its own 32-lane slab; a 16-row instruction may select
+either half explicitly. Warp scope therefore requires a statically selected
+warp. Warpgroup/CTA scope exposes the warp as a group axis. Register slots
+may be separated as long as their owners match this table.
+
+Example: Layout B without changing its physical mapping
+--------------------------------------------------------
+
+Layout B represents logical ``(64,128)`` values across 128 physical lanes
+and 64 columns. Logical ``[i,j]`` is at lane ``i+64*(j//64)``, column
+``j%64``. The following round trip isolates transfers from the MMA producer:
+it uses one CTA and writes/reads this layout with explicit st/ld instructions.
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+ from tvm.backend.cuda.layout import tmem_datapath_layout
+
+ layout_b = tmem_datapath_layout("B", 64, 128)
+
+
+ @T.function
+ def tmem_roundtrip(A: T.Tensor((128, 64), "float32"), Out: T.Tensor((128, 64), "float32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=128))
+ tid = T.cuda.thread_idx("x")
+ warp = T.cuda.warp_id()
+ address = T.alloc_shared((1,), "uint32", layout="default")
+ if warp == 0:
+ T.ptx.tcgen05.alloc.cta_group__1.sync.aligned.shared__cta.b32(
+ T.address_of(address), T.uint32(64)
+ )
+ T.cuda.cta_sync()
+ mem = T.cuda.decl_tmem(
+ address[0], ty_args=[T.Tensor((64, 128), "float32", scope="tmem", layout=layout_b)]
+ )
+ src = T.alloc_tcgen05_ldst_frag("32x32b", (64, 128), "float32")
+ dst = T.alloc_tcgen05_ldst_frag("32x32b", (64, 128), "float32")
+ local_src = src.local()
+ for j in range(64):
+ local_src[j] = A[tid, j]
+ mv = mem.rearrange("(wh r) (wl c) -> (wl wh) r c", wh=2, r=32, wl=2)
+ sv = src.rearrange("(wh r) (wl c) -> (wl wh) r c", wh=2, r=32, wl=2)
+ dv = dst.rearrange("(wh r) (wl c) -> (wl wh) r c", wh=2, r=32, wl=2)
+ T.cuda.tile.apply(T.ptx["tcgen05.st.sync.aligned.32x32b.x64.b32"], mv, sv, scope="warpgroup")
+ T.ptx.tcgen05.wait__st.sync.aligned()
+ T.cuda.cta_sync()
+ T.cuda.tile.apply(T.ptx["tcgen05.ld.sync.aligned.32x32b.x64.b32"], dv, mv, scope="warpgroup")
+ T.ptx.tcgen05.wait__ld.sync.aligned()
+ local_dst = dst.local()
+ for j in range(64):
+ Out[tid, j] = local_dst[j]
+ T.cuda.cta_sync()
+ if warp == 0:
+ T.ptx.tcgen05.relinquish_alloc_permit.cta_group__1.sync.aligned()
+ T.ptx.tcgen05.dealloc.cta_group__1.sync.aligned.b32(address[0], T.uint32(64))
+
+The views have shape ``(4,32,64)`` with ``warp=2*wl+wh``. There are no repeat
+axes: each of the four warps issues one x64 store and one x64 load. No
+physical rearrangement or fragment copy is inserted. The helper allocates
+the register image; it does not select the opcode.
+
+.. code-block:: python
+
+ import numpy as np
+ import tvm
+
+ x = np.arange(8192, dtype="float32").reshape(128, 64)
+ compiled = tvm.tirx.compile(tmem_roundtrip, target="cuda", pipeline="tirx")
+ out = tvm.runtime.tensor(np.zeros_like(x), tvm.cuda())
+ compiled(tvm.runtime.tensor(x, tvm.cuda()), out)
+ np.testing.assert_array_equal(out.numpy(), x)
+
+Pack/unpack is a different TMEM image
+-------------------------------------
+
+``tcgen05.ld...pack::16b.b32`` and ``tcgen05.st...unpack::16b.b32`` keep
+the ordinary b32 register image. The TMEM operand instead has suffix
+``(row,column,half)`` with ``half=2`` and a 16-bit storage type. The two
+halves of register word ``[row,column]`` correspond to the **low 16 bits of
+two different TMEM cells**, at columns ``2*column`` and ``2*column+1``.
+
+In 16-bit layout units, the column and half strides are respectively
+``4 @ TCol`` and ``2 @ TCol``. Using strides 2 and 1 would instead select
+both halves of one dense cell and would not satisfy the instruction. Hardware
+performs the pack/unpack; the tile emits no software conversion. The upper
+16 bits left by an unpack are not part of its defined 16-bit payload.
+
+Split accesses: 16x32bx2
+------------------------
+
+This form uses register suffix ``(access,row,column)=(2,16,N)``; lane ``l``
+owns ``(l//16,l%16,r)``. Supply two separate TMEM regions and the explicit
+native immediate ``imm_half_splitoff``. Each TMEM suffix is ``(16,N)``
+or ``(16,N,2)`` for packing. For example, for an already prepared
+warpgroup view ``mem`` with shape ``(4,32,16)``::
+
+ T.cuda.tile.apply(
+ T.ptx["tcgen05.ld.sync.aligned.16x32bx2.x4.b32"],
+ registers, mem[:, :16, :4], mem[:, 16:32, 8:12],
+ imm_half_splitoff=(16 << 16) + 8, scope="warpgroup",
+ )
+
+Both halves share the same TMEM backing allocation. The second encoded
+address equals the first plus the immediate; no extra sixteen-row shift is
+implied by lane 16. For stores, operand order is ``mem_low,mem_high,registers``.
+One instruction handles both accesses for each repeat.
+
+Completion and producer dependencies
+-------------------------------------
+
+Call ``tcgen05.wait::ld`` before consuming loaded registers and
+``tcgen05.wait::st`` before relying on stores. For an MMA/cp producer,
+complete its commit/mbarrier protocol and required thread-sync fence before
+issuing a TMEM load; ``wait::ld`` does not complete a preceding MMA.
+Allocation, deallocation, fences and synchronization remain explicit.
+
+Implementation: ``python/tvm/backend/cuda/tile/patterns/tmem.py``.
diff --git a/docs/tirx/tile_ops/tcgen05_mma.rst b/docs/tirx/tile_ops/tcgen05_mma.rst
new file mode 100644
index 0000000000..6a8a966fc0
--- /dev/null
+++ b/docs/tirx/tile_ops/tcgen05_mma.rst
@@ -0,0 +1,230 @@
+.. 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.
+
+tcgen05.mma
+===========
+
+Choose dense, weight-stationary ``.ws``, or ``.block_scale`` in the PTX
+opcode, together with CTA group and kind. Supply
+``instruction_shape=(M,N,K)`` explicitly. M is the total for the CTA group,
+so M=128 with ``cta_group::2`` gives 64 local M rows per CTA. The instruction
+is issued by an explicitly elected thread; group scope does not elect one.
+
+Operands and repeat axes
+------------------------
+
+Dense/.ws calls take ``D,A,B``. Block scale adds ``SFA,SFB``. After any
+execution-group and batch prefix, the canonical shapes are:
+
+.. list-table::
+ :header-rows: 1
+
+ * - Operand
+ - Shape, with g = CTA group
+ - Storage
+ * - D
+ - ``(M_repeat,N_repeat,M/g,N)``
+ - TMEM
+ * - A
+ - ``(M_repeat,K_repeat,M/g,K)``
+ - Shared or TMEM
+ * - B
+ - ``(N_repeat,K_repeat,N/g,K)``
+ - Shared
+ * - SFA
+ - ``(M_repeat,K_repeat,M/g,scale_vec)``
+ - TMEM
+ * - SFB
+ - ``(N_repeat,K_repeat,N,scale_vec)``
+ - TMEM, including full N in each CTA of a group-2 operation
+
+M64 .ws with TMEM A adds an atom bank axis: A ends in ``(2,64,K)`` rather
+than ``(64,K)``. D and A use the selected native TMEM datapath; the pure
+``tmem_mma_operand_layout`` helper constructs that layout from explicit
+shape/type/mode parameters. It does not choose the MMA opcode.
+
+D stores one accumulator per b32 TMEM cell. Float16 accumulation uses only
+the low half, hence a TCol stride of two in 16-bit element units. A in TMEM
+is packed according to its input format instead. ``a_format``/``b_format``
+default to tensor dtype; supply ``"tf32"`` explicitly for TF32 instructions
+whose storage dtype is float32.
+
+Shared descriptors are byte-stride contracts
+--------------------------------------------
+
+Supply ``a_descriptor=(ldo_bytes,sdo_bytes,swizzle)`` for shared A and
+``b_descriptor=...`` for B. ``a_major``/``b_major`` are explicitly ``"k"``
+or ``"mn"``. For swizzle modes 0..3, a descriptor row is ``16 << swizzle``
+bytes and groups contain eight rows. K-major rows enumerate M/N; MN-major
+rows enumerate K. The next row group and next contiguous chunk use:
+
+.. list-table::
+ :header-rows: 1
+
+ * - Mode
+ - Row-group stride
+ - Contiguous-chunk stride
+ * - K-major
+ - SDO
+ - LDO; swizzled instruction atoms fit within one row
+ * - MN-major, swizzled
+ - SDO
+ - LDO
+ * - MN-major, unswizzled
+ - LDO
+ - SDO
+
+The final physical address is the selected swizzle of this byte walk. The
+origin is inverse-swizzled before encoding so it is not permuted twice.
+Descriptor strides are 16-byte multiples below 2^18, origins are 16-byte
+aligned, and backing allocations align to the full swizzle period. Layouts
+are checked against these fields; no descriptor-building policy is inferred.
+
+``swizzle=4`` is the separate 128B-span/32B-atomicity mode for 32-bit
+MN-major operands. It groups **four** K rows of 128 bytes per SDO and uses
+LDO for the next 128-byte M/N chunk. Its permutation is
+``offset ^ ((offset >> 2) & 96)``, with a 512-byte period. It is not an
+alternative row pitch for ordinary mode 3, and is not this pattern's
+K-major mode. TF32 MN-major requires this mode; other MN-major element
+widths use modes 0..3.
+
+Example: four K16 issues for a 128x64x64 product
+------------------------------------------------
+
+The input B is stored as ``(N,K)``. Shared initialization uses ordinary
+scalar stores so the layout and MMA behavior can be inspected independently
+of a TMA pipeline. The descriptor fields and the K decomposition are visible
+at the call site.
+
+.. code-block:: python
+
+ from tvm.script import tirx as T
+ from tvm.backend.cuda.layout import mma_shared_layout, tmem_mma_operand_layout
+
+ al = mma_shared_layout("float16", 3, (128, 64))
+ bl = mma_shared_layout("float16", 3, (64, 64))
+ dl = tmem_mma_operand_layout("D", (128, 64), "float32", M=128, cta_group=1)
+
+
+ @T.function
+ def tensor_mma(
+ A: T.Tensor((128, 64), "float16"),
+ B: T.Tensor((64, 64), "float16"),
+ Out: T.Tensor((128, 64), "float32"),
+ ):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=128))
+ tid = T.cuda.thread_idx("x")
+ warp = T.cuda.warp_id()
+ a = T.alloc_shared((128, 64), "float16", layout=al, align=1024)
+ b = T.alloc_shared((64, 64), "float16", layout=bl, align=1024)
+ address = T.alloc_shared((1,), "uint32", layout="default")
+ mbar = T.alloc_shared((1,), "uint64", layout="default")
+ if warp == 0:
+ T.ptx.tcgen05.alloc.cta_group__1.sync.aligned.shared__cta.b32(
+ T.address_of(address), T.uint32(64)
+ )
+ if tid == 0:
+ T.ptx.mbarrier.init.shared.b64(mbar.ptr_to([0]), T.uint32(1))
+ for j in range(64):
+ a[tid, j] = A[tid, j]
+ if tid < 64:
+ b[tid, j] = B[tid, j]
+ T.cuda.cta_sync()
+ T.ptx.fence.proxy.async_.shared__cta()
+ d = T.cuda.decl_tmem(
+ address[0], ty_args=[T.Tensor((128, 64), "float32", scope="tmem", layout=dl)]
+ )
+ av = a.rearrange("m (kr k) -> kr m k", k=16).view(1, 4, 128, 16)
+ bv = b.rearrange("n (kr k) -> kr n k", k=16).view(1, 4, 64, 16)
+ if tid == 0:
+ T.cuda.tile.apply(
+ T.ptx["tcgen05.mma.cta_group::1.kind::f16"],
+ d.view(1, 1, 128, 64),
+ av,
+ bv,
+ instruction_shape=(128, 64, 16),
+ a_descriptor=(16, 1024, 3),
+ b_descriptor=(16, 1024, 3),
+ accum=False,
+ )
+ T.ptx.tcgen05.commit.cta_group__1.mbarrier__arrive__one.shared__cluster.b64(
+ mbar.ptr_to([0])
+ )
+ T.cuda.mbarrier_wait(mbar.ptr_to([0]), 0)
+ T.ptx.tcgen05.fence__after_thread_sync()
+ frag = T.alloc_tcgen05_ldst_frag("32x32b", (128, 64), "float32")
+ T.cuda.tile.apply(
+ T.ptx["tcgen05.ld.sync.aligned.32x32b.x64.b32"],
+ frag.view(4, 32, 64),
+ d.view(4, 32, 64),
+ scope="warpgroup",
+ )
+ T.ptx.tcgen05.wait__ld.sync.aligned()
+ local = frag.local()
+ for j in range(64):
+ Out[tid, j] = local[j]
+ T.cuda.cta_sync()
+ if warp == 0:
+ T.ptx.tcgen05.relinquish_alloc_permit.cta_group__1.sync.aligned()
+ T.ptx.tcgen05.dealloc.cta_group__1.sync.aligned.b32(address[0], T.uint32(64))
+
+.. code-block:: python
+
+ import numpy as np
+ import tvm
+
+ rng = np.random.default_rng(0)
+ a = (rng.integers(-4, 5, (128, 64)) / 4).astype("float16")
+ b = (rng.integers(-4, 5, (64, 64)) / 4).astype("float16")
+ compiled = tvm.tirx.compile(tensor_mma, target="cuda", pipeline="tirx")
+ out = tvm.runtime.tensor(np.zeros((128, 64), dtype="float32"), tvm.cuda())
+ compiled(tvm.runtime.tensor(a, tvm.cuda()), tvm.runtime.tensor(b, tvm.cuda()), out)
+ np.testing.assert_allclose(
+ out.numpy(), a.astype("float32") @ b.astype("float32").T, rtol=2e-4, atol=2e-4
+ )
+
+The issuing thread emits four identical-kind MMA instructions, one for
+each explicit K repeat. ``accum=False`` applies to the first K issue only;
+later issues enable D input. One commit covers that sequence. The four warps
+then each execute one x64 load and wait before consuming their registers.
+
+Weight-stationary and block scaling
+-----------------------------------
+
+``.ws`` is selected in the opcode, not inferred from packed C. M64 .ws uses
+Layout E with an explicit bank dimension for TMEM A; M128 .ws uses Layout D.
+The current pattern requires CTA group 1 and N=64/128/256 for .ws. Splitting
+an M64 N128 operation into two N64 issues changes each instruction's bank
+boundary from 64 to 32 columns. Preserve the intended result with explicit
+B preparation and a matching C view; cutting two consecutive logical N
+slices is insufficient. ``test_tcgen05_mma_ws_split_n.py`` gives the
+complete numerical example with distinct data in the two A banks.
+
+Block-scaled opcodes specify ``.kind`` and ``.scale_vec::1X/2X/4X``. Prepare
+SFA/SFB in TMEM: scale row r uses lane ``r%32`` and byte-column
+``4*(r//32)+sf``, with four replicas spaced 32 lanes apart. Scale-factor IDs
+are encoded from each tensor's own byte-column origin. Supplying scales
+neither allocates them nor changes the explicit input kind/scale-vector mode.
+
+``neg_a``, ``neg_b``, ``sat_d``, dense ``disable_output_lane`` masks,
+``zero_column_mask`` for .ws, and ``pred`` are explicit parameters where
+applicable. The tile does not initialize barriers, elect a thread, perform
+TMA, prepare scale factors, commit or wait. Keep the shared/TMEM allocations
+alive until their consumers complete.
+
+Implementation: ``patterns/tcgen05_mma.py`` and ``patterns/descriptors.py``
+under ``python/tvm/backend/cuda/tile``.
diff --git a/docs/tirx/tile_primitives.rst b/docs/tirx/tile_primitives.rst
deleted file mode 100644
index 4dd6a64ea5..0000000000
--- a/docs/tirx/tile_primitives.rst
+++ /dev/null
@@ -1,140 +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.
-
-Tensor Instructions
-===================
-
-Tensor instructions expose a backend instruction over tensors with layouts::
-
- from tvm.script import tirx as T
-
- 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)
-
-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.
-
-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.
-
-``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.
-
-CUDA instructions
------------------
-
-* ``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
-------------------
-
-.. toctree::
- :maxdepth: 1
-
- 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
deleted file mode 100644
index 000d2cec2b..0000000000
--- a/docs/tirx/tile_primitives/copy.rst
+++ /dev/null
@@ -1,36 +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 Tensor Loads, Stores, and Moves
-====================================
-
-``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.
-
-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/reg
- copy/ldstmatrix
diff --git a/docs/tirx/tile_primitives/copy/ldstmatrix.rst b/docs/tirx/tile_primitives/copy/ldstmatrix.rst
deleted file mode 100644
index 1f3f68b645..0000000000
--- a/docs/tirx/tile_primitives/copy/ldstmatrix.rst
+++ /dev/null
@@ -1,241 +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.
-
-ldmatrix / stmatrix
-===================
-
-``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**. An incompatible layout is rejected; use
-:doc:`reg` for ordinary memory instructions. Source:
-``python/tvm/backend/cuda/tile/copy/ld_stmatrix.py``.
-
-What it accepts
----------------
-
-The predicate is lean — scope, a valid copy, and a register↔shared pair:
-
-.. code-block:: python
-
- def _is_ldstmatrix(op_call, sctx):
- if not sctx.is_target("cuda"):
- return False, "non-cuda target"
- if sctx.scope_kind not in ("warp", "warpgroup", "cta"):
- return False, f"unsupported exec_scope {sctx.scope_kind} (need warp, warpgroup, or cta)"
- for check in (
- lambda: _all_threads_active(sctx),
- lambda: _is_valid_copy(op_call, sctx),
- lambda: _scope_allowed(op_call, sctx, allowed_pairs=_REG_SMEM_PAIRS), # (local, shared*)
- ):
- ok, msg = check()
- if not ok:
- return False, msg
- return True, None
-
-The layout fit is validated during emission. The selected matrix instruction
-must match an m8n8 fragment, or lowering reports an error:
-
-.. code-block:: python
-
- # _emit: try the widest matrix count that fits, else decline
- for num in (4, 2, 1):
- chosen = _try_num(r, s, num)
- if chosen is not None:
- break
- if chosen is None:
- fail("ldstmatrix layout doesn't fit any num ∈ {4,2,1}")
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target / scope
- - ``cuda``; ``warp`` / ``warpgroup`` / ``cta`` (needs a full warp), all active
- * - memory pair
- - ``_REG_SMEM_PAIRS`` = ``(local, shared*)`` / ``(shared*, local)``
- * - dtype
- - the predicate only requires matching source/destination dtypes; it does
- not inspect their bit width. The emitter always uses ``.b16`` and its
- layout math assumes eight 16-bit elements per shared row, so callers use
- 16-bit element layouts. Each ``.x1`` tile gives every lane one 32-bit
- register containing two 16-bit elements
- * - layout fit
- - both operands regroup to ``[T/32, 8, 4, M/(2·num), num, 2]`` with the
- register side equal to the m8n8 fragment pattern and the shared side row- or
- column-major with 16-B-aligned tile strides (``_try_num``), for some
- ``num ∈ {4, 2, 1}``
- * - thread layout
- - neither side may have replica axes; the register side has exactly one
- thread-axis name from ``laneid``, ``tid_in_wg``, or ``tx``. Its total
- thread extent is divisible by 32; an outer multi-warp iter, when present,
- has register stride 32
- * - shared swizzle
- - a ``ComposeLayout`` is accepted only when ``per_element >= 3``, preserving
- 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,
-``S[(8,4,2,2):(4@laneid, 1@laneid, 2, 1)]``):
-
-.. code-block:: python
-
- from tvm.tirx.layout import S, TileLayout
- from tvm.backend.cuda.axis import laneid
-
- num = 2
- M, N = 8, num * 8
- r_layout = TileLayout(S[(8, 4, num, 2) : (4 @ laneid, 1 @ laneid, 2, 1)])
- s_layout = TileLayout(S[(8, 4, num, 2) : (num * 8, 2, 8, 1)]) # row-major
- full = (slice(0, 8), slice(0, 4), slice(0, num), slice(0, 2))
-
-
- @Tx.function
- def kernel(A: Tx.Tensor((M, N), "float16"), B: Tx.Tensor((M, N), "float16")):
-
- Tx.device_entry(launch=Tx.cuda.LaunchConfig(grid=(1,), block=(32,)))
- _index = Tx.cuda.block_idx("x")
- _lane = Tx.cuda.lane_id()
- tid = Tx.cuda.thread_idx("x")
- A_smem = Tx.alloc_tensor(
- ty_args=[Tx.Tensor((8, 4, num, 2), "float16", scope="shared", layout=s_layout)]
- )
- # ... stage A into A_smem (row = tid//4, cp = tid%4) ...
- Tx.cuda.cta_sync()
- R = Tx.alloc_tensor(
- ty_args=[Tx.Tensor((8, 4, num, 2), "float16", scope="local", layout=r_layout)]
- )
- Tx.cuda.tile.ldmatrix(R[full], A_smem[full], scope="warp") # shared -> register (ldmatrix)
- # ... write R back out to B ...
-
-Algorithm
----------
-
-**1. Regroup both layouts to the matrix geometry.** ``_try_num(r, s, num)`` groups
-each layout's iters into ``[T/32, 8, 4, M/(2·num), num, 2]``: the warp-replication
-outer, the **8** rows of a tile, the **4** lane-column-pairs, ``m_outer`` tiles
-along M, the ``num`` tiles, and the inner **2** (the ``.b16`` element pair). If the
-group fails, the layout isn't a fragment → ``None``.
-
-**2. The register side must be the exact m8n8 fragment.** The 8/4/2 register
-strides must be ``(4, 1, 1)`` — i.e. the canonical ldmatrix fragment where lane
-``i`` holds row ``i//4``, column-pair ``i%4``:
-
-.. code-block:: python
-
- r8, r4, _r_num_iters, r2 = rs
- if (r8, r4, r2) != (4, 1, 1):
- return None
-
-**3. The shared side decides** ``trans`` **and the per-tile stride** ``p``.
-Row-major shared (``s4, s2 == 2, 1``, ``s8`` a positive multiple of 8) → plain
-``ldmatrix`` with ``p = s8``; column-major (``s8 == 1``, ``s4 == 2·s2``,
-``s2`` a multiple of 8) → the ``.trans`` form:
-
-.. code-block:: python
-
- if (s4, s2) == (2, 1) and s8 > 0 and s8 % 8 == 0:
- return (rg, rsep, sg, ssep, False, s8, num) # trans=False, p=s8
- if s8 == 1 and s2 > 0 and s2 % 8 == 0 and s4 == 2 * s2:
- return (rg, rsep, sg, ssep, True, s2, num) # trans=True, p=s2
-
-The 8-multiple checks keep every shared-memory row start and every ``m_outer``
-advance aligned to 16 bytes (8 fp16). The warp distributes that row across
-lanes; each lane's destination or source register word is 32 bits.
-
-**4. Emit one instruction per** ``m_outer`` **tile group.** Each lane contributes
-its shared address (tile offset + ``(laneid % 8) · p``) and ``num`` register
-words:
-
-.. code-block:: python
-
- for mm in Tx.unroll(m_outer):
- smem_ptr = _ptr_off(s_buf.ptr_to(s_zero), _smem_off(mm, tile_off + (laneid % 8) * p))
- words = [r_local[...] for i in range(num)]
- chain = f"{direction}matrix.sync.aligned.m8n8.x{num}{trans_seg}.shared.b16"
- if direction == "ld":
- Tx.ptx[chain](*words, smem_ptr)
- else:
- Tx.ptx[chain](smem_ptr, *words) # stmatrix takes the address first
-
-This lowerer unrolls the small ``m_outer`` loop.
-
-Generated TIRx IR
------------------
-
-For the demo (``num = 2``, ``M = 8`` ⇒ ``m_outer = 1``):
-
-.. code-block:: python
-
- for mm in Tx.unroll(1):
- Tx.ptx["ldmatrix.sync.aligned.m8n8.x2.shared.b16"](
- r_local[0], r_local[2], smem_ptr)
-
-Generated CUDA
---------------
-
-.. code-block:: c++
-
- __forceinline__ __device__ void tvm_builtin_ptx_ldmatrix_sync_aligned_m8n8_x2_shared_b16(
- uint32_t& __d0, uint32_t& __d1, uint32_t __a) {
- asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];"
- : "=r"(__d0), "=r"(__d1) : "r"(__a));
- }
- // call site (per lane):
- tvm_builtin_ptx_ldmatrix_sync_aligned_m8n8_x2_shared_b16(
- r_local_ptr[0], r_local_ptr[2], smem_addr);
-
-``num = 2`` becomes ``.x2`` with two destination registers; the warp's 32 lanes
-cooperatively supply the 8 source rows and receive the shuffled fragment.
-
-How inputs change the algorithm
--------------------------------
-
-``num`` (the matrix count that fits) selects the instruction width and the number
-of register words; ``trans`` (set by the shared layout) selects the transposing
-form:
-
-.. list-table::
- :header-rows: 1
- :widths: 20 20 60
-
- * - input
- - emitted
- - PTX
- * - ``num = 1``
- - ``.x1``
- - ``ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];``
- * - ``num = 2``
- - ``.x2``
- - ``ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];``
- * - ``num = 4``
- - ``.x4``
- - ``ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];``
- * - ``trans = True``
- - ``.trans``
- - ``ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0, %1}, [%2];``
-
-A larger M raises ``m_outer`` (more unrolled instructions per lane); the ``st``
-direction emits ``stmatrix`` with the same width/trans logic. If no ``num`` fits,
-the copy is handled by the :doc:`reg` path in ``vec_auto`` instead.
diff --git a/docs/tirx/tile_primitives/copy/reg.rst b/docs/tirx/tile_primitives/copy/reg.rst
deleted file mode 100644
index bf20f53341..0000000000
--- a/docs/tirx/tile_primitives/copy/reg.rst
+++ /dev/null
@@ -1,222 +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.
-
-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/copy/vec_auto_reg.py``.
-
-What it accepts
----------------
-
-.. code-block:: python
-
- def _is_reg_copy(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),
- lambda: _is_valid_copy(op_call, sctx),
- lambda: _scope_allowed(op_call, sctx, allowed_pairs=_REG_PAIRS),
- lambda: _r_side_layout_valid(op_call, sctx), # the register operand
- lambda: _s_side_slice_ok(op_call), # the other operand
- ):
- ok, msg = check()
- if not ok:
- return False, msg
- return True, None
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target / scope
- - ``cuda``; ``thread`` / ``warp`` / ``warpgroup`` / ``cta`` with all threads
- active
- * - memory pair
- - ``_REG_PAIRS`` = ``(local, shared*)`` / ``(shared*, local)`` /
- ``(local, global)`` / ``(global, local)`` — exactly one side is ``local``
- * - register layout
- - ``_r_side_layout_valid``: the ``local`` operand is a non-swizzle
- ``TileLayout`` whose thread axes have a register-level subscope no wider
- than the exec scope and a **zero sliced thread offset** (the region does
- not start partway through a thread partition)
- * - other side
- - ``_s_side_slice_ok``: the ``shared*`` / ``global`` operand slices cleanly to
- 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
-contiguous elements). From ``test_reg.py``:
-
-.. code-block:: python
-
- from tvm.tirx.layout import S, TileLayout
- from tvm.backend.cuda.axis import laneid
-
- shape, dtype = (32, 8), "float32"
- r_layout = TileLayout(S[shape : (1 @ laneid, 1)]) # lane i -> row i, 8 regs
- s_layout = TileLayout(S[shape])
- fs = (slice(0, 32), slice(0, 8))
-
-
- @Tx.function
- def kernel(B: Tx.Tensor(shape, dtype)):
-
- Tx.device_entry(launch=Tx.cuda.LaunchConfig(grid=(1,), block=(32,)))
- _index = Tx.cuda.block_idx("x")
- _lane = Tx.cuda.lane_id()
- tid = Tx.cuda.thread_idx("x")
- A_smem = Tx.alloc_tensor(ty_args=[Tx.Tensor(shape, dtype, scope="shared", layout=s_layout)])
- for kk in range(8):
- A_smem[tid, kk] = Tx.cast(tid * 100 + kk + 1, dtype)
- Tx.cuda.cta_sync()
- R = Tx.alloc_tensor(ty_args=[Tx.Tensor(shape, dtype, scope="local", layout=r_layout)])
- Tx.cuda.tile.ld(R[fs], A_smem[fs], scope="warp") # shared -> register (this dispatch)
- # ... clear A_smem, cta_sync ...
- 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 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.
-
-**2. Linearize and choose the vector width.** The per-thread elements are
-flattened into ``(extent, stride)`` atoms; the vector width is chosen widest-first
-(``128 → … → 8`` bits) so the contiguous tail divides it and the outer atom
-strides + base offsets are aligned. Crucially the **thread-axis strides are
-excluded** from this alignment check (they live in *partition-coordinate* space —
-which thread owns which element — and never appear in a single thread's physical
-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.cuda.lane_id()`` etc.),
-and the register side is a flat per-thread ``local`` tensor. The emit is a serial
-loop to limit generated code size:
-
-.. code-block:: python
-
- r_local = r_buf.local() # raw per-thread physical span
- r_words = r_local.view(reg_dtype) # PTX register-container view
- for f in range(total_outer):
- ds, dr = _outer_const_offsets(outer, f) # shared / reg deltas
- s_ptr = _ptr_off(s_buf.ptr_to(s_zero_indices), _s_iter_off(f, ds, s_off))
- r_w = (r_off_base + dr) * words_per_elem // elems_per_word
- if r_is_src:
- Tx.ptx[st_chain](s_ptr, *[r_words[r_w + i] for i in range(lanes)])
- else:
- Tx.ptx[ld_chain](*[r_words[r_w + i] for i in range(lanes)], s_ptr)
-
-``dr`` and ``r_off_base`` are physical storage offsets, so the register alias
-must use the raw-span ``local()`` view. This remains correct when storage
-iterators are permuted or leave gaps; every physical slot through
-``layout.storage().span()`` is directly addressable.
-
-Generated TIRx IR
------------------
-
-``LowerTIRx`` turns the shared → register copy into a per-thread loop over the
-8-element local bundle (trimmed):
-
-.. code-block:: python
-
- r_local = Tx.decl_tensor(
- R.data, ty_args=[Tx.Tensor((8,), scope="local")]
- ) # 8 fp32 elements / lane
- r_words = r_local.view("uint32")
- for f in range(2): # outer = 8 / vec 4
- s_ptr = pointer_offset(A_smem, ...) # this lane's row
- r_w = f * 4
- Tx.ptx.ld.shared.v4.u32(
- r_words[r_w], r_words[r_w + 1], r_words[r_w + 2], r_words[r_w + 3], s_ptr
- )
-
-(The register-to-shared copy uses ``Tx.ptx.st.shared.v4.u32`` with the pointer
-first, followed by the four register-container operands.)
-
-Generated PTX instruction
--------------------------
-
-The CUDA code generator emits the vector load directly::
-
- ld.shared.v4.u32 {r0, r1, r2, r3}, [s_ptr];
-
-The reverse direction emits ``st.shared.v4.u32 [s_ptr], {r0, r1, r2, r3};``.
-
-Each lane copies its own 8 elements as 2 × 128-bit transfers; no cross-lane
-addressing appears because the thread partition was resolved away at lowering.
-
-How inputs change the algorithm
--------------------------------
-
-The register layout's **per-thread element count** (the non-thread extents — here
-``k``) and the **dtype** set the local element count, PTX container count, vector
-width, and round count:
-
-.. list-table::
- :header-rows: 1
- :widths: 18 18 22 22 20
-
- * - dtype
- - ``k``
- - elements / lane
- - ``vec``
- - ``outer = k / vec``
- * - ``float32``
- - 8
- - 8
- - 4
- - 2
- * - ``float32``
- - 16
- - 16
- - 4
- - 4
- * - ``float16``
- - 8
- - 8
- - 8
- - 1
- * - ``float16``
- - 16
- - 16
- - 8
- - 2
-
-The copy uses a 128-bit ``v4.u32`` load or store when the contiguous tail allows.
-The **scope** sets the thread axis (``warp`` → ``laneid``, ``cta`` → ``tx``, …) the
-register layout must use; a different R layout (e.g. a strided or multi-row
-ownership) changes which elements each lane holds and therefore the atom list and
-``outer``.
diff --git a/docs/tirx/tile_primitives/copy_async.rst b/docs/tirx/tile_primitives/copy_async.rst
deleted file mode 100644
index aa4b64b705..0000000000
--- a/docs/tirx/tile_primitives/copy_async.rst
+++ /dev/null
@@ -1,45 +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.
-
-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
-
- copy_async/ldgsts
- copy_async/tma
- copy_async/dsmem
- copy_async/tcgen05_cp
- copy_async/tcgen05_ldst
diff --git a/docs/tirx/tile_primitives/copy_async/dsmem.rst b/docs/tirx/tile_primitives/copy_async/dsmem.rst
deleted file mode 100644
index 73874f3179..0000000000
--- a/docs/tirx/tile_primitives/copy_async/dsmem.rst
+++ /dev/null
@@ -1,222 +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.
-
-cp_async_bulk: DSMEM
-====================
-
-``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
-(``cp.async.bulk.shared::cluster``); the hardware decrements the *destination* CTA's
-mbarrier when the bytes land. Source:
-``python/tvm/backend/cuda/tile/copy_async/dsmem.py``.
-
-What it accepts
----------------
-
-Three predicates: a valid copy, a single-thread scope, and a shared → shared pair:
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target
- - ``cuda``
- * - scope
- - the call must already be in a **single-thread** scope; the caller normally
- selects that thread with control flow
- * - memory pair
- - both ``shared*`` (``_is_shared_to_shared``); the copy targets a *remote* CTA
- via ``remote_cta_id``
- * - chunk size
- - the contiguous chunk must be **≥ 16 bytes and a multiple of 16**
- (``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
-mbarrier and writes the result out (from ``test_dsmem.py``):
-
-.. code-block:: python
-
- from tvm.backend.cuda.lang import MBarrier, SMEMPool
-
- shape, dtype, CLUSTER_N = (128, 64), "float16", 2
- src_layout = dst_layout = TileLayout(S[128, 64])
- copy_bytes = 128 * 64 * 2
- r = (slice(0, 128), slice(0, 64))
-
-
- @Tx.function
- def dsmem_copy(A: Tx.Tensor(shape, dtype), B: Tx.Tensor(shape, dtype)):
-
- Tx.device_entry(launch=Tx.cuda.LaunchConfig(grid=(CLUSTER_N,), block=(1,), cluster=(CLUSTER_N,)))
- cbx = Tx.cuda.cluster_cta_id("x")
- _index = Tx.cuda.block_idx("x")
- tid = Tx.cuda.thread_idx("x")
- pool = SMEMPool()
- src_raw = pool.alloc([8192], dtype, align=128)
- src_smem = Tx.decl_tensor(
- src_raw.data,
- ty_args=[
- Tx.Tensor(
- list(shape),
- dtype,
- byte_offset=src_raw.byte_offset,
- scope="shared.dyn",
- layout=src_layout,
- )
- ],
- )
- dst_raw = pool.alloc([8192], dtype, align=128)
- dst_smem = Tx.decl_tensor(
- dst_raw.data,
- ty_args=[
- Tx.Tensor(
- list(shape),
- dtype,
- byte_offset=dst_raw.byte_offset,
- scope="shared.dyn",
- layout=dst_layout,
- )
- ],
- )
- mbar = MBarrier(pool, 1)
- pool.commit()
- mbar.init(1)
- Tx.ptx.fence.mbarrier_init.release.cluster()
- Tx.cuda.cluster_sync()
- if tid == 0:
- if cbx == 0: # source CTA
- 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.cuda.tile.cp_async_bulk(
- dst_smem[r],
- src_smem[r],
- 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)
- 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 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:
-
-.. code-block:: python
-
- chunk_bytes = chunk_elements * dtype_bytes
- if chunk_bytes < 16 or chunk_bytes % 16 != 0:
- fail(...)
-
-**2. Map the remote address.** ``Tx.ptx.mapa.u64`` translates a local shared
-pointer into the destination CTA's window — applied to both the destination
-tensor pointer and the mbarrier (``mapa`` writes into a declared register, so
-the mapped addresses live in a small local scratch tensor):
-
-.. code-block:: python
-
- mapped = Tx.alloc_local([2], "uint64")
- Tx.ptx.mapa.u64(mapped[0], mbar, Tx.uint32(remote_cta_id)) # remote_mbar
- Tx.ptx.mapa.u64(mapped[1], dst_buf.ptr_to(dst_st), Tx.uint32(remote_cta_id)) # cluster_dst
-
-**3. Issue one bulk copy per chunk.** Fully contiguous → a single instruction; a
-strided region loops over the outer (non-contiguous) extents, re-deriving the
-chunk's offsets each step:
-
-.. code-block:: python
-
- if Tx.constexpr(not outer_extents): # one contiguous chunk
- Tx.ptx["cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes"](
- Tx.cast(mapped[1], "uint32"), src_buf.ptr_to(src_st),
- Tx.cast(chunk_bytes, "uint32"), Tx.cast(mapped[0], "uint32"))
- else:
- for loop_vars in T.grid(*outer_extents): # one chunk per outer coord
- ... # re-decl src/dst views at the per-chunk offset
- Tx.ptx["cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes"](
- Tx.cast(mapped[1], "uint32"), src_ptr,
- 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 lowerer emits no wait — the caller arms the
-mbarrier (``arrive.expect_tx``) and waits.
-
-Generated TIRx IR
------------------
-
-The fully contiguous ``128×64`` fp16 tile (``16384`` bytes) is a **single chunk**:
-
-.. code-block:: python
-
- Tx.ptx["cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes"](
- Tx.cast(mapped[1], "uint32"), src_ptr[0], Tx.uint32(16384), Tx.cast(mapped[0], "uint32"))
-
-Generated CUDA
---------------
-
-.. code-block:: c++
-
- // map local shared addresses into CTA 1's window (mapa)
- tvm_builtin_ptx_mapa_shared__cluster_u64(remote_mbar, &mbar, /*rank=*/1);
- tvm_builtin_ptx_mapa_shared__cluster_u64(cluster_dst, &dst_smem, /*rank=*/1);
- // bulk-copy 16384 bytes local shared -> CTA 1 shared, signalling its mbarrier
- "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes ..."
-
-The selected thread on CTA 0 launches the whole 16 KB transfer; CTA 1's mbarrier
-fires when it lands.
-
-How inputs change the algorithm
--------------------------------
-
-.. list-table::
- :header-rows: 1
- :widths: 30 70
-
- * - input
- - effect
- * - layout contiguity
- - fully contiguous (matching row-major both sides) → **one** ``cp.async.bulk``;
- a stride gap or mismatched outer stride → a loop of **N** chunks (one per
- outer coord)
- * - dtype / chunk size
- - sets ``chunk_bytes`` (must stay ≥ 16 and a multiple of 16); smaller
- contiguous runs mean smaller, more numerous chunks
- * - ``remote_cta_id``
- - the ``mapa`` rank — which cluster CTA receives the data
- * - incompatible layouts
- - e.g. row-major source vs column-major destination → no matching contiguous
- chain → lowering fails
diff --git a/docs/tirx/tile_primitives/copy_async/ldgsts.rst b/docs/tirx/tile_primitives/copy_async/ldgsts.rst
deleted file mode 100644
index d8c3026f71..0000000000
--- a/docs/tirx/tile_primitives/copy_async/ldgsts.rst
+++ /dev/null
@@ -1,195 +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.
-
-cp_async
-========
-
-``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. Its layout planner constructs an
-``[outer, threads, vec]`` partition. Source:
-``python/tvm/backend/cuda/tile/copy_async/ldgsts.py``.
-
-What it accepts
----------------
-
-.. code-block:: python
-
- _LDGSTS_PAIRS = [("global", "shared*")] # cp.async is unidirectional
- _LDGSTS_VEC_BITS = (128, 64, 32) # cp_size ∈ {16, 8, 4} bytes
-
- def _is_ldgsts(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),
- lambda: _is_valid_copy(op_call, sctx),
- lambda: _scope_allowed(op_call, sctx, allowed_pairs=_LDGSTS_PAIRS),
- lambda: _divides_thread_cnt_ldgsts(op_call, sctx),
- ):
- ok, msg = check()
- if not ok:
- return False, msg
- return True, None
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target / scope
- - ``cuda``; ``thread`` / ``warp`` / ``warpgroup`` / ``cta``, all active
- * - direction
- - **global → shared only** (``_LDGSTS_PAIRS``) — ``cp.async`` is unidirectional
- * - dtype / shape
- - ``_is_valid_copy``; region element count is divisible by the thread count
- * - 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
-
-Configuration
--------------
-
-``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
-zero-fill the destination. ``direct=True`` bypasses cooperative partitioning
-and issues one exact 4-, 8-, or 16-byte copy; that mode requires thread scope
-and a constant-size region. Only the exact string ``"zero"`` changes the
-operand form; other ``fill_mode`` values follow the ordinary predicated or
-unpredicated branch. The implementation converts ``prefetch_size`` to an
-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``):
-
-.. code-block:: python
-
- shape, dtype = (128, 32), "float16"
- s_layout = TileLayout(S[shape])
- full = (slice(0, 128), slice(0, 32))
-
-
- @Tx.function
- def copy_async(A: Tx.Tensor(shape, dtype), B: Tx.Tensor(shape, dtype)):
-
- Tx.device_entry(launch=Tx.cuda.LaunchConfig(grid=(1,), block=(128,)))
- _index = Tx.cuda.block_idx("x")
- _index = Tx.cuda.warp_id()
- _lane = Tx.cuda.lane_id()
- tid = Tx.cuda.thread_idx("x")
- A_smem = Tx.alloc_tensor(ty_args=[Tx.Tensor(shape, dtype, scope="shared", layout=s_layout)])
- 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()
- for k in Tx.serial(32):
- value = A_smem[tid, k]
- B[tid, k] = value
-
-Algorithm
----------
-
-**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
-widest legal width is ``vec = 8`` (``8 × 2 B = 16 B``), giving ``outer = 4``.
-
-**2. Emit** ``cp.async`` **instead of a synchronous copy**, and do **not** sync:
-
-.. code-block:: python
-
- for f in Tx.unroll(total_outer):
- s_lin = s_p.apply(f, tid, v0, shape=apply_shape)["m"]
- g_lin = g_p.apply(f, tid, v0, shape=apply_shape)["m"]
- s_ptr = _ptr_off(s_buf.ptr_to(s_zero), _s_off(f, s_lin))
- g_ptr = _ptr_off(g_buf.ptr_to(g_zero), g_lin)
- Tx.ptx[chain](s_ptr, g_ptr, cp_size) # async; chain includes cache/prefetch modifiers
- # NO cta_sync — commit_group / wait_group / cta_sync are the caller's job
-
-Completion is the caller's responsibility
-(``Tx.ptx.cp.async_.commit_group()`` then
-``Tx.ptx.cp.async_.wait_group(0)``); the lowerer only issues the in-flight
-loads.
-
-Generated TIRx IR
------------------
-
-.. code-block:: python
-
- for f in Tx.unroll(4): # outer = 4
- s_ptr = pointer_offset(A_smem, ...)
- g_ptr = pointer_offset(A_1, ...)
- Tx.ptx["cp.async.cg.shared.global"](s_ptr, g_ptr, 16) # cp_size = 16 B
-
-Generated CUDA
---------------
-
-.. code-block:: c++
-
- // Core instructions (the generated source wraps these in inline helpers):
- asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ...); // x4 per thread
- asm volatile("cp.async.commit_group;");
- asm volatile("cp.async.wait_group 0;");
-
-Each thread issues four asynchronous 16-byte copies; nothing blocks until the
-caller's ``wait_group``.
-
-How inputs change the algorithm
--------------------------------
-
-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
- :widths: 24 18 20 38
-
- * - case
- - ``vec``
- - ``cp_size``
- - ``outer = 4096 / (128 · vec)``
- * - ``float16``, aligned
- - 8
- - 16 B
- - 4
- * - ``float32``, aligned
- - 4
- - 16 B
- - 8
- * - 8-B-aligned only
- - (clamped)
- - 8 B
- - (doubles)
- * - 4-B-aligned only
- - (clamped)
- - 4 B
- - (×4)
-
-If the region can't satisfy even a 4-byte (32-bit) cp_size, ``align_layouts_gs``
-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
deleted file mode 100644
index 167bd834eb..0000000000
--- a/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst
+++ /dev/null
@@ -1,222 +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.
-
-tcgen05.cp
-==========
-
-``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 tensor layouts. The lowerer issues only the copy; the caller signals
-completion with ``tcgen05.commit``. Source:
-``python/tvm/backend/cuda/tile/copy_async/tcgen05_cp.py``.
-
-Shape selection
----------------
-
-- ``shape=`` qualifier: forces that PTX shape (with ``multicast=`` where the shape
- has more than one legal qualifier).
-- 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
- 256b/128b pair are mutually exclusive by the tmem (lane, replica) pattern, so
- a bare warpx4 copy still resolves to ``32x128b.warpx4``.
-
-Each shape pins a tmem row→lane mapping and replica (multicast) pattern
-(verified bit-exactly on B200 by ``test_tcgen05_cp.py``):
-
-.. list-table::
- :header-rows: 1
- :widths: 18 26 30 26
-
- * - shape
- - multicast
- - t lane pattern
- - t replica
- * - ``128x256b``
- - (none)
- - ``(128, 1@TLane)``
- - —
- * - ``4x256b``
- - (none)
- - ``(4, 32@TLane)``
- - —
- * - ``128x128b``
- - (none)
- - ``(128, 1@TLane)``
- - —
- * - ``64x128b``
- - ``warpx2::02_13``
- - ``(64, 1@TLane)``
- - ``(2, 64@TLane)``
- * - ``64x128b``
- - ``warpx2::01_23``
- - ``(2,32):(64,1)@TLane``
- - ``(2, 32@TLane)``
- * - ``32x128b``
- - ``warpx4``
- - ``(32, 1@TLane)``
- - ``(4, 32@TLane)``
-
-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:
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target
- - ``cuda`` target with ``tcgen05`` support (tested with ``sm_100a``)
- * - scope
- - **single thread** issues the copy
- * - CTA group
- - ``cta_group`` is ``1`` (default) or ``2`` and is forwarded to the PTX
- instruction
- * - memory pair
- - source ``shared*`` → destination ``tmem`` (with its backing address supplied by
- a prior ``tcgen05.alloc``); both tensors carry layouts and their element
- bit widths match. Equal-width reinterpretation is allowed
- * - tmem layout
- - must slice to one shape's (lane, replica) pattern from the table above
- * - smem layout
- - rows in 8-row descriptor core-matrix groups; the atom row width derives
- the swizzle mode (K-byte ∈ {16, 32, 64, 128} → sw 0..3) and must match
- the tensor'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 qualifier, the planner infers
-``32x128b.warpx4`` from the layouts (from ``test_tcgen05_cp.py``; readback /
-dealloc tail elided):
-
-.. code-block:: python
-
- from tvm.tirx.layout import R, S, TileLayout
- from tvm.backend.cuda.axis import TCol, TLane
-
- A_smem = Tx.alloc_tensor(
- ty_args=[
- Tx.Tensor(
- [32, 16],
- "uint8",
- scope="shared",
- layout=TileLayout(S[(32, 16) : (16, 1)]),
- align=1024,
- )
- ]
- )
- tmem_addr = Tx.alloc_shared([1], "uint32")
- cp_mbar = Tx.alloc_shared([1], "uint64")
- if warp_id == 0:
- Tx.ptx["tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32"](
- Tx.address_of(tmem_addr), Tx.uint32(16))
- # ... mbarrier.init, fence, cta_sync, fill A_smem from global ...
- tmem = Tx.cuda.decl_tmem(
- tmem_addr[0],
- ty_args=[Tx.Tensor(
- [32, 16], "uint8", scope="tmem",
- layout=TileLayout(S[(32, 16) : (1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane]),
- )],
- )
- if tid_in_wg == 0:
- 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]))
- Tx.cuda.mbarrier_wait(cp_mbar.ptr_to([0]), 0)
- # ... readback via tcgen05.ld, then tcgen05.dealloc ...
-
-Algorithm
----------
-
-**1. Resolve the shape** — explicit ``shape=`` or layout inference (above).
-
-**2. Validate the plan.** Slice both layouts to the region, permute to
-TLane-first / TCol-stride-descending, isolate broadcast iters, and split each
-side into ``(lane, middle, col)`` segments against the atom: ``lane`` = one
-instruction's rows (must match the shape's lane pattern), ``col`` = one
-instruction's columns, ``middle`` = the loop dims — the outer columns and, for
-lane-tiled atoms (e.g. 16 ``4x256b`` cps filling each warp's first 16 rows, the
-M=64 Layout-F scatter), the extra row iters. The smem side derives the
-descriptor fields: 8-row group strides → sdo + swizzle mode, 16B-unit stride →
-ldo (256b atoms).
-
-**3. Encode the matrix descriptor once.** A 64-bit shared descriptor is encoded
-at the smem tensor base right after its allocation, cached per
-``(smem_buf, ldo, sdo, swizzle)``; each cp patches only the 14-bit address
-field.
-
-**4. Issue the copies** — the middle dims flatten into one unrolled loop; each
-step bumps the descriptor's 16-byte shared offset and the tmem address (column
-bits, plus the lane half-word for lane-tiled atoms):
-
-.. code-block:: python
-
- for flat in Tx.unroll(total):
- t_off, s_off = Tx.meta_var(compute_offsets(flat))
- Tx.ptx[f"tcgen05.cp.cta_group::{cta_group}.{shape}{multicast_seg}"](
- 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 lowerer emits **no** ``tcgen05.commit`` / ``wait`` — the caller commits
-against an mbarrier (as in the demo).
-
-Generated CUDA
---------------
-
-.. code-block:: c++
-
- // one warpx4 copy: shared (named by the matrix descriptor) -> tensor memory
- "tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" // [%0]=tmem addr, %1=descriptor
-
-(Compiled for ``sm_100a``. End-to-end correctness — including the tmem readback
-— is covered by ``test_tcgen05_cp.py``.)
-
-How inputs change the algorithm
--------------------------------
-
-.. list-table::
- :header-rows: 1
- :widths: 28 72
-
- * - input
- - effect
- * - tmem layout
- - selects the shape under inference via its (lane, replica) pattern; a
- ``R[4:32@TLane]`` replica resolves to ``32x128b.warpx4``
- * - dtype
- - sets ``elem_per_atom = atom_bits / dtype_bits`` and, with the smem row
- strides, the descriptor swizzle mode
- * - region size
- - middle dims: one atom → a single ``tcgen05.cp``; wider regions → an
- unrolled loop over columns; more rows than one atom → lane-tiled cps
- stepping the tmem lane half-word
- * - shared swizzle layout
- - changes the encoded ``swizzle`` mode (must match the derived atom K-byte)
- * - cta_group
- - forwarded to the instruction; ``cta_group::2`` is pair-collective — one
- even-CTA issue makes each CTA copy from its own smem into its own tmem
- (B200-pinned by the cta_group=2 round-trip tests)
diff --git a/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst b/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst
deleted file mode 100644
index 261ed574d2..0000000000
--- a/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst
+++ /dev/null
@@ -1,263 +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.
-
-tcgen05.ld / tcgen05.st
-=======================
-
-``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. ``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
-instructions, completion (``tcgen05.wait.ld`` / ``tcgen05.wait.st``) is the caller's. Source:
-``python/tvm/backend/cuda/tile/copy_async/tcgen05_ldst.py``.
-
-What it accepts
----------------
-
-The named instruction fixes the transfer direction. Both require warpgroup scope.
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target
- - ``cuda`` target with ``tcgen05`` support (tested with ``sm_100a``)
- * - scope
- - **warpgroup** (``exec_scope_ok(expected_scopes=["warpgroup"])``) — the four
- warps act together
- * - memory pair
- - ``(tmem, local)`` or ``(local, tmem)`` — exactly one side is tensor memory
- * - register layout
- - matched against a ``tcgen05_atom_layout`` (``.16x64b`` / ``.16x128b`` /
- ``.16x256b``) first. If no atom matches, a compatible 128-row Layout D
- may use ``.32x32b``. Layout B instead uses its dedicated ``.32x32b``
- logical ``(64, N)`` image; incompatible layouts decline
- * - tmem datapath
- - classified ``D`` (M=128 identity), ``F`` (M=64 scattered), or ``B``
- (per-CTA M=64, ``cta_group=2`` column split); an F layout also selects
- 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``,
-fp16):
-
-.. code-block:: python
-
- from tvm.tirx.layout import S, TileLayout
- from tvm.backend.cuda.axis import TCol, TLane
- from tvm.backend.cuda.axis import tid_in_wg as axis_tid_in_wg
-
- local_view = TileLayout(S[(128, WIDTH) : (1 @ axis_tid_in_wg, 1)])
-
-
- @Tx.function
- def copy_async_test(
- A: Tx.Tensor((128, WIDTH), "float16"), B: Tx.Tensor((128, WIDTH), "float16")
- ):
-
- Tx.device_entry(launch=Tx.cuda.LaunchConfig(grid=(1,), block=(128,)))
- warp_id = Tx.cuda.warp_id()
- wg_id = Tx.cuda.warpgroup_id()
- tid = Tx.cuda.thread_idx("x")
- tmem_addr = Tx.alloc_shared([1], "uint32")
- if wg_id == 0:
- if warp_id == 0:
- Tx.ptx["tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32"](
- Tx.address_of(tmem_addr), Tx.uint32(32)
- )
- Tx.gpu_storage_sync("shared")
- tmem = Tx.cuda.decl_tmem(
- tmem_addr[0],
- ty_args=[Tx.Tensor(
- (128, WIDTH), "float16", scope="tmem",
- layout=TileLayout(S[(128, WIDTH) : (1 @ TLane, 1 @ TCol)]),
- )],
- )
- A_reg = Tx.alloc_local((WIDTH,), "float16")
- B_reg = Tx.alloc_local((WIDTH,), "float16")
- 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.cuda.tile.tcgen05.st(tmem[:, :], A_local[:, :], scope="warpgroup")
- Tx.ptx.tcgen05.wait__st.sync.aligned() # store (local -> tmem)
- Tx.cuda.cta_sync()
- 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 ...
-
-Algorithm
----------
-
-**1. Infer direction.** ``tmem → local`` is a load (``tcgen05.ld``); ``local → tmem``
-is a store (``tcgen05.st``).
-
-**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
-tries ``.32x32b`` only for a compatible 128-row Layout D and probes its legal
-``num`` values against the column width. Other layouts decline instead of
-silently taking that path.
-
-**3. Issue per datapath slab.** For an M=128 Layout D ``.16x*b`` copy the
-fragment spans two 16-row slabs, so the warps issue the atom twice
-(``row = 0`` and ``row = 16``). An M=64 Layout F copy issues once, at
-``row = 0`` for ``sub_slab=0`` or ``row = 16`` for ``sub_slab=1``. The
-``.32x32b`` path covers M=128 in a single issue at ``row = 0``:
-
-.. code-block:: python
-
- chain = f"tcgen05.{'ld' if load else 'st'}.sync.aligned.{shape}.x{num}.b32"
- classified = _check_tmem_layout_for_atom(tmem_buf, "16x*b", frag_rows)
- datapath, sub_slab = classified if classified is not None else (None, 0)
- for slab in range(n_slabs): # 1 for M=64; 2 for M=128
- op(Tx.cast(Tx.reinterpret(tmem_buf.data, ty="uint64"), "uint32"),
- *[local_32b[reg_base + i] for i in range(regs_eff)],
- shape=shape, num=num_eff,
- row=(sub_slab + slab) * 16, col=col_off_32b)
-
-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 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
-64-row tensors alias the lower and upper halves of one 128-row allocation:
-
-.. code-block:: python
-
- from tvm.tirx.layout import tmem_datapath_layout
-
- lower = Tx.cuda.decl_tmem(
- tmem_addr[0],
- ty_args=[Tx.Tensor(
- (64, cols), "float32", scope="tmem",
- layout=tmem_datapath_layout("F", 64, cols, sub_slab=0),
- )],
- )
- upper = Tx.cuda.decl_tmem(
- tmem_addr[0],
- ty_args=[Tx.Tensor(
- (64, cols), "float32", scope="tmem",
- layout=tmem_datapath_layout("F", 64, cols, sub_slab=1),
- )],
- )
-
- Tx.cuda.tile.tcgen05.ld(lower_frag, lower, scope="warpgroup")
- Tx.ptx.tcgen05.wait__ld.sync.aligned()
- 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
-``.16x64b``, ``.16x128b``, and ``.16x256b`` atoms. Layout D has 128 rows and
-already spans both sub-slabs, so it only accepts ``sub_slab=0``.
-
-Layout B readback
------------------
-
-Layout B is produced by an M=64-per-CTA ``tcgen05.mma`` with
-``cta_group=2``. Its logical ``(64, N)`` columns are split into two
-``N/2`` halves: the low half uses physical TMEM lanes 0–63, and the high
-half uses lanes 64–127. It therefore occupies all 128 lanes and ``N/2``
-``TCol`` values.
-
-Use the public allocation and fragment APIs together:
-
-.. code-block:: python
-
- accumulator = tmem_pool.alloc_tcgen05_mma_D(
- (64, N), "float32", M=128, cta_group=2)
- frag = Tx.alloc_tcgen05_ldst_frag("32x32b", (64, N), "float32")
-
- 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.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
-even, ``N/2`` must be a valid PTX ``num``, and the fragment is fp32-only.
-The first implementation intentionally requires the full logical
-``(64, N)`` region: a partial logical-column slice is not contiguous after
-the two-way lane split. A conventional ``.16x*b`` fragment is rejected with
-an actionable error.
-
-Generated TIRx IR
------------------
-
-For the ``128×8`` fp16 tile the layout takes the ``.32x32b`` path with ``num = 4``
-(4 registers per thread), one issue each way:
-
-.. code-block:: python
-
- Tx.ptx["tcgen05.st.sync.aligned.32x32b.x4.b32"]( # local -> tmem
- Tx.uint32(tmem_addr[0]), local_32b[0], local_32b[1], local_32b[2], local_32b[3])
- Tx.ptx["tcgen05.ld.sync.aligned.32x32b.x4.b32"]( # tmem -> local
- local_32b_1[0], local_32b_1[1], local_32b_1[2], local_32b_1[3],
- Tx.uint32(tmem_addr[0]))
-
-Generated CUDA
---------------
-
-.. code-block:: c++
-
- "tcgen05.st.sync.aligned.32x32b.x4.b32 ..." // 4 registers -> tmem
- "tcgen05.ld.sync.aligned.32x32b.x4.b32 ..." // tmem -> 4 registers
-
-Verified end-to-end on ``sm_100a`` (the round trip reproduces the input exactly).
-
-How inputs change the algorithm
--------------------------------
-
-.. list-table::
- :header-rows: 1
- :widths: 28 72
-
- * - input
- - effect
- * - 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 lowering fails
- * - column width / dtype
- - sets ``num`` (the ``.xN`` count) and the registers per thread
- (``elem_per_32b = 32 / dtype_bits``)
- * - direction
- - ``tmem → local`` → ``tcgen05.ld``; ``local → tmem`` → ``tcgen05.st`` (same
- shape/num logic)
- * - datapath D vs F vs B
- - ``D`` (M=128) covers all 128 rows; an M=128 ``.16x*b`` copy issues two slabs
- (``row = 0`` / ``row = 16``). ``F`` (M=64) scatters rows to lanes and
- its layout selects one issue at ``row = 0`` or ``row = 16``. ``B``
- (per-CTA M=64, ``cta_group=2``) splits N across two 64-lane halves and
- emits one logical-64-row ``.32x32b`` image
diff --git a/docs/tirx/tile_primitives/copy_async/tma.rst b/docs/tirx/tile_primitives/copy_async/tma.rst
deleted file mode 100644
index cf72443c22..0000000000
--- a/docs/tirx/tile_primitives/copy_async/tma.rst
+++ /dev/null
@@ -1,243 +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.
-
-TMA tensor instructions
-=======================
-
-``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);
-* construct and cache ``cuTensorMap`` descriptors on the host with
- ``cuTensorMapEncodeTiled``;
-* use the same hardware validator, descriptor cache, prefetch path, and PTX
- emitter; and
-* require the source and destination regions to contain the same total number
- of bytes. Their ranks and per-dimension shapes need not match.
-
-The descriptor modes differ only in how the TensorMap and issue count are planned.
-
-``tma_auto``
-------------
-
-``tma_auto`` derives the largest legal TMA box from the shared-memory iteration
-order. It is intended for ordinary full-tile copies whose layout relationship
-can be proven statically:
-
-.. code-block:: python
-
- Tx.cuda.tile.cp_async_bulk_tensor_load(
- A_smem[:, :],
- A[tile_m : tile_m + 64, tile_k : tile_k + 64],
- descriptor_mode="auto",
- mbar=mbar.ptr_to([0]),
- )
-
-The shared slice must contain only memory iterators and must form one complete
-contiguous stride chain. The planner groups the global layout against that
-chain, selects the maximum hardware-legal shared prefix for the TensorMap box,
-and emits mixed-radix issue loops for the remaining dimensions.
-
-Every raw candidate runs the same address-preserving dimension
-canonicalization: address-free unit dimensions are removed where legal, and
-adjacent contiguous dimensions are merged when the inner dimension is copied
-in full from coordinate zero. When an innermost contiguous, fully boxed,
-coordinate-zero pair is blocked only because its merged box would exceed 256,
-canonicalization may first apply one byte-preserving descriptor-unit promotion
-and retry that same boundary, but only if that single promotion makes the merge
-legal and a farther outer dimension exists. This avoids skipping outward and
-changing the TensorMap tiling of the contiguous inner chain. Promotion is
-otherwise a repair step used only when the canonical candidate fails a
-repairable hardware rule. It preserves byte strides, payload size,
-transaction size, global base, and shared pointer, and is not used for
-reductions, TF32, packed dtypes, interleave, OOB fill, gather4, or issue-driven
-innermost axes.
-
-``tma_auto`` does not accept ``gather4`` or ``src_selector`` and only accepts the
-default no-OOB contract (``oob=None``). Symbolic facts that affect layout
-mapping, prefix selection, or repair must be proven. A dynamic ``globalDim``
-whose range is otherwise unknown is the exception: the runtime TensorMap
-encoder checks it is in ``(0, 2^32]`` immediately before the CUDA Driver call.
-Use ``tma_explicit`` for explicit OOB behavior or when other descriptor facts
-are only known at runtime.
-
-``tma_explicit``
-----------------
-
-``tma_explicit`` maps the supplied global Tensor or view directly:
-
-* Tensor/view shape becomes ``globalDim``;
-* layout strides become byte ``globalStrides``;
-* region start becomes the instruction coordinates;
-* region extent becomes ``boxDim``; and
-* Tensor data plus ``elem_offset`` becomes the TensorMap base.
-
-It never regroups, compresses, promotes, shrinks, or splits a copy. One
-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.cuda.tile.cp_async_bulk_tensor_store(
- O[:, atom * 64 : (atom + 1) * 64],
- O_smem[:, atom * 64 : (atom + 1) * 64],
- descriptor_mode="explicit",
- )
-
-The sliced shared layout must canonicalize to a trivial box after its pointer
-offset and swizzle are extracted. Global rank and memory-layout rank must
-match. A statically illegal value is rejected; symbolic global shapes,
-strides, and alignment are retained for validation by the runtime encoder.
-
-Gather4
-~~~~~~~
-
-Gather4 is an explicit global-to-shared operation on SM100 or newer. It
-requires a two-dimensional TensorMap and exactly four absolute row
-coordinates. Public axis zero is the four-row payload, and PTX receives
-coordinates in ``{column, row0, row1, row2, row3}`` order:
-
-.. code-block:: python
-
- Tx.cuda.tile.cp_async_bulk_tensor_load(
- K_smem[0:4, :],
- K[0:1, :],
- descriptor_mode="explicit",
- mbar=mbar.ptr_to([0]),
- gather4=[row0, row1, row2, row3],
- )
-
-Longer gathers must be written as multiple four-row calls. Each destination
-slice must be four-row box-linear and each source row must have the descriptor's
-declared byte stride.
-
-Descriptor selection
-~~~~~~~~~~~~~~~~~~~~
-
-``src_selector`` selects among alternate global Tensors or views while reusing
-the main operand's region and gather coordinates:
-
-.. code-block:: python
-
- Tx.cuda.tile.cp_async_bulk_tensor_load(
- K_smem[0:4, :],
- K_main[0:1, :],
- descriptor_mode="explicit",
- mbar=mbar.ptr_to([0]),
- gather4=[row0, row1, row2, row3],
- src_selector=[
- (use_extra, K_extra.sub[base_row:, :]),
- (use_backup, K_backup),
- ],
- )
-
-Candidates must cover the entire tensor view; use a subview to change its base.
-Conditions use first-true priority and the main Tensor 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
-dtype, rank, box, swizzle, and transfer byte count. Lowering selects a
-pointer-typed descriptor and emits one TMA instruction; it does not select
-coordinates or generate instruction-level branches.
-
-``prefetch_tensormap=True`` deduplicates and prefetches only the main
-descriptor. Selector candidates are not prefetched automatically.
-
-Common configuration
---------------------
-
-.. list-table::
- :header-rows: 1
- :widths: 28 72
-
- * - Option
- - Meaning
- * - ``mbar`` / ``mbarrier_addr``
- - ``mbar`` is the required completion barrier for global-to-shared copies.
- ``mbarrier_addr=False`` leaves it in generic-pointer form; ``True`` or
- any TIR expression requests the PTX shared-address operand form, so
- lowering converts ``mbar`` once before the instruction. The expression
- is a form-selection marker, not a replacement barrier operand. Neither
- option is valid for shared-to-global.
- * - ``cta_group`` / ``cta_mask``
- - ``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`` / ``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.
- * - ``tensormap_l2_promotion``
- - ``None`` (the default, meaning ``L2::128B``), integer ``0`` through
- ``3``, ``none``, ``L2::none``, ``L2::64B``, ``L2::128B``, or
- ``L2::256B``.
- * - ``tma_dtype``
- - ``tf32`` or ``tfloat32`` for a float32 descriptor conversion.
- * - ``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
- encoding; an explicit ``zero`` selects the same encoding, while ``nan``
- is limited to non-packed floating-point descriptors. ``tma_auto``
- requires ``oob=None`` and rejects an explicitly supplied mode.
-
-Hardware validation
--------------------
-
-TensorMap dimensions use CUDA API order: dimension zero is innermost and
-``globalStrides`` omits its unit stride. Layout order is reversed once when
-the descriptor specification is constructed.
-
-The shared validator checks, among other rules:
-
-* rank 1 through 5 and target SM90 or newer;
-* legal scalar descriptor dtype and conversion dtype;
-* 16-byte global base alignment, raised where packed or interleaved modes
- require 32 bytes, and 64-byte TensorMap object alignment;
-* ``globalDim`` in ``(0, 2^32]``;
-* each explicit byte stride non-negative, aligned, and less than ``2^40``;
-* ``boxDim`` in ``[1, 256]`` and element stride in ``[1, 8]``;
-* unit innermost stride and ordinary inner-box bytes divisible by 16;
-* legal interleave, swizzle, L2 promotion, and OOB combinations;
-* swizzled inner boxes no larger than their 32-, 64-, or 128-byte atom; and
-* direction, reduction, mbarrier, CTA group, coordinate count, and load-mode
- combinations.
-
-Packed sub-byte descriptors are validated in their actual TensorMap encoding
-units rather than by truncating ``bits // 8``.
-
-For ``tma_auto``, a statically disproven rule always rejects the candidate.
-Unknown planner-sensitive rules also reject it; only the dynamic
-``globalDim`` range is deferred to the runtime encoder. ``tma_explicit``
-retains unknown descriptor values for runtime validation because it does not
-choose among alternative boxes or issue loops.
-
-Completion
-----------
-
-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
-their surrounding algorithm.
diff --git a/docs/tirx/tile_primitives/elementwise.rst b/docs/tirx/tile_primitives/elementwise.rst
deleted file mode 100644
index 2e98b7559c..0000000000
--- a/docs/tirx/tile_primitives/elementwise.rst
+++ /dev/null
@@ -1,43 +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.
-
-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
-
- elementwise/reg
- elementwise/smem
diff --git a/docs/tirx/tile_primitives/elementwise/reg.rst b/docs/tirx/tile_primitives/elementwise/reg.rst
deleted file mode 100644
index ae3cde18f3..0000000000
--- a/docs/tirx/tile_primitives/elementwise/reg.rst
+++ /dev/null
@@ -1,159 +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.
-
-elementwise → reg
-=================
-
-The register lowerer expands an elementwise instruction (``sqrt``, ``exp``, ``add``,
-``fma``, …) when **all tensor operands use ``local`` scope**. Scalar inputs are
-also accepted where the operation's authoring API permits them. Like the copy
-:doc:`../copy/reg` instruction the partition is *induced* by the operands'
-local-tensor layout — the thread axes are dropped, leaving each thread its
-private bundle — and the op is applied to every element in that bundle. The
-lowering path uses local operands; final register allocation
-remains a CUDA compiler decision. Source:
-``python/tvm/backend/cuda/tile/elementwise/reg.py``.
-
-What it accepts
----------------
-
-``is_reg_ewise(spec)`` builds the predicate:
-
-.. code-block:: python
-
- def check(op_call, sctx):
- if not sctx.is_target("cuda"): return False, "non-cuda target"
- if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"): ...
- ok, reason = _all_threads_active(sctx)
- plan, msg = spec.parse(op_call)
- for br in tensor_regions(plan):
- if br.tensor.scope() != "local": # every tensor operand local
- return False, f"operand scope {br.tensor.scope()} != local"
- if br.tensor.layout is None: ...
- # + spec.check_extras (dtype rules), pick_anchor + _validate_anchor_layout,
- # _validate_scope_level_anchor, NumPy-style shape broadcast checks,
- # and agreement of operand thread/local/replica layout signatures
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target / scope
- - ``cuda``; ``thread`` / ``warp`` / ``warpgroup`` / ``cta`` (all active)
- * - operands
- - **every tensor operand** in ``local``; scalar sources are allowed by
- ``fill``, the binary ops, and ``fma``
- * - op
- - any CUDA elementwise ``OpSpec`` listed on the parent page (unary / binary /
- ``fma``); ``spec.check_extras`` validates the dtype combo
- * - local-tensor layout
- - the anchor layout must validate, and its thread axis must match the
- complete scope chain (it induces the partition); all operands must agree
- on their thread, per-thread storage, and replica signatures
- * - shapes
- - every input region is NumPy-style right-aligned broadcast-compatible with
- 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``):
-
-.. code-block:: python
-
- from tvm.tirx.layout import S, TileLayout
- from tvm.backend.cuda.axis import laneid
-
- r_layout = TileLayout(S[(32, 8) : (1 @ laneid, 1)])
- fs = (slice(0, 32), slice(0, 8))
-
-
- @Tx.function
- def k(A: Tx.Tensor((32, 8), "float32"), B: Tx.Tensor((32, 8), "float32")):
-
- Tx.device_entry(launch=Tx.cuda.LaunchConfig(grid=(1,), block=(32,)))
- _index = Tx.cuda.block_idx("x")
- _lane = Tx.cuda.lane_id()
- tid = Tx.cuda.thread_idx("x")
- A_smem = Tx.alloc_tensor(
- ty_args=[Tx.Tensor((32, 8), "float32", scope="shared", layout=TileLayout(S[(32, 8)]))]
- )
- for k in Tx.serial(8):
- value = A[tid, k]
- A_smem[tid, k] = value
- Tx.cuda.cta_sync()
- R = Tx.alloc_tensor(ty_args=[Tx.Tensor((32, 8), "float32", scope="local", layout=r_layout)])
- 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()
- for k in Tx.serial(8):
- value = A_smem[tid, k]
- B[tid, k] = value
-
-Algorithm
----------
-
-**1. Parse and check.** ``spec.parse`` builds the op plan; the predicate confirms
-every tensor operand is local, validates NumPy-style broadcasting, checks
-the anchor against the complete scope chain, and requires compatible
-thread/local/replica layout signatures across operands.
-
-**2. Induce the partition** from the anchor's thread axis (``laneid`` here): drop
-the thread iters, leaving each thread its private bundle (8 elements per lane).
-
-**3. Apply the op to the per-thread bundle.** Registered packed implementations
-are tried widest-first after every operand proves a matching contiguous tail;
-otherwise the lowering uses the scalar operation in a per-thread loop.
-
-Generated TIRx IR
------------------
-
-.. code-block:: python
-
- tensor[f] = Tx.sqrt(tensor_1[f]) # over each element f in the lane's bundle
-
-Generated CUDA
---------------
-
-.. code-block:: c++
-
- r_local_ptr[f_2] = sqrtf(r_local_ptr[f_2]); // per-thread local element
-
-(Verified on ``sm_100a`` — the result equals ``sqrt(A)``.)
-
-How inputs change the algorithm
--------------------------------
-
-.. list-table::
- :header-rows: 1
- :widths: 28 72
-
- * - input
- - effect
- * - op
- - unary → ``sqrtf`` / ``expf`` / … per element; binary → ``a + b`` per
- element; ``fma`` → ``a * b + c``
- * - dtype
- - selects the scalar operation and, where registered, a packed form such as
- ``f32x2`` or a two-element cast
- * - local-tensor layout
- - the anchor's thread axis sets the partition; a wider per-lane bundle means a
- longer loop
diff --git a/docs/tirx/tile_primitives/elementwise/smem.rst b/docs/tirx/tile_primitives/elementwise/smem.rst
deleted file mode 100644
index 09696a1979..0000000000
--- a/docs/tirx/tile_primitives/elementwise/smem.rst
+++ /dev/null
@@ -1,160 +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.
-
-elementwise → smem
-==================
-
-The shared-memory lowerer expands an elementwise instruction (``sqrt``, ``exp``, ``add``,
-``fma``, …) when **all tensor operands are in shared memory**. Scalar inputs are
-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:
-``python/tvm/backend/cuda/tile/elementwise/smem.py``.
-
-What it accepts
----------------
-
-``is_smem_ewise(spec)`` builds the predicate:
-
-.. code-block:: python
-
- def check(op_call, sctx):
- if not sctx.is_target("cuda"): return False, "non-cuda target"
- if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"): ...
- ok, reason = _all_threads_active(sctx) # full scope
- plan, msg = spec.parse(op_call) # parse the op's operands
- for br in tensor_regions(plan):
- if not br.tensor.scope().startswith("shared"): # every tensor operand shared*
- return False, f"operand scope {br.tensor.scope()} != shared*"
- if br.tensor.layout is None: ...
- # + spec.check_extras (dtype rules) and anchor-layout validation
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target / scope
- - ``cuda``; ``thread`` / ``warp`` / ``warpgroup`` / ``cta`` (all active)
- * - operands
- - **every tensor operand** (including the output) in ``shared*``; scalar
- sources are allowed by ``fill``, the binary ops, and ``fma``
- * - op
- - any CUDA elementwise ``OpSpec`` listed on the parent page (unary
- ``sqrt``/``exp``/``zero``…, binary ``add``/``mul``…, ``fma``);
- ``spec.check_extras`` validates the dtype combo
- * - layout
- - operands have layouts. The partition is synthesized from the scope's
- thread count; dtype, logical innermost region extents, and elements per
- 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
-round):
-
-.. code-block:: python
-
- s_layout = TileLayout(S[(32, 32)])
- full = (slice(0, 32), slice(0, 32))
-
-
- @Tx.function
- def unary_op(A: Tx.Tensor((32, 32), "float32", layout=s_layout)):
-
- Tx.device_entry(launch=Tx.cuda.LaunchConfig(grid=(1,), block=(256,)))
- _index = Tx.cuda.block_idx("x")
- _index = Tx.cuda.warp_id()
- _lane = Tx.cuda.lane_id()
- tid = Tx.cuda.thread_idx("x")
- A_smem = Tx.alloc_tensor(
- ty_args=[Tx.Tensor((32, 32), "float32", scope="shared", layout=s_layout)]
- )
- 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
----------
-
-**1. Parse the op and check operands.** ``spec.parse`` turns the call into a plan
-(inputs, output, the op); the predicate confirms every tensor operand is shared.
-
-**2. Synthesize the partition** from the scope's **thread count** (as
-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
-example, ``32×32 = 1024`` ``float32`` over 256 threads gives ``vec = 4`` and
-``outer = 1``.
-
-**3. Apply the op per element.** Instead of a copy, each (thread, round) reads its
-``vec`` elements, applies the op, and writes back — vectorized:
-
-Generated TIRx IR
------------------
-
-.. code-block:: python
-
- for f in Tx.serial(1): # outer = 1
- for vec in Tx.vectorized(4):
- A_smem[tid * 4 + vec] = Tx.sqrt(A_smem[tid * 4 + vec])
-
-Generated CUDA
---------------
-
-The ``vec = 4`` element bundle becomes a ``float4`` and the op is applied per
-component:
-
-.. code-block:: c++
-
- float4 v_ = *(float4*)(&A_smem_ptr[tid * 4]);
- __1.x = sqrtf(v_.x); __1.y = sqrtf(v_.y);
- __1.z = sqrtf(v_.z); __1.w = sqrtf(v_.w);
-
-(Verified on ``sm_100a`` — the tile equals ``sqrt(A)``.)
-
-How inputs change the algorithm
--------------------------------
-
-.. list-table::
- :header-rows: 1
- :widths: 28 72
-
- * - input
- - effect
- * - op
- - unary → ``sqrtf`` / ``expf`` / … per component; binary → the two inputs
- combined (``a + b``); ``fma`` → ``a * b + c``
- * - dtype
- - bounds the candidate width; elements per thread and logical innermost
- region extents can reduce it, changing the round count. The current
- width selection does not inspect physical layout contiguity
- * - scope
- - sets the thread axis and count, hence the synthesized partition
diff --git a/docs/tirx/tile_primitives/gemm.rst b/docs/tirx/tile_primitives/gemm.rst
deleted file mode 100644
index 0ba6d17609..0000000000
--- a/docs/tirx/tile_primitives/gemm.rst
+++ /dev/null
@@ -1,195 +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.
-
-mma_sync
-========
-
-``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 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/gemm/mma_m16n8k_.py``. (For the
-Blackwell async tensor-core path see :doc:`gemm_async`.)
-
-What it accepts
----------------
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - 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
- such as a cluster axis. ``mma.sync`` remains warp-collective, so callers
- use a warp or a wider scope made of complete warps
- * - operand scope
- - **A, B, C, D all in registers** (``local``); a shared operand makes the
- dispatch ``fail`` (stage with ldmatrix first)
- * - no replica
- - none of D/A/B/C may carry a broadcast/replica axis (``_no_replica``)
- * - shape
- - ``M % 16 == 0``, ``N % 8 == 0``, and ``K % 8 == 0``. The dispatcher
- first tries ``m16n8k16`` and then ``m16n8k8``; the selected instruction
- must tile all operand layouts exactly
- * - dtype
- - A and B are both ``float16`` or both ``bfloat16``; C and D are
- ``float32``
- * - alpha / beta
- - ``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``):
-
-.. code-block:: python
-
- from tvm.tirx.layout import S, TileLayout
- from tvm.backend.cuda.axis import laneid
-
- D_FRAG = TileLayout(S[(2, 8, 4, 2) : (2, 4 @ laneid, 1 @ laneid, 1)])
- A_FRAG_K8 = TileLayout(S[(2, 8, 4, 2) : (2, 4 @ laneid, 1 @ laneid, 1)])
- B_FRAG_K8 = TileLayout(S[(4, 2, 8) : (1 @ laneid, 1, 4 @ laneid)])
- A_FRAG = A_FRAG_K8.tile_to([16, 16], [16, 8])
- B_FRAG = B_FRAG_K8.tile_to([16, 8], [8, 8])
-
-
- @Tx.function
- def gemm(
- A_g: Tx.Tensor((16, 16), "float16"),
- B_g: Tx.Tensor((16, 8), "float16"),
- D_g: Tx.Tensor((16, 8), "float32"),
- ):
-
- Tx.device_entry(launch=Tx.cuda.LaunchConfig(grid=(1,), block=((1) * 32,)))
- _index = Tx.cuda.block_idx("x")
- _index = Tx.cuda.warp_id()
- lane = Tx.cuda.lane_id()
- A_f = Tx.alloc_tensor(
- ty_args=[Tx.Tensor((16, 16), "float16", scope="local", layout=A_FRAG)]
- )
- B_f = Tx.alloc_tensor(ty_args=[Tx.Tensor((16, 8), "float16", scope="local", layout=B_FRAG)])
- D_f = Tx.alloc_tensor(ty_args=[Tx.Tensor((16, 8), "float32", scope="local", layout=D_FRAG)])
- A_reg = A_f.local(8) # stage A into the lane's 8 regs
- for s in Tx.unroll(8):
- kp, rM, kHi = s % 2, (s // 2) % 2, s // 4
- A_reg[s] = A_g[lane // 4 + 8 * rM, 2 * (lane % 4) + kp + 8 * kHi]
- B_reg = B_f.local(4) # stage B into the lane's 4 regs
- 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.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):
- rN, rM = s % 2, s // 2
- D_g[lane // 4 + 8 * rM, 2 * (lane % 4) + rN] = D_reg[s]
-
-Algorithm
----------
-
-**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
-instruction that fits, with matching warp-tiling, wins.
-
-**2. Derive register layouts.** Each operand gets a mediated per-lane view:
-D/C as ``[Mo, No, rM, rN]`` (4 f32), A as
-``[Mo, Ko, rM, kHi, k_pack]``, and B as
-``[Ko, No, kHi, k_pack]``. The corresponding raw physical register order
-exposed by the default ``local`` view is D/C ``[Mo, No, rM, rN]``, A
-``[Mo, Ko, kHi, rM, k_pack]``, and B ``[Ko, No, kHi, k_pack]`` — the PTX
-operand enumeration that ``mma.sync`` expects.
-
-**3. Emit the unrolled nest** — initialize D (from C if ``beta==1``, else 0), then
-accumulate over K in place, one ``mma`` per (m, n) tile:
-
-.. code-block:: python
-
- for m in Tx.unroll(M_tiles):
- for n in Tx.unroll(N_tiles):
- for rM, rN in ...: d_local[m, n, rM, rN] = c_local[...] if use_c else Tx.float32(0)
- for k in Tx.unroll(K_tiles):
- d_regs = [d_local[m, n, rM, rN] for rM in range(2) for rN in range(2)] # 4 f32
- a_regs = [a_words[m, k, rM, kHi, 0] for kHi in range(n_kHi) for rM in range(2)]
- b_regs = [b_words[k, n, kHi, 0] for kHi in range(n_kHi)]
- mma_chain = (f"mma.sync.aligned.{shape_str}.row.col"
- f".f32.{a_elem}.{b_elem}.f32")
- Tx.ptx[mma_chain](*d_regs, *a_regs, *b_regs, *d_regs) # d = a·b + d
-
-Generated TIRx IR
------------------
-
-The single 16×8×16 tile lowers to one ``mma`` (4 D regs, 4 A regs, 2 B regs):
-
-.. code-block:: python
-
- Tx.ptx["mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32"](
- d_local[0], d_local[1], d_local[2], d_local[3], a_local[0], ...)
-
-Generated CUDA
---------------
-
-.. code-block:: c++
-
- "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
- "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
-
-The accumulator ``{%0..%3}`` is both the C input and the D output (in-place
-accumulate); ``{%4..%7}`` are A's four ``b32`` registers, ``{%8, %9}`` B's two.
-Verified on ``sm_100a`` (``D == A@B`` within fp16 tolerance).
-
-How inputs change the algorithm
--------------------------------
-
-.. list-table::
- :header-rows: 1
- :widths: 28 72
-
- * - input
- - effect
- * - input dtype
- - ``float16`` → ``…f32.f16.f16.f32``; ``bfloat16`` → ``…f32.bf16.bf16.f32``
- (register counts unchanged — 2 elems per ``b32``)
- * - K instruction
- - ``k16`` → A 4 ``b32`` / B 2 ``b32``; ``k8`` → A 2 / B 1
- (``mma.…m16n8k8.…``)
- * - M / N / K extents
- - set the ``M_tiles`` / ``N_tiles`` / ``K_tiles`` unrolled loop counts (one
- ``mma`` per (m, n), K accumulated in place)
- * - beta
- - ``0`` → D zero-initialized; ``1`` → D initialized from C (the ``mma`` itself
- is identical)
- * - transpose_A / transpose_B
- - transpose the logical A or B region before shape and layout matching;
- the transformed operands must still fit the selected ``m16n8k`` frame
- * - operand scope
- - A/B **must** be register fragments; a shared operand makes the dispatch
- ``fail`` (stage via :doc:`copy/ldstmatrix` first)
diff --git a/docs/tirx/tile_primitives/gemm_async.rst b/docs/tirx/tile_primitives/gemm_async.rst
deleted file mode 100644
index 7ac12370e6..0000000000
--- a/docs/tirx/tile_primitives/gemm_async.rst
+++ /dev/null
@@ -1,298 +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.
-
-tcgen05.mma / tcgen05.mma_block_scale
-=====================================
-
-``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
-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
-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/gemm_async/tcgen05.py``. (For the
-synchronous warp-register path see :doc:`gemm`.)
-
-What it accepts
----------------
-
-A single predicate — single-thread or warp scope:
-
-.. list-table::
- :header-rows: 1
- :widths: 22 78
-
- * - Property
- - Requirement
- * - target / scope
- - ``cuda`` target with ``tcgen05`` support (this implementation is tested
- 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``)
- * - dtype
- - without ``is_AB_tf32``, dense A/B use the same dtype: ``float16``,
- ``bfloat16``, ``float8_e4m3fn``, or ``float8_e5m2``. Typical TF32
- authoring uses ``float32`` tensors with ``is_AB_tf32=True``; the flag
- replaces both semantic A/B dtypes with ``tf32`` before validation, so
- the current implementation does not separately validate the underlying
- storage dtypes or their equality in that mode.
- Block-scaled A/B: ``float8_e4m3fn`` with tmem ``SFA`` / ``SFB`` of
- ``float8_e8m0fnu``, or ``float4_e2m1fn`` with matching
- ``float8_e8m0fnu`` or ``float8_e4m3fn`` scale factors.
- The accumulator is always ``float32``
- * - shape
- - the per-CTA M region is tiled by 64- or 128-row instruction blocks;
- ``N`` is divisible by 8 (cta_group=1) or 16 (cta_group=2), and ``K`` is
- divisible by ``MMA_K`` = 8 (tf32) / 16 (f16/bf16) / 32 (fp8) / 64
- (fp4). With cta_group=2, the CTA pair covers twice the per-CTA M
- * - cta_group
- - ``1`` (one CTA) or ``2`` (two CTAs split the operand)
- * - descriptor mode
- - optional ``smem_desc`` controls shared matrix-descriptor construction:
- ``"hoist"`` (default), ``"local_hoist"``, ``"encode"``, or
- ``"recompute"``. ``"encode"`` is currently dense-only. Other strings
- are not rejected and take the descriptor-add path used by hoisting
- * - explicit instruction tile
- - ``mma_m`` and ``mma_n`` must be supplied together. ``mma_m`` must equal
- ``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 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
-into shared (from
-``test_gemm_async.py``; setup/readback abbreviated):
-
-.. code-block:: python
-
- from tvm.tirx.layout import S, TileLayout
- from tvm.backend.cuda.axis import TCol, TLane, tid_in_wg as axis_tid_in_wg
- from tvm.backend.cuda.tile.tma_utils import mma_shared_layout
-
- A_smem = Tx.alloc_tensor(
- ty_args=[
- Tx.Tensor(
- (3, 128, 64),
- "float16",
- scope="shared",
- layout=mma_shared_layout("float16", 3, (3, 128, 64)),
- )
- ]
- )
- B_smem = Tx.alloc_tensor(
- ty_args=[
- Tx.Tensor(
- (3, 128, 64),
- "float16",
- scope="shared",
- layout=mma_shared_layout("float16", 3, (3, 128, 64)),
- )
- ]
- )
- tmem_addr = Tx.alloc_shared([1], "uint32"); mma_mbar = Tx.alloc_shared([1], "uint64")
- # ... mbarrier.init, cta_sync ...
- if warp_id == 0:
- Tx.ptx["tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32"](
- Tx.address_of(tmem_addr), Tx.uint32(512))
- Tx.cuda.cta_sync()
- tmem = Tx.cuda.decl_tmem(
- tmem_addr[0],
- ty_args=[Tx.Tensor(
- (128, 512), "float32", scope="tmem",
- layout=TileLayout(S[(128, 512) : (1 @ TLane, 1 @ TCol)]),
- )],
- )
- # ... TMA-load A_smem, B_smem from global, wait ...
- if tid_in_wg == 0:
- 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]))
- Tx.cuda.mbarrier_wait(mma_mbar.ptr_to([0]), 0)
- # ... tcgen05.fence.after_thread_sync(); read tmem back via tcgen05.ld; dealloc ...
-
-Algorithm
----------
-
-**1. Encode or synthesize shared matrix descriptors.** Each shared operand is named
-by a 64-bit descriptor (leading-dim offset ``ldo``, stride-dim offset ``sdo``,
-swizzle mode). ``smem_desc`` selects where that descriptor comes from:
-
-* ``"hoist"`` (default): encode one uniform descriptor per operand after shared
- allocation and add each per-MMA 16-byte offset.
-* ``"local_hoist"``: encode at this ``gemm_async`` call site, under the caller's
- control flow, then add offsets. This is for call sites where only the elected
- issue thread should construct the descriptor.
-* ``"encode"``: encode the exact shared pointer for each MMA issue.
-* ``"recompute"``: synthesize the descriptor value inline per MMA without a local
- descriptor cell.
-
-.. code-block:: python
-
- Tx.cuda.tcgen05.encode_matrix_descriptor(descA.data, A_smem.ptr_to([0]), ldo, sdo, swizzle)
- Tx.cuda.tcgen05.encode_matrix_descriptor(descB.data, B_smem.ptr_to([0]), ldo, sdo, swizzle)
-
-**2. Choose the MMA tile.** ``M_mma × N_mma`` are chosen to tile ``M``/``N``
-(with ``MMA_K`` set by dtype: 8 tf32, 16 f16/bf16, 32 fp8, 64 fp4); a
-compile-time *instruction descriptor* packs the shape and dtypes.
-
-**3. Issue the async MMA** in an unrolled ``(mi, ni, ki)`` nest. The
-``enable_input_d`` input is ``accum or ki != 0``: with ``accum=False``, the first
-K issue overwrites the destination; with ``accum=True``, it also accumulates the
-destination's existing value.
-
-.. code-block:: python
-
- mma_chain = f"tcgen05.mma.cta_group::{cta_group}.kind::{kind}"
- Tx.ptx[mma_chain](
- Tx.cast(Tx.cuda.get_tmem_addr(tmem_addr, mi * M_mma, tmem_col), "uint32"), # C in tmem
- smem_desc_add_16B_offset(descA, a_off), descB_val, descI, # A / B descriptors
- *zero_masks, # disabled output lanes
- accum or ki != 0, # use existing C/D input
- )
-
-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 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
-``base + col``. This keeps the generated issue loop close to hand-written
-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
-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``.
-
-Accumulator datapaths and readback
-----------------------------------
-
-The accumulator layout must match the MMA's row placement:
-
-* Layout D is the M=128 identity placement.
-* Layout F is the single-CTA M=64 scattered placement.
-* Layout B is the per-CTA M=64 placement for ``cta_group=2``. Its logical
- N columns split across physical lane halves 0–63 and 64–127, so it
- occupies all 128 lanes and ``N/2`` tensor-memory columns.
-
-Allocate and read a Layout B result as follows:
-
-.. code-block:: python
-
- accumulator = tmem_pool.alloc_tcgen05_mma_D(
- (64, N), "float32", M=128, cta_group=2)
- Tx.cuda.tile.tcgen05.mma(
- accumulator[:, :],
- A_smem[:, :],
- B_smem[:, :],
- cta_group=2,
- )
-
- frag = Tx.alloc_tcgen05_ldst_frag("32x32b", (64, N), "float32")
- 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
-``.32x32b`` transfer over all 128 lanes. The gemm write-side layout and
-the allocation/readback layout are produced by the same
-``tmem_datapath_layout("B", 64, N)`` factory.
-
-Generated TIRx IR
------------------
-
-For the ``128×64 × 64×128`` fp16 tile (swizzle mode 3):
-
-.. code-block:: python
-
- Tx.cuda.tcgen05.encode_matrix_descriptor(Tx.address_of(descA[0]), Tx.address_of(A_smem[0]), 64, 64, 3)
- Tx.cuda.tcgen05.encode_matrix_descriptor(Tx.address_of(descB[0]), Tx.address_of(B_smem[0]), 64, 64, 3)
- Tx.ptx["tcgen05.mma.cta_group::1.kind::f16"](
- Tx.cast(Tx.cuda.get_tmem_addr(tmem_addr[0], mi * 128, 256 + ni * 128), "uint32"), ...)
-
-Generated CUDA
---------------
-
-.. code-block:: c++
-
- // async tensor-core MMA: A,B (shared, via descriptors) -> C (tmem)
- "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, ...;"
- // [%0] = C tmem address; %1 = A descriptor; %2 = B descriptor; %3 = instr descriptor
-
-``kind::f16`` selects the fp16/bf16 datapath. Verified on ``sm_100a`` (the tmem
-result, read back, equals ``A@B`` within fp16 tolerance).
-
-How inputs change the algorithm
--------------------------------
-
-.. list-table::
- :header-rows: 1
- :widths: 28 72
-
- * - input
- - effect
- * - dtype
- - ``tf32`` → ``MMA_K = 8``; ``float16``/``bfloat16`` → ``kind::f16`` and
- ``MMA_K = 16``; ``fp8`` → ``MMA_K = 32``; ``fp4`` → ``MMA_K = 64`` and
- the **block-scaled** path
- * - block scaling (SFA/SFB)
- - present → ``tcgen05.mma.block_scale`` with SFA/SFB tmem scale-factor
- addresses and a runtime-encoded or caller-supplied instruction descriptor
- * - cta_group
- - ``1`` → one CTA; ``2`` → two CTAs split the operand and each sees half
- of B's logical N extent. The selected physical instruction has
- cluster-wide M of 64/128 for one CTA or 128/256 for two CTAs; larger
- operand regions are tiled by those instruction blocks. The per-CTA
- M=64 output uses Layout B
- * - M / N / K extents
- - set the ``(mi, ni, ki)`` unrolled loop counts; K iterations accumulate into
- the same tmem accumulator. ``accum=True`` preserves and accumulates the
- value already in that accumulator on the first K iteration
- * - shared swizzle
- - sets the ``swizzle`` mode + ``ldo``/``sdo`` in the matrix descriptors;
- no-swizzle packed layouts are accepted when the selected tile has a
- hardware-compatible 16-byte packed stride
- * - ``smem_desc``
- - selects hoisted, call-site-hoisted, per-MMA encoded, or recomputed shared
- descriptor construction. Per-MMA ``"encode"`` is dense-only. The
- choice changes code shape only; the MMA operands still describe the same
- selected shared tiles.
- * - packed tmem accumulator
- - layouts of the form ``TileLayout(S[(M, 2, N//2) : (1@TLane, 64@TLane,
- 1@TCol)])`` are treated as packed ``N/2`` physical columns, matching
- FlashMLA-style low/high accumulator placement.
- * - ``weight_stationary``
- - on the dense path with ``cta_group=1``, selects the ``tcgen05.mma.ws``
- form when explicitly true; it is also inferred for the packed M=64
- Layout-E accumulator. The accumulator and A-operand layouts must match
- that datapath. The block-scaled path does not append ``.ws``, and the PTX
- table does not register a ``.ws.cta_group::2`` form.
diff --git a/python/tvm/backend/cuda/cpp/descriptors.py b/python/tvm/backend/cuda/cpp/descriptors.py
index 52b77d82af..e6c451e58e 100644
--- a/python/tvm/backend/cuda/cpp/descriptors.py
+++ b/python/tvm/backend/cuda/cpp/descriptors.py
@@ -23,7 +23,7 @@ asm, and nothing that maps to a PTX instruction.
The tcgen05 MMA validation tables live here too: they decide the dtype
*kind* and legal (M, N, K) shape that the instruction descriptor encodes,
-and ``tile/gemm_async/tcgen05.py`` reuses them when it folds the
+and ``tile/patterns/tcgen05_mma.py`` reuses them when it folds the
same constraints into its dispatch.
"""
@@ -526,7 +526,7 @@ def _get_tcgen05_mma_scale_vec_size(kind, scale_dtype):
# =============================================================================
# tcgen05 address / descriptor patch helpers — used by the dispatch wrappers
-# in ``tile/gemm_async/tcgen05.py``. They are tcgen05-specific rather
+# in ``tile/patterns/tcgen05_mma.py``. They are tcgen05-specific rather
# than generic address arithmetic:
# - get_tmem_addr packs a TMEM (taddr, row, col) tuple into the uint32 the
# PTX asm slots expect.
diff --git a/python/tvm/backend/cuda/lang/alloc_pool.py b/python/tvm/backend/cuda/lang/alloc_pool.py
index 1e9b3c24c6..d32364c8f2 100644
--- a/python/tvm/backend/cuda/lang/alloc_pool.py
+++ b/python/tvm/backend/cuda/lang/alloc_pool.py
@@ -52,7 +52,7 @@ def _shape_product(shape):
def _auto_swizzle_mode(dtype):
"""Select the default MMA swizzle mode for a shared-memory allocation."""
- from tvm.backend.cuda.tile.tma_utils import SwizzleMode
+ from tvm.backend.cuda.layout import SwizzleMode
del dtype
return SwizzleMode.SWIZZLE_128B_ATOM
@@ -60,7 +60,7 @@ def _auto_swizzle_mode(dtype):
def _swizzle_atom_bytes(swizzle_mode):
"""Return the row width (in bytes) of one swizzle atom for *swizzle_mode*."""
- from tvm.backend.cuda.tile.tma_utils import SwizzleMode
+ from tvm.backend.cuda.layout import SwizzleMode
return {
SwizzleMode.SWIZZLE_NONE: 0,
@@ -95,7 +95,7 @@ def _validate_mma_alloc_shape(shape, dtype, swizzle_mode):
Validation skipped when *swizzle_mode* is ``SWIZZLE_NONE`` (no atom).
"""
- from tvm.backend.cuda.tile.tma_utils import SwizzleMode
+ from tvm.backend.cuda.layout import SwizzleMode
if swizzle_mode == SwizzleMode.SWIZZLE_NONE:
return
@@ -146,7 +146,7 @@ def _validate_mma_alloc_shape(shape, dtype, swizzle_mode):
f"to a multiple of {min_cols} elements (>= {atom_bytes}B at {dtype})."
)
- # Atom rows is always 8 (see ``mma_atom_shape`` in tma_utils.py).
+ # Atom rows is always 8 (see ``mma_atom_shape`` in backend/cuda/layout.py).
atom_rows = 8
if rows < atom_rows or rows % atom_rows != 0:
raise ValueError(
@@ -234,7 +234,7 @@ class SMEMPool:
def alloc_tcgen05_mma_AB(self, shape, dtype="float16", swizzle_mode="auto", align=1024):
"""Allocate MMA-compatible shared memory with an inferred swizzle layout."""
- from tvm.backend.cuda.tile.tma_utils import (
+ from tvm.backend.cuda.layout import (
SwizzleMode,
mma_shared_layout,
)
diff --git a/python/tvm/backend/cuda/layout.py b/python/tvm/backend/cuda/layout.py
index 330afed506..28d5b026c3 100644
--- a/python/tvm/backend/cuda/layout.py
+++ b/python/tvm/backend/cuda/layout.py
@@ -16,9 +16,12 @@
# under the License.
"""CUDA-specific memory and instruction layout factories."""
+import copy
+from enum import Enum
+
import tvm
from tvm.backend.cuda.axis import Axis as CUDAAxis
-from tvm.tirx.layout import Axis, Iter, R, S, TileLayout
+from tvm.tirx.layout import Axis, ComposeLayout, Iter, Layout, R, S, TileLayout
# ============================================================================
# TMEM datapath layouts (PTX ISA §9.7.16.10.5)
@@ -174,15 +177,15 @@ def _mma_datapath_letter(M, cta_group, ws=False, sparse=False):
"""Map (instruction ``M``, ``cta_group``, ``ws``, ``sparse``) to the PTX
tcgen05 datapath letter. Raises for families rejected in the v1 semantic
allocator API (sparse Layout C, M=32 Layout G). See
- ``localdoc/claude_plan.txt`` S3a."""
+ the documented PTX TMEM datapaths."""
if sparse:
raise ValueError(
f"alloc_tcgen05_mma_AB: sparse A (Layout C) not supported in v1 "
f"(M={M}, cta_group={cta_group})"
)
- if ws and not (cta_group == 1 and M == 64):
+ if ws and not (cta_group == 1 and M in (64, 128)):
raise ValueError(
- f"alloc_tcgen05_mma_AB: ws=True only valid for cta_group=1, M=64 (Layout E); "
+ f"alloc_tcgen05_mma_AB: ws=True requires cta_group=1, M=64 or M=128; "
f"got M={M}, cta_group={cta_group}"
)
if cta_group == 1:
@@ -212,8 +215,7 @@ def tmem_mma_operand_layout(
per-CTA buffer rows; ``per_cta_rows = M // cta_group``. Physical 32-bit
column count is left to the pool's ``_resolve_cols`` (layout-aware).
- Reproduces the hand-written layouts the three FlashMLA kernels use; the
- supported/reject domain is spec'd in ``localdoc/claude_plan.txt`` S3a.
+ Selects the documented PTX datapath from explicit instruction parameters.
"""
tlane = Axis.get("TLane")
tcol = Axis.get("TCol")
@@ -231,7 +233,7 @@ def tmem_mma_operand_layout(
if operand == "D":
if group is not None:
# flat buffer (m, N) + explicit grouping -> (m, s, 2, n) layout
- # (codex plan option b): keeps kernel O buffer 2D, layout 4-axis.
+ # Keep the kernel output tensor 2D and its physical layout 4-axis.
if len(ext) != 2:
raise ValueError(f"alloc_tcgen05_mma_D group= requires 2D (m,N), got {ext}")
m, N = ext
@@ -310,6 +312,18 @@ def tmem_mma_operand_layout(
layout = TileLayout(S[(m, K) : (1 @ tlane, 1 @ tcol)])
else:
raise ValueError(f"tmem_mma_operand_layout: operand must be 'A' or 'D', got {operand!r}")
+ # PTX matrix D keeps sub-word accumulators in the low bits of a b32 cell.
+ # A remains packed. TCol strides are expressed in tensor element units.
+ from tvm.runtime import DataType
+
+ if operand == "D" and DataType(dtype).bits == 16:
+ layout = TileLayout.from_iters(
+ shard=[
+ Iter(i.extent, i.stride * (2 if i.axis.name == "TCol" else 1), i.axis)
+ for i in layout.shard
+ ],
+ replica=layout.replica,
+ )
return layout.canonicalize()
@@ -588,3 +602,200 @@ def tcgen05_atom_layout(instr_shape: str, tensor_shape: tuple[int, int], dtype)
iters = [Iter(ext, stride, axis) for ext, stride, axis in row_iters + col_iters]
return TileLayout.from_iters(iters, [], {})
+
+
+class SwizzleMode(Enum):
+ """The swizzle mode of the TMA."""
+
+ SWIZZLE_NONE = 0
+ SWIZZLE_32B_ATOM = 1
+ SWIZZLE_64B_ATOM = 2
+ SWIZZLE_128B_ATOM = 3
+
+
+def mma_atom_layout(dtype: str, swizzle_mode: SwizzleMode | int) -> ComposeLayout:
+ """Generate the MMA-compatible shared-memory atom layout."""
+ bits = tvm.DataType(dtype).bits
+ if isinstance(swizzle_mode, int):
+ swizzle_mode = SwizzleMode(swizzle_mode)
+ per_element = (128 // bits).bit_length() - 1
+ swizzle_len = swizzle_mode.value
+ atom_len = 3
+ period = 1 << (per_element + swizzle_len + atom_len)
+ return ComposeLayout(per_element, swizzle_len, atom_len, TileLayout(S[(period,)]))
+
+
+def mma_atom_shape(dtype: str, swizzle_mode: SwizzleMode | int, shape: list[int] | None = None):
+ """Generate the MMA-compatible shared-memory atom shape."""
+ bits = tvm.DataType(dtype).bits
+ if isinstance(swizzle_mode, int):
+ swizzle_mode = SwizzleMode(swizzle_mode)
+ atom_shape = {
+ SwizzleMode.SWIZZLE_32B_ATOM: [8, 256],
+ SwizzleMode.SWIZZLE_64B_ATOM: [8, 512],
+ SwizzleMode.SWIZZLE_128B_ATOM: [8, 1024],
+ }[swizzle_mode]
+ atom_shape[-1] //= bits
+ if shape is None:
+ return atom_shape
+ atom_shape = [1] * (len(shape) - len(atom_shape)) + atom_shape
+ return atom_shape
+
+
+def mma_shared_layout(dtype: str, swizzle_mode: SwizzleMode | int, shape) -> Layout:
+ """Generate the MMA-compatible shared-memory layout for shape and dtype.
+
+ It uses a default tiling strategy to tile the TMA atom layout into the shared memory.
+ """
+ if isinstance(swizzle_mode, int):
+ swizzle_mode = SwizzleMode(swizzle_mode)
+ if swizzle_mode == SwizzleMode.SWIZZLE_NONE:
+ # No-swizzle MMA smem is the packed-16B atom: offset e16*m +
+ # M*e16*(k//e16) + k%e16 (e16=128/bits); plain tile if K % e16 != 0.
+ bits = tvm.DataType(dtype).bits
+ e16 = 128 // bits
+ m, k = int(shape[-2]), int(shape[-1])
+ if k % e16 == 0:
+ lead = [int(s) for s in shape[:-2]]
+ extents = [*lead, m, k // e16, e16]
+ strides = []
+ stride = m * k
+ for e in reversed(lead):
+ strides.insert(0, stride)
+ stride *= e
+ strides += [e16, m * e16, 1]
+ return TileLayout(S[tuple(extents) : tuple(strides)]).canonicalize()
+ return TileLayout(S[tuple(shape)]).canonicalize()
+ atom_shape = mma_atom_shape(dtype, swizzle_mode, shape)
+ layout = mma_atom_layout(dtype, swizzle_mode)
+ tile_to_shape = copy.copy(atom_shape)
+ tile_to_shape[-2] = shape[-2]
+ return layout.tile_to(tile_to_shape, atom_shape).tile_to(shape, tile_to_shape).canonicalize()
+
+
+def sf_smem_layout(rows, SF_K, sf_per_mma, sf_reuse=1, pipe_depth=None):
+ """SMEM-side layout for SF in tcgen05.cp scale-factor copy.
+
+ The hardware reads SFs in 128-row super-blocks: 32 lanes x 16 bytes/lane.
+ The 16 bytes per lane row decompose as
+ ``M_SF_INNER (=4) x sf_per_mma x in_lane_K`` where
+ ``in_lane_K = epc / sf_per_mma`` and ``epc = 4`` (32-bit TMEM cell / 8-bit
+ SF). The remaining ``K_outer = SF_K / epc`` super-blocks march along K
+ with stride 512. ``sf_reuse > 1`` appends a stride-0 broadcast dim.
+
+ Buffer shape: ``(rows, SF_K * sf_reuse)`` (or ``(pipe_depth, rows,
+ SF_K * sf_reuse)``). Mirrors :func:`sf_tmem_layout` parameterization.
+
+ Args:
+ rows: Multiple of 128; M-direction rows.
+ SF_K: Number of unique SFs along K per row (multiple of 4).
+ sf_per_mma: Atom inner SFs per MMA-K step (must divide 4).
+ nvfp4=4, mxfp4=2, fp8=1.
+ sf_reuse: Broadcast factor (stride-0 dim). 1 = no broadcast.
+ pipe_depth: Optional pipeline depth as outermost dim.
+ """
+ epc = 4
+ M_SUPER_ROWS = 128
+ LANE = 32
+ M_SF_INNER = M_SUPER_ROWS // LANE
+ if rows % M_SUPER_ROWS != 0:
+ raise ValueError(f"rows={rows} must be a multiple of {M_SUPER_ROWS}")
+ if epc % sf_per_mma != 0:
+ raise ValueError(f"sf_per_mma={sf_per_mma} must divide epc={epc}")
+ if SF_K % epc != 0:
+ raise ValueError(f"SF_K={SF_K} must be a multiple of epc={epc}")
+
+ in_lane_K = epc // sf_per_mma
+ K_outer = SF_K // epc
+ M_super = rows // M_SUPER_ROWS
+ LANE_BYTES = epc * M_SF_INNER # 16
+ SUPER_BYTES = LANE_BYTES * LANE # 512
+ K_TOTAL_BYTES = SUPER_BYTES * K_outer
+ STAGE_BYTES = K_TOTAL_BYTES * M_super
+
+ raw_shape = [M_super, M_SF_INNER, LANE, K_outer, sf_per_mma, in_lane_K]
+ raw_strides = [K_TOTAL_BYTES, epc, LANE_BYTES, SUPER_BYTES, in_lane_K, 1]
+ if sf_reuse > 1:
+ raw_shape.append(sf_reuse)
+ raw_strides.append(0)
+ # Drop unit (extent-1) dims for cleaner canonical form.
+ shape = [s for s in raw_shape if s != 1]
+ strides = [st for s, st in zip(raw_shape, raw_strides) if s != 1]
+ if pipe_depth is not None:
+ shape = [pipe_depth, *shape]
+ strides = [STAGE_BYTES, *strides]
+ return TileLayout(S[tuple(shape) : tuple(strides)])
+
+
+def sf_tmem_layout(rows, SF_K, sf_per_mma, sf_reuse=1, pipe_depth=None):
+ """Create a TileLayout for SFA/SFB TMEM via atom direct_sum outer (+ optional reuse dim).
+
+ Args:
+ rows: CTA M-direction row count (multiple of 32).
+ SF_K: Number of *unique* SFs along K per row (loaded from gmem).
+ sf_per_mma: Atom inner SFs — number of SFs one MMA reads in K.
+ Equals ``mma_k // sf_vec``: nvfp4=4 (mma_k=64,sf_vec=16),
+ mxfp4=2 (64,32), fp8=1 (32,32).
+ sf_reuse: Number of MMAs that reuse one physical SF group via a
+ stride-0 broadcast dim. Equals ``quant_size // mma_k``;
+ default 1 (no reuse). fp8 blockwise with quant=128 and
+ mma_k=32 → ``sf_reuse=4``.
+ pipe_depth: Optional outer pipe-depth dim for double-buffered TMEM SF
+ allocations. Stride is ``M*epc @ CUDAAxis.TCol`` (one stage spans
+ ``M*epc`` cols). When ``None`` no pipe dim is added.
+
+ Buffer shape: ``(rows, SF_K * sf_reuse)`` (or ``(pipe_depth, rows,
+ SF_K * sf_reuse)``). Gemm dispatch iterates the last dim
+ ``SF_K * sf_reuse`` MMA times; only ``SF_K`` distinct SFs are physically
+ stored due to broadcast. Scale factor dtype is assumed 8-bit (epc=4);
+ all current SF formats (e8m0fnu, e4m3fn) fit.
+ """
+ if SF_K % sf_per_mma != 0:
+ raise ValueError(f"SF_K={SF_K} must be a multiple of sf_per_mma={sf_per_mma}")
+ K = SF_K // sf_per_mma # outer K iterations of unique SFs
+
+ M = rows // 32
+ epc = 4 # 32-bit TMEM column / 8-bit SF
+
+ # Atom: one 32-row chunk, one MMA's worth of SF.
+ atom = TileLayout(
+ S[(32, sf_per_mma) : (1 @ CUDAAxis.TLane, 1 @ CUDAAxis.TCol)] + R[4 : 32 @ CUDAAxis.TLane]
+ )
+
+ if K == 1:
+ outer = TileLayout(S[M : epc @ CUDAAxis.TCol])
+ else:
+ # Pack consecutive ki's within one uint32 TMEM column when possible.
+ pack_factor = epc // sf_per_mma
+ while pack_factor > 1 and K % pack_factor != 0:
+ pack_factor //= 2
+ if pack_factor > 1:
+ K_outer = K // pack_factor
+ if K_outer == 1:
+ outer = TileLayout(
+ S[(M, pack_factor) : (epc @ CUDAAxis.TCol, sf_per_mma @ CUDAAxis.TCol)]
+ )
+ else:
+ outer = TileLayout(
+ S[
+ (M, K_outer, pack_factor) : (
+ epc @ CUDAAxis.TCol,
+ M * epc @ CUDAAxis.TCol,
+ sf_per_mma @ CUDAAxis.TCol,
+ )
+ ]
+ )
+ else:
+ outer = TileLayout(S[(M, K) : (epc @ CUDAAxis.TCol, M * epc @ CUDAAxis.TCol)])
+
+ base = atom.direct_sum(outer, left_shape=[M, K], right_shape=[32, sf_per_mma])
+ if sf_reuse == 1 and pipe_depth is None:
+ return base
+ shard = list(base.shard)
+ if sf_reuse > 1:
+ # Append a stride-0 reuse dim on CUDAAxis.TCol for fp8 blockwise (vec_NX) mode.
+ shard.append(Iter(sf_reuse, 0, shard[0].axis))
+ if pipe_depth is not None:
+ # Prepend a pipe-depth dim that strides one stage (M*epc CUDAAxis.TCols).
+ shard.insert(0, Iter(pipe_depth, M * epc, shard[0].axis))
+ return TileLayout.from_iters(shard, list(base.replica), dict(base.offset))
diff --git a/python/tvm/backend/cuda/ptx/__init__.py b/python/tvm/backend/cuda/ptx/__init__.py
index e30175828a..f757a93b57 100644
--- a/python/tvm/backend/cuda/ptx/__init__.py
+++ b/python/tvm/backend/cuda/ptx/__init__.py
@@ -23,10 +23,17 @@ registers every table entry as a TVM Op with a generic codegen; the
``register_backend()`` via ``script_namespaces()``.
"""
-from .engine import PTXNamespace, register_addr, register_table
+from .engine import PTXNamespace, instruction_forms, instruction_name, register_addr, register_table
from .table import TABLE
register_addr()
register_table(TABLE)
-__all__ = ["TABLE", "PTXNamespace", "register_addr", "register_table"]
+__all__ = [
+ "TABLE",
+ "PTXNamespace",
+ "instruction_forms",
+ "instruction_name",
+ "register_addr",
+ "register_table",
+]
diff --git a/python/tvm/backend/cuda/ptx/engine.py b/python/tvm/backend/cuda/ptx/engine.py
index c17e243f77..97e0cbf667 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -666,7 +666,7 @@ def _coerce_pred(entry, pred):
raise ValueError(f"{entry.name}: pred must be a bool/uint32/int32 expression")
-def _emit(entry, filled, operands, pred=None, preserve_dst=False):
+def _resolve_modifiers(entry, filled):
# Modifiers resolve before the operands are looked at: the attribute chain
# is parsed before the call happens, and a register group's length may be a
# function of the modifiers, so the expected arity needs the modifier map.
@@ -680,6 +680,49 @@ def _emit(entry, filled, operands, pred=None, preserve_dst=False):
error = entry.check(mod_map)
if error:
raise ValueError(f"{entry.name}: {error}")
+ return mod_map
+
+
+def instruction_forms(instruction):
+ """Return table entries and resolved modifiers for a fully spelled opcode.
+
+ Operand overloads remain separate entries. Consumers select an operand form,
+ never inspect frontend candidate chains or infer missing opcode qualifiers.
+ """
+ if isinstance(instruction, str):
+ instruction = PTXNamespace()[instruction]
+ if not isinstance(instruction, _InstrChain):
+ raise TypeError("Expected a T.ptx instruction or PTX instruction string")
+ forms, errors = [], []
+ for entry, filled in instruction._cands:
+ try:
+ modifiers = _resolve_modifiers(entry, filled)
+ except ValueError as error:
+ errors.append(str(error))
+ else:
+ forms.append((entry, modifiers))
+ names = {
+ ".".join([entry.ptx_name, *(v for v in modifiers.values() if v)])
+ for entry, modifiers in forms
+ }
+ if len(names) != 1:
+ raise ValueError("PTX instruction is not fully specified: " + "; ".join(errors))
+ return tuple(forms)
+
+
+def instruction_name(instruction):
+ """Return validated PTX text in table slot order, without binding operands.
+
+ Accept a ``T.ptx`` chain or its string spelling. Operand overloads remain
+ unresolved: this identifies the opcode and modifiers, not an operand form.
+ Tile patterns can use it without inspecting the frontend's candidates.
+ """
+ entry, modifiers = instruction_forms(instruction)[0]
+ return ".".join([entry.ptx_name, *(v for v in modifiers.values() if v)])
+
+
+def _emit(entry, filled, operands, pred=None, preserve_dst=False):
+ mod_map = _resolve_modifiers(entry, filled)
layout = operand_layout(entry, mod_map)
n_args = sum(n for _, _, n in layout)
if len(operands) != n_args:
diff --git a/python/tvm/backend/cuda/tile/common.py b/python/tvm/backend/cuda/tile/common.py
deleted file mode 100644
index 49f5a6e694..0000000000
--- a/python/tvm/backend/cuda/tile/common.py
+++ /dev/null
@@ -1,278 +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.
-
-"""Common utilities for CUDA operator scheduling (basic helpers and copy ops)."""
-
-import functools
-import operator
-import re
-from enum import Enum
-
-from tvm.ir import TensorRegion
-from tvm.runtime import DataType, MemoryScope
-from tvm.script import tirx as T
-from tvm.sym.analyzer import Analyzer
-from tvm.tirx import Function, Var
-from tvm.tirx.tile import DispatchContext, fail
-from tvm.tirx.tile.instruction import TensorCall
-
-_SCOPE_GLOBAL = MemoryScope("global")
-_SCOPE_LOCAL = MemoryScope("local")
-_SCOPE_SHARED = MemoryScope("shared")
-_SCOPE_SHARED_DYN = MemoryScope("shared.dyn")
-
-
-def next_power_of_2(x: int) -> int:
- """Return the smallest power of 2 greater than or equal to x."""
- if x <= 1:
- return 1
- return 1 << (x - 1).bit_length()
-
-
-def get_st_extent(buffer_region: TensorRegion):
- """Get the start and extent of a buffer region."""
- region = buffer_region.region
- return [r.min for r in region], [r.extent for r in region]
-
-
-def get_indices(nth, start, extent):
- """Convert a fused index into multi-dimensional indices."""
- assert len(start) == len(extent)
- if len(start) == 1:
- return [start[0] + nth]
- relative = []
- for e in reversed(extent):
- relative.append(nth % e)
- nth //= e
- return [r + s for r, s in zip(reversed(relative), start)]
-
-
-def smem_desc_replace_lo(desc_val, desc_lo):
- """Replace the lower address lane of a 64-bit SMEM descriptor."""
- desc_halves = T.reinterpret(desc_val, ty="uint32x2")
- desc_hi = T.Shuffle([desc_halves], [1])
- return T.reinterpret(T.Shuffle([T.cast(desc_lo, "uint32"), desc_hi], [0, 1]), ty="uint64")
-
-
-def smem_desc_add_16B_offset(desc_val, offset):
- """Add a 16B-aligned byte offset to the lower 32 bits of a SMEM descriptor.
-
- The address lane wraps as uint32 without carrying into the descriptor's
- upper control bits.
- """
- desc_halves = T.reinterpret(desc_val, ty="uint32x2")
- desc_lo = T.Shuffle([desc_halves], [0]) + T.cast(offset, "uint32")
- return smem_desc_replace_lo(desc_val, desc_lo)
-
-
-class CopyInstType(Enum):
- """Enumeration of instruction types for memory operations."""
-
- NORMAL = 0
- CP_ASYNC = 1
-
-
-def validate_copy_op(
- op_call: TensorCall,
- sctx: DispatchContext, # pylint: disable=unused-argument
-) -> bool:
- """Sanity check for copy op"""
- dst_buffer_region, src_buffer_region = op_call.args[:2]
- src: Var = src_buffer_region.source
- dst: Var = dst_buffer_region.source
- if not (src.ty.effective_layout() and dst.ty.effective_layout() and src.dtype == dst.dtype):
- return False
- # Extract regions and validate dimensions
- analyzer = Analyzer()
- src_region, dst_region = src_buffer_region.region, dst_buffer_region.region
- # Extract extents and validate non-unit dimensions match
- src_extent_ = [r.extent for r in src_region if r.extent != 1]
- dst_extent_ = [r.extent for r in dst_region if r.extent != 1]
- if len(src_extent_) != len(dst_extent_) or not all(
- analyzer.can_prove_equal(s, d) for s, d in zip(src_extent_, dst_extent_)
- ):
- return False
- return True
-
-
-def get_vec_len(
- dst_buffer_region: TensorRegion,
- src_buffer_region: TensorRegion,
- vec_candidates: list[int],
- thread_cnt=1,
-) -> int | None:
- """Get the vector length for the copy operation."""
-
- dst: Var = dst_buffer_region.source
- src: Var = src_buffer_region.source
- # layout=None (flat local buffer) is treated as trivial for vectorization purposes
- if not (
- (dst.ty.effective_layout() is None or dst.ty.effective_layout().is_trivial())
- and (src.ty.effective_layout() is None or src.ty.effective_layout().is_trivial())
- ):
- return None
-
- # Extract regions and validate dimensions
- analyzer = Analyzer()
- src_st, src_extent = get_st_extent(src_buffer_region)
- dst_st, dst_extent = get_st_extent(dst_buffer_region)
-
- # Thread and vectorization setup
- DataType(src.dtype).bits # in bits
- n_elements = functools.reduce(operator.mul, src_extent, 1)
- if n_elements % thread_cnt != 0:
- return None
-
- # Find valid vector length
- for vec_len in vec_candidates:
- if vec_len > 0 and all(
- analyzer.can_prove_equal(x % vec_len, 0)
- for x in [
- src_st[-1],
- dst_st[-1],
- src.shape[-1] if len(src.shape) > 1 else 0,
- dst.shape[-1] if len(dst.shape) > 1 else 0,
- src_extent[-1],
- dst_extent[-1],
- n_elements // thread_cnt,
- ]
- ):
- return vec_len
- else:
- return None
-
-
-def copy_vec_load_impl(
- 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
- threads in a CTA/using a single thread.
- """
- dst_buffer_region, src_buffer_region = op_call.args[:2]
- src: Var = src_buffer_region.source
- dst: Var = dst_buffer_region.source
- if not (
- (src.scope() == _SCOPE_GLOBAL and (dst.scope() in (_SCOPE_SHARED, _SCOPE_SHARED_DYN)))
- or ((src.scope() in (_SCOPE_SHARED, _SCOPE_SHARED_DYN)) and dst.scope() == _SCOPE_GLOBAL)
- or (src.scope() == _SCOPE_GLOBAL and dst.scope() == _SCOPE_LOCAL)
- or (src.scope() == _SCOPE_LOCAL and dst.scope() == _SCOPE_GLOBAL)
- or ((src.scope() in (_SCOPE_SHARED, _SCOPE_SHARED_DYN)) and dst.scope() == _SCOPE_LOCAL)
- or ((dst.scope() in (_SCOPE_SHARED, _SCOPE_SHARED_DYN)) and src.scope() == _SCOPE_LOCAL)
- ):
- fail(f"unsupported memory scopes src={src.scope()} dst={dst.scope()}")
-
- # Thread and vectorization setup
- if sctx.is_cta:
- tx = sctx.launch_params["threadIdx.x"][1]
- assert "threadIdx.y" not in sctx.launch_params and "threadIdx.z" not in sctx.launch_params
- elif sctx.is_thread:
- tx = 1
- else:
- fail(f"unsupported exec_scope {sctx.scope_kind}")
-
- elem_size = DataType(src.dtype).bits # in bits
- vec_len = op_call.options.get("vec_len", None)
- if vec_len is None:
- vec_len = get_vec_len(
- dst_buffer_region,
- src_buffer_region,
- [128 // elem_size, 64 // elem_size, 32 // elem_size, 1],
- thread_cnt=tx,
- )
- if vec_len is None:
- fail("no valid vector length; check alignment/extents/thread-count")
-
- # cp-size (the size of data in bytes) can only be 4, 8 and 16 for cp.async
- if inst_type == CopyInstType.CP_ASYNC:
- cp_size = vec_len * elem_size // 8 # in bytes
- if cp_size not in [4, 8, 16]:
- fail("invalid cp.async cp_size; expected 4, 8 or 16 bytes")
-
- src_st, src_extent = get_st_extent(src_buffer_region)
- dst_st, dst_extent = get_st_extent(dst_buffer_region)
- n_elements = functools.reduce(operator.mul, src_extent, 1)
-
- if sctx.is_cta:
- # fmt: off
- @T.function
- def impl():
- """Implement copy operation with vectorized loads/stores."""
- for s in T.serial(0, n_elements // (tx * vec_len)):
- with T.launch_thread("threadIdx.x", tx) as tid_x:
- if T.constexpr(inst_type == CopyInstType.NORMAL):
- for vec in T.vectorized(vec_len):
- fused = T.meta_var((s * tx + tid_x) * vec_len + vec)
- dst_indices = T.meta_var(get_indices(fused, dst_st, dst_extent))
- src_indices = T.meta_var(get_indices(fused, src_st, src_extent))
- dst[tuple(dst_indices)] = src[tuple(src_indices)]
- elif T.constexpr(inst_type == CopyInstType.CP_ASYNC):
- fused = T.meta_var((s * tx + tid_x) * vec_len)
- dst_indices = T.meta_var(get_indices(fused, dst_st, dst_extent))
- src_indices = T.meta_var(get_indices(fused, src_st, src_extent))
- T.evaluate(T.ptx[f"cp.async.{'cg' if T.constexpr(cp_size == 16) else 'ca'}.shared.global"](dst.ptr_to(dst_indices), src.ptr_to(src_indices), cp_size)) # noqa: E501
- if T.constexpr(
- dst.scope() in (_SCOPE_SHARED, _SCOPE_SHARED_DYN)
- and inst_type == CopyInstType.NORMAL
- ):
- T.gpu_storage_sync("shared")
- # fmt: on
- elif sctx.is_thread:
- # fmt: off
- @T.function(check_well_formed=False)
- def impl():
- for s in T.serial(0, n_elements // (vec_len)):
- if T.constexpr(inst_type == CopyInstType.NORMAL):
- for vec in T.vectorized(vec_len):
- fused = T.meta_var(s * vec_len + vec)
- dst_indices = T.meta_var(get_indices(fused, dst_st, dst_extent))
- src_indices = T.meta_var(get_indices(fused, src_st, src_extent))
- dst[tuple(dst_indices)] = src[tuple(src_indices)]
- elif T.constexpr(inst_type == CopyInstType.CP_ASYNC):
- fused = T.meta_var(s * vec_len)
- dst_indices = T.meta_var(get_indices(fused, dst_st, dst_extent))
- src_indices = T.meta_var(get_indices(fused, src_st, src_extent))
- T.evaluate(T.ptx[f"cp.async.{'cg' if T.constexpr(cp_size == 16) else 'ca'}.shared.global"](dst.ptr_to(dst_indices), src.ptr_to(src_indices), cp_size)) # noqa: E501
- # fmt: on
- else:
- fail(f"unsupported exec_scope {sctx.scope_kind}")
- return impl
-
-
-def get_thread_cnt(sctx: DispatchContext) -> int | None:
- """Get thread count for the current execution scope."""
- scope_name = sctx.scope_kind
- if scope_name == "cta":
- return sctx.launch_params["threadIdx.x"][1]
- if scope_name == "warpgroup":
- return 128
- if scope_name == "warp":
- return 32
- if scope_name == "thread":
- return 1
- return None
-
-
-def sm_version_ok(
- 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 ""
- sm_match = re.match(r"sm_(\d+)", target_arch)
- sm_version = int(sm_match.group(1)) if sm_match else 0
- ok = sm_version >= min_version
- return (ok, None if ok else f"sm_version {sm_version} < {min_version}")
diff --git a/python/tvm/backend/cuda/tile/copy/__init__.py b/python/tvm/backend/cuda/tile/copy/__init__.py
deleted file mode 100644
index 56059157be..0000000000
--- a/python/tvm/backend/cuda/tile/copy/__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 .utils import _is_valid_copy, _scope_allowed, _single_thread_exec
diff --git a/python/tvm/backend/cuda/tile/copy/_common.py b/python/tvm/backend/cuda/tile/copy/_common.py
deleted file mode 100644
index b2668a3900..0000000000
--- a/python/tvm/backend/cuda/tile/copy/_common.py
+++ /dev/null
@@ -1,590 +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 partition / layout algorithm for synthesized-partition copy
-dispatches (currently ``gmem_smem`` and ``ldgsts``).
-
-``gmem_smem`` (sync ``Tx.copy`` global ↔ shared) and ``ldgsts`` (async
-``Tx.copy_async`` global → shared via cp.async / SASS LDGSTS) share the
-same algorithm to pick a vec-isolating + thread-distributing layout for
-``G ↔ S`` copies. Only emit-time details differ (which copy instruction
-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.tile.dispatch import DispatchContext
-
-from ..layout_utils import strip_swizzle_to_tile
-
-
-def _alignment_ok(vec_len: int, terms) -> bool:
- """Every term must be a multiple of ``vec_len``. Constants checked
- directly; Expr / symbolic terms checked via ``sym.Analyzer``.
-
- ``vec_len=1`` always passes (the scalar fallback). When a symbolic
- term can't be proved divisible, returns ``False`` conservatively —
- the candidate loop will then try a smaller ``vec_len``.
- """
- if vec_len <= 1:
- return True
- analyzer = sym.Analyzer()
- for t in terms:
- if isinstance(t, int):
- if t % vec_len != 0:
- return False
- else:
- if not analyzer.can_prove_equal(t % vec_len, 0):
- return False
- return True
-
-
-# scope_kind → CUDA index call that decomposes the scope into individual threads.
-_TID_AXIS_FOR_SCOPE = {
- "warp": "laneid",
- "warpgroup": "tid_in_wg",
- "cta": "tx",
-}
-
-
-def _thread_cnt(sctx: DispatchContext) -> int:
- """Total threads active in the current scope = ∏ intra-axis extents.
-
- For thread scope ``sctx.intra`` is empty → returns 1.
- """
- n = 1
- for ext, _off in sctx.intra.values():
- n *= int(ext)
- return n
-
-
-# -----------------------------------------------------------------------------
-# Layout primitives
-# -----------------------------------------------------------------------------
-
-
-def _contig_group(iters: list) -> list[int]:
- """Indices (in iters) of the maximal physical-contiguous chain starting
- at the stride=1 iter, ordered stride-ascending.
-
- Returns [] if no stride=1 iter exists.
- """
- one_idx = next(
- (i for i, it in enumerate(iters) if int(it.stride) == 1),
- None,
- )
- if one_idx is None:
- return []
- chain = [one_idx]
- acc = int(iters[one_idx].extent)
- used = {one_idx}
- while True:
- nxt = next(
- (i for i, it in enumerate(iters) if i not in used and int(it.stride) == acc),
- None,
- )
- if nxt is None:
- break
- chain.append(nxt)
- acc *= int(iters[nxt].extent)
- used.add(nxt)
- return chain
-
-
-def _try_split_vec(iters: list, vec_len: int):
- """Try to walk ``vec_len`` consecutive elements along the contig chain.
-
- Returns ``(new_iters, selected_positions)`` on success, ``None`` on
- failure. ``new_iters`` may contain a freshly-split iter (replacing one
- entry with its "outer" half, with the "inner" half appended at the end);
- ``selected_positions`` are positions in ``new_iters`` that together
- cover the ``vec_len`` contig elements.
- """
- chain = _contig_group(iters)
- if not chain:
- return None
- rem = vec_len
- new_iters = list(iters)
- selected: list[int] = []
- for orig_idx in chain:
- if rem == 0:
- break
- it = new_iters[orig_idx]
- ext = int(it.extent)
- if ext <= rem:
- if rem % ext != 0:
- return None
- selected.append(orig_idx)
- rem //= ext
- else:
- if ext % rem != 0:
- return None
- stride = int(it.stride)
- outer = Iter(ext // rem, stride * rem, it.axis)
- inner = Iter(rem, stride, it.axis)
- new_iters[orig_idx] = outer
- new_iters.append(inner)
- selected.append(len(new_iters) - 1)
- rem = 0
- break
- if rem != 0:
- return None
- return new_iters, selected
-
-
-def _isolated_shape(iters: list, selected: list[int]) -> tuple[list[int], list[tuple[int, int]]]:
- """Build the isolated shape: each selected iter is its own segment;
- adjacent unselected iters are merged into a single segment.
-
- Returns ``(shape, segments)`` where ``segments[i] = (start, end)`` is
- the half-open range in ``iters`` covered by shape entry ``i``.
- """
- sel_set = set(selected)
- shape: list[int] = []
- segments: list[tuple[int, int]] = []
- cur_start = None
- cur_ext = 1
- for i, it in enumerate(iters):
- if i in sel_set:
- if cur_start is not None:
- shape.append(cur_ext)
- segments.append((cur_start, i))
- cur_start = None
- cur_ext = 1
- shape.append(int(it.extent))
- segments.append((i, i + 1))
- else:
- if cur_start is None:
- cur_start = i
- cur_ext *= int(it.extent)
- if cur_start is not None:
- shape.append(cur_ext)
- segments.append((cur_start, len(iters)))
- return shape, segments
-
-
-def _vec_perm(iters: list, selected: list[int]) -> list[int]:
- """Reorder ``iters`` into ``[outer, vec]``, both ordered by stride
- descending so the stride=1 iter ends up at the very last position."""
- sel_set = set(selected)
- unsel_sorted = sorted(
- (i for i in range(len(iters)) if i not in sel_set),
- key=lambda i: -int(iters[i].stride),
- )
- sel_sorted = sorted(selected, key=lambda i: -int(iters[i].stride))
- return list(unsel_sorted) + sel_sorted
-
-
-def _try_split_thread(iters: list, vec_selected: list[int], thread_cnt: int):
- """After ``_try_split_vec``, carve ``thread_cnt`` from the OUTER tail
- (smallest-stride outer iter, then towards bigger stride if needed).
-
- Unlike vec split, this doesn't require physical contiguity — T
- consecutive fused indices map to per-thread offsets via the layout's
- stride (which may be > 1).
-
- Returns ``(new_iters, thread_selected_positions)`` on success, ``None``
- on failure (outer doesn't divide T cleanly, or no outer iters left).
- """
- if thread_cnt == 1:
- return list(iters), []
- vec_set = set(vec_selected)
- outer = [i for i in range(len(iters)) if i not in vec_set]
- if not outer:
- return None
- outer_by_stride_desc = sorted(outer, key=lambda i: -int(iters[i].stride))
- rem = thread_cnt
- new_iters = list(iters)
- thread_selected: list[int] = []
- for orig_idx in reversed(outer_by_stride_desc):
- if rem == 0:
- break
- it = new_iters[orig_idx]
- ext = int(it.extent)
- if ext <= rem:
- if rem % ext != 0:
- return None
- thread_selected.append(orig_idx)
- rem //= ext
- else:
- if ext % rem != 0:
- return None
- stride = int(it.stride)
- new_iters[orig_idx] = Iter(ext // rem, stride * rem, it.axis)
- new_iters.append(Iter(rem, stride, it.axis))
- thread_selected.append(len(new_iters) - 1)
- rem = 0
- break
- if rem != 0:
- return None
- return new_iters, thread_selected
-
-
-def _three_segment_perm(iters: list, t_selected: list[int], vec_selected: list[int]) -> list[int]:
- """Reorder ``iters`` into ``[outer, T, vec]`` segments. Within each
- segment, stride descending so stride=1 sits at the very end."""
- t_set = set(t_selected)
- vec_set = set(vec_selected)
- outer = sorted(
- (i for i in range(len(iters)) if i not in t_set and i not in vec_set),
- key=lambda i: -int(iters[i].stride),
- )
- t_sorted = sorted(t_selected, key=lambda i: -int(iters[i].stride))
- vec_sorted = sorted(vec_selected, key=lambda i: -int(iters[i].stride))
- return list(outer) + t_sorted + vec_sorted
-
-
-def _shape_perm_for_isolated(
- shape_segments: list[tuple[int, int]], iter_perm: list[int]
-) -> list[int]:
- """Given segments (one per shape entry, each = (start, end) in
- pre-perm iter positions) and the iter permutation, compute the
- corresponding shape permutation."""
- seg_of = [0] * sum(end - start for start, end in shape_segments)
- for seg_idx, (start, end) in enumerate(shape_segments):
- for k in range(start, end):
- seg_of[k] = seg_idx
- seen: set[int] = set()
- perm: list[int] = []
- for orig_idx in iter_perm:
- seg_idx = seg_of[orig_idx]
- if seg_idx not in seen:
- seen.add(seg_idx)
- perm.append(seg_idx)
- return perm
-
-
-def _verify_s_tail_contig(s_p: TileLayout, vec_len: int) -> bool:
- """Check the last iters of ``s_p`` form a stride=1 contig chain whose
- extent product equals ``vec_len``."""
- iters = list(s_p.shard)
- if not iters:
- return vec_len == 1
- last = iters[-1]
- if int(last.stride) != 1:
- return False
- acc = int(last.extent)
- if acc == vec_len:
- return True
- for k in range(len(iters) - 2, -1, -1):
- it = iters[k]
- if int(it.stride) != acc:
- break
- acc *= int(it.extent)
- if acc == vec_len:
- return True
- if acc > vec_len:
- return False
- return acc >= vec_len and acc % vec_len == 0
-
-
-_VEC_BITS_CANDIDATES = (128, 64, 32, 16, 8)
-
-
-def _vec_len_candidates(elem_bits: int, allowed_bits: tuple | None = None) -> list[int]:
- """Vec-length candidates (in elements) for the given element width.
-
- ``allowed_bits`` optionally filters the per-instruction allowed widths
- (e.g. cp.async only accepts {128, 64, 32} bits = 16/8/4 bytes).
- Defaults to ``_VEC_BITS_CANDIDATES`` if not specified.
- """
- bits_tuple = allowed_bits if allowed_bits is not None else _VEC_BITS_CANDIDATES
- out: list[int] = []
- for vb in bits_tuple:
- if vb < elem_bits or vb % elem_bits != 0:
- continue
- n = vb // elem_bits
- if n not in out:
- out.append(n)
- if 1 not in out and allowed_bits is None:
- # Scalar fallback is only added for the unrestricted candidate set;
- # an instruction-specific list (cp.async etc.) keeps its strictness.
- out.append(1)
- return out
-
-
-def _extract_tile(layout, region):
- """Strip swizzle so we can perm/group as a TileLayout."""
- # Region bounds may be constant-valued but remain as unfolded expressions
- # after substitution. Simplify before converting to a Python integer;
- # genuinely symbolic extents fall back inside ``strip_swizzle_to_tile``.
- analyzer = sym.Analyzer()
- return strip_swizzle_to_tile(
- layout, lambda: [int(analyzer.simplify(end - start)) for (start, end) in region]
- )
-
-
-def _sort_by_stride_desc(layout: TileLayout) -> TileLayout:
- """Reorder shard so list order = traversal order (outer first, stride=1
- last). Required before canonicalize() can fuse non-adjacent-but-contig
- iters."""
- iters = list(layout.shard)
- perm = sorted(range(len(iters)), key=lambda i: -int(iters[i].stride))
- if perm == list(range(len(iters))):
- return layout
- return layout.permute_dims(perm)
-
-
-def _carve_tail(iters: list, chunk: int):
- """Carve ``chunk`` elements off the tail. Walk back across multiple
- iters as needed; at most one iter is split.
-
- Per iter (from last to first), let ``ext`` = iter extent and ``rem``
- = remaining chunk to fill:
-
- * ``ext == rem``: eat this iter whole, done.
- * ``ext < rem``: must divide ``rem``; eat whole, ``rem //= ext``.
- * ``ext > rem``: must divide ``ext``; split into
- ``(ext/rem, stride*rem) + (rem, stride)``, take the inner. Done.
-
- Returns the new iter list on success, ``None`` on failure.
- """
- if not iters or chunk <= 0:
- return None
- rem = chunk
- work = list(iters)
- for idx in range(len(work) - 1, -1, -1):
- it = work[idx]
- ext = int(it.extent)
- if ext == rem:
- return work
- if ext < rem:
- if rem % ext != 0:
- return None
- rem //= ext
- continue
- if ext % rem != 0:
- return None
- stride = int(it.stride)
- work[idx] = Iter(ext // rem, stride * rem, it.axis)
- work.insert(idx + 1, Iter(rem, stride, it.axis))
- return work
- return None
-
-
-def align_layouts_gs(
- g_layout,
- g_shape,
- g_region,
- s_layout,
- s_shape,
- s_region,
- elem_bits,
- thread_cnt: int,
- vec_bits_candidates: tuple | None = None,
- debug: bool = False,
-):
- """Align G and S layouts for a synthesized G↔S copy.
-
- Algorithm:
- 1. Sort G iters by stride desc + canonicalize (fuses anything physically
- contig). Same for S.
- 2. Group S by G's iter shape (per-iter shape list) + permute_by_groups
- with identity so S's groups line up with G's iters.
- 3. Carve ``vec_len`` off G's tail (a single iter split if needed).
- 4. Carve ``T = thread_cnt`` off the iter before vec (multi-iter walk,
- single split).
- 5. Re-group S (already permuted) by the new finer per-iter shape. No
- further permute needed because steps 3-4 only refine the tail.
- 6. Verify S's tail (vec segment) is physically contig.
-
- ``vec_bits_candidates`` optionally restricts the per-instruction allowed
- widths (e.g. cp.async only accepts {128, 64, 32} bits). When ``None``
- (default), uses the full {128, 64, 32, 16, 8} set plus a scalar (1)
- fallback.
-
- Returns ``(g_p, s_p, vec_len)``. ``g_p.shard`` ends as
- ``[outer iters..., T iter, vec iter]``; ``s_p.shard`` has the same iter
- count and matching iter-by-iter extents.
- """
- g = g_layout.slice(list(g_shape), g_region)
- s = s_layout.slice(list(s_shape), s_region)
- # Detect a swizzle on the S side BEFORE _extract_tile strips it.
- # vec_len must fit inside one swizzle chunk (C = 2^per_element elements);
- # otherwise the vec ld/st crosses a swizzle XOR boundary and hits the
- # wrong physical bytes mid-vec.
- s_swizzle_chunk_elems = None
- if isinstance(s_layout, ComposeLayout):
- s_swizzle_chunk_elems = 1 << int(s_layout.per_element)
- g = _extract_tile(g, g_region)
- s = _extract_tile(s, s_region)
-
- # Only G drives the canonical form. S's iter order is derived from G's
- # pre-sort iter extents (used as the grouping shape) and then permuted
- # by G's stride-desc permutation. Independently sorting/canonicalizing
- # S would fuse iters whose strides happen to chain into one contiguous
- # range — that loses the layout's logical `(i, j) → addr` mapping for
- # layout-permuting copies (e.g. row-major GMEM → K-tiled SMEM, where
- # both layouts cover the same byte range but with different coord maps).
- g_pre_sort_extents = [int(it.extent) for it in g.shard]
- g_perm = sorted(range(len(g.shard)), key=lambda i: -int(g.shard[i].stride))
- try:
- s_grp1, seps1 = s.group(g_pre_sort_extents)
- except Exception as e:
- if debug:
- print(f" Step-2 S.group({g_pre_sort_extents}) failed: {e}")
- return _sort_by_stride_desc(g).canonicalize(), s, 1
- # S iters keep their original strides; only the *group order* is
- # rearranged to follow G's sorted iter order. Canonicalize after the
- # permute is safe — it only fuses iters whose strides genuinely chain
- # in the post-permute order, which preserves the logical structure.
- s_aligned = s_grp1.permute_by_groups(list(seps1), g_perm).canonicalize()
- g = _sort_by_stride_desc(g).canonicalize()
- if debug:
- print(f" g (sort+canon): shard={[(int(it.extent), int(it.stride)) for it in g.shard]}")
- print(
- f" s (grouped+permuted by G shape {g_pre_sort_extents}): shard="
- f"{[(int(it.extent), int(it.stride)) for it in s_aligned.shard]}"
- )
- print(f" thread_cnt: {thread_cnt}")
-
- dummy_axis = g.shard[-1].axis if g.shard else None
-
- for vec_len in _vec_len_candidates(elem_bits, vec_bits_candidates):
- if vec_len == 1:
- g_after_vec = [*list(g.shard), Iter(1, 1, dummy_axis)]
- else:
- # Swizzle chunk-size cap: vec must fit in one swizzle chunk.
- if s_swizzle_chunk_elems is not None and vec_len > s_swizzle_chunk_elems:
- if debug:
- print(
- f" vec_len={vec_len}: exceeds swizzle chunk "
- f"({s_swizzle_chunk_elems} elements)"
- )
- continue
- g_after_vec = _carve_tail(list(g.shard), vec_len)
- if g_after_vec is None:
- if debug:
- print(f" vec_len={vec_len}: G vec carve failed")
- continue
- outer_for_t = g_after_vec[:-1]
- outer_after_t = _carve_tail(outer_for_t, thread_cnt) if thread_cnt > 1 else outer_for_t
- if outer_after_t is None:
- if debug:
- print(f" vec_len={vec_len}: T={thread_cnt} carve failed")
- continue
- g_final_iters = [*outer_after_t, g_after_vec[-1]]
- new_shape = [int(it.extent) for it in g_final_iters]
- try:
- s_final_grp, _ = s_aligned.group(new_shape)
- except Exception as e:
- if debug:
- print(f" vec_len={vec_len}: S.group({new_shape}) failed: {e}")
- continue
- g_p = TileLayout.from_iters(g_final_iters, list(g.replica), dict(g.offset))
- s_p = s_final_grp
-
- ok = _verify_s_tail_contig(s_p, vec_len)
- if debug:
- print(
- f" vec_len={vec_len}: shape={new_shape}, "
- f"g_p.shard={[(int(it.extent), int(it.stride)) for it in g_p.shard]}, "
- f"s_p.shard={[(int(it.extent), int(it.stride)) for it in s_p.shard]}, "
- f"s_tail_contig={ok}"
- )
- if not ok:
- continue
- # Alignment: per-thread starting addr = base_ptr + sizeof(elem) *
- # (region_base + tid*t_stride + outer_iter_strides). For the
- # vec_bits/8 byte vector op to be naturally aligned, every one of
- # those element-count terms must be a multiple of vec_len.
- align_terms = []
- for it in s_p.shard[:-1]:
- align_terms.append(int(it.stride))
- for it in g_p.shard[:-1]:
- align_terms.append(int(it.stride))
- align_terms.extend(s_p.offset.values())
- align_terms.extend(g_p.offset.values())
- if not _alignment_ok(vec_len, align_terms):
- if debug:
- print(f" vec_len={vec_len}: alignment check failed")
- continue
- return g_p, s_p, vec_len
-
- return g, s_aligned, 1
-
-
-def _flat_outer_coords(outer_exts: list[int], flat_idx: int) -> list[int]:
- """Decode a row-major flat index into per-iter coords. ``outer_exts`` is
- in stride-desc order (outermost first), so the first coord changes
- slowest."""
- coords: list[int] = []
- rem = flat_idx
- for ext in reversed(outer_exts):
- coords.append(rem % ext)
- rem //= ext
- coords.reverse()
- return coords
-
-
-def _outer_offsets(outer_iters_s, outer_iters_g, flat_idx):
- """Returns ``(ds, dg)``: constant offsets on S and G sides for the
- given flat outer-loop iteration."""
- outer_exts = [int(it.extent) for it in outer_iters_s]
- coords = _flat_outer_coords(outer_exts, flat_idx)
- ds = sum(c * int(it.stride) for c, it in zip(coords, outer_iters_s))
- dg = sum(c * int(it.stride) for c, it in zip(coords, outer_iters_g))
- return ds, dg
-
-
-def copy_ptx_form(num_bytes: int) -> tuple[str, int, str]:
- """Map copy width (bytes) to ``(ptx chain tail, lanes, register dtype)``.
-
- The chain tail is the ``{vec.}type`` suffix; ``lanes`` is how many operands
- the group takes; the register dtype is the container the lanes are read and
- written through, which follows the PTX type token rather than the element
- type -- these copies move bits, not values.
- """
- return {
- 32: ("v8.u32", 8, "uint32"),
- 16: ("v4.u32", 4, "uint32"),
- 8: ("v2.u32", 2, "uint32"),
- 4: ("u32", 1, "uint32"),
- 2: ("u16", 1, "uint16"),
- 1: ("u8", 1, "uint8"),
- }[num_bytes]
-
-
-def copy_ptx_ld_chain(
- space: str,
- tail: str,
- *,
- nc: bool = False,
- l1_evict: str = "",
- l2_evict: str = "",
- prefetch_size: str = "",
-) -> str:
- """Assemble an ld chain in the ISA's qualifier order.
-
- Order matters for the printed instruction, not for filling the slots:
- ``ld.global.nc.L1::no_allocate.L2::evict_first.L2::256B.v8.u32``.
- """
- toks = ["ld", space]
- if nc:
- toks.append("nc")
- if l1_evict:
- toks.append(l1_evict)
- if l2_evict:
- if not tail.startswith("v"):
- # The ISA spells .level2::eviction_priority only on the 256-bit
- # vector lines, so a scalar copy has nowhere to put it.
- raise ValueError("L2 eviction priorities are vector-only")
- toks.append(l2_evict)
- if prefetch_size:
- toks.append(prefetch_size)
- toks.append(tail)
- return ".".join(toks)
diff --git a/python/tvm/backend/cuda/tile/copy/ld_stmatrix.py b/python/tvm/backend/cuda/tile/copy/ld_stmatrix.py
deleted file mode 100644
index b8d3f69987..0000000000
--- a/python/tvm/backend/cuda/tile/copy/ld_stmatrix.py
+++ /dev/null
@@ -1,401 +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 variant: ldmatrix / stmatrix (TBD algorithm).
-
-Handles register ↔ shared copies on CUDA via PTX ``ldmatrix`` / ``stmatrix``.
-Direction (ld vs st) and exec scope (warp / warpgroup) are decided inside
-``_emit`` from the src/dst scopes and ``sctx.scope_kind``.
-"""
-
-from math import prod
-
-from tvm.runtime import MemoryScope
-from tvm.script import tirx as T
-from tvm.tirx import Function
-from tvm.tirx.layout import ComposeLayout, S, TileLayout
-from tvm.tirx.tile.dispatch import DispatchContext, fail
-from tvm.tirx.tile.instruction import TensorCall
-
-from ._common import ( # noqa: F401 (_carve_tail reserved for future variants)
- _carve_tail,
- _extract_tile,
-)
-from .utils import _is_valid_copy, _scope_allowed
-from .vec_auto_reg import _all_threads_active, _ptr_off
-
-_SCOPE_SHARED = MemoryScope("shared")
-_SCOPE_SHARED_DYN = MemoryScope("shared.dyn")
-
-_SCOPE_LOCAL = MemoryScope("local")
-
-_REG_SMEM_PAIRS = (
- (_SCOPE_LOCAL, _SCOPE_SHARED),
- (_SCOPE_LOCAL, _SCOPE_SHARED_DYN),
- (_SCOPE_SHARED, _SCOPE_LOCAL),
- (_SCOPE_SHARED_DYN, _SCOPE_LOCAL),
-)
-
-_VALID_R_LANE_AXES = {"laneid", "tid_in_wg", "tx"}
-
-
-def _compute_r_perm(r):
- """Permutation: thread iters first (stride-desc), then memory iters (stride-desc)."""
-
- def key(p):
- it = p[1]
- return (0 if it.axis.is_thread() else 1, -int(it.stride))
-
- return [i for i, _ in sorted(enumerate(r.shard), key=key)]
-
-
-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"):
- return False, f"unsupported exec_scope {sctx.scope_kind} (need warp, warpgroup, or cta)"
- for check in (
- lambda: _all_threads_active(sctx),
- lambda: _is_valid_copy(op_call, sctx),
- lambda: _scope_allowed(op_call, sctx, allowed_pairs=_REG_SMEM_PAIRS),
- ):
- ok, msg = check()
- if not ok:
- return False, msg
- return True, None
-
-
-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
- dst_br = op_call.dst
- if src_br.source.scope() == _SCOPE_LOCAL:
- r_br, s_br = src_br, dst_br
- direction = "st" # reg -> smem (stmatrix)
- else:
- r_br, s_br = dst_br, src_br
- direction = "ld" # smem -> reg (ldmatrix)
- r_buf = r_br.source
- s_buf = s_br.source
- r_shape = list(r_buf.shape)
- r_layout = r_buf.ty.effective_layout()
- s_shape = list(s_buf.shape)
- s_layout = s_buf.ty.effective_layout()
-
- r_region = [(r.min, r.min + r.extent) for r in r_br.region]
- s_region = [(r.min, r.min + r.extent) for r in s_br.region]
- with sctx.target:
- r_sliced = r_layout.slice(r_shape, r_region)
- s_sliced = s_layout.slice(s_shape, s_region)
- r = r_sliced.canonicalize()
- s = s_sliced.canonicalize()
-
- # Peel the S-side wrapper before structural checks, retaining its
- # parameters so the selected TileLayout can be wrapped again at emit.
- s_swizzle = (
- s_buf.ty.effective_layout()
- if isinstance(s_buf.ty.effective_layout(), ComposeLayout)
- else None
- )
- if s_swizzle is not None and s_swizzle.per_element < 3:
- # ldmatrix/stmatrix .b16 reads/writes 8 fp16 = 128b per lane in one
- # contiguous chunk. The swizzle preserves the lowest ``per_element``
- # bits of the address (in-chunk offset). For the per-lane 128b unit
- # to stay contiguous post-swizzle, ``2^per_element >= 8`` ⇒ p >= 3.
- fail(
- f"swizzle per_element={s_swizzle.per_element} < 3 incompatible "
- f"with .b16 ldmatrix/stmatrix (need 8-fp16 chunk integrity)"
- )
- s = _extract_tile(s, s_region)
-
- # Step 3: ldstmatrix doesn't broadcast — require zero replica on both sides.
- if len(r.replica) != 0:
- fail(f"R layout has replica {list(r.replica)}; ldstmatrix requires no replica")
- if len(s.replica) != 0:
- fail(f"S layout has replica {list(s.replica)}; ldstmatrix requires no replica")
-
- # Step 4: R must have exactly one kind of lane axis from the valid set.
- r_thread_axes = {it.axis.name for it in r.shard if it.axis.is_thread()}
- if len(r_thread_axes) != 1:
- fail(f"R must have exactly one thread axis name; got {sorted(r_thread_axes)}")
- r_lane_axis = next(iter(r_thread_axes))
- if r_lane_axis not in _VALID_R_LANE_AXES:
- fail(f"R thread axis {r_lane_axis!r} not in {sorted(_VALID_R_LANE_AXES)}")
-
- # Step 5: group S by R's iter extents (one S group per R iter, outer→inner).
- r_group_shape = [int(it.extent) for it in r.shard]
- s_grp, s_seps = s.group(r_group_shape)
-
- # Step 6: permute R so thread iters come first (stride-desc), then memory
- # iters (stride-desc).
- r_perm = _compute_r_perm(r)
- r = r.permute_dims(r_perm)
-
- # Step 7: apply R's perm to S in group units (1-to-1 with R's iters), and
- # rebuild s_seps to track group boundaries in the new order.
- s = s_grp.permute_by_groups(list(s_seps), r_perm)
- old_sizes = [s_seps[i + 1] - s_seps[i] for i in range(len(s_seps) - 1)]
- s_seps = [0]
- for pi in r_perm:
- s_seps.append(s_seps[-1] + old_sizes[pi])
-
- # Step 7.5: canonicalize both R and S after permute. Fuses adjacent
- # contig iters — keeps step 8's group input clean. Push target so
- # scope-aware fusers run (laneid+wid_in_wg → tid_in_wg, etc.).
- with sctx.target:
- r = r.canonicalize()
- s = s.canonicalize()
-
- t_total = prod(int(it.extent) for it in r.shard if it.axis.is_thread())
- m_total = prod(int(it.extent) for it in r.shard if not it.axis.is_thread())
- if t_total % 32 != 0:
- fail(f"R thread section total {t_total} not divisible by 32")
-
- def _strs(lay, seps):
- # Atoms 8 / 4 / 2 (segs 1, 2, 5) must be single iters — their strides
- # feed downstream stride checks (lane partition + fragment 2-fp16
- # contig). The num atom (seg 4) may be MULTI-ITER: we return its iter
- # list and let layout.apply handle the decomposition at emit time.
- fixed_segs = [list(lay.shard[seps[i] : seps[i + 1]]) for i in (1, 2, 5)]
- if not all(len(g) == 1 for g in fixed_segs):
- return None
- num_iters = list(lay.shard[seps[4] : seps[5]])
- return (
- int(fixed_segs[0][0].stride), # 8 atom stride
- int(fixed_segs[1][0].stride), # 4 atom stride
- num_iters, # num atom iter list (multi-iter OK)
- int(fixed_segs[2][0].stride), # 2 atom stride
- )
-
- def _try_num(r_in, s_in, num):
- """Try grouping (r_in, s_in) with [T/32, 8, 4, M/(2num), num, 2].
-
- Returns (rg, rsep, sg, ssep, trans, p, num) if structural checks pass,
- else None. ``trans`` is the ldmatrix .trans flag; ``p`` is the
- per-tile-row S stride used at emit.
- """
- gs = [t_total // 32, 8, 4, m_total // (num * 2), num, 2]
- try:
- rg, rsep = r_in.group(gs)
- sg, ssep = s_in.group(gs)
- except Exception:
- return None
- # R seg 0 (T/32 outer): require single iter with stride 32. When
- # T/32 == 1 the segment is trivial — skip.
- if t_total > 32:
- seg0 = list(rg.shard[rsep[0] : rsep[1]])
- if len(seg0) != 1 or int(seg0[0].stride) != 32:
- return None
- rs, ss = _strs(rg, rsep), _strs(sg, ssep)
- if rs is None or ss is None:
- return None
- r8, r4, _r_num_iters, r2 = rs
- s8, s4, s_num_iters, s2 = ss
- if (r8, r4, r2) != (4, 1, 1):
- return None
- # S num atom: every iter must have stride > 0 and multiple of 8 (the
- # per-tile spacing geometry of ldmatrix m8n8; 8 fp16 = 16 bytes = one
- # tile column dimension).
- if num > 1 and not all(
- int(it.stride) > 0 and int(it.stride) % 8 == 0 for it in s_num_iters
- ):
- return None
- # m_outer (seg 3) iters: each per-mm advance must keep the per-lane
- # SMEM address 16-byte aligned (ldmatrix .b16 reads 8 fp16 = 16 bytes
- # per lane), so the m_outer S-stride must also be a multiple of 8.
- # Without this, mm > 0 iterations land at unaligned addresses and
- # silently read garbage even though the layout group succeeds.
- # Skip extent-1 trivial iters — they contribute no per-mm advance,
- # so their (placeholder) stride is irrelevant.
- m_outer_iters = list(sg.shard[ssep[3] : ssep[4]])
- if not all(int(it.extent) == 1 or int(it.stride) % 8 == 0 for it in m_outer_iters):
- return None
- if (s4, s2) == (2, 1) and s8 > 0 and s8 % 8 == 0:
- return (rg, rsep, sg, ssep, False, s8, num)
- if s8 == 1 and s2 > 0 and s2 % 8 == 0 and s4 == 2 * s2:
- return (rg, rsep, sg, ssep, True, s2, num)
- return None
-
- # Try the **sorted** variant: 5D-group, sub-group R's M/2 by S's M/2
- # extents, sort the sub-groups by descending S-stride, rebuild. If
- # anything in the rebuild raises (e.g. M/2 can't be sub-grouped by S's
- # extents), silently fall back to the no-sort path below.
- r_sort = s_sort = None
- try:
- gs5 = [t_total // 32, 8, 4, m_total // 2, 2]
- rg5, rsep5 = r.group(gs5)
- sg5, ssep5 = s.group(gs5)
- r_m_iters = list(rg5.shard[rsep5[3] : rsep5[4]])
- s_m_iters = list(sg5.shard[ssep5[3] : ssep5[4]])
- s_m_extents = [int(it.extent) for it in s_m_iters]
- # Sub-group R's M/2 iters by S's M/2 iter extents. This 1-to-1's
- # the R sub-groups with the S iters so we can permute them together.
- r_m_sub = TileLayout.from_iters(r_m_iters)
- r_m_grouped, r_m_seps = r_m_sub.group(s_m_extents)
- # Sort S iters by S-stride descending; permute R sub-groups in lockstep.
- perm = sorted(range(len(s_m_iters)), key=lambda i: -int(s_m_iters[i].stride))
- if perm != list(range(len(perm))):
- r_m_permuted = r_m_grouped.permute_by_groups(list(r_m_seps), perm)
- s_m_permuted = [s_m_iters[i] for i in perm]
- r_sort = TileLayout.from_iters(
- list(rg5.shard[: rsep5[3]])
- + list(r_m_permuted.shard)
- + list(rg5.shard[rsep5[4] :]),
- offset=dict(rg5.offset),
- )
- s_sort = TileLayout.from_iters(
- list(sg5.shard[: ssep5[3]]) + list(s_m_permuted) + list(sg5.shard[ssep5[4] :]),
- offset=dict(sg5.offset),
- )
- # If perm is identity, sorted == unsorted; no need to build duplicate layouts.
- except Exception:
- r_sort = s_sort = None
-
- # Enumerate num largest-first; for each num try sorted then unsorted.
- chosen = None
- for num in (4, 2, 1):
- if m_total % (num * 2):
- continue
- if r_sort is not None:
- res = _try_num(r_sort, s_sort, num)
- if res is not None:
- chosen = res
- break
- res = _try_num(r, s, num)
- if res is not None:
- chosen = res
- break
-
- if chosen is None:
- fail("ldstmatrix layout doesn't fit any num ∈ {4,2,1}")
- r, r_seps, s, s_seps, trans, _row_stride, num = chosen
-
- # Order the num atom (seg 4) iters by R m-stride descending. The dst of
- # each ldmatrix/stmatrix must be 4 consecutive registers; when the seg-4
- # decomposition order matches R's fragment word order (largest m-stride
- # slowest), the frag words land in consecutive dst registers and ptxas
- # needs no register-shuffle MOVs. R and S are permuted in lockstep, so
- # the R↔S element pairing (and thus numerics) is unchanged.
- r_num_iters = list(r.shard[r_seps[4] : r_seps[5]])
- if len(r_num_iters) > 1:
- num_order = sorted(range(len(r_num_iters)), key=lambda i: -int(r_num_iters[i].stride))
- if num_order != list(range(len(r_num_iters))):
- r_iters = list(r.shard)
- s_iters = list(s.shard)
- r_seg = [r_iters[r_seps[4] : r_seps[5]][i] for i in num_order]
- s_seg = [s_iters[s_seps[4] : s_seps[5]][i] for i in num_order]
- r = TileLayout.from_iters(
- r_iters[: r_seps[4]] + r_seg + r_iters[r_seps[5] :],
- offset=dict(r.offset),
- )
- s = TileLayout.from_iters(
- s_iters[: s_seps[4]] + s_seg + s_iters[s_seps[5] :],
- offset=dict(s.offset),
- )
-
- # Step 10: emit one ldmatrix/stmatrix per mm, per warp.
-
- def _get_warp_idx_in_T():
- # Pick the CUDA index call while parsing the instruction implementation.
- if r_lane_axis == "laneid":
- return 0
- if r_lane_axis == "tid_in_wg":
- return T.cuda.warp_in_warpgroup()
- return T.cuda.warp_id() # "tx"
-
- def _seg4_coord(laneid_expr):
- # num=1: seg 4 trivially extent-1, pass 0. num>1: use lane//8 (tile
- # index in ldmatrix lane convention); layout.apply decomposes through
- # the seg's iter structure (single or multi-iter).
- if num > 1:
- return laneid_expr // 8
- return 0
-
- # Built out here: binding a Python string inside the traced body is not
- # something the parser can carry.
- _trans_seg = ".trans" if trans else ""
- ld_chain = f"ldmatrix.sync.aligned.m8n8.x{num}{_trans_seg}.shared.b16"
- st_chain = f"stmatrix.sync.aligned.m8n8.x{num}{_trans_seg}.shared.b16"
-
- apply_shape = [t_total // 32, 8, 4, m_total // (num * 2), num, 2]
- r_mem_axis = r.shard[r_seps[5]].axis.name
- s_mem_axis = s.shard[s_seps[5]].axis.name
- s_mem_axis_obj = s.shard[s_seps[5]].axis
- m_outer = m_total // (num * 2)
- s_zero = [0] * len(s_buf.shape)
-
- def _apply_s_layout(warp_idx, lane_idx, mm_idx):
- # The PTX lane-row contribution is not always one of the grouped S
- # coordinates (the transposed form takes its stride from another
- # group). Add it to the TileLayout offset so ComposeLayout sees the
- # complete structured address and can prove its bounded low part.
- offset = dict(s.offset)
- row_offset = (lane_idx % 8) * _row_stride
- offset[s_mem_axis_obj] = offset.get(s_mem_axis_obj, 0) + row_offset
- row_tile = TileLayout.from_iters(list(s.shard), list(s.replica), offset)
- row_layout = row_tile
- if s_swizzle is not None:
- row_layout = ComposeLayout(
- s_swizzle.per_element,
- s_swizzle.swizzle_len,
- s_swizzle.atom_len,
- row_tile,
- s_swizzle.swizzle_inner,
- )
- return row_layout.apply(
- warp_idx,
- 0,
- 0,
- mm_idx,
- _seg4_coord(lane_idx),
- 0,
- shape=apply_shape,
- )[s_mem_axis]
-
- # fmt: off
- @T.function(check_well_formed=False)
- def impl():
- r_local = r_buf.local(m_total, layout=TileLayout(S[(m_total,)]))
- laneid = T.cuda.lane_id()
- warp_idx_in_T = _get_warp_idx_in_T()
- for mm in T.unroll(m_outer):
- smem_off = _apply_s_layout(warp_idx_in_T, laneid, mm)
- smem_ptr = _ptr_off(s_buf.ptr_to(s_zero), smem_off)
- # Both instructions move num b32 registers; the fragment buffer is
- # b16, so they land through a uint32 view, two elements per word
- # (the tcgen05_ldst 16-bit pattern).
- r_words = r_local.view("uint32")
- words = [
- r_words[r.apply(0, 0, 0, mm, i, 0, shape=apply_shape)[r_mem_axis] // 2]
- for i in range(num)
- ]
- if T.constexpr(direction == "ld"):
- T.ptx[ld_chain](*words, smem_ptr)
- else:
- # stmatrix reverses ldmatrix's operand order: address first.
- T.ptx[st_chain](smem_ptr, *words)
- # fmt: on
- return impl
-
-
-def copy_schedule_ldstmatrix(op_call: TensorCall, sctx: DispatchContext) -> Function:
- return _emit(op_call, sctx)
-
-
-__all__ = ["copy_schedule_ldstmatrix"]
diff --git a/python/tvm/backend/cuda/tile/copy/utils.py b/python/tvm/backend/cuda/tile/copy/utils.py
deleted file mode 100644
index 61cccca029..0000000000
--- a/python/tvm/backend/cuda/tile/copy/utils.py
+++ /dev/null
@@ -1,74 +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 copy operator dispatches on CUDA targets."""
-
-from collections.abc import Iterable
-
-from tvm.runtime import MemoryScope
-from tvm.tirx.tile.dispatch import DispatchContext
-from tvm.tirx.tile.instruction import TensorCall
-
-from ..common import validate_copy_op
-
-_SCOPE_GLOBAL = MemoryScope("global")
-_SCOPE_LOCAL = MemoryScope("local")
-_SCOPE_SHARED = MemoryScope("shared")
-_SCOPE_SHARED_DYN = MemoryScope("shared.dyn")
-
-
-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"
- return (ok, None if ok else f"expected thread exec_scope, got {exec_scope}")
-
-
-DEFAULT_ALLOWED_PAIRS = (
- (_SCOPE_GLOBAL, _SCOPE_SHARED),
- (_SCOPE_GLOBAL, _SCOPE_SHARED_DYN),
- (_SCOPE_SHARED, _SCOPE_GLOBAL),
- (_SCOPE_SHARED_DYN, _SCOPE_GLOBAL),
- (_SCOPE_GLOBAL, _SCOPE_LOCAL),
- (_SCOPE_LOCAL, _SCOPE_GLOBAL),
- (_SCOPE_SHARED, _SCOPE_LOCAL),
- (_SCOPE_SHARED_DYN, _SCOPE_LOCAL),
- (_SCOPE_LOCAL, _SCOPE_SHARED),
- (_SCOPE_LOCAL, _SCOPE_SHARED_DYN),
-)
-
-
-def _scope_allowed(
- op_call: TensorCall,
- sctx: DispatchContext,
- allowed_pairs: Iterable[tuple[MemoryScope, MemoryScope]] = DEFAULT_ALLOWED_PAIRS,
-):
- 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()
- ok = (src_scope, dst_scope) in allowed_pairs
- if not ok:
- allowed_str = ", ".join((f"{a.name}->{b.name}" for a, b in allowed_pairs))
- return (
- False,
- f"unsupported memory scopes src={src_scope} dst={dst_scope}; allowed: {allowed_str}",
- )
- return (True, None)
-
-
-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/copy/vec_auto_reg.py b/python/tvm/backend/cuda/tile/copy/vec_auto_reg.py
deleted file mode 100644
index 254f74488b..0000000000
--- a/python/tvm/backend/cuda/tile/copy/vec_auto_reg.py
+++ /dev/null
@@ -1,633 +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.
-
-"""Non-ldmatrix copy dispatch for register ↔ memory.
-
-This file owns every copy where one side is per-thread local (``R`` =
-register). That R side carries the partition: its ``TileLayout`` ``shard``
-has thread-axis iters telling us which thread owns which logical coordinate.
-The other side (``S``) can be ``shared*`` or ``global`` — the algorithm is
-identical either way.
-
-Slice/canonicalize both sides, align via perm+group, then emit a per-thread
-vectorized copy loop. Direction-symmetric: covers R2S / S2R / R2G / G2R.
-"""
-
-import tvm_ffi
-
-import tvm
-from tvm.runtime import DataType, MemoryScope
-from tvm.script import tirx as T
-from tvm.sym import Analyzer, ConstIntBound
-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.tile.dispatch import DispatchContext
-from tvm.tirx.tile.instruction import TensorCall
-
-from ..layout_utils import recompose_swizzle, strip_swizzle_to_tile
-from ._common import _alignment_ok, copy_ptx_form, copy_ptx_ld_chain
-from .utils import _is_valid_copy, _scope_allowed
-from .vec_forced import _ld_cache_config
-
-_SCOPE_GLOBAL = MemoryScope("global")
-
-_SCOPE_LOCAL = MemoryScope("local")
-_SCOPE_SHARED = MemoryScope("shared")
-_SCOPE_SHARED_DYN = MemoryScope("shared.dyn")
-
-
-def _extract_tile(layout, region):
- """Strip swizzle off ``layout`` so it can be perm/grouped as a TileLayout."""
- return strip_swizzle_to_tile(layout, lambda: [int(end - start) for (start, end) in region])
-
-
-_REG_PAIRS = (
- (_SCOPE_LOCAL, _SCOPE_SHARED),
- (_SCOPE_LOCAL, _SCOPE_SHARED_DYN),
- (_SCOPE_SHARED, _SCOPE_LOCAL),
- (_SCOPE_SHARED_DYN, _SCOPE_LOCAL),
- (_SCOPE_LOCAL, _SCOPE_GLOBAL),
- (_SCOPE_GLOBAL, _SCOPE_LOCAL),
-)
-_SCOPE_RANK = {"thread": 0, "warp": 1, "warpgroup": 2, "cta": 3}
-_VALID_R_SUBSCOPES = {"thread", "warp", "warpgroup"}
-
-
-def _all_threads_active(sctx: DispatchContext) -> tuple[bool, str | None]:
- if sctx.scope_kind == "thread":
- return True, None
- required: dict[str, int] = {}
- if sctx.scope_kind in ("warp", "warpgroup", "cta"):
- required["laneid"] = 32
- if sctx.scope_kind == "warpgroup":
- required["wid_in_wg"] = 4
- if sctx.scope_kind == "cta":
- tx_iv = sctx.launch_params.get("threadIdx.x")
- if tx_iv is None:
- return False, "cta scope missing threadIdx.x launch_params"
- try:
- required["warpid"] = int(tx_iv[1]) // 32
- except (TypeError, ValueError):
- return False, f"non-static threadIdx.x extent: {tx_iv[1]}"
- for axis_name, expected in required.items():
- if axis_name not in sctx.intra:
- return False, f"sctx.intra missing {axis_name!r}"
- ext_raw, off_raw = sctx.intra[axis_name]
- try:
- ext, off = int(ext_raw), int(off_raw)
- except (TypeError, ValueError):
- return False, f"non-static range for {axis_name}: ({ext_raw}, {off_raw})"
- if ext != expected or off != 0:
- return False, f"{axis_name} narrowed to [{off}, {off + ext}) vs full [0, {expected})"
- return True, None
-
-
-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() == _SCOPE_LOCAL else dst
- layout = r_buf.ty.effective_layout()
- if layout is None:
- return False, "R has no layout"
- if layout.is_swizzle():
- return False, "R layout is swizzle"
- if not isinstance(layout, TileLayout):
- return False, f"R layout is {type(layout).__name__}, not TileLayout"
-
- scope_rank = _SCOPE_RANK[sctx.scope_kind]
- for it in layout.shard:
- ax = it.axis
- if not ax.is_thread():
- continue
- ax_scope = ax.get_scope()
- ax_sub = ax.get_subscope()
- if ax_scope is None or ax_sub is None:
- return False, f"R thread axis {ax.name!r} missing scope/subscope"
- if ax_sub.name not in _VALID_R_SUBSCOPES:
- return False, f"R thread axis {ax.name!r} subscope={ax_sub.name!r} (not register-level)"
- if ax_scope.name not in _SCOPE_RANK or _SCOPE_RANK[ax_scope.name] > scope_rank:
- return (
- False,
- f"R thread axis {ax.name!r} scope={ax_scope.name!r} > exec {sctx.scope_kind!r}",
- )
- r_br = op_call.src if src.scope() == _SCOPE_LOCAL else op_call.dst
- region = [(r.min, r.min + r.extent) for r in r_br.region]
- sliced = layout.slice(list(r_buf.shape), region)
- if sliced is None:
- return False, "R layout slice failed"
- analyzer = Analyzer()
- for axis, off in sliced.offset.items():
- if axis.is_thread() and not analyzer.can_prove_equal(off, 0):
- return False, f"R sliced offset on thread axis {axis.name!r} = {off}"
- return True, 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 = TensorCall.decode(op_call)
- src_br = op_call.src
- dst_br = op_call.dst
- s_br = dst_br if src_br.source.scope() == _SCOPE_LOCAL else src_br
- s_buf: Var = s_br.source
- layout = s_buf.ty.effective_layout()
- if layout is None:
- return False, "S has no layout"
- region = [(r.min, r.min + r.extent) for r in s_br.region]
- if layout.slice(list(s_buf.shape), region) is None:
- return False, "S layout slice failed"
- return True, 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"):
- 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=_REG_PAIRS),
- lambda: _r_side_layout_valid(op_call, sctx),
- lambda: _s_side_slice_ok(op_call),
- ):
- ok, msg = check()
- if not ok:
- return False, msg
- return True, None
-
-
-def _compute_perm_r(r):
- # thread axes first, then by stride descending
- def key(p):
- it = p[1]
- return (0 if it.axis.is_thread() else 1, -int(it.stride))
-
- return [i for i, _ in sorted(enumerate(r.shard), key=key)]
-
-
-def align_layouts_raw(r_sliced, s_sliced, s_region):
- """Returns (r_p, s_p, s_seps, r_perm).
-
- ``r_p`` for ``_s_thread_offset``; ``r_perm`` for ``_split_thread_loop`` pairing.
- """
- r = r_sliced.canonicalize()
- s = s_sliced.canonicalize()
- s = _extract_tile(s, s_region)
- perm = _compute_perm_r(r)
- r_shape_for_group = [int(it.extent) for it in r.shard]
- s_grp, seps = s.group(r_shape_for_group)
- s_p = s_grp.permute_by_groups(list(seps), perm)
- r_perm = r.permute_dims(perm)
- r_p = r_perm.canonicalize()
- sizes = [seps[i + 1] - seps[i] for i in range(len(seps) - 1)]
- s_seps = [0]
- for p in perm:
- s_seps.append(s_seps[-1] + sizes[p])
- return r_p, s_p, s_seps, r_perm
-
-
-def _split_thread_loop(r_perm, s_p, s_seps):
- """Drop R's thread-axis positions and return per-R-position bundles:
- (r_iters, s_groups) — same length lists; s_groups[k] is the list of S
- iters belonging to the k-th kept R position.
-
- ``r_perm`` is the permuted-but-uncanonicalized R layout: its iter list lines
- up one-to-one with ``s_seps`` (which is built from the same ``perm``), so a
- multi-group memory axis keeps each of its groups paired with the matching S
- group instead of collapsing into one (see ``align_layouts_raw``)."""
- r_iters = []
- s_groups = []
- for k, r_it in enumerate(r_perm.shard):
- if r_it.axis.is_thread():
- continue
- r_iters.append(r_it)
- s_groups.append(list(s_p.shard[s_seps[k] : s_seps[k + 1]]))
- return r_iters, s_groups
-
-
-def _build_atoms(r_iters, s_groups):
- """One atom per (R position, intra-group S iter): (extent, s_stride, r_mul).
- r_mul = R_stride_at_position * (product of S group extents to the right
- of this intra-position index) — i.e. how much R address advances per unit
- of this iter's loop input."""
- atoms = []
- for r_it, s_group in zip(r_iters, s_groups, strict=True):
- rs = int(r_it.stride)
- extents = [int(it.extent) for it in s_group]
- for j, s_it in enumerate(s_group):
- inner_prod = 1
- for e in extents[j + 1 :]:
- inner_prod *= e
- atoms.append((int(s_it.extent), int(s_it.stride), rs * inner_prod))
- return atoms
-
-
-def _atoms_contiguous_tail_extent(atoms) -> int:
- """Like _contiguous_tail_extent but on atoms (uses s_stride for chaining)."""
- if not atoms or atoms[-1][1] != 1:
- return 0
- acc = atoms[-1][0]
- for k in range(len(atoms) - 2, -1, -1):
- if atoms[k][1] == acc:
- acc *= atoms[k][0]
- else:
- break
- return acc
-
-
-def _split_atoms_for_vec(atoms, vec_len):
- """Returns outer atoms (the inner vec_len-element tail is consumed by one
- vec ld/st and dropped). Splits the boundary atom if needed."""
- outer = list(atoms)
- acc = 1
- while outer:
- ext, ss, rm = outer[-1]
- new_acc = acc * ext
- if new_acc == vec_len:
- outer.pop()
- return outer
- if new_acc > vec_len:
- inner_factor = vec_len // acc
- outer[-1] = (ext // inner_factor, ss * inner_factor, rm * inner_factor)
- return outer
- acc = new_acc
- outer.pop()
- raise ValueError(f"tail too short for vec_len {vec_len}")
-
-
-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() == _SCOPE_LOCAL:
- r_br, s_br = src_br, dst_br
- else:
- r_br, s_br = dst_br, src_br
- r_buf = r_br.source
- s_buf = s_br.source
- r_region = [(r.min, r.min + r.extent) for r in r_br.region]
- s_region = [(r.min, r.min + r.extent) for r in s_br.region]
- with sctx.target:
- r_sliced = r_buf.ty.effective_layout().slice(list(r_buf.shape), r_region)
- s_sliced = s_buf.ty.effective_layout().slice(list(s_buf.shape), s_region)
- return align_layouts_raw(r_sliced, s_sliced, s_region)
-
-
-def _make_thread_placeholders(r_p) -> dict[str, _TirVar]:
- placeholders: dict[str, _TirVar] = {}
- for it in r_p.shard:
- name = it.axis.name
- if it.axis.is_thread() and name not in placeholders:
- placeholders[name] = _TirVar(name, "int32")
- return placeholders
-
-
-def _s_thread_contributions(r_p, s_p, placeholders: dict[str, _TirVar]):
- """Return the per-S-iter coordinates contributing to a thread's address."""
- coord = [
- placeholders[it.axis.name] if it.axis.is_thread() else _IntImm("int32", 0)
- for it in r_p.shard
- ]
- input_shape = [int(it.extent) for it in r_p.shard]
- per_iter = s_p.apply_to_shape(coord, input_shape)
- return list(zip(per_iter, (it.stride for it in s_p.shard), strict=True))
-
-
-_VEC_BITS_CANDIDATES = (128, 64, 32, 16, 8)
-
-
-def _vec_len_candidates(elem_bits: int) -> list[int]:
- """Widest-first element counts to try; always ends with scalar (1)."""
- out: list[int] = []
- for vb in _VEC_BITS_CANDIDATES:
- if vb < elem_bits or vb % elem_bits != 0:
- continue
- n = vb // elem_bits
- if n not in out:
- out.append(n)
- if 1 not in out:
- out.append(1)
- return out
-
-
-def _choose_vec_len(elem_bits: int, atoms, r_p, s_p, max_vec_len: int | None = None) -> int:
- """Widest candidate that:
- 1. divides the atom contiguous-tail extent (so vec_len consecutive
- R-side regs map to vec_len contiguous S-side elements), AND
- 2. keeps every per-thread / per-round address-offset term a
- multiple of vec_len, so the resulting vec ld/st pointer is
- naturally aligned to vec_bits/8 bytes.
-
- Only **mem-axis** strides contribute to physical address. Thread-axis
- iter strides live in partition-coord space (which thread owns which
- logical position), not in the buffer's storage space — they're
- redistributed through ``apply_to_shape`` into the mem iters and don't
- appear directly in the per-thread address. So neither r-side nor
- s-side thread-axis strides belong in the alignment check.
-
- The contig-tail atoms (whose extents the vec ld/st consumes) have
- stride 1 by definition; they live entirely inside the vec and
- contribute nothing to the per-round address delta. Only the
- **post-vec-split** outer atom strides matter for the per-round delta.
- """
- t = _atoms_contiguous_tail_extent(atoms)
- # Region-base offsets are real address contributions; thread-iter
- # strides are partition-virtual, not storage-physical — exclude them.
- shared_terms = list(s_p.offset.values()) + list(r_p.offset.values())
- for n in _vec_len_candidates(elem_bits):
- if max_vec_len is not None and n > max_vec_len:
- continue
- if n == 1:
- return n
- if t % n != 0 or t < n:
- continue
- # Post-vec-split outer atoms: these are the strides that contribute
- # to per-round address deltas after the vec consumes the inner tail.
- outer = _split_atoms_for_vec(atoms, n)
- outer_atom_terms = [a[1] for a in outer] + [a[2] for a in outer]
- if not _alignment_ok(n, outer_atom_terms + shared_terms):
- continue
- return n
- return 1
-
-
-def _axis_decl(axis_name: str, sctx: DispatchContext):
- """Return the ordinary CUDA index expression for one thread axis."""
- if axis_name == "tx":
- return sctx.launch_params["threadIdx.x"][0]
- if axis_name == "laneid":
- return T.cuda.lane_id()
- if axis_name == "wid_in_wg":
- return T.cuda.warp_in_warpgroup()
- if axis_name == "tid_in_wg":
- return T.cuda.thread_in_warpgroup()
- if axis_name == "warpid":
- return T.cuda.warp_id()
- if axis_name == "wgid":
- return T.cuda.warpgroup_id()
- raise ValueError(f"unsupported thread axis {axis_name}")
-
-
-def _thread_axis_extent(axis_name: str, sctx: DispatchContext) -> int | None:
- """Return the full hardware range of a scope-id axis when it is static."""
- if axis_name in sctx.intra:
- active = sctx.intra[axis_name]
- try:
- extent, offset = int(active[0]), int(active[1])
- except (TypeError, ValueError):
- return None
- if offset != 0:
- return None
- return extent
-
- fixed = {"laneid": 32, "wid_in_wg": 4, "tid_in_wg": 128}
- if axis_name in fixed:
- return fixed[axis_name]
-
- tx = sctx.launch_params.get("threadIdx.x")
- if tx is None:
- return None
- try:
- tx_extent = int(tx[1])
- except (TypeError, ValueError):
- return None
- divisor = {"tx": 1, "warpid": 32, "wgid": 128}.get(axis_name)
- if divisor is None or tx_extent % divisor != 0:
- return None
- return tx_extent // divisor
-
-
-def _build_s_apply_layout(original_layout, r_p, s_p, outer, placeholders, sctx):
- """Build a structured tile whose ``m`` is exactly old ``s_off + ds``.
-
- Each bounded per-thread coordinate becomes its own Iter with a proven
- hardware-derived extent. Outer atoms retain their exact extent/stride and
- are addressed by one flattened loop coordinate through shape-aware apply.
- Slice offsets and any thread term without a finite non-negative bound stay
- in the tile's ``m`` offset, which makes ComposeLayout conservatively use
- its full-swizzle fallback when needed.
- """
- analyzer = Analyzer()
- for name, placeholder in placeholders.items():
- extent = _thread_axis_extent(name, sctx)
- if extent is not None:
- analyzer.bind(placeholder, tvm.ir.Range.from_min_extent(0, extent))
-
- shard = []
- thread_coords = []
- apply_shape = []
- offset = _IntImm("int32", 0)
- for coord, stride in _s_thread_contributions(r_p, s_p, placeholders):
- coord = analyzer.simplify(coord)
- if analyzer.can_prove_equal(coord, 0):
- continue
- bound = analyzer.const_int_bound(coord)
- if (
- bound.min_value >= 0
- and bound.max_value != ConstIntBound.POS_INF
- and bound.max_value < (1 << 31) - 1
- ):
- extent = int(bound.max_value) + 1
- extent_expr = _IntImm("int32", extent)
- shard.append(Iter(extent_expr, stride, "m"))
- thread_coords.append(coord)
- apply_shape.append(extent_expr)
- else:
- offset = analyzer.simplify(offset + coord * stride)
-
- for _axis, value in s_p.offset.items():
- offset = analyzer.simplify(offset + value)
-
- total_outer = 1
- for extent, stride, _r_mul in outer:
- shard.append(Iter(extent, stride, "m"))
- total_outer *= extent
- apply_shape.append(_IntImm("int32", total_outer))
-
- tile = TileLayout.from_iters(shard, [], {"m": offset})
- return recompose_swizzle(original_layout, tile), thread_coords, apply_shape
-
-
-def _axis_substitution(placeholders: dict[str, _TirVar], sctx: DispatchContext):
- """Declare the real scope ids once and map placeholder Vars to them."""
- return {placeholders[name]: _axis_decl(name, sctx) for name in placeholders}
-
-
-def _apply_s_layout(layout, thread_coords, flat_outer, shape, placeholders, sctx):
- coord = [*thread_coords, flat_outer]
- mapped = layout.apply(*coord, shape=shape)["m"]
- var_map = _axis_substitution(placeholders, sctx)
- return tvm_ffi.structural_map(
- mapped,
- (_TirVar, lambda var: var_map.get(var, var)),
- order="post",
- )
-
-
-def _flat_coords(outer_atoms, flat_idx: int) -> list[int]:
- coords = []
- rem = flat_idx
- for a in reversed(outer_atoms):
- coords.append(rem % a[0])
- rem //= a[0]
- coords.reverse()
- return coords
-
-
-_POINTER_OFFSET_SRC = (
- "\ntemplate <typename T>\n"
- "__forceinline__ __device__ T* tvm_builtin_pointer_offset(T* ptr, int offset) {\n"
- " return ptr + offset;\n"
- "}\n"
-)
-
-
-def _ptr_off(base_ptr, off):
- return T.cuda.func_call(
- "tvm_builtin_pointer_offset",
- base_ptr,
- off,
- _POINTER_OFFSET_SRC,
- ty="handle",
- )
-
-
-def _outer_const_offsets(outer_atoms, flat_idx: int) -> tuple[int, int]:
- """Returns (s_offset_const, r_offset_const) for one outer-loop flat index."""
- coords = _flat_coords(outer_atoms, flat_idx)
- ds = sum(c * a[1] for c, a in zip(coords, outer_atoms))
- dr = sum(c * a[2] for c, a in zip(coords, outer_atoms))
- return ds, dr
-
-
-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() == _SCOPE_LOCAL:
- r_buf, s_buf, r_is_src = src, dst, True
- else:
- r_buf, s_buf, r_is_src = dst, src, False
-
- with sctx.target:
- r_p, s_p, s_seps, r_perm = _align_layouts(op_call, sctx)
- r_iters, s_groups = _split_thread_loop(r_perm, s_p, s_seps)
- atoms = _build_atoms(r_iters, s_groups)
- elem_bits = DataType(src.dtype).bits
- max_vec_len = None
- if (
- isinstance(s_buf.ty.effective_layout(), ComposeLayout)
- and int(s_buf.ty.effective_layout().swizzle_len) > 0
- ):
- max_vec_len = 1 << int(s_buf.ty.effective_layout().per_element)
- vec_len = _choose_vec_len(elem_bits, atoms, r_p, s_p, max_vec_len)
- vec_bits = vec_len * elem_bits
- outer = _split_atoms_for_vec(atoms, vec_len)
-
- placeholders = _make_thread_placeholders(r_p)
- s_apply_layout, thread_coords, s_apply_shape = _build_s_apply_layout(
- s_buf.ty.effective_layout(), r_p, s_p, outer, placeholders, sctx
- )
- # A plain TileLayout is additive, so keep its loop-invariant thread base
- # outside the serial copy loop. A real swizzle must still apply to the
- # complete base + outer coordinate inside each iteration.
- has_swizzle = isinstance(s_apply_layout, ComposeLayout) and int(s_apply_layout.swizzle_len) > 0
-
- # R-side base offset from slicing (e.g. ``R[i*8:i*8+8]`` ⇒ ``i*8``), from
- # ``r_p.offset``; sum across all axes on R's local stride-1 storage.
- r_off_base = _IntImm("int32", 0)
- for _ax, val in r_p.offset.items():
- r_off_base = r_off_base + val
-
- s_is_shared = s_buf.scope() in (_SCOPE_SHARED, _SCOPE_SHARED_DYN)
- num_bytes = vec_bits // 8
- tail, lanes, reg_dtype = copy_ptx_form(num_bytes)
- space = "shared" if s_is_shared else "global"
-
- # Honor the load cache hint (cache="nc") only on global -> register loads;
- # ``nc`` is a global qualifier, so stores and shared copies stay plain.
- _cache, _cache_hints = _ld_cache_config(op_call)
- _use_nc = _cache == "nc" and space == "global" and not r_is_src
- _l1_evict = _cache_hints.get("l1_evict", "")
- _l2_evict = _cache_hints.get("l2_evict", "")
- _prefetch = _cache_hints.get("prefetch_size", "")
- # Element offsets index the element buffer; the instruction moves whole
- # containers, so they are scaled to container index. The vector alignment
- # rule guarantees the division is exact.
- container_bits = num_bytes * 8 // lanes
- elems_per_word = container_bits // elem_bits if container_bits >= elem_bits else 1
- words_per_elem = elem_bits // container_bits if elem_bits > container_bits else 1
- # Chains are built here: a Python string bound inside the traced body is
- # not something the parser can carry. `nc` folds into the chain, so the
- # three-way branch at the call site collapses to load-or-store.
- st_chain = f"st.{space}.{tail}"
- ld_chain = copy_ptx_ld_chain(
- space,
- tail,
- nc=_use_nc,
- l1_evict=_l1_evict,
- l2_evict=_l2_evict,
- prefetch_size=_prefetch,
- )
-
- total_outer = 1
- for a in outer:
- total_outer *= a[0]
-
- # fmt: off
- s_zero_indices = [0] * len(s_buf.shape)
-
- @T.function(check_well_formed=False)
- def impl():
- if T.constexpr(not has_swizzle):
- s_base = _apply_s_layout(
- s_apply_layout, thread_coords, 0, s_apply_shape, placeholders, sctx
- )
- # ``dr`` and ``r_off_base`` below are physical storage offsets, so use
- # the raw-span local view rather than storage-iterator coordinates.
- r_local = r_buf.local()
- # The registers the instruction moves are containers, not elements.
- r_words = r_local.view(reg_dtype)
- # Keep a serial TIR loop and let ptxas unroll; explicit ``T.unroll``
- # replicates per-iter scratch arrays and pressures registers.
- for f in range(total_outer):
- ds, dr = _outer_const_offsets(outer, f)
- if T.constexpr(has_swizzle):
- s_off = _apply_s_layout(
- s_apply_layout, thread_coords, f, s_apply_shape, placeholders, sctx
- )
- s_ptr = _ptr_off(s_buf.ptr_to(s_zero_indices), s_off)
- else:
- s_ptr = _ptr_off(s_buf.ptr_to(s_zero_indices), s_base + ds)
- # The registers move through the reg_dtype word view, not a raw
- # element pointer: the container index is the element offset scaled
- # between element and word granularity.
- r_w = (r_off_base + dr) * words_per_elem // elems_per_word
- if T.constexpr(r_is_src):
- T.ptx[st_chain](s_ptr, *[r_words[r_w + i] for i in range(lanes)])
- else:
- T.ptx[ld_chain](*[r_words[r_w + i] for i in range(lanes)], s_ptr)
- # fmt: on
- import os
-
- if os.environ.get("R2S_DUMP"):
- print("=== emitted impl ===")
- print(impl.script())
- return impl
diff --git a/python/tvm/backend/cuda/tile/copy/vec_forced.py b/python/tvm/backend/cuda/tile/copy/vec_forced.py
deleted file mode 100644
index 815eba999b..0000000000
--- a/python/tvm/backend/cuda/tile/copy/vec_forced.py
+++ /dev/null
@@ -1,205 +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.
-
-"""Explicit fixed-width vector copy dispatches.
-
-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.
-- ``l1_evict`` / ``l2_evict`` / ``prefetch_size``: L1/L2 cache hint kwargs
- forwarded verbatim to the emitted load (same values ``T.ptx.ld`` /
- ``T.ptx.ld.global_.nc`` accept, e.g. ``"L1::no_allocate"``,
- ``"L2::evict_first"``, ``"L2::256B"``). Requires a global src.
-
-Example::
-
- Tx.copy(dst_local[0:8], src_global[base : base + 8], dispatch="vec_256b",
- cache="nc", l1_evict="L1::no_allocate", prefetch_size="L2::256B")
-"""
-
-from tvm.ir import TensorRegion
-from tvm.runtime import DataType, MemoryScope
-from tvm.script import tirx as T
-from tvm.sym.analyzer import Analyzer
-from tvm.tirx import Var
-from tvm.tirx.tile.dispatch import DispatchContext
-from tvm.tirx.tile.instruction import TensorCall
-
-from ._common import copy_ptx_form, copy_ptx_ld_chain
-from .utils import _scope_allowed
-
-_SCOPE_GLOBAL = MemoryScope("global")
-_SCOPE_LOCAL = MemoryScope("local")
-_SCOPE_SHARED = MemoryScope("shared")
-_SCOPE_SHARED_DYN = MemoryScope("shared.dyn")
-
-
-def _region_start(buffer_region: TensorRegion):
- return [r.min for r in buffer_region.region]
-
-
-def _region_elements(buffer_region: TensorRegion):
- product = 1
- for r in buffer_region.region:
- product *= r.extent
- return product
-
-
-def _can_prove_equal(lhs, rhs) -> bool:
- if isinstance(lhs, int) and isinstance(rhs, int):
- return lhs == rhs
- return Analyzer().can_prove_equal(lhs, rhs)
-
-
-def _ptx_space(scope: MemoryScope) -> str:
- if scope in (_SCOPE_SHARED, _SCOPE_SHARED_DYN):
- return "shared"
- return scope.name
-
-
-_LD_CACHE_HINT_KEYS = ("l1_evict", "l2_evict", "prefetch_size")
-
-
-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.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.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: TensorCall,
- sctx: DispatchContext,
- *,
- variant: str,
- num_bytes: int,
-):
- if sctx.scope_kind != "thread":
- return False, f"expected thread exec_scope, got {sctx.scope_kind}"
-
- scope_ok, scope_reason = _scope_allowed(op_call, sctx)
- if not scope_ok:
- return False, scope_reason
-
- op_call = TensorCall.decode(op_call)
- src: Var = op_call.src.source
- dst: Var = op_call.dst.source
- if src.dtype != dst.dtype:
- return False, f"dtype mismatch: src={src.dtype}, dst={dst.dtype}"
-
- elem_bits = DataType(src.dtype).bits
- width_bits = num_bytes * 8
- if width_bits % elem_bits != 0:
- return False, f"{variant} is not an integral number of {src.dtype} elements"
- expected_elements = width_bits // elem_bits
-
- src_elements = _region_elements(op_call.src)
- dst_elements = _region_elements(op_call.dst)
- if not _can_prove_equal(src_elements, expected_elements):
- return False, f"src region does not contain exactly {expected_elements} elements"
- if not _can_prove_equal(dst_elements, expected_elements):
- return False, f"dst region does not contain exactly {expected_elements} elements"
-
- cache, hints = _ld_cache_config(op_call)
- if cache not in (None, "nc"):
- return False, f"unsupported cache config {cache!r}; expected None or 'nc'"
- if (cache is not None or hints) and src.scope() != _SCOPE_GLOBAL:
- return False, (
- "cache/L1/L2 hint config applies to global loads; "
- f"requires src scope 'global', got {src.scope()!r}"
- )
- return True, None
-
-
-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()
- dst_scope = dst.scope()
- src_is_local = src_scope == _SCOPE_LOCAL
- dst_is_local = dst_scope == _SCOPE_LOCAL
- src_space = _ptx_space(src_scope)
- dst_space = _ptx_space(dst_scope)
- src_ptr = src.ptr_to(_region_start(op_call.src))
- dst_ptr = dst.ptr_to(_region_start(op_call.dst))
- elem_bits = DataType(src.dtype).bits
- tail, lanes, reg_dtype = copy_ptx_form(num_bytes)
- container_bits = num_bytes * 8 // lanes
-
- def _words(buf, region):
- """The container lvalues a local region's registers land in.
-
- The instruction moves whole containers, so the element region start is
- scaled to a container index; the vector alignment rule the caller
- already enforced makes the division exact.
- """
- flat = buf.view(-1).view(reg_dtype)
- start, stride = 0, 1
- for idx, extent in zip(reversed(_region_start(region)), reversed(list(buf.shape))):
- start = start + idx * stride
- stride = stride * extent
- base = (
- start * (elem_bits // container_bits)
- if elem_bits > container_bits
- else start // (container_bits // elem_bits)
- )
- return [flat[base + i] for i in range(lanes)]
-
- # Chains are built here: a Python string bound inside the traced body is
- # not something the parser can carry. `nc` folds into the chain, so the
- # use_nc branches collapse.
- st_chain = f"st.{dst_space}.{tail}"
- cache, hints = _ld_cache_config(op_call)
- use_nc = cache == "nc"
- l1_evict = hints.get("l1_evict", "")
- l2_evict = hints.get("l2_evict", "")
- prefetch_size = hints.get("prefetch_size", "")
- ld_chain = copy_ptx_ld_chain(
- src_space,
- tail,
- nc=use_nc,
- l1_evict=l1_evict,
- l2_evict=l2_evict,
- prefetch_size=prefetch_size,
- )
-
- # fmt: off
- @T.function(check_well_formed=False)
- def impl():
- if T.constexpr(src_is_local):
- T.ptx[st_chain](dst_ptr, *_words(src, op_call.src))
- elif T.constexpr(dst_is_local):
- T.ptx[ld_chain](*_words(dst, op_call.dst), src_ptr)
- else:
- tmp = T.alloc_local((lanes,), reg_dtype, layout="default")
- T.ptx[ld_chain](*[tmp[i] for i in range(lanes)], src_ptr)
- T.ptx[st_chain](dst_ptr, *[tmp[i] for i in range(lanes)])
- # fmt: on
- return impl
diff --git a/python/tvm/backend/cuda/tile/copy_async/__init__.py b/python/tvm/backend/cuda/tile/copy_async/__init__.py
deleted file mode 100644
index 37cf03d8c0..0000000000
--- a/python/tvm/backend/cuda/tile/copy_async/__init__.py
+++ /dev/null
@@ -1,24 +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 asynchronous tensor transfer instruction lowerers."""
-
-from .dsmem import *
-from .ldgsts import *
-from .tcgen05_cp import *
-from .tcgen05_ldst import *
-from .tma import *
diff --git a/python/tvm/backend/cuda/tile/copy_async/dsmem.py b/python/tvm/backend/cuda/tile/copy_async/dsmem.py
deleted file mode 100644
index 2ba6839ad6..0000000000
--- a/python/tvm/backend/cuda/tile/copy_async/dsmem.py
+++ /dev/null
@@ -1,197 +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_async dispatch variant: dsmem (shared::cta -> shared::cluster)."""
-
-import functools
-import operator
-
-import tvm
-from tvm.runtime import MemoryScope
-from tvm.script import tirx as T
-from tvm.tirx import Function, Var
-from tvm.tirx.tile import DispatchContext, fail
-from tvm.tirx.tile.instruction import TensorCall
-
-from .utils import find_contiguous_region, to_tile_layout
-
-_SCOPE_SHARED = MemoryScope("shared")
-_SCOPE_SHARED_DYN = MemoryScope("shared.dyn")
-
-
-def _is_shared_to_shared(op_call: TensorCall) -> bool:
- """Check if both src and dst are in shared memory."""
- op_call = TensorCall.decode(op_call)
- src_scope = op_call.src.source.scope()
- dst_scope = op_call.dst.source.scope()
- return (src_scope in (_SCOPE_SHARED, _SCOPE_SHARED_DYN)) and (
- dst_scope in (_SCOPE_SHARED, _SCOPE_SHARED_DYN)
- )
-
-
-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
- to copy data from the executing CTA's shared memory to a remote CTA's shared
- memory within the same cluster.
-
- The copy region is decomposed into contiguous byte chunks based on layout
- analysis of both src and dst buffers. Non-contiguous dimensions are iterated
- over, emitting one cp.async.bulk instruction per contiguous chunk.
- """
- op_call = TensorCall.decode(op_call)
-
- # Extract config
- 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.options.get("mbar", None)
- if mbar is None:
- fail("mbar not set in config")
-
- # Extract buffer regions
- dst_buffer_region = op_call.dst
- src_buffer_region = op_call.src
- src_buf: Var = src_buffer_region.source
- dst_buf: Var = dst_buffer_region.source
-
- src_st = [r.min for r in src_buffer_region.region]
- src_ext = [r.extent for r in src_buffer_region.region]
- dst_st = [r.min for r in dst_buffer_region.region]
- dst_ext = [r.extent for r in dst_buffer_region.region]
-
- dtype_bytes = tvm.DataType(src_buf.dtype).bits // 8
-
- # Get tile layouts for both buffers
- src_tile_layout = to_tile_layout(src_buf.ty.effective_layout(), src_buf.shape)
- dst_tile_layout = to_tile_layout(dst_buf.ty.effective_layout(), dst_buf.shape)
-
- # Slice layouts to copy region
- src_region_tuples = [(src_st[i], src_st[i] + src_ext[i]) for i in range(len(src_st))]
- sliced_src = src_tile_layout.slice([s for s in src_buf.shape], src_region_tuples)
- if sliced_src is None:
- fail("Cannot slice src layout for DSMEM copy")
-
- dst_region_tuples = [(dst_st[i], dst_st[i] + dst_ext[i]) for i in range(len(dst_st))]
- sliced_dst = dst_tile_layout.slice([s for s in dst_buf.shape], dst_region_tuples)
- if sliced_dst is None:
- fail("Cannot slice dst layout for DSMEM copy")
-
- # Group src layout by region extents, then group dst by src's shard extents
- # This creates 1:1 shard correspondence between the two layouts
- grouped_src, src_seps = sliced_src.canonicalize().group(src_ext)
- src_shard_extents = [s.extent for s in grouped_src.shard]
- grouped_dst, dst_seps = sliced_dst.canonicalize().group(src_shard_extents)
-
- # Find contiguous regions in both layouts
- src_contig_indices, _ = find_contiguous_region(grouped_src)
- dst_contig_indices, _ = find_contiguous_region(grouped_dst)
-
- # Intersect: walk from innermost outward, include only matching shard indices
- shared_contig_indices = []
- for s_idx, d_idx in zip(src_contig_indices, dst_contig_indices):
- if s_idx != d_idx:
- break
- shared_contig_indices.append(s_idx)
-
- # Compute chunk size
- if shared_contig_indices:
- chunk_elements = functools.reduce(
- operator.mul, [grouped_src.shard[i].extent for i in shared_contig_indices], 1
- )
- else:
- chunk_elements = 1
-
- chunk_bytes = chunk_elements * dtype_bytes
- if chunk_bytes < 16 or chunk_bytes % 16 != 0:
- fail(
- f"Layouts not compatible for bulk DSMEM copy: "
- f"chunk_bytes={chunk_bytes} (need >= 16 and multiple of 16)"
- )
-
- # Build iteration space over non-contiguous (outer) shards
- shared_contig_set = set(shared_contig_indices)
- outer_shard_indices = [i for i in range(len(grouped_src.shard)) if i not in shared_contig_set]
- outer_extents = [grouped_src.shard[i].extent for i in outer_shard_indices]
- outer_src_strides = [grouped_src.shard[i].stride for i in outer_shard_indices]
- outer_dst_strides = [grouped_dst.shard[i].stride for i in outer_shard_indices]
-
- # Helper to compute element offsets from loop variables (called via T.meta_var)
- def compute_offsets(loop_vars):
- if isinstance(loop_vars, Var):
- loop_vars = (loop_vars,)
- src_off = 0
- dst_off = 0
- for j, v in enumerate(loop_vars):
- src_off = src_off + v * outer_src_strides[j]
- dst_off = dst_off + v * outer_dst_strides[j]
- return src_off, dst_off
-
- src_tile = to_tile_layout(src_buf.ty.effective_layout(), src_buf.shape)
- dst_tile = to_tile_layout(dst_buf.ty.effective_layout(), dst_buf.shape)
-
- # fmt: off
- @T.function(check_well_formed=False)
- def impl():
- # Map mbar to remote CTA (complete_tx targets the destination's mbar).
- # mapa writes its result into a declared register, so the mapped
- # addresses live in a small local scratch buffer.
- mapped = T.alloc_local([2], "uint64", layout="default")
- T.ptx.mapa.u64(mapped[0], mbar, T.uint32(remote_cta_id))
- remote_mbar = mapped[0]
-
- if T.constexpr(not outer_extents):
- # Single contiguous chunk — no iteration needed
- src_ptr = src_buf.ptr_to(src_st)
- T.ptx.mapa.u64(mapped[1], dst_buf.ptr_to(dst_st), T.uint32(remote_cta_id))
- cluster_dst = mapped[1]
- T.ptx["cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes"](
- T.cast(cluster_dst, "uint32"), src_ptr, T.cast(chunk_bytes, "uint32"),
- T.cast(remote_mbar, "uint32"),
- )
- else:
- for loop_vars in T.grid(*outer_extents):
- src_elem_offset, dst_elem_offset = T.meta_var(compute_offsets(loop_vars))
-
- src_buf_w = T.decl_tensor(
- src_buf.data,
- ty_args=[T.Tensor(
- src_buf.shape, src_buf.dtype,
- elem_offset=src_buf.elem_offset + src_elem_offset,
- scope=src_buf.scope(), layout=src_tile,
- )],
- )
- dst_buf_w = T.decl_tensor(
- dst_buf.data,
- ty_args=[T.Tensor(
- dst_buf.shape, dst_buf.dtype,
- elem_offset=dst_buf.elem_offset + dst_elem_offset,
- scope=dst_buf.scope(), layout=dst_tile,
- )],
- )
-
- src_ptr = src_buf_w.ptr_to(src_st)
- T.ptx.mapa.u64(mapped[1], dst_buf_w.ptr_to(dst_st), T.uint32(remote_cta_id))
- cluster_dst = mapped[1]
- T.ptx["cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes"](
- T.cast(cluster_dst, "uint32"), src_ptr, T.cast(chunk_bytes, "uint32"),
- T.cast(remote_mbar, "uint32"),
- )
- # fmt: on
-
- return impl
diff --git a/python/tvm/backend/cuda/tile/copy_async/ldgsts.py b/python/tvm/backend/cuda/tile/copy_async/ldgsts.py
deleted file mode 100644
index e440019759..0000000000
--- a/python/tvm/backend/cuda/tile/copy_async/ldgsts.py
+++ /dev/null
@@ -1,269 +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_async`` dispatch for ``global → shared`` via ``cp.async``
-(SASS: ``LDGSTS``).
-
-Shares the partition / layout-alignment algorithm with
-``cuda/copy/vec_auto_gmem_smem.py`` (sync ``T.copy`` global ↔ shared); differs at
-emit time only:
-
-* direction: ``cp.async`` is global → shared only (hardware restriction).
-* cp_size: PTX ``cp.async`` only accepts 4 / 8 / 16 bytes, so the vec-width
- candidate set is restricted to ``{32, 64, 128}`` bits.
-* emit: one ``T.ptx["cp.async..."]`` per slice instead of the synchronous
- ``T.cuda.copy_{vec_bits}b(dst, src)``.
-* config: ``prefetch_size`` lands in the instruction spelling, ``predicate``
- rides ``pred=`` (@p), and ``fill_mode="zero"`` selects the src-size arity
- (zero-filling the tail). This covers predicated zero-fill loads such as
- FlashMLA RoPE KV staging while preserving the same layout partitioner.
-* config: ``direct=True`` is an explicit thread-scope fast path for callers
- that already selected a physically contiguous 4/8/16-byte slice. It emits
- one ``cp.async`` from the region starts and bypasses the synthesized
- partitioner, matching hand-written per-thread copies.
-
-Note: ``cp.async`` does **not** sync at emit time — caller is responsible
-for ``commit_group`` / ``wait_group`` / ``cta_sync`` plumbing around the
-async pipeline.
-"""
-
-from tvm.runtime import DataType, MemoryScope
-from tvm.script import tirx as T
-from tvm.tirx import Function, Var
-from tvm.tirx.expr import IntImm as _IntImm
-from tvm.tirx.tile.dispatch import DispatchContext
-from tvm.tirx.tile.instruction import TensorCall
-
-from ..copy._common import (
- _TID_AXIS_FOR_SCOPE,
- _thread_cnt,
- align_layouts_gs,
-)
-from ..copy.utils import _is_valid_copy, _scope_allowed
-from ..copy.vec_auto_reg import _all_threads_active, _axis_decl, _ptr_off
-from ..layout_utils import recompose_swizzle
-
-_SCOPE_SHARED = MemoryScope("shared")
-_SCOPE_SHARED_DYN = MemoryScope("shared.dyn")
-
-_SCOPE_GLOBAL = MemoryScope("global")
-
-# cp.async is unidirectional: global → shared.
-_LDGSTS_PAIRS = (
- (_SCOPE_GLOBAL, _SCOPE_SHARED),
- (_SCOPE_GLOBAL, _SCOPE_SHARED_DYN),
-)
-# cp.async cp_size ∈ {4, 8, 16} bytes ⇒ vec_bits ∈ {32, 64, 128}.
-_LDGSTS_VEC_BITS = (128, 64, 32)
-
-
-def _emit_cp_async(dst_ptr, src_ptr, cp_size, prefetch_size, predicate_expr, fill_mode):
- """Emit one ptx cp.async for the ldgsts configuration.
-
- fill_mode="zero" is the src-size arity (src-size = pred ? cp-size : 0,
- zero-filling the tail); a bare predicate rides pred= (@p); otherwise the
- plain form. The instruction spelling carries the prefetch qualifier.
- """
- from tvm.tirx.op import if_then_else
-
- cop = "cg" if int(cp_size) == 16 else "ca"
- pref = "" if int(prefetch_size) == -1 else f".L2::{int(prefetch_size)}B"
- chain = f"cp.async.{cop}.shared.global{pref}"
- is_default = (isinstance(predicate_expr, int) and predicate_expr == -1) or (
- hasattr(predicate_expr, "value") and int(predicate_expr.value) == -1
- )
- has_pred = not is_default
- if fill_mode == "zero":
- src_size = T.cast(if_then_else(predicate_expr != 0, int(cp_size), 0), "uint32")
- T.evaluate(T.ptx[chain](dst_ptr, src_ptr, int(cp_size), src_size))
- elif has_pred:
- T.evaluate(T.ptx[chain](dst_ptr, src_ptr, int(cp_size), pred=predicate_expr))
- else:
- T.evaluate(T.ptx[chain](dst_ptr, src_ptr, int(cp_size)))
-
-
-def _config_bool(value) -> bool:
- if isinstance(value, bool):
- return value
- if isinstance(value, _IntImm):
- return bool(int(value.value))
- return bool(value)
-
-
-def _divides_thread_cnt_ldgsts(
- 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 = 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)"
- g_br = op_call.src if op_call.src.source.scope() == _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_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"):
- 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=_LDGSTS_PAIRS),
- lambda: _divides_thread_cnt_ldgsts(op_call, sctx),
- ):
- ok, msg = check()
- if not ok:
- return False, msg
- return True, None
-
-
-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.
- g_buf, g_br = src, op_call.src
- s_buf, s_br = dst, op_call.dst
-
- elem_bits = DataType(src.dtype).bits
- 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.options.get("direct", False)):
- if sctx.scope_kind != "thread":
- raise ValueError("ldgsts direct=True is only valid in thread scope")
- n_elements = 1
- for r in s_br.region:
- try:
- n_elements *= int(r.extent)
- except (TypeError, ValueError) as err:
- raise ValueError(
- f"ldgsts direct=True requires constant extent, got {r.extent}"
- ) from err
- cp_size = n_elements * elem_bits // 8
- if n_elements * elem_bits % 8 != 0 or cp_size not in (4, 8, 16):
- raise ValueError(
- f"ldgsts direct=True requires a 4/8/16-byte region, got {n_elements} "
- f"elements of {elem_bits} bits"
- )
- s_start = [r.min for r in s_br.region]
- g_start = [r.min for r in g_br.region]
-
- @T.function(check_well_formed=False)
- def impl():
- _emit_cp_async(
- s_buf.ptr_to(s_start),
- g_buf.ptr_to(g_start),
- cp_size,
- prefetch_size,
- predicate_expr,
- fill_mode,
- )
-
- return impl
-
- 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]
-
- thread_cnt = _thread_cnt(sctx)
-
- with sctx.target:
- g_p, s_p, vec_len = align_layouts_gs(
- g_buf.ty.effective_layout(),
- g_buf.shape,
- g_region,
- s_buf.ty.effective_layout(),
- s_buf.shape,
- s_region,
- elem_bits,
- thread_cnt,
- vec_bits_candidates=_LDGSTS_VEC_BITS,
- )
- s_apply_layout = recompose_swizzle(s_buf.ty.effective_layout(), s_p)
-
- vec_bits = vec_len * elem_bits
- cp_size = vec_bits // 8 # cp.async cp_size is in bytes
- if cp_size not in (4, 8, 16):
- # align_layouts_gs already restricted candidates to _LDGSTS_VEC_BITS,
- # so reaching here means no candidate worked at all.
- from tvm.tirx.tile.dispatch import fail
-
- fail(f"ldgsts: cannot find a cp.async-compatible vec_len for elem_bits={elem_bits}")
-
- # Mirror vec_auto_gmem_smem.py: build 3D `(f, tid, 0)` against
- # `[total_outer, thread_cnt, vec_len]` and let `s_p.apply(coord, shape)`
- # flatten + resplit into whatever multi-iter T / outer-iter structure
- # `align_layouts_gs` picked. Emit is oblivious to how many shard iters
- # cover T.
- 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()
- for f in T.unroll(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)
- _emit_cp_async(s_ptr, g_ptr, cp_size, prefetch_size, predicate_expr, fill_mode)
- # cp.async is caller-synced — no cta_sync here (commit_group /
- # wait_group / cta_sync are the caller's responsibility).
- # fmt: on
- return impl
-
-
-def copy_schedule_ldgsts(op_call: TensorCall, sctx: DispatchContext) -> Function:
- return _emit_ldgsts(op_call, sctx)
diff --git a/python/tvm/backend/cuda/tile/copy_async/tcgen05_cp.py b/python/tvm/backend/cuda/tile/copy_async/tcgen05_cp.py
deleted file mode 100644
index 8f6d1c438b..0000000000
--- a/python/tvm/backend/cuda/tile/copy_async/tcgen05_cp.py
+++ /dev/null
@@ -1,823 +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.
-
-"""smem->tmem dispatch via ``tcgen05.cp`` — generic planner for all PTX shapes.
-
-``tcgen05.cp`` is inherently async; this dispatch emits the cp loop only and
-leaves completion signaling (``tcgen05.commit`` against a barrier) to the
-caller. Callers who want sync semantics should issue ``tcgen05.commit``
-themselves after the copy.
-
-Routing
--------
-- ``shape=`` config: forces that shape (see ``_CP_SHAPE_MULTICASTS`` below;
- atom rows/bits are parsed from the shape name itself).
-- No ``shape`` config: the shape is inferred from the buffer layouts — each
- candidate (widest atom first) is planned until one validates. Bare warpx4
- copies resolve to ``32x128b.warpx4`` exactly as the historical dispatch.
-
-Descriptor fields are always derived from the layouts.
-
-Shape table (PTX ISA 8.8 §9.7.16.9.2)
--------------------------------------
-Each ``tcgen05.cp`` shape is ``rows x bits-per-row``; some shapes REQUIRE a
-multicast qualifier. Multicast replicates the copied data across warp lane
-slabs of the destination TMEM:
-
-=========== ========================== ====================== ==============
-shape multicast t lane pattern t replica
-=========== ========================== ====================== ==============
-128x256b (none) (128, 1@TLane) —
-4x256b (none) (4, 32@TLane) —
-128x128b (none) (128, 1@TLane) —
-64x128b warpx2::02_13 (required) (64, 1@TLane) (2, 64@TLane)
-64x128b warpx2::01_23 (required) (2,32):(64,1)@TLane (2, 32@TLane)
-32x128b warpx4 (required) (32, 1@TLane) (4, 32@TLane)
-=========== ========================== ====================== ==============
-
-The ``warpx2`` row→lane mappings follow PTX ISA 8.8 p675 ("each warp in the
-warp pair receives half of the data", pairs 02_13 = {0,2},{1,3} and 01_23 =
-{0,1},{2,3}): the two warps of a pair mirror each other and the pairs split
-the 64 rows in listed order. For 02_13 this is exactly the Layout E / B
-datapath organization (Figures 213/207, §9.7.16.10.5: M 0-31 in warp-ranks 0
-and 2, M 32-63 in warp-ranks 1 and 3). ``4x256b`` writes one row per
-warp-quadrant datapath (lanes 0/32/64/96). All row→lane mappings and replica
-structures are verified bit-exactly on B200 hardware by the round-trip tests
-in ``test_tcgen05_cp.py``.
-
-Algorithm
----------
-Given ``Tx.copy_async(t_region, s_region)`` where t is in tmem (declared with
-the replica pattern of the requested (shape, multicast)), and s is in shared
-memory:
-
-A. Slice + canonicalize both layouts at the given regions.
-B. Verify ``t.replica`` matches the shape table entry (empty for
- non-multicast shapes).
-C. Compute permutation that puts TLane first, then TCol stride-descending;
- apply to t.permute_dims and to s via group + permute_by_groups.
-D. Canonicalize again.
-E. Isolate broadcast: split-by-stride-zero on both t and s; their split
- sequences must match (same distinct prefix prods + broadcast extents).
- Drop stride-0 iters → ``t_iso`` and ``s_iso``.
-F. Split each side's iters into three segments by grouping the flattened
- element space as ``atom_rows x n_mid x elem_per_atom``
- (``elem_per_atom = atom_bits / dtype_bits``). ``*_lane`` covers exactly one
- instruction's rows (product == atom_rows); ``*_col`` = one instruction's
- own columns; ``*_middle`` = everything between — the outer columns (which
- instruction) and, when the region has more rows than one atom (e.g. 4x256b
- tiling each warp's first 16 rows: the M=64 Layout-F scatter), the extra row
- iters, stepped through the taddr lane half-word. ``t_*`` = tmem side,
- ``s_*`` = smem side.
- Validate:
- - t_lane (row → TMEM lane) matches the shape table lane pattern (above)
- - t_col: one row's elements land in contiguous TMEM columns
- - s_col: one row = contiguous 16B units in smem; for 256b atoms the two
- units may sit at a non-contiguous stride when unswizzled (→ LDO field)
- - s_lane: rows sit in smem as (atom_rows/8, 8) groups with strides
- (SDO_stride, atom_K_stride); 4x256b is a single 4-row group (SDO=0)
- - atom_K_byte ∈ {16, 32, 64, 128} → swizzle_mode 0..3, which must match
- s_buf.layout's swizzle (if any)
-G. Alignment checks:
- - t_iso TCol offset ≡ 0 (mod 128-bit)
- - t_iso TLane offset folds into the taddr lane half-word
- - s_iso m offset ≡ 0 (mod 16B for sw=0; mod atom_size for sw>0)
- - middle iter strides 16B-aligned
-H. middle 1-1 correspondence (simple-mode): t_middle and s_middle have same
- iter count and matching extents per position.
-I. Emit:
- - SmemDescriptor encoded once at SMEM base (hoisted via post_tensor_def_stmt).
- - Loop over middle iters; each cp uses ``desc.add_16B_offset(init + loop)``
- and writes to ``tmem_addr + t_addr_off + Σ i_j * t_step_j``.
-"""
-
-import functools
-import operator
-from functools import cache
-
-import tvm
-from tvm.backend.cuda.axis import TCol, TLane
-from tvm.backend.cuda.op import _tmem_address
-from tvm.ir import Bind, Call, DictAttrs, Evaluate, SeqStmt, StringImm
-from tvm.runtime import DataType, MemoryScope
-from tvm.script import tirx as T
-from tvm.sym import Analyzer
-from tvm.tirx import Function, Var
-from tvm.tirx.layout import ComposeLayout, TileLayout
-from tvm.tirx.layout import m as m_axis
-from tvm.tirx.tile import DispatchContext
-from tvm.tirx.tile.instruction import TensorCall
-
-_SCOPE_SHARED = MemoryScope("shared")
-_SCOPE_SHARED_DYN = MemoryScope("shared.dyn")
-
-
-@cache
-def _scope_tmem():
- return MemoryScope("tmem")
-
-
-# Allowed .multicast qualifiers per shape (PTX ISA 8.8 §9.7.16.9.2): "" = none,
-# 64x128b needs an explicit warpx2 pick.
-_CP_SHAPE_MULTICASTS = {
- "128x256b": ("",),
- "4x256b": ("",),
- "128x128b": ("",),
- "64x128b": ("warpx2::02_13", "warpx2::01_23"),
- "32x128b": ("warpx4",),
-}
-
-
-def _shape_dims(shape):
- """(atom_rows, atom_bits) parsed from the PTX shape name "<rows>x<bits>b"."""
- rows, bits = shape.split("x")
- return int(rows), int(bits[:-1])
-
-
-# Inference order for bare copies: widest atom first (a 256b source also fits
-# 128b at 2x instructions); other candidates are mutually exclusive.
-_CP_SHAPE_CANDIDATES = (
- ("128x256b", ""),
- ("4x256b", ""),
- ("128x128b", ""),
- ("64x128b", "warpx2::02_13"),
- ("64x128b", "warpx2::01_23"),
- ("32x128b", "warpx4"),
-)
-
-
-def _cp_lane_replica_pattern(shape: str, multicast: str):
- """Expected (t_lane pattern, t.replica pattern) for a (shape, multicast).
-
- Both are lists of ``(extent, stride)`` on the TLane axis. The replica
- strides are the warp-slab offsets that receive multicast copies; the lane
- pattern is the row → lane mapping of one copy.
-
- Multicast semantics (PTX ISA 8.8 §9.7.16.9.2, p675): for warpx2 the data
- is "multicasted into warp pairs and each warp in the warp pair receives
- half of the data" — the two warps OF a pair MIRROR each other (hold the
- same half), and the pairs split the 64 rows in listed order:
-
- - ``warpx2::02_13`` (pairs {0,2}, {1,3}): rows 0-31 to warps 0/2, rows
- 32-63 to warps 1/3. One copy occupies lanes 0-63 in row order (warp0 +
- warp1), replicated at lane offset +64 (warp2 + warp3) → lane
- (64, 1@TLane), replica (2, 64@TLane). This is exactly the Layout E /
- Layout B datapath organization (PTX ISA 8.8 Figures 213 and 207:
- M 0-31 in warp-ranks 0 and 2, M 32-63 in warp-ranks 1 and 3).
- - ``warpx2::01_23`` (pairs {0,1}, {2,3}): rows 0-31 to warps 0/1, rows
- 32-63 to warps 2/3 → one copy at lanes 0-31 (rows 0-31) + lanes 64-95
- (rows 32-63), replicated at +32 → lane (2, 32):(64, 1)@TLane, replica
- (2, 32@TLane).
-
- ``4x256b`` writes one row per warp-quadrant datapath: rows 0..3 land at
- lanes 0, 32, 64, 96 → lane (4, 32@TLane), no replica.
-
- All mappings are verified bit-exactly on B200 hardware by the round-trip
- tests in ``test_tcgen05_cp.py``.
- """
- rows, _ = _shape_dims(shape)
- if multicast == "":
- if rows == 4:
- return [(4, 32)], []
- return [(rows, 1)], []
- if multicast == "warpx4":
- return [(32, 1)], [(4, 32)]
- if multicast == "warpx2::02_13":
- return [(64, 1)], [(2, 64)]
- if multicast == "warpx2::01_23":
- return [(2, 64), (32, 1)], [(2, 32)]
- raise ValueError(f"unknown tcgen05.cp multicast {multicast!r}")
-
-
-def _resolve_cp_shape(op_call: TensorCall):
- """Resolve (shape, multicast) from an explicit ``shape=`` config."""
- 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)
- if allowed is None:
- raise ValueError(
- f"unknown tcgen05.cp shape {shape!r}; expected one of {sorted(_CP_SHAPE_MULTICASTS)}"
- )
- if multicast is None:
- if len(allowed) == 1:
- multicast = allowed[0]
- else:
- raise ValueError(
- f"tcgen05.cp shape {shape!r} requires an explicit multicast config; "
- f"choose one of {list(allowed)}"
- )
- else:
- multicast = str(multicast)
- if multicast not in allowed:
- raise ValueError(
- f"illegal multicast {multicast!r} for tcgen05.cp shape {shape!r}; "
- f"allowed: {list(allowed)}"
- )
- return shape, multicast
-
-
-# -----------------------------------------------------------------------------
-# Helpers
-# -----------------------------------------------------------------------------
-def _compute_perm(t):
- def key(p):
- it = p[1]
- return (0 if it.axis == TLane else 1, -int(it.stride))
-
- return [i for i, _ in sorted(enumerate(t.shard), key=key)]
-
-
-def _split_by_zero(lay):
- """Split lay.shard into segments at stride==0 positions.
- Returns (split_seq, kept_iters_with_nonzero_stride)."""
- new_seq = []
- keep = []
- cur = 1
- for it in lay.shard:
- e, st = int(it.extent), int(it.stride)
- if st == 0:
- if cur > 1:
- new_seq.append(cur)
- new_seq.append(e)
- cur = 1
- else:
- cur *= e
- keep.append(it)
- if cur > 1:
- new_seq.append(cur)
- return new_seq, keep
-
-
-def _align_middles(t_middle, s_middle):
- """Sub-group both middles by union-of-boundaries so they become 1-1.
-
- Both inputs must be post-canonicalize iter lists with equal extent products.
- The shape is the consecutive ratios of sorted(B_t U B_s) where B_x is the
- set of cumulative extent boundaries of x_middle. Each segment then contains
- at most one iter per side (whole or sub-divided), so trivially single iter.
-
- Returns (new_t_middle, new_s_middle) with len() == len() == k segments,
- each segment a single Iter on each side.
- """
-
- def cum_bounds(iters):
- b, p = [], 1
- for it in iters:
- p *= int(it.extent)
- b.append(p)
- return b
-
- t_bounds = cum_bounds(t_middle)
- s_bounds = cum_bounds(s_middle)
- if not t_bounds and not s_bounds:
- return t_middle, s_middle
- N = t_bounds[-1] if t_bounds else s_bounds[-1]
- if (s_bounds and s_bounds[-1] != N) or (t_bounds and t_bounds[-1] != N):
- raise ValueError(f"middle extent mismatch: t={N} s={s_bounds[-1] if s_bounds else 0}")
-
- cuts = sorted(set(t_bounds) | set(s_bounds))
- shape, prev = [], 1
- for c in cuts:
- if c % prev != 0:
- raise ValueError(
- f"middle align failed: cut {c} not divisible by prev cut {prev} "
- f"(t_bounds={t_bounds}, s_bounds={s_bounds})"
- )
- shape.append(c // prev)
- prev = c
-
- def subgroup(iters):
- if len(iters) == 1 and shape == [int(iters[0].extent)]:
- return iters
- lay, _seps = TileLayout.from_iters(iters, [], {}).group(shape)
- seps = list(_seps)
- out = []
- for i in range(len(shape)):
- seg = list(lay.shard[seps[i] : seps[i + 1]])
- seg_canon = list(TileLayout.from_iters(seg, [], {}).canonicalize().shard)
- if len(seg_canon) != 1:
- raise ValueError(
- f"middle sub-group seg[{i}] not single iter after canon: {seg_canon}"
- )
- out.append(seg_canon[0])
- return out
-
- return subgroup(t_middle), subgroup(s_middle)
-
-
-# -----------------------------------------------------------------------------
-# Plan (state object)
-# -----------------------------------------------------------------------------
-def _build_plan(op_call: TensorCall):
- """Run A..H and return a dispatch plan.
-
- Plan fields:
- - s_buf, t_buf
- - dtype, dtype_bits
- - shape, multicast (PTX qualifier strings)
- - elem_per_atom, elem_per_32b
- - SmemSwizzleMode (int)
- - LDO_field, SDO_field, atom_K_byte
- - middle_iters: list of (extent, s_step_16B, t_step_32bcol)
- - init_off_16B (Expr)
- - 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 = 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.options.get("multicast")
- if isinstance(multicast_cfg, tvm.ir.StringImm):
- multicast_cfg = multicast_cfg.value
- errors = []
- for shape, multicast in _CP_SHAPE_CANDIDATES:
- if multicast_cfg is not None and str(multicast_cfg) != multicast:
- continue
- try:
- return _plan_for_shape(op_call, shape, multicast)
- except ValueError as err:
- errors.append(f"{shape}{('.' + multicast) if multicast else ''}: {err}")
- raise ValueError(
- "no tcgen05.cp shape fits the given layouts; candidates rejected:\n " + "\n ".join(errors)
- )
-
-
-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
- t_buf: Var = dst_region.source
- dtype = s_buf.dtype
... 37557 lines suppressed ...