Mojo module
tile_pipeline
Tile pipeline for SM100 producer-consumer synchronization.
Provides staged tile storage with producer-consumer barrier synchronization for TMA-MMA pipeline coordination.
Key abstractions:
- TilePipeline: Manages A/B tiles across pipeline stages with sync barriers
- OutputTilePipeline: Manages TMEM accumulator stages for MMA→Epilogue pipeline
Usage: var tile_pipeline = TilePipeline[...](storage_ptr, a_tiles, b_tiles)
# Producer (TMA Load warp):
with tile_pipeline.producer() as producer:
with producer.acquire() as tiles:
tiles.expect_bytes(expected_bytes)
for j in range(k_group_size):
var a_tile, b_tile = tiles.get_tile(j)
# TMA load into a_tile, b_tile
# Automatically signals completion on exit
# Consumer (MMA warp):
with tile_pipeline.consumer() as consumer:
with consumer.acquire() as tiles:
for j in range(k_group_size):
var a_tile, b_tile = tiles.get_tile(j)
# MMA using a_tile, b_tile
# Automatically signals completion on exitcomptime values
MbarPtr
comptime MbarPtr = LegacyUnsafePointer[SharedMemBarrier, address_space=AddressSpace.SHARED]
Structs
-
ConsumerStage: Context manager for consumer tile access with encapsulated stage indexing. -
InputConsumer: Consumer view for MMA warp (input pipeline). -
InputProducer: Producer view for TMA Load warp (input pipeline). -
OutputConsumer: Consumer view for epilogue warp (output pipeline). -
OutputProducer: Producer view for MMA warp (output pipeline). -
OutputStage: Acquired output stage with TMEM handle and pipeline reference. -
OutputTilePipeline: Pipeline for MMA→Epilogue TMEM stage synchronization. -
ProducerStage: Context manager for producer tile access with encapsulated stage indexing. -
TilePipeline: Staged tile storage with producer-consumer synchronization for SM100.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!