Mojo struct
WorkIterator
@register_passable(trivial)
struct WorkIterator[num_stages: Int, cluster_shape: IndexList[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() while work_iter.has_work(): with work_iter.next() as current: work_iter.throttle_signal(ctx.is_first_cta_in_cluster) do_work(current)
Fields
- scheduler (
WorkIterator[num_stages, cluster_shape, rasterize_order, block_swizzle_size].SchedulerType): - work_info (
WorkInfo): - consumer_state (
PipelineState[num_stages]): - throttle_pipeline (
WorkIterator[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 = WorkIterator[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 work iterator with initial work_info. Throttle 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 (advance AFTER work pattern).
Returns:
AdvanceAfterWorkContext
next_prefetch
next_prefetch[state_origin: MutOrigin, //](ref [state_origin] self) -> PrefetchBeforeWorkContext[origin_of(state_origin._mlir_origin.work_info)]
Get next work item with prefetch (advance BEFORE work pattern).
Returns:
PrefetchBeforeWorkContext
throttle_signal
throttle_signal(mut self, is_first_cta_in_cluster: Bool)
Signal CLC throttle if this is the first CTA in cluster.
The Load warp acts as producer for CLC throttle, signaling that it has started processing a new work item. This prevents the scheduler from getting too far ahead.
Args:
- is_first_cta_in_cluster (
Bool): Only first CTA signals to avoid duplicates.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!