Mojo struct
InputProducerStage
struct InputProducerStage[origin: MutOrigin, Payload: TilePayload, num_group_stages: Int, k_group_size: Int]
Linear type handle for producer tile access.
Compiler-enforced: must call release() to advance the producer stage.
Usage (linear): var tiles = producer.acquire_stage() load_tiles(tiles.payload(), tiles.stage(), tiles.barrier()) tiles^.release() # Compiler enforces this call
For scoped/automatic usage, use ProducerTiles via producer.acquire() instead.
Parameters
- origin (
MutOrigin): Origin of the pipeline reference. - Payload (
TilePayload): The tile payload type. - num_group_stages (
Int): Number of synchronization stages. - k_group_size (
Int): Number of tiles per synchronization stage.
Fields
- pipeline_ptr (
Pointer[InputTilePipeline[Payload, num_group_stages, k_group_size], origin]):
Implemented traits
comptime members
PipelineType
comptime PipelineType = InputTilePipeline[Payload, num_group_stages, k_group_size]
Methods
__init__
__init__(out self, pipeline_ptr: Pointer[InputTilePipeline[Payload, num_group_stages, k_group_size], origin], stage: UInt32, barrier: UnsafePointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED])
payload
payload(self) -> Payload
Get the tile payload for direct access.
Returns:
Payload
stage
expect_bytes
expect_bytes(self, num_bytes: Int)
Set expected bytes on the barrier for TMA loads.
barrier
barrier(self) -> MbarPtr
Get the barrier pointer for TMA multicast loads.
Returns:
MbarPtr
release
release(deinit self)
Advance producer to next stage.
This is the only way to destroy this linear type. The compiler will error if you don't call this.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!