.. 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 → tcgen05_ldst ========================= The ``tcgen05_ldst`` variant lowers a ``copy_async`` between **tensor memory and registers** (Blackwell ``tcgen05.ld`` / ``tcgen05.st``). It is warpgroup-collective: the four warps cooperatively move a tensor-memory tile to/from their per-thread registers. One registration handles both directions — ``tmem → local`` lowers to ``tcgen05.ld``, ``local → tmem`` to ``tcgen05.st`` — and the dispatch picks the widest instruction shape the register layout matches. As with the other async variants, completion (``tcgen05.wait.ld`` / ``wait.st``) is the caller's. Source: ``python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py``. What it accepts --------------- A single registration (``variant="tmem<->local"``); direction is inferred at lowering: .. code-block:: python @register_dispatch("copy_async", "cuda", variant="tmem<->local", priority=10, when=[ predicate("validate_copy_op", _is_valid_copy), predicate("exec_scope", exec_scope_ok, expected_scopes=["warpgroup"]), predicate("storage_scope", _scope_allowed, allowed_pairs=[("tmem", "local"), ("local", "tmem")]), ]) # direction inferred in copy_tmem_local_impl: # src tmem + dst local -> "tmem2local" (ld); else "local2tmem" (st) .. list-table:: :header-rows: 1 :widths: 22 78 * - Property - Requirement * - target / priority - ``cuda`` (Blackwell, sm_100+); priority ``10`` * - 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``) for the fast path; otherwise the ``.32x32b`` fallback * - tmem datapath - classified ``D`` (M=128 identity) or ``F`` (M=64 scattered) — sets how fragment rows map to lanes 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, TCol, TileLayout, TLane from tvm.tirx.layout import tid_in_wg as axis_tid_in_wg local_view = TileLayout(S[(128, WIDTH) : (1 @ axis_tid_in_wg, 1)]) @T.prim_func def copy_async_test(A_ptr: T.handle, B_ptr: T.handle): A = T.match_buffer(A_ptr, (128, WIDTH), "float16"); B = T.match_buffer(B_ptr, (128, WIDTH), "float16") T.device_entry() warp_id = T.warp_id([4]); wg_id = T.warpgroup_id([1]); tid = T.thread_id([128]) tmem_addr = T.alloc_shared([1], "uint32") if wg_id == 0: if warp_id == 0: T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=32, cta_group=1) T.tvm_storage_sync("shared") tmem = T.decl_buffer((128, WIDTH), "float16", scope="tmem", allocated_addr=tmem_addr[0], layout=TileLayout(S[(128, WIDTH) : (1 @ TLane, 1 @ TCol)])) A_reg = T.alloc_local((WIDTH,), "float16"); B_reg = T.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.wg.copy_async(tmem[:, :], A_local[:, :]); T.ptx.tcgen05.wait.st() # store (local -> tmem) T.cuda.cta_sync() Tx.wg.copy_async(B_local[:, :], tmem[:, :]); T.ptx.tcgen05.wait.ld() # 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 dispatch 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 it falls back to ``.32x32b`` and probes ``num ∈ {1, 2, 4, 8, …}`` against the column width. **3. Issue per datapath slab.** For an M=128 ``.16x*b`` copy the fragment spans two 16-row slabs, so the warps issue the atom twice (``row = 0`` and ``row = 16``); the ``.32x32b`` path covers M=128 in a single issue (``row = 0``): .. code-block:: python op = T.ptx.tcgen05.ld if load else T.ptx.tcgen05.st for slab in range(n_slabs): # 1 for .32x32b / M=64; 2 for .16x*b M=128 op(tmem_buf.allocated_addr[0], *[local_32b[reg_base + i] for i in range(regs_eff)], shape=shape, num=num_eff, row=slab * 16, col=col_off_32b) The dispatch emits **no** wait — the caller issues ``tcgen05.wait.ld()`` / ``wait.st()`` (as in the demo). 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 T.ptx.tcgen05.st(tmem_addr[0], 0, 0, "32x32b", 4, False, local_32b[0], local_32b[1], local_32b[2], local_32b[3]) # local -> tmem T.ptx.tcgen05.ld(tmem_addr[0], 0, 0, "32x32b", 4, False, local_32b_1[0], local_32b_1[1], local_32b_1[2], local_32b_1[3]) # tmem -> local 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 → the ``.32x32b`` fallback (this demo) * - 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 - ``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