idesc = sm100_desc.mma_op_to_idesc(mma_op)
# Encodes: shape (M, N, K), dtype (A, B, C), operand source (SMEM/TMEM)

def gemm_w_idx(tiled_mma, acc, tCrA, tCrB, A_idx=None, B_idx=None, zero_init=False, swap_AB=False):
    """High-level GEMM using cute.gemm().
    
    Uses CuTe's gemm abstraction. No PTX inline asm.
    A_idx/B_idx index into the stage dimension.
    """

def gemm(tiled_mma, acc, tCrA, tCrB, zero_init=False):
    """Simple high-level GEMM using cute.gemm().
    
    Iterates over K dimension with loop.
    """

def gemm_ptx_w_idx(tiled_mma, acc, tCrA, tCrB, sA, sB, A_idx=None, B_idx=None, 
                   zero_init=False, cta_group=1, **kwargs):
    """PTX-based GEMM, delegates to gemm_ptx_partial.
    
    Converts tiled_mma to mma_atom, then calls gemm_ptx_partial.
    """

def gemm_ptx(op, acc, tCrA, tCrB, sA, sB, zero_init=False):
    """Full PTX GEMM with inline asm.
    
    Computes smem descriptors inside the function.
    Emits one tcgen05.mma per K tile.
    """

def gemm_ptx_loop(op, acc, tCrA, tCrB, sA, sB, zero_init=False):
    """PTX GEMM with precomputed offset differences.
    
    Precomputes offset_a_diff and offset_b_diff for efficient loop.
    Uses add.u32 instead of full descriptor recomputation.
    """

def gemm_ptx_partial(op, acc_tmem_addr, tCrA, tCrB, sA, sB, mbar_ptr=None, 
                     mbar_phase=None, split_arrive=None, zero_init=False, cta_group=1):
    """Partial PTX GEMM with split_arrive support.
    
    split_arrive: Number of K tiles to process before mbarrier wait.
    Used for pipelining: process 3/4 of K tiles, wait for softmax, process 1/4.
    """

# Given inputs for P×V:
op                  = pv_mma_op          # MMA operation
acc_tmem_addr       = 0x1000             # $3: TMEM address for O
tCrA                = P register frag    # A is in TMEM
tCrB                = V register frag    # B is from SMEM
sA                  = None               # A is in TMEM, not SMEM
sB                  = V in SMEM          # B is in SMEM
mbar_ptr            = 0x5000             # $4: mbarrier address
mbar_phase          = 1                  # $5: mbarrier phase
split_arrive        = 96                 # Split at 96 elements (3/4 of 128)
zero_init           = False
cta_group           = 2

def gemm_ptx_partial1(op, acc_tmem_addr, tCrA, tCrB, sA_base_addr_for_desc, 
                      sA_addr_offset_for_desc, sA_stage, sB_base_addr_for_desc,
                      sB_addr_offset_for_desc, sB_stage, sA_layout, sB_layout,
                      sA_swizzle, sB_swizzle, zero_init=False):
    """Variant with precomputed base addresses and stage indices.
    
    Uses mad.lo.u32 to compute descriptor: base + stage * offset.
    """

def gemm_ptx_precomputed(acc_tmem_addr, smem_desc_start_a, smem_desc_start_b, idesc,
                         smem_desc_base_a, smem_desc_base_b, tCrA_layout, tCrB_layout,
                         mbar_ptr=None, mbar_phase=None, zero_init=False, cta_group=1):
    """PTX GEMM with precomputed descriptors.
    
    Descriptors are passed as arguments instead of computed inside.
    Supports mbarrier wait for pipelining.
    """

def gemm_ptx_precomputed_varname(acc_tmem_addr, smem_desc_start_b, smem_desc_base_b,
                                 tCrB_layout, smem_var_name_prefix, idesc_var_name,
                                 smem_offset, zero_init=False, cta_group=1):
    """PTX GEMM using pre-declared PTX register variables.
    
    Requires declare_ptx_smem_desc and declare_ptx_idesc to be called first.
    Uses PTX variables: {smem_var_name_prefix}_<k> and {idesc_var_name}.
    """

# Without precompute: Compute A descriptor from scratch for each stage
smem_desc_base_a = smem_desc_base_from_tensor(sQ[stage])  # Every call
smem_desc_start_a = make_smem_desc_start_addr(sQ[stage])  # Every call

# With precompute: Compute A descriptor once, adjust with stride
# declare_ptx_smem_desc(sQ[stage=1])  # Once at kernel entry
# Inside gemm_ptx_precomputed_varname:
mov.b64 {smem_desc_a_lo, smem_desc_a_hi}, fa_fwd_q_smem_desc_0;  # LOAD precomputed
add.s32 smem_desc_a_lo, smem_desc_a_lo, -16;                     # Just adjust stride
mov.b64 fa_fwd_q_smem_desc_0, {smem_desc_a_lo, smem_desc_a_hi};  # STORE back

