For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

InputConsumer

struct InputConsumer[origin: MutOrigin, Payload: TilePayload, num_group_stages: Int, k_group_size: Int]

Consumer view for MMA warp. Use acquire() to get stages.

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, Deinitable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

PipelineType

comptime PipelineType = InputTilePipeline[Payload, num_group_stages, k_group_size]

Methods

__enter__

def __enter__(mut self) -> Self

__exit__

def __exit__(mut self)

acquire

def acquire(mut self) -> ConsumerTiles[origin, Payload, num_group_stages, k_group_size]

Acquire next stage, waiting for tiles to be ready.

Returns a context manager for processing tiles.

Returns:

ConsumerTiles[origin, Payload, num_group_stages, k_group_size]

acquire_stage

def acquire_stage(mut self) -> InputConsumerStage[origin, Payload, num_group_stages, k_group_size]

Acquire next stage as a linear type, waiting for tiles to be ready.

Returns a compiler-enforced linear type that must be explicitly released via tiles^.release().

Returns:

InputConsumerStage[origin, Payload, num_group_stages, k_group_size]

try_acquire

def try_acquire(mut self) -> Bool

Non-blocking check if next consumer stage has data.

Use with acquire_if_needed() for the try-acquire pattern:

var ready = consumer.try_acquire()
# ... do other work ...
with consumer.acquire_if_needed(ready) as tiles:
    process_tiles()

Returns:

Bool: True if the stage has data, False if waiting is needed.

acquire_if_needed

def acquire_if_needed(mut self, already_ready: Bool) -> ConsumerTiles[origin, Payload, num_group_stages, k_group_size]

Acquire stage, only waiting if not already ready.

Args:

• ​already_ready (Bool): Result from try_acquire(). Skips wait if True.

Returns:

ConsumerTiles[origin, Payload, num_group_stages, k_group_size]: Context manager wrapping the consumer stage.