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
RolePipeline
struct RolePipeline[number_of_stages: Int, is_producer: Bool = True, producer_sub_stages: Int = Int(1), consumer_sub_stages: Int = Int(1), cta_group: Int = Int(1)]
Unified producer/consumer pipeline for barrier synchronization.
Producer role: Starts with phase=1, uses acquire/commit methods. Consumer role: Starts with phase=0, uses wait/release methods.
Sub-stages allow multiple barriers per stage:
- Total producer barriers: num_stages * producer_sub_stages
- Total consumer barriers: num_stages * consumer_sub_stages
Synchronization behavior (example with num_stages=1):
Producer: p0. consumer_mbar.wait(phase=1) # 1 != 0: falls through p1. producer_mbar.commit() # producer_mbar.phase=1 p2. step() # phase = 0 p3. consumer_mbar.wait(phase=0) # 0 == 0: blocked until c1 ...
Consumer: c0. producer_mbar.wait(phase=0) # 0 == 0: blocked until p1 c1. consumer.release() # consumer_mbar.phase=1 c2. step() # phase = 1 ...
Parameters
- number_of_stages (
Int): Number of double-buffered pipeline stages. - is_producer (
Bool): Whether this instance is the producer role (defaults toTrue). - producer_sub_stages (
Int): Number of producer mbarriers per stage (defaults to 1). - consumer_sub_stages (
Int): Number of consumer mbarriers per stage (defaults to 1). - cta_group (
Int): Number of cooperating CTAs for MMA commit arrivals (defaults to 1).
Fields
- producer_mbar_base (
MBarType): - consumer_mbar_base (
MBarType): - state (
PipelineState[number_of_stages]):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
num_stages
comptime num_stages = number_of_stages
Methods
__init__
def __init__(producer_mbar_base: Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], consumer_mbar_base: Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False]) -> Self
producer_mbar
def producer_mbar[sub_stage_idx: Int = Int(0)](self) -> MBarType
Get producer mbar for current stage and optional sub-stage.
Parameters:
- sub_stage_idx (
Int): Sub-stage index (0 to producer_sub_stages-1).
Returns:
MBarType
consumer_mbar
def consumer_mbar[sub_stage_idx: Int = Int(0)](self) -> MBarType
Get consumer mbar for current stage and optional sub-stage.
Parameters:
- sub_stage_idx (
Int): Sub-stage index (0 to consumer_sub_stages-1).
Returns:
MBarType
acquire
def acquire[sub_stage_idx: Int = Int(0)](self)
Wait until consumer has released the buffer. Producer-only.
Parameters:
- sub_stage_idx (
Int): Consumer sub-stage barrier to wait on (defaults to 0).
commit
def commit(mut self)
Commit production and step. Producer-only.
commit_mma
def commit_mma(self)
Commit via MMA arrive using elected thread. Producer-only.
def commit_mma(self, elect: Int32)
Commit via MMA arrive with explicit elect value. Producer-only.
Args:
- elect (
Int32):elect()result selecting the single arriving thread.
wait
def wait(self)
Wait for producer to complete. Consumer-only.
release
def release[sub_stage_idx: Int = Int(0)](mut self)
Release buffer at sub-stage and step. Consumer-only.
Parameters:
- sub_stage_idx (
Int): Consumer sub-stage barrier to arrive on (defaults to 0).
release_no_step
def release_no_step[sub_stage_idx: Int = Int(0)](self)
Release buffer without stepping. For multi-sub-stage release.
Parameters:
- sub_stage_idx (
Int): Consumer sub-stage barrier to arrive on (defaults to 0).
step
def step(mut self)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!