For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

TileScheduler

struct TileScheduler[problem_shape: IndexList[Int(3)], tile_shape: IndexList[Int(3)], grid_shape: IndexList[Int(2)], cluster: IndexList[Int(3)] = Index[Int, Int, Int](Int(1), Int(1), Int(1)), raster_dim: UInt32 = UInt32(1), schedule: MatmulSchedule = MatmulSchedule.TILE2D]

Persistent tile scheduler for GPU matmul output tiles.

Maps thread blocks to (m, n) output tile coordinates using 1D, 2D-wave, or DeepSeek scheduling strategies. Supports block clusters with multicast and advances through tiles in a persistent kernel loop via fetch_next_work.

Parameters

• ​problem_shape (IndexList[Int(3)]): Static (M, N, K) dimensions of the matmul problem.
• ​tile_shape (IndexList[Int(3)]): Static (BM, BN, BK) tile sizes.
• ​grid_shape (IndexList[Int(2)]): (grid_x, grid_y) number of thread blocks in the launch grid.
• ​cluster (IndexList[Int(3)]): Block cluster shape (default 1×1×1 for no clustering).
• ​raster_dim (UInt32): Axis along which to rasterize (0 = M, 1 = N).
• ​schedule (MatmulSchedule): Tile scheduling strategy to use.

Fields

• ​idx (UInt32):
• ​prob_shape (IndexList[Int(3)]):
• ​num_waves_m (UInt32):
• ​num_waves_n (UInt32):
• ​log_num_waves_n (FastDiv[DType.uint32]):
• ​current_iter (Int):
• ​num_aligned_m_blocks (UInt32):
• ​num_blocks (UInt32):

Implemented traits

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

kNum1DBlocksPerGroup

comptime kNum1DBlocksPerGroup = UInt32(16)

kNumNBlocks

comptime kNumNBlocks = SIMD(ceildiv(problem_shape[Int(1)], tile_shape[Int(1)]))

num_grids

comptime num_grids = SIMD(Int((mul grid_shape[Int(0)], grid_shape[Int(1)])))

wave_shape

comptime wave_shape = Index[Int, Int, dtype=DType.uint32](Int((mul grid_shape[Int(1)], tile_shape[Int(0)])), Int((mul grid_shape[Int(0)], tile_shape[Int(1)])))

Methods

__init__

def __init__(prob_shape: IndexList[Int(3)]) -> Self

get_current_work_info

def get_current_work_info(mut self) -> WorkInfo

Returns:

WorkInfo

advance

def advance(mut self)

fetch_next_work

def fetch_next_work(mut self) -> WorkInfo

Returns:

WorkInfo

num_output_tiles

def num_output_tiles(self) -> Int

Returns:

Int

fetch_next_work_ds

def fetch_next_work_ds(mut self) -> WorkInfo

Returns:

WorkInfo