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.
Fields
- pipeline_ptr (
Pointer[InputTilePipeline[Payload, num_group_stages, k_group_size], origin]):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
PipelineType
comptime PipelineType = InputTilePipeline[Payload, num_group_stages, k_group_size]
Methods
__enter__
__enter__(mut self) -> Self
__exit__
__exit__(mut self)
drain
drain(mut self)
Drain pipeline to prevent CTA exit while peer is still working.
acquire
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
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
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
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.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!