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
WorkIterator
struct WorkIterator[num_stages: Int, cluster_shape: IndexList[Int(3), element_type=DType.uint32], rasterize_order: RasterOrder, block_swizzle_size: Int]
Per-warp work iterator that owns work_info and pipeline state.
Each warp creates its own WorkIterator which internally manages both the current work item and the CLC pipeline consumer state. Throttle pipeline is obtained from the scheduler.
Usage: var work_iter = scheduler.work_iterator() for current in work_iter: scheduler.throttle_signal(ctx.is_first_cta_in_cluster) do_work(current)
Parameters
- num_stages (
Int): Number of CLC pipeline stages for work distribution. - cluster_shape (
IndexList[Int(3), element_type=DType.uint32]): Cluster tile counts as(m, n, k). - rasterize_order (
RasterOrder): Order CLC rasterizes tiles across the cluster grid. - block_swizzle_size (
Int): Block swizzle factor for tile remapping, one of 0, 1, 2, 4, or 8.
Fields
- scheduler (
WorkIterator[num_stages, cluster_shape, rasterize_order, block_swizzle_size].SchedulerType): - work_info (
WorkInfo): - consumer_state (
PipelineState[num_stages]): - needs_fetch (
Bool):
Implemented traits
AnyType,
Copyable,
ImplicitlyDeletable,
Iterable,
Iterator,
Movable,
RegisterPassable
comptime members
Element
comptime Element = WorkInfo
IteratorType
comptime IteratorType[iterable_mut: Bool, //, iterable_origin: Origin[mut=iterable_mut]] = WorkIterator[num_stages, cluster_shape, rasterize_order, block_swizzle_size]
Parameters
- iterable_mut (
Bool): - iterable_origin (
Origin[mut=iterable_mut]):
SchedulerType
comptime SchedulerType = TileScheduler[num_stages, cluster_shape, rasterize_order, block_swizzle_size]
Methods
__init__
def __init__(scheduler: TileScheduler[num_stages, cluster_shape, rasterize_order, block_swizzle_size], work_info: WorkInfo) -> Self
Create work iterator with initial work_info.
Args:
- scheduler (
TileScheduler[num_stages, cluster_shape, rasterize_order, block_swizzle_size]): TileScheduler owning the CLC and throttle state. - work_info (
WorkInfo): Initial work item for the iterator.
__iter__
def __iter__(ref self) -> Self
__next__
def __next__(mut self) -> WorkInfo
Return current work item, deferring fetch to next call.
On the first call, returns the initial work_info without fetching. On subsequent calls, fetches next work (deferred from previous iteration) before returning. This allows the caller to signal throttle between iterations, avoiding deadlock with the scheduler warp.
Returns:
WorkInfo
Raises:
StopIteration: When there is no more work to process.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!