"""
Native Blackwell Stage-C drop-in implementation for FmhaV3.

This file is intentionally a patch module rather than a second full copy of the
whole test.  It contains the exact role layout, pipeline definitions, and the
native correction helpers to splice into your current test_fmha_v3.py / Stage-B
kernel.

Architecture:
  softmax warps     0-3  : S(TMEM) -> P(TMEM), vec(TMEM)
  correction warps  4-7  : vec(TMEM) + O(TMEM) -> corrected O(SMEM)
  MMA warp          8    : QK and PV
  TMA/load warp     9    : Q/K/V load
  epilogue warp     10   : corrected O SMEM -> GMEM
  empty warp        11   : tmem dealloc mbar init/helper

The key design decision: softmax NEVER rescales O.  Correction is the only code
path that touches O for online rescale/final normalization.  Epilogue never
reads TMEM.
"""

# -------------------------
# 1) role layout
# -------------------------
# Put this in __init__
SOFTMAX_WARP_IDS = (0, 1, 2, 3)
CORRECTION_WARP_IDS = (4, 5, 6, 7)
MMA_WARP_ID = 8
TMA_WARP_ID = 9
EPILOGUE_WARP_ID = 10
EMPTY_WARP_ID = 11
THREADS_PER_CTA = 32 * 12

# Pipeline stages to add to your class
MMA_SOFTMAX_STAGE = 1
SOFTMAX_CORR_STAGE = 1
MMA_CORR_STAGE = 2
EPI_STAGE = 2

# TMEM offsets for a single 128x128 S/P tile and 128x64 O tile.
# vec reuses S region only after the S tile has been loaded by softmax.
TMEM_S0_OFFSET = 0
TMEM_VEC0_OFFSET = 0
TMEM_P0_OFFSET = 32
TMEM_O0_OFFSET = 128


# -------------------------
# 2) shared storage fields
# -------------------------
SHARED_STORAGE_FIELDS = r'''
@cute.struct
class SS:
    q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage * 2]
    kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage * 2]
    mma_s_bar: cute.struct.MemRange[cutlass.Int64, self.mma_softmax_stage * 2]
    s_corr_bar: cute.struct.MemRange[cutlass.Int64, self.softmax_corr_stage * 2]
    mma_corr_bar: cute.struct.MemRange[cutlass.Int64, self.mma_corr_stage * 2]
    corr_epi_bar: cute.struct.MemRange[cutlass.Int64, self.epi_stage * 2]
    tmem_dealloc: cute.struct.MemRange[cutlass.Int64, 1]
    holding: cutlass.Int32
'''


# -------------------------
# 3) pipeline creation
# -------------------------
PIPELINE_CREATION = r'''
def cg_threads(n: int):
    return pipeline.CooperativeGroup(pipeline.Agent.Thread, n)

q_prod, q_cons = pipeline.PipelineTmaUmma.create(
    barrier_storage=st.q_bar.data_ptr(),
    num_stages=self.q_stage,
    producer_group=cg_threads(1),
    consumer_group=cg_threads(1),
    tx_count=self.q_tx_bytes,
    cta_layout_vmnk=cl_vmnk,
    defer_sync=True,
).make_participants()

kv_prod, kv_cons = pipeline.PipelineTmaUmma.create(
    barrier_storage=st.kv_bar.data_ptr(),
    num_stages=self.kv_stage,
    producer_group=cg_threads(1),
    consumer_group=cg_threads(1),
    tx_count=self.kv_tx_bytes,
    cta_layout_vmnk=cl_vmnk,
    defer_sync=True,
).make_participants()

# MMA publishes S; softmax consumes S and releases the handle only after P is ready.
mma_s_prod, mma_s_cons = pipeline.PipelineUmmaAsync.create(
    barrier_storage=st.mma_s_bar.data_ptr(),
    num_stages=self.mma_softmax_stage,
    producer_group=cg_threads(1),
    consumer_group=cg_threads(32 * len(self.softmax_warp_ids)),
    cta_layout_vmnk=cl_vmnk,
    defer_sync=True,
).make_participants()

# Softmax publishes vec[row] = [old_max,new_max] each tile and [row_sum,row_max] at final.
s_corr_prod, s_corr_cons = pipeline.PipelineAsync.create(
    barrier_storage=st.s_corr_bar.data_ptr(),
    num_stages=self.softmax_corr_stage,
    producer_group=cg_threads(32 * len(self.softmax_warp_ids)),
    consumer_group=cg_threads(32 * len(self.correction_warp_ids)),
).make_participants()

# MMA publishes O after each PV.  Correction consumes O for online rescale / final epilog.
mma_corr_prod, mma_corr_cons = pipeline.PipelineUmmaAsync.create(
    barrier_storage=st.mma_corr_bar.data_ptr(),
    num_stages=self.mma_corr_stage,
    producer_group=cg_threads(1),
    consumer_group=cg_threads(32 * len(self.correction_warp_ids)),
    cta_layout_vmnk=cl_vmnk,
    defer_sync=True,
).make_participants()

# Correction publishes final converted O in SMEM; epilogue warp TMA-stores it.
corr_epi_prod, corr_epi_cons = pipeline.PipelineAsync.create(
    barrier_storage=st.corr_epi_bar.data_ptr(),
    num_stages=self.epi_stage,
    producer_group=cg_threads(32 * len(self.correction_warp_ids)),
    consumer_group=cg_threads(32),
).make_participants()

# TMEM allocation/retrieve participants are softmax + correction + MMA only.
tmem_bar = pipeline.NamedBarrier(
    barrier_id=2,
    num_threads=32 * len((*self.softmax_warp_ids, *self.correction_warp_ids, self.mma_warp_id)),
)
tmem = utils.TmemAllocator(
    st.holding.ptr,
    barrier_for_retrieve=tmem_bar,
    allocator_warp_id=self.softmax_warp_ids[0],
    is_two_cta=cute.size(qk_mma.thr_id.shape) == 2,
    two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr,
)

# Deallocation mbarrier is NOT a named barrier.  Only softmax+correction arrive;
# MMA waits and deallocates after relinquish_tmem_alloc_permit().
if warp_idx == self.empty_warp_id:
    cute.arch.mbarrier_init(
        st.tmem_dealloc.data_ptr(),
        32 * len((*self.softmax_warp_ids, *self.correction_warp_ids)),
    )
    cute.arch.mbarrier_init_fence()
'''


