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
-
BlockScaledTilePayload: Tile payload for block-scaled matmul (A, B, SFA, SFB tiles). -
ConsumerStage: Context manager for consumer tile access with encapsulated stage indexing. -
InputConsumer: Consumer view for MMA warp (input pipeline). -
InputConsumerStage: Consumer stage context manager. Released on scope exit. -
InputProducer: Producer view for TMA Load warp (input pipeline). -
InputProducerStage: Producer stage context manager. Released on scope exit. -
InputTilePipeline: Tile pipeline with configurable payload type. -
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. -
StandardTilePayload: Tile payload for standard matmul (A and B tiles). -
TileConsumer: Consumer view for MMA warp. Use acquire() to get stages. -
TilePipeline: Staged tile storage with producer-consumer synchronization for SM100. -
TileProducer: Producer view for TMA Load warp. Use acquire() to get stages.
Traits
-
TilePayload: Trait for tile payload types. Must be @register_passable("trivial").
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!