IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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, _safe=False]):
  • ​full_mbar (Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False]):
  • ​empty_mbar (Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False]):

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, 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, _safe=False], full_mbar_ptr: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED, _safe=False], empty_mbar_ptr: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED, _safe=False]) -> Self

work_info_from_clc_response​

static def work_info_from_clc_response(result: Pointer[UInt128, address_space=AddressSpace.SHARED, _safe=False]) -> 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]

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact