wehub-resource-sync
207 lines
8.0 KiB
ReStructuredText
207 lines
8.0 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 → reg
|
||
==========
|
||
|
||
The ``reg`` variant lowers a synchronous ``copy`` where **exactly one side is a
|
||
register** (``local``) buffer and the other is ``shared*`` or ``global``. Unlike
|
||
:doc:`gmem_smem`, the partition is **not synthesized** — it is *induced* by the
|
||
register operand's layout: that layout's thread-axis iters already say which thread
|
||
owns which logical coordinate, so the dispatch drops those axes, leaves each thread
|
||
its private bundle of elements, and copies them in a vectorized serial loop. Source:
|
||
``python/tvm/backend/cuda/operator/tile_primitive/copy/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-axis iters have **stride 1**, a register-level
|
||
subscope no wider than the exec scope, and a **zero sliced thread offset**
|
||
(the region doesn't split a thread axis)
|
||
* - 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, 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))
|
||
|
||
@T.prim_func
|
||
def kernel(B_ptr: T.handle):
|
||
B = T.match_buffer(B_ptr, shape, dtype)
|
||
T.device_entry(); T.cta_id([1]); T.lane_id([32]); tid = T.thread_id([32])
|
||
A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
|
||
for kk in range(8): A_smem[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
|
||
T.cuda.cta_sync()
|
||
R = T.alloc_buffer(shape, dtype, scope="local", layout=r_layout)
|
||
Tx.warp.copy(R[fs], A_smem[fs]) # shared -> register (this dispatch)
|
||
# ... clear A_smem, cta_sync ...
|
||
Tx.warp.copy(A_smem[fs], R[fs]) # 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 dispatch 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 ``T.lane_id()`` etc.),
|
||
and the register side is a flat per-thread ``local`` buffer. The emit is a serial
|
||
loop (not ``T.unroll`` — same flooding rationale as :doc:`gmem_smem`):
|
||
|
||
.. code-block:: python
|
||
|
||
r_local = r_buf.local(*per_thread_r_shape) # flat per-thread registers
|
||
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_ptr = _ptr_off(r_local.ptr_to([0]), r_off_base + dr)
|
||
if r_is_src:
|
||
copy_op(s_ptr, r_ptr) # register -> shared/global
|
||
else:
|
||
copy_op(r_ptr, s_ptr) # shared/global -> register
|
||
|
||
Generated TIRx IR
|
||
-----------------
|
||
|
||
``LowerTIRx`` turns the shared → register copy into a per-thread loop over the
|
||
8-element register bundle (trimmed):
|
||
|
||
.. code-block:: python
|
||
|
||
r_local = T.decl_buffer((8,), data=R.data, scope="local") # 8 regs / lane
|
||
for f in range(2): # outer = 8 / vec 4
|
||
s_ptr = pointer_offset(A_smem, ...) # this lane's row
|
||
r_ptr = pointer_offset(r_local, dr)
|
||
T.cuda.copy_bytes(r_ptr, s_ptr, 16) # 16 B = vec 4 × 4 B
|
||
|
||
(The register → shared copy is the mirror: ``copy_bytes(s_ptr, r_ptr, 16)``.)
|
||
|
||
Generated CUDA
|
||
--------------
|
||
|
||
.. code-block:: c++
|
||
|
||
alignas(64) float r_local_ptr[8]; // 8 registers, private to the lane
|
||
for (int f = 0; f < 2; ++f) {
|
||
void* r_ptr = tvm_builtin_pointer_offset(&r_local_ptr[0], dr);
|
||
void* s_ptr = tvm_builtin_pointer_offset(&A_smem_ptr[0], /* lane row + f*4 */);
|
||
tvm_builtin_copy_128b(r_ptr, s_ptr); // shared -> register
|
||
}
|
||
// ... register -> shared mirror writes A_smem back ...
|
||
|
||
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 register count, vector width, and round count:
|
||
|
||
.. list-table::
|
||
:header-rows: 1
|
||
:widths: 18 18 22 22 20
|
||
|
||
* - dtype
|
||
- ``k``
|
||
- regs / 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 is always a 128-bit transfer (``copy_bytes = 16``) 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``.
|