# -------------------------
# 4) correction helpers
# -------------------------
CORRECTION_HELPERS = r'''
@cute.jit
def correction_rescale(self, pv_thr, tOtO, scale: Float32):
    """Correction warpgroup: O[row,:] *= scale[row].  O stays in TMEM."""
    cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
    tOcO = pv_thr.partition_C(cO)

    corr_tile_size = 16
    tmem_load_atom = cute.make_copy_atom(
        tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)),
        self.pv_acc_dtype,
    )
    tmem_store_atom = cute.make_copy_atom(
        tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)),
        self.pv_acc_dtype,
    )

    tOtO_i_layout = cute.composition(tOtO.layout, cute.make_layout((128, corr_tile_size)))
    tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size)))
    tOtO_i = cute.make_tensor(tOtO.iterator, tOtO_i_layout)
    tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout)

    tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_i)
    tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_i)

    tidx, _, _ = cute.arch.thread_idx()
    thread_idx = tidx % (32 * len(self.correction_warp_ids))
    thr_load = tiled_tmem_load.get_slice(thread_idx)
    thr_store = tiled_tmem_store.get_slice(thread_idx)

    tTMEM_LOADtO = thr_load.partition_S(tOtO_i)
    tTMEM_LOADcO = thr_load.partition_D(tOcO_i)
    tTMEM_STOREtO = thr_store.partition_D(tOtO_i)

    o_col_tiles = self.pv_mma_tiler[1] // corr_tile_size
    tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, o_col_tiles), self.pv_acc_dtype)
    for i in range(o_col_tiles):
        tTMrO_i_ = tTMrO[None, i]
        tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0]))
        tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout)
        tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout)
        tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout)
        cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i)
        for j in cutlass.range(cute.size(tTMrO_i), vectorize=True):
            tTMrO_i[j] = tTMrO_i[j] * scale
        cute.copy(tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i)


@cute.jit
def correction_epilog(self, pv_thr, tOtO, scale: Float32, sO):
    """Correction final: load O from TMEM, normalize/convert, write SMEM."""
    cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
    corr_tile_size = 32 * 8 // self.o_dtype.width

    tOsO = pv_thr.partition_C(sO)
    tOcO = pv_thr.partition_C(cO)
    tOtO_i = cute.logical_divide(tOtO, cute.make_layout((128, corr_tile_size)))
    tOcO_i = cute.logical_divide(tOcO, cute.make_layout((128, corr_tile_size)))
    tOsO_i = cute.logical_divide(tOsO, cute.make_layout((128, corr_tile_size)))

    tidx, _, _ = cute.arch.thread_idx()
    thread_idx = tidx % (32 * len(self.correction_warp_ids))
    epi_subtile = (self.epi_tile[0], corr_tile_size)

    tmem_copy_atom = utils.sm100.get_tmem_load_op(
        self.pv_mma_tiler,
        self.c_layout,
        self.o_dtype,
        self.pv_acc_dtype,
        epi_subtile,
        use_2cta_instrs=False,
    )
    tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_i[(None, None), 0])
    thr_tmem_load = tiled_tmem_load.get_slice(thread_idx)

    smem_copy_atom = utils.sm100.get_smem_store_op(
        self.c_layout,
        self.o_dtype,
        self.pv_acc_dtype,
        tiled_tmem_load,
    )
    tiled_smem_store = cute.make_tiled_copy_D(smem_copy_atom, tiled_tmem_load)

    tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_i[(None, None), None])
    tTMEM_LOADsO = thr_tmem_load.partition_D(tOsO_i[(None, None), None])
    tTMEM_LOADoO = thr_tmem_load.partition_D(tOcO_i[(None, None), None])

    for i in range(self.pv_mma_tiler[1] // corr_tile_size):
        tTMrO = cute.make_rmem_tensor(tTMEM_LOADoO[None, 0, 0, i].shape, self.pv_acc_dtype)
        cute.copy(tiled_tmem_load, tTMEM_LOADtO[None, 0, 0, i], tTMrO)
        for j in cutlass.range(cute.size(tTMrO), vectorize=True):
            tTMrO[j] = tTMrO[j] * scale
        tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype)
        tSMrO.store(tTMrO.load().to(self.o_dtype))
        cute.copy(tiled_smem_store, tSMrO, tTMEM_LOADsO[None, 0, 0, i])

    cute.arch.fence_proxy("async.shared", space="cta")
'''


