For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo struct
InputTilePipeline
struct InputTilePipeline[Payload: TilePayload, num_group_stages: Int, k_group_size: Int]
Tile pipeline with configurable payload type.
Separates synchronization from tile storage. The Payload parameter (e.g., StandardTilePayload or BlockScaledTilePayload) holds tile arrays.
Parameters
- Payload (
TilePayload): Tile payload type holding the tile arrays. - num_group_stages (
Int): Number of synchronization stages in the pipeline. - k_group_size (
Int): Number of K-group tiles stored per synchronization stage.
Fields
- pipeline (
InputTilePipeline[Payload, num_group_stages, k_group_size].Pipeline): - payload (
Payload):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
BarrierArray
comptime BarrierArray = SMemArray[SharedMemBarrier, (num_group_stages * Int(2))]
Pipeline
comptime Pipeline = ProducerConsumerPipeline[num_group_stages]
Methods
__init__
def __init__(barriers: SMemArray[SharedMemBarrier, (num_group_stages * Int(2))], payload: Payload) -> Self
Initialize from typed barrier array and payload.
Args:
- barriers (
SMemArray[SharedMemBarrier, (num_group_stages * Int(2))]): Shared memory barrier array for producer-consumer synchronization. - payload (
Payload): Tile payload holding the tile arrays.
init_barriers
static def init_barriers(storage_ptr: Pointer[SharedMemBarrier, MutUntrackedOrigin, address_space=AddressSpace.SHARED, _safe=False], producer_arv_count: Int32, consumer_arv_count: Int32)
Initialize pipeline barriers. Called once by elect_one thread.
Args:
- storage_ptr (
Pointer[SharedMemBarrier, MutUntrackedOrigin, address_space=AddressSpace.SHARED, _safe=False]): Pointer to shared memory barrier storage. - producer_arv_count (
Int32): Number of producer threads expected to arrive at each barrier. - consumer_arv_count (
Int32): Number of consumer threads expected to arrive at each barrier.
try_acquire_producer
def try_acquire_producer(self) -> Bool
Non-blocking check if next producer stage is available.
Example (TMA Load warp):
var ready = pipeline.try_acquire_producer() # ... do other work while potentially waiting ... pipeline.wait_producer_if_needed(ready) var stage = pipeline.producer_stage() # ... load tiles ...
Returns:
Bool: True if consumer has freed the stage, False otherwise.
try_acquire_consumer
def try_acquire_consumer(self) -> Bool
Non-blocking check if next consumer stage has data.
Example (MMA warp):
var ready = pipeline.try_acquire_consumer() # ... do other work while potentially waiting ... pipeline.wait_consumer_if_needed(ready) var stage = pipeline.consumer_stage() # ... process tiles ...
Returns:
Bool: True if producer has filled the stage, False otherwise.
wait_producer_if_needed
def wait_producer_if_needed(self, already_ready: Bool)
Conditionally wait for producer stage if not already ready.
Args:
- already_ready (
Bool): Result from try_acquire_consumer().
wait_consumer_if_needed
def wait_consumer_if_needed(self, already_ready: Bool)
Conditionally wait for consumer to free stage if not already ready.
Args:
- already_ready (
Bool): Result from try_acquire_producer().
producer_stage
consumer_stage
producer_mbar
def producer_mbar(self, stage: UInt32) -> MbarPtr
Returns:
MbarPtr
consumer_mbar
def consumer_mbar(self, stage: UInt32) -> MbarPtr
Returns:
MbarPtr
producer
def producer[mut_origin: MutOrigin](ref[Payload] self) -> InputProducer[mut_origin, Payload, num_group_stages, k_group_size]
Get producer view for TMA Load warp.
Parameters:
- mut_origin (
MutOrigin): Origin of the pipeline reference.
Returns:
InputProducer[mut_origin, Payload, num_group_stages, k_group_size]
consumer
def consumer[mut_origin: MutOrigin](ref[Payload] self) -> InputConsumer[mut_origin, Payload, num_group_stages, k_group_size]
Get consumer view for MMA warp.
Parameters:
- mut_origin (
MutOrigin): Origin of the pipeline reference.
Returns:
InputConsumer[mut_origin, Payload, num_group_stages, k_group_size]
acquire_producer
def acquire_producer[mut_origin: MutOrigin](ref[Payload] self) -> InputProducerStage[mut_origin, Payload, num_group_stages, k_group_size]
Acquire a producer stage handle (linear type).
Waits for the consumer to free the current stage, then returns a linear type handle that MUST be released (compiler-enforced).
Usage: var tiles = pipeline.acquire_producer() load_tiles(tiles.payload(), tiles.stage(), tiles.barrier()) tiles^.release() # Advances to next stage
Parameters:
- mut_origin (
MutOrigin): Origin of the pipeline reference.
Returns:
InputProducerStage[mut_origin, Payload, num_group_stages, k_group_size]: An InputProducerStage handle that must be released.
acquire_consumer
def acquire_consumer[mut_origin: MutOrigin](ref[Payload] self) -> InputConsumerStage[mut_origin, Payload, num_group_stages, k_group_size]
Acquire a consumer stage handle (linear type).
Waits for the producer to fill the current stage, then returns a linear type handle that MUST be released (compiler-enforced).
Usage: var tiles = pipeline.acquire_consumer() process_tiles(tiles.payload(), tiles.stage()) tiles^.release() # Signals complete and advances
Parameters:
- mut_origin (
MutOrigin): Origin of the pipeline reference.
Returns:
InputConsumerStage[mut_origin, Payload, num_group_stages, k_group_size]: An InputConsumerStage handle that must be released.
drain_producer
def drain_producer(mut self)
Drain pipeline to prevent CTA exit while peer is still working.
Call this after all producer iterations are complete. This is the linear type equivalent of InputProducer.drain().
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!