# Given inputs:
acc_tmem_addr       = 0x1000        # $2 in PTX (TMEM address)
smem_desc_start_b   = 0x2000        # $0 in PTX (B start address)
smem_desc_base_b    = 0xBEEF_DEAD_0000_5678  # 64-bit B base
tCrB_layout         # shape[2] = 4 → num_k_tile = 4
smem_var_name_prefix= "fa_fwd_q_smem_desc"
idesc_var_name      = "fa_fwd_qk_mma_idesc"
smem_offset         = -16           # hardcoded offset for stage 0
zero_init           = True          # $1 in PTX
cta_group           = 2             # hardcoded in instruction

# Step 1: Split B base
smem_desc_base_b_lo = 0x0000_5678
smem_desc_b_hi      = 0xBEEF_DEAD

# Step 2: Compute B offsets from layout
offset_b = [crd2idx((0,0,k), tCrB_layout) for k in range(4)]
# = [0, 32, 64, 96]

# Step 3: Combine B start + base_lo
smem_desc_start_b_lo = 0x0000_5678 | 0x2000 = 0x0000_7678

def declare_ptx_smem_desc(smem_desc_start_a, smem_desc_base_a, tCrA_layout, var_name_prefix):
    """Declare PTX registers for smem descriptors.
    
    Inputs:
    - smem_desc_start_a: 32-bit start address offset (from iterator)
    - smem_desc_base_a:  64-bit base descriptor (from layout + swizzle + major)
    - tCrA_layout:       Layout to compute per-K-tile offsets
    - var_name_prefix:   PTX variable name prefix
    """

# Given inputs:
smem_desc_start_a = 0x1234        # $0 in inline asm (32-bit)
smem_desc_base_a  = 0xDEAD_BEEF_CAFE_0000  # 64-bit, split to lo/hi
tCrA_layout       # Used to compute offset_a[k] for each K tile
var_name_prefix   = "fa_fwd_q_smem_desc"

# Step 1: Split base into lo/hi
smem_desc_base_a_lo = 0xCAFE_0000   # lower 32 bits
smem_desc_a_hi      = 0xDEAD_BEEF   # upper 32 bits

# Step 2: Combine start + base_lo
smem_desc_start_a_lo = smem_desc_base_a_lo | smem_desc_start_a
# = 0xCAFE_0000 | 0x1234 = 0xCAFE_1234

# Step 3: Compute offsets from layout
offset_a = [crd2idx((0,0,k), tCrA_layout) for k in range(num_k_tile)]
# = [0, 16, 32, 48, ...]  # depends on layout stride

# Generated PTX:
.reg .b32 fa_fwd_q_smem_desc_lo;              # temp register
.reg .b64 fa_fwd_q_smem_desc_<4>;             # array of 4 descriptors

# k=0: combine $0 (start_lo) with base_hi
mov.b64 fa_fwd_q_smem_desc_0, {$0, 0xDEADBEEF};
# fa_fwd_q_smem_desc_0 = {0xCAFE1234, 0xDEADBEEF}

# k=1: add offset[1] to start, combine with base_hi
add.s32 fa_fwd_q_smem_desc_lo, $0, 0x10;      # offset[1] = 16 = 0x10
mov.b64 fa_fwd_q_smem_desc_1, {fa_fwd_q_smem_desc_lo, 0xDEADBEEF};
# fa_fwd_q_smem_desc_1 = {0xCAFE1244, 0xDEADBEEF}

# k=2: add offset[2] to start, combine with base_hi
add.s32 fa_fwd_q_smem_desc_lo, $0, 0x20;      # offset[2] = 32 = 0x20
mov.b64 fa_fwd_q_smem_desc_2, {fa_fwd_q_smem_desc_lo, 0xDEADBEEF};
# fa_fwd_q_smem_desc_2 = {0xCAFE1254, 0xDEADBEEF}

# k=3: add offset[3] to start, combine with base_hi
add.s32 fa_fwd_q_smem_desc_lo, $0, 0x30;      # offset[3] = 48 = 0x30
mov.b64 fa_fwd_q_smem_desc_3, {fa_fwd_q_smem_desc_lo, 0xDEADBEEF};
# fa_fwd_q_smem_desc_3 = {0xCAFE1264, 0xDEADBEEF}

def declare_ptx_idesc(op, var_name):
    """Declare PTX register for idesc.
    
    Emits:
      .reg .b32 {var_name};
      mov.b32 {var_name}, {idesc};
    """

def i64_to_i32x2(i):
    """Split 64-bit int into (lo, hi) 32-bit tuple."""
    return i & 0xFFFF_FFFF, (i >> 32) & 0xFFFF_FFFF

# Step 1: Get MMA operations from TiledMma
qk_mma_op, pv_mma_op = tiled_mma_qk.op, tiled_mma_pv.op

