@dataclass(frozen=True)
class PipelineAsync:
    """Generic producer-consumer pipeline with barrier synchronization.
    
    State transitions for one pipeline entry (mbarrier):
    
    +-----------+-----------+-----------+-----------+-----------+-----------+
    | Barrier   | State     | p.acquire | p.commit  | c.wait    | c.release |
    +===========+===========+===========+===========+===========+===========+
    | empty_bar | empty     | <Return>  | n/a       | n/a       | -         |
    +-----------+-----------+-----------+-----------+-----------+-----------+
    | empty_bar | wait      | <Block>   | n/a       | n/a       | -> empty  |
    +-----------+-----------+-----------+-----------+-----------+-----------+
    | full_bar  | wait      | n/a       | -> full   | <Block>   | n/a       |
    +-----------+-----------+-----------+-----------+-----------+-----------+
    | full_bar  | full      | n/a       | -         | <Return>  | n/a       |
    +-----------+-----------+-----------+-----------+-----------+-----------+
    """
    sync_object_full: SyncObject    # Barrier for "data ready"
    sync_object_empty: SyncObject   # Barrier for "buffer empty"
    num_stages: int                # Number of pipeline stages
    producer_mask: Optional[Int32] # Threads participating in producer
    consumer_mask: Optional[Int32] # Threads participating in consumer

# Producer side:
producer.acquire(state)     # Wait for buffer to be empty
producer.commit(state)# Signal buffer is full

# Consumer side:
consumer.wait(state)        # Wait for buffer to be full
consumer.release(state)     # Signal buffer is empty

# From dense_gemm_persistent.py

# Create pipeline with num_stages buffers
ab_producer, ab_consumer = pipeline.PipelineAsync.create(
    num_stages=self.num_ab_stage,# e.g., 2 or 3
    producer_group=producer_warp,
    consumer_group=consumer_warp,
)

# Initialize state for producer and consumer
producer_state = pipeline.make_pipeline_state(
    pipeline.PipelineUserType.Producer, self.num_ab_stage
)
consumer_state = pipeline.make_pipeline_state(
    pipeline.PipelineUserType.Consumer, self.num_ab_stage
)

# Producer loop (e.g., TMA load warp)
for k_tile in range(k_tile_cnt):
    # Wait for buffer to be empty
    ab_producer.acquire(producer_state)
    
    # Get buffer index
    buffer_idx = producer_state.index()
    
    # Write data to shared memory
    cute.copy(tma_atom_A, gA_slice, sA[(None, buffer_idx)])
    cute.copy(tma_atom_B, gB_slice, sB[(None, buffer_idx)])
    
    # Signal buffer is full
    ab_producer.commit(producer_state)
    
    # Advance to next stage
    producer_state.advance()

# Consumer loop (e.g., MMA warps)
for k_tile in range(k_tile_cnt):
    # Wait for buffer to be full
    ab_consumer.wait(consumer_state)
    
    # Get buffer index
    buffer_idx = consumer_state.index()
    
    # Read data from shared memory
    cute.cp_async_wait(ab_buffer.barrier)  # Ensure TMA arrived
    # ... compute MMA ...
    
    # Signal buffer is empty
    ab_consumer.release(consumer_state)
    
    # Advance to next stage
    consumer_state.advance()

# Pattern is the same for all pipelines:
# Producer: acquire → write → commit
# Consumer: wait → read → release

# PipelineTmaAsync example (TMA is producer):
producer, consumer = PipelineTmaAsync.create(num_stages=2, ...)
producer.acquire(state)      # Wait for buffer empty
# TMA loads to smem (hardware arrives automatically)
producer.commit(state)         # Signal buffer full

consumer.wait(state)          # Wait for buffer full
# Read from smem
consumer.release(state)      # Signal buffer empty

def make_pipeline_state(type: PipelineUserType, stages: int):
    """Create initial state for producer or consumer.
    
    Args:
        type: Producer or Consumer
        stages: Number of pipeline stages
    
    Returns:
        PipelineState with initial index and phase
    """
    if type is PipelineUserType.Producer:
        return PipelineState(stages, Int32(stages))# Start at stages, phase = 1
    elif type is PipelineUserType.Consumer:
        return PipelineState(stages, Int32(0))    # Start at 0, phase = 0

