wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tvm/backend/cuda/lang/alloc_pool.py
+529
View File
@@ -0,0 +1,529 @@
# 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__()