wehub-resource-sync
478 lines
19 KiB
ReStructuredText
478 lines
19 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.
|
||
|
||
Buffers and memory
|
||
==================
|
||
|
||
Parameter buffers are bound with ``T.match_buffer``; scratch buffers are created
|
||
in the body with one of two declaration APIs (below). Index a buffer with
|
||
``A[i, j]``, slice it with ``A[m0:m0+BM, 0:BK]`` (a ``BufferRegion``), and take a
|
||
pointer with ``A.ptr_to([i, j])`` or the raw data pointer ``A.data``.
|
||
|
||
Declaring buffers
|
||
-----------------
|
||
|
||
Two fundamental APIs create a buffer:
|
||
|
||
- ``T.alloc_buffer(shape, dtype, scope=..., ...)`` — **allocates new storage**
|
||
(emits an ``AllocBuffer`` node) and returns the ``Buffer``. ``T.alloc_shared`` /
|
||
``T.alloc_local`` are just ``alloc_buffer`` with ``scope="shared"`` /
|
||
``scope="local"``.
|
||
- ``T.decl_buffer(shape, dtype, data=..., ...)`` — **declares a view** over an
|
||
existing pointer ``data`` (no allocation); use it to alias or reinterpret
|
||
storage — a sub-region of a pool, or a tensor-memory address. With ``data=None``
|
||
it allocates, like ``alloc_buffer``.
|
||
|
||
A buffer's ``data`` pointer is an immutable ``Var`` (``alloc_buffer`` defines it;
|
||
``decl_buffer`` takes one). To back a buffer with a pointer *expression*, bind it
|
||
first — see :doc:`data_types`.
|
||
|
||
Both share one descriptor; the parameters that matter most:
|
||
|
||
.. list-table::
|
||
:header-rows: 1
|
||
:widths: 28 72
|
||
|
||
* - Parameter
|
||
- Meaning
|
||
* - ``dtype``
|
||
- element type — ``"float32"``, ``"float16"``, ``"float4_e2m1fn"``, …
|
||
* - ``shape``
|
||
- logical shape (a tuple of extents)
|
||
* - ``layout``
|
||
- physical mapping (:doc:`TileLayout <../../layout>`); ``"default"`` = dense
|
||
row-major
|
||
* - ``elem_offset`` / ``allocated_addr``
|
||
- ``elem_offset`` (or ``byte_offset``) places a *view* at an offset into
|
||
``data``; ``allocated_addr`` carries a pre-assigned address (tensor memory)
|
||
* - ``align``
|
||
- alignment of the data pointer, in bytes
|
||
|
||
The ``scope`` argument selects the memory space:
|
||
|
||
.. list-table::
|
||
:header-rows: 1
|
||
:widths: 26 22 52
|
||
|
||
* - Scope
|
||
- Shorthand
|
||
- Memory
|
||
* - ``"global"``
|
||
- (default)
|
||
- device global memory
|
||
* - ``"shared"``
|
||
- ``T.alloc_shared``
|
||
- static shared memory (``__shared__``)
|
||
* - ``"shared.dyn"``
|
||
- (pool)
|
||
- dynamic shared memory (pooled — see below)
|
||
* - ``"local"``
|
||
- ``T.alloc_local``
|
||
- per-thread registers
|
||
* - ``"tmem"``
|
||
- (TMEM pool)
|
||
- Blackwell tensor memory (see below)
|
||
|
||
.. code-block:: python
|
||
|
||
A = T.match_buffer(A_ptr, (M, K), "float16", align=16) # parameter buffer
|
||
As = T.alloc_shared((BM, BK), "float16") # new shared tile
|
||
acc = T.alloc_local((4,), "float32") # register accumulator
|
||
view = T.decl_buffer((BM, BK), "float16", data=As.data) # a view over As
|
||
|
||
**A ptr-based buffer is just metadata over a pointer.** For any non-tmem buffer,
|
||
the declaration is a pointer plus a layout, and indexing resolves to an address::
|
||
|
||
addr(buffer[coord]) = buffer.data + elem_offset + layout.apply(coord, shape=shape)["m"]
|
||
|
||
(``layout.apply`` returns the per-axis mapping; its ``"m"`` component is the
|
||
element offset.) So the *same* logical access compiles to different address
|
||
arithmetic depending purely on the buffer's metadata. Writing
|
||
``B[i, j] = A[i, j] + 1`` over a 4×8 region, with ``B`` declared four ways:
|
||
|
||
.. code-block:: python
|
||
|
||
from tvm.tirx.layout import TileLayout, S
|
||
|
||
B = T.match_buffer(p, (4, 8), "float32") # row-major
|
||
B = T.match_buffer(p, (4, 8), "float32", layout=TileLayout(S[(4, 8):(1, 4)])) # column-major
|
||
B = T.match_buffer(p, (4, 8), "float32", elem_offset=64) # shifted view
|
||
B = T.match_buffer(p, (4, 8), "float32", layout=TileLayout(S[(4, 8):(16, 1)])) # row stride 16
|
||
|
||
each makes ``B[i, j]`` lower to a different index in the generated CUDA (the
|
||
``A[i, j]`` load stays ``i*8 + j`` — only ``B``'s metadata changed):
|
||
|
||
.. code-block:: c++
|
||
|
||
B_ptr[((i * 8) + j)] = ...; // row-major: i*8 + j
|
||
B_ptr[((j * 4) + i)] = ...; // column-major: j*4 + i
|
||
B_ptr[(((i * 8) + j) + 64)] = ...; // elem_offset=64: i*8 + j + 64
|
||
B_ptr[((i * 16) + j)] = ...; // row stride 16: i*16 + j
|
||
|
||
Shared memory
|
||
-------------
|
||
|
||
Shared memory comes in two flavors — **static** (fixed at compile time) and
|
||
**dynamic** (sized at launch) — plus a pool helper that manages the dynamic case.
|
||
|
||
Static
|
||
~~~~~~
|
||
|
||
The simplest shared buffer is a **static** one — ``T.alloc_shared`` (that is,
|
||
``scope="shared"``), sized at compile time. Stage data into it, ``cta_sync`` so the
|
||
whole block sees the writes, then read it back:
|
||
|
||
.. code-block:: python
|
||
|
||
@T.prim_func
|
||
def smem_demo(A_ptr: T.handle, B_ptr: T.handle):
|
||
A = T.match_buffer(A_ptr, (128,), "float32")
|
||
B = T.match_buffer(B_ptr, (128,), "float32")
|
||
T.device_entry()
|
||
bx = T.cta_id([1])
|
||
tx = T.thread_id([128])
|
||
sm = T.alloc_shared((128,), "float32") # static shared memory
|
||
sm[tx] = A[tx]
|
||
T.cuda.cta_sync()
|
||
B[tx] = sm[tx] * T.float32(2.0)
|
||
|
||
It lowers to a plain ``__shared__`` array (generated CUDA, boilerplate elided):
|
||
|
||
.. code-block:: c++
|
||
|
||
extern "C" __global__ void __launch_bounds__(128)
|
||
smem_demo_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr) {
|
||
int tx = ((int)threadIdx.x);
|
||
__shared__ alignas(64) float sm_ptr[128]; // T.alloc_shared
|
||
sm_ptr[tx] = A_ptr[tx];
|
||
__syncthreads(); // T.cuda.cta_sync()
|
||
B_ptr[tx] = sm_ptr[tx] * 2.0f;
|
||
}
|
||
|
||
Dynamic
|
||
~~~~~~~
|
||
|
||
**Dynamic** shared memory (``scope="shared.dyn"``) is sized per launch (the
|
||
``sharedMemBytes`` launch parameter), not at compile time. A kernel may have **only
|
||
one** dynamic-shared allocation — the *arena*. So you allocate it once and ``decl``
|
||
each buffer as a view into it: ``T.decl_buffer`` with ``data=`` the arena pointer
|
||
and an ``elem_offset``:
|
||
|
||
.. code-block:: python
|
||
|
||
arena = T.alloc_buffer((128,), "float32", scope="shared.dyn") # the one arena
|
||
As = T.decl_buffer((64,), "float32", data=arena.data, scope="shared.dyn") # offset 0
|
||
Bs = T.decl_buffer((64,), "float32", data=arena.data, elem_offset=64, scope="shared.dyn") # offset 64
|
||
As[tx] = A[tx]
|
||
Bs[tx] = B[tx]
|
||
T.cuda.cta_sync()
|
||
C[tx] = As[tx] + Bs[tx]
|
||
|
||
Both views share the single ``extern __shared__`` arena (generated CUDA,
|
||
boilerplate elided; arena named ``smem`` for clarity):
|
||
|
||
.. code-block:: c++
|
||
|
||
extern __shared__ __align__(64) float smem[]; // the one dynamic-shared arena
|
||
smem[tx] = A_ptr[tx]; // As — view at offset 0
|
||
smem[tx + 64] = B_ptr[tx]; // Bs — view at offset 64
|
||
__syncthreads();
|
||
C_ptr[tx] = smem[tx] + smem[tx + 64];
|
||
|
||
(Two separate ``alloc_buffer(scope="shared.dyn")`` is an error — *only one dynamic
|
||
shared memory allocation is allowed*.) So static shared memory is sized at compile
|
||
time (``__shared__ T x[N];``); dynamic shared memory is this one launch-sized arena
|
||
with views decl'd at offsets inside it.
|
||
|
||
.. note::
|
||
|
||
**How TVM annotates the dynamic-shared size.** The arena's size is known at
|
||
compile time (here ``128`` floats = ``512`` bytes). During lowering TVM appends
|
||
a ``"tirx.use_dyn_shared_memory"`` tag to the device kernel's
|
||
``tirx.kernel_launch_params``, and the host launcher computes the total bytes and
|
||
passes them as the last launch argument:
|
||
|
||
.. code-block:: python
|
||
|
||
# device kernel attribute:
|
||
"tirx.kernel_launch_params": ["blockIdx.x", "threadIdx.x", "tirx.use_dyn_shared_memory"]
|
||
|
||
# host-side launch call (..., gridDim.x, blockDim.x, dyn_shared_bytes):
|
||
T.call_packed("dyn_kernel", A.data, B.data, C.data, 1, 64, 512)
|
||
|
||
At run time that ``512`` becomes ``config.sharedMemBytes`` in the
|
||
``cuLaunchKernelEx`` call. You never set it by hand — it is derived from the
|
||
``shared.dyn`` allocation's size.
|
||
|
||
Pool sugar
|
||
~~~~~~~~~~
|
||
|
||
``T.SMEMPool`` automates that arena bookkeeping — it bump-allocates the offsets so
|
||
you don't ``decl`` views by hand. Beyond ``alloc`` / ``commit``, it offers
|
||
per-buffer ``align=``, an ``alloc_mma`` helper that builds an MMA-compatible
|
||
swizzle layout for you, and ``move_base_to`` to rewind the cursor and reuse space:
|
||
|
||
.. code-block:: python
|
||
|
||
pool = T.SMEMPool() # bump allocator over shared.dyn
|
||
As = pool.alloc((BM, BK), "float16", align=128) # carve a tile
|
||
Bs = pool.alloc((BK, BN), "float16", align=128)
|
||
Cs = pool.alloc_mma((BM, BN), "float16") # MMA-compatible, swizzle inferred
|
||
pool.commit() # finalize the pool's size
|
||
# pool.move_base_to(offset) rewinds the cursor to reuse space
|
||
|
||
The TMEM pool (`Tensor memory`_, below) is layered on top of an ``SMEMPool``.
|
||
|
||
Registers
|
||
---------
|
||
|
||
Per-thread scratch lives in registers. Allocate it with ``T.alloc_local(shape,
|
||
dtype)`` (i.e. ``scope="local"``): it is private to each thread and lowers to a
|
||
local array kept in registers.
|
||
|
||
.. code-block:: python
|
||
|
||
r = T.alloc_local((4,), "float32") # per-thread register array
|
||
for k in T.unroll(4):
|
||
r[k] = A[tx, k]
|
||
# ... compute on r[0..3] ...
|
||
|
||
.. code-block:: c++
|
||
|
||
alignas(64) float r_ptr[4]; // per-thread, register-resident
|
||
r_ptr[0] = A_ptr[tx * 4 + 0];
|
||
r_ptr[1] = A_ptr[tx * 4 + 1];
|
||
// ...
|
||
|
||
.. note::
|
||
|
||
The ``alignas(64)`` is the *default* buffer alignment — a buffer's
|
||
``data_alignment`` defaults to ``runtime::kAllocAlignment`` (64 bytes), and the
|
||
CUDA codegen stamps it onto every allocation, including per-thread ``local``
|
||
arrays where it is meaningless. For these register-resident arrays it has **no
|
||
performance impact**: a thread-local array with statically-resolvable indices is
|
||
promoted to registers by nvcc/ptxas (scalar replacement of aggregates, SROA), so
|
||
it never lives in addressable local memory and the alignment is a no-op. (A
|
||
dynamically-indexed array that spilled to local memory would actually pick up the
|
||
over-alignment, but that is the unusual case.) This over-alignment of register
|
||
locals is a known rough edge we plan to fix (use the dtype's natural alignment
|
||
for ``local`` scope).
|
||
|
||
Scalar
|
||
~~~~~~
|
||
|
||
A scalar is just a register array with **one element** — strictly, you don't need a
|
||
separate concept. You can allocate a size-1 ``local`` buffer and index ``[0]``:
|
||
|
||
.. code-block:: python
|
||
|
||
phase = T.alloc_local((1,), "int32") # 1-element register array
|
||
phase[0] = 0
|
||
while phase[0] < 4:
|
||
acc = acc + A[tx, phase[0]]
|
||
phase[0] += 1
|
||
|
||
But writing ``phase[0]`` everywhere is clumsy, so a **scalar** is sugar for exactly
|
||
this — a one-element register buffer you read and write **by name**:
|
||
|
||
.. code-block:: python
|
||
|
||
phase: T.int32 = 0 # mutable scalar (sugar for the above)
|
||
while phase < 4:
|
||
acc = acc + A[tx, phase]
|
||
phase += 1
|
||
|
||
s = T.local_scalar("int32") # explicit form; assign by name (s = ..., not s[0])
|
||
acc: T.float32 = 0.0 # a type-annotated assignment also makes one
|
||
|
||
The two are not just similar — they parse to **structurally identical TIRx**. The
|
||
sugar is resolved entirely in the parser: ``phase: T.int32`` *is* that one-element
|
||
``local`` buffer, and ``phase`` / ``phase += 1`` *are* ``phase[0]`` /
|
||
``phase[0] += 1``. ``tvm.ir.assert_structural_equal`` on the two kernels passes, and
|
||
the printer even renders the explicit ``alloc_local`` + ``[0]`` form **back** as the
|
||
scalar form — so once parsing is done there is no difference at all. Both therefore
|
||
lower to the same ``alignas(64) int phase_ptr[1];``; the scalar just lets you drop
|
||
the ``[0]``. (``T.local_scalar`` / ``T.shared_scalar`` / ``T.alloc_scalar`` choose
|
||
the scope explicitly.)
|
||
|
||
.. note::
|
||
|
||
**Why not a** ``Var``\ **?** A TIRx ``Var`` is *immutable* — a single static
|
||
binding (it is exactly what ``T.let`` produces, below). A scalar needs to be
|
||
*mutable* — you reassign it in loops and accumulators — so it must be backed by a
|
||
one-element buffer you can store into repeatedly, not a ``Var``.
|
||
|
||
``let``
|
||
~~~~~~~
|
||
|
||
A ``T.let`` binding is **immutable** — a single ``LetStmt`` (a named value, not a
|
||
buffer). Use it for derived constants:
|
||
|
||
.. code-block:: python
|
||
|
||
n: T.let = M * K # immutable binding (LetStmt)
|
||
half: T.let[T.int32] = N // 2 # ... with an explicit type
|
||
|
||
It lowers to a **plain scalar C variable** — not a buffer (no array, no ``[0]``).
|
||
For ``half: T.let = m * 2`` (with a runtime ``m``):
|
||
|
||
.. code-block:: c++
|
||
|
||
int half = m * 2; // the `let` -> a const-like local
|
||
|
||
Because the value is immutable, the simplifier is free to propagate and CSE it, so
|
||
at the use sites you often see ``m * 2`` substituted directly (or shared through a
|
||
common-subexpression temporary) rather than a reference to ``half``.
|
||
|
||
.. note::
|
||
|
||
**Why have an immutable binding at all?** Because the value cannot change, the
|
||
arithmetic analyzer binds the var to it (``analyzer.Bind(var, value)`` when it
|
||
simplifies a ``LetStmt``), so facts proven about the value — constant bounds, the
|
||
modular set (divisibility / alignment), ranges — **propagate through every use**.
|
||
That feeds index simplification, bounds-check elimination, and
|
||
alignment/vectorization decisions. A *mutable* scalar is a memory load
|
||
(``buf[0]``): the analyzer cannot assume it stays constant, so none of those
|
||
properties carry through. A ``let`` is also a pure value — no allocation, and
|
||
free to inline / substitute / CSE — whereas a scalar is a one-element buffer with
|
||
load/store semantics.
|
||
|
||
Tensor memory
|
||
-------------
|
||
|
||
Blackwell *tensor memory* is not a plain scratch scope: it must be explicitly
|
||
reserved and freed with the warp-uniform ``T.ptx.tcgen05.alloc`` /
|
||
``tcgen05.dealloc`` intrinsics, and each tensor is a view into it declared with
|
||
``T.decl_buffer(..., scope="tmem", allocated_addr=<column>, layout=<tmem layout>)``.
|
||
The ``allocated_addr`` (a column offset) is mandatory — the tensor-core dispatch
|
||
asserts it — so ``T.alloc_buffer(scope="tmem")`` (which does **not** set it) will not
|
||
work. Unlike shared memory, tensor memory is not directly addressable: it is read
|
||
and written only through ``tcgen05`` ``mma`` / ``ld`` / ``st`` / ``cp``.
|
||
|
||
By hand, one warp issues the allocation into a shared slot, you ``decl`` each
|
||
tensor as a view at a column offset, and one warp frees it at the end:
|
||
|
||
.. code-block:: python
|
||
|
||
addr = T.alloc_shared((1,), "uint32") # slot for the allocated base
|
||
if warp_id == alloc_warp: # tcgen05.alloc is warp-uniform
|
||
T.ptx.tcgen05.alloc(T.address_of(addr), n_cols=512, cta_group=cta_group)
|
||
acc = T.decl_buffer((CTA_M, 512), "float32", scope="tmem",
|
||
allocated_addr=0, layout=tmem_layout) # view at column 0
|
||
# ... use acc as a gemm_async / copy_async operand ...
|
||
if warp_id == alloc_warp:
|
||
T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group)
|
||
T.ptx.tcgen05.dealloc(addr, n_cols=512, cta_group=cta_group)
|
||
|
||
You manage the column offsets and the ``tmem_layout`` (a datapath D/F layout)
|
||
yourself. This is exactly the sequence the pool below emits.
|
||
|
||
Pool
|
||
~~~~
|
||
|
||
``T.TMEMPool`` wraps all of that — the warp-uniform alloc/dealloc, the column
|
||
bump-allocation, and the datapath layout:
|
||
|
||
.. code-block:: python
|
||
|
||
tmem_addr = pool.alloc((1,), "uint32") # pool = the kernel's smem pool
|
||
tmem_pool = T.TMEMPool(pool, total_cols=512, cta_group=cta_group,
|
||
tmem_addr=tmem_addr)
|
||
acc = tmem_pool.alloc((CTA_M, 512), "float32") # allocated_addr set for you
|
||
tmem_pool.commit() # emits tcgen05.alloc (one warp)
|
||
# ... use acc ...
|
||
tmem_pool.dealloc() # emits tcgen05.dealloc (one warp)
|
||
|
||
See the :doc:`../../tile_primitives` walkthroughs for full examples.
|
||
|
||
Buffer APIs
|
||
-----------
|
||
|
||
A ``Buffer`` is metadata over a pointer (see *Declaring buffers* above), so most of
|
||
its methods are *compile-time* reshapes/reinterprets that change index arithmetic
|
||
or hand you a pointer — they emit no runtime op of their own. The common ones:
|
||
|
||
.. list-table::
|
||
:header-rows: 1
|
||
:widths: 34 66
|
||
|
||
* - Method
|
||
- What it is
|
||
* - ``B.data``
|
||
- the raw data pointer (a ``Var``); prints as ``B_ptr``
|
||
* - ``B.ptr_to([i, j])``
|
||
- a typed pointer to an element (``address_of``); prints as ``&B_ptr[…]``
|
||
* - ``B.vload([i], dtype="float32x4")`` / ``B.vstore([i], v)``
|
||
- a vectorized load / store; prints as ``*(float4*)(B_ptr + …)``
|
||
* - ``B.view(*shape, layout=…)``
|
||
- reinterpret the same storage under a new shape/layout (no copy)
|
||
* - ``B.local(*shape, layout=…)``
|
||
- the calling thread's private register slice of a ``local`` buffer
|
||
* - ``B.permute(*dims)``
|
||
- a view with axes permuted (a transposed layout)
|
||
* - ``B.access_ptr(mask, …)``
|
||
- a masked access pointer (the ``tvm_access_ptr`` builtin), for passing a
|
||
region to an intrinsic
|
||
|
||
**Pointers — ``ptr_to`` / ``data``.** ``ptr_to`` is how you hand an element address
|
||
to an intrinsic or inline function; ``data`` is the base pointer:
|
||
|
||
.. code-block:: python
|
||
|
||
B[tx] = T.cuda.func_call("ld", A.ptr_to([tx]), source_code=SRC, return_type="float32")
|
||
|
||
.. code-block:: c++
|
||
|
||
B_ptr[tx] = ld(&A_ptr[tx]); // ptr_to([tx]) -> &A_ptr[tx]; A.data -> A_ptr
|
||
|
||
**Vectorized access — ``vload`` / ``vstore``.** Move several elements as one wide
|
||
transfer (see also :doc:`data_types`):
|
||
|
||
.. code-block:: python
|
||
|
||
B.vstore([tx * 4], A.vload([tx * 4], dtype="float32x4"))
|
||
|
||
.. code-block:: c++
|
||
|
||
*(float4*)(B_ptr + tx * 4) = *(float4*)(A_ptr + tx * 4);
|
||
|
||
**Reshape / reinterpret — ``view`` / ``permute``.** Both are pure metadata; the
|
||
data pointer is unchanged, only the index arithmetic differs. ``A.view(64, 4)``
|
||
sees the 256-element buffer as ``64×4``; ``A.permute(1, 0)`` transposes the axes:
|
||
|
||
.. code-block:: python
|
||
|
||
A2 = A.view(64, 4); y = A2[tx, 0] + A2[tx, 3] # A2[tx, j] -> A_ptr[tx*4 + j]
|
||
At = A.permute(1, 0); z = At[i, j] # At[i, j] -> A_ptr[j*4 + i]
|
||
|
||
.. code-block:: c++
|
||
|
||
A2_ptr[tx * 4] /* +3 */ // view: row-major 64x4 index
|
||
At_ptr[(j * 4) + i] // permute: swapped strides
|
||
|
||
**Registers — ``local``.** Decomposes a thread-axis ``local`` layout into the
|
||
calling thread's flat register bundle (used pervasively by the tile primitives):
|
||
|
||
.. code-block:: python
|
||
|
||
R = T.alloc_buffer((32, 8), "float32", scope="local", layout=TileLayout(S[(32, 8) : (1 @ laneid, 1)]))
|
||
Rl = R.local(8) # this lane's 8 registers
|
||
|
||
.. code-block:: c++
|
||
|
||
alignas(64) float Rl_ptr[8]; // the lane's private registers
|