# Step 2: Convert ops to idesc (instruction descriptors)
qk_mma_idesc, pv_mma_idesc = (
    sm100_desc.mma_op_to_idesc(qk_mma_op),   # For Q×K→S
    sm100_desc.mma_op_to_idesc(pv_mma_op),   # For P×V→O
)

# Step 3: Compute smem descriptor bases
q_smem_base = sm100_desc.smem_desc_base_from_tensor(sQ, sm100_desc.Major.K)
k_smem_base = sm100_desc.smem_desc_base_from_tensor(sK, sm100_desc.Major.K)
v_smem_base = sm100_desc.smem_desc_base_from_tensor(sV, sm100_desc.Major.MN)

# Step 4: Compute smem descriptor start addresses for each Q stage
q_smem_start = [
    sm100_desc.make_smem_desc_start_addr(sQ[None, None, None, stage].iterator)
    for stage in range(self.q_stage)  # stage 0 and 1
]

# Declare PTX registers for Q smem descriptors (one per K tile)
sm100_utils.declare_ptx_smem_desc(
    q_smem_start[self.q_stage - 1],  # Start address for last Q stage
    q_smem_base,                      # Base address for Q
    tSrQ[None, None, None, 0].layout, # Q layout for descriptor computation
    var_name_prefix="fa_fwd_q_smem_desc",
)

# Declare PTX registers for idesc
sm100_utils.declare_ptx_idesc(qk_mma_op, var_name="fa_fwd_qk_mma_idesc")
sm100_utils.declare_ptx_idesc(pv_mma_op, var_name="fa_fwd_pv_mma_idesc")

# Compute Q stage stride for descriptor offset
sQ_stage_stride = (sQ.layout.stride[-1] * sQ.element_type.width // 8) >> 4

# Setup GEMM functions for Q×K→S (one per Q stage)
gemm_Si = [
    partial(
        sm100_utils.gemm_ptx_precomputed_varname,
        self.tmem_s_offset[stage],           # TMEM address for S[stage]
        smem_desc_base_b=k_smem_base,        # K smem base
        tCrB_layout=tSrK[None, None, None, 0].layout,  # K layout
        smem_var_name_prefix="fa_fwd_q_smem_desc",     # Pre-declared Q descriptors
        idesc_var_name="fa_fwd_qk_mma_idesc",          # Pre-declared idesc
        smem_offset=-sQ_stage_stride if stage == 0 else sQ_stage_stride,  # Stage offset
        zero_init=True,                      # First K tile zeros accumulator
        cta_group=self.cta_group_size,       # 2-CTA instruction
    )
    for stage in range(self.q_stage)
]

# Setup GEMM functions for P×V→O (one per Q stage)
gemm_Pi = [
    partial(
        sm100_utils.gemm_ptx_partial,
        pv_mma_op,                           # MMA operation
        self.tmem_o_offset[stage],           # TMEM address for O[stage]
        tOrP[None, None, None, stage],       # P register fragment
        sA=None,                             # A is in TMEM (P), not SMEM
        split_arrive=self.split_P_arrive,    # Split P arrival for pipelining
        cta_group=self.cta_group_size,       # 2-CTA instruction
    )
    for stage in range(self.q_stage)
]

# Inside the MMA loop, for each K/V block:
for k_block in range(n_block_min, n_block_max):
    # Wait for Q and K to be ready
    pipeline_q.consumer_wait(q_consumer_state)
    pipeline_kv.consumer_wait(kv_consumer_state)
    
    # Execute Q×K→S for current Q stage
    stage = q_mma_phase
    gemm_Si[stage]()  # Calls gemm_ptx_precomputed_varname
    
    # Signal S is ready for softmax
    pipeline_s_p_o.producer_commit(P_full_O_rescaled_phase)

# After softmax computes P:
for k_block in range(n_block_min, n_block_max):
    # Wait for V to be ready
    pipeline_kv.consumer_wait(kv_consumer_state)
    
    # Execute P×V→O for current Q stage
    stage = q_mma_phase
    gemm_Pi[stage]()  # Calls gemm_ptx_partial
    
    # Signal O is ready
    pipeline_o_acc.producer_commit(...)

# 1. Setup (once at kernel entry)
idesc = sm100_desc.mma_op_to_idesc(mma_op)
smem_base = sm100_desc.smem_desc_base_from_tensor(tensor, major)
smem_start = sm100_desc.make_smem_desc_start_addr(iterator)

# 2. Declare PTX registers
declare_ptx_smem_desc(smem_start, smem_base, layout, prefix)
declare_ptx_idesc(mma_op, var_name)

# 3. Create GEMM function
gemm_fn = partial(gemm_ptx_precomputed_varname, acc_addr, smem_start_b, 
                  smem_base_b, layout_b, prefix, var_name, offset, ...)

# 4. Execute GEMM (in loop)
gemm_fn()  # Emits tcgen05.mma per K tile