wehub-resource-sync
201 lines
8.6 KiB
ReStructuredText
201 lines
8.6 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 → tma
|
||
================
|
||
|
||
The ``tma`` variant lowers ``copy_async`` between **global and shared** to the
|
||
hardware **Tensor Memory Accelerator**: a single elected thread issues a
|
||
descriptor-driven bulk copy (``cp.async.bulk.tensor``), and the hardware walks the
|
||
multi-dimensional tile described by a ``cuTensorMap``. The descriptor is built once
|
||
on the host (``cuTensorMapEncodeTiled``); the device only *issues* the copy — the
|
||
hardware signals the caller's mbarrier when the transfer completes (the dispatch
|
||
itself emits no completion op). Source:
|
||
``python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py``.
|
||
|
||
What it accepts
|
||
---------------
|
||
|
||
The dispatch registers two predicates — a valid copy and a **single-thread** scope:
|
||
|
||
.. code-block:: python
|
||
|
||
# register_dispatch(..., priority=10, when=[
|
||
predicate("validate_copy_op", lambda op, sctx: (validate_copy_op(op, sctx), "not a valid copy op")),
|
||
predicate("single_thread", lambda op, sctx: (single_thread(op, sctx), "expected single thread")),
|
||
# ])
|
||
|
||
def single_thread(op_call, sctx):
|
||
return sctx.is_thread # exactly one elected thread issues the TMA
|
||
|
||
.. list-table::
|
||
:header-rows: 1
|
||
:widths: 22 78
|
||
|
||
* - Property
|
||
- Requirement
|
||
* - target / priority
|
||
- ``cuda``; priority ``10`` (the bulk path for ``copy_async`` global ↔ shared)
|
||
* - scope
|
||
- **single thread** (``sctx.is_thread``) — TMA is issued by one thread, not a
|
||
partitioned warp
|
||
* - direction
|
||
- ``global → shared`` (g2s) or ``shared → global`` (s2g), inferred from the
|
||
buffer scopes at lowering
|
||
* - dtype / shape
|
||
- ``validate_copy_op``: both sides have layouts, equal dtype, equal non-unit
|
||
extents
|
||
* - layout
|
||
- must form a legal descriptor: rank ≤ 5, innermost stride 1, innermost box
|
||
fits the shared swizzle atom (else the plan search shrinks / declines)
|
||
|
||
Demonstration program
|
||
----------------------
|
||
|
||
One thread bulk-copies an ``8×256`` ``float16`` tile global → shared (with a
|
||
128-byte swizzled shared layout), signals an mbarrier, waits, then reads it back
|
||
(mirrors ``test_tma.py``'s G2S smoke test):
|
||
|
||
.. code-block:: python
|
||
|
||
from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout
|
||
|
||
g_shape = s_shape = (8, 256); dtype = "float16"
|
||
shared_layout = mma_shared_layout(dtype, 3, (8, 256)) # 128-B swizzle
|
||
smem_bytes = 8 * 256 * 2
|
||
|
||
@T.prim_func
|
||
def copy_async(A_ptr: T.handle, B_ptr: T.handle):
|
||
A = T.match_buffer(A_ptr, g_shape, dtype, layout=TileLayout(S[8, 256]))
|
||
B = T.match_buffer(B_ptr, g_shape, dtype, layout=TileLayout(S[8, 256]))
|
||
T.device_entry(); T.cta_id([1]); tid = T.thread_id([8])
|
||
dyn = T.alloc_buffer([smem_bytes + 8], "uint8", scope="shared.dyn") # arena
|
||
A_smem = T.decl_buffer(s_shape, dtype, dyn.data, elem_offset=0, layout=shared_layout)
|
||
mbarrier = T.decl_buffer([1], "uint64", dyn.data, elem_offset=smem_bytes // 8)
|
||
phase: T.int32 = 0
|
||
if tid == 0:
|
||
T.ptx.mbarrier.init(mbarrier.ptr_to([0]), 1)
|
||
T.ptx.fence.proxy_async("shared::cta"); T.cuda.cta_sync()
|
||
if tid == 0:
|
||
Tx.copy_async(A_smem[0:8, 0:256], A[0:8, 0:256], dispatch="tma", mbar=mbarrier.ptr_to([0]))
|
||
T.ptx.mbarrier.arrive.expect_tx(mbarrier.ptr_to([0]), smem_bytes)
|
||
T.ptx.mbarrier.try_wait(mbarrier.ptr_to([0]), phase)
|
||
T.ptx.fence.proxy_async("shared::cta"); T.cuda.cta_sync()
|
||
Tx.cta.copy(B[0:8, 0:256], A_smem[0:8, 0:256])
|
||
|
||
Algorithm
|
||
---------
|
||
|
||
**1. Infer direction from scopes.** ``global → shared`` is g2s, ``shared → global``
|
||
is s2g (anything else is an error):
|
||
|
||
.. code-block:: python
|
||
|
||
if src.scope() == "global" and dst.scope().startswith("shared"):
|
||
direction, s_buf, g_buf = "g2s", dst, src
|
||
elif src.scope().startswith("shared") and dst.scope() == "global":
|
||
direction, s_buf, g_buf = "s2g", src, dst
|
||
|
||
**2. Plan the descriptor (L1 → L2 → L3).** The dispatch canonicalizes both
|
||
layouts (L1), then for each global iter finds the maximal contiguous stride-1 shard
|
||
chain and cuts the axis into descriptor **box** segments (L2), then stacks those
|
||
into a ``cuTensorMap`` and validates the hardware constraints — rank ≤ 5, innermost
|
||
stride 1, innermost box fits the shared swizzle atom — shrinking the chain prefix
|
||
and retrying if a constraint fails (L3). Adjacent fully-boxed contiguous dims are
|
||
merged, and an over-256 box may trigger element-type promotion.
|
||
|
||
**3. Emit the host descriptor once,** keyed by a cache so a repeated copy reuses it:
|
||
|
||
.. code-block:: python
|
||
|
||
T.call_packed("runtime.cuTensorMapEncodeTiled", tensormap, dtype_str, rank,
|
||
tensor_ptr, *reversed(shape), *reversed(strides[:-1]),
|
||
*reversed(box_dim), *element_strides, 0, swizzle_mode, 2, oob_fill)
|
||
|
||
**4. Emit the device issue loop** — an unrolled loop over the issue axes, one
|
||
``cp.async.bulk.tensor`` per step, direction-specific:
|
||
|
||
.. code-block:: python
|
||
|
||
if direction == "g2s":
|
||
T.ptx.cp_async.bulk.tensor.g2c(plan.rank, s_buf.ptr_to(s_st), mbar,
|
||
T.address_of(tensor_map), cta_mask, cta_group,
|
||
cache_hint, *tma_coords)
|
||
else:
|
||
T.ptx.cp_async.bulk.tensor.s2g(plan.rank, s_buf.ptr_to(s_st),
|
||
T.address_of(tensor_map), cache_hint, *tma_coords)
|
||
|
||
Like all ``copy_async`` variants the dispatch emits no completion — the caller's
|
||
mbarrier ``arrive.expect_tx`` / ``try_wait`` (g2s) close the loop.
|
||
|
||
Generated TIRx IR
|
||
-----------------
|
||
|
||
The ``8×256`` swizzled tile produces a **rank-3** descriptor and a single issue:
|
||
|
||
.. code-block:: python
|
||
|
||
# host (once): encode the tensor map (rank 3, reversed shape/box/strides, swizzle 3)
|
||
T.call_packed("runtime.cuTensorMapEncodeTiled", A_ptr_tensormap, "float16", 3,
|
||
A.data, 64, 8, 4, 512, 128, 64, 8, 4, 1, 1, 1, 0, 3, 2, 0)
|
||
# device:
|
||
for loop_vars in T.unroll(1):
|
||
T.ptx.cp_async.bulk.tensor.g2c(3, T.address_of(s_buf_w_offset[0]),
|
||
T.address_of(mbarrier[0]),
|
||
T.address_of(A_ptr_tensormap), 0, 1, ..., 0, 0, 0)
|
||
|
||
Generated CUDA
|
||
--------------
|
||
|
||
.. code-block:: c++
|
||
|
||
// one TMA instruction copies the whole rank-3 tile, async, into shared
|
||
"cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes"
|
||
".cta_group::1 [%0], [%1, {%3, %4, %5}], [%2];"
|
||
// call: ptx_cp_async_bulk_tensor_g2cluster_tile_3d(smem, mbar, tensormap, coords...)
|
||
|
||
The three ``{%3, %4, %5}`` are the descriptor coordinates; ``[%1]`` is the
|
||
tensor-map address, ``[%2]`` the mbarrier. One thread launches the entire 8×256
|
||
copy. (This was compiled for ``sm_100a`` — Blackwell — so the instruction carries
|
||
the ``.cta_group::1`` qualifier; on Hopper the qualifier is omitted.)
|
||
|
||
How inputs change the algorithm
|
||
-------------------------------
|
||
|
||
.. list-table::
|
||
:header-rows: 1
|
||
:widths: 26 74
|
||
|
||
* - input
|
||
- effect
|
||
* - direction
|
||
- ``g2s`` → ``cp.async.bulk.tensor.*.g2c``; ``s2g`` → ``…s2g``; with a reduce
|
||
op → ``…s2g_reduce`` (e.g. ``add``)
|
||
* - shared swizzle mode
|
||
- sets the ``swizzle_mode`` in the descriptor and the innermost-box constraint;
|
||
a 128-B swizzle on a 2-D tile yields a **rank-3** descriptor (the inner axis
|
||
splits into swizzle atoms), as in the demo
|
||
* - box shape / chain prefix
|
||
- more selected stride-1 shards → more box>1 descriptor dims; merge collapses
|
||
contiguous full-box dims; box > 256 triggers dtype promotion (1→2→4→8 B)
|
||
* - dtype
|
||
- sets element size and the descriptor's element strides / box byte width
|
||
|
||
A copy whose layout cannot form a legal descriptor (rank > 5 after shrinking, or no
|
||
swizzle-atom-aligned innermost box) makes the plan search fail and the variant
|
||
declines.
|