Mojo struct
SchedulerWorkIteratorSplitK
struct SchedulerWorkIteratorSplitK[num_stages: Int, reduction_tile_shape: IndexList[3], cluster_shape: IndexList[3, element_type=DType.uint32], rasterize_order: RasterOrder, block_swizzle_size: Int, num_split_k: Int]
Work iterator for Scheduler warp (split-K) using next-style iteration.
Usage: var sched_iter = scheduler.scheduler_iterator() for _ in sched_iter: sched_iter.signal_and_advance() sched_iter.drain()
Fields
- scheduler (
SchedulerWorkIteratorSplitK[num_stages, reduction_tile_shape, cluster_shape, rasterize_order, block_swizzle_size, num_split_k].SchedulerType): - work_info (
WorkInfo): - consumer_state (
PipelineState[num_stages]): - producer_state (
PipelineState[num_stages]): - throttle_pipeline (
SchedulerWorkIteratorSplitK[num_stages, reduction_tile_shape, cluster_shape, rasterize_order, block_swizzle_size, num_split_k].ThrottlePipeline): - needs_fetch (
Bool):
Implemented traits
AnyType,
Copyable,
ImplicitlyDestructible,
Iterable,
Iterator,
Movable,
RegisterPassable
comptime members
Element
comptime Element = WorkInfo
IteratorType
comptime IteratorType[iterable_mut: Bool, //, iterable_origin: Origin[mut=iterable_mut]] = SchedulerWorkIteratorSplitK[num_stages, reduction_tile_shape, cluster_shape, rasterize_order, block_swizzle_size, num_split_k]
Parameters
SchedulerType
comptime SchedulerType = TileScheduler[num_stages, reduction_tile_shape, cluster_shape, rasterize_order, block_swizzle_size, num_split_k]
ThrottlePipeline
comptime ThrottlePipeline = SchedulerWorkIteratorSplitK[num_stages, reduction_tile_shape, cluster_shape, rasterize_order, block_swizzle_size, num_split_k].SchedulerType.ThrottlePipeline
Methods
__init__
__init__(scheduler: TileScheduler[num_stages, reduction_tile_shape, cluster_shape, rasterize_order, block_swizzle_size, num_split_k], work_info: WorkInfo) -> Self
Create scheduler iterator. Throttle pipeline from scheduler.
__iter__
__iter__(ref self) -> Self
__next__
__next__(mut self) -> WorkInfo
Return current work item, deferring fetch to next call.
Returns:
WorkInfo
Raises:
StopIteration: When there is no more work to process.
signal_and_advance
signal_and_advance(mut self)
Signal CLC throttle consumer and advance to next work request.
drain
drain(mut self)
Drain all pending CLC requests before kernel exit.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!