Mojo struct
WorkIteratorSplitK
struct WorkIteratorSplitK[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]
Per-warp work iterator for split-K using next-style iteration.
Usage: var work_iter = scheduler.work_iterator() for current in work_iter: scheduler.throttle_signal(ctx.is_first_cta_in_cluster) do_work(current)
Fields
- scheduler (
WorkIteratorSplitK[num_stages, reduction_tile_shape, cluster_shape, rasterize_order, block_swizzle_size, num_split_k].SchedulerType): - work_info (
WorkInfo): - consumer_state (
PipelineState[num_stages]): - 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]] = WorkIteratorSplitK[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]
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 work iterator with initial work_info.
__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.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!