Mojo struct
PerKConsumerStage
struct PerKConsumerStage[origin: MutOrigin, opc: OutputPipelineConfig]
Context manager for per-K epilogue consumption.
enter: Acquires stage, waits for MMA to complete this K iteration exit: Signals consumer barrier to release stage for MMA reuse
IMPORTANT: Unlike standard per-tile consumption, per-K consumption must signal the consumer barrier explicitly. The MMA warp waits on this barrier before each K iteration, so we must signal after each K iteration.
Fields
- pipeline_ptr (
Pointer[OutputTilePipeline[opc], origin]): - stage (
PerKConsumerStage[origin, opc].Stage):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
cta_group
comptime cta_group = opc.cta_group
num_stages
comptime num_stages = opc.num_stages
Stage
comptime Stage = OutputStage[opc]
TilePipelineType
comptime TilePipelineType = OutputTilePipeline[opc]
Methods
__init__
__init__(pipeline_ptr: Pointer[OutputTilePipeline[opc], origin]) -> Self
__enter__
__enter__(mut self) -> PerKConsumerStage[origin, opc].Stage
Returns:
PerKConsumerStage
__exit__
__exit__(mut self)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!