class PipelineStateSimple:
    """Optimized pipeline state using single Int32 for index + phase.
    
    Encoding: phase_index = index + phase * stages
    - index = phase_index % stages
    - phase = phase_index // stages
    """
    def __init__(self, stages: int, phase_index: Int32):
        self._stages = stages
        self._phase_index = phase_index
    
    @property
    def index(self) -> Int32:
        if self._stages == 1:
            return Int32(0)# Optimized: no modulo needed
        else:
            return self._phase_index % self._stages
    
    @property
    def phase(self) -> Int32:
        if self._stages == 1:
            return self._phase_index  # Optimized: no division needed
        else:
            return self._phase_index // self._stages
    
    def advance(self):
        if self._stages == 1:# Single stage: flip phase bit
            self._phase_index ^= 1
        else:
            self._phase_index += 1

class _PipelineIndexPhaseMixin:
    """Mixin providing *_w_index_phase methods."""
    
    def producer_acquire_w_index_phase(self, index: Int32, phase: Int32, ...):
        state = _make_state(index, phase)  # Create PipelineState
        self.producer_acquire(state, ...)   # Call parent method
    
    def producer_commit_w_index(self, index: Int32, ...):
        state = _make_state(index, Int32(0))  # Phase not needed for commit
        self.producer_commit(state, ...)
    
    def consumer_wait_w_index_phase(self, index: Int32, phase: Int32, ...):
        state = _make_state(index, phase)
        self.consumer_wait(state, ...)
    
    def consumer_release_w_index(self, index: Int32, ...):
        state = _make_state(index, Int32(0))
        self.consumer_release(state, ...)

@dataclass(frozen=True)
class PipelineAsync(_PipelineIndexPhaseMixin, PipelineAsyncOg):
    """PipelineAsync with optional elect_one for barrier operations.
    
    Args:
        elect_one_commit: If True, only elected thread signals producer_commit
        syncwarp_before_commit: If True, syncwarp before elect_one
        elect_one_release: If True, only elected thread signals consumer_release
        syncwarp_before_release: If True, syncwarp before elect_one
    """
    _elect_one_commit: bool = False
    _syncwarp_before_commit: bool = True
    _elect_one_release: bool = False
    _syncwarp_before_release: bool = True

# All threads in consumer_group call barrier.arrive_and_release
# Good for: barrel equals thread count

# Only elected thread signals barrier
# Good for: barrier count is 1 per warp
with cute.arch.elect_one():
    pipeline.consumer_release(state)# Only one thread executes

# From flash_fwd_sm90.py

# Create pipeline for KV loading
pipeline_k = pipeline.PipelineTmaAsync.create(
    num_stages=self.num_stages,# e.g., 2
    producer_group=load_warp,
    consumer_group=mma_warps,
)

# Initialize state
kv_producer_state = pipeline.make_pipeline_state(
    pipeline.PipelineUserType.Producer, self.num_stages
)
kv_consumer_state = pipeline.make_pipeline_state(
    pipeline.PipelineUserType.Consumer, self.num_stages
)

# Producer (TMA load warp)
for k_tile in range(k_tile_cnt):
    # Wait for buffer to be empty (TMA waits on empty barrier)
    pipeline_k.producer_acquire(kv_producer_state)
    
    # Issue TMA load
    # TMA writes to sK[(None, kv_producer_state.index)] directly
    cute.copy(tma_atom_K, gK, sK[(None, kv_producer_state.index)])
    
    # Signal buffer is full (sets mbarrier arrival count)
    pipeline_k.producer_commit(kv_producer_state)
    
    kv_producer_state.advance()

# Consumer (MMA warps)
for k_tile in range(k_tile_cnt):
    # Wait for buffer to be full
    pipeline_k.consumer_wait(kv_consumer_state)
    
    # Compute MMA
    # ... uses sK[(None, kv_consumer_state.index)]
    
    # Signal buffer is empty
    pipeline_k.consumer_release(kv_consumer_state)
    
    kv_consumer_state.advance()

def producer_acquire(
    self,
    state: PipelineState,
    try_acquire_token: Optional[Boolean] = None,
    extra_tx_count: int = 0,  # Additional TMA transactions to wait for
    ...
):
    # Wait for buffer to be empty
    if try_acquire_token is None or try_acquire_token == 0:
        self.sync_object_empty.wait(state.index, state.phase)
    
    # Arrive on full barrier with expected transaction count
    if extra_tx_count == 0:
        self.sync_object_full.arrive(state.index, self.producer_mask)
    else:
        # Total transactions = base_tx_count + extra_tx_count
        tx_count = self.sync_object_full.tx_count + extra_tx_count
        self.sync_object_full.arrive_and_expect_tx(state.index, tx_count)

