Mojo struct
ProducerTiles
struct ProducerTiles[origin: MutOrigin, Payload: TilePayload, num_group_stages: Int, k_group_size: Int]
Context manager for producing one input pipeline stage.
Provides the same accessor interface as InputProducerStage (payload, stage, barrier, expect_bytes) but automatically releases the producer on scope exit.
Usage: with producer.acquire() as tiles: tiles.expect_bytes(num_bytes) load_tiles(tiles.payload(), tiles.stage(), tiles.barrier()) # release called automatically
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β
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime membersβ
PipelineTypeβ
comptime PipelineType = InputTilePipeline[Payload, num_group_stages, k_group_size]
Methodsβ
__init__β
__init__(pipeline_ptr: Pointer[InputTilePipeline[Payload, num_group_stages, k_group_size], origin], stage: UInt32, barrier: UnsafePointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED]) -> Self
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
__enter__β
__enter__(self) -> Self
__exit__β
__exit__(self)
Release the producer (advances to next stage).
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!