Mojo struct
ProducerTiles
@register_passable(trivial)
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[ProducerTiles[origin, Payload, num_group_stages, k_group_size].PipelineType, origin]):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
__copy_ctor_is_trivial
comptime __copy_ctor_is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__move_ctor_is_trivial
comptime __move_ctor_is_trivial = True
PipelineType
comptime PipelineType = InputTilePipeline[Payload, num_group_stages, k_group_size]
Methods
__init__
__init__(pipeline_ptr: Pointer[ProducerTiles[origin, Payload, num_group_stages, k_group_size].PipelineType, origin], stage: UInt32, barrier: LegacyUnsafePointer[SharedMemBarrier, 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!