# -------------------------
# 5) native Stage-C flow
# -------------------------
NATIVE_STAGE_C_FLOW = r'''
# MMA loop pseudocode, replacing softmax_done_bar/pv_done_bar:
if is_mma:
    tmem.wait_for_alloc()
    ... load Q/K/V pipeline setup ...
    s_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_softmax_stage)
    o_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_corr_stage)
    mma_s_prod.producer_acquire(s_state)
    mma_corr_prod.producer_acquire(o_state)
    for kt in range(n_kv_tiles):
        # QK -> S
        ... cute.gemm(qk_mma, ..., tStS0) ...
        cute.arch.fence_view_async_tmem_store()
        mma_s_prod.producer_commit(s_state)
        s_state.advance()

        # IMPORTANT: there is no named barrier here.  The softmax consumer releases
        # the same S handle only after P has been stored, so this is P-ready.
        if kt + 1 < n_kv_tiles:
            mma_s_prod.producer_acquire(s_state)

        # PV -> O
        ... cute.gemm(pv_mma, tOtO0, tOrP0, V, tOtO0) ...
        cute.arch.fence_view_async_tmem_store()
        mma_corr_prod.producer_commit(o_state)
        o_state.advance()
        if kt + 1 < n_kv_tiles:
            mma_corr_prod.producer_acquire(o_state)
    mma_s_prod.producer_tail(s_state)
    mma_corr_prod.producer_tail(o_state)

    cute.arch.relinquish_tmem_alloc_permit()
    cute.arch.mbarrier_wait(st.tmem_dealloc.data_ptr(), 0)
    tmem_ptr = cute.arch.retrieve_tmem_ptr(self.qk_acc_dtype, alignment=16, ptr_to_buffer_holding_addr=st.holding.ptr)
    cute.arch.dealloc_tmem(tmem_ptr, Int32(self.num_tmem_alloc_cols))


# Softmax loop:
if is_softmax:
    tmem.allocate(self.num_tmem_alloc_cols)
    tmem.wait_for_alloc()
    vec_handle = s_corr_prod.acquire_and_advance()
    row_max = -Float32.inf
    row_sum = Float32(0.0)
    for kt in range(n_kv_tiles):
        si = mma_s_cons.wait_and_advance()
        # load S, compute old/new row max, store vec=[old_max,new_max]
        # compute P=exp2((S-new_max)*scale), store P to TMEM
        # update row_sum
        vec_handle.commit()
        si.release()              # P-ready signal to MMA
        vec_handle = s_corr_prod.acquire_and_advance()
    # final vec=[row_sum,row_max]
    final_si = mma_s_cons.wait_and_advance()
    ... store final vec ...
    vec_handle.commit()
    s_corr_prod.acquire()         # balance final pipe step
    final_si.release()
    cute.arch.mbarrier_arrive(st.tmem_dealloc.data_ptr())


# Correction loop:
if is_correction:
    tmem.wait_for_alloc()
    first_vec = s_corr_cons.wait_and_advance()
    first_vec.release()           # first tile has no previous O to rescale
    for kt in range(n_kv_tiles - 1):
        vec = s_corr_cons.wait_and_advance()
        ... load [old_max,new_max] ...
        scale = exp2(scale_log2 * (old_max - new_max))
        o = mma_corr_cons.wait_and_advance()
        self.correction_rescale(pv_thr, tOtO0, scale)
        cute.arch.fence_view_async_tmem_store()
        o.release()
        vec.release()
    final_vec = s_corr_cons.wait_and_advance()
    ... load [row_sum,row_max] ...
    final_vec.release()
    final_o = mma_corr_cons.wait_and_advance()
    epi = corr_epi_prod.acquire_and_advance()
    self.correction_epilog(pv_thr, tOtO0, self.scale_output / row_sum, sC[(None, None, epi.index)])
    final_o.release()
    epi.commit()
    cute.arch.mbarrier_arrive(st.tmem_dealloc.data_ptr())


# Epilogue warp:
if is_epi:
    h = corr_epi_cons.wait_and_advance()
    cute.copy(tma_c, tCsC[(None, h.index)], tCgC_epi[(None, 0)])
    cute.arch.cp_async_bulk_commit_group()
    cute.arch.cp_async_bulk_wait_group(0, read=True)
    h.release()
'''
