wehub-resource-sync
184 lines
7.5 KiB
ReStructuredText
184 lines
7.5 KiB
ReStructuredText
.. 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_cp
|
||
=======================
|
||
|
||
The ``tcgen05_cp`` variant lowers a ``copy_async`` from **shared memory to tensor
|
||
memory** (Blackwell ``tmem``). One elected thread issues
|
||
``tcgen05.cp.32x128b.warpx4``: a shared **matrix descriptor** names the source tile,
|
||
and the ``warpx4`` multicast routes 32 lanes × 128 bits into the tensor-memory lanes
|
||
owned by all four warps. The dispatch issues only the copy; the caller signals
|
||
completion with ``tcgen05.commit``. Source:
|
||
``python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py``.
|
||
|
||
What it accepts
|
||
---------------
|
||
|
||
Two predicates — a valid shared→tmem copy and a single-thread scope:
|
||
|
||
.. code-block:: python
|
||
|
||
# register_dispatch(..., variant="smem->tmem", priority=10, when=[
|
||
predicate("validate_smem_tmem_copy", _is_valid_smem_tmem_copy),
|
||
predicate("exec_scope", _single_thread_exec), # exec_scope == "thread"
|
||
# ])
|
||
|
||
def _is_valid_smem_tmem_copy(op, sctx):
|
||
if not (src.scope().startswith("shared") and dst.scope() == "tmem"): ...
|
||
if not (src.layout and dst.layout): ...
|
||
if dst.allocated_addr is None: ...
|
||
rep = dst.layout.replica # the warpx4 router
|
||
if not (len(rep) == 1 and int(rep[0].extent) == 4
|
||
and int(rep[0].stride) == 32 and "TLane" in str(rep[0].axis)):
|
||
return False, f"requires R[4:32@TLane] on tmem, got {list(rep)}"
|
||
return True, None
|
||
|
||
.. list-table::
|
||
:header-rows: 1
|
||
:widths: 22 78
|
||
|
||
* - Property
|
||
- Requirement
|
||
* - target / priority
|
||
- ``cuda`` (Blackwell, sm_100+); priority ``10``
|
||
* - scope
|
||
- **single thread** issues the copy
|
||
* - memory pair
|
||
- source ``shared*`` → destination ``tmem`` (with ``allocated_addr`` set by a
|
||
prior ``tcgen05.alloc``)
|
||
* - tmem layout
|
||
- the replica **must be exactly** ``R[4:32@TLane]`` — the warpx4 router that
|
||
fans the copy across all four warps' tensor-memory lanes
|
||
* - dtype
|
||
- sets ``elem_per_128b = 128 / dtype_bits`` (uint8 → 16) and the descriptor
|
||
swizzle mode
|
||
|
||
Demonstration program
|
||
----------------------
|
||
|
||
A warpgroup allocates 16 tmem columns, fills a ``32×16`` ``uint8`` shared tile, and
|
||
copies it into tmem with ``tcgen05_cp`` (from ``test_smem_tmem.py``; the readback /
|
||
dealloc tail is elided):
|
||
|
||
.. code-block:: python
|
||
|
||
from tvm.tirx.layout import R, S, TCol, TileLayout, TLane
|
||
|
||
A_smem = T.alloc_buffer([32, 16], "uint8", scope="shared",
|
||
layout=TileLayout(S[(32, 16) : (16, 1)]), align=1024)
|
||
tmem_addr = T.alloc_shared([1], "uint32")
|
||
cp_mbar = T.alloc_shared([1], "uint64")
|
||
if warp_id == 0:
|
||
T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=16, cta_group=1)
|
||
# ... mbarrier.init, fence, cta_sync, fill A_smem from global ...
|
||
tmem = T.decl_buffer([32, 16], "uint8", scope="tmem", allocated_addr=tmem_addr[0],
|
||
layout=TileLayout(S[(32, 16) : (1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane]))
|
||
if tid_in_wg == 0:
|
||
Tx.copy_async(tmem[0:32, 0:16], A_smem[0:32, 0:16], cta_group=1) # smem -> tmem
|
||
T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=1) # caller signals
|
||
T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0)
|
||
# ... readback via tcgen05.ld, then tcgen05.dealloc ...
|
||
|
||
(The ``copy_async`` is auto-dispatched to the ``smem->tmem`` variant — the source is
|
||
shared and the destination is the ``R[4:32@TLane]`` tmem buffer.)
|
||
|
||
Algorithm
|
||
---------
|
||
|
||
**1. Verify the warpx4 router and re-order.** After slicing both layouts to the
|
||
region, the dispatch confirms the tmem replica is ``R[4:32@TLane]``, permutes to
|
||
TLane-first / TCol-stride-descending, isolates the broadcast, and groups the
|
||
remaining iters into ``(32, middle, elem_per_128b)`` — the 32×128-bit atom plus a
|
||
list of *middle* tiles to loop over.
|
||
|
||
**2. Encode the matrix descriptor once.** A 64-bit shared descriptor (leading-dim
|
||
offset ``ldo``, stride-dim offset ``sdo``, swizzle mode) is encoded right after the
|
||
shared buffer is allocated, cached per ``(smem_buf, ldo, sdo, swizzle)``:
|
||
|
||
.. code-block:: python
|
||
|
||
desc_buf = decl_buffer((1,), "uint64", scope="local")
|
||
T.ptx.tcgen05.encode_matrix_descriptor(desc_buf.data, s_buf.ptr_to([0, 0]), ldo, sdo, swizzle)
|
||
|
||
**3. Issue the copy** — one ``tcgen05.cp`` for a single atom, or an unrolled loop
|
||
that bumps the tmem column offset and the descriptor's 16-byte shared offset per
|
||
middle tile:
|
||
|
||
.. code-block:: python
|
||
|
||
if total == 1:
|
||
T.ptx.tcgen05.cp(t_addr[0] + t_col0,
|
||
smem_desc_add_16B_offset(desc_buf[0], init_off_16B),
|
||
shape="32x128b", cta_group=cta_group, multicast="warpx4")
|
||
else:
|
||
for flat in T.unroll(total):
|
||
t_off, s_off = T.meta_var(compute_offsets(flat))
|
||
T.ptx.tcgen05.cp(t_addr[0] + t_col0 + t_off,
|
||
smem_desc_add_16B_offset(desc_buf[0], init_off_16B + s_off),
|
||
shape="32x128b", cta_group=cta_group, multicast="warpx4")
|
||
|
||
The dispatch emits **no** ``tcgen05.commit`` / ``wait`` — the caller commits against
|
||
an mbarrier (as in the demo).
|
||
|
||
Generated TIRx IR
|
||
-----------------
|
||
|
||
The ``32×16`` uint8 tile is a single atom (ldo=16, sdo=8, swizzle=0):
|
||
|
||
.. code-block:: python
|
||
|
||
T.ptx.tcgen05.encode_matrix_descriptor(cp_desc.data, T.address_of(A_smem[0]), 16, 8, 0)
|
||
T.ptx.tcgen05.cp(tmem_addr[0],
|
||
smem_desc_add_16B_offset(cp_desc[0], 0),
|
||
shape="32x128b", cta_group=1, multicast="warpx4")
|
||
|
||
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_smem_tmem.py``.)
|
||
|
||
How inputs change the algorithm
|
||
-------------------------------
|
||
|
||
.. list-table::
|
||
:header-rows: 1
|
||
:widths: 28 72
|
||
|
||
* - input
|
||
- effect
|
||
* - dtype
|
||
- sets ``elem_per_128b = 128 / dtype_bits`` (uint8 → 16, uint32 → 4) and the
|
||
descriptor swizzle mode (atom K-byte ∈ {16, 32, 64, 128} → swizzle 0/1/2/3)
|
||
* - number of tiles
|
||
- ``total`` middle tiles: ``1`` → a single ``tcgen05.cp``; ``> 1`` → an
|
||
unrolled loop, each step bumping the tmem column and the descriptor's 16-B
|
||
shared offset
|
||
* - shared swizzle layout
|
||
- changes the encoded ``swizzle`` mode (must match the shared buffer's swizzle)
|
||
* - tmem layout (D vs F) / cta_group
|
||
- the permutation order sets per-tile column steps; ``cta_group`` selects the
|
||
multicast routing (``cta_group::1`` vs ``::2``)
|
||
* - atom shape
|
||
- fixed at ``32x128b`` ``warpx4`` — a different atom would need a new variant
|