Mojo struct
InputTilePipeline
@register_passable(trivial)
struct InputTilePipeline[Payload: TilePayload, num_group_stages: Int, k_group_size: Int]
Tile pipeline with configurable payload type.
Separates synchronization from tile storage. The Payload parameter (e.g., StandardTilePayload or BlockScaledTilePayload) holds tile arrays.
Fields
- pipeline (
InputTilePipeline[Payload, num_group_stages, k_group_size].Pipeline): - payload (
Payload):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable
comptime members
__copyinit__is_trivial
comptime __copyinit__is_trivial = Payload.__copyinit__is_trivial
__del__is_trivial
comptime __del__is_trivial = Payload.__del__is_trivial
__moveinit__is_trivial
comptime __moveinit__is_trivial = Payload.__moveinit__is_trivial
BarrierArray
comptime BarrierArray = SMemArrayType[SharedMemBarrier, (num_group_stages * 2)]
Pipeline
comptime Pipeline = ProducerConsumerPipeline[num_group_stages]
Methods
__init__
__init__(barriers: SMemArrayType[SharedMemBarrier, (num_group_stages * 2)], payload: Payload) -> Self
Initialize from typed barrier array and payload.
init_barriers
static init_barriers(storage_ptr: LegacyUnsafePointer[SharedMemBarrier, address_space=AddressSpace.SHARED], producer_arv_count: Int32, consumer_arv_count: Int32)
Initialize pipeline barriers. Called once by elect_one thread.
acquire_producer
acquire_producer(mut self) -> Tuple[UInt32, MbarPtr]
Wait for slot availability and return (stage, barrier).
Returns:
release_producer
release_producer(mut self)
Signal completion and advance producer stage.
acquire_consumer
acquire_consumer(mut self) -> Tuple[UInt32, MbarPtr]
Wait for data availability and return (stage, barrier).
Returns:
release_consumer
release_consumer(mut self)
Signal completion and advance consumer stage.
producer_stage
consumer_stage
producer_mbar
producer_mbar(self, stage: UInt32) -> MbarPtr
Returns:
MbarPtr
consumer_mbar
consumer_mbar(self, stage: UInt32) -> MbarPtr
Returns:
MbarPtr
producer
producer[mut_origin: MutOrigin](ref [mut_origin] self) -> TileProducer[mut_origin, Payload, num_group_stages, k_group_size]
Get producer view for TMA Load warp.
Returns:
TileProducer
consumer
consumer[mut_origin: MutOrigin](ref [mut_origin] self) -> TileConsumer[mut_origin, Payload, num_group_stages, k_group_size]
Get consumer view for MMA warp.
Returns:
TileConsumer
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!