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

InputProducer

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

Producer view for TMA Load 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)

drain

def drain(mut self)

Drain pipeline to prevent CTA exit while peer is still working.

acquire

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

Acquire next stage, waiting for slot availability.

Returns a context manager for loading tiles.

Returns:

ProducerTiles[origin, Payload, num_group_stages, k_group_size]

acquire_stage

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

Acquire next stage as a linear type, waiting for slot availability.

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

Returns:

InputProducerStage[origin, Payload, num_group_stages, k_group_size]

try_acquire

def try_acquire(mut self) -> Bool

Non-blocking check if next producer stage is available.

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

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

Returns:

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

acquire_if_needed

def acquire_if_needed(mut self, already_ready: Bool) -> ProducerTiles[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:

ProducerTiles[origin, Payload, num_group_stages, k_group_size]: Context manager wrapping the producer stage.