wehub-resource-sync
946 lines
36 KiB
Python
946 lines
36 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.
|
|
"""Reusable tile scheduler helpers for TIR tests/kernels.
|
|
|
|
These classes emit TIR via @T.inline. Decorate with @T.meta_class so that
|
|
instances are automatically treated as meta values inside @T.prim_func.
|
|
"""
|
|
|
|
from tvm.backend.cuda.lang.pipeline import Pipeline, PipelineState
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
@T.meta_class
|
|
class BaseTileScheduler:
|
|
"""Base class for tile schedulers with common state and macros."""
|
|
|
|
def __init__(self, prefix: str):
|
|
self.m_idx = T.local_scalar("int32")
|
|
self.n_idx = T.local_scalar("int32")
|
|
self.linear_idx = T.local_scalar("int32")
|
|
|
|
@T.inline
|
|
def update_current_m_n_idx(self, linear_idx):
|
|
# To be implemented by subclasses
|
|
pass
|
|
|
|
@T.inline
|
|
def init(self, linear_init):
|
|
self.linear_idx = linear_init
|
|
self.update_current_m_n_idx(linear_init)
|
|
|
|
@T.inline
|
|
def next_tile(self, step):
|
|
self.linear_idx = self.linear_idx + step
|
|
self.update_current_m_n_idx(self.linear_idx)
|
|
|
|
def valid(self, total_tiles):
|
|
return self.linear_idx < total_tiles
|
|
|
|
|
|
class ClusterPersistentScheduler2D(BaseTileScheduler):
|
|
"""
|
|
Tile scheduler for cluster-based persistent kernels.
|
|
|
|
Distributes a 2D tile grid across persistent clusters using group-major ordering
|
|
for L2 cache locality. Each cluster starts at its cluster_id and strides by
|
|
num_clusters to process tiles.
|
|
|
|
Tile Ordering (group-major for L2 locality):
|
|
- Tiles are grouped into "L2 groups" of `l2_group_size` rows
|
|
- Within a group, tiles are visited in column-major order within the group
|
|
- Groups are processed in row-major order
|
|
|
|
Example with 4x4 tiles, l2_group_size=2:
|
|
Group 0 (rows 0-1): 0 2 4 6
|
|
1 3 5 7
|
|
Group 1 (rows 2-3): 8 10 12 14
|
|
9 11 13 15
|
|
|
|
Serpentine Mode (serpentine=True):
|
|
- Uses CUTLASS-style 2D block swizzle with serpentine traversal
|
|
- Grid is divided into swizzle_size x swizzle_size blocks
|
|
- Within each block, tiles are visited in row-major order
|
|
- Blocks are traversed in serpentine order (even block-rows forward, odd backward)
|
|
- This provides better L2 locality by reusing both A and B tiles
|
|
|
|
Example with 4x4 tiles, swizzle_size=2, serpentine=True:
|
|
Block layout:
|
|
Block(0,0) Block(0,1)
|
|
Block(1,0) Block(1,1)
|
|
|
|
Tile numbering with serpentine:
|
|
n=0 n=1 n=2 n=3
|
|
m=0 0 1 14 15
|
|
m=1 2 3 12 13
|
|
m=2 4 5 10 11
|
|
m=3 6 7 8 9
|
|
|
|
Traversal: Block(0,0) -> Block(1,0) -> Block(1,1) -> Block(0,1)
|
|
(serpentine: down in col 0, then up in col 1)
|
|
|
|
Parameters
|
|
----------
|
|
prefix : str
|
|
Prefix for TIR variable names
|
|
num_m_tiles : int | T.ExprLike
|
|
Total number of tiles in M dimension (can be runtime expression)
|
|
num_n_tiles : int
|
|
Total number of tiles in N dimension
|
|
num_clusters : int
|
|
Number of persistent clusters (determines stride)
|
|
l2_group_size : int
|
|
Number of M-tile rows per L2 locality group (default: 8)
|
|
When serpentine=True, this is used as swizzle_size for 2D blocks
|
|
cluster_m : int
|
|
Cluster dimension in M for hierarchical scheduling (default: 1)
|
|
cluster_n : int
|
|
Cluster dimension in N for hierarchical scheduling (default: 1)
|
|
serpentine : bool
|
|
If True, use CUTLASS-style 2D block swizzle with serpentine traversal (default: False)
|
|
|
|
Attributes
|
|
----------
|
|
m_idx : T.local_scalar
|
|
Current M tile index (output)
|
|
n_idx : T.local_scalar
|
|
Current N tile index (output)
|
|
work_idx : T.local_scalar
|
|
Global work item index for this cluster
|
|
tile_count : T.local_scalar
|
|
Number of tiles processed by this cluster so far
|
|
|
|
Usage
|
|
-----
|
|
```python
|
|
scheduler = ClusterPersistentScheduler2D(
|
|
"sched", num_m_tiles=M_TILES, num_n_tiles=N_TILES,
|
|
num_clusters=NUM_CLUSTERS, l2_group_size=8
|
|
)
|
|
scheduler.init(cluster_id) # cluster_id = cta_idx // CLUSTER_SIZE
|
|
|
|
while scheduler.valid():
|
|
m = T.meta_var(scheduler.m_idx) # current M tile
|
|
n = T.meta_var(scheduler.n_idx) # current N tile
|
|
# ... process tile (m, n) ...
|
|
scheduler.next_tile()
|
|
```
|
|
|
|
Examples
|
|
--------
|
|
Example 1: Basic persistent kernel
|
|
```
|
|
num_m_tiles=4, num_n_tiles=4, num_clusters=3, l2_group_size=2
|
|
cluster_m=1, cluster_n=1 (default, no tile subdivision)
|
|
|
|
Group-major tile numbering (l2_group_size=2):
|
|
n=0 n=1 n=2 n=3
|
|
m=0 0 2 4 6 ┐ L2 group 0
|
|
m=1 1 3 5 7 ┘
|
|
m=2 8 10 12 14 ┐ L2 group 1
|
|
m=3 9 11 13 15 ┘
|
|
|
|
Work distribution (cluster starts at cluster_id, strides by num_clusters=3):
|
|
cluster 0: work_idx 0,3,6,9,12,15 -> tiles 0,3,6,9,12,15
|
|
cluster 1: work_idx 1,4,7,10,13 -> tiles 1,4,7,10,13
|
|
cluster 2: work_idx 2,5,8,11,14 -> tiles 2,5,8,11,14
|
|
|
|
Tile grid (which cluster handles each tile):
|
|
n=0 n=1 n=2 n=3
|
|
m=0 C0 C2 C1 C0 ┐ L2 group 0
|
|
m=1 C1 C0 C2 C1 ┘
|
|
m=2 C2 C1 C0 C2 ┐ L2 group 1
|
|
m=3 C0 C2 C1 C0 ┘
|
|
|
|
Tile sequence per cluster (in execution order):
|
|
cluster 0: (0,0)->(1,1)->(0,3)->(2,0)->(2,3)->(3,3)
|
|
cluster 1: (1,0)->(0,2)->(1,3)->(2,1)->(3,2)
|
|
cluster 2: (0,1)->(1,2)->(2,0)->(3,1)->(2,3)
|
|
```
|
|
|
|
Example 2: 2SM GEMM (typical B200 config)
|
|
```
|
|
M=1024, N=512, CTA_M=128, MMA_N=128, CLUSTER_M=2, CLUSTER_N=1
|
|
=> M_TILES=8, N_TILES=4
|
|
=> CLUSTER_M_TILES=4, CLUSTER_N_TILES=4 (scheduler at cluster granularity)
|
|
|
|
Scheduler params:
|
|
num_m_tiles=4, num_n_tiles=4, num_clusters=74, l2_group_size=8
|
|
cluster_m=1, cluster_n=1
|
|
|
|
Key: Scheduler outputs CLUSTER-level tiles.
|
|
All CTAs in same cluster get SAME (m_idx, n_idx) from scheduler.
|
|
CTAs differentiate via cluster_rank (computed OUTSIDE scheduler):
|
|
cluster_rank = cta_idx % CLUSTER_SIZE
|
|
cb_m = cluster_rank % CLUSTER_M # 0 or 1 for 2SM
|
|
cb_n = cluster_rank // CLUSTER_M # 0 for 2SM
|
|
|
|
Final CTA tile:
|
|
cta_m = m_idx * CLUSTER_M + cb_m
|
|
cta_n = n_idx * CLUSTER_N + cb_n
|
|
|
|
Example: cluster 5 gets scheduler tile (1,2)
|
|
CTA rank=0 (cb_m=0): actual tile (2,2)
|
|
CTA rank=1 (cb_m=1): actual tile (3,2)
|
|
```
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
prefix: str,
|
|
num_m_tiles,
|
|
num_n_tiles: int,
|
|
num_clusters: int,
|
|
l2_group_size: int = 8,
|
|
cluster_m: int = 1,
|
|
cluster_n: int = 1,
|
|
serpentine: bool = False,
|
|
):
|
|
super().__init__(prefix)
|
|
self._num_m_tiles = num_m_tiles
|
|
self._num_n_tiles = num_n_tiles
|
|
self._num_clusters = num_clusters
|
|
self._l2_group_size = l2_group_size
|
|
self._cluster_m = cluster_m
|
|
self._cluster_n = cluster_n
|
|
self._serpentine = serpentine
|
|
|
|
# Rename internal state for clarity
|
|
self.work_idx = self.linear_idx # alias: global work item index
|
|
self.tile_count = T.local_scalar("int32")
|
|
self.tile_idx = self.tile_count # alias for backward compatibility
|
|
|
|
is_static_m = isinstance(num_m_tiles, int)
|
|
|
|
# Number of tile columns after accounting for cluster_n
|
|
n_tile_cols = (num_n_tiles + cluster_n - 1) // cluster_n
|
|
self._N_TILE_COLS = n_tile_cols
|
|
|
|
if is_static_m:
|
|
self._M_TILE_ROWS = (num_m_tiles + cluster_m - 1) // cluster_m
|
|
self._FULL_GROUPS = self._M_TILE_ROWS // l2_group_size
|
|
else:
|
|
# Dynamic expressions for runtime M
|
|
self._M_TILE_ROWS = T.truncdiv(self._num_m_tiles + self._cluster_m - 1, self._cluster_m)
|
|
self._FULL_GROUPS = T.truncdiv(self._M_TILE_ROWS, self._l2_group_size)
|
|
|
|
self._TAIL_ROWS = self._M_TILE_ROWS - self._FULL_GROUPS * l2_group_size
|
|
self._TOTAL_TILES = self._M_TILE_ROWS * n_tile_cols * cluster_m * cluster_n
|
|
|
|
# For serpentine mode: precompute block counts
|
|
if serpentine:
|
|
self._N_BLOCKS = n_tile_cols // l2_group_size # full blocks in N
|
|
self._M_BLOCKS = (
|
|
self._M_TILE_ROWS // l2_group_size
|
|
if is_static_m
|
|
else T.truncdiv(self._M_TILE_ROWS, l2_group_size)
|
|
)
|
|
self._BLOCK_SIZE = l2_group_size * l2_group_size # tiles per block
|
|
self._FULL_BLOCK_TILES = self._M_BLOCKS * self._N_BLOCKS * self._BLOCK_SIZE
|
|
# Residual tiles (not covered by full blocks)
|
|
self._RESIDUAL_N = n_tile_cols - self._N_BLOCKS * l2_group_size
|
|
self._RESIDUAL_M = self._M_TILE_ROWS - self._M_BLOCKS * l2_group_size
|
|
|
|
# fmt: off
|
|
@T.inline
|
|
def update_current_m_n_idx(self, work_idx):
|
|
"""Convert global work index to (m_idx, n_idx) tile coordinates."""
|
|
CLUSTER_M = T.meta_var(self._cluster_m)
|
|
CLUSTER_N = T.meta_var(self._cluster_n)
|
|
|
|
# Extract hierarchical cluster-local offsets
|
|
cluster_m_offset = T.meta_var(work_idx % CLUSTER_M)
|
|
t = T.meta_var(work_idx // CLUSTER_M)
|
|
cluster_n_offset = T.meta_var(t % CLUSTER_N)
|
|
tile_linear = T.meta_var(t // CLUSTER_N)
|
|
|
|
@T.inline
|
|
def set_tile_coords(tile_row, tile_col):
|
|
self.m_idx = tile_row * CLUSTER_M + cluster_m_offset
|
|
self.n_idx = tile_col * CLUSTER_N + cluster_n_offset
|
|
|
|
if self._serpentine:
|
|
self._update_serpentine(tile_linear, set_tile_coords)
|
|
else:
|
|
self._update_group_major(tile_linear, set_tile_coords)
|
|
|
|
def _update_group_major(self, tile_linear, set_tile_coords):
|
|
"""Group-major ordering with parse-time pruning of statically-dead branches.
|
|
|
|
The TIR script parser does not constant-fold ``if False: ...``, so a
|
|
Python-literal ``FULL_GROUPS == 0`` would otherwise produce
|
|
``T.bitwise_and(T.bool(False), tile_linear < 0)`` IR plus the dead
|
|
then-leg. Branch in plain Python here and only invoke the inline
|
|
emitter that can actually fire.
|
|
"""
|
|
full_zero = isinstance(self._FULL_GROUPS, int) and self._FULL_GROUPS == 0
|
|
tail_zero = isinstance(self._TAIL_ROWS, int) and self._TAIL_ROWS == 0
|
|
if full_zero and tail_zero:
|
|
self._gm_emit_zero(set_tile_coords)
|
|
elif full_zero:
|
|
self._gm_emit_tail_only(tile_linear, set_tile_coords)
|
|
elif tail_zero:
|
|
self._gm_emit_full_only(tile_linear, set_tile_coords)
|
|
else:
|
|
self._gm_emit_full_and_tail(tile_linear, set_tile_coords)
|
|
|
|
@T.inline
|
|
def _gm_emit_zero(self, set_tile_coords):
|
|
set_tile_coords(0, 0)
|
|
|
|
@T.inline
|
|
def _gm_emit_full_only(self, tile_linear, set_tile_coords):
|
|
FULL_GROUPS = T.meta_var(self._FULL_GROUPS)
|
|
GROUP_SIZE = T.meta_var(self._l2_group_size)
|
|
GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS)
|
|
if (FULL_GROUPS > 0) & (tile_linear < FULL_GROUPS * GROUP_SPAN):
|
|
group_id: T.let = tile_linear // GROUP_SPAN
|
|
within_group: T.let = tile_linear % GROUP_SPAN
|
|
tile_row: T.let = group_id * GROUP_SIZE + (within_group % GROUP_SIZE)
|
|
tile_col: T.let = within_group // GROUP_SIZE
|
|
set_tile_coords(tile_row, tile_col)
|
|
else:
|
|
set_tile_coords(0, 0)
|
|
|
|
@T.inline
|
|
def _gm_emit_tail_only(self, tile_linear, set_tile_coords):
|
|
FULL_GROUPS = T.meta_var(self._FULL_GROUPS)
|
|
TAIL_ROWS = T.meta_var(self._TAIL_ROWS)
|
|
GROUP_SIZE = T.meta_var(self._l2_group_size)
|
|
GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS)
|
|
if TAIL_ROWS > 0:
|
|
rem: T.let = tile_linear - FULL_GROUPS * GROUP_SPAN
|
|
tile_row: T.let = FULL_GROUPS * GROUP_SIZE + (rem % TAIL_ROWS)
|
|
tile_col: T.let = rem // TAIL_ROWS
|
|
set_tile_coords(tile_row, tile_col)
|
|
else:
|
|
set_tile_coords(0, 0)
|
|
|
|
@T.inline
|
|
def _gm_emit_full_and_tail(self, tile_linear, set_tile_coords):
|
|
FULL_GROUPS = T.meta_var(self._FULL_GROUPS)
|
|
TAIL_ROWS = T.meta_var(self._TAIL_ROWS)
|
|
GROUP_SIZE = T.meta_var(self._l2_group_size)
|
|
GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS)
|
|
if (FULL_GROUPS > 0) & (tile_linear < FULL_GROUPS * GROUP_SPAN):
|
|
group_id: T.let = tile_linear // GROUP_SPAN
|
|
within_group: T.let = tile_linear % GROUP_SPAN
|
|
tile_row: T.let = group_id * GROUP_SIZE + (within_group % GROUP_SIZE)
|
|
tile_col: T.let = within_group // GROUP_SIZE
|
|
set_tile_coords(tile_row, tile_col)
|
|
elif TAIL_ROWS > 0:
|
|
rem: T.let = tile_linear - FULL_GROUPS * GROUP_SPAN
|
|
tile_row: T.let = FULL_GROUPS * GROUP_SIZE + (rem % TAIL_ROWS)
|
|
tile_col: T.let = rem // TAIL_ROWS
|
|
set_tile_coords(tile_row, tile_col)
|
|
else:
|
|
set_tile_coords(0, 0)
|
|
|
|
@T.inline
|
|
def _update_serpentine(self, tile_linear, set_tile_coords):
|
|
"""CUTLASS-style 2D block swizzle with serpentine traversal.
|
|
|
|
Algorithm:
|
|
1. Divide grid into swizzle_size x swizzle_size blocks
|
|
2. Within each block, visit tiles in row-major order
|
|
3. Blocks are traversed column by column (along N)
|
|
4. Within each column of blocks, use serpentine:
|
|
- Even columns: top to bottom
|
|
- Odd columns: bottom to top
|
|
|
|
This maximizes L2 reuse for both A and B matrices.
|
|
"""
|
|
S = T.meta_var(self._l2_group_size) # swizzle_size
|
|
M_BLOCKS = T.meta_var(self._M_BLOCKS)
|
|
N_BLOCKS = T.meta_var(self._N_BLOCKS)
|
|
BLOCK_SIZE = T.meta_var(self._BLOCK_SIZE) # S * S
|
|
FULL_BLOCK_TILES = T.meta_var(self._FULL_BLOCK_TILES)
|
|
M_TILE_ROWS = T.meta_var(self._M_TILE_ROWS)
|
|
T.meta_var(self._N_TILE_COLS)
|
|
RESIDUAL_N = T.meta_var(self._RESIDUAL_N)
|
|
RESIDUAL_M = T.meta_var(self._RESIDUAL_M)
|
|
|
|
# Check if we're in the full block region
|
|
if (M_BLOCKS > 0) & (N_BLOCKS > 0) & (tile_linear < FULL_BLOCK_TILES):
|
|
# Which block (in linear order along columns of blocks)
|
|
block_linear: T.let = tile_linear // BLOCK_SIZE
|
|
within_block: T.let = tile_linear % BLOCK_SIZE
|
|
|
|
# Block column and row
|
|
block_col: T.let = block_linear // M_BLOCKS
|
|
block_row_raw: T.let = block_linear % M_BLOCKS
|
|
|
|
# Serpentine: odd columns go bottom-to-top
|
|
block_row: T.let = T.Select(
|
|
block_col % 2 == 0,
|
|
block_row_raw,
|
|
M_BLOCKS - 1 - block_row_raw
|
|
)
|
|
|
|
# Position within block (row-major within block)
|
|
local_row: T.let = within_block // S
|
|
local_col: T.let = within_block % S
|
|
|
|
tile_row: T.let = block_row * S + local_row
|
|
tile_col: T.let = block_col * S + local_col
|
|
set_tile_coords(tile_row, tile_col)
|
|
|
|
elif RESIDUAL_N > 0:
|
|
# Residual tiles in the rightmost partial column of blocks
|
|
# These are tiles where n >= N_BLOCKS * S
|
|
rem: T.let = tile_linear - FULL_BLOCK_TILES
|
|
|
|
# First handle the right residual strip (full M height, partial N width)
|
|
right_strip_tiles: T.let = M_TILE_ROWS * RESIDUAL_N
|
|
if rem < right_strip_tiles:
|
|
# Row-major within the right strip
|
|
tile_row: T.let = rem // RESIDUAL_N
|
|
tile_col: T.let = N_BLOCKS * S + (rem % RESIDUAL_N)
|
|
set_tile_coords(tile_row, tile_col)
|
|
elif RESIDUAL_M > 0:
|
|
# Bottom residual strip (already covered in right strip overlap)
|
|
# This handles corner case - shouldn't normally reach here
|
|
# as right strip already covers full M height
|
|
set_tile_coords(0, 0)
|
|
else:
|
|
set_tile_coords(0, 0)
|
|
|
|
elif RESIDUAL_M > 0:
|
|
# Bottom residual strip only (no right residual)
|
|
rem: T.let = tile_linear - FULL_BLOCK_TILES
|
|
bottom_strip_tiles: T.let = RESIDUAL_M * (N_BLOCKS * S)
|
|
if rem < bottom_strip_tiles:
|
|
tile_row: T.let = M_BLOCKS * S + (rem % RESIDUAL_M)
|
|
tile_col: T.let = rem // RESIDUAL_M
|
|
set_tile_coords(tile_row, tile_col)
|
|
else:
|
|
set_tile_coords(0, 0)
|
|
else:
|
|
# Fallback
|
|
set_tile_coords(0, 0)
|
|
|
|
@T.inline
|
|
def init(self, cluster_id):
|
|
"""Initialize scheduler for a given cluster.
|
|
|
|
Parameters
|
|
----------
|
|
cluster_id : int
|
|
The cluster's index (typically cta_idx // CLUSTER_SIZE)
|
|
"""
|
|
self.linear_idx = cluster_id
|
|
self.tile_count = 0
|
|
self.update_current_m_n_idx(cluster_id)
|
|
|
|
@T.inline
|
|
def next_tile(self):
|
|
"""Advance to the next tile for this cluster."""
|
|
self.linear_idx = self.linear_idx + self._num_clusters
|
|
self.tile_count = self.tile_count + 1
|
|
self.update_current_m_n_idx(self.linear_idx)
|
|
|
|
@T.inline
|
|
def next_tile_stride(self, stride: int):
|
|
"""Advance by a custom stride (for non-standard scheduling)."""
|
|
self.linear_idx = self.linear_idx + stride
|
|
self.tile_count = self.tile_count + 1
|
|
self.update_current_m_n_idx(self.linear_idx)
|
|
# fmt: on
|
|
|
|
def valid(self):
|
|
"""Check if this cluster has more tiles to process."""
|
|
return self.linear_idx < self._TOTAL_TILES
|
|
|
|
|
|
class GroupMajor3D(BaseTileScheduler):
|
|
"""
|
|
3D grouped-row scheduler (M,N,K) with tail handling on M.
|
|
|
|
Args
|
|
----
|
|
prefix: str
|
|
m_tiles: int | T Expr # tiles along M (static or runtime)
|
|
n_tiles: int # tiles along N (static)
|
|
k_tiles: int # tiles along K (static)
|
|
group_rows: int # rows per group along M
|
|
step: int = 1 # default stride for next_tile()
|
|
"""
|
|
|
|
def __init__(
|
|
self, prefix: str, m_tiles, n_tiles: int, k_tiles: int, group_rows: int, step: int = 1
|
|
):
|
|
super().__init__(prefix)
|
|
self._step = step
|
|
self.tile_idx = T.local_scalar("int32")
|
|
self.k_idx = T.local_scalar("int32")
|
|
|
|
# ---- constants / primexprs baked once ----
|
|
self._G = group_rows
|
|
self._N = n_tiles
|
|
self._K = k_tiles
|
|
|
|
if isinstance(m_tiles, int):
|
|
self._GROUPS = m_tiles // group_rows
|
|
self._FINAL_ROWS = m_tiles - self._GROUPS * group_rows
|
|
self._SAFE_FINAL_ROWS = max(self._FINAL_ROWS, 1)
|
|
self._GROUP_SIZE = group_rows * n_tiles * k_tiles
|
|
self._TOTAL = m_tiles * n_tiles * k_tiles
|
|
else:
|
|
self._GROUPS = T.truncdiv(m_tiles, group_rows)
|
|
self._FINAL_ROWS = m_tiles - self._GROUPS * group_rows
|
|
self._SAFE_FINAL_ROWS = T.max(self._FINAL_ROWS, 1)
|
|
self._GROUP_SIZE = self._G * self._N * self._K
|
|
self._TOTAL = m_tiles * n_tiles * k_tiles
|
|
|
|
# handy composites used in macro
|
|
self._FULL_BOUND = self._GROUPS * self._GROUP_SIZE
|
|
self._HAS_FULL = self._GROUPS > 0
|
|
self._HAS_TAIL = self._FINAL_ROWS > 0
|
|
|
|
# fmt: off
|
|
@T.inline
|
|
def update_current_m_n_idx(self, linear_idx):
|
|
# full-group formulas
|
|
full_m: T.let = T.floordiv(linear_idx, self._GROUP_SIZE) * self._G + T.floormod(
|
|
linear_idx, self._G
|
|
)
|
|
full_n: T.let = T.floormod(T.floordiv(linear_idx, self._G), self._N)
|
|
full_k: T.let = T.floordiv(T.floormod(linear_idx, self._GROUP_SIZE), self._G * self._N)
|
|
|
|
# tail formulas (relative to FULL_BOUND)
|
|
# Use _SAFE_FINAL_ROWS (max(FINAL_ROWS, 1)) to avoid divide-by-zero when there is no tail
|
|
rem: T.let = linear_idx - self._FULL_BOUND
|
|
tail_m: T.let = self._GROUPS * self._G + T.floormod(rem, self._SAFE_FINAL_ROWS)
|
|
tail_n: T.let = T.floordiv(rem, self._SAFE_FINAL_ROWS) % self._N
|
|
tail_k: T.let = T.floordiv(rem, self._SAFE_FINAL_ROWS * self._N)
|
|
|
|
# choose phase
|
|
if self._HAS_FULL & (linear_idx < self._FULL_BOUND):
|
|
self.m_idx = full_m
|
|
self.n_idx = full_n
|
|
self.k_idx = full_k
|
|
elif self._HAS_TAIL:
|
|
self.m_idx = tail_m
|
|
self.n_idx = tail_n
|
|
self.k_idx = tail_k
|
|
else:
|
|
self.m_idx = 0
|
|
self.n_idx = 0
|
|
self.k_idx = 0
|
|
|
|
@T.inline
|
|
def init(self, linear_init):
|
|
self.linear_idx = linear_init
|
|
self.tile_idx = 0
|
|
self.update_current_m_n_idx(linear_init)
|
|
|
|
@T.inline
|
|
def next_tile(self):
|
|
self.linear_idx = self.linear_idx + self._step
|
|
self.tile_idx = self.tile_idx + 1
|
|
self.update_current_m_n_idx(self.linear_idx)
|
|
|
|
@T.inline
|
|
def next_tile_stride(self, stride: int):
|
|
self.linear_idx = self.linear_idx + stride
|
|
self.tile_idx = self.tile_idx + 1
|
|
self.update_current_m_n_idx(self.linear_idx)
|
|
# fmt: on
|
|
|
|
def valid(self):
|
|
return self.linear_idx < self._TOTAL
|
|
|
|
|
|
class RankAwareGroupMajorTileScheduler(BaseTileScheduler):
|
|
"""
|
|
Group-major scheduler that applies a rank-aware remapping (remote rows first).
|
|
Kept as a thin adapter because it depends on NVSHMEM rank at device-side.
|
|
"""
|
|
|
|
def __init__(
|
|
self, prefix: str, m_clusters: int, n_clusters: int, group_size: int, world_size: int
|
|
):
|
|
super().__init__(prefix)
|
|
self._m_clusters = m_clusters
|
|
self._n_clusters = n_clusters
|
|
self._group_size = group_size
|
|
self._world_size = world_size
|
|
|
|
@T.inline
|
|
def update_current_m_n_idx(self, linear_idx):
|
|
my_rank: T.let = T.nvshmem.my_pe()
|
|
remote_m_clusters: T.let = self._m_clusters - self._m_clusters // self._world_size
|
|
group_rows: T.let = (remote_m_clusters // self._group_size) * self._group_size
|
|
final_rows: T.let = remote_m_clusters - group_rows
|
|
group_repeat: T.let = self._group_size * self._n_clusters
|
|
if linear_idx < group_rows * self._n_clusters and group_rows > 0:
|
|
self.m_idx = (
|
|
(linear_idx // group_repeat) * self._group_size
|
|
+ (linear_idx % self._group_size)
|
|
+ (my_rank + 1) * self._m_clusters // self._world_size
|
|
) % self._m_clusters
|
|
self.n_idx = (linear_idx % group_repeat) // self._group_size
|
|
elif linear_idx < remote_m_clusters * self._n_clusters:
|
|
remainder_idx: T.let = linear_idx - group_rows * self._n_clusters
|
|
self.m_idx = (
|
|
group_rows
|
|
+ remainder_idx % final_rows
|
|
+ (my_rank + 1) * self._m_clusters // self._world_size
|
|
) % self._m_clusters
|
|
self.n_idx = remainder_idx // final_rows
|
|
else:
|
|
remainder_idx: T.let = linear_idx - remote_m_clusters * self._n_clusters
|
|
self.m_idx = (
|
|
remote_m_clusters
|
|
+ remainder_idx % (self._m_clusters // self._world_size)
|
|
+ (my_rank + 1) * self._m_clusters // self._world_size
|
|
) % self._m_clusters
|
|
self.n_idx = remainder_idx // (self._m_clusters // self._world_size)
|
|
|
|
@T.inline
|
|
def next_tile(self, stride: int):
|
|
self.linear_idx = self.linear_idx + stride
|
|
self.update_current_m_n_idx(self.linear_idx)
|
|
|
|
def valid(self):
|
|
return self.linear_idx < self._m_clusters * self._n_clusters
|
|
|
|
|
|
class IndexedTripleTileScheduler(BaseTileScheduler):
|
|
"""Scheduler that maps linear_idx to (b_idx, h_idx, q_idx) via index lists."""
|
|
|
|
def __init__(self, prefix: str, b_indices, h_indices, q_indices, tiles_indptr):
|
|
super().__init__(prefix)
|
|
self.b_indices = b_indices
|
|
self.h_indices = h_indices
|
|
self.q_indices = q_indices
|
|
self.tiles_indptr = tiles_indptr
|
|
self.q_idx = T.local_scalar("int32")
|
|
self.h_idx = T.local_scalar("int32")
|
|
self.b_idx = T.local_scalar("int32")
|
|
self.linear_lim = T.local_scalar("int32")
|
|
|
|
@T.inline
|
|
def _load(self):
|
|
self.q_idx = self.q_indices[self.linear_idx]
|
|
self.h_idx = self.h_indices[self.linear_idx]
|
|
self.b_idx = self.b_indices[self.linear_idx]
|
|
|
|
@T.inline
|
|
def init(self, sm):
|
|
self.linear_idx = self.tiles_indptr[sm]
|
|
self.linear_lim = self.tiles_indptr[sm + 1]
|
|
self._load()
|
|
|
|
@T.inline
|
|
def next_tile(self):
|
|
self.linear_idx = self.linear_idx + 1
|
|
self._load()
|
|
|
|
def valid(self):
|
|
return self.linear_idx < self.linear_lim
|
|
|
|
|
|
class FlashAttentionLinearScheduler(BaseTileScheduler):
|
|
"""Linear 3D scheduler for flash attention (batch, head, m_block).
|
|
|
|
Used for non-causal attention with simple linear decomposition.
|
|
Maps linear_idx -> (batch_idx, head_idx, m_block_idx) using:
|
|
batch = linear_idx // (num_heads * num_m_blocks)
|
|
head = (linear_idx % (num_heads * num_m_blocks)) // num_m_blocks
|
|
m_block = linear_idx % num_m_blocks
|
|
|
|
Parameters
|
|
----------
|
|
prefix : str
|
|
Prefix for TIR variable names
|
|
num_batches : int
|
|
Number of batches
|
|
num_heads : int
|
|
Number of KV heads
|
|
num_m_blocks : int
|
|
Number of Q blocks (M dimension tiles)
|
|
num_ctas : int
|
|
Number of CTAs for persistent kernel stride
|
|
"""
|
|
|
|
def __init__(
|
|
self, prefix: str, num_batches: int, num_heads: int, num_m_blocks: int, num_ctas: int
|
|
):
|
|
super().__init__(prefix)
|
|
self._num_batches = num_batches
|
|
self._num_heads = num_heads
|
|
self._num_m_blocks = num_m_blocks
|
|
self._num_ctas = num_ctas
|
|
self._total_tasks = num_batches * num_heads * num_m_blocks
|
|
|
|
# Output indices
|
|
self.batch_idx = T.local_scalar("int32")
|
|
self.head_idx = T.local_scalar("int32")
|
|
self.m_block_idx = T.local_scalar("int32")
|
|
|
|
# fmt: off
|
|
@T.inline
|
|
def update_current_m_n_idx(self, linear_idx):
|
|
"""Convert linear index to (batch, head, m_block) coordinates."""
|
|
NUM_HEADS = T.meta_var(self._num_heads)
|
|
NUM_M_BLOCKS = T.meta_var(self._num_m_blocks)
|
|
HEAD_M_PRODUCT = T.meta_var(NUM_HEADS * NUM_M_BLOCKS)
|
|
|
|
self.batch_idx = linear_idx // HEAD_M_PRODUCT
|
|
self.head_idx = (linear_idx % HEAD_M_PRODUCT) // NUM_M_BLOCKS
|
|
self.m_block_idx = linear_idx % NUM_M_BLOCKS
|
|
|
|
@T.inline
|
|
def init(self, cta_id):
|
|
"""Initialize scheduler with CTA ID."""
|
|
self.linear_idx = cta_id
|
|
self.update_current_m_n_idx(cta_id)
|
|
|
|
@T.inline
|
|
def next_tile(self):
|
|
"""Advance to next tile by striding by num_ctas."""
|
|
self.linear_idx = self.linear_idx + self._num_ctas
|
|
self.update_current_m_n_idx(self.linear_idx)
|
|
# fmt: on
|
|
|
|
def valid(self):
|
|
"""Check if there are more tiles to process."""
|
|
return self.linear_idx < self._total_tasks
|
|
|
|
|
|
class FlashAttentionLPTScheduler(BaseTileScheduler):
|
|
"""LPT scheduler with L2 swizzle for causal flash attention.
|
|
|
|
Processes high-work Q blocks (with more KV blocks to attend to) first using
|
|
Longest Processing Time (LPT) scheduling. Also applies L2 cache swizzle
|
|
for better cache locality across batch*head dimensions.
|
|
|
|
The LPT aspect comes from reversing m_block order: lower Q blocks have more
|
|
KV blocks to process due to causal masking, so processing them first balances load.
|
|
|
|
The scheduler is only applied to non-persistent kernels.
|
|
|
|
L2 Swizzle: Groups consecutive batch*head indices together for L2 locality.
|
|
|
|
Parameters
|
|
----------
|
|
prefix : str
|
|
Prefix for TIR variable names
|
|
num_batches : int
|
|
Number of batches
|
|
num_heads : int
|
|
Number of KV heads
|
|
num_m_blocks : int
|
|
Number of Q blocks (M dimension tiles)
|
|
num_ctas : int
|
|
Number of CTAs (should equal total_tasks for causal)
|
|
l2_swizzle : int
|
|
L2 swizzle factor for cache locality
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
prefix: str,
|
|
num_batches: int,
|
|
num_heads: int,
|
|
num_m_blocks: int,
|
|
l2_swizzle: int,
|
|
num_ctas: int | None = None,
|
|
):
|
|
super().__init__(prefix)
|
|
self._num_batches = num_batches
|
|
self._num_heads = num_heads
|
|
self._num_m_blocks = num_m_blocks
|
|
self._l2_swizzle = l2_swizzle
|
|
self._num_ctas = num_ctas
|
|
self._total_tasks = num_batches * num_heads * num_m_blocks
|
|
|
|
# Derived constants for L2 swizzle
|
|
self._num_hb = num_batches * num_heads
|
|
self._l2_major = l2_swizzle * num_m_blocks
|
|
self._num_hb_quotient = self._num_hb // l2_swizzle
|
|
|
|
# Output indices
|
|
self.batch_idx = T.local_scalar("int32")
|
|
self.head_idx = T.local_scalar("int32")
|
|
self.m_block_idx = T.local_scalar("int32")
|
|
|
|
# fmt: off
|
|
@T.inline
|
|
def update_current_m_n_idx(self, linear_idx):
|
|
"""Convert linear index to (batch, head, m_block) with LPT + L2 swizzle."""
|
|
L2_SWIZZLE = T.meta_var(self._l2_swizzle)
|
|
L2_MAJOR = T.meta_var(self._l2_major)
|
|
NUM_HB_QUOTIENT = T.meta_var(self._num_hb_quotient)
|
|
NUM_HB = T.meta_var(self._num_hb)
|
|
NUM_HEADS = T.meta_var(self._num_heads)
|
|
NUM_M_BLOCKS = T.meta_var(self._num_m_blocks)
|
|
|
|
# L2 swizzle decomposition
|
|
bidhb: T.let = linear_idx // L2_MAJOR
|
|
l2_mod: T.let = linear_idx % L2_MAJOR
|
|
|
|
# Handle residual section (last partial swizzle group)
|
|
num_hb_remainder: T.let = T.max(NUM_HB % L2_SWIZZLE, 1)
|
|
m_block_raw: T.let = T.Select(bidhb < NUM_HB_QUOTIENT, l2_mod // L2_SWIZZLE, l2_mod // num_hb_remainder) # noqa: E501
|
|
bidhb_residual: T.let = T.Select(bidhb < NUM_HB_QUOTIENT, l2_mod % L2_SWIZZLE, l2_mod % num_hb_remainder) # noqa: E501
|
|
bidhb_actual: T.let = bidhb * L2_SWIZZLE + bidhb_residual
|
|
|
|
self.batch_idx = bidhb_actual // NUM_HEADS
|
|
self.head_idx = bidhb_actual % NUM_HEADS
|
|
|
|
# LPT: Reverse block order so high-work blocks are processed first
|
|
self.m_block_idx = (NUM_M_BLOCKS - 1) - m_block_raw
|
|
|
|
@T.inline
|
|
def init(self, cta_id):
|
|
"""Initialize scheduler with CTA ID."""
|
|
self.linear_idx = cta_id
|
|
self.update_current_m_n_idx(cta_id)
|
|
|
|
@T.inline
|
|
def next_tile(self):
|
|
"""Advance to the next tile.
|
|
|
|
Single-tile mode (``num_ctas=None``, the default): each CTA owns one
|
|
task; terminate. Persistent mode (``num_ctas=N``): stride by N, like
|
|
:class:`FlashAttentionLinearScheduler`, while keeping the LPT + L2
|
|
swizzle index mapping.
|
|
"""
|
|
if self._num_ctas is None:
|
|
self.linear_idx = self._total_tasks
|
|
else:
|
|
self.linear_idx = self.linear_idx + self._num_ctas
|
|
self.update_current_m_n_idx(self.linear_idx)
|
|
# fmt: on
|
|
|
|
def valid(self):
|
|
"""Check if there are more tiles to process."""
|
|
return self.linear_idx < self._total_tasks
|
|
|
|
|
|
class _CLCWorker(ClusterPersistentScheduler2D):
|
|
"""Per-role CLC handle: IS-A ClusterPersistentScheduler2D (so m_idx / n_idx work as
|
|
usual) plus the role-local barrier phase and handshake. A coord-free role (e.g. an
|
|
MMA warp consuming whatever a loader staged) arms the loop with reset() not init().
|
|
"""
|
|
|
|
def __init__(self, clc, prefix):
|
|
super().__init__(
|
|
prefix,
|
|
num_m_tiles=clc._num_m_tiles,
|
|
num_n_tiles=clc._num_n_tiles,
|
|
num_clusters=clc._num_m_tiles * clc._num_n_tiles,
|
|
l2_group_size=clc._l2_group_size,
|
|
)
|
|
self._clc = clc
|
|
self._sa = PipelineState(1, 0)
|
|
self._done = T.local_scalar("int32")
|
|
self._nxt = T.local_scalar("uint32")
|
|
|
|
@T.inline
|
|
def reset(self):
|
|
self._done = 0
|
|
|
|
@T.inline
|
|
def init(self, cluster_id):
|
|
# Explicit base call: TVMScript's parser has no zero-arg super().
|
|
ClusterPersistentScheduler2D.init(self, cluster_id)
|
|
self._done = 0
|
|
|
|
def valid(self):
|
|
return self._done == 0
|
|
|
|
@T.inline
|
|
def consume(self):
|
|
# Single-elected-thread scope: wait for the handle, decode, release the slot.
|
|
self._clc.sched_arr.full.wait(0, self._sa.phase)
|
|
self._sa.advance()
|
|
self._nxt = T.ptx.clc_query_cancel(T.address_of(self._clc.clc_handle[0]))
|
|
self._clc.sched_fin.empty.arrive(0, cta_id=0, pred=True)
|
|
|
|
@T.inline
|
|
def consume_wg(self, wg_id, warp_id, lane_id):
|
|
# Warpgroup scope: all threads decode; one elected lane releases the slot.
|
|
self._clc.sched_arr.full.wait(0, self._sa.phase)
|
|
self._sa.advance()
|
|
self._nxt = T.ptx.clc_query_cancel(T.address_of(self._clc.clc_handle[0]))
|
|
T.cuda.warpgroup_sync(wg_id + 1)
|
|
if (warp_id == 0) & (lane_id == 0):
|
|
self._clc.sched_fin.empty.arrive(0, cta_id=0, pred=True)
|
|
|
|
@T.inline
|
|
def advance_coords(self):
|
|
if self._nxt != 0xFFFFFFFF:
|
|
self.update_current_m_n_idx(self._nxt // self._clc._cta_group)
|
|
|
|
@T.inline
|
|
def mark_done_if_drained(self):
|
|
if self._nxt == 0xFFFFFFFF:
|
|
self._done = 1
|
|
|
|
|
|
@T.meta_class
|
|
class ClusterLaunchControlScheduler:
|
|
"""Blackwell Cluster Launch Control (CLC) tile scheduler.
|
|
|
|
A scheduler warp runs ``run_scheduler`` (issues ``try_cancel`` to steal the next
|
|
cluster); worker roles each take a ``worker()`` handle and pull the stolen tile
|
|
through the shared smem handshake. Owns the CLC smem: the 16B response handle, the
|
|
arrival barrier (handle ready), and the finished barrier (slot consumed;
|
|
``finish_arrivals`` arrivals per round). Tile-coord mapping is delegated to
|
|
``ClusterPersistentScheduler2D`` (group-major L2 ordering).
|
|
"""
|
|
|
|
def __init__(self, pool, num_m_tiles, num_n_tiles, l2_group_size, cta_group, finish_arrivals):
|
|
self._num_m_tiles = num_m_tiles
|
|
self._num_n_tiles = num_n_tiles
|
|
self._l2_group_size = l2_group_size
|
|
self._cta_group = cta_group
|
|
self.sched_arr = Pipeline(pool, 1, full="tma", empty="mbar", init_empty=1)
|
|
self.sched_fin = Pipeline(pool, 1, full="mbar", empty="mbar", init_empty=finish_arrivals)
|
|
self.clc_handle = pool.alloc((4,), "uint32", align=16)
|
|
self._s_done = T.local_scalar("int32")
|
|
self._s_nxt = T.local_scalar("uint32")
|
|
|
|
def worker(self, prefix):
|
|
return _CLCWorker(self, prefix)
|
|
|
|
@T.inline
|
|
def run_scheduler(self, cbx):
|
|
# cta0 drives try_cancel; both CTAs expect_bytes + consume the handle so the
|
|
# finished-barrier count is met and the slot can be reissued.
|
|
if T.ptx.elect_sync():
|
|
sa = PipelineState(1, 0)
|
|
sf = PipelineState(1, 1)
|
|
self._s_done = 0
|
|
while self._s_done == 0:
|
|
if cbx == 0:
|
|
self.sched_fin.empty.wait(0, sf.phase)
|
|
sf.advance()
|
|
T.ptx.clc_try_cancel(
|
|
T.address_of(self.clc_handle[0]), T.address_of(self.sched_arr.full.buf[0])
|
|
)
|
|
self.sched_arr.full.arrive(0, 16) # expect_bytes for the 16B handle
|
|
self.sched_arr.full.wait(0, sa.phase)
|
|
sa.advance()
|
|
self._s_nxt = T.ptx.clc_query_cancel(T.address_of(self.clc_handle[0]))
|
|
self.sched_fin.empty.arrive(0, cta_id=0, pred=True)
|
|
if self._s_nxt == 0xFFFFFFFF:
|
|
self._s_done = 1
|