Mojo struct
SchedulerWorkIterator
@register_passable(trivial)
struct SchedulerWorkIterator[num_stages: Int, cluster_shape: IndexList[3, element_type=DType.uint32], rasterize_order: RasterOrder, block_swizzle_size: Int]
Work iterator for Scheduler warp - owns work_info and both pipeline states.
The Scheduler warp uniquely needs to:
- Consume work responses (like other warps) via next()
- Signal throttle and produce new work requests via signal_and_advance()
- Drain pending requests at exit via drain()
Usage: var sched_iter = scheduler.scheduler_iterator() while sched_iter.has_work(): with sched_iter.next(): sched_iter.signal_and_advance() sched_iter.drain()
Fields
- scheduler (
SchedulerWorkIterator[num_stages, cluster_shape, rasterize_order, block_swizzle_size].SchedulerType): - work_info (
WorkInfo): - consumer_state (
PipelineState[num_stages]): - producer_state (
PipelineState[num_stages]): - throttle_pipeline (
SchedulerWorkIterator[num_stages, cluster_shape, rasterize_order, block_swizzle_size].ThrottlePipeline):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
Movable,
UnknownDestructibility
comptime members
__copyinit__is_trivial
comptime __copyinit__is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__moveinit__is_trivial
comptime __moveinit__is_trivial = True
SchedulerType
comptime SchedulerType = TileScheduler[num_stages, cluster_shape, rasterize_order, block_swizzle_size]
ThrottlePipeline
comptime ThrottlePipeline = SchedulerWorkIterator[num_stages, cluster_shape, rasterize_order, block_swizzle_size].SchedulerType.ThrottlePipeline
Methods
__init__
__init__(scheduler: TileScheduler[num_stages, cluster_shape, rasterize_order, block_swizzle_size], work_info: WorkInfo) -> Self
Create scheduler iterator. Throttle pipeline from scheduler.
has_work
next
next[state_origin: MutOrigin, //](ref [state_origin] self) -> AdvanceAfterWorkContext[origin_of(state_origin._mlir_origin.work_info), origin_of(state_origin._mlir_origin.consumer_state), num_stages, cluster_shape, rasterize_order, block_swizzle_size]
Get next work item.
Returns:
AdvanceAfterWorkContext
signal_and_advance
signal_and_advance(mut self)
Signal CLC throttle consumer and advance to next work request.
Combines two operations that always happen together in Scheduler warp:
- Signal throttle consumer (tells Load warp we've consumed a response)
- Issue next CLC work request (producer side)
drain
drain(mut self)
Drain all pending CLC requests before kernel exit.
Must be called after the work loop completes to ensure all CLC pipeline stages are properly synchronized before exit.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!