Mojo module
pipeline
Producer-consumer pipeline utilities for SM100 structured kernels.
This module provides pipeline synchronization primitives for warp-specialized GPU kernels, enabling efficient producer-consumer patterns between warps.
Key abstraction:
- ProducerConsumerPipeline: Low-level barrier management for N-stage pipelines
Context manager API (recommended): # Producer side (e.g., MMA warp producing to epilogue): with pipeline.produce() as stage: # stage.index() - current stage index # stage.mbar() - barrier for signaling (use with mma_arrive) # ... do work ... # Must signal via stage.mbar() before exit # exit calls producer_step()
# Consumer side (e.g., epilogue consuming from MMA):
with pipeline.consume() as stage:
# stage.index() - current stage index
# ... do work ...
# __exit__ signals consumption complete and calls consumer_step()Direct API (for special cases): pipeline.wait_producer() / wait_consumer() pipeline.producer_step() / consumer_step() pipeline.producer_mbar(stage) / consumer_mbar(stage)
comptime values
MbarPtr
comptime MbarPtr = LegacyUnsafePointer[SharedMemBarrier, address_space=AddressSpace.SHARED]
UnsafePointer
comptime UnsafePointer = LegacyUnsafePointer[?, address_space=?, origin=?]
Structs
-
ConsumeContext: Context for consuming one pipeline stage. -
ConsumerStage: Stage info returned by ConsumeContext.enter. -
ExplicitConsumeContext: Context for consuming one pipeline stage with EXPLICIT arrive. -
ProduceContext: Context for producing one pipeline stage. -
ProducerConsumerPipeline: A producer-consumer pipeline using shared memory barriers to enforce synchronization (between producer and consumer warps). -
ProducerStage: Stage info returned by ProduceContext.enter.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!