Skip to main content

@register_passable(trivial) struct InputConsumerStage[origin: MutOrigin, Payload: TilePayload, num_group_stages: Int, k_group_size: Int]

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

Two usage patterns:

  1. Context manager (scoped): with consumer.acquire() as tiles: process_tiles(tiles.payload(), tiles.stage())

    release() called automatically by exit

  2. Linear-style (flat): var tiles = consumer.acquire() process_tiles(tiles.payload(), tiles.stage()) tiles.release() # Manual release

Lifecycle:

  1. Created via consumer.acquire() - waits for producer
  2. Use payload(), stage() for tile access
  3. Call release() or let __exit__ signal and advance

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

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, RegisterType, TrivialRegisterType

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[InputConsumerStage[origin, Payload, num_group_stages, k_group_size].PipelineType, origin], stage: UInt32, mbar: 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

mbar

mbar(self) -> MbarPtr

Get the barrier pointer.

Returns:

MbarPtr

release

release(mut self)

Signal consumption and advance 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?