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[num_stages: Int, cluster_shape: IndexList[Int(3), element_type=DType.uint32] = Index[Int, Int, Int, dtype=DType.uint32](Int(1), Int(1), Int(1)), rasterize_order: RasterOrder = RasterOrder.AlongM, block_swizzle_size: Int = Int(8)]

Schedules output-tile work across SM100 clusters for persistent matmul kernels using Cluster Launch Control (CLC).

Applies a configurable rasterization order and optional block swizzling to map cluster-level work tiles to global output-tile coordinates, and drives producer/consumer synchronization through shared-memory barriers.

Parameters

• ​num_stages (Int): Number of pipeline stages used for producer/consumer synchronization across waves; 0 means only the initial wave is valid and subsequent work is invalid.
• ​cluster_shape (IndexList[Int(3), element_type=DType.uint32]): Three-element IndexList giving the number of CTA blocks per cluster along the M, N, and K output-tile axes (defaults to (1, 1, 1)).
• ​rasterize_order (RasterOrder): Direction the scheduler rasterizes clusters across output tiles, either along M or N (defaults to RasterOrder.AlongM).
• ​block_swizzle_size (Int): Size of the block-swizzle group used to remap cluster coordinates for improved L2 reuse; one of 0, 1, 2, 4, or 8, where 0 disables swizzling (defaults to 8).

Fields

• ​cluster_dim (StaticTuple[Int32, Int(3)]):
• ​log_cluster_dim_m (FastDiv[DType.uint32]):
• ​log_cluster_dim_n (FastDiv[DType.uint32]):
• ​log_cluster_dim_k (FastDiv[DType.uint32]):
• ​clc_response (Pointer[UInt128, MutAnyOrigin, address_space=AddressSpace.SHARED]):
• ​full_mbar (Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED]):
• ​empty_mbar (Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED]):

Implemented traits

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

comptime members

cluster_size

comptime cluster_size = (Int((mul cluster_shape[Int(0)], cluster_shape[Int(1)])) * cluster_shape[Int(2)])

log_cluster_k

comptime log_cluster_k = FastDiv(cluster_shape[Int(2)])

log_cluster_m

comptime log_cluster_m = FastDiv(cluster_shape[Int(0)])

log_cluster_n

comptime log_cluster_n = FastDiv(cluster_shape[Int(1)])

Methods

__init__

def __init__(cluster_dim: StaticTuple[Int32, Int(3)], clc_response_ptr: Pointer[UInt128, address_space=AddressSpace.SHARED], full_mbar_ptr: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED], empty_mbar_ptr: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED]) -> Self

work_info_from_clc_response

static def work_info_from_clc_response(result: Pointer[UInt128, address_space=AddressSpace.SHARED]) -> WorkInfo

Returns:

WorkInfo

work_info_from_cluster

static def work_info_from_cluster(work_info: WorkInfo, cluster_dim: StaticTuple[Int32, Int(3)], log_cluster_dim_m: FastDiv[DType.uint32], log_cluster_dim_n: FastDiv[DType.uint32]) -> WorkInfo

Returns:

WorkInfo

initial_work_info

def initial_work_info(self) -> WorkInfo

Returns:

WorkInfo

fetch_next_work

def fetch_next_work(self, work_info: WorkInfo, consumer_state: PipelineState[num_stages]) -> WorkInfo

Returns:

WorkInfo

advance_to_next_work

def advance_to_next_work(self, mut clc_state: PipelineState[num_stages]) -> PipelineState[num_stages]

Returns:

PipelineState[num_stages]