Mojo struct
SchedulerWorkIteratorSplitK
@register_passable(trivial)
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) - owns work_info and both states. Throttle pipeline is obtained from the scheduler.
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):
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, 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.
has_work
next
next[state_origin: MutOrigin, //](ref [state_origin] self) -> AdvanceAfterWorkContextSplitK[origin_of(state_origin._mlir_origin.work_info), origin_of(state_origin._mlir_origin.consumer_state), num_stages, reduction_tile_shape, cluster_shape, rasterize_order, block_swizzle_size, num_split_k]
Get next work item.
Returns:
AdvanceAfterWorkContextSplitK
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!