wehub-resource-sync
530 lines
20 KiB
Python
530 lines
20 KiB
Python
# 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.
|
|
"""SMEM and TMEM bump-allocator pools for TIRX kernels."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
import operator
|
|
|
|
from tvm import DataType
|
|
from tvm.tirx.layout import S, TCol, TileLayout, TLane
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ir_builder helpers — imported lazily to avoid circular deps at module level
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_ir = None
|
|
|
|
|
|
def _get_ir():
|
|
global _ir
|
|
if _ir is None:
|
|
from tvm.tirx.script.builder import ir as _mod
|
|
|
|
_ir = _mod
|
|
return _ir
|
|
|
|
|
|
def _get_frame():
|
|
from tvm.tirx.script.builder import frame
|
|
|
|
return frame
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Shared utilities
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_POOL_UNSET = object()
|
|
|
|
|
|
def _default_tmem_layout(rows, cols):
|
|
return TileLayout(S[(rows, cols) : (1 @ TLane, 1 @ TCol)])
|
|
|
|
|
|
def _emit_stmt(expr):
|
|
ir = _get_ir()
|
|
ir.add_to_parent(ir.evaluate(expr))
|
|
|
|
|
|
def _shape_product(shape):
|
|
return functools.reduce(operator.mul, shape, 1)
|
|
|
|
|
|
def _auto_swizzle_mode(dtype):
|
|
"""Select the default MMA swizzle mode for a shared-memory allocation."""
|
|
from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
|
|
|
|
del dtype
|
|
return SwizzleMode.SWIZZLE_128B_ATOM
|
|
|
|
|
|
def _swizzle_atom_bytes(swizzle_mode):
|
|
"""Return the row width (in bytes) of one swizzle atom for *swizzle_mode*."""
|
|
from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
|
|
|
|
return {
|
|
SwizzleMode.SWIZZLE_NONE: 0,
|
|
SwizzleMode.SWIZZLE_32B_ATOM: 32,
|
|
SwizzleMode.SWIZZLE_64B_ATOM: 64,
|
|
SwizzleMode.SWIZZLE_128B_ATOM: 128,
|
|
}[swizzle_mode]
|
|
|
|
|
|
def _suggest_swizzle_for_row_bytes(row_bytes):
|
|
"""Pick the largest valid swizzle mode whose atom row fits within *row_bytes*."""
|
|
|
|
for atom_bytes, mode in (
|
|
(128, "SWIZZLE_128B_ATOM"),
|
|
(64, "SWIZZLE_64B_ATOM"),
|
|
(32, "SWIZZLE_32B_ATOM"),
|
|
):
|
|
if row_bytes >= atom_bytes and row_bytes % atom_bytes == 0:
|
|
return mode
|
|
return "SWIZZLE_NONE"
|
|
|
|
|
|
def _validate_mma_alloc_shape(shape, dtype, swizzle_mode):
|
|
"""Validate that *shape* / *dtype* / *swizzle_mode* are mutually compatible.
|
|
|
|
``mma_shared_layout`` tiles a swizzle atom of shape ``[8, swizzle_bytes / dtype_bytes]``
|
|
over the last two logical dimensions of *shape*. If the row width or row count of
|
|
the request is smaller than (or not a multiple of) the atom, the underlying
|
|
``Layout.tile_to`` lowers to a ``floordiv``/``floormod`` by zero and raises an
|
|
opaque internal "Divide by zero" diagnostic from ``tile_tile_ops.cc``. Catch the
|
|
misconfiguration here so callers see *what* is wrong and *how* to fix it.
|
|
|
|
Validation skipped when *swizzle_mode* is ``SWIZZLE_NONE`` (no atom).
|
|
"""
|
|
from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
|
|
|
|
if swizzle_mode == SwizzleMode.SWIZZLE_NONE:
|
|
return
|
|
|
|
if len(shape) < 2:
|
|
raise ValueError(
|
|
f"alloc_mma shape={tuple(shape)} has fewer than 2 dimensions; "
|
|
f"swizzled MMA layouts tile over the last two dims (rows, cols). "
|
|
f"Use swizzle_mode='none' for 1-D allocations."
|
|
)
|
|
|
|
# Only validate concrete int dims; symbolic dims fall through (the analyzer
|
|
# in C++ will still ICHECK on them, but at least we don't false-positive).
|
|
rows = shape[-2]
|
|
cols = shape[-1]
|
|
if not (isinstance(rows, int) and isinstance(cols, int)):
|
|
return
|
|
|
|
dtype_bytes = DataType(dtype).bits // 8
|
|
if dtype_bytes == 0:
|
|
# Sub-byte dtype (e.g. float4); ``cols`` is already in element units, so
|
|
# use a fractional check expressed via bits.
|
|
col_bits = cols * DataType(dtype).bits
|
|
atom_bits = _swizzle_atom_bytes(swizzle_mode) * 8
|
|
if col_bits < atom_bits or col_bits % atom_bits != 0:
|
|
row_bytes = col_bits // 8 if col_bits % 8 == 0 else col_bits / 8
|
|
atom_bytes = _swizzle_atom_bytes(swizzle_mode)
|
|
suggestion = _suggest_swizzle_for_row_bytes(col_bits // 8 if col_bits >= 8 else 0)
|
|
raise ValueError(
|
|
f"alloc_mma shape={tuple(shape)} with dtype={dtype!r} produces "
|
|
f"{row_bytes}B rows, which is incompatible with the {atom_bytes}B "
|
|
f"swizzle atom selected by {swizzle_mode.name}. "
|
|
f"Use swizzle_mode=SwizzleMode.{suggestion}, or widen shape[-1] "
|
|
f"to a multiple of "
|
|
f"{(atom_bits + DataType(dtype).bits - 1) // DataType(dtype).bits} elements."
|
|
)
|
|
else:
|
|
row_bytes = cols * dtype_bytes
|
|
atom_bytes = _swizzle_atom_bytes(swizzle_mode)
|
|
if row_bytes < atom_bytes or row_bytes % atom_bytes != 0:
|
|
suggestion = _suggest_swizzle_for_row_bytes(row_bytes)
|
|
min_cols = atom_bytes // dtype_bytes
|
|
raise ValueError(
|
|
f"alloc_mma shape={tuple(shape)} with dtype={dtype!r} produces "
|
|
f"{row_bytes}B rows, which is incompatible with the {atom_bytes}B "
|
|
f"swizzle atom selected by {swizzle_mode.name}. "
|
|
f"Use swizzle_mode=SwizzleMode.{suggestion}, or widen shape[-1] "
|
|
f"to a multiple of {min_cols} elements (>= {atom_bytes}B at {dtype})."
|
|
)
|
|
|
|
# Atom rows is always 8 (see ``mma_atom_shape`` in tma_utils.py).
|
|
atom_rows = 8
|
|
if rows < atom_rows or rows % atom_rows != 0:
|
|
raise ValueError(
|
|
f"alloc_mma shape={tuple(shape)} has shape[-2]={rows}, but the "
|
|
f"{swizzle_mode.name} atom requires shape[-2] to be a positive "
|
|
f"multiple of {atom_rows}. Use swizzle_mode='none', or widen shape[-2] "
|
|
f"to a multiple of {atom_rows}."
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TMEMStages
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _meta_class(cls):
|
|
"""Apply @meta_class decorator from ir_builder."""
|
|
return _get_ir().meta_class(cls)
|
|
|
|
|
|
@_meta_class
|
|
class TMEMStages:
|
|
"""Parse-time staged view over a TMEM buffer.
|
|
|
|
Parameters
|
|
----------
|
|
buf : Buffer
|
|
The underlying TMEM buffer (e.g. f32 or f16 view).
|
|
col_start : int
|
|
First column of stage 0 in *buf*'s column space.
|
|
width : int
|
|
Number of columns per stage.
|
|
stages : int
|
|
Number of pipeline stages (default 1).
|
|
stride : int or None
|
|
Column distance between consecutive stages. When *None* (default),
|
|
equals *width* (stages are packed back-to-back).
|
|
"""
|
|
|
|
def __init__(self, buf, col_start, width, stages=1, stride=None):
|
|
self.buf = buf
|
|
self.col_start = col_start
|
|
self.width = width
|
|
self.stages = stages
|
|
self.stride = width if stride is None else stride
|
|
|
|
def _stage_base(self, stage):
|
|
return self.col_start + stage * self.stride
|
|
|
|
def __getitem__(self, item):
|
|
if isinstance(item, tuple):
|
|
assert len(item) == 2, "TMEMStages expects region[stage] or region[stage, start:stop]"
|
|
stage, col_slice = item
|
|
assert isinstance(col_slice, slice), "TMEMStages tuple indexing requires a slice"
|
|
base = self._stage_base(stage)
|
|
start = 0 if col_slice.start is None else col_slice.start
|
|
stop = self.width if col_slice.stop is None else col_slice.stop
|
|
return self.buf[:, base + start : base + stop : col_slice.step]
|
|
base = self._stage_base(item)
|
|
return self.buf[:, base : base + self.width]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TMEMPool
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@_meta_class
|
|
class TMEMPool:
|
|
"""Bump allocator over TMEM columns."""
|
|
|
|
def __init__(
|
|
self,
|
|
pool,
|
|
total_cols=512,
|
|
*,
|
|
cta_group=1,
|
|
alloc_warp=0,
|
|
dealloc_warp=None,
|
|
tmem_addr=None,
|
|
sync_after_alloc=True,
|
|
):
|
|
# tcgen05 alloc/dealloc are warp-uniform PTX instructions: every lane
|
|
# in the chosen warp must participate, and exactly one warp in the
|
|
# CTA must execute them. The pool emits its own
|
|
# ``if warp_id() == target_warp: tcgen05.alloc(...)``
|
|
# guard, using the cta->warp scope id ``T.warp_id()``.
|
|
# NOTE: synccheck currently false-deadlocks on kernels that declare a
|
|
# second warp-scope id (cpusim binds only one warp var); the generated
|
|
# CUDA is equivalent to ``thread_rank() // 32 == target_warp``.
|
|
self.pool = pool
|
|
self.total_cols = total_cols
|
|
self.cta_group = cta_group
|
|
self.alloc_warp = alloc_warp
|
|
self.dealloc_warp = alloc_warp if dealloc_warp is None else dealloc_warp
|
|
self.sync_after_alloc = sync_after_alloc
|
|
self.offset = 0
|
|
self.max_offset = 0
|
|
self._committed = False
|
|
self._deallocated = False
|
|
self._addr_buf = pool.alloc([1], "uint32", align=4) if tmem_addr is None else tmem_addr
|
|
|
|
def _addr_slot(self):
|
|
try:
|
|
return self._addr_buf[0]
|
|
except TypeError:
|
|
return self._addr_buf
|
|
|
|
@property
|
|
def addr(self):
|
|
return self._addr_slot()
|
|
|
|
def _emit_warp_guard(self, target_warp, emit):
|
|
from tvm.script import tirx as T
|
|
|
|
warp_id = T.warp_id()
|
|
with T.If(warp_id == target_warp):
|
|
with T.Then():
|
|
emit()
|
|
|
|
def _resolve_cols(self, shape, dtype, cols, layout=None):
|
|
if cols is not None:
|
|
return cols
|
|
bits = DataType(dtype).bits
|
|
if layout is not None:
|
|
# span("TCol") is in *element* (buffer dtype) units; one TMEM cell
|
|
# holds 32 bits regardless of the element type.
|
|
tcol_elems = int(layout.span("TCol"))
|
|
tcol_bits = tcol_elems * bits
|
|
assert tcol_bits % 32 == 0, (
|
|
f"layout TCol span={tcol_elems} elems x {bits}b is not 32-bit aligned"
|
|
)
|
|
return tcol_bits // 32
|
|
assert len(shape) == 2, "TMEMPool.alloc() requires cols= for non-2D TMEM buffers"
|
|
total_bits = _shape_product(shape) * bits
|
|
rows = shape[0]
|
|
assert total_bits % (32 * rows) == 0, (
|
|
f"Cannot infer TMEM columns from shape={shape}, dtype={dtype!r}; "
|
|
"please pass cols= explicitly"
|
|
)
|
|
return total_bits // (32 * rows)
|
|
|
|
def alloc(self, shape, dtype="float32", *, layout=None, cols=None, datapath=None):
|
|
"""Allocate a TMEM buffer.
|
|
|
|
Parameters
|
|
----------
|
|
shape, dtype, cols
|
|
Standard buffer shape / dtype / column count.
|
|
layout
|
|
Explicit ``TileLayout``. Mutually exclusive with ``datapath``.
|
|
datapath : str | None
|
|
Optional tcgen05 datapath letter (``"D"`` for M=128 full datapath,
|
|
``"F"`` for M=64 non-``.ws`` scattered). When provided, the buffer's
|
|
layout is derived from ``tmem_datapath_layout(datapath, *shape)``
|
|
so the row index reflects the *physical* TMEM lane occupation
|
|
(PTX ISA §9.7.16.10.5). The downstream ``.16x*b`` / ``.32x32b``
|
|
dispatches structurally check this layout to catch mismatched
|
|
atoms (e.g. a ``.16x*b`` M=128 read against a Layout F buffer).
|
|
Defaults to ``None``, which means Layout D's identity row→lane
|
|
mapping — keep this for shape ``(128, X)`` buffers that hold
|
|
an M=128 MMA accumulator.
|
|
"""
|
|
from tvm.tirx.layout import tmem_datapath_layout
|
|
|
|
if layout is not None and datapath is not None:
|
|
raise ValueError("TMEMPool.alloc: pass at most one of layout= and datapath=")
|
|
if datapath is not None:
|
|
assert len(shape) == 2, "TMEMPool.alloc: datapath= requires a 2-D shape"
|
|
layout = tmem_datapath_layout(datapath, shape[0], shape[1])
|
|
|
|
ir = _get_ir()
|
|
cols = self._resolve_cols(shape, dtype, cols, layout)
|
|
col_start = self.offset
|
|
col_end = col_start + cols
|
|
assert col_end <= self.total_cols, f"TMEM overflow: {col_end} > {self.total_cols}"
|
|
if layout is None:
|
|
assert len(shape) == 2, "TMEMPool.alloc() requires layout= for non-2D TMEM buffers"
|
|
layout = _default_tmem_layout(shape[0], shape[1])
|
|
res = ir.decl_buffer(shape, dtype, scope="tmem", allocated_addr=col_start, layout=layout)
|
|
self.offset = col_end
|
|
self.max_offset = max(self.max_offset, self.offset)
|
|
return res
|
|
|
|
def alloc_sf(self, shape, dtype, *, sf_per_mma, sf_reuse=1):
|
|
"""Allocate a tcgen05 block-scaled SF TMEM buffer with an inferred layout.
|
|
|
|
``shape`` last two dims are ``(rows, SF_K * sf_reuse)`` (the last dim is
|
|
what gemm dispatch iterates over). When ``shape`` has 3 dims, the first
|
|
is treated as a pipe-depth outer.
|
|
"""
|
|
from tvm.backend.cuda.operator.tile_primitive.gemm_async.tcgen05 import sf_tmem_layout
|
|
|
|
if len(shape) == 2:
|
|
pipe_depth, rows, last = None, shape[0], shape[1]
|
|
elif len(shape) == 3:
|
|
pipe_depth, rows, last = shape[0], shape[1], shape[2]
|
|
else:
|
|
raise ValueError(
|
|
f"alloc_sf expects 2D (rows, SF_K*sf_reuse) or 3D "
|
|
f"(pipe_depth, rows, SF_K*sf_reuse); got shape={shape}"
|
|
)
|
|
assert last % sf_reuse == 0, (
|
|
f"alloc_sf: shape last dim {last} must be divisible by sf_reuse={sf_reuse}"
|
|
)
|
|
SF_K = last // sf_reuse
|
|
layout = sf_tmem_layout(
|
|
rows=rows, SF_K=SF_K, sf_per_mma=sf_per_mma, sf_reuse=sf_reuse, pipe_depth=pipe_depth
|
|
)
|
|
return self.alloc(shape, dtype, layout=layout)
|
|
|
|
def move_base_to(self, col):
|
|
self.offset = col
|
|
self.max_offset = max(self.max_offset, self.offset)
|
|
|
|
def commit(self):
|
|
assert not self._committed, "TMEMPool.commit() can only be called once"
|
|
from tvm.script import tirx as T
|
|
|
|
def emit_alloc():
|
|
_emit_stmt(
|
|
T.ptx.tcgen05.alloc(
|
|
T.address_of(self.addr), n_cols=self.total_cols, cta_group=self.cta_group
|
|
)
|
|
)
|
|
if self.sync_after_alloc:
|
|
_emit_stmt(T.cuda.warp_sync())
|
|
|
|
self._emit_warp_guard(self.alloc_warp, emit_alloc)
|
|
self._committed = True
|
|
|
|
def dealloc(self):
|
|
assert self._committed, "TMEMPool.dealloc() called before commit()"
|
|
assert not self._deallocated, "TMEMPool.dealloc() can only be called once"
|
|
self._deallocated = True
|
|
from tvm.script import tirx as T
|
|
|
|
def emit_dealloc():
|
|
_emit_stmt(T.ptx.tcgen05.relinquish_alloc_permit(cta_group=self.cta_group))
|
|
_emit_stmt(
|
|
T.ptx.tcgen05.dealloc(self.addr, n_cols=self.total_cols, cta_group=self.cta_group)
|
|
)
|
|
|
|
self._emit_warp_guard(self.dealloc_warp, emit_dealloc)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# SMEMPool
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@_meta_class
|
|
class SMEMPool:
|
|
"""Bump allocator over a contiguous shared memory region.
|
|
|
|
Parameters
|
|
----------
|
|
ptr : Var or None, optional
|
|
If omitted, an ``alloc_buffer([0], "uint8", scope="shared.dyn")`` is
|
|
created automatically and ``commit()`` must be called after all
|
|
allocations to emit the size annotation.
|
|
If a ``Var`` is provided, the caller manages the backing buffer and
|
|
``commit()`` is a no-op.
|
|
"""
|
|
|
|
def __init__(self, ptr=_POOL_UNSET):
|
|
ir = _get_ir()
|
|
if ptr is _POOL_UNSET:
|
|
self.buf = ir.alloc_buffer([0], "uint8", scope="shared.dyn")
|
|
self.ptr = self.buf.data
|
|
self._owns_buffer = True
|
|
else:
|
|
self.buf = None
|
|
self.ptr = ptr
|
|
self._owns_buffer = False
|
|
self.offset = 0
|
|
self.max_offset = 0
|
|
|
|
def alloc(
|
|
self,
|
|
shape,
|
|
dtype="float32",
|
|
strides=None,
|
|
scope="shared.dyn",
|
|
align=0,
|
|
buffer_type="",
|
|
axis_separators=None,
|
|
layout="default",
|
|
):
|
|
ir = _get_ir()
|
|
if align > 0:
|
|
self.offset = (self.offset + align - 1) // align * align
|
|
res = ir.decl_buffer(
|
|
shape,
|
|
dtype,
|
|
data=self.ptr,
|
|
strides=strides,
|
|
byte_offset=self.offset,
|
|
scope=scope,
|
|
align=align,
|
|
buffer_type=buffer_type,
|
|
axis_separators=axis_separators,
|
|
layout=layout,
|
|
)
|
|
# Advance in bits then round up to bytes so sub-byte dtypes (e.g.
|
|
# float4_e2m1fn = 4 bits) still bump the cursor instead of leaving it
|
|
# at 0 (bits // 8) and silently overlapping the next allocation.
|
|
self.offset += (_shape_product(shape) * DataType(dtype).bits + 7) // 8
|
|
if self._owns_buffer:
|
|
self.max_offset = max(self.max_offset, self.offset)
|
|
return res
|
|
|
|
def alloc_mma(self, shape, dtype="float16", swizzle_mode="auto", align=1024):
|
|
"""Allocate MMA-compatible shared memory with an inferred swizzle layout."""
|
|
from tvm.backend.cuda.operator.tile_primitive.tma_utils import (
|
|
SwizzleMode,
|
|
mma_shared_layout,
|
|
)
|
|
|
|
if isinstance(swizzle_mode, str):
|
|
if swizzle_mode == "auto":
|
|
swizzle_mode = _auto_swizzle_mode(dtype)
|
|
elif swizzle_mode == "none":
|
|
swizzle_mode = SwizzleMode.SWIZZLE_NONE
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported swizzle_mode={swizzle_mode!r}; expected 'auto', 'none', "
|
|
"or SwizzleMode"
|
|
)
|
|
_validate_mma_alloc_shape(shape, dtype, swizzle_mode)
|
|
layout = mma_shared_layout(dtype, swizzle_mode, shape)
|
|
return self.alloc(shape, dtype, align=align, layout=layout)
|
|
|
|
def move_base_to(self, offset):
|
|
self.offset = offset
|
|
if self._owns_buffer:
|
|
self.max_offset = max(self.max_offset, self.offset)
|
|
|
|
def commit(self, size=None):
|
|
"""Emit pool size annotation into the IR.
|
|
|
|
Must be called after all ``alloc()`` / ``move_base_to()`` calls.
|
|
|
|
Parameters
|
|
----------
|
|
size : int, optional
|
|
Explicit shared memory size in bytes. When *None* (the default),
|
|
the high-water mark ``max_offset`` tracked by the allocator is used.
|
|
"""
|
|
if not self._owns_buffer:
|
|
return
|
|
ir = _get_ir()
|
|
frame_mod = _get_frame()
|
|
resolved = size if size is not None else self.max_offset
|
|
assert resolved >= self.max_offset, (
|
|
f"Specified smem size ({resolved}) is smaller than "
|
|
f"the pool high-water mark ({self.max_offset})"
|
|
)
|
|
attr_frame = ir.attr(self.ptr, "tirx.pool_max_bytes", resolved)
|
|
if isinstance(attr_frame, frame_mod.AttrFrame):
|
|
from functools import partial
|
|
|
|
attr_frame.add_callback(partial(attr_frame.__exit__, None, None, None))
|
|
attr_frame.__enter__()
|