# Example: TMA load Q + another async operation
pipeline.producer_acquire(state, extra_tx_count=additional_bytes)
# Consumer waits until BOTH TMA bytes + async operation arrive

# FlashAttention imports CuTe native classes with "Og" suffix (Original)
from cutlass.pipeline import PipelineAsync as PipelineAsyncOg
from cutlass.pipeline import PipelineTmaAsync as PipelineTmaAsyncOg
from cutlass.pipeline import PipelineCpAsync as PipelineCpAsyncOg
# ... etc

@dataclass(frozen=True)
class PipelineAsync(_PipelineIndexPhaseMixin, PipelineAsyncOg):
    """Adds elect_one and syncwarp flags."""
    _elect_one_commit: bool = False
    _syncwarp_before_commit: bool = True
    _elect_one_release: bool = False
    _syncwarp_before_release: bool = True

# From flash_fwd_sm90.py

# Setup for Q, K, V loading
pipeline_q = pipeline.PipelineTmaAsync.create(
    num_stages=1,
    producer_group=tma_warp,
    consumer_group=consumer_warps,
)

pipeline_k = pipeline.PipelineTmaAsync.create(
    num_stages=self.num_stages,
    producer_group=tma_warp,
    consumer_group=consumer_warps,
)

pipeline_v = pipeline.PipelineTmaAsync.create(
    num_stages=self.num_stages,
    producer_group=tma_warp,
    consumer_group=consumer_warps,
)

# Shared memory barriers
@mbar_load_Q: MemRange[Int64, num_stages * 2]# *2 for full and empty
@mbar_load_KV: MemRange[Int64, num_stages * 2]

# Create
pipeline = PipelineAsync.create(
    num_stages=2,
    producer_group=producer_warp,
    consumer_group=consumer_warp,
)

# State
producer_state = make_pipeline_state(PipelineUserType.Producer, 2)
consumer_state = make_pipeline_state(PipelineUserType.Consumer, 2)

# Producer
producer.acquire(producer_state)
# ... write to buffer ...
producer.commit(producer_state)
producer_state.advance()

# Consumer
consumer.wait(consumer_state)
# ... read from buffer ...
consumer.release(consumer_state)
consumer_state.advance()

# Create (same)
pipeline = PipelineAsync.create(
    num_stages=2,
    producer_group=producer_warp,
    consumer_group=consumer_warp,
    elect_one_commit=True,# New: only elected thread commits
    syncwarp_before_commit=True,  # New: sync before elect_one
)

# State (optimized)
producer_state = make_pipeline_state(PipelineUserType.Producer, 2)
consumer_state = make_pipeline_state(PipelineUserType.Consumer, 2)

# Producer (same pattern)
pipeline.producer_acquire(producer_state)
# ... write to buffer ...
pipeline.producer_commit(producer_state)  # Uses elect_one internally
producer_state.advance()

# Consumer (same pattern)
pipeline.consumer_wait(consumer_state)
# ... read from buffer ...
pipeline.consumer_release(consumer_state)  # Uses elect_one internally
consumer_state.advance()

# 2-stage pipeline overlaps load and compute
# While computing buffer 0, loading buffer 1

Time | Producer        | Consumer
-----|-----------------|------------------
t0| acq(0), load(0) | wait(prev)
t1| commit(0)       | wait(0)
t2| acq(1), load(1) | read(0), compute
t3| commit(1)        | release(0)
t4| acq(0), load(0) | wait(1)
t5| ...              | read(1), compute

# 1-stage pipeline with phase alternation
# Producer and consumer alternate on same buffer

Time | Producer      | Consumer
-----|---------------|--------
t0   | acq, load     | -
t1   | commit        | wait
t2   | -             | read, compute
t3   | acq (block)   | release
t4   | load          | -

# 3-stage pipeline for complex kernels
# More overlap, more smem, more latency hiding

Producer: acq(2) → load(2) → commit(2) → acq(0) → ...
Consumer:             wait(0) → read(0) → release(0) → wait(1) → ...