wehub-resource-sync
219 lines
8.4 KiB
ReStructuredText
219 lines
8.4 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 → ldstmatrix
|
||
=================
|
||
|
||
The ``ldstmatrix`` variant lowers a ``copy`` between **register and shared** memory
|
||
to the warp-collective PTX ``ldmatrix`` / ``stmatrix`` instructions: one
|
||
instruction moves ``num`` 8×8 16-bit matrix tiles between shared memory and the
|
||
warp's registers, with the hardware performing the lane↔element shuffle that an
|
||
MMA fragment needs. It only applies when the register and shared **layouts match
|
||
the m8n8 fragment geometry**; otherwise the copy falls back to :doc:`reg`. Source:
|
||
``python/tvm/backend/cuda/operator/tile_primitive/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 **real** gate is the layout fit, applied during emit. Both this variant and
|
||
:doc:`reg` are priority 10 and both accept ``local ↔ shared``; ``ldstmatrix`` is
|
||
tried first and **declines** (via ``fail(...)``) if the layouts are not ldmatrix
|
||
fragments, leaving ``reg`` to handle the copy:
|
||
|
||
.. 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
|
||
- 16-bit (``.b16``) — ldmatrix/stmatrix move 8 fp16 = 16 B per lane per tile
|
||
* - 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}``
|
||
|
||
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, 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))
|
||
|
||
@T.prim_func
|
||
def kernel(A_ptr: T.handle, B_ptr: T.handle):
|
||
A = T.match_buffer(A_ptr, (M, N), "float16")
|
||
B = T.match_buffer(B_ptr, (M, N), "float16")
|
||
T.device_entry(); T.cta_id([1]); T.lane_id([32]); tid = T.thread_id([32])
|
||
A_smem = T.alloc_buffer((8, 4, num, 2), "float16", scope="shared", layout=s_layout)
|
||
# ... stage A into A_smem (row = tid//4, cp = tid%4) ...
|
||
T.cuda.cta_sync()
|
||
R = T.alloc_buffer((8, 4, num, 2), "float16", scope="local", layout=r_layout)
|
||
Tx.warp.copy(R[full], A_smem[full]) # 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 enforce 16-byte alignment (8 fp16) for every tile and every
|
||
``m_outer`` advance, since each lane's ``.b16`` access reads 16 bytes.
|
||
|
||
**4. Emit one instruction per** ``m_outer`` **tile group.** Each lane contributes
|
||
its shared address (tile offset + ``(laneid % 8) · p``) and ``num`` register
|
||
handles:
|
||
|
||
.. code-block:: python
|
||
|
||
for mm in T.unroll(m_outer):
|
||
smem_ptr = _ptr_off(s_buf.ptr_to(s_zero), _smem_off(mm, tile_off + (laneid % 8) * p))
|
||
handles = [r_local.ptr_to([...]) for i in range(num)]
|
||
if direction == "ld":
|
||
T.ptx.ldmatrix(trans, num, ".b16", smem_ptr, *handles)
|
||
else:
|
||
T.ptx.stmatrix(trans, num, ".b16", smem_ptr, *handles, shape="m8n8", space="shared")
|
||
|
||
(This is the one copy variant that **does** use ``T.unroll`` — ``m_outer`` is tiny.)
|
||
|
||
Generated TIRx IR
|
||
-----------------
|
||
|
||
For the demo (``num = 2``, ``M = 8`` ⇒ ``m_outer = 1``):
|
||
|
||
.. code-block:: python
|
||
|
||
for mm in T.unroll(1):
|
||
T.ptx.ldmatrix(T.bool(False), 2, ".b16", smem_ptr,
|
||
T.address_of(r_local[0]), T.address_of(r_local[2]))
|
||
|
||
Generated CUDA
|
||
--------------
|
||
|
||
.. code-block:: c++
|
||
|
||
__forceinline__ __device__ void ptx_ldmatrix_2_b16_0(void* smem_ptr, void* dst0, void* dst1) {
|
||
// ...
|
||
"ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];"
|
||
// ...
|
||
}
|
||
// call site (per lane):
|
||
ptx_ldmatrix_2_b16_0(smem_ptr, &r_local_ptr[0], &r_local_ptr[2]);
|
||
|
||
``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 handles; ``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 :doc:`reg` instead.
|