@register_passable(trivial) struct InputProducerStage[_mlir_origin: LITMutOrigin, //, origin: MutOrigin, Payload: TilePayload, num_group_stages: Int, k_group_size: Int]

Handle for producer tile access - works as context manager or linear-style.

Two usage patterns:

1. Context manager (scoped): with producer.acquire() as tiles: load_tiles(tiles.payload(), tiles.stage(), tiles.barrier())

   release() called automatically by exit

2. Linear-style (flat): var tiles = producer.acquire() load_tiles(tiles.payload(), tiles.stage(), tiles.barrier()) tiles.release() # Manual release

Lifecycle:

1. Created via producer.acquire() - waits for consumer
2. Use payload(), stage(), barrier() for TMA operations
3. Call release() or let __exit__ advance producer stage

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[InputProducerStage[origin, Payload, num_group_stages, k_group_size].PipelineType, origin]):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

__copyinit__is_trivial

comptime __copyinit__is_trivial = True

__del__is_trivial

comptime __del__is_trivial = True

__moveinit__is_trivial

comptime __moveinit__is_trivial = True

PipelineType

comptime PipelineType = InputTilePipeline[Payload, num_group_stages, k_group_size]

Methods

__init__

__init__(pipeline_ptr: Pointer[InputProducerStage[origin, Payload, num_group_stages, k_group_size].PipelineType, origin], stage: UInt32, barrier: LegacyUnsafePointer[SharedMemBarrier, address_space=AddressSpace.SHARED]) -> Self

__enter__

__enter__(mut self) -> Self

__exit__

__exit__(mut self)

payload

payload(self) -> Payload

Get the tile payload for direct access.

Returns:

Payload

stage

stage(self) -> UInt32

Get the current stage index.

Returns:

UInt32

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(mut self)

Advance producer to next stage (linear-style API).

Use this for flat code structure instead of context manager. Equivalent to what exit does.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit