################################################################################
#
# Copyright (C) 2022-2026 Advanced Micro Devices, Inc. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell cop-
# ies of the Software, and to permit persons to whom the Software is furnished
# to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IM-
# PLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNE-
# CTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
################################################################################

from rocisa.code import Label, Module, RegSet, TextBlock
from rocisa.container import SMEMModifiers, VOP3PModifiers, MUBUFModifiers, \
  SDWAModifiers, replaceHolder, EXEC, VCC, vgpr, sgpr, ContinuousRegister, mgpr
from rocisa.enum import CvtType, HighBitSel, RoundType, SaturateCastType, SelectBit
from rocisa.instruction import BufferAtomicAddF32, BufferAtomicCmpswapB32, \
  BufferAtomicCmpswapB64, BufferStoreB16, BufferStoreB32, BufferStoreB64, BufferStoreB128, DSBPermuteB32, FlatAtomicCmpswapB32, \
  SAddCU32, SAddU32, SAndB32, \
  SAndB64, SAtomicDec, SBarrier, SBranch, SCBranchExecNZ, SCBranchExecZ, \
  SCBranchSCC0, SCBranchSCC1, SCmpGtU32, SCmpKGtU32, SCSelectB32, SCmpEQI32, SCmpEQU32, SCmpGtI32, SCmpLeI32, SMinU32, SEndpgm, \
  SLShiftLeftB32, SLShiftLeftB64, SLShiftRightB32, SLShiftRightB64, SMovB32, SMovB64, SMulI32, \
  SNop, SOrB32, SOrB64, SOrSaveExecB32, SOrSaveExecB64, SSleep, SSubI32, SSubU32, \
  SSwapPCB64, SWaitCnt, SWaitAlu, VAShiftRightI32, VAddCCOU32, VAddCOU32, VAddF32, VAddF64, \
  VAddI32, VAddPKF16, VAddPKF32, VAddU32, VBfeI32, VCmpEQU32, VCmpGEI32, VCmpGtU32, \
  VCmpNeU32, VCmpNeU64, VCndMaskB32, VCvtBF8toF32, VCvtF16toF32, VCvtF32toF16, VCvtF32toI32, \
  VCvtFP8toF32, VCvtI32toF32, VCvtPkBF8toF32, VCvtPkF32toBF16, VCvtPkF32toFP16, VCvtPkFP8toF32, \
  VFmaF32, VFmaF64, VFmaMixF32, VAndB32, VLShiftLeftB32, VPermlane16SwapB32, VPermlane32SwapB32, \
  VLShiftRightB32, VMacF32, VMadMixF32, VMaxF32, VMovB32, VMovB64, VMulF32, VMulF64, \
  VMulLOU32, VMulPKF16, VMulPKF32, VPackF16toB32, VReadfirstlaneB32, VRndneF32, VCvtBF16toFP32
from rocisa.functions import vectorStaticMultiply

from ..Common import DataDirection, SemanticVersion, isSubtileMultiDU
from ..Common.DataType import DataType
from ..Component import GlobalWriteComponents
from ..Component import Component
from ..SolutionStructs import Solution
from ..Activation import ActivationModule
from ..AsmStoreState import StoreState
from ..AsmAddressCalculation import AddrCalculation
from ..Components.PackData import formatting, PackData_F16, PackData_BF16, PackData_FLOAT8, PackData_FLOAT8_fnuz
from rocisa.instruction import ECvtF16toF32, ECvtPkFP8toF32, ECvtPkBF8toF32
from ..KernelWriterModules import hasSequentialValuC

from math import ceil, log2


def _scmpGtU32(writer, src, imm, comment=""):
    """ISA-aware scalar compare: s_cmpk_gt_u32 when available, else s_cmp_gt_u32 via temp SGPR."""
    if writer.states.asmCaps["HasSCMPK"]:
        return SCmpKGtU32(src=src, simm16=imm, comment=comment)
    else:
        module = Module("scmpGtU32")
        tmpSgpr = writer.sgprPool.checkOut(1, preventOverflow=False)
        module.add(SMovB32(dst=sgpr(tmpSgpr), src=imm))
        module.add(SCmpGtU32(src0=src, src1=sgpr(tmpSgpr), comment=comment))
        writer.sgprPool.checkIn(tmpSgpr)
        return module

class GlobalWriteBatchComponent(GlobalWriteComponents):
  kernel = {"ProblemType": {"OperationType": "GEMM" }}
  def __call__(self, kernel: Solution, tPA, tPB, activation: ActivationModule, ss: StoreState, \
    batchIdx, applyAlpha, beta, edge, atomic, gwvw, atomicW, \
    batchElements, addrE, addrD, addrC, addrBias, addrScaleAVec, addrScaleBVec, addrScaleAlphaVec, isLocalBarrierInit: bool, \
    tmpVgpr, tmpVgprDynamic, cvtVgprStruct, activationSetPCStruct, activationTypeStr, batchElementSgprs, tmpSgpr, codeAccVgprRead, \
    codeMulAlpha, packdata, parentWriter, factorDim, amdClangVersion: SemanticVersion, numBatches: int,
    inter_iter_rowInc: int = 0, direct_next_rowInc: int = 0) -> Module:
    return GlobalWriteBatchWriter(kernel, tPA, tPB, activation, ss, batchIdx, applyAlpha, \
      beta, edge, atomic, gwvw, atomicW, \
      batchElements, addrE, addrD, addrC, addrBias, addrScaleAVec, addrScaleBVec, addrScaleAlphaVec, isLocalBarrierInit, \
      tmpVgpr, tmpVgprDynamic, cvtVgprStruct, activationSetPCStruct, activationTypeStr, batchElementSgprs, tmpSgpr, \
      codeAccVgprRead, codeMulAlpha, packdata, parentWriter, factorDim, amdClangVersion, numBatches,
      inter_iter_rowInc, direct_next_rowInc).emit()

class GlobalWriteBatchWriter:
  def __init__(self, kernel: Solution, tPA, tPB, activation: ActivationModule, ss: StoreState, \
    batchIdx, applyAlpha, beta, edge, atomic, gwvw, atomicW, \
    batchElements, addrE, addrD, addrC, addrBias, addrScaleAVec, addrScaleBVec, addrScaleAlphaVec, isLocalBarrierInit: bool, \
    tmpVgpr, tmpVgprDynamic, cvtVgprStruct, activationSetPCStruct, activationTypeStr, batchElementSgprs, tmpSgpr, codeAccVgprRead, \
    codeMulAlpha, packdata, parentWriter, factorDim, amdClangVersion: SemanticVersion, numBatches: int,
    inter_iter_rowInc: int = 0, direct_next_rowInc: int = 0):
    self.kernel = kernel
    self.tPA    = tPA
    self.tPB    = tPB
    self.activation = activation
    self.ss = ss
    self.batchIdx = batchIdx
    self.applyAlpha = applyAlpha
    self.beta = beta
    self.edge = edge
    self.atomic = atomic
    self.gwvw = gwvw
    self.atomicW = atomicW
    self.batchElements = batchElements
    self.addrE    = addrE
    self.addrD    = addrD
    self.addrC    = addrC
    self.addrBias = addrBias
    self.addrScaleAVec = addrScaleAVec
    self.addrScaleBVec = addrScaleBVec
    self.addrScaleAlphaVec = addrScaleAlphaVec
    self.isLocalBarrierInit  = isLocalBarrierInit
    self.activationSetPCStruct = activationSetPCStruct
    self.activationTypeStr     = activationTypeStr
    self.tmpVgpr = tmpVgpr.idx
    self.tmpVgprSize = tmpVgpr.size
    self.tmpVgprDynamic = None
    if tmpVgprDynamic:
      self.tmpVgprDynamic = tmpVgprDynamic.idx
      self.tmpVgprDynamicSize = tmpVgprDynamic.size
    self.cvtVgprStruct = cvtVgprStruct
    self.batchElementSgprs = batchElementSgprs
    self.tmpSgpr = tmpSgpr
    self.codeAccVgprRead = codeAccVgprRead
    self.codeMulAlpha = codeMulAlpha
    self.packdata     = packdata
    self.parentWriter = parentWriter
    self.storesIssued = 0
    self.factorDim = factorDim
    self.amdClangVersion = amdClangVersion

    # Stateful tracking for N-group OOB guard deduplication (_emitSubtileOobGuard).
    # The outer loop iterates N-outer / M-inner, so all M elements within a fixed N
    # group share the same N guard result.  We emit the N s_cmp/s_cbranch only once
    # per N group and skip it for subsequent M elements in the same group.
    self._subtilePrevBlockIdxN = -1       # sentinel: no group seen yet
    self._subtileNGroupSkipLabel = None   # end-of-N-group label (M cbranch target)
    self._subtileAllStoresEndLabel = None # end-of-all-stores label (N cbranch target)
    self._subtileCloadPrevD1 = -1         # sentinel: last d1 group seen in C load guard
    self._subtilePendingSrdDInc = None    # deferred SrdD incToNextRow (emitted after N-group label)
    self._align8NMaskBlockIdxN = -1       # last blockIdxN for which N mask was computed
    self.numBatches = numBatches

    # CompactLoopStore: stash the "next batch's first elt rowInc" passed in via
    # `inter_iter_rowInc`. Used as the look-ahead fall-through value at the
    # LAST emitting elt of the batch (cross-batch transition), where the
    # intra-batch forward scan finds no further emitting elt.
    # 0 means no override (non-CLS / last batch / atomic).
    self.inter_iter_rowInc = inter_iter_rowInc
    self.direct_next_rowInc = direct_next_rowInc

    # Internal state for GlobalWriteBatch
    # 0 for None, 1 for WorkGroupReduction = False, 2 for WorkGroupReduction = True
    self.storeBiasD = 0
    if self.parentWriter.states.useBias == DataDirection.WRITE and \
      (not self.kernel["WorkGroupReduction"]) and \
      self.kernel["ProblemType"]["BiasSrc"] == "D":
      self.storeBiasD = 1

  @property
  def needsAccumToDestConversion(self) -> bool:
    """
    Check if accumulation values need to be converted to destination type:
    1. HighPrecisionAccumulate is enabled (accumulator precision > output precision)
       e.g., F32 accumulator -> FP16/BF16/FP8/BF8/I32/I8 output
    2. _GlobalAccumulation is not 'MultipleBuffer'

    When True, the pack/convert module will be generated to perform:
    - F32 -> FP16/BF16 packing
    - F32 -> FP8/BF8 conversion (with optional stochastic rounding)
    - F32 -> I32/I8 conversion and packing
    """
    return self.kernel["ProblemType"]["HighPrecisionAccumulate"] and \
           (self.kernel["_GlobalAccumulation"] != 'MultipleBuffer')

  @property
  def skipRearrangement(self) -> bool:
    """
    Check if we can skip v_mov_b32 rearrangement and use WMMA output registers directly.

    skipRearrangement changes store to read from (elementSumIdx - elementSumIdx[0]),
    but other operations (bias, activation, alpha) work on (elementSumIdx - startVgprValu).
    These positions are only the same when no operations modify the data.

    Safe to skip rearrangement only when:
    1. MIArchVgpr is True (MFMA/WMMA writes directly to arch VGPRs)
    2. hasSequentialValuC is True (WMMA output is already sequential)
    3. needsAccumToDestConversion is False (no pack module to do the rearrangement)
    4. Not a Beta path (Beta paths need separate alpha multiply + beta*C)
    5. Not complex (real/imag are interleaved in acc VGPRs, so the relative
       offset elementSumIdx[i]-elementSumIdx[0] does not locate the imag half)
    6. Single output tile (MIWaveTile == [1,1] and VectorWidth == 1). The
       i*gwvw offset only matches the arch-VGPR layout for one contiguous block;
       with multiple wave-tiles the acc registers are grouped per tile (e.g. TN +
       2x2 read the wrong registers).
    7. CompactLoopStore is off. CLS-loop store reads acc VGPRs via
       v_movrelsd_2_b32 indexing, which relies on the rearrangement layout.
    """
    if self.parentWriter.states.useBias == DataDirection.READ or \
       self.kernel.get("CompactLoopStore", False) or \
       self.kernel.get("ActivationFuncCall", False) or \
       self.applyAlpha or \
       self.kernel["ProblemType"].get("UseScaleAlphaVec", 0) or \
       self.kernel["ProblemType"].get("UseScaleAB", "") == "Vector" or \
       self.kernel["ProblemType"].get("UseScaleCD", False) or \
       self.kernel["ProblemType"]["DataType"].isComplex():
      return False

    miWaveTile = self.kernel.get("MIWaveTile", [1, 1])
    if miWaveTile[0] != 1 or miWaveTile[1] != 1 or \
       self.kernel.get("VectorWidthA", 1) != 1 or self.kernel.get("VectorWidthB", 1) != 1:
      return False

    return self.kernel["MIArchVgpr"] and \
           hasSequentialValuC(self.kernel) and \
           not self.needsAccumToDestConversion and \
           not self.beta

  @property
  def wavelen(self) -> int:
    return self.kernel["WavefrontSize"]

  @property
  def laneSGPRC(self) -> int:
    return self.parentWriter.states.laneSGPRCount

  @property
  def tmpS01(self):
    return self.tmpSgpr

  @property
  def tmpS23(self):
    return self.tmpS01 + self.laneSGPRC

  @property
  def debugConfig(self):
    return self.parentWriter.db

  @property
  def computeDataType(self) -> DataType:
    return self.kernel["ProblemType"]["ComputeDataType"]

  @property
  def destDataType(self) -> DataType:
    return self.kernel["ProblemType"]["DestDataType"]

  @property
  def moduleName(self):
    return "globalWriteBatch (Atomic)" if self.atomic else "globalWriteBatch (Non atomic)"

  def getEdgeMovInstType(self):
    return SMovB32 if self.wavelen == 32 else SMovB64

  def getEdgeOrInstType(self):
    return SOrB32 if self.wavelen == 32 else SOrB64

  def getEdgeAndInstType(self):
    return SAndB32 if self.wavelen == 32 else SAndB64

  def getSOrSaveExecType(self):
    return SOrSaveExecB32 if self.wavelen == 32 else SOrSaveExecB64

  @staticmethod
  def computeCLSLayout(kernel, numBatches: int):
    """Single source of truth for the CLS loop layout math.

    Returns (batchesPerCLSBody, iterCount, m0Step), all derived from ONE shared
    read of the MI output layout so the loop trip count and the M0 src stride
    can never drift apart:
      - batchesPerCLSBody: how many batches one CLS loop body covers. The body
        must align to the SrdD-advance period, else the loop would run s_add
        SrdD too many times.
      - iterCount: numBatches / batchesPerCLSBody (>= 1).
      - m0Step: stride of the CLS iter dim in the v_movrelsd_2_b32 src formula
        (added to M0 each iteration).
    Which dim is the CLS iter dim depends on the output layout:
      (a) outerTT1  > 1, VW1 == 1            : iter = wgIdx1, step = OPM*BM*BN*VW0*outerTT0*VW1
      (b) outerTT1 == 1, VW1 > 1, SS=False   : iter = vw1,    step = OPM*BM*BN*VW0*outerTT0
      (c) outerTT1 == 1, VW1 == 1, SS=True   : iter = tIdx,   step = NEPBS / inner_dims
    Store paths not covered by the CLS loop (non-MI, StoreRemap, StreamK) and
    non-divisible / non-regular layouts fall back to a single iteration
    (batchesPerCLSBody = numBatches), where m0Step is dead.
    """
    batchesPerBody = numBatches
    m0Step = 1
    if not kernel["EnableMatrixInstruction"]:
      return batchesPerBody, 1, m0Step

    VW0 = kernel["VectorWidthA"]
    VW1 = kernel["VectorWidthB"]
    outerTT0 = kernel["MIWaveTile"][0] // VW0
    outerTT1 = kernel["MIWaveTile"][1] // VW1
    miT  = min(kernel["MatrixInstM"], kernel["MatrixInstN"])
    miM_ = (kernel["MatrixInstM"] * kernel["MatrixInstBM"]) if (kernel["MatrixInstM"] == 4) else miT
    miN_ = (kernel["MatrixInstN"] * kernel["MatrixInstBN"]) if (kernel["MatrixInstN"] == 4) else miT
    matrixInstBM = 1 if (kernel["MatrixInstM"] == 4) else kernel["MatrixInstBM"]
    matrixInstBN = 1 if (kernel["MatrixInstN"] == 4) else kernel["MatrixInstBN"]
    OPM   = miM_ * miN_ // kernel["WavefrontSize"]
    NEPBS = kernel["NumElementsPerBatchStore"]

    if outerTT1 > 1 and VW1 == 1:
      m0Step = OPM * matrixInstBM * matrixInstBN * VW0 * outerTT0 * VW1
      if numBatches % outerTT1 == 0:
        batchesPerBody = max(1, numBatches // outerTT1)
    elif outerTT1 == 1 and VW1 > 1 and not kernel["SourceSwap"]:
      m0Step = OPM * matrixInstBM * matrixInstBN * VW0 * outerTT0
      if numBatches % VW1 == 0:
        batchesPerBody = max(1, numBatches // VW1)
    elif outerTT1 == 1 and VW1 == 1 and kernel["SourceSwap"]:
      inner_dims = VW0 * outerTT0 * matrixInstBM * matrixInstBN
      if NEPBS % inner_dims == 0:
        m0Step = max(1, NEPBS // inner_dims)
        batchesPerBody = max(1, ceil(numBatches / NEPBS))

    # StoreRemap / StreamK store paths are not covered by the CLS loop: force a
    # single iteration (body = all batches). m0Step is left as derived above to
    # preserve the original emitted SrdD step (it is dead when iterCount == 1).
    if kernel.get("StoreRemapVectorWidth", 0) != 0 or kernel.get("StreamK", 0) != 0:
      batchesPerBody = numBatches

    iterCount = max(1, numBatches // batchesPerBody)
    return batchesPerBody, iterCount, m0Step

  @staticmethod
  def computeBatchesPerCLSBody(kernel, numBatches: int) -> int:
    return GlobalWriteBatchWriter.computeCLSLayout(kernel, numBatches)[0]

  def _computeBatchesPerCLSBody(self) -> int:
    return GlobalWriteBatchWriter.computeCLSLayout(self.kernel, self.numBatches)[0]

  @staticmethod
  def computeCLSIterCount(kernel, numBatches: int) -> int:
    """CLS loop iter count = numBatches / batchesPerCLSBody. Minimum 1."""
    return GlobalWriteBatchWriter.computeCLSLayout(kernel, numBatches)[1]

  def _computeCLSIterCount(self) -> int:
    return GlobalWriteBatchWriter.computeCLSLayout(self.kernel, self.numBatches)[1]

  def _computeCLSLayout(self):
    return GlobalWriteBatchWriter.computeCLSLayout(self.kernel, self.numBatches)

  def emit(self) -> Module:
    assert self._checkAtomicPreconditions()
    module = Module(self.moduleName)
    self._prolog(module)
    # The bias/SAV drain barrier ordering is a multi-DU-only hardening that
    # prevents cross-wave LDS corruption from ds_bpermute. Multi-DU emits the
    # drain+barrier before the _emitAdd subtile stores; non-multi-DU emits
    # _emitAdd first.
    isMultiDU = isSubtileMultiDU(self.kernel)
    drainBiasSav = self.kernel.get("UseSubtileImpl") and \
       (self.parentWriter.states.useBias != DataDirection.NONE or \
        self.kernel["ProblemType"].get("UseScaleAlphaVec", 0))
    if not isMultiDU:
      self._emitAdd(module)
    if drainBiasSav:
      module.add(SWaitCnt(dscnt=0, comment="drain bias/SAV LDS reads"))
      module.add(SBarrier(comment="sync waves before subtile paired stores"))
    if isMultiDU:
      self._emitAdd(module)
    self._epilog(module)
    # CompactLoopStore CLS countdown tail: emit countdown + branch + s_endpgm at
    # END of the CLS-loop body (= last batch of batchesPerCLSBody). Gated by
    # CompactLoopStore so non-CLS .s matches baseline (no CLS tail emit).
    if self.kernel["CompactLoopStore"]:
      #     updateCoord1 = (edge or multi-packed)
      if self.direct_next_rowInc != 0 and self.ss.elementAddr:
        kw = self.parentWriter
        rowInc = self.direct_next_rowInc
        bufferStore = self.kernel["BufferStore"]
        updateCoord1 = self.edge or len(self.kernel["PackedC1IndicesX"]) > 1
        emit_coord1 = (not bufferStore) or updateCoord1
        emit_rowptr = (not self.ss.optSrdIncForRow) and bufferStore
        if emit_coord1 or emit_rowptr:
          module.addComment0("CLS look-ahead: next batch's row advance at end of this batch")
          # Reuse the canonical per-element advance helpers (emitCoord1Advance +
          # emitRowPtrAdvance) on a representative AddrCalculation instead of
          # re-implementing the coord1 / rowPtr advance inline. The look-ahead is just
          # the next batch's emitAddressSetupCode row-advance, produced from the same
          # source. Same spirit as the delayed-primer reuse of incrementToNextRow.
          addrCalc = self.ss.elementAddr[0]
        if emit_coord1:
          module.add(addrCalc.emitCoord1Advance(rowInc, self.tmpS01,
                     comment="coord1.la: coord1Vgpr += rowInc (look-ahead)",
                     scomment="rowInc look-ahead"))
        if emit_rowptr:
          # Reuse the canonical row-pointer advance (the SAME helper the per-element
          # advance uses in emitAddressSetupCode) instead of re-implementing the
          # cinRowPtr/coutRowPtrD adds inline. This also makes the look-ahead cover
          # coutRowPtrE / coutRowPtrBias / packed-C1 for free (same conditions).
          module.add(addrCalc.emitRowPtrAdvance(self.kernel, self.ss, self.tmpS01, rowInc, lookahead=True))

      clsLabel = getattr(self.ss, "_clsLoopLabel", None)
      if clsLabel is not None and (self._computeBatchesPerCLSBody() - 1 == self.batchIdx) and self.ss.elementAddr:
        module.add(SSubI32(dst=sgpr("CLSLoopCounter"), src0=sgpr("CLSLoopCounter"), src1=1))
        module.add(SCmpEQU32(src0=sgpr("CLSLoopCounter"), src1=0))
        module.add(SCBranchSCC0(clsLabel.getLabelName(), "loop while counter != 0"))
        # if not self.kernel["StreamK"] == 3:
        #   module.add(SEndpgm(comment="stop here after CLS loop"))
        self.ss._clsLoopLabel = None
    return module

  def globalStoreWait(self, elementIdx, waitCnter, vlcntTotalIssued, dscntTotalIssued, interleaveStoreVmcnt: bool):
    vlcnt = -1
    dscnt = -1
    vscnt = -1
    isSingleKernel = ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel") or self.kernel["StreamK"] > 0
    if interleaveStoreVmcnt:
      waitLocalLoadCnt = 0
      waitLocalLoadCntStrList = []
      waitLoadCnt = 0
      waitLoadCntStrList = []
      # Calculate global loads
      if self.beta:
        waitLoadCnt += self.betaLoadIssued[elementIdx]
        waitLoadCntStrList.append("%d (beta)"%self.betaLoadIssued[elementIdx])
      if self.loadE:
        waitLoadCnt += self.eLoadIssued[elementIdx]
        waitLoadCntStrList.append("%d (load E)"%self.eLoadIssued[elementIdx])
      if self.parentWriter.states.useGateResidual and elementIdx < len(self.gateLoadIssued):
        waitLoadCnt += self.gateLoadIssued[elementIdx]
        waitLoadCntStrList.append("%d (load Gate)"%self.gateLoadIssued[elementIdx])
      # Calculate local loads
      # UseSubtileImpl with bias/SAV: skip bias/SAV LDS loads from interleaved
      # waitcnt and rely on the batch-start barrier for LDS synchronization.
      subtileBarrierDrains = self.kernel.get("UseSubtileImpl") and \
        (self.parentWriter.states.useBias != DataDirection.NONE or \
         self.kernel["ProblemType"].get("UseScaleAlphaVec", 0))
      if self.parentWriter.states.useBias == DataDirection.READ and not subtileBarrierDrains:
        waitLocalLoadCnt += self.biasLoadIssued[elementIdx]
        waitLocalLoadCntStrList.append("%d (bias)"%self.biasLoadIssued[elementIdx])
      if (self.kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
        waitLocalLoadCnt += self.scaleAVecLoadIssued[elementIdx]
        waitLocalLoadCntStrList.append("%d (scaleAVec)"%self.scaleAVecLoadIssued[elementIdx])
        waitLocalLoadCnt += self.scaleBVecLoadIssued[elementIdx]
        waitLocalLoadCntStrList.append("%d (scaleBVec)"%self.scaleBVecLoadIssued[elementIdx])
      # Skip scaleAlphaVec when subtileBarrierDrains
      if self.kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel and not subtileBarrierDrains:
        waitLocalLoadCnt += self.scaleAlphaVecLoadIssued[elementIdx]
        waitLocalLoadCntStrList.append("%d (scaleAlphaVec)"%self.scaleAlphaVecLoadIssued[elementIdx])
      # Get vlcnt and dscnt
      vlcnt = vlcntTotalIssued - waitLoadCnt
      if waitCnter[0] > 0  or vlcnt != waitCnter[0] : # Check if global load issued > 0
        if waitCnter[0] == vlcnt: # No need to wait if the global load cnt doesn't change
          vlcnt = -1
        else:
          waitCnter[0] = vlcnt
      else:
        vlcnt = -1

      dscnt = dscntTotalIssued - waitLocalLoadCnt
      if waitCnter[1] > 0 or dscnt != waitCnter[1]: # Check if local load issued > 0
        if waitCnter[1] == dscnt: # No need to wait if the local load cnt doesn't change
          dscnt = -1
        else:
          waitCnter[1] = dscnt
      else:
        dscnt = -1
      # Get vscnt
      if vlcnt != -1 and not (self.parentWriter.states.asmCaps["SeparateVscnt"] or self.parentWriter.states.asmCaps["SeparateVMcnt"]):
          if self.kernel.get("UseSubtileImpl") and not self.kernel["GroupLoadStore"]:
            vscnt = 0
          else:
            vscnt = self.storesIssued if not self.kernel["GroupLoadStore"] else 0
      if (vlcnt != -1) or (vscnt != -1) or (dscnt != -1):
        # Get comment
        comment = ""
        if vlcnt != -1:
          tmp = ""
          for cntStr in waitLoadCntStrList:
            tmp += " - %s"%cntStr
          comment = "vlcnt(%s) = %d%s"%(vlcnt, vlcntTotalIssued, tmp)
        if vscnt != -1:
          comment = comment + (" " if comment else "") + "vscnt(%s)"%(vscnt)
        if dscnt != -1:
          tmp = ""
          for cntStr in waitLocalLoadCntStrList:
            tmp += " - %s"%cntStr
          comment = comment + (" " if comment else "") + "dscnt(%d) = %d%s"%(dscnt, dscntTotalIssued, tmp)
        # if not self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
        return SWaitCnt(dscnt=dscnt, vlcnt=vlcnt, vscnt=vscnt, comment="%s (interleaved)"%comment)
    else:
      commentList = []
      # Global read wait
      if self.beta:
        vlcnt = 0
        commentList.append("Beta")
      if self.loadE:
        vlcnt = 0
        commentList.append("E")
      if self.parentWriter.states.useGateResidual:
        vlcnt = 0
        commentList.append("Gate")
      # Local read wait
      if self.parentWriter.states.useBias == DataDirection.READ:
        dscnt = 0
        commentList.append("Bias LDS")
      if (self.kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
        dscnt = 0
        commentList.append("ScaleABVec")
      if self.kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel:
        dscnt = 0
        commentList.append("ScaleAlphaVec")
      if (vlcnt != -1) or (dscnt != -1):
        # Get comment
        comment = "wait for " + commentList[0]
        for c in commentList[1:]:
          comment += ", %s"%c
        return SWaitCnt(dscnt=dscnt, vlcnt=vlcnt, vscnt=vscnt, comment=comment)
    return None

  ##############################################################################
  # choose the ADD instruction for combining external C with internal C
  # used in atomic=1 case to compute expected external data
  ##############################################################################
  def _chooseAddForAtomic(self, kernel, dst, src0, src1, comment):
    module = Module("chooseAddForAtomic")
    if kernel["ProblemType"]["DataType"].isBFloat16():
      if kernel["_GlobalAccumulation"]:
        module.add(VAddF32(dst, src0, src1, comment=comment))
    elif kernel["ProblemType"]["DataType"].isHalf():
      if kernel["_GlobalAccumulation"]:
        module.add(VAddF32(dst, src0, src1, comment=comment))
      elif kernel["ProblemType"]["HighPrecisionAccumulate"]:
        if self.parentWriter.states.asmCaps["v_fma_mix_f32"]:
          module.add(VFmaMixF32(dst, src0, 1, src1, comment=comment))
        elif self.parentWriter.states.asmCaps["v_mad_mix_f32"]:
          module.add(VMadMixF32(dst, src0, 1, src1, comment=comment))
        else:
          assert False, "No valid v_mad_mix_f32 equivalent"
      else:
        module.add(VAddPKF16(dst, src0, src1, comment))
    elif kernel["ProblemType"]["DataType"].isInt8x4() or kernel["ProblemType"]["DataType"].isInt8():
      # assume v_add_i32 can be used in place of v_add_f32
      # need to add saturation directive to v_add_i32 instruction to clamp integer arithmetic
      module.add(VAddI32(dst, src0, src1, comment=comment))
    elif kernel["ProblemType"]["DataType"].isSingle():
      module.add(VAddF32(dst, src0, src1, comment=comment))
    else:
       #support for double
      module.add(VAddF64(dst, src0, src1, comment=comment))

    return module

  def _emitLdsBarrierIfNeeded(self, targetModule: Module, isSingleKernel: bool):
    """Emit the LDS write barrier once per batch; idempotent via
    self.isLocalBarrierInit so repeat callers (preamble/body) are no-ops.
    """
    if isSingleKernel and (not self.isLocalBarrierInit):
      targetModule.add(SWaitCnt(dscnt=0, comment="Wait for LDS write"))
      targetModule.add(SBarrier(comment="LDS write barrier"))
      self.isLocalBarrierInit = True

  def _emitElt0EpilogueLoads(self, module: Module, addrCalc: 'AddrCalculation',
                             mask, elementIdx: int, bufferOOB,
                             loadInputCode: Module, factor_gwvw: int,
                             preamble: bool = False):
    """Emit one element's epilogue LDS loads (Bias -> ScaleAlphaVec ->
    ScaleAVec/ScaleBVec, then reorder). Shared by the CLS preamble
    (preamble=True: address compute only) and the per-element loop
    (preamble=False: also ds_load). xxxLoadIssued lists are appended in the
    body only, so the preamble adds no extra entry.
    """
    addrBiasVgpr          = addrCalc.addrBiasVgpr
    addrScaleAVecVgpr     = addrCalc.addrScaleAVecVgpr
    addrScaleBVecVgpr     = addrCalc.addrScaleBVecVgpr
    addrScaleAlphaVecVgpr = addrCalc.addrScaleAlphaVecVgpr
    dataBias              = self.ss.elementDataBias[elementIdx]
    dataScaleAVec         = self.ss.elementDataScaleAVec[elementIdx]
    dataScaleBVec         = self.ss.elementDataScaleBVec[elementIdx]
    dataScaleAlphaVec     = self.ss.elementDataScaleAlphaVec[elementIdx]
    skipLoad = True if self.factorDim else False
    isSingleKernel = ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel") or self.kernel["StreamK"] > 0

    def addEpilogueLoad(modGwvw, ldName: str, addrVecVgpr, addrVec, dataVec, loadedDataVec,
                        vecOffset, gwvw, referenceVgpr, dim, referenceDim,
                        skipLoad: bool = False, comment: str = "") -> int:
      """One vector's epilogue load: emitLdChange (address compute) always runs;
      captured `preamble` gates ONLY the ds_load (preamble=True = address only).
      """
      loadsIssued = 0
      tmpInrSgpr = self._epilogScratchSgpr(1)
      module.add(addrCalc.emitLdChange(self.kernel, self.ss, ldName, self.edge, self.beta, mask, bufferOOB, (elementIdx == 0), self.tmpVgpr, tmpInrSgpr, addrVecVgpr, addrVec, dim))
      self._epilogScratchFree(tmpInrSgpr)
      ldsAddrVgpr = referenceVgpr if (referenceVgpr and (dim == referenceDim)) else addrVecVgpr
      if dataVec not in loadedDataVec:
        # GroupLoadStore routes barrier+ds_load into `loadInputCode` so they
        # are grouped with the C input; otherwise both go into `module`.
        targetModule = loadInputCode if self.kernel["GroupLoadStore"] else module
        self._emitLdsBarrierIfNeeded(targetModule, isSingleKernel)
        if not preamble:
          targetModule.add(self.parentWriter.addLdsLoad(self.kernel["ProblemType"]["ComputeDataType"], dataVec, ldsAddrVgpr, vecOffset, gwvw, comment=comment))
          loadedDataVec[dataVec] = ceil(self.kernel["ProblemType"]["ComputeDataType"].numBytes() * gwvw / 16)
          loadsIssued = ceil(self.kernel["ProblemType"]["ComputeDataType"].numBytes() * gwvw / 16)
          if (self.ss.cfg.gwvw != gwvw) and (not skipLoad):
            remain_load = self.ss.cfg.gwvw - 1
            #For below ds_read instruction do not add bias issued , because of all ds_load instructions need to be completed at the same time in this batch.
            for r in range(remain_load):
              modGwvw.add(self.parentWriter.addLdsLoad(self.kernel["ProblemType"]["ComputeDataType"], dataVec, ldsAddrVgpr, vecOffset, factor_gwvw, comment=comment))
      return loadsIssued

    modGwvwScale = []
    localReferenceVgpr = None
    if self.parentWriter.states.useBias == DataDirection.READ:
      modGwvwBias = Module("GwvwBias")
      self.localLoadsBiasIssued += addEpilogueLoad(modGwvwBias, 'Bias', addrBiasVgpr, self.addrBias, dataBias, self.loadedDataBias, addrCalc.biasOffset[self.factorDim], factor_gwvw, localReferenceVgpr, self.factorDim, self.factorDim, skipLoad=skipLoad, comment="load Bias")
      localReferenceVgpr = addrBiasVgpr
      modGwvwScale.append(modGwvwBias)
    if not preamble:
      self.biasLoadIssued.append(len(self.loadedDataBias) * ceil(self.kernel["ProblemType"]["ComputeDataType"].numBytes() * factor_gwvw / 16))

    if self.kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel:
      modGwvwScaleAlpha = Module("GwvwScaleAlpha")
      # For multi-DU, the subtile ScaleAlphaVec epilogue load passes None as the
      # LDS reference vgpr; non-multi-DU uses localReferenceVgpr.
      savIsMultiDU = isSubtileMultiDU(self.kernel)
      if savIsMultiDU:
        savLdsRefVgpr = None if (self.kernel.get("UseSubtileImpl") and addrScaleAlphaVecVgpr is not None) else localReferenceVgpr
      else:
        savLdsRefVgpr = localReferenceVgpr
      self.loadsScaleAlphaVecIssued += addEpilogueLoad(modGwvwScaleAlpha, "ScaleAlphaVec", addrScaleAlphaVecVgpr, self.addrScaleAlphaVec, dataScaleAlphaVec, self.loadedDataScaleAlphaVec, addrCalc.scaleAlphaVecOffset[self.factorDim], factor_gwvw, savLdsRefVgpr, self.factorDim, self.factorDim, skipLoad=skipLoad, comment="load scaleAlpha")
      if localReferenceVgpr == None:
        localReferenceVgpr = addrScaleAlphaVecVgpr
      modGwvwScale.append(modGwvwScaleAlpha)
    if not preamble:
      self.scaleAlphaVecLoadIssued.append(len(self.loadedDataScaleAlphaVec) if self.factorDim else len(self.loadedDataScaleAlphaVec) * ceil(self.kernel["ProblemType"]["ComputeDataType"].numBytes() * factor_gwvw / 16))

    if (self.kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
      modGwvwScaleA = Module("GwvwScaleA")
      modGwvwScaleB = Module("GwvwScaleB")
      self.loadsScaleAVecIssued += addEpilogueLoad(modGwvwScaleA, "ScaleAVec", addrScaleAVecVgpr, self.addrScaleAVec, dataScaleAVec, self.loadedDataScaleAVec, addrCalc.scaleAVecOffset, self.ss.cfg.gwvw, localReferenceVgpr, 0, self.factorDim, comment="load scaleA")
      self.loadsScaleBVecIssued += addEpilogueLoad(modGwvwScaleB, "ScaleBVec", addrScaleBVecVgpr, self.addrScaleBVec, dataScaleBVec, self.loadedDataScaleBVec, addrCalc.scaleBVecOffset, 1, localReferenceVgpr, 1, self.factorDim, skipLoad=True, comment="load scaleB")
      if localReferenceVgpr == None:
        localReferenceVgpr = addrScaleAVecVgpr if self.factorDim == 0 else addrScaleBVecVgpr
      modGwvwScale.append(modGwvwScaleA)
      modGwvwScale.append(modGwvwScaleB)
    if not preamble:
      self.scaleAVecLoadIssued.append(len(self.loadedDataScaleAVec) * ceil(self.kernel["ProblemType"]["ComputeDataType"].numBytes() * self.ss.cfg.gwvw / 16))
      self.scaleBVecLoadIssued.append(len(self.loadedDataScaleBVec))

    # Reorder scale
    length = 0
    for mod in modGwvwScale:
      length = max(length, len(mod.items()))

    for index in range(0, length):
      for mod in modGwvwScale:
        if len(mod.items()) > index:
          module.add(mod.items()[index])

  def _emitElt0LdsPreambleBeforeBanner(self, module: Module, bufferOOB,
                                       loadInputCode: Module, factor_gwvw: int):
    """CompactLoopStore preamble (batch 0): hoist elt-0 LDS setup out of the
    per-element loop so the CLS countdown loop need not re-emit it each iter. The
    body dedups via state flags (ss.singleCol*AddrUpdated, isLocalBarrierInit).
    Per-section gating (addr compute only if optSingleColVgpr; D scaleToBpe /
    sgpr-offset init / MSB prewarm NonEdge-only) is documented inline below.
    """
    elementIdx = 0
    addrCalc: AddrCalculation = self.ss.elementAddr[elementIdx]
    addrBiasVgpr          = addrCalc.addrBiasVgpr
    addrCVgpr             = addrCalc.addrCVgpr
    addrDVgpr             = addrCalc.addrDVgpr
    addrScaleAlphaVecVgpr = addrCalc.addrScaleAlphaVecVgpr
    addrScaleAVecVgpr     = addrCalc.addrScaleAVecVgpr
    addrScaleBVecVgpr     = addrCalc.addrScaleBVecVgpr
    mask                  = self.ss.elementMask[elementIdx]
    bufferOOB             = None

    isSingleKernel = ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or \
                      self.kernel["GlobalSplitUAlgorithm"] == "MultipleBufferSingleKernel") or \
                     self.kernel["StreamK"] > 0

    # emitAddressSetupCode for elt 0 is safe: coordOffset0=0 and rowInc=0, so it
    # only emits the (d1,vc1,d0,vc0) comment + sets coord0Vgpr state -- no real
    # instructions. The body re-runs it for elt 0 (just re-emits the comment).
    tmpInrSgpr = self._epilogScratchSgpr(1)
    module.add(addrCalc.emitAddressSetupCode(self.kernel, self.tPB, self.ss, self.tmpVgpr, \
        tmpInrSgpr, self.edge, self.beta, self.atomic, elementIdx, addrDVgpr))
    self._epilogScratchFree(tmpInrSgpr)

    # elt-0 epilogue hoist: optSingleColVgpr -> hoist the address compute
    # (preamble=True; body folds it via ss.singleCol*AddrUpdated). Else (Edge /
    # non-optSingleCol) -> hoist only the LDS barrier.
    if self.ss.optSingleColVgpr:
      self._emitElt0EpilogueLoads(module, addrCalc, mask, elementIdx, bufferOOB,
                                   loadInputCode, factor_gwvw,
                                   True)
    else:
      targetModule = loadInputCode if self.kernel["GroupLoadStore"] else module
      self._emitLdsBarrierIfNeeded(targetModule, isSingleKernel)

    # D scaleToBpe -- NonEdge only (Edge has its own per-elt path that is not
    # safe to hoist).
    if not self.edge:
      tmpInrSgpr = self._epilogScratchSgpr(1)
      if self.kernel["GlobalSplitU"] == 1 or (self.kernel["GlobalSplitUAlgorithm"] != "MultipleBufferSingleKernel"):
        module.add(addrCalc.emitLdChange(self.kernel, self.ss, 'D', self.edge, self.beta, \
            mask, bufferOOB, True, self.tmpVgpr, tmpInrSgpr, addrDVgpr, self.addrD, 0))
      if self.beta:
        module.add(addrCalc.emitLdChange(self.kernel, self.ss, 'C', self.edge, self.beta, \
            mask, bufferOOB, True, self.tmpVgpr, tmpInrSgpr, addrCVgpr, self.addrC, 0))
      self._epilogScratchFree(tmpInrSgpr)

    # Init sgpr offset -- NonEdge only. Needed for the delayed-pattern
    # incrementToNextRow when optSrdIncForRow=1 (NonEdge). Edge path
    # (optSrdIncForRow=0) can skips the delayed pattern entirely so these inits
    # might not needed.

    module.add(SMovB32(dst=sgpr(self.tmpS01),   src=0, comment="Init sgpr offset"))
    module.add(SMovB32(dst=sgpr(self.tmpS01+1), src=0, comment="Init sgpr offset"))

    # Prewarm VGPR MSB bank -- NonEdge only. Forces the upper-bank toggle
    # for the upcoming ds_load to settle ahead of time so the loop body
    # avoids a stall on first use.
    prewarmVgpr = None
    if self.parentWriter.states.useBias == DataDirection.READ and addrBiasVgpr is not None:
      prewarmVgpr = addrBiasVgpr
    elif self.kernel["ProblemType"]["UseScaleAlphaVec"] and addrScaleAlphaVecVgpr is not None:
      prewarmVgpr = addrScaleAlphaVecVgpr
    elif (self.kernel["ProblemType"]["UseScaleAB"] == "Vector") and addrScaleAVecVgpr is not None:
      prewarmVgpr = addrScaleAVecVgpr
    if prewarmVgpr is not None:
      module.add(VMovB32(dst=vgpr(prewarmVgpr), src=vgpr(prewarmVgpr),
                          comment="prewarm VGPR MSB bank for upcoming ds_load"))

  def _lookaheadRowInc(self, elementIdx: int) -> int:
    """CompactLoopStore look-ahead: the rowInc that the NEXT EMITTING elt's
    s_add will consume, so this elt's delayed AFTER-primer can encode it
    (removes the off-by-one in the CLS primer chain). Forward-scan from
    elementIdx+1 for the first elt with rowInc != 0; if none remain in this
    batch, fall through to self.inter_iter_rowInc (the cross-batch advance,
    precomputed by KernelWriterAssembly; 0 on the last batch). Returns 0 when
    there is no override (non-CLS, no elementAddr, or no further advance).
    """
    if not (self.kernel["CompactLoopStore"] and self.ss.elementAddr):
      return 0
    for _j in range(elementIdx + 1, len(self.batchElements)):
      _ri = self.ss.elementAddr[_j].rowInc
      if _ri != 0:
        return _ri
    return self.inter_iter_rowInc

  def _epilogScratchSgpr(self, n: int = 1):
    """Scratch sgpr for epilogue address math.
    CompactLoopStore: borrow `n` aligned sgpr(s) from the pool so the address
    math does NOT clobber the row-increment primer chain (which lives in
    tmpS01 / tmpS01+1). Non-CLS: reuse self.tmpSgpr -- this is the baseline, so
    CLS=0 codegen is bit-for-bit unchanged. Pair with _epilogScratchFree.
    """
    if self.kernel["CompactLoopStore"]:
      return self.parentWriter.sgprPool.checkOutAligned(n, 1)
    return self.tmpSgpr

  def _epilogScratchFree(self, sgprIdx):
    """Release a scratch sgpr from _epilogScratchSgpr (no-op for non-CLS, which
    reused self.tmpSgpr and owns nothing).
    """
    if self.kernel["CompactLoopStore"]:
      self.parentWriter.sgprPool.checkIn(sgprIdx)

  def _prolog(self, module: Module):
    module.addComment0("optSingleColVgpr=%u optSharedColVgpr=%u optSGPRUsage=%s optSrdIncForRow=%u factorDim=%u" % \
              (self.ss.optSingleColVgpr, self.ss.optSharedColVgpr, self.ss.optSGPRUsage, self.ss.optSrdIncForRow, self.factorDim))

    if self.kernel["StoreSyncOpt"]:
      self._storeSyncOpt(module)

    # comment tt1, tt0, vc1, vc0
    # tt = thread tile, vc=vector component
    commentStr = "Global Write%s%s Batch #%u (d1,d0,vc1,vc0) =\n   " \
        % (" Beta" if self.beta else "", " Edge" if self.edge else "", self.batchIdx)

    commentStr = ''.join([commentStr] \
                            + ["(%u,%u,%u,%u:vw%u%s)%s" % \
                               (element[0], element[1], element[2], element[3], self.gwvw,
                               ":vaw:%u"%self.atomicW if self.atomic else "",
                               "" if idx == len(self.batchElements) -1 else "; ")
                               for idx, element in enumerate(self.batchElements)])
    # setupStoreElementsForBatch must run before the CLS preamble call so
    # ss.elementAddr[0] is populated. For non-CLS the batch banner is emitted
    # right after in the else branch (matches baseline ordering); for CLS the
    # banner is moved into the CLS header block below (after sgpr setup, before
    # the CLS label) so the per-iter store body is visually a single labelled
    # unit.
    if self.kernel["_GlobalAccumulation"] != "MultipleBufferSingleKernel":
      self.ss.setupStoreElementsForBatch(self.kernel, self.gwvw, self.batchElements, self.batchElementSgprs, isOptNLL=False, factorDim=self.factorDim)
    else:
      self.ss.setupStoreElementsForBatch(self.kernel, self.gwvw, self.batchElements, self.batchElementSgprs, isOptNLL=True, factorDim=self.factorDim)

    self.localLoadsBiasIssued = 0
    self.storesIssued    = 0
    self.loadsBetaIssued   = 0
    self.loadsEIssued      = 0
    self.loadsGateIssued   = 0
    self.loadsScaleAVecIssued = 0
    self.loadsScaleBVecIssued = 0
    self.loadsScaleAlphaVecIssued     = 0

    ########################################
    # calculate addr and masks
    # On input, coord0 and coord1 are VGPRs computed in the pre-batch code, based
    # on the thread and tid number.  These are ELEMENT offsets from start of tensor C
    # for the top-left corner this thread will write.  These are not changed
    # across all the store loop iters.
    if self.debugConfig["ConservativeWaitCnt"] & 0x10:
      module.add(SBarrier(comment="debug"))
      module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="ConservativeWaitCnt"))
      module.add(SBarrier(comment="debug"))

    if not self.edge and self.debugConfig["ForceEdgeStores"] >= 2:
      module.add(self.parentWriter.getBomb()) # should not get here
    if self.edge and self.debugConfig["AssertNoEdge"]:
      module.add(self.parentWriter.getBomb()) # should not get here

    ########################################
    # rC *= alpha
    if not self.kernel["InterleaveAlpha"] and self.applyAlpha and self.parentWriter.alphaBeforeLoadC:
      module.addComment1("rC *= alpha batchElements=%s"%self.batchElements)
      if self.codeMulAlpha is None:
        elementIdx = 0
        while elementIdx < len(self.batchElements):
          isEnd = (elementIdx == len(self.batchElements) - 1)
          if not isEnd and (self.ss.elementSumIdx[elementIdx] + 1 == self.ss.elementSumIdx[elementIdx + 1]) and (self.ss.elementSumIdx[elementIdx] % 2 == 0):
            module.add(self._applyAlpha(self.kernel, self.gwvw, self.ss.elementSumIdx, elementIdx, self.tmpS01, usePK=True))
            elementIdx += 2
          else:
            module.add(self._applyAlpha(self.kernel, self.gwvw, self.ss.elementSumIdx, elementIdx, self.tmpS01))
            elementIdx += 1
      else:
          regsPerScalar = self.parentWriter.states.bpeCinternal // self.parentWriter.states.bpr # register per scalar
          for elementIdx in range(len(self.batchElements)):
            for vi in range(self.gwvw):
              rh = replaceHolder(self.codeMulAlpha.popFirstItem(), self.ss.elementSumIdx[elementIdx]*regsPerScalar + regsPerScalar*vi)
              if (self.kernel["GlobalSplitU"] == 1) and (self.kernel["ProblemType"]["ComputeDataType"].isSingle() and self.kernel["ProblemType"]["DataType"].isInt8()):
                srcRegName = rh.getParams()[2].getCompleteRegName()
                module.add(VCvtI32toF32(dst=vgpr(srcRegName), src=vgpr(srcRegName), comment="Convert MI out reg to fp32"))
              module.add(rh)


    loadInputCode    = Module("loadInputCode")

    self.betaLoadIssued = []
    self.eLoadIssued = []
    self.gateLoadIssued = []
    self.biasLoadIssued = []
    self.scaleAVecLoadIssued = []
    self.scaleBVecLoadIssued = []
    self.scaleAlphaVecLoadIssued = []

    self.loadedDataBeta = {}
    self.loadedDataE = {}
    self.loadedDataGate = {}
    self.loadedDataBias = {}
    self.loadedDataScaleAVec = {}
    self.loadedDataScaleBVec = {}
    self.loadedDataScaleAlphaVec = {}

    #when factorDim = 1 the bias's gwvw is alwasy be 1.
    factor_gwvw = 1 if self.factorDim else self.ss.cfg.gwvw

    # CompactLoopStore CLS header (batch 0 only): preamble + sgpr setup + CLS
    # label + M0 assignment + M0 step. Wrapped in `if CompactLoopStore` as a
    # unit so non-CLS .s emits ONLY the original module.addComment2(commentStr)
    # in the else branch (matches baseline)
    if self.kernel["CompactLoopStore"] and self.batchIdx == 0:
      self._emitElt0LdsPreambleBeforeBanner(module, None, #bufferOOB,
                                             loadInputCode,
                                             factor_gwvw)

      module.add(SMovB32(dst=sgpr("CLSm0Base"), src=hex(0x0), comment="CLS M0 base = 0"))
      module.add(SMovB32(dst=sgpr("CLSLoopCounter"), src=hex(self._computeCLSIterCount()), comment="CLS loop iter count"))
      module.addComment2(commentStr)
      self.ss._clsLoopLabel = Label(self.parentWriter.labels.getNameInc("CLS"), "")
      module.add(self.ss._clsLoopLabel)
      module.add(SMovB32(dst=mgpr(0), src=sgpr("CLSm0Base"),
          comment="LDS clamp at sgpr(CLSm0Base)"))
      # CLS M0 step: M0 indirectly offsets the src VGPR index of v_movrelsd_2_b32.
      # Step = the stride of the CLS iter dim in the SRC formula; taken from the
      # same layout derivation as the loop iter count (computeCLSLayout) so the
      # two can never drift apart.
      _, _, cls_m0_step = self._computeCLSLayout()
      module.add(SAddU32(dst=sgpr("CLSm0Base"), src0=sgpr("CLSm0Base"), src1=cls_m0_step,
                         comment="CLS M0 step (src coef of CLS iter dim)"))
    else:
      module.addComment2(commentStr)

    module.addComment1("calc coords, apply mask, and issue loads (if necessary)")

    if self.kernel["BufferStore"] and (self.edge or (self.kernel["NumWaveSplitK"] > 1)):
      bufferOOB = self.tmpVgpr + self.tmpVgprSize - 1
      module.add(VMovB32(dst=vgpr(bufferOOB), src="BufferOOB"))
    else:
      bufferOOB = None

    for elementIdx, element in enumerate(self.batchElements):
      addrCalc: AddrCalculation = self.ss.elementAddr[elementIdx]
      addrCVgpr    = addrCalc.addrCVgpr
      addrDVgpr    = addrCalc.addrDVgpr
      addrEVgpr    = addrCalc.addrEVgpr
      addrBiasVgpr = addrCalc.addrBiasVgpr
      addrScaleAVecVgpr = addrCalc.addrScaleAVecVgpr
      addrScaleBVecVgpr = addrCalc.addrScaleBVecVgpr
      addrScaleAlphaVecVgpr = addrCalc.addrScaleAlphaVecVgpr
      data     = self.ss.elementData[elementIdx]
      dataBeta = self.ss.elementData[elementIdx]
      dataE    = self.ss.elementDataE[elementIdx]
      dataGate = self.ss.elementDataGate[elementIdx] if self.parentWriter.states.useGateResidual else 0
      dataBias = self.ss.elementDataBias[elementIdx]
      dataScaleAVec = self.ss.elementDataScaleAVec[elementIdx]
      dataScaleBVec = self.ss.elementDataScaleBVec[elementIdx]
      dataScaleAlphaVec = self.ss.elementDataScaleAlphaVec[elementIdx]
      mask     = self.ss.elementMask[elementIdx]
      vc0 = element[3]
      sumIdxGSUSYNC = self.ss.elementSumIdx[elementIdx]

      # CompactLoopStore look-ahead override: this elt's delayed AFTER-primer
      # encodes the NEXT EMITTING elt's rowInc (see _lookaheadRowInc).
      _emitOverrideRows = self._lookaheadRowInc(elementIdx)

      tmpInrSgpr = self._epilogScratchSgpr(1)
      _skipCrossBatchAdv = (self.kernel["CompactLoopStore"] and elementIdx == 0 and self.batchIdx > 0)
      module.add(addrCalc.emitAddressSetupCode(self.kernel, self.tPB, self.ss, self.tmpVgpr, tmpInrSgpr, self.edge, self.beta, self.atomic, elementIdx, addrDVgpr,
                                               skipCrossBatchAdvance=_skipCrossBatchAdv))
      self._epilogScratchFree(tmpInrSgpr)

      if self.edge:
        tmpInrSgpr = self._epilogScratchSgpr(2)
        module.add(addrCalc.edgeProtectCode(self.kernel, self.edge, self.beta, self.atomic, mask, tmpInrSgpr))
        self._epilogScratchFree(tmpInrSgpr)
        if self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
          module.addComment1("edge Protect")
      # create code Module to push mov vgpr,acc instructions
      if self.beta:
        tmpInrSgpr = self._epilogScratchSgpr(1)
        module.add(addrCalc.emitLdChange(self.kernel, self.ss, 'C', self.edge, self.beta, mask, bufferOOB, (elementIdx == 0), self.tmpVgpr, tmpInrSgpr, addrCVgpr, self.addrC, 0))
        self._epilogScratchFree(tmpInrSgpr)
        if dataBeta not in self.loadedDataBeta:
          # In the UseSubtileImpl NonEdge path the workgroup-level edge check is relaxed
          # (subtile-aligned remainder is allowed into NonEdge), so individual waves may
          # own rows/columns beyond the valid output region.  Gate each C load by writing
          # SrdC+2 (num_records): BufferOOB → normal load, 0 → hardware returns zero.
          #
          # element loop is N(d1)-outer / M(d0)-inner.
          # d1 (N) check: emitted once per d1 group — sets SrdC+2 = BufferOOB if N valid, else 0.
          # d0 (M) check: emitted per element — overwrites SrdC+2 = SrdC+2 if M valid, else 0.
          #   (AND semantics: SrdC+2 = BufferOOB only when both M and N are valid.)
          # d0 is monotone within each d1 group: once OOB, remaining d0s are also OOB.
          mGuardSgpr = self.parentWriter.states.subtileM32ValidBlocksSgpr
          nGuardSgpr = self.parentWriter.states.subtileN16ValidBlocksSgpr
          if not self.edge and (mGuardSgpr is not None or nGuardSgpr is not None):
            d1, d0 = element[0], element[1]
            # N guard: emit once per d1 group.
            if nGuardSgpr is not None and d1 != self._subtileCloadPrevD1:
              d1Cmp = d1 * 16 if self.parentWriter.states.storeAlign8 else d1
              module.add(_scmpGtU32(self.parentWriter, sgpr("SubtileNGuard"), d1Cmp,
                                    comment="subtile C load: clamped > %d?" % d1Cmp))
              module.add(SCSelectB32(dst=sgpr("SrdC+2"), src0="BufferOOB", src1=0,
                                     comment="SrdC+2 = BufferOOB if N valid, else 0"))
              self._subtileCloadPrevD1 = d1
            # M guard: emit per element, AND into SrdC+2.
            if mGuardSgpr is not None:
              module.add(_scmpGtU32(self.parentWriter, sgpr("SubtileMGuard"), d0,
                                    comment="subtile C load: numMBlocks > d0=%d?" % d0))
              if nGuardSgpr is not None:
                module.add(SCSelectB32(dst=sgpr("SrdC+2"), src0=sgpr("SrdC+2"), src1=0,
                                       comment="SrdC+2 = SrdC+2 if M valid, else 0 (AND with N result)"))
              else:
                module.add(SCSelectB32(dst=sgpr("SrdC+2"), src0="BufferOOB", src1=0,
                                       comment="SrdC+2 = BufferOOB if M valid, else 0"))
          # Pass `elementIdx` so the readInput's incrementToNextRow gate can
          # open the CompactLoopStore chain-seed emit for elt-0. Pass
          # `overrideAfterPrimerRows=_emitOverrideRows` so the AFTER-primer
          # encodes the NEXT EMITTING elt's rowInc (look-ahead). Non-CLS:
          # `elementIdx` is unused and `_emitOverrideRows`==0 means no override.
          if self.kernel["GroupLoadStore"]:
            loadInputCode.add(self.parentWriter.readInput(self.kernel, self.ss, 'C', self.kernel["ProblemType"]["DestDataType"], addrCalc, vc0, data, self.gwvw, addrCVgpr, self.tmpS01, elementIdx, self.batchIdx,
                                                          overrideAfterPrimerRows=_emitOverrideRows))
          else:
            module.add(self.parentWriter.readInput(self.kernel, self.ss, 'C', self.kernel["ProblemType"]["DestDataType"], addrCalc, vc0, data, self.gwvw, addrCVgpr, self.tmpS01, elementIdx, self.batchIdx,
                                                  overrideAfterPrimerRows=_emitOverrideRows))
          self.loadedDataBeta[dataBeta] = ceil(self.kernel["ProblemType"]["DestDataType"].numBytes() * self.ss.cfg.gwvw / 16)
          self.loadsBetaIssued += ceil(self.kernel["ProblemType"]["DestDataType"].numBytes() * self.gwvw / 16)
      self.betaLoadIssued.append(len(self.loadedDataBeta) * ceil(self.kernel["ProblemType"]["DestDataType"].numBytes() * self.ss.cfg.gwvw / 16))

      if (self.kernel["ProblemType"]["UseE"] and self.kernel["ProblemType"]["Gradient"] and self.kernel["ProblemType"]["ActivationType"] != 'none') and ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["StreamK"] > 0):
        tmpInrSgpr = self._epilogScratchSgpr(1)
        module.add(addrCalc.emitLdChange(self.kernel, self.ss, 'E', self.edge, self.beta, mask, bufferOOB, (elementIdx == 0), self.tmpVgpr, tmpInrSgpr, addrEVgpr, self.addrE, 0))
        self._epilogScratchFree(tmpInrSgpr)
        if dataE not in self.loadedDataE:
          loadOffset = int((self.kernel["ProblemType"]["ComputeDataType"].numRegisters() - self.kernel["ProblemType"]["DataTypeE"].numRegisters()) * self.ss.cfg.gwvw)
          if self.kernel["GroupLoadStore"]:
            loadInputCode.add(self.parentWriter.readInput(self.kernel, self.ss, 'E', self.kernel["ProblemType"]["DataTypeE"], addrCalc, vc0, dataE + loadOffset, self.gwvw, addrEVgpr, self.tmpS01))
          else:
            module.add(self.parentWriter.readInput(self.kernel, self.ss, 'E', self.kernel["ProblemType"]["DataTypeE"], addrCalc, vc0, dataE + loadOffset, self.gwvw, addrEVgpr, self.tmpS01))
          self.loadedDataE[dataE] = ceil(self.kernel["ProblemType"]["DataTypeE"].numBytes() * self.ss.cfg.gwvw / 16)
          self.loadsEIssued += ceil(self.kernel["ProblemType"]["DataTypeE"].numBytes() * self.gwvw / 16)
        self.loadE = True
      else:
        self.loadE = False
      self.eLoadIssued.append(len(self.loadedDataE) * ceil(self.kernel["ProblemType"]["DataTypeE"].numBytes() * self.ss.cfg.gwvw / 16))

      # Per-element epilogue LDS loads (same helper the CLS preamble used for
      # elt 0); if the preamble already loaded it, the ds_load is skipped here.
      self._emitElt0EpilogueLoads(module, addrCalc, mask, elementIdx, bufferOOB,
                                   loadInputCode, factor_gwvw,
                                   False)
      tmpInrSgpr = self._epilogScratchSgpr(1)
      if (self.kernel["ProblemType"]["UseE"] and not self.kernel["ProblemType"]["Gradient"]) and ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["StreamK"] > 0):
        module.add(addrCalc.emitLdChange(self.kernel, self.ss, 'E', self.edge, self.beta, mask, bufferOOB, (elementIdx == len(self.batchElements) - 1), self.tmpVgpr, tmpInrSgpr, addrEVgpr, self.addrE, 0))
      if self.storeBiasD == 1:
        module.add(addrCalc.emitLdChange(self.kernel, self.ss, 'Bias', self.edge, self.beta, mask, bufferOOB, (elementIdx == len(self.batchElements) - 1), self.tmpVgpr, tmpInrSgpr, addrBiasVgpr, self.addrBias, self.factorDim))

      if self.parentWriter.states.useGateResidual and \
         (self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1 or self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"):
        _gateList = self.kernel["ProblemType"].get("GateResidualDataTypeList", [])
        _destDtype = self.kernel["ProblemType"]["DestDataType"]
        if not _gateList:
          _prologLoadDtype = self.kernel["ProblemType"]["DestDataType"]
        elif len(_gateList) == 1:
          _prologLoadDtype = _gateList[0]
        else:
          _prologLoadDtype = _destDtype if _destDtype in _gateList else _gateList[0]

        _multiDtypeGate = (_gateList and len(_gateList) > 1)

        addrGateVgpr = addrCalc.addrGateVgpr if addrCalc.addrGateVgpr != None else addrDVgpr
        if not _multiDtypeGate and not self.ss.optSingleColVgpr:
          module.add(addrCalc.emitLdChange(
              self.kernel, self.ss, 'Gate', self.edge, self.beta, mask, bufferOOB,
              (elementIdx == 0), self.tmpVgpr, tmpInrSgpr, addrGateVgpr, self.addrD, 0))
        if dataGate not in self.loadedDataGate:
          if self.ss.optSingleColVgpr:
            # opt (single + multi dtype): hoisted to ONE branch after the prolog loop
            # (see _emitHoistedGateLoadPhase).
            pass
          elif not _multiDtypeGate:
            # no-opt single-dtype: unconditional load.
            # Null gate -> SRD num_records==0 -> buffer_load returns 0;
            _glTgt = loadInputCode if self.kernel["GroupLoadStore"] else module
            gateLoadMod = self.parentWriter.readInput(
                self.kernel, self.ss, 'Gate',
                _prologLoadDtype,
                addrCalc, vc0, dataGate, self.gwvw, addrGateVgpr, self.tmpS01)
            _glTgt.add(gateLoadMod)
          else:
            # no-opt (edge) multi-dtype: per-dtype dispatcher per element (gate
            # borrows D's per-element addr, so it must stay interleaved here).
            # all using this element's fresh mask. The dispatcher tail-branches
            # to a common "GateLoad_End" label per element.
            labels = self.parentWriter.labels
            gateLoadEndLabel = Label(
                labels.getNameInc("GateLoad_End_%u"%elementIdx), "")
            gateLoadTypeLabels = [
                Label(labels.getNameInc(
                    "GateLoad_%s_%u"%(g.toNameAbbrev(), elementIdx)), "")
                for g in _gateList]
            gateLoadTypeLabels.append(gateLoadEndLabel)

            # Snapshot mutables we'll restore after the dispatcher chain.
            _savedBpeGate = self.parentWriter.states.bpeGate
            _savedGlobalOffsetGate = addrCalc.globalOffsetGate

            # Per-element skip when gate pointer is null (no-gate problem): jump past
            # the whole dtype dispatcher straight to GateLoad_End for this element.
            module.add(self.parentWriter.getSCMPKInstruction(
                "EQU32", "SrdGate+2", 0, comment="gate disabled? (SrdGate num_records==0)"))
            module.add(SCBranchSCC1(gateLoadEndLabel.getLabelName(), "skip gate load if disabled (null gate)"))

            for i, nextLabel in enumerate(gateLoadTypeLabels[1:]):
              typeValue = _gateList[i].value
              gDtype = _gateList[i]
              module.add(gateLoadTypeLabels[i])
              module.add(self.parentWriter.getSCMPKInstruction(
                  "LGU32", "GateType", typeValue,
                  comment="GateType != %u"%typeValue))
              module.add(SCBranchSCC1(nextLabel.getLabelName(),
                                      "Branch if true (try next gate dtype)"))
              # per-dtype bpe. no-opt path only (opt is handled by the hoisted phase),
              # so gate uses per-element runtime addr -> globalOffsetGate fixed 0.
              perBranchBpe = int(self.parentWriter.states.bpr * gDtype.numRegisters())
              perBranchBpe = max(1, perBranchBpe)
              self.parentWriter.states.bpeGate = perBranchBpe
              addrCalc.globalOffsetGate = 0
              # Force fresh per-dtype gate address setup.
              self.ss.singleColGateAddrUpdated = False
              module.add(addrCalc.emitLdChange(
                  self.kernel, self.ss, 'Gate', self.edge, self.beta, mask, bufferOOB,
                  (elementIdx == 0), self.tmpVgpr, tmpInrSgpr, addrGateVgpr, self.addrD, 0))
              module.add(self.parentWriter.readInput(
                  self.kernel, self.ss, 'Gate', gDtype,
                  addrCalc, vc0, dataGate, self.gwvw, addrGateVgpr, self.tmpS01))
              # Restore bpe/offset for the next branch in this elem.
              self.parentWriter.states.bpeGate = _savedBpeGate
              addrCalc.globalOffsetGate = _savedGlobalOffsetGate
              module.add(SBranch(labelName=gateLoadEndLabel.getLabelName(),
                                 comment="Branch to GateLoad end"))
            module.add(gateLoadEndLabel)
          gateBytes = ceil(_prologLoadDtype.numBytes() * self.ss.cfg.gwvw / 16)
          self.loadedDataGate[dataGate] = gateBytes
          self.loadsGateIssued += gateBytes
        self.gateLoadIssued.append(len(self.loadedDataGate) * ceil(_prologLoadDtype.numBytes() * self.ss.cfg.gwvw / 16))
      else:
        self.gateLoadIssued.append(0)

      if self.kernel["GlobalSplitU"] == 1 or (self.kernel["_GlobalAccumulation"] != "MultipleBufferSingleKernel"): # "SingleBuffer" or "MultipleBuffer"
        module.add(addrCalc.emitLdChange(self.kernel, self.ss, 'D', self.edge, self.beta, mask, bufferOOB, (elementIdx == len(self.batchElements) - 1), self.tmpVgpr, tmpInrSgpr, addrDVgpr, self.addrD, 0))
      if self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
        module.add(addrCalc.emitLdChange(self.kernel, self.ss, 'TD', self.edge, self.beta, mask, bufferOOB, (elementIdx == len(self.batchElements) - 1), self.tmpVgpr, tmpInrSgpr, addrCalc.addrGSUSyncVgprs, self.addrD, 0))
      self._epilogScratchFree(tmpInrSgpr)
      if self.atomic and (not self.parentWriter.states.useAtomicAdd):
        # load c into data+1 because of CAS structure
        # TODO - Fix for double here, would need bigger load
        # FIXME
        # gwvw is the number of elements in the batch
        # iterate over number of atomic operations to perform, each of width atomicW
        for avi in range(self.gwvw // self.atomicW):
          dataV = self.ss.elementData[elementIdx] + int(avi*self.ss.cfg.numVgprsPerDataPerVI)
          bpm = self.parentWriter.states.bpeCexternal * self.atomicW
          useBuffer = self.kernel["BufferStore"]
          if self.kernel["BufferStore"]: # yes, BufferStore here - use same addressing regs for this load
            addr0 = vgpr(addrDVgpr)
            addr1 = sgpr("SrdD", 4)
          else:
            addr0 = vgpr(addrDVgpr, 2)
            addr1 = ""
          # Calculate vgpr Index for 32-bit/64-bit instruction
          # DGEMM use SRCS[2] register
          vgprIdx = bpm // 4
          module.add(self.parentWriter.chooseGlobalRead(useBuffer, bpm, dataV + vgprIdx, \
                    addr0, addr1, soffset=0, offset=addrCalc.globalOffset,
                    comment="load D (atomic) bpm=%u vaw=%u"%(bpm,self.atomicW)))

      if self.kernel["InterleaveAlpha"] and self.applyAlpha:
        module.add(self._applyAlpha(self.kernel, self.gwvw, self.ss.elementSumIdx, elementIdx, self.tmpS01))

      if not self.kernel["BufferStore"]:
        # emitLdChange for 'D' (above) already computed addrDVgpr = addrD + offset
        # Only need to build the address here if emitLdChange for 'D' was NOT called
        if not (self.kernel["GlobalSplitU"] == 1 or (self.kernel["GlobalSplitUAlgorithm"] != "MultipleBufferSingleKernel")):
          offsetSrc = (self.tmpVgpr + 2) if self.beta else addrDVgpr

          module.add(VAddCOU32(vgpr(addrDVgpr+0), VCC(), vgpr(self.addrD+0), \
              vgpr(offsetSrc+0), "addrDVgpr = D + index*bytes (lo)"))
          module.add(VAddCCOU32(vgpr(addrDVgpr+1), VCC(), vgpr(self.addrD+1), \
              vgpr(offsetSrc+1), VCC(), "addrDVgpr = D + index*bytes (hi)"))

        # restore full exec mask for calculating addr of next element
        if self.edge and (self.beta or self.loadE or self.atomic):
          module.add(self.getEdgeMovInstType()(EXEC(), -1, "full mask -1 -> exec"))

      if self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
        if self.ss.optSrdIncForRow and (addrCalc.rowInc or (self.kernel["CompactLoopStore"] and elementIdx == 0 and self.batchIdx == 0)) and self.kernel["StoreRemapVectorWidth"] > 0:
          module.addComment1("StoreRemap: shift coord1 address MultipleBufferSingleKernel")
          if self.kernel["ProblemType"]["UseE"] and (self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1):
            # TODO Check if works with StreamK
            printExit("Use E does not support StoreRemapVectorWidth if GSU == 1.")
            # module.add(addrCalc.incrementToNextRow(self.kernel, "E", self.ss, self.tmpS01, isCompute=True))
          module.add(addrCalc.incrementToNextRow(self.kernel, "D", self.ss, self.tmpS01, forceinitrow0=1, overrideAfterPrimerRows=_emitOverrideRows))
          module.add(VMovB32(vgpr(self.tmpVgpr), addrCalc.rowInc, comment="set shift rows"))
          module.add(VAddU32(vgpr(self.parentWriter.vgprs.storeRemapCoord1), vgpr(self.parentWriter.vgprs.storeRemapCoord1), vgpr(self.tmpVgpr), "shift storeRemap coord1"))

    module.add(loadInputCode)

    # opt + multi-dtype gate: emit the gate prolog load as ONE GateType dispatch
    # (the per-element loop above skipped it on the opt path).
    self._emitHoistedGateLoadPhase(module, bufferOOB)

    # Restore SrdC+2 = BufferOOB after subtile NonEdge C-load OOB gating (which may have set it to 0).
    if self.beta and not self.edge:
      mGuardSgpr = self.parentWriter.states.subtileM32ValidBlocksSgpr
      nGuardSgpr = self.parentWriter.states.subtileN16ValidBlocksSgpr
      if mGuardSgpr is not None or nGuardSgpr is not None:
        module.add(SMovB32(dst=sgpr("SrdC+2"), src="BufferOOB",
                           comment="restore SrdC+2 after subtile NonEdge C-load OOB gating"))

    if self.beta and self.kernel["StoreSyncOpt"]:
      self._storeSyncOpt(module)

    ########################################
    # AccVgpr read
    if self.codeAccVgprRead is not None and (self.kernel["LocalSplitU"] == 1 or self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"):
      regsPerScalar = self.parentWriter.states.bpeCinternal // self.parentWriter.states.bpr # register per scalar
      #TODOBS: Need to change this, for LSU>1 + subtile impl case
      if self.kernel["MIArchVgpr"] and self.kernel["LocalSplitU"] > 1:
        tmpStartVgprValuC = self.parentWriter.states.c.startVgprValu
        self.parentWriter.states.c.startVgprValu = 0
        module.add(RegSet("v", "vgprValuC", 0))
      # loop over store instructions within one batch
      for elementIdx in range(len(self.batchElements)):
        # loop over scalars within one store instruction
        for vi in range(self.gwvw):
          # loop over registers within one scalar
          for rIdx in range(0, regsPerScalar):
            module.add(replaceHolder(self.codeAccVgprRead.popFirstItem(), self.ss.elementSumIdx[elementIdx]*regsPerScalar + regsPerScalar*vi + rIdx - self.parentWriter.states.c.startVgprValu))

      if self.kernel["MIArchVgpr"] and self.kernel["LocalSplitU"] > 1:
        self.parentWriter.states.c.startVgprValu = tmpStartVgprValuC
        module.add(RegSet("v", "vgprValuC", tmpStartVgprValuC))

    elif self.kernel["LocalSplitU"] > 1:
      # read from LSU VGPRs
      regsPerScalar = self.parentWriter.states.bpeCinternal // self.parentWriter.states.bpr # register per scalar
      if self.ss.lsuStartVgprOffset > 0:
        for elementIdx in range(len(self.batchElements)):
          for vi in range(self.gwvw):
            for rIdx in range(0, regsPerScalar):
              idx = self.ss.elementSumIdx[elementIdx]*regsPerScalar + regsPerScalar*vi + rIdx - self.parentWriter.states.c.startVgprValu
              module.add(VMovB32(vgpr("ValuC+%u"%(idx)), vgpr("ValuC+%u"%(idx + self.ss.lsuStartVgprOffset)), comment="load from "+str(idx + self.ss.lsuStartVgprOffset)+" to "+str(idx) ))
      self.ss.lsuStartVgprOffset += len(self.batchElements) * self.gwvw * regsPerScalar

      if not self.kernel["MIArchVgpr"]:
        module.add(SNop(1, "2 wait states required before reading vgpr"))

    if self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
      module.addComment1("store after Acc, "+"GSU: "+str(self.kernel["GlobalSplitU"]))

    storeCodeGSUSK = Module("GroupLoadStore")
    if self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":#GSUGSU
      storeWidth = self.kernel["StoreVectorWidth"]
      for elementIdx in range(0, len(self.batchElements)):
        addrCalc: AddrCalculation = self.ss.elementAddr[elementIdx]
        if self.batchIdx == 0 and elementIdx == 0:
          addrDVgpr = addrCalc.addrDVgpr
          storeCodeGSUSK.add(vectorStaticMultiply(vgpr(addrDVgpr), vgpr("Serial"), storeWidth * self.parentWriter.states.bpeCinternal, ContinuousRegister(self.tmpS01, 1)))
          storeCodeGSUSK.add(SMovB32(dst=sgpr(self.tmpS01), src=0, comment="Init sgpr offset"))
          storeCodeGSUSK.addSpaceLine()
        if (self.kernel["ProblemType"]["UseE"] and not self.kernel["ProblemType"]["Gradient"]) and ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["StreamK"] > 0):
          vgprIdx = self.ss.elementSumIdx[elementIdx] - self.parentWriter.states.c.startVgprValu
          vgprDst = self.activationSetPCStruct.vgprActCopy if mergeActFuncCall else "ValuC+%d"%vgprIdx
          module.add(self.parentWriter.addStore(self.kernel, self.ss, 'E', addrCalc, vgprDst, self.tmpS01, self.edge, comment="store E"))

        sumIdx = self.ss.elementSumIdx[elementIdx]
        if self.kernel["StoreRemapVectorWidth"]:
          rpe = self.parentWriter.states.bpeCinternal // self.parentWriter.states.bpr
          module.add(self.parentWriter.storeRemapAddLocalWrite(self.kernel, self.ss, addrCalc, sumIdx*rpe))
          # Column Block Shape has been written to LDS
          # Now read back and write out to global memory
      # module.add(storeCodeGSUSK)

    if self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel" and self.kernel["StoreRemapVectorWidth"]:
      if self.parentWriter.StoreRemapLastBatch == 1:
        module.addComment1("Handle local read and global write")
        tmpInrSgpr = self._epilogScratchSgpr(2)
        storeModule, numNewStores = self.parentWriter.storeRemapAddStore(self.kernel, self.tmpVgpr, tmpInrSgpr, self.edge, self.parentWriter.StoreRemapLastBatch)
        self._epilogScratchFree(tmpInrSgpr)
        module.add(storeModule)
        self.storesIssued += numNewStores

    gsuComponent = Component.GSU.find(self.parentWriter)
    module.add(gsuComponent.globalWriteBatchProlog(self.parentWriter, self.kernel, self.tmpVgpr, self.tmpVgprSize, self.tmpVgprDynamic, \
                                                   self.batchIdx, self.ss, self.gwvw, self.batchElements, \
                                                   self.beta, self.edge, sumIdxGSUSYNC, addrCalc))

    # rC *= alpha
    if not self.kernel["InterleaveAlpha"] and self.applyAlpha and not self.parentWriter.alphaBeforeLoadC:
      module.addComment1("rC *= alpha batchElements=%s"%self.batchElements)
      if self.codeMulAlpha is None:
        elementIdx = 0
        while elementIdx < len(self.batchElements):
          isEnd = (elementIdx == len(self.batchElements) - 1)
          if not isEnd and (self.ss.elementSumIdx[elementIdx] + 1 == self.ss.elementSumIdx[elementIdx + 1]) and (self.ss.elementSumIdx[elementIdx] % 2 == 0):
            module.add(self._applyAlpha(self.kernel, self.gwvw, self.ss.elementSumIdx, elementIdx, self.tmpS01, usePK=True))
            elementIdx += 2
          else:
            module.add(self._applyAlpha(self.kernel, self.gwvw, self.ss.elementSumIdx, elementIdx, self.tmpS01))
            elementIdx += 1
      else:
          regsPerScalar = self.parentWriter.states.bpeCinternal // self.parentWriter.states.bpr # register per scalar
          for elementIdx in range(len(self.batchElements)):
            for vi in range(self.gwvw):
              rh = replaceHolder(self.codeMulAlpha.popFirstItem(), self.ss.elementSumIdx[elementIdx]*regsPerScalar + regsPerScalar*vi - self.parentWriter.states.c.startVgprValu)
              if (self.kernel["GlobalSplitU"] == 1) and (self.kernel["ProblemType"]["ComputeDataType"].isSingle() and self.kernel["ProblemType"]["DataType"].isInt8()):
                srcRegName = rh.getParams()[2].getCompleteRegName()
                module.add(VCvtI32toF32(dst=vgpr(srcRegName), src=vgpr(srcRegName), comment="Convert MI out reg to fp32"))
              module.add(rh)

  def _epilog(self, module: Module):
    # return registers to pool:
    lastDataD       = -1
    lastDataE       = -1
    checkedDataBias = {}
    checkedDataGate = {}
    checkedDataScaleAVec = {}
    checkedDataScaleBVec = {}
    checkedDataScaleAlphaVec = {}
    for elementIdx in range(len(self.batchElements)):
      sumIdxGSUSYNC = self.ss.elementSumIdx[elementIdx]
      if not self.ss.sharedColDVgprs:
        addrCalc: AddrCalculation = self.ss.elementAddr[elementIdx]
        addrEVgpr    = addrCalc.addrEVgpr
        addrDVgpr    = addrCalc.addrDVgpr
        addrGSUSyncVgprs    = addrCalc.addrGSUSyncVgprs
        addrCVgpr    = addrCalc.addrCVgpr
        addrBiasVgpr = addrCalc.addrBiasVgpr
        addrGateVgpr = addrCalc.addrGateVgpr
        addrScaleAVecVgpr = addrCalc.addrScaleAVecVgpr
        addrScaleBVecVgpr = addrCalc.addrScaleBVecVgpr
        addrScaleAlphaVecVgpr = addrCalc.addrScaleAlphaVecVgpr
        if addrEVgpr != None:
          self.parentWriter.vgprPool.checkIn(addrEVgpr)
        self.parentWriter.vgprPool.checkIn(addrDVgpr)
        if addrCVgpr != addrDVgpr:
          self.parentWriter.vgprPool.checkIn(addrCVgpr)
        if addrGSUSyncVgprs != None:
          self.parentWriter.vgprPool.checkIn(addrGSUSyncVgprs)
        if addrBiasVgpr != None:
          self.parentWriter.vgprPool.checkIn(addrBiasVgpr)
        # No-opt path: addrGateVgpr aliases addrDVgpr
        if addrGateVgpr != None and addrGateVgpr != addrDVgpr:
          self.parentWriter.vgprPool.checkIn(addrGateVgpr)
        if addrScaleAVecVgpr != None:
          self.parentWriter.vgprPool.checkIn(addrScaleAVecVgpr)
        if addrScaleBVecVgpr != None:
          self.parentWriter.vgprPool.checkIn(addrScaleBVecVgpr)
        if addrScaleAlphaVecVgpr != None:
          self.parentWriter.vgprPool.checkIn(addrScaleAlphaVecVgpr)

      data = self.ss.elementData[elementIdx]
      if data != 0:
        if data != lastDataD:
          self.parentWriter.vgprPool.checkIn(data)
        lastDataD = data

      dataBias = self.ss.elementDataBias[elementIdx]
      if dataBias != 0:
        if dataBias not in checkedDataBias:
          self.parentWriter.vgprPool.checkIn(dataBias)
        checkedDataBias[dataBias] = 1

      dataE = self.ss.elementDataE[elementIdx]
      if dataE != 0:
        if dataE != lastDataE:
          self.parentWriter.vgprPool.checkIn(dataE)
        lastDataE = dataE

      if self.parentWriter.states.useGateResidual:
        dataGate = self.ss.elementDataGate[elementIdx]
        if dataGate != 0:
          if dataGate not in checkedDataGate:
            self.parentWriter.vgprPool.checkIn(dataGate)
          checkedDataGate[dataGate] = 1

      def checkScaleVec(dataScaleVec, checkedDataScaleVec):
        if dataScaleVec != 0:
          if dataScaleVec not in checkedDataScaleVec:
            self.parentWriter.vgprPool.checkIn(dataScaleVec)
          checkedDataScaleVec[dataScaleVec] = 1

      checkScaleVec(self.ss.elementDataScaleAVec[elementIdx], checkedDataScaleAVec)
      checkScaleVec(self.ss.elementDataScaleBVec[elementIdx], checkedDataScaleBVec)
      checkScaleVec(self.ss.elementDataScaleAlphaVec[elementIdx], checkedDataScaleAlphaVec)

    self.ss.firstBatch = False
    self.ss.checkInTempVgprC()
    if self.kernel["_GlobalAccumulation"] != "MultipleBufferSingleKernel" and self.kernel["StoreRemapVectorWidth"]:
      if self.parentWriter.StoreRemapLastBatch == 1:
        module.addComment1("Handle local read and global write")
        # this seems buggy? it's possible to issue more than one stores for SR
        # module.add(self.storeRemapAddStore(kernel, tmpVgpr, tmpS01, edge))
        # storesIssued += 1
        tmpInrSgpr = self._epilogScratchSgpr(2)
        storeModule, numNewStores = self.parentWriter.storeRemapAddStore(self.kernel, self.tmpVgpr, tmpInrSgpr, self.edge, self.parentWriter.StoreRemapLastBatch)
        self._epilogScratchFree(tmpInrSgpr)
        module.add(storeModule)
        self.storesIssued += numNewStores

    if self.parentWriter.states.serializedStore:
      module.add(SNop(0, "1 wait state required when next inst writes vgprs held by previous dwordx4 store inst"))

    if self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":
      module.addCommentAlign("GW end") #GSUSYNC

  def _emitAdd(self, module: Module):
    if self.atomic:
      del self.tmpVgpr # catch bugs
      if self.parentWriter.states.useAtomicAdd:
        self._emitAtomicAdd(module)
      else:
        self._emitCasAdd(module)
    else:
      self._emitNonatomicAdd(module)

  def _emitGateCvt(self, dataGate, gateDtype):
    cvtMod = Module("GateCvt_%s" % gateDtype.toNameAbbrev())
    if gateDtype.isSingle():
      pass
    elif gateDtype.isHalf():
      for vi in range(self.gwvw - 1, -1, -1):
        sel = HighBitSel.LOW if (vi % 2) == 0 else HighBitSel.HIGH
        cvtMod.add(ECvtF16toF32(dst=vgpr(dataGate + vi), src=vgpr(dataGate + vi // 2), sel=sel,
          comment="GateCvt[h]: f16 -> f32 (vi=%d)"%vi))
    elif gateDtype.isBFloat16():
      # bf16 = top 16 bits of an f32; lo half shifts up, hi half masks.
      for vi in range(self.gwvw - 1, -1, -1):
        if (vi % 2) == 0:
          cvtMod.add(VLShiftLeftB32(dst=vgpr(dataGate + vi), shiftHex=16, src=vgpr(dataGate + vi // 2),
            comment="GateCvt[bf16]: lo bf16 -> f32 (vi=%d)"%vi))
        else:
          cvtMod.add(VAndB32(dst=vgpr(dataGate + vi), src0=hex(0xffff0000), src1=vgpr(dataGate + vi // 2),
            comment="GateCvt[bf16]: hi bf16 -> f32 (vi=%d)"%vi))
    elif gateDtype.isAnyFloat8() or gateDtype.isAnyBFloat8():
      # PK-cvt 2 elements per instruction into a tmp, then move to dataGate.
      isFP8 = gateDtype.isAnyFloat8()
      CvtPk = ECvtPkFP8toF32 if isFP8 else ECvtPkBF8toF32
      cvtTmp = self.tmpVgpr
      for vi in range(((self.gwvw - 1) // 2) * 2, -1, -2):
        sel = HighBitSel.LOW if (vi % 4) == 0 else HighBitSel.HIGH
        cvtMod.add(CvtPk(dst=vgpr(cvtTmp, 2), src=vgpr(dataGate + vi // 4), sel=sel,
          comment="GateCvt[%s]: 2x8bit -> 2xf32 (vi=%d)"%(gateDtype.toNameAbbrev(), vi)))
        for j in range(2):
          if (vi + j) < self.gwvw:
            cvtMod.add(VMovB32(dst=vgpr(dataGate + vi + j), src=vgpr(cvtTmp + j),
              comment="GateCvt[%s]: store f32 (vi=%d)"%(gateDtype.toNameAbbrev(), vi + j)))
    elif gateDtype.isInt8():
      # int8 = 4 signed bytes/dword; extract+sign-extend byte (vi%4) to i32, then -> f32.
      for vi in range(self.gwvw - 1, -1, -1):
        dst = dataGate + vi
        if (vi % 4) != 3:
          cvtMod.add(VMovB32(dst=vgpr(self.tmpVgpr), src=hex((vi % 4) * 8), comment="byte offset"))
          cvtMod.add(VBfeI32(dst=vgpr(dst), src0=vgpr(dataGate + vi // 4), src1=vgpr(self.tmpVgpr), src2=8,
            comment="GateCvt[i8]: int8 -> int32 (vi=%d)"%vi))
        else:
          cvtMod.add(VAShiftRightI32(dst=vgpr(dst), shiftHex=24, src=vgpr(dataGate + vi // 4),
            comment="GateCvt[i8]: int8(byte3) -> int32 (vi=%d)"%vi))
        cvtMod.add(VCvtI32toF32(dst=vgpr(dst), src=vgpr(dst), comment="GateCvt[i8]: int32 -> f32 (vi=%d)"%vi))
    elif gateDtype.isInt32():
      # int32 = 1 value/dword; cvt in place to f32.
      for vi in range(self.gwvw - 1, -1, -1):
        cvtMod.add(VCvtI32toF32(dst=vgpr(dataGate + vi), src=vgpr(dataGate + vi),
          comment="GateCvt[i32]: int32 -> f32 (vi=%d)"%vi))
    else:
      raise RuntimeError(
          "GateResidual: unsupported gate dtype %s" % str(gateDtype))
    return cvtMod

  def _emitHoistedGateLoadPhase(self, module: Module, bufferOOB):
    """opt path: emit the gate prolog load wrapped in ONE null-gate skip branch."""
    if not (self.parentWriter.states.useGateResidual and
            (self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1 or self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel")):
      return
    if not self.ss.optSingleColVgpr:
      return  # no-opt path keeps the per-element load in the prolog loop
    gateList = list(self.kernel["ProblemType"].get("GateResidualDataTypeList", []))
    if not gateList:
      return
    multi = len(gateList) > 1

    labels = self.parentWriter.labels
    endLabel = Label(labels.getNameInc("GateLoadAll_End"), "")
    savedBpe = self.parentWriter.states.bpeGate

    def _emitLoadsFor(gDtype):
      bpe = max(1, int(self.parentWriter.states.bpr * gDtype.numRegisters()))
      self.parentWriter.states.bpeGate = bpe
      self.ss.singleColGateAddrUpdated = False
      seen = set()
      for ei, element in enumerate(self.batchElements):
        addrCalc = self.ss.elementAddr[ei]
        dataGate = self.ss.elementDataGate[ei]
        if dataGate == 0 or dataGate in seen:
          continue
        seen.add(dataGate)
        addrGateVgpr = addrCalc.addrGateVgpr  # opt: gate's own sharedColGateVgprs
        addrCalc.globalOffsetGate = addrCalc.coordOffset0 * bpe
        module.add(addrCalc.emitLdChange(
            self.kernel, self.ss, 'Gate', self.edge, self.beta, self.ss.elementMask[ei],
            bufferOOB, (ei == 0), self.tmpVgpr, self.tmpSgpr, addrGateVgpr, self.addrD, 0))
        module.add(self.parentWriter.readInput(
            self.kernel, self.ss, 'Gate', gDtype, addrCalc, element[3], dataGate,
            self.gwvw, addrGateVgpr, self.tmpS01))

    if not multi:
      # single-dtype: no GateType dispatch needed, just the loads.
      _emitLoadsFor(gateList[0])
    else:
      typeLabels = [Label(labels.getNameInc("GateLoadAll_%s"%g.toNameAbbrev()), "") for g in gateList]
      typeLabels.append(endLabel)
      for i, nextLabel in enumerate(typeLabels[1:]):
        gDtype = gateList[i]
        module.add(typeLabels[i])
        module.add(self.parentWriter.getSCMPKInstruction(
            "LGU32", "GateType", gDtype.value, comment="GateType != %u"%gDtype.value))
        module.add(SCBranchSCC1(nextLabel.getLabelName(), "Branch if true (try next gate dtype)"))
        _emitLoadsFor(gDtype)
        module.add(SBranch(labelName=endLabel.getLabelName(), comment="Branch to GateLoadAll end"))
    module.add(endLabel)
    self.parentWriter.states.bpeGate = savedBpe

  def _emitGateCvtAllPhase(self, module: Module):
    """Multi-dtype Gate: convert ALL loaded gate values to f32 in one GateType
    dispatch."""
    if not (self.parentWriter.states.useGateResidual and
            (self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1 or self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel")):
      return
    gateList = list(self.kernel["ProblemType"].get("GateResidualDataTypeList", []))
    if len(gateList) <= 1:
      if self.ss.optSingleColVgpr:
        module.add(SWaitCnt(vlcnt=0, comment="GateResidual: wait for hoisted single-dtype gate loads"))
      return

    labels = self.parentWriter.labels
    cvtEndLabel = Label(labels.getNameInc("GateCvtAll_End"), "")
    cvtTypeLabels = [Label(labels.getNameInc("GateCvtAll_%s"%g.toNameAbbrev()), "") for g in gateList]
    cvtTypeLabels.append(cvtEndLabel)

    # Skip the whole gate cvt phase when the gate pointer is null (no-gate problem).
    module.add(self.parentWriter.getSCMPKInstruction(
        "EQU32", "SrdGate+2", 0, comment="gate disabled? (SrdGate num_records==0)"))
    module.add(SCBranchSCC1(cvtEndLabel.getLabelName(), "skip gate cvt if disabled (null gate)"))

    # One wait for all prolog gate loads (multi-dtype already paid a blunt wait
    # per element previously; this hoists it to a single wait).
    module.add(SWaitCnt(vlcnt=0, comment="GateResidual: wait for prolog gate loads (branch-once cvt)"))
    for i, nextLabel in enumerate(cvtTypeLabels[1:]):
      gDtype = gateList[i]
      module.add(cvtTypeLabels[i])
      module.add(self.parentWriter.getSCMPKInstruction(
          "LGU32", "GateType", gDtype.value, comment="GateType != %u"%gDtype.value))
      module.add(SCBranchSCC1(nextLabel.getLabelName(), "Branch if true (try next gate dtype)"))
      seen = set()
      for ei in range(len(self.batchElements)):
        dg = self.ss.elementDataGate[ei]
        if dg != 0 and dg not in seen:
          seen.add(dg)
          module.add(self._emitGateCvt(dg, gDtype))
      module.add(SBranch(labelName=cvtEndLabel.getLabelName(), comment="Branch to GateCvtAll end"))
    module.add(cvtEndLabel)

  def _emitNonatomicAdd(self, module: Module):
    ########################################
    # Not Atomic
    ########################################
    # edge has v_cndmask so loads or stores may not issue, hard to track vmcnt:
    interleaveStoreVmcnt = self.parentWriter.states.interleaveStoreVmcnt and not self.edge

    for elementIdx in range(len(self.batchElements)):
      for vi in range(self.gwvw):
        sumIdxV = self.ss.elementSumIdx[elementIdx] + vi
        newSumIdxV = sumIdxV - self.parentWriter.states.c.startVgprValu
        # covers sgemm, gemm_ex(HHS/HSS/BBS/BSS (HPA=T)), int8 (int8x4?)
        if self.kernel["ProblemType"]["ComputeDataType"].isInt32() or \
            self.kernel["ProblemType"]["ComputeDataType"].isSingle(): # covers sgemm/gemm_ex(HHS/HSS/BBS/BSS)
            if self.debugConfig["ForceExpectedValue"]:
              module.add(VMovB32(vgpr("ValuC+%u"%newSumIdxV), self.debugConfig["ValueCExpectedValue"], "force expected value" ))
            if self.parentWriter.db["ForceVSerial"]:
              module.add(VMovB32(vgpr("ValuC+%u"%newSumIdxV), vgpr("Serial"), "force expected value to serial" ))
            if self.parentWriter.db["CheckValueC"]:
              module.add(SMovB32(sgpr(self.tmpS01), self.debugConfig["ValueCExpectedValue"], "Move expected value"))
              module.add(self.parentWriter.getCmpAssert(self.parentWriter.asmAssert.eq, vgpr("ValuC+%u"%newSumIdxV), sgpr(self.tmpS01)))

    ########################################
    # wait for batched load
    # Here we wait all
    if not interleaveStoreVmcnt:
      waitcntInst = self.globalStoreWait(0, [], 0, 0, False)
      if waitcntInst:
        module.add(waitcntInst)

    if self.kernel["ProblemType"]["StochasticRounding"]:
      if self.parentWriter.states.asmCaps["v_prng_b32"]:
        vgprRND = self.parentWriter.vgprPool.checkOut(1, tag="_emitNonatomicAdd_vgprRND")
      else:
        # legacy PRNG approach needs extra 2 VGPRs
        # Ref.: Module("StochasticRoundingCvt")
        vgprRND = self.parentWriter.vgprPool.checkOut(3, tag="_emitNonatomicAdd_vgprRND2")

    # Multi-dtype Gate: branch-once cvt of all loaded gate values to f32.
    self._emitGateCvtAllPhase(module)

    module.addComment1("apply mask, calc new C and issue writes")
    # module.add(self.getBomb()) # can see store addresses just before the store inst

    activationCDataType = self.kernel["ProblemType"]["ActivationComputeDataType"]

    if self.kernel["_GlobalAccumulation"] != 'MultipleBuffer':
      if self.kernel["ProblemType"]["DestDataType"].isBFloat16() and self.kernel["ProblemType"]["HighPrecisionAccumulate"]:
        module.add(VMovB32(vgpr(self.cvtVgprStruct.vgprBf16Mask), "0xffff0000", comment="mask for pack two bfloat16 element to 32bit" ))
        module.add(VMovB32(vgpr(self.cvtVgprStruct.vgprFp32Nan), "0x7fff0000", comment="fp32 Nan" ))
        module.add(VMovB32(vgpr(self.cvtVgprStruct.vgprBf16Inc), "0x7fff", comment="rounding bias for bfloat16" ))
    # is16bitSubtile: controls partner-lane address setup for dwordx4 paired-subtile stores.
    # Must match is16bitSubtilePaired (per-element store dispatch) exactly.
    # Excluded for "MultipleBufferSingleKernel" and "MultipleBuffer" (StreamK partial-tile
    # workspace path) — both write float32 to workspace, not 16bit to D output.
    isSubtileNonEdge = (
      self.kernel.get("UseSubtileImpl") and not self.edge
      and self.kernel["_GlobalAccumulation"] not in ("MultipleBufferSingleKernel", "MultipleBuffer")
    )
    is16bitSubtile = (
      isSubtileNonEdge
      and (self.kernel["ProblemType"]["DestDataType"].isBFloat16() or
           self.kernel["ProblemType"]["DestDataType"].isHalf())
      and self.kernel["ProblemType"]["HighPrecisionAccumulate"]
      and self.kernel["WavefrontSize"] != 32  # wave32: skip permute-based packed store (uses wave64-only ops)
    )
    if is16bitSubtile:
      assert self.kernel["BufferStore"], \
        "UseSubtileImpl 16bit optimized store requires BufferStore=1"
      # Compute ds_permute partner-lane address for 16bit dwordx4 paired-subtile stores.
      # vtmp1 = lane_id, vtmp2 = partner_lane_id (for ds_permute forward scatter)
      # After v_permlane32_swap + v_permlane16_swap + exec masking:
      #   each lane ends up with the lane_id of the partner that will scatter data to it.
      # vPermAddr = partner_lane_id * 4  (byte address for ds_permute_b32)
      vPermAddr = self.cvtVgprStruct.vgprPermAddr
      vTmp = self.cvtVgprStruct.vgprBf16Temp  # reuse scratch temp before it's used for mask init
      module.addComment1("16bit dwordx4 UseSubtileImpl: compute ds_permute partner-lane address")
      module.add(VAndB32(dst=vgpr(vTmp),     src0=self.kernel["WavefrontSize"]-1, src1=vgpr("Serial"), comment="lane_id & (WS-1)"))
      module.add(VAndB32(dst=vgpr(vPermAddr), src0=self.kernel["WavefrontSize"]-1, src1=vgpr("Serial"), comment="copy of lane_id"))
      module.add(VPermlane32SwapB32(dst=vgpr(vTmp), src=vgpr(vTmp), comment="lane XOR 32 swap"))
      module.add(SNop(waitState=0, comment="delay after v_permlane32_swap"))
      module.add(VPermlane16SwapB32(dst=vgpr(vTmp), src=vgpr(vTmp), comment="lane XOR 16 swap"))
      # Exec mask: lanes where both XOR swaps changed the value (i.e., the 'first' half of each pair)
      # selects lanes 0-15 and 32-47 within the wave.
      stmp = self.parentWriter.sgprPool.checkOutAligned(2,2, tag="_emitNonatomicAdd_stmp")
      module.add(SMovB32(dst=sgpr(stmp), src="0x0000ffff", comment="select lanes 0-15, 32-47"))
      module.add(SMovB32(dst=sgpr(stmp+1), src="0xffff0000"))
      module.add(VCndMaskB32(dst=vgpr(vTmp), src0=vgpr(vTmp), src1=vgpr(vPermAddr), src2=sgpr(stmp,2), comment="restore original lane_id for selected lanes"))
      self.parentWriter.sgprPool.checkIn(stmp)
      module.add(VLShiftLeftB32(dst=vgpr(vPermAddr), shiftHex=2, src=vgpr(vTmp), comment="partner_lane * 4 = ds_permute byte addr"))
      # Pre-compute lane_group*8 once; reused as the row-byte address correction in every
      # paired dwordx4 store (addrDVgpr encodes lane_group*8 but we need lane_group*16).
      vLGDelta = self.cvtVgprStruct.vgprLaneGroupDelta
      module.addComment1("16bit dwordx4: pre-compute lane_group*8 row-byte correction")
      module.add(VAndB32(dst=vgpr(vLGDelta), src0=self.kernel["WavefrontSize"]-1, src1=vgpr("Serial"),
                         comment="lane_id = Serial & (WS-1)"))
      module.add(VLShiftRightB32(dst=vgpr(vLGDelta), shiftHex=4, src=vgpr(vLGDelta),
                                 comment="lane_group = lane_id >> 4"))
      module.add(VLShiftLeftB32(dst=vgpr(vLGDelta), shiftHex=3, src=vgpr(vLGDelta),
                                comment="vgprLaneGroupDelta = lane_group * 8"))
      # Compute bpe scale shift once (compile-time constant); used inside
      # _emit16bitSubtilePairedStore to adjust addrDVgpr inline without
      # modifying it in place, so no restore loop is needed after the stores.
    elif self.kernel["_GlobalAccumulation"] != 'MultipleBuffer':
      if self.kernel["ProblemType"]["DestDataType"].isFloat8_fnuz() and self.kernel["ProblemType"]["HighPrecisionAccumulate"]:
        module.add(VMovB32(vgpr(self.cvtVgprStruct.vgprFp8NanInf), "0x207", comment="Nan and +/- inf" ))
        module.add(VMovB32(vgpr(self.cvtVgprStruct.vgprFp8Max), "0x43700000", comment="Fp8 Max value 240 as float32" ))
        module.add(VMovB32(vgpr(self.cvtVgprStruct.vgprFp8Min), "0xc3700000", comment="Fp8 Min value -240 as float32" ))
      elif self.kernel["ProblemType"]["DestDataType"].isFloat8() and self.kernel["ProblemType"]["HighPrecisionAccumulate"]:
        module.add(VMovB32(vgpr(self.cvtVgprStruct.vgprFp8NanInf), "0x207", comment="Nan and +/- inf" ))
        module.add(VMovB32(vgpr(self.cvtVgprStruct.vgprFp8Max), "0x43E00000", comment="Fp8 Max value 448 as float32" ))
        module.add(VMovB32(vgpr(self.cvtVgprStruct.vgprFp8Min), "0xc3E00000", comment="Fp8 Min value -448 as float32" ))
      elif self.kernel["ProblemType"]["DestDataType"].isAnyBFloat8() and self.kernel["ProblemType"]["HighPrecisionAccumulate"]:
        module.add(VMovB32(vgpr(self.cvtVgprStruct.vgprBF8NanInf), "0x207", comment="Nan and +/- inf" ))
        module.add(VMovB32(vgpr(self.cvtVgprStruct.vgprBF8Max), "0x47600000", comment="BF8 Max value 57344 as float32" ))
        module.add(VMovB32(vgpr(self.cvtVgprStruct.vgprBF8Min), "0xc7600000", comment="BF8 Min value -57344 as float32" ))

    storeCode = Module("GroupLoadStore")
    vlcntTotalIssued = self.loadsBetaIssued + self.loadsEIssued + self.loadsGateIssued
    dscntTotalIssued = self.localLoadsBiasIssued + self.loadsScaleAVecIssued + self.loadsScaleBVecIssued + self.loadsScaleAlphaVecIssued
    waitCnter = [vlcntTotalIssued, dscntTotalIssued]
    for elementIdx in range(0, len(self.batchElements)):
      element = self.batchElements[elementIdx]
      addrCalc: AddrCalculation = self.ss.elementAddr[elementIdx]
      addr = addrCalc.addrDVgpr
      dataE = self.ss.elementDataE[elementIdx]
      dataGate = self.ss.elementDataGate[elementIdx] if self.parentWriter.states.useGateResidual else 0
      dataBias = self.ss.elementDataBias[elementIdx]
      dataScaleAVec = self.ss.elementDataScaleAVec[elementIdx]
      dataScaleBVec = self.ss.elementDataScaleBVec[elementIdx]
      dataScaleAlphaVec = self.ss.elementDataScaleAlphaVec[elementIdx]
      mask = self.ss.elementMask[elementIdx]
      vc0 = element[3]

      # When skipRearrangement is True:
      # - Data is at WMMA output (v[0:N]), not at elementSumIdx (v[144:N])
      # - Store should read from WMMA output directly
      # Note: Beta paths need the rearrangement because they load C and compute rC = alpha*rC + beta*C
      #
      # skipRearrangement requires elementSumIdx to be contiguous so that
      # (elementSumIdx[i] - elementSumIdx[0]) == i.  When checkOutAligned
      # returns non-contiguous blocks (gaps in the pool), the relative offset
      # no longer matches the sequential MFMA output layout, causing stores
      # to read from wrong VGPRs.
      esIdx = self.ss.elementSumIdx
      contiguous = all(esIdx[i] - esIdx[0] == i for i in range(len(esIdx)))
      if self.skipRearrangement and contiguous:
        sumIdx = esIdx[elementIdx] - esIdx[0]
      else:
        sumIdx = self.ss.elementSumIdx[elementIdx]

      # CompactLoopStore look-ahead override for this store loop (see
      # _lookaheadRowInc). Separate for-loop pass from the _prolog one, so it is
      # recomputed here.
      _emitOverrideRows = self._lookaheadRowInc(elementIdx)

      # print(str(element)+" rowInc="+str(addrCalc.rowInc))
      # Already write wave column block into LDS
      # Now read lds data back to registers and write to global memroy
      if self.kernel["_GlobalAccumulation"] != "MultipleBufferSingleKernel":
        if self.ss.optSrdIncForRow and (addrCalc.rowInc or (self.kernel["CompactLoopStore"] and elementIdx == 0 and self.batchIdx == 0)) and self.kernel["StoreRemapVectorWidth"] > 0:
          module.addComment1("StoreRemap: shift coord1 address")
          if self.kernel["ProblemType"]["UseE"] and (self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1):
            # TODO Check if works with StreamK
            printExit("Use E does not support StoreRemapVectorWidth if GSU == 1.")
            # module.add(addrCalc.incrementToNextRow(self.kernel, "E", self.ss, self.tmpS01, isCompute=True))
          module.add(addrCalc.incrementToNextRow(self.kernel, "D", self.ss, self.tmpS01, forceinitrow0=1, overrideAfterPrimerRows=_emitOverrideRows))
          module.add(VMovB32(vgpr(self.tmpVgpr), addrCalc.rowInc, comment="set shift rows"))
          module.add(VAddU32(vgpr(self.parentWriter.vgprs.storeRemapCoord1), vgpr(self.parentWriter.vgprs.storeRemapCoord1), vgpr(self.tmpVgpr), "shift storeRemap coord1"))

      # When stores are interleaved (GLS=0) with subtile NonEdge guards, the
      # M-guard branch for the last store in N-group K targets the N-group end
      # label.  That label must be placed BEFORE the beta*C fmacs for N-group K+1,
      # otherwise the M-guard branch skips the fmacs and the next N-group stores zeros.
      if isSubtileNonEdge and not self.kernel["GroupLoadStore"]:
        blockIdxN = element[0]
        if blockIdxN != self._subtilePrevBlockIdxN and self._subtileNGroupSkipLabel is not None:
          module.add(self._subtileNGroupSkipLabel)
          self._subtileNGroupSkipLabel = None
          if self._subtilePendingSrdDInc is not None:
            module.add(self._subtilePendingSrdDInc)
            self._subtilePendingSrdDInc = None

      # apply in-bounds exec mask
      if self.edge and not self.kernel["BufferStore"]:
        module.add(self.getEdgeMovInstType()(EXEC(), sgpr(mask, self.laneSGPRC), "sgprs -> exec"))

      if interleaveStoreVmcnt:
        waitcntInst = self.globalStoreWait(elementIdx, waitCnter, vlcntTotalIssued, dscntTotalIssued, True)
        if waitcntInst:
          module.addSpaceLine()
          module.add(waitcntInst)

      def applyScaleVec(vecModule, addressStr, dataScaleVec, factorDim, isGlobal=True):
        if not self.beta and not self.applyAlpha: # case for beta-0 and alpha == 1,(OptNLL)
          if (self.kernel["ProblemType"]["DestDataType"].isInt8() or self.kernel["ProblemType"]["DestDataType"].isInt32() or \
              (self.kernel["ProblemType"]["DataType"].isInt8() and self.kernel["ProblemType"]["DestDataType"].isHalf()) or \
              (self.kernel["ProblemType"]["DataType"].isInt8() and self.kernel["ProblemType"]["DestDataType"].isBFloat16())) and \
            self.kernel["ProblemType"]["ComputeDataType"].isSingle():
            module.add(convertData(self.gwvw, self.ss.elementSumIdx[elementIdx], cvtType=CvtType.CVT_I32_to_F32, \
                                        inputPrefix="ValuC+", prefixOffset=self.parentWriter.states.c.startVgprValu))

        if self.kernel["ProblemType"]["ComputeDataType"].isSingle():
          maskConst = 1.0
        elif self.kernel["ProblemType"]["ComputeDataType"].isInt32():
          maskConst = 1

        gwvw = 1 if factorDim else self.gwvw
        if isGlobal:
          vecModule.add(VCmpGtU32(dst=sgpr("Address%s"%addressStr, self.parentWriter.states.laneSGPRCount), src0=sgpr("Srd%s+2"%addressStr), src1=0, comment=" == 0 ?"))
          for vi2 in range(0, gwvw):
            vecModule.add(VCndMaskB32(
              dst=vgpr(dataScaleVec + vi2), \
              src1=vgpr(dataScaleVec + vi2), \
              src0=maskConst, \
              src2=sgpr("Address%s"%addressStr, self.parentWriter.states.laneSGPRCount), \
              comment="1. mul 1 if 0"))
        if factorDim and self.gwvw > 1:
          vecModule.add(VMovB32(dst=vgpr(dataScaleVec+1), src=vgpr(dataScaleVec), comment="copy data%s to data%s+1"%(addressStr, addressStr)))

        for vi in range(0, self.gwvw):
          inputScaleVecVgpr = dataScaleVec + (0 if factorDim else vi)
          sumIdxV   = self.ss.elementSumIdx[elementIdx] + vi
          if self.kernel["ProblemType"]["ComputeDataType"].isSingle():
            vgprIdx = sumIdxV - self.parentWriter.states.c.startVgprValu
            # Generate single f32 code if edge is detected.
            if ((vi + 1) == self.gwvw) and ((self.gwvw % 2) == 1):
              vecModule.add(VMulF32(dst=vgpr("ValuC+%d"%vgprIdx), src0=vgpr(inputScaleVecVgpr), src1=vgpr("ValuC+%d"%vgprIdx), comment="*= %sVMul"%addressStr ))
            # Original packed route
            elif vi%2 == 1:
              assert (self.gwvw % 2 == 0)
            else:
              vecModule.add(VMulPKF32(dst=vgpr("ValuC+%d"%vgprIdx, 2), src0=vgpr(inputScaleVecVgpr, 2), src1=vgpr("ValuC+%d"%vgprIdx, 2), comment="*= %sVMulPK(%d)(%d)"%(addressStr, dataScaleVec,vi)))
          elif self.kernel["ProblemType"]["ComputeDataType"].isInt32():
            vgprIdx = sumIdxV - self.parentWriter.states.c.startVgprValu
            # Generate single i32 code if edge is detected.
            if ((vi + 1) == self.gwvw) and ((self.gwvw % 2) == 1):
              vecModule.add(VMulLOU32(dst=vgpr("ValuC+%d"%vgprIdx), src0=vgpr(inputScaleVecVgpr), src1=vgpr("ValuC+%d"%vgprIdx), comment="*= %sVMul"%addressStr ))
            elif vi%2 == 1:
              assert (self.gwvw % 2 == 0)
            else:
              vecModule.add(VMulLOU32(dst=vgpr("ValuC+%d"%vgprIdx), src0=vgpr(inputScaleVecVgpr), src1=vgpr("ValuC+%d"%vgprIdx), comment="*= %sVMulPK(%d)(%d)"%(addressStr, dataScaleAlphaVec,vi)))
              vecModule.add(VMulLOU32(dst=vgpr("ValuC+%d"%(vgprIdx+1)), src0=vgpr(inputScaleVecVgpr+1), src1=vgpr("ValuC+%d"%(vgprIdx+1)), comment="*= %sVMulPK(%d)(%d)"%(addressStr, dataScaleAlphaVec,vi)))
          else:
            raise RuntimeError("Unsupported %s compute data type %s."%(addressStr, str(self.kernel["ProblemType"]["ComputeDataType"])))

      isSingleKernel = ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel") or self.kernel["StreamK"] > 0

      scaleAVecModule = Module("ScaleAVecModule")
      scaleBVecModule = Module("ScaleBVecModule")
      if (self.kernel["ProblemType"]["UseScaleAB"] == "Vector") and isSingleKernel:
        applyScaleVec(scaleAVecModule, "ScaleA", dataScaleAVec, 0, isGlobal=False)
        applyScaleVec(scaleBVecModule, "ScaleB", dataScaleBVec, 1, isGlobal=False)
      module.add(scaleAVecModule)
      module.add(scaleBVecModule)

      scaleAlphaVecModule = Module("scaleAlphaVecModule")
      if self.kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel:
        applyScaleVec(scaleAlphaVecModule, "ScaleAlphaVec", dataScaleAlphaVec, self.factorDim, isGlobal=False)
      module.add(scaleAlphaVecModule)

      if self.beta:
        module.add(self._addSumAlphaWithCBeta(self.kernel, self.ss, self.gwvw, elementIdx, vc0, self.tmpVgpr, self.cvtVgprStruct))
      elif ((self.parentWriter.states.useBias == DataDirection.READ) or self.kernel["ActivationFuncCall"]) and not self.applyAlpha \
        and not ( self.kernel["ProblemType"]["UseScaleAlphaVec"] and isSingleKernel): # case of alpha=1 and beta=0
        if (self.kernel["ProblemType"]["DestDataType"].isInt8() or self.kernel["ProblemType"]["DestDataType"].isInt32() or \
            (self.kernel["ProblemType"]["DataType"].isInt8() and self.kernel["ProblemType"]["DestDataType"].isHalf()) or \
            (self.kernel["ProblemType"]["DataType"].isInt8() and self.kernel["ProblemType"]["DestDataType"].isBFloat16())) and \
           self.kernel["ProblemType"]["ComputeDataType"].isSingle():
          module.add(convertData(self.gwvw, self.ss.elementSumIdx[elementIdx], cvtType=CvtType.CVT_I32_to_F32, \
                                      inputPrefix="ValuC+", prefixOffset=self.parentWriter.states.c.startVgprValu))

      # Add bias
      mergeActFuncCall = False
      if self.parentWriter.states.useBias == DataDirection.READ:
        if activationCDataType == self.kernel["ProblemType"]["ComputeDataType"] and self.kernel["ActivationFuncCall"]:
          mergeActFuncCall = True
        if (self.kernel["ProblemType"]["Gradient"] and self.kernel["ProblemType"]["ActivationType"] != 'none' and self.kernel["ProblemType"]["UseE"]) and ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["StreamK"] > 0):
          mergeActFuncCall = False

        if self.factorDim and self.gwvw > 1:
          module.add(VMovB32(dst=vgpr(dataBias+1), src=vgpr(dataBias), comment="copy dataBias to dataBIas+1"))

        for vi in range(0, self.gwvw):
          inputVgpr = dataBias + + (0 if self.factorDim else vi)
          sumIdxV   = self.ss.elementSumIdx[elementIdx] + vi
          if self.kernel["ProblemType"]["ComputeDataType"].isSingle():
            vgprIdx = sumIdxV - self.parentWriter.states.c.startVgprValu
            vgprDst = (self.activationSetPCStruct.vgprActCopy + vi) if mergeActFuncCall else "ValuC+%d"%vgprIdx
            # Generate single f32 code if edge is detected.
            if ((vi + 1) == self.gwvw) and ((self.gwvw % 2) == 1):
              module.add(VAddF32(dst=vgpr(vgprDst), src0=vgpr(inputVgpr), src1=vgpr("ValuC+%d"%vgprIdx), \
                                 comment="C += bias"))

            # Original packed route
            elif vi%2 == 1:
              assert (self.gwvw % 2 == 0)
            else:
              module.add(VAddPKF32(dst=vgpr(vgprDst, 2), src0=vgpr(inputVgpr, 2), \
                                   src1=vgpr("ValuC+%d"%vgprIdx, 2), comment="C += bias"))
          else:
            raise RuntimeError("Unsupported bias compute data type %s."%str(self.kernel["ProblemType"]["ComputeDataType"]))

      # Gate Residual: acc = gate*acc + gate. Built here but DEFERRED — added after
      # activation+scaleD (see module.add(gateModule) below), on the f32 ValuC.
      gateModule = Module("Empty GateResidual")
      if self.parentWriter.states.useGateResidual and \
         (self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1 or self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel"):
        # TODO: I8 input (int32 compute) — needs int32 acc -> cvt f32 -> FMA ->
        # saturate-cast. Not supported yet; reject below.
        if not self.kernel["ProblemType"]["ComputeDataType"].isSingle():
          raise RuntimeError(
              "GateResidual currently requires ComputeDataType=f32 (HPA); got %s" %
              str(self.kernel["ProblemType"]["ComputeDataType"]))

        gateList = list(self.kernel["ProblemType"].get("GateResidualDataTypeList", []))
        if not gateList:
          # Default fallback: use DestDataType as the gate dtype.
          gateList = [self.kernel["ProblemType"]["DestDataType"]]

        def _emit_gate_fma():
          """Identity (branchless null): ValuC = (gate+s)*ValuC + gate.
          GateNullOne (s) = 1.0 if gate null else 0.0 -> null:(0+1)*acc+0=acc; real:gate*acc+gate."""
          fmaMod = Module("GateFMA")
          for vi in range(0, self.gwvw):
            sumIdxV = self.ss.elementSumIdx[elementIdx] + vi
            vgprIdx = sumIdxV - self.parentWriter.states.c.startVgprValu
            fmaMod.add(VAddF32(dst=vgpr(self.tmpVgpr), src0=vgpr(dataGate + vi),
              src1=sgpr("GateNullOne"), comment="gate + s (vi=%d)"%vi))
            fmaMod.add(VFmaF32(
              dst=vgpr("ValuC+%d"%vgprIdx),
              src0=vgpr(self.tmpVgpr),
              src1=vgpr("ValuC+%d"%vgprIdx),
              src2=vgpr(dataGate + vi),
              comment="GateResidual identity: acc = (gate+s)*acc + gate (vi=%d)"%vi))
          return fmaMod

        if len(gateList) == 1:
          gateModule.add(self._emitGateCvt(dataGate, gateList[0]))
          gateModule.add(_emit_gate_fma())
        else:
          gateModule.add(_emit_gate_fma())

      if (self.kernel["ProblemType"]["UseE"] and not self.kernel["ProblemType"]["Gradient"]) and ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["StreamK"] > 0):
        vgprIdx   = self.ss.elementSumIdx[elementIdx] - self.parentWriter.states.c.startVgprValu
        vgprDst   = self.activationSetPCStruct.vgprActCopy if mergeActFuncCall else vgprIdx
        prefixStr = "" if mergeActFuncCall else "ValuC+"
        prefixOffset = 0 if mergeActFuncCall else self.parentWriter.states.c.startVgprValu
        # Packdata if needed
        tmpVgpr = self.tmpVgpr
        if mergeActFuncCall:
          tmpVgpr += self.gwvw * self.kernel["ProblemType"]["ComputeDataType"].numRegisters()
        if self.kernel["ProblemType"]["ComputeDataType"].isSingle():
          if self.kernel["ProblemType"]["DataTypeE"].isHalf():
            packdata = PackData_F16()
            module.add(packdata(self.gwvw, tmpVgpr, vgprDst, tmpVgpr=tmpVgpr, inputPrefix=prefixStr, prefixOffset=prefixOffset))
            vgprDst = tmpVgpr
          elif self.kernel["ProblemType"]["DataTypeE"].isBFloat16():
            packdata = PackData_BF16()
            module.add(packdata(self.gwvw, tmpVgpr, vgprDst, self.cvtVgprStruct, self.tmpS01, self.laneSGPRC,
                                tmpVgpr=tmpVgpr, inputPrefix=prefixStr, prefixOffset=prefixOffset))
            vgprDst = tmpVgpr
          elif self.kernel["ProblemType"]["DataTypeE"].isSingle():
            if not mergeActFuncCall:
              vgprDst = "ValuC+%d" % vgprDst
          elif self.kernel["ProblemType"]["DataTypeE"].isFloat8():
            packdata = PackData_FLOAT8()
            module.add(packdata(self.gwvw, tmpVgpr, vgprDst, self.cvtVgprStruct, self.tmpS01, self.laneSGPRC,
                                inputPrefix=prefixStr, prefixOffset=prefixOffset))
            vgprDst = tmpVgpr
          elif self.kernel["ProblemType"]["DataTypeE"].isFloat8_fnuz():
            packdata = PackData_FLOAT8_fnuz()
            module.add(packdata(self.gwvw, tmpVgpr, vgprDst, self.cvtVgprStruct, self.tmpS01, self.laneSGPRC,
                                inputPrefix=prefixStr, prefixOffset=prefixOffset))
            vgprDst = tmpVgpr
          else:
            printExit("Unsupport type for E output. (%s)"%self.kernel["ProblemType"]["DataTypeE"].toEnum())
        else:
          printExit("Unsupport compute type for E output. (%s)"%self.kernel["ProblemType"]["ComputeDataType"].toEnum())

        module.add(self.parentWriter.addStore(self.kernel, self.ss, 'E', addrCalc, vgprDst, self.tmpS01, self.edge, comment="store E"))

      SaturateTypeInt8 = SaturateCastType.NORMAL

      gradientCvtModule = Module("gradientCvtModule")
      if (self.kernel["ProblemType"]["UseE"] and self.kernel["ProblemType"]["Gradient"]) and ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["StreamK"] > 0):
        loadOffset = int((self.kernel["ProblemType"]["ComputeDataType"].numRegisters() - self.kernel["ProblemType"]["DataTypeE"].numRegisters()) * self.ss.cfg.gwvw)
        if activationCDataType != self.kernel["ProblemType"]["DataTypeE"]:
          if activationCDataType.isSingle() and self.kernel["ProblemType"]["DataTypeE"].isHalf():
            for vi in range(0, self.gwvw):
              dataEV  = dataE + vi
              dataEV2 = dataE + vi // 2
              selectbit = HighBitSel.LOW if (self.gwvw != 1 and vi % 2 == 0) or (self.gwvw == 1 and elementIdx % 2 == 0) else HighBitSel.HIGH
              gradientCvtModule.add(ECvtF16toF32(dst=vgpr(dataEV), src=vgpr(dataEV2+loadOffset), sel=selectbit, comment="gwvw %d, elementIdx %d"%(self.gwvw, elementIdx)))
          elif activationCDataType.isSingle() and self.kernel["ProblemType"]["DataTypeE"].isBFloat16():
            for vi in range(0, self.gwvw):
              dataEV  = dataE + vi
              dataEV2 = dataE + vi // 2
              # Consider bf16 without packing (gwvw==1)
              # TODO: check correctness for gfx950
              selectWord = 0 if (self.gwvw != 1 and vi % 2 == 0) or (self.gwvw == 1) else 1
              module.add(VCvtBF16toFP32(dst=vgpr(dataEV), src=vgpr(dataEV2+loadOffset), vgprMask=vgpr(self.cvtVgprStruct.vgprBf16Mask), vi=(selectWord), comment="gwvw %d, elementIdx %d"%(self.gwvw, elementIdx)))
          else:
            printExit("[Gradient input] Unsupported conversion.")

      # Activation
      activationModule = None
      isActivationInsertAfter = False
      if self.kernel["ProblemType"]["Gradient"] and ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["StreamK"] > 0):
        gradientInput = dataE
        enableValuC   = False
      else:
        gradientInput = self.ss.elementSumIdx[elementIdx]
        enableValuC   = True
        if self.kernel["LocalSplitU"] > 1:
          # When LSU > 1, the VGPRs are from LSU output.
          # the elementSumIdx has indicated the VGPRs from LSU.
          # Don't use the ValuC prefix here.
          enableValuC = False
        # gradientInput is an absolute vgpr index (= elementSumIdx). The
        # ActivationFuncCall/copyData and getActivationDestDataType consumers below
        # address it directly as an absolute vgpr, so keep it absolute here. Only
        # getActivationActivationComputeType() uses the "ValuC+N" relative namespace
        # (when enableValuC is set); that startVgprValu offset is applied at its
        # call site below.
      if self.kernel["ActivationFuncCall"]:
        if (activationCDataType == self.kernel["ProblemType"]["DestDataType"]) and \
          (activationCDataType != self.kernel["ProblemType"]["ComputeDataType"]) and ((self.kernel["ProblemType"]["UseScaleCD"] == False) or (self.kernel["ProblemType"]["UseScaleAlphaVec"] == False)):
          isActivationInsertAfter = True
        activationModule = Module("ActivationFuncCall")
        if (not mergeActFuncCall) and (not isActivationInsertAfter):
          activationModule.appendModule (copyData(activationCDataType, gradientInput, self.gwvw, \
            self.activationSetPCStruct.vgprActCopy))
        swappc = SSwapPCB64(dst=sgpr(self.activationSetPCStruct.sgprOffsetBack, 2), \
          src=sgpr(self.activationSetPCStruct.sgprOffsetActivation, 2))
        calleeLabelsByGwvw = getattr(self.activationSetPCStruct, "calleeLabelsByGwvw", None)
        if calleeLabelsByGwvw:
          calleeFuncs = list(calleeLabelsByGwvw.get(self.gwvw, ()))
          if calleeFuncs:
            swappc.calleeFuncs = calleeFuncs
        activationModule.add(swappc)
        activationModule.appendModule (copyData(activationCDataType, gradientInput, self.gwvw, \
          self.activationSetPCStruct.vgprActCopy, 1))
      elif self.parentWriter.insertActivationAfterPacked(self.kernel, self.activationTypeStr) and (self.kernel["ProblemType"]["UseScaleAlphaVec"] == False):
        isActivationInsertAfter = True
        activationModule = self.parentWriter.getActivationDestDataType(self.kernel, self.activation, \
          self.activationTypeStr, self.gwvw, gradientInput , gradientInput, self.tmpVgpr, self.tmpSgpr)
      else:
        satInt8 = False
        if self.kernel["ProblemType"]["DestDataType"].isInt8():
          if (self.activationTypeStr == 'abs') or (self.activationTypeStr == 'relu'):
            SaturateTypeInt8 = SaturateCastType.DO_NOTHING
            satInt8 = True
        # getActivationActivationComputeType emits on the "ValuC+N" relative
        # register namespace when enableValuC is set, so make its input index
        # relative to startVgprValu. The offset is scoped to this consumer only;
        # the absolute-vgpr consumers above are left untouched.
        actComputeInput = gradientInput
        if enableValuC:
          actComputeInput = gradientInput - self.parentWriter.states.c.startVgprValu
        activationModule = self.parentWriter.getActivationActivationComputeType(self.kernel, self.activation, \
          self.activationTypeStr, self.gwvw, actComputeInput, actComputeInput, self.tmpVgpr, self.tmpSgpr, satInt8, enableValuC)
      # Add C *= GradientAct
      if self.kernel["ProblemType"]["ActivationType"] != 'none' and self.kernel["ProblemType"]["Gradient"] and ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["StreamK"] > 0):
        if isActivationInsertAfter:
          assert 0, "Gradient does not support isActivationInsertAfter."
        for vi in range(0, self.gwvw):
          sumIdxV = self.ss.elementSumIdx[elementIdx] + vi
          dataEV  = dataE + vi
          if self.kernel["ProblemType"]["ComputeDataType"].isSingle():
            vgprIdx = sumIdxV - self.parentWriter.states.c.startVgprValu
            # Generate single f32 code if edge is detected.
            if ((vi + 1) == self.gwvw) and ((self.gwvw % 2) == 1):
              activationModule.add(VMulF32(dst=vgpr("ValuC+%d"%vgprIdx), src0=vgpr("ValuC+%d"%vgprIdx), src1=vgpr(dataEV), comment="C *= GradAct"))
            # Original packed route
            elif vi%2 == 1:
              assert (self.gwvw % 2 == 0)
            else:
              activationModule.add(VMulPKF32(dst=vgpr("ValuC+%d"%vgprIdx, 2), src0=vgpr("ValuC+%d"%vgprIdx, 2), src1=vgpr(dataEV, 2), comment="C *= GradAct"))
          else:
            assert 0, "Unsupported gradient type"

      scaleDModule = Module("Empty scaleDModule")
      if self.kernel["ProblemType"]["UseScaleCD"] and ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["StreamK"] > 0):
        for vi in range(0, self.gwvw):
          sumIdxV = self.ss.elementSumIdx[elementIdx] + vi
          if self.kernel["ProblemType"]["ComputeDataType"].isSingle():
            vgprIdx = sumIdxV - self.parentWriter.states.c.startVgprValu
            # Generate single f32 code if edge is detected.
            if ((vi + 1) == self.gwvw) and ((self.gwvw % 2) == 1):
              if self.kernel["ProblemType"]["OutputAmaxD"]:
                if self.edge:
                  activationModule.add(VCmpEQU32(dst=VCC(), src0="BufferOOB", src1=(vgpr(addrCalc.addrDVgpr)), comment =""))
                  activationModule.add(VCndMaskB32(dst=vgpr("AmaxOutB"), src0=vgpr("ValuC+%d"%vgprIdx), src1=0, src2=VCC(), comment="Check If OOB, put zero if OOB"))
                  activationModule.add(VMaxF32(dst=vgpr("AmaxOut"), src0=vgpr("AmaxOut"), src1=vgpr("AmaxOutB", isAbs=True), comment="absmax"))
                else:
                  activationModule.add(VMaxF32(dst=vgpr("AmaxOut"), src0=vgpr("AmaxOut"), src1=vgpr("ValuC+%d"%vgprIdx, isAbs=True), comment="absmax"))
              activationModule.add(VMulF32(dst=vgpr("ValuC+%d"%vgprIdx), src0=vgpr("ValuC+%d"%vgprIdx), src1=sgpr("ScaleD"), comment="result *= ScaleD"))
            # Original packed route
            elif vi%2 == 1:
              assert (self.gwvw % 2 == 0)
            else:
              activationModule.add(VMulPKF32(dst=vgpr("ValuC+%d"%vgprIdx, 2), src0=vgpr("ValuC+%d"%vgprIdx, 2), src1=sgpr("ScaleD", 2), vop3=VOP3PModifiers(op_sel_hi=[1,0,1]), comment="result *= ScaleD"))
          else:
            assert 0, "Unsupported scaleD type"


      # pack stores, beta and non-beta reach here:
      packModule = Module("Empty pack module")
      convertModule = Module("Empty convert module")
      if self.needsAccumToDestConversion:
        if self.kernel["ActivationFuncCall"] and activationCDataType == self.kernel["ProblemType"]["DestDataType"]:
          destIdx = self.activationSetPCStruct.vgprActCopy
        else:
          destIdx = self.ss.elementSumIdx[elementIdx]
        packTmpS01 = self._epilogScratchSgpr(self.laneSGPRC)
        if self.kernel["ProblemType"]["DestDataType"].isHalf():
          # For UseSubtileImpl non-edge: paired dwordx4 path handles packing in _emit16bitSubtilePairedStore.
          if not is16bitSubtile:
            packModule = self.packdata(self.gwvw, destIdx, self.ss.elementSumIdx[elementIdx], inputPrefix="ValuC+", prefixOffset=self.parentWriter.states.c.startVgprValu)
        elif self.kernel["ProblemType"]["DestDataType"].isBFloat16():
          # For UseSubtileImpl non-edge: paired dwordx4 path handles packing in _emit16bitSubtilePairedStore.
          if not is16bitSubtile:
            packModule = self.packdata(self.gwvw, destIdx, self.ss.elementSumIdx[elementIdx], bf16CVTVgprStruct=self.cvtVgprStruct,
                                       tmpS01=packTmpS01, laneSGPRC=self.laneSGPRC, inputPrefix="ValuC+", prefixOffset=self.parentWriter.states.c.startVgprValu)
        elif self.kernel["ProblemType"]["DestDataType"].isAnyFloat8():
          if self.kernel["ProblemType"]["StochasticRounding"]:
            packModule = self.packdata(self.gwvw, destIdx, self.ss.elementSumIdx[elementIdx], fp8CVTVgprStruct=self.cvtVgprStruct, \
                                       tmpS01=packTmpS01, laneSGPRC=self.laneSGPRC, vgprTmp=vgprRND, inputPrefix="ValuC+", prefixOffset=self.parentWriter.states.c.startVgprValu, alphaScale=1.0)
          else:
            packModule = self.packdata(self.gwvw, destIdx, self.ss.elementSumIdx[elementIdx], fp8CVTVgprStruct=self.cvtVgprStruct, \
                                       tmpS01=packTmpS01, laneSGPRC=self.laneSGPRC, inputPrefix="ValuC+", prefixOffset=self.parentWriter.states.c.startVgprValu)
        elif self.kernel["ProblemType"]["DestDataType"].isAnyBFloat8():
          # TODO: BF8 stochastic rounding is not yet supported here.
          #       VCvtSRF32toBF8 instruction exists but stochasticRoundingCvt() only emits VCvtSRF32toFP8.
          #       To support BF8 SR: add SR branch here, generalize stochasticRoundingCvt() to accept bf8CVTVgprStruct,
          #       and select VCvtSRF32toBF8 based on DestDataType.
          packModule = self.packdata(self.gwvw, destIdx, self.ss.elementSumIdx[elementIdx], bf8CVTVgprStruct=self.cvtVgprStruct, \
                                     tmpS01=packTmpS01, laneSGPRC=self.laneSGPRC, inputPrefix="ValuC+", prefixOffset=self.parentWriter.states.c.startVgprValu)
        elif self.kernel["ProblemType"]["DestDataType"].isInt32():
          if self.kernel["ProblemType"]["ComputeDataType"].isSingle() and ((self.parentWriter.states.useBias == DataDirection.READ) or self.kernel["ActivationFuncCall"] or self.applyAlpha or self.beta):
            convertModule = convertData(self.gwvw, self.ss.elementSumIdx[elementIdx], cvtType=CvtType.CVT_F32_to_I32, \
                                        inputPrefix="ValuC+", prefixOffset=self.parentWriter.states.c.startVgprValu)
        elif self.kernel["ProblemType"]["DestDataType"].isInt8():
          if self.kernel["ProblemType"]["ComputeDataType"].isSingle() and ((self.parentWriter.states.useBias == DataDirection.READ) or self.kernel["ActivationFuncCall"] or self.applyAlpha or self.beta):
            convertModule = convertData(self.gwvw, self.ss.elementSumIdx[elementIdx], cvtType=CvtType.CVT_F32_to_I32, roundType=RoundType.ROUND_TO_NEAREST_EVEN, \
                                        inputPrefix="ValuC+", prefixOffset=self.parentWriter.states.c.startVgprValu)
          packModule = self.packdata(self.gwvw, destIdx, self.ss.elementSumIdx[elementIdx], self.cvtVgprStruct, self.tmpS01,
                                     SaturateTypeInt8=SaturateTypeInt8, inputPrefix="ValuC+", prefixOffset=self.parentWriter.states.c.startVgprValu)
        self._epilogScratchFree(packTmpS01)

      if self.parentWriter.states.asmCaps["HasWMMA_V1"] and self.kernel["EnableMatrixInstruction"] and self.kernel["ProblemType"]["DestDataType"].isHalf() and (not self.kernel["ProblemType"]["HighPrecisionAccumulate"]):
        for vi in range(0, self.gwvw):
          sumIdxV = self.ss.elementSumIdx[elementIdx] + vi
          if vi%2 == 1:
            formatVgpr = formatting(sumIdxV, "ValuC+", self.parentWriter.states.c.startVgprValu)
            d = self.ss.elementSumIdx[elementIdx] + vi//2
            dVgpr = formatting(d, "ValuC+", self.parentWriter.states.c.startVgprValu)
            packModule.add(VPackF16toB32(dst=vgpr(dVgpr), src0=vgpr(formatting(sumIdxV-1, "ValuC+", self.parentWriter.states.c.startVgprValu)), src1=vgpr(formatVgpr), \
                          comment="Pack with neighbor"))

      if self.kernel["ExpertSchedulingMode"] > 0:
        packModule.add(SWaitAlu(va_vdst=0, comment="wait for writes to complete"))

      biasReductionModule = Module("biasReductionModule")
      if self.storeBiasD == 1:
        vgprIdx = self.ss.elementSumIdx[elementIdx] - self.parentWriter.states.c.startVgprValu
        biasReductionModule.add(self.parentWriter.addStore(self.kernel, self.ss, 'Bias', addrCalc, "ValuC+%d"%vgprIdx, self.tmpS01, self.edge, comment="store Bias"))

      if isActivationInsertAfter:
        if self.parentWriter.states.useGateResidual and \
           (self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1):
          raise RuntimeError(
              "GateResidual is not supported with isActivationInsertAfter "
              "(activation applied after pack on dest-dtype value).")
        module.add(convertModule)
        module.add(packModule)
        module.add(gradientCvtModule)
        module.add(activationModule)
      else:
        module.add(gradientCvtModule)
        module.add(activationModule)
        module.add(scaleDModule)
        module.add(biasReductionModule)
        module.add(gateModule)
        module.add(convertModule)
        module.add(packModule)

      if not self.kernel["StoreRemapVectorWidth"]:
        # 16bit UseSubtileImpl non-edge: emit paired dwordx4 stores combining sba=0
        # with sba=1 subtile data into one buffer_store_dwordx4.  Works for both
        # bf16 and fp16 HPA output types.
        #
        # UseSubtileImpl splits MIWaveTile[0] into two subtile groups:
        #   sba=0 owns even tt0 values (0, 2, 4, ...)
        #   sba=1 owns odd  tt0 values (1, 3, 5, ...)
        # The element list interleaves them as consecutive (even, odd) tt0 pairs:
        #   element 0: tt0=0 (sba=0)
        #   element 1: tt0=1 (sba=1)   <- pair with element 0
        #   element 2: tt0=2 (sba=0)   (if MIWaveTile[0]>2)
        #   ...
        # Pairing key: tt0 % 2 — even tt0 is sba=0, odd tt0 is sba=1.
        storeCodeModule = storeCode if self.kernel["GroupLoadStore"] else module
        if is16bitSubtile:
          tt0 = element[1]  # d0: thread-tile index along M
          # Epilogue (bias/activation) is applied per-element in iteration order.
          # The paired store must be emitted AFTER both sba=0 and sba=1 elements have
          # had their epilogue applied, so we defer it to the sba=1 (odd tt0) iteration.
          if tt0 % 2 == 1:
            # sba=1 element (odd tt0): both sba=0 and sba=1 epilogues are done — emit paired store.
            # Find the sba=0 partner: the immediately preceding element with tt0-1.
            partnerElementIdx = elementIdx - 1
            partnerExists = (partnerElementIdx >= 0 and
                             self.batchElements[partnerElementIdx][1] == tt0 - 1)
            if partnerExists:
              # Paired dwordx4 store for (sba=0 at tt0-1, sba=1 at tt0).
              partnerAddrCalc: AddrCalculation = self.ss.elementAddr[partnerElementIdx]
              sumIdx0 = self.ss.elementSumIdx[partnerElementIdx]
              sumIdx1 = self.ss.elementSumIdx[elementIdx]
              prefixOffset = self.parentWriter.states.c.startVgprValu
              blockIdxN = element[0]
              # Guard with tt0-1 (lower block): skip if even the lower M-block is OOB.
              # This also handles N-group transitions.
              blockIdxM = tt0 - 1
              skipLabel = self._emitSubtileOobGuard(storeCodeModule, blockIdxM, blockIdxN,
                                                    labelPrefix="subtile_skip_store")
              # Additional check: paired store needs BOTH blocks valid (MGuard > tt0).
              # When only the lower block is valid, fall through to a scalar fallback.
              guardMSgpr = self.parentWriter.states.subtileM32ValidBlocksSgpr
              if guardMSgpr is not None:
                afterPairedLabel = Label(self.parentWriter.labels.getNameInc("subtile_after_paired"),
                                        f"after paired/fallback store tt0={tt0}")
                fallbackLabelName = self.parentWriter.labels.getNameInc("subtile_scalar_fallback")
                fallbackLabel = Label(fallbackLabelName,
                                      f"scalar fallback for d0={tt0-1} when d0={tt0} is OOB")
                storeCodeModule.add(_scmpGtU32(self.parentWriter, sgpr("SubtileMGuard"), tt0,
                                               comment=f"paired store: both M-blocks valid? (MGuard > {tt0})"))
                storeCodeModule.add(SCBranchSCC0(labelName=fallbackLabel.getLabelName(),
                                                 comment=f"only d0={tt0-1} valid -> scalar fallback"))
                tmpStoreCode = self._emit16bitSubtilePairedStore(partnerAddrCalc, sumIdx0, sumIdx1, prefixOffset, tt0 - 1, blockIdxM=blockIdxM, blockIdxN=blockIdxN)
                storeCodeModule.add(tmpStoreCode)
                storeCodeModule.add(SBranch(labelName=afterPairedLabel.getLabelName(),
                                            comment="skip scalar fallback"))
                storeCodeModule.add(fallbackLabel)
                tmpFallbackCode = self._emit16bitSubtileScalarStore(partnerAddrCalc, sumIdx0, prefixOffset, tt0 - 1, blockIdxM=blockIdxM, blockIdxN=blockIdxN)
                storeCodeModule.add(tmpFallbackCode)
                storeCodeModule.add(afterPairedLabel)
              else:
                tmpStoreCode = self._emit16bitSubtilePairedStore(partnerAddrCalc, sumIdx0, sumIdx1, prefixOffset, tt0 - 1, blockIdxM=blockIdxM, blockIdxN=blockIdxN)
                storeCodeModule.add(tmpStoreCode)
              if skipLabel is not None:
                storeCodeModule.add(skipLabel)
              self.storesIssued += 1
            else:
              # sba=1 orphan (no sba=0 partner in this batch — split by batch boundary).
              blockIdxM = tt0
              blockIdxN = element[0]
              orphanSkipLabel = self._emitSubtileOobGuard(storeCodeModule, blockIdxM, blockIdxN,
                                                          labelPrefix="subtile_skip_orphan")
              sumIdx0 = self.ss.elementSumIdx[elementIdx]
              prefixOffset = self.parentWriter.states.c.startVgprValu
              tmpStoreCode = self._emit16bitSubtileScalarStore(addrCalc, sumIdx0, prefixOffset, tt0, blockIdxM=blockIdxM, blockIdxN=blockIdxN)
              storeCodeModule.add(tmpStoreCode)
              if orphanSkipLabel is not None:
                storeCodeModule.add(orphanSkipLabel)
              self.storesIssued += 1
          else:
            # sba=0 element (even tt0): defer SRD row increment until after N-group label.
            if self.ss.optSrdIncForRow and addrCalc.rowInc:
              self._subtilePendingSrdDInc = addrCalc.incrementToNextRow(self.kernel, "D", self.ss, self.tmpS01)
            partnerElementIdx = elementIdx + 1
            partnerExists = (partnerElementIdx < len(self.batchElements) and
                             self.batchElements[partnerElementIdx][1] == tt0 + 1)
            if not partnerExists:
              # Orphan element (no sba=1 partner in this batch): scalar 16bit store now.
              # Guard against OOB wave groups (same as paired store path).
              blockIdxM = tt0
              blockIdxN = element[0]
              # Early exit: skip this orphan scalar store if the wave group is outside the valid M/N tile bounds.
              orphanSkipLabel = self._emitSubtileOobGuard(storeCodeModule, blockIdxM, blockIdxN,
                                                          labelPrefix="subtile_skip_orphan")
              sumIdx0 = self.ss.elementSumIdx[elementIdx]
              prefixOffset = self.parentWriter.states.c.startVgprValu
              tmpStoreCode = self._emit16bitSubtileScalarStore(addrCalc, sumIdx0, prefixOffset, tt0, blockIdxM=blockIdxM, blockIdxN=blockIdxN)
              storeCodeModule.add(tmpStoreCode)
              if orphanSkipLabel is not None:
                storeCodeModule.add(orphanSkipLabel)
              self.storesIssued += 1
        elif self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":#GSUGSU
          tmpStoreCode = self.parentWriter.addStore(self.kernel, self.ss, 'TD', addrCalc, sumIdx, self.tmpS01, self.edge, comment="store TD not StoreRemapVectorWidth")
          storeCodeModule.add(tmpStoreCode)
          self.storesIssued += 1
        else:
          # Regular store path. If UseSubtileImpl NonEdge, guard against OOB wave groups.
          skipLabel = None
          if isSubtileNonEdge:
            tt0 = element[1]
            blockIdxM = tt0  # each tt0 maps to one mBlockSize-row block
            blockIdxN = element[0]
            # Early exit: skip this store if the wave group is outside the valid M/N tile bounds.
            skipLabel = self._emitSubtileOobGuard(storeCodeModule, blockIdxM, blockIdxN,
                                                  labelPrefix="subtile_skip_store")
          # Apply exec mask for partial M/N blocks (regular fp32 store path)
          if self.parentWriter.states.storeAlign8 and isSubtileNonEdge:
            self._emitAlign8ExecMask(storeCodeModule, self.tmpS01, self.tmpS23, blockIdxM, blockIdxN,
                                     mGuardOffset=1, rowScaleShift=2)
            storeCodeModule.add(self.getEdgeMovInstType()(EXEC(), sgpr(self.tmpS01, self.laneSGPRC), "apply exec mask"))
          # _emitOverrideRows reused from the top of this store loop (see _lookaheadRowInc).
          tmpStoreCode = self.parentWriter.addStore(self.kernel, self.ss, 'D', addrCalc, sumIdx, self.tmpS01, self.edge, elementIdx, self.batchIdx,
                                                   overrideAfterPrimerRows=_emitOverrideRows, comment="store D")
          storeCodeModule.add(tmpStoreCode)
          if self.parentWriter.states.storeAlign8 and isSubtileNonEdge:
            storeCodeModule.add(self.getEdgeMovInstType()(EXEC(), -1, "restore exec"))
          if skipLabel is not None:
            storeCodeModule.add(skipLabel)
          self.storesIssued += 1

        if (self.kernel["ProblemType"]["UseE"] and not self.kernel["ProblemType"]["Gradient"]) and ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["StreamK"] > 0):
          self.storesIssued += 1
        if self.storeBiasD == 1:
          self.storesIssued += 1

      else:
        if not self.kernel["_GlobalAccumulation"] == "MultipleBufferSingleKernel":#GSUGSU
          rpe = self.parentWriter.states.bpeCinternal // self.parentWriter.states.bpr
          module.add(self.parentWriter.storeRemapAddLocalWrite(self.kernel, self.ss, addrCalc, sumIdx*rpe))
          # Column Block Shape has been written to LDS
          # Now read back and write out to global memory
        else:
          tmpStoreCode = self.parentWriter.addStore(self.kernel, self.ss, 'TD', addrCalc, sumIdx, self.tmpS01, self.edge, comment="store TD StoreRemapVectorWidth")

          if self.kernel["GroupLoadStore"]:
            storeCode.add(tmpStoreCode)

          module.add(tmpStoreCode)

          self.storesIssued += 1
          if (self.kernel["ProblemType"]["UseE"] and not self.kernel["ProblemType"]["Gradient"]) and ((self.kernel["GlobalSplitU"] == 1 or self.kernel["GlobalSplitU"] == -1) or self.kernel["StreamK"] > 0):
            self.storesIssued += 1
          if self.storeBiasD == 1:
            self.storesIssued += 1

    # Close the last N-group OOB skip label (if any) opened by _emitSubtileOobGuard.
    self._finalizeSubtileOobGuards(storeCode if self.kernel["GroupLoadStore"] else module)
    if self.kernel["ProblemType"]["StochasticRounding"]:
      self.parentWriter.vgprPool.checkIn(vgprRND)

    module.add(storeCode)

    if self.parentWriter.db["CheckStoreC"]>=0:
      useBuffer = self.kernel["BufferStore"]
      # Note - CheckStoreC won't work for EDGE store cases since they load 0 for OOB, would need more sophisticated check
      # Note - TODO- CheckStoreC also won't work for StoreRemap
      module.add(SWaitCnt(vscnt=0, comment="CheckStoreC, wait for stores to complete"))
      for elementIdx in range(0, len(self.batchElements)):
        addr = self.ss.elementAddr[elementIdx].addrDVgpr
        sumIdx = self.ss.elementSumIdx[elementIdx]

        bps = self.kernel["ProblemType"]["DestDataType"].numBytes() * self.gwvw
        if self.kernel["BufferStore"]:
          addr0 = vgpr(addr)
          addr1 = sgpr("SrdC", 4)
        else:
          addr0 = vgpr(addr,2)
          addr1 = ""

        if self.kernel["ProblemType"]["DestDataType"].isHalf() or self.kernel["ProblemType"]["DestDataType"].isBFloat16():
          if not self.kernel["ProblemType"]["HighPrecisionAccumulate"]:
            module.add(self.parentWriter.chooseGlobalRead(useBuffer, bps, sumIdx//2, \
                                  addr0, addr1, soffset=0, offset=0, hi16=sumIdx%2))
          else:
            module.add(self.parentWriter.chooseGlobalRead(useBuffer, bps, sumIdx, \
                                  addr0, addr1, soffset=0, offset=0, hi16=0))
        elif self.kernel["ProblemType"]["DestDataType"].isInt32() or self.kernel["ProblemType"]["DestDataType"].isSingle():
          module.add(self.parentWriter.chooseGlobalRead(useBuffer, bps, sumIdx, \
                                addr0, addr1, soffset=0, offset=0))
        elif self.kernel["ProblemType"]["DestDataType"].isDouble() or self.kernel["ProblemType"]["DestDataType"].isSingleComplex() :
          module.add(self.parentWriter.chooseGlobalRead(useBuffer, bps, sumIdx*2, \
                                addr0, addr1, soffset=0, offset=0))
        elif self.kernel["ProblemType"]["DestDataType"].isDoubleComplex():
          module.add(self.parentWriter.chooseGlobalRead(useBuffer, bps, sumIdx*4, \
                                addr0, addr1, soffset=0, offset=0))
      module.add(SWaitCnt(vscnt=0, comment="CheckStoreC, wait for stores to complete"))
      # Add checks for expected values:
      module.add(SMovB32(sgpr(self.tmpS01), self.parentWriter.db["CheckStoreC"], "expected value"))
      for elementIdx in range(0, len(self.batchElements)):
        sumIdx = self.ss.elementSumIdx[elementIdx]
        # Need to fix for other types:
        assert (self.kernel["ProblemType"]["DestDataType"].isSingle() or self.kernel["ProblemType"]["DestDataType"].isInt32())
        module.add(self.parentWriter.getCmpAssert(self.parentWriter.asmAssert.eq, vgpr(sumIdx), sgpr(self.tmpS01)))


    if self.edge and (self.atomic or not self.kernel["BufferStore"]):
      # subsequent batch must start with full exec mask
      # BufferStore doesn't need exec since it used buffer range checking when
      # possible
      module.add(self.getEdgeMovInstType()(EXEC(), -1, "full mask -> exec"))

    if self.parentWriter.db["ConservativeWaitCnt"] & 0x40:
      module.add(SBarrier(comment="debug"))
      module.add(SWaitCnt(vscnt=0, comment="ConservativeWaitCnt"))
      module.add(SBarrier(comment="debug"))

  def _emitSubtilePackedPermute(self, vPack: int, vPermAddr: int, addrWhilePermuting=None) -> Module:
    """Shuffle four packed dwords across wave halves for a subtile dwordx4 store.

    After the caller packs 8 f32 accumulator values into four 16bit dwords
    (vPack+0..+3), this routine performs the two-step permute that assembles
    eight consecutive M-rows owned by a pair of lane-groups into a contiguous
    dwordx4 payload:

      Step 1 — ds_bpermute (in-place, 4×): each lane fetches vPack+k from its
               partner lane l' (= the lane at LG±1 distance, pre-encoded as a
               byte address in vPermAddr).  The LDS pipe latches vPermAddr at
               issue time, so vPermAddr can be repurposed as soon as all four
               ds_bpermute instructions are issued.

      Step 2 — v_permlane32_swap_b32 (2×): exchange (vPack+0 ↔ vPack+2) and
               (vPack+1 ↔ vPack+3) across the 32-lane boundary so that lanes
               0-31 end up with rows LG*8+0..LG*8+7 in ascending order.

    The caller may supply an optional `addrWhilePermuting` callable that adds
    address-preparation instructions to the same module *between* the four
    ds_bpermute issues and the SWaitCnt.  This overlaps address arithmetic
    with the LDS round-trip latency at no extra cost.

    Args:
      vPack:              Base VGPR index of the four packed dwords (must be
                          2-aligned to satisfy dwordx4 store alignment).
      vPermAddr:          VGPR holding the partner-lane byte address (pre-computed
                          once per batch in the vgprPermAddr slot).
      addrWhilePermuting: Optional callable() that appends address instructions
                          to `module` while the ds_bpermute results are in-flight.

    Returns:
      Module containing ds_bpermute × 4, optional address code, SWaitCnt,
      and v_permlane32_swap_b32 × 2.  Leaves vPack+0..+3 holding the
      correctly ordered dwords ready for buffer_store_dwordx4.
    """
    module = Module("SubtilePackedPermute")

    module.addComment1("ds_bpermute in-place: gather packed dwords from partner lane-group")
    for k in range(4):
      module.add(DSBPermuteB32(dst=vgpr(vPack+k), src0=vgpr(vPermAddr), src1=vgpr(vPack+k),
                               comment=f"perm dword {k}"))

    if addrWhilePermuting is not None:
      addrWhilePermuting(module)

    module.add(SWaitCnt(dscnt=0, comment="wait for ds_bpermute (lgkmcnt=0)"))

    module.addComment1("v_permlane32_swap_b32: swap across lane-32 boundary")
    module.add(VPermlane32SwapB32(dst=vgpr(vPack+0), src=vgpr(vPack+2), comment="swap dwords 0↔2"))
    module.add(VPermlane32SwapB32(dst=vgpr(vPack+1), src=vgpr(vPack+3), comment="swap dwords 1↔3"))

    return module

  def _emitSubtileOobGuard(self, targetModule, blockIdxM: int, blockIdxN: int, labelPrefix: str = "subtile_skip_store"):
    """Emit M/N OOB guard branches for UseSubtileImpl NonEdge stores.

    Background
    ----------
    UseSubtileImpl assigns each wave a fixed subtile region of the output matrix.
    In the NonEdge path the macro-tile fits entirely within the output bounds, but
    individual wave groups within the macro-tile may still be out-of-bounds when
    the problem size is not a multiple of the macro-tile.  The SGPRs
    subtileM32ValidBlocksSgpr and subtileN16ValidBlocksSgpr count how many M/N
    blocks (in units of mBlockSize rows / nBlockSize columns) belong to valid
    output for this wave, and are set to None when no guard is needed (edge path
    or problem is tile-aligned).

    Logic
    -----
    For a store at (blockIdxM, blockIdxN):
      - If numValidNBlocks <= blockIdxN → N OOB: jump past ALL remaining stores.
        Valid because N is monotone: subsequent N groups (blockIdxN+1, ...) are also OOB.
      - If numValidMBlocks <= blockIdxM → M OOB: jump to the end of the current N group.
        Valid because M is monotone: remaining M elements in this N group are also OOB.

    The N guard is emitted ONCE per N group (when blockIdxN changes).  It branches to
    _subtileAllStoresEndLabel (past all stores).  _subtileNGroupSkipLabel marks the
    boundary between N groups; M guards branch there to skip the rest of the current
    N group without re-testing the remaining M elements.
    Both labels are placed by _finalizeSubtileOobGuards (called after the element loop).

    Returns a per-element skip Label only when there is no N guard (M-only case); the
    caller must add it after the store.  Returns None in all other cases.
    """
    guardMSgpr = self.parentWriter.states.subtileM32ValidBlocksSgpr
    guardNSgpr = self.parentWriter.states.subtileN16ValidBlocksSgpr
    # No guard SGPRs means the store is always in-bounds for this path; nothing to emit.
    if guardMSgpr is None and guardNSgpr is None:
      return None

    # --- N-group guard (emitted once per unique blockIdxN) ---
    # Branches to _subtileAllStoresEndLabel when N OOB, skipping all remaining stores.
    # Because N is monotone (blockIdxN increases each group), if this group is OOB
    # then every subsequent group is also OOB — no need to test them.
    if guardNSgpr is not None and blockIdxN != self._subtilePrevBlockIdxN:
      # Place the previous N group's end label before starting a new group.
      if self._subtileNGroupSkipLabel is not None:
        targetModule.add(self._subtileNGroupSkipLabel)
        self._subtileNGroupSkipLabel = None
      if self._subtilePendingSrdDInc is not None:
        targetModule.add(self._subtilePendingSrdDInc)
        self._subtilePendingSrdDInc = None
      # Create the single end-of-all-stores label on the first N group.
      if self._subtileAllStoresEndLabel is None:
        endLabelName = self.parentWriter.labels.getNameInc("subtile_all_stores_end")
        self._subtileAllStoresEndLabel = Label(endLabelName, "end of all subtile NonEdge D stores")
      nGroupEndLabelName = self.parentWriter.labels.getNameInc(
        f"{labelPrefix}_N{blockIdxN}_end")
      nGroupEndLabel = Label(nGroupEndLabelName,
                             f"end of N group blockIdxN={blockIdxN} (M cbranch target)")
      nGuardCmp = blockIdxN * 16 if self.parentWriter.states.storeAlign8 else blockIdxN
      targetModule.add(_scmpGtU32(self.parentWriter, sgpr("SubtileNGuard"), nGuardCmp,
                                   comment=f"quick-exit: clamped > {nGuardCmp}? (OOB -> skip all stores)"))
      targetModule.add(SCBranchSCC0(labelName=self._subtileAllStoresEndLabel.getLabelName(),
                                     comment=f"quick-exit: N OOB at blockIdxN={blockIdxN}, skip all remaining stores"))
      self._subtileNGroupSkipLabel = nGroupEndLabel
      self._subtilePrevBlockIdxN = blockIdxN

    # --- M guard (emitted per element) ---
    # Branches to end of current N group when M OOB, skipping remaining M elements.
    # Because M is monotone (blockIdxM increases within the N group), if this element
    # is OOB then all subsequent M elements in this N group are also OOB.
    if guardMSgpr is None:
      return None
    targetModule.add(_scmpGtU32(self.parentWriter, sgpr("SubtileMGuard"), blockIdxM,
                                 comment=f"quick-exit: numValidMBlocks > {blockIdxM}? (OOB -> skip N group)"))
    if guardNSgpr is not None and self._subtileNGroupSkipLabel is not None:
      # M OOB → jump to end of this N group (no per-element label needed).
      targetModule.add(SCBranchSCC0(labelName=self._subtileNGroupSkipLabel.getLabelName(),
                                     comment=f"quick-exit: M OOB at blockIdxM={blockIdxM}, skip rest of N group"))
      return None
    else:
      # No N guard → fall back to a per-element skip label (caller places it after the store).
      skipLabelName = self.parentWriter.labels.getNameInc(
        f"{labelPrefix}_M{blockIdxM}_N{blockIdxN}")
      skipLabel = Label(skipLabelName,
                        f"skip OOB store blockIdxM={blockIdxM} blockIdxN={blockIdxN}")
      targetModule.add(SCBranchSCC0(labelName=skipLabel.getLabelName(),
                                     comment=f"quick-exit: M OOB at blockIdxM={blockIdxM}, skip store"))
      return skipLabel

  def _finalizeSubtileOobGuards(self, targetModule):
    """Place the pending N-group end label and end-of-all-stores label after the element loop.

    Must be called once after all elements have been emitted to close out the last
    N group and anchor the N-cbranch target past all stores.
    """
    if self._subtileNGroupSkipLabel is not None:
      targetModule.add(self._subtileNGroupSkipLabel)
      self._subtileNGroupSkipLabel = None
    if self._subtilePendingSrdDInc is not None:
      targetModule.add(self._subtilePendingSrdDInc)
      self._subtilePendingSrdDInc = None
    if self._subtileAllStoresEndLabel is not None:
      targetModule.add(self._subtileAllStoresEndLabel)
      self._subtileAllStoresEndLabel = None

  def _emitAlign8ExecMask(self, module, tmpS, tmpS2, blockIdxM, blockIdxN, mGuardOffset, rowScaleShift):
    """Emit exec mask for partial M/N blocks in the NonEdge store path.

    Computes a 64-bit exec mask that disables OOB lanes for partial M and N
    blocks at the tile boundary.  The mask is the AND of an M component and
    an N component, each computed independently.

    MMA output layout (MI16x16, wavefront=64):
      64 lanes = 4 lane-groups (LGs) of 16 lanes each.
      Each LG owns consecutive M-rows:
        - Paired bf16 store: 8 M-rows per LG (32 rows total, 2 MMA tiles)
        - Scalar/fp32 store: 4 M-rows per LG (16 rows total, 1 MMA tile)
      Within each LG, lane_id % 16 selects the N-column (0..15).

    M mask algorithm:
      Given validRows (number of valid M-rows in this block), the number of
      active lane-groups is validRows / rowsPerLG.  Each LG occupies 16
      consecutive lanes, so the mask is the bottom (numValidLGs * 16) bits of
      a 64-bit word.  This is computed as:
        shiftAmt = 64 - validRows * (16 / rowsPerLG)
        mask = (uint64_t)-1 >> shiftAmt
      validM_wave is precomputed once per tile in _emitSubtileGuards.

    N mask algorithm:
      SubtileNGuard holds the clamped valid-N-column count for this wave.
      partialN = SubtileNGuard % 16 gives the number of valid columns within
      the last 16-column MMA tile.  The N mask disables columns >= partialN:
        nMask = (1 << partialN) - 1        e.g. partialN=5 -> 0x001F
        nMask = nMask | (nMask << 16)      replicate across both LG halves
      This 32-bit word is ANDed into both lo and hi halves of the exec mask.

    This sequence is designed to fit within the ds_bpermute latency window
    (~88 cycles at 4 cyc/issue) for zero-overhead execution on interior blocks.

    Args:
        module:         Target module to emit instructions into.
        tmpS:           SGPR index for mask result (2 consecutive, even-aligned).
        tmpS2:          SGPR index for scratch (2 consecutive).
        blockIdxM:      M block index for this store element.
        blockIdxN:      N block index for this store element.
        mGuardOffset:   Number of MMA tiles this store spans (2=paired, 1=single).
        rowScaleShift:  Left-shift to convert validRows to lane-count
                        (1 for paired/8-rows-per-LG, 2 for scalar-fp32/4-rows-per-LG).
    """
    miM = self.kernel["MatrixInstM"]
    blockStartRow = blockIdxM * miM
    validMWaveSgpr = self.parentWriter.states.subtileTotalMOffsetSgpr
    mMaskDone = Label(self.parentWriter.labels.getNameInc("align8_m_done"), "")

    # --- M mask: per-lane-group selection via right-shift ---
    isWave32 = self.wavelen == 32
    if isWave32:
      module.add(SMovB32(dst=sgpr(tmpS), src=-1, comment="mask = full"))
    else:
      module.add(SMovB64(dst=sgpr(tmpS, 2), src=-1, comment="mask = full"))
    module.add(_scmpGtU32(self.parentWriter, sgpr("SubtileMGuard"), blockIdxM + mGuardOffset,
                          comment=f"SubtileMGuard > {blockIdxM + mGuardOffset}? (block fully interior)"))
    module.add(SCBranchSCC1(labelName=mMaskDone.getLabelName(),
                            comment="interior M block -> mask stays -1"))
    if blockStartRow > 0:
      module.add(SSubU32(dst=sgpr(tmpS2), src0=sgpr(validMWaveSgpr), src1=blockStartRow,
                         comment=f"validRows = validM_wave - {blockStartRow}"))
      module.add(SCSelectB32(dst=sgpr(tmpS2), src0=0, src1=sgpr(tmpS2), comment="clamp to 0"))
      module.add(SLShiftLeftB32(dst=sgpr(tmpS2), src=sgpr(tmpS2), shiftHex=rowScaleShift,
                                comment=f"validRows * {1 << rowScaleShift}"))
    else:
      module.add(SLShiftLeftB32(dst=sgpr(tmpS2), src=sgpr(validMWaveSgpr), shiftHex=rowScaleShift,
                                comment=f"validM_wave * {1 << rowScaleShift}"))
    module.add(SSubU32(dst=sgpr(tmpS2), src0=self.wavelen, src1=sgpr(tmpS2),
                       comment=f"shiftAmt = {self.wavelen} - validRows * scale"))
    if isWave32:
      module.add(SLShiftRightB32(dst=sgpr(tmpS), src=-1, shiftHex=sgpr(tmpS2),
                                 comment="M mask = -1 >> shiftAmt"))
    else:
      module.add(SLShiftRightB64(dst=sgpr(tmpS, 2), src=-1, shiftHex=sgpr(tmpS2),
                                 comment="M mask = -1 >> shiftAmt"))
    module.add(mMaskDone)

    # --- N mask: per-column bit-mask for arbitrary partial N ---
    nCmpVal = (blockIdxN + 1) * 16
    nMaskDone = Label(self.parentWriter.labels.getNameInc("align8_n_done"), "")
    module.add(_scmpGtU32(self.parentWriter, sgpr("SubtileNGuard"), nCmpVal,
                          comment=f"clamped > {nCmpVal}? (not last N block)"))
    module.add(SCBranchSCC1(labelName=nMaskDone.getLabelName(),
                            comment="interior N block -> mask unchanged"))
    module.add(SAndB32(dst=sgpr(tmpS2), src0=sgpr("SubtileNGuard"), src1=0xF,
                       comment="partialN = clamped %% 16, SCC=1 if non-zero"))
    nFullLabel = Label(self.parentWriter.labels.getNameInc("align8_n_full"), "")
    module.add(SCBranchSCC0(labelName=nFullLabel.getLabelName(),
                            comment="partialN==0 -> full 16-col block"))
    module.add(SLShiftLeftB32(dst=sgpr(tmpS2), src=1, shiftHex=sgpr(tmpS2),
                              comment="1 << partialN"))
    module.add(SSubU32(dst=sgpr(tmpS2), src0=sgpr(tmpS2), src1=1,
                       comment="(1 << partialN) - 1"))
    if isWave32:
      # wave32: replicate lo16 to both halves without extra scratch SGPR
      module.add(SMulI32(dst=sgpr(tmpS2), src0=sgpr(tmpS2), src1=hex(0x10001),
                         comment="replicate lo16 to both halves"))
      module.add(SAndB32(dst=sgpr(tmpS), src0=sgpr(tmpS), src1=sgpr(tmpS2),
                         comment="mask &= N mask"))
    else:
      module.add(SLShiftLeftB32(dst=sgpr(tmpS2+1), src=sgpr(tmpS2), shiftHex=16,
                                comment="replicate to hi16"))
      module.add(SOrB32(dst=sgpr(tmpS2), src0=sgpr(tmpS2), src1=sgpr(tmpS2+1),
                        comment="N mask word = lo16 | hi16"))
      module.add(SAndB32(dst=sgpr(tmpS), src0=sgpr(tmpS), src1=sgpr(tmpS2),
                         comment="mask_lo &= N mask"))
      module.add(SAndB32(dst=sgpr(tmpS+1), src0=sgpr(tmpS+1), src1=sgpr(tmpS2),
                         comment="mask_hi &= N mask"))
    module.add(nFullLabel)
    module.add(nMaskDone)

  def _emit16bitSubtilePairedStore(self, addrCalc, sumIdx0: int, sumIdx1: int, prefixOffset: int, tt0: int = 0, blockIdxM: int = 0, blockIdxN: int = 0) -> Module:
    """Emit a paired 16bit store combining sba=0 and sba=1 subtile data.

    Works for both bf16 and fp16 HPA output types.

    sba = subtile block index along A (M dimension).  UseSubtileImpl iterates over
    two subtile groups (sba=0, sba=1) that share the same (tt1, tt0) element
    coordinates but draw from different accumulator registers.  The element list
    therefore contains consecutive pairs with identical (tt1, tt0): sba=0 first
    (even elementIdx), sba=1 second (odd elementIdx).

    Converts 8 f32 accvgprs (4 from sba=0, 4 from sba=1) to 16bit, shuffles them
    across wave halves via ds_bpermute + v_permlane32_swap_b32, then issues
    1 × buffer_store_dwordx4 at the sba=0 element's address.  The cvtVgpr block
    is 2-aligned (64-bit) in KWA so vgprBf16Temp satisfies the dwordx4 alignment.

    Args:
      addrCalc:     AddrCalculation for the sba=0 element.
      sumIdx0:      elementSumIdx for the sba=0 element.
      sumIdx1:      elementSumIdx for the sba=1 element.
      prefixOffset: parentWriter.states.c.startVgprValu (offset into ValuC).
      tt0:          thread-tile M index (same for both sba=0 and sba=1).
    """
    module = Module("16bitSubtilePairedStore")
    isFp16 = self.kernel["ProblemType"]["DestDataType"].isHalf()

    ntd = self.kernel["NonTemporalD"]
    isGlc = bool(ntd & 0x1)
    isSlc = bool(ntd & 0x2)
    isNT  = bool(ntd & 0x4)

    # Reuse cvtVgprStruct.vgprBf16Temp..vgprBf16Inc (+0..+3) as 4 scratch vgprs.
    # The cvtVgpr block is allocated with 2-alignment (64-bit aligned) in KWA so that
    # vgprBf16Temp is at an even VGPR index, satisfying buffer_store_dwordx4's
    # alignment requirement.  The +0..+3 slots are safely overwritten here as pack/perm
    # staging for each pair.
    vPack = self.cvtVgprStruct.vgprBf16Temp  # +0..3: packed 16bit dwords, 2-aligned

    vPermAddr    = self.cvtVgprStruct.vgprPermAddr
    vLGDelta     = self.cvtVgprStruct.vgprLaneGroupDelta
    vAddrScratch = self.cvtVgprStruct.vgprAddrScratch
    addrDVgpr    = addrCalc.addrDVgpr

    typeStr = "fp16" if isFp16 else "bf16"
    VCvtPkF32to16 = VCvtPkF32toFP16 if isFp16 else VCvtPkF32toBF16
    module.addComment1(f"{typeStr} paired dwordx4 store tt0={tt0} (sba=0+sba=1): pack 8 f32 accvgprs -> 4 {typeStr} dwords")

    # Pack sba=0 subtile: ValuC+sumIdx0+{0,1} → vPack+0; ValuC+sumIdx0+{2,3} → vPack+1
    # Pack sba=1 subtile: ValuC+sumIdx1+{0,1} → vPack+2; ValuC+sumIdx1+{2,3} → vPack+3
    def vc(sumIdx, vi):
      idx = sumIdx + vi - prefixOffset
      return vgpr("ValuC+" + str(idx))

    def packF32pair(dst, src0, src1, comment):
      """Pack two f32 VGPRs into one dword of two 16bit values."""
      module.add(VCvtPkF32to16(dst=vgpr(dst), src0=src0, src1=src1, comment=f"{comment} -> {typeStr}"))

    packF32pair(vPack+0, vc(sumIdx0, 0), vc(sumIdx0, 1), f"sba=0 tt0={tt0}[0:1]")
    packF32pair(vPack+1, vc(sumIdx0, 2), vc(sumIdx0, 3), f"sba=0 tt0={tt0}[2:3]")
    packF32pair(vPack+2, vc(sumIdx1, 0), vc(sumIdx1, 1), f"sba=1 tt0={tt0}[0:1]")
    packF32pair(vPack+3, vc(sumIdx1, 2), vc(sumIdx1, 3), f"sba=1 tt0={tt0}[2:3]")

    # Compute adjusted D address into vgprAddrScratch while ds_bpermute results are in-flight.
    # addrDVgpr holds the M-byte offset in bpeCexternal units; scale to bpeCexternalGSU1
    # (16bit=2 bytes) then add lane_group*8 so the dwordx4 store lands at the correct row.
    # addrDVgpr and vgprPermAddr are left unchanged — vgprAddrScratch is dedicated scratch
    # for this purpose so no restore is needed.
    bpeCurr = self.parentWriter.states.bpeCexternal
    bpeDest = self.parentWriter.states.bpeCexternalGSU1
    globalOffset = addrCalc.globalOffset * bpeDest // bpeCurr
    addrScaleShift = int(log2(bpeCurr // bpeDest)) if bpeCurr > bpeDest else 0

    useAlign8 = self.parentWriter.states.storeAlign8

    def emitAddrWhilePermuting(module):
      """Callback emitted between ds_bpermute and s_waitcnt lgkmcnt(0).

      The ds_bpermute has ~88 cycles of LDS latency.  We overlap SALU/VALU
      work in this window to hide the cost:
        1. Compute the adjusted D store address (VALU).
        2. Compute the exec mask for partial M/N blocks (SALU) — this
           suppresses OOB lanes at tile boundaries without adding any
           latency to the critical path.
      `module` is the permute Module, so instructions emitted here land
      between the ds_bpermute issue and the s_waitcnt that consumes results."""
      if addrScaleShift:
        module.add(VLShiftRightB32(dst=vgpr(vAddrScratch), shiftHex=addrScaleShift,
                                   src=vgpr(addrDVgpr), comment=f"scale addrDVgpr bpe {bpeCurr}->{bpeDest}"))
        module.add(VAddU32(dst=vgpr(vAddrScratch), src0=vgpr(vAddrScratch), src1=vgpr(vLGDelta),
                           comment="adjusted D addr = scaled addrDVgpr + lane_group*8"))
      else:
        module.add(VAddU32(dst=vgpr(vAddrScratch), src0=vgpr(addrDVgpr), src1=vgpr(vLGDelta),
                           comment="adjusted D addr = addrDVgpr + lane_group*8"))
      if useAlign8:
        self._emitAlign8ExecMask(module, self.tmpS01, self.tmpS23, blockIdxM, blockIdxN,
                                 mGuardOffset=2, rowScaleShift=1)

    module.add(self._emitSubtilePackedPermute(vPack, vPermAddr, addrWhilePermuting=emitAddrWhilePermuting))

    if useAlign8:
      tmpS = self.tmpS01
      module.add(self.getEdgeMovInstType()(EXEC(), sgpr(tmpS, self.laneSGPRC), "apply exec mask"))

    module.addComment1("buffer_store_dwordx4: write 8 16bit values (4 dwords, 2-aligned src)")
    module.add(BufferStoreB128(
      src=vgpr(vPack, 4),
      vaddr=vgpr(vAddrScratch),
      saddr=sgpr("SrdD", 4),
      soffset=0,
      mubuf=MUBUFModifiers(offen=True, offset12=globalOffset, glc=isGlc, slc=isSlc, nt=isNT),
      comment=f"16bit paired dwordx4 store tt0={tt0},{tt0+1}"
    ))

    if useAlign8:
      module.add(self.getEdgeMovInstType()(EXEC(), -1, "restore exec"))

    # WAR hazard: buffer_store_dwordx4 reads vPack[0:3] as source operands.
    # The next paired store's v_cvt_pk_bf16_f32 will overwrite vPack.
    # Insert nop to ensure the store has latched its source VGPRs.
    module.add(SNop(waitState=0, comment="1 wait state: WAR hazard between store src and next pack dst"))

    return module

  def _emit16bitSubtileScalarStore(self, addrCalc, sumIdx0: int, prefixOffset: int, tt0: int = 0, blockIdxM: int = 0, blockIdxN: int = 0) -> Module:
    """Emit a 16bit store for an orphan subtile element with no partner.

    Used when MIWaveTile[0] is odd and the last sba=0 element has no sba=1
    partner, or when batch boundaries split an (sba=0, sba=1) pair.

    The layout below is specific to the mfma instruction used here: lane l = LG*16 + r
    owns 4 output values at M-rows (LG*4 + 0..3) and a single N-column
    (l % 16 = r = lane_id & 15).  In column-major (row-first in memory) layout
    these 4 values ARE contiguous
    in memory (consecutive M-rows at fixed N-col), so we use 2x buffer_store_dwordx2
    after packing all 4 16bit values into 2 dwords.

    The per-lane vaddr encodes:
      vaddr = (lane_id & 15) * StrideD1J * bpe   [N-col byte offset within wave tile]
            + vLGDelta                            [LG*4 M-rows * bpe = LG*8 bytes]
            + wg0*MT0*bpe                         [workgroup M byte base]
            + waveId0 * waveM_stride * bpe        [M-wave byte offset within WG]
            + waveId1 * waveN_stride * StrideD1J * bpe  [N-wave byte offset]
    and a constant offset12 = globalOffset (encodes d0 M-tile position within wave).

    The SRD base encodes only wg1*MT1*StrideD1J*bpe (N workgroup offset).
    The M workgroup offset (wg0*MT0*bpe) and wave-within-WG offsets must be
    included in the vaddr explicitly.

    Args:
      addrCalc:     AddrCalculation for the element.
      sumIdx0:      elementSumIdx for the element.
      prefixOffset: parentWriter.states.c.startVgprValu (offset into ValuC).
    """
    module = Module("16bitSubtileScalarStore")
    isFp16 = self.kernel["ProblemType"]["DestDataType"].isHalf()

    ntd = self.kernel["NonTemporalD"]
    isGlc = bool(ntd & 0x1)
    isSlc = bool(ntd & 0x2)
    isNT  = bool(ntd & 0x4)

    # Scratch vgprs from the cvtVgprStruct block (overwritten each call):
    #   vPack+0  : 16bit packed dword (vc=0,1)
    #   vPack+1  : wave ID scratch / 16bit packed dword (vc=2,3)
    #   vPack+2  : per-lane vaddr (N-col byte offset + M offsets)
    #   vPack+3  : temp for N-col byte offset computation
    vPack    = self.cvtVgprStruct.vgprBf16Temp
    vLGDelta = self.cvtVgprStruct.vgprLaneGroupDelta  # LG*4*bpe = LG*8 bytes (pre-computed)

    # addrCalc.globalOffset was computed with bpeCexternal (may be 4 for _GlobalAccumulation kernels),
    # but the 16bit orphan store always targets the final 16bit output (bpeCexternalGSU1=2).
    bpeCurr = self.parentWriter.states.bpeCexternal
    bpe     = self.parentWriter.states.bpeCexternalGSU1  # always 2 for 16bit dest
    globalOffset = addrCalc.globalOffset * bpe // bpeCurr

    def vc(vi):
      idx = sumIdx0 + vi - prefixOffset
      return vgpr("ValuC+" + str(idx))

    # Derive the D-stride sgpr name (e.g. "StrideDJ") the same way incrementToNextRow does.
    packedC1  = self.kernel["PackedC1IndicesX"]
    indexChar = self.parentWriter.states.indexChars[packedC1[0]]
    strideD1J = "StrideD%s" % indexChar

    ws     = self.kernel["WavefrontSize"]
    miwg0  = self.kernel["MIWaveGroup"][0]
    miwg1  = self.kernel["MIWaveGroup"][1]
    matM   = self.kernel["MatrixInstM"]
    matN   = self.kernel["MatrixInstN"]

    typeStr = "fp16" if isFp16 else "bf16"
    VCvtPkF32to16 = VCvtPkF32toFP16 if isFp16 else VCvtPkF32toBF16
    module.addComment1(f"{typeStr} orphan subtile tt0={tt0}: pack 4 M-rows (vc=0..3) at fixed N-col, store as 2x dwordx2")

    # Build per-lane vaddr:
    #   vaddr = (lane_id & 15) * StrideD1J * bpe   [N-col]
    #         + vLGDelta                            [LG*4 M-rows = LG*8 bytes]
    #         + wg0*MT0*bpe                         [M-WG base]
    #         + waveId0*waveM_stride*bpe            [M-wave offset, if miwg0>1]
    #         + waveId1*waveN_stride*StrideD1J*bpe  [N-wave offset, if miwg1>1]
    # The SRD already encodes wg1*MT1*StrideD1J*bpe (N-WG offset).
    tmpS = self.tmpS01
    mt0bpe = self.kernel["MacroTile0"] * bpe

    module.addComment1("compute per-lane orphan vaddr = N_col_off + LG_M_off + wg0_M_off [+ wave offsets]")

    # N-col byte offset: (lane_id & 15) * StrideD1J * bpe
    module.add(VAndB32(dst=vgpr(vPack+2), src0=15, src1=vgpr("Serial"),
                       comment="col_in_wave = lane_id & 15  (N-column index)"))
    module.add(VMulLOU32(dst=vgpr(vPack+3), src0=vgpr(vPack+2), src1=sgpr(strideD1J),
                         comment="col_in_wave * StrideD1J"))
    if bpe == 2:
      module.add(VLShiftLeftB32(dst=vgpr(vPack+2), shiftHex=1, src=vgpr(vPack+3),
                                comment="N_col_off = col_in_wave * StrideD1J * 2"))
    else:
      module.add(VMulLOU32(dst=vgpr(vPack+2), src0=vgpr(vPack+3), src1=bpe,
                           comment="N_col_off = col_in_wave * StrideD1J * bpe"))

    # Add LG M-row offset: vLGDelta = LG*4*bpe = LG*8 bytes (pre-computed at batch start)
    module.add(VAddU32(dst=vgpr(vPack+2), src0=vgpr(vPack+2), src1=vgpr(vLGDelta),
                       comment="vaddr += LG_M_off (= vLGDelta = LG*4*bpe)"))

    # Add M-WG offset: wg0 * MT0 * bpe
    module.add(SMulI32(dst=sgpr(tmpS), src0=sgpr("WorkGroup0"), src1=mt0bpe,
                       comment="wg0_M_off = WorkGroup0 * MT0 * bpe"))
    module.add(VAddU32(dst=vgpr(vPack+2), src0=vgpr(vPack+2), src1=sgpr(tmpS),
                       comment="vaddr += wg0_M_off"))

    # Add M-wave offset: waveId0 * MIWaveTile[0] * matM * bpe.
    if miwg0 > 1:
      wsLog2 = int(log2(ws))
      waveM_stride_bpe = self.kernel["MIWaveTile"][0] * matM * bpe
      module.add(VLShiftRightB32(dst=vgpr(vPack+3), shiftHex=wsLog2, src=vgpr("Serial"),
                                 comment=f"waveId = Serial >> {wsLog2}"))
      if miwg0 & (miwg0 - 1) == 0:  # power of 2 — use AND mask
        module.add(VAndB32(dst=vgpr(vPack+3), src0=miwg0 - 1, src1=vgpr(vPack+3),
                           comment=f"waveId0 = waveId & {miwg0-1}"))
      else:
        raise NotImplementedError(f"Non-power-of-2 MIWaveGroup[0]={miwg0} not supported in orphan store")
      module.add(SMovB32(dst=sgpr(tmpS), src=waveM_stride_bpe,
                         comment=f"waveM_stride_bpe={waveM_stride_bpe}"))
      module.add(VMulLOU32(dst=vgpr(vPack+3), src0=vgpr(vPack+3), src1=sgpr(tmpS),
                           comment=f"wave_M_off = waveId0 * {waveM_stride_bpe}"))
      module.add(VAddU32(dst=vgpr(vPack+2), src0=vgpr(vPack+2), src1=vgpr(vPack+3),
                         comment="vaddr += wave_M_off"))

    # Add N-wave offset: waveId1 * MIWaveTile[1] * matN * StrideD1J * bpe.
    if miwg1 > 1:
      wsLog2 = int(log2(ws))
      waveN_stride_bpe = self.kernel["MIWaveTile"][1] * matN * bpe
      module.add(VLShiftRightB32(dst=vgpr(vPack+3), shiftHex=wsLog2, src=vgpr("Serial"),
                                 comment=f"waveId = Serial >> {wsLog2}"))
      if miwg0 & (miwg0 - 1) == 0:  # miwg0 is power of 2
        module.add(VLShiftRightB32(dst=vgpr(vPack+3), shiftHex=int(log2(miwg0)),
                                   src=vgpr(vPack+3),
                                   comment=f"waveId1 = waveId / {miwg0}"))
      else:
        raise NotImplementedError(f"Non-power-of-2 MIWaveGroup[0]={miwg0} not supported in orphan store")
      module.add(SMovB32(dst=sgpr(tmpS), src=waveN_stride_bpe,
                         comment=f"waveN_stride_bpe={waveN_stride_bpe}"))
      module.add(VMulLOU32(dst=vgpr(vPack+3), src0=vgpr(vPack+3), src1=sgpr(tmpS),
                           comment=f"waveId1 * {waveN_stride_bpe}"))
      module.add(VMulLOU32(dst=vgpr(vPack+3), src0=vgpr(vPack+3), src1=sgpr(strideD1J),
                           comment=f"wave_N_off = waveId1 * {waveN_stride_bpe} * StrideD1J"))
      module.add(VAddU32(dst=vgpr(vPack+2), src0=vgpr(vPack+2), src1=vgpr(vPack+3),
                         comment="vaddr += wave_N_off"))

    # Pack all 4 16bit values (consecutive M-rows at fixed N-col) into 2 dwords.
    # vc=0 → M-row+0 (lo16 of dword0), vc=1 → M-row+1 (hi16 of dword0)
    # vc=2 → M-row+2 (lo16 of dword1), vc=3 → M-row+3 (hi16 of dword1)
    #
    module.add(VCvtPkF32to16(dst=vgpr(vPack+0), src0=vc(0), src1=vc(1), comment=f"M-row+0/+1 -> {typeStr}"))
    module.add(VCvtPkF32to16(dst=vgpr(vPack+1), src0=vc(2), src1=vc(3), comment=f"M-row+2/+3 -> {typeStr}"))
    module.add(SNop(waitState=0, comment=f"delay after pk_{typeStr}"))

    useAlign8 = self.parentWriter.states.storeAlign8
    if useAlign8:
      self._emitAlign8ExecMask(module, self.tmpS01, self.tmpS23, blockIdxM, blockIdxN,
                               mGuardOffset=1, rowScaleShift=2)
      module.add(self.getEdgeMovInstType()(EXEC(), sgpr(self.tmpS01, self.laneSGPRC), "apply exec mask"))

    module.addComment1(f"buffer_store_b64: write 4 {typeStr} M-rows at fixed N-col (orphan subtile)")
    module.add(BufferStoreB64(
      src=vgpr(vPack+0, 2),
      vaddr=vgpr(vPack+2),
      saddr=sgpr("SrdD", 4),
      soffset=0,
      mubuf=MUBUFModifiers(offen=True, offset12=globalOffset, glc=isGlc, slc=isSlc, nt=isNT),
      comment=f"orphan tt0={tt0} vc=0..3: 4 consecutive M-rows at fixed N-col"
    ))

    if useAlign8:
      module.add(self.getEdgeMovInstType()(EXEC(), -1, "restore exec"))

    return module

  def _emitAtomicAdd(self, module: Module):
    ########################################
    # first attempt write
    module.addComment1("issue first atomic writes")
    for elementIdx in range(len(self.batchElements)):
      addrCalc = self.ss.elementAddr[elementIdx]
      mask     = self.ss.elementMask[elementIdx]

      # apply in-bounds exec mask
      if self.edge:
        module.add(self.getEdgeMovInstType()(EXEC(), sgpr(mask, self.laneSGPRC), "sgprs -> exec (before atomic)"))

      for avi in range(0, self.gwvw // self.atomicW):
        sumIdxV = self.ss.elementSumIdx[elementIdx] + avi
        newSumIdxV = sumIdxV - self.parentWriter.states.c.startVgprValu
        if self.parentWriter.do["GlobalWrite"]:
          if self.kernel["BufferStore"]:
            module.add(BufferAtomicAddF32(vgpr("ValuC+%u"%newSumIdxV), \
                         vgpr(addrCalc.addrDVgpr,1), \
                         sgpr("SrdD", 4), \
                         0,
                         MUBUFModifiers(offen=True, offset12=addrCalc.globalOffset),
                         "attempt write avi=%u" % (avi)))
          else:
            pass # TODO:

    if self.edge:
      module.add(self.getEdgeMovInstType()(EXEC(), -1, "full mask -> exec"))

  def _emitCasAdd(self, module: Module):
    # TODO for atomic GWVW:
    #  - Use vi to compute addresses, sumIdx.
    #  - Need a solution for the mask.  Can move to all buffer or can fix?
    element = self.batchElements[0]
    d1 = element[0]
    d0 = element[1]
    vc1 = element[2]
    vc0 = element[3]
    labels = self.parentWriter.labels
    labelString = "Global_Write%s%s_%u_%u_%u_%u" % ("_Beta" if self.beta else "", "_Edge" if self.edge else "", vc0, vc1, d0, d1 )
    labelComment = "Global_Write (Beta) (Edge) vc0 vc1 d0 d1"
    label = Label(labels.getName(labelString), labelComment)
    labelString += "_EarlyExit"
    labelAfterAtomicLoop = Label(labels.getName(labelString), labelComment)

    ########################################
    # wait for batched load
    # TODO - we are always atomic here?
    module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="wait C (atomic)"))
    ########################################
    # first attempt write
    module.addComment1("issue first atomic writes")
    for elementIdx, element in enumerate(self.batchElements):
      addrCalc = self.ss.elementAddr[elementIdx]
      mask = self.ss.elementMask[elementIdx]

      # apply in-bounds exec mask
      if self.edge:
        module.add(self.getEdgeMovInstType()(EXEC(), sgpr(mask, self.parentWriter.states.laneSGPRCount), "sgprs -> exec (before atomic)"))

      for avi in range(0, self.gwvw//self.atomicW):
        dataV = self.ss.elementData[elementIdx] + int(avi*self.ss.cfg.numVgprsPerDataPerVI)
        sumIdxV = self.ss.elementSumIdx[elementIdx] + avi
        ## number of src[s]/dst[s] register for DGEMM / SGEMM HGEMM
        vgprCnt = 2 if self.kernel["ProblemType"]["DestDataType"].isDouble() else 1
        if self.kernel["ProblemType"]["DestDataType"].numRegisters() < 1 and not self.kernel["_GlobalAccumulation"]:
          sumIdxV //= 2
        if self.kernel["ProblemType"]["DestDataType"].isDouble(): sumIdxV = sumIdxV * 2
        newSumIdxV = sumIdxV - self.parentWriter.states.c.startVgprValu
        bpm = self.parentWriter.states.bpeCexternal * self.atomicW
        # Calculate vgpr Index for 32-bit/64-bit instruction
        # DGEMM use SRCS[2] register
        vgprIdx = 1*(bpm//4)
        # for atomic, data[1] = original c, data[0] = new c
        module.add(self._chooseAddForAtomic(self.kernel, \
                  vgpr(dataV+0,vgprCnt), vgpr(dataV+1*vgprIdx,vgprCnt), vgpr("ValuC+%u"%newSumIdxV,vgprCnt), \
                  "desired value avi=%u"%avi))

        # attempt write
        atomicDestVgpr = dataV if self.kernel["BufferStore"] else dataV+2
        if self.parentWriter.do["GlobalWrite"]:
          if self.kernel["BufferStore"]:
            # use cmpswap_x2 for DGEMM in CAS loop
            if self.kernel["ProblemType"]["DestDataType"].isDouble():
              module.add(BufferAtomicCmpswapB64(vgpr(dataV,4), \
                              vgpr(addrCalc.addrDVgpr,1), \
                              sgpr("SrdD", 4),  \
                              0,
                              MUBUFModifiers(offen=True, offset12=addrCalc.globalOffset, glc=True),
                              "attempt write avi=%u"%(avi)))
            else:
            # use cmpswap for SGEMM in CAS loop
              module.add(BufferAtomicCmpswapB32(vgpr(dataV,2), \
                           vgpr(addrCalc.addrDVgpr,1), \
                           sgpr("SrdD", 4), \
                           0, \
                           MUBUFModifiers(offen=True, offset12=addrCalc.globalOffset, glc=True), \
                           "attempt write avi=%u"%(avi)))
          else:
            module.add(FlatAtomicCmpswapB32(vgpr(atomicDestVgpr), \
                                            vgpr(addrCalc.addrDVgpr,2), \
                                            vgpr(dataV,2),
                                            FLATModifiers(glc=True),
                                            "attempt write"))
        else:
            # Fake successful CAS swap
            module.add(VMovB32(vgpr(atomicDestVgpr), vgpr(dataV+1), "Fake successful CAS" ))

    ########################################
    # wait for first attempt write
    module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="wait for atomic writes"))
    ########################################
    # check first attempt
    module.addComment1("check success of writes, update masks")
    for elementIdx, element in enumerate(self.batchElements):
      mask = self.ss.elementMask[elementIdx]

      # calculate new masks
      if self.edge:
        module.add(self.getEdgeMovInstType()(EXEC(), sgpr(mask, self.laneSGPRC), "sgprs -> exec"))
        for avi in range(0, self.gwvw // self.atomicW):
          dataV = self.ss.elementData[elementIdx] + int(avi * self.ss.cfg.numVgprsPerDataPerVI)
          atomicDestVgpr = dataV if self.kernel["BufferStore"] else dataV+2
          # need to apply element mask before comparison
          # so that all valid lanes are doing the cmp
          if avi == 0:
            # use u64 for DGEMM
            if self.kernel["ProblemType"]["DestDataType"].isDouble():
              module.add(VCmpNeU64(sgpr(self.tmpS01, self.laneSGPRC), vgpr(atomicDestVgpr,2), \
                  vgpr(dataV+2,2), comment="c read during atomic == c read during prior load (avi=%u, first)"%avi))
            else:
              module.add(VCmpNeU32(sgpr(self.tmpS01, self.laneSGPRC), vgpr(atomicDestVgpr), \
                  vgpr(dataV+1), comment="c read during atomic == c read during prior load (avi=%u, first)"%avi))
          else:
            if self.kernel["ProblemType"]["DestDataType"].isDouble():
              module.add(VCmpNeU64(sgpr(self.tmpS23, self.laneSGPRC), vgpr(atomicDestVgpr,2), \
                  vgpr(dataV+2,2), comment="c read during atomic != c read during prior load"))
            else:
              module.add(VCmpNeU32(sgpr(self.tmpS23, self.laneSGPRC), vgpr(atomicDestVgpr), \
                  vgpr(dataV+1), comment="c read during atomic == c read during prior load (avi=%u)"%avi))
            module.add(self.getEdgeOrInstType()(sgpr(self.tmpS01, self.laneSGPRC), \
                  sgpr(self.tmpS01, self.laneSGPRC), sgpr(self.tmpS23, self.laneSGPRC), "combine with tmp mask"))

        module.add(self.getEdgeAndInstType()(sgpr(mask, self.laneSGPRC), sgpr(self.tmpS01, self.laneSGPRC), sgpr(mask,self.laneSGPRC), "inBounds & must try again"))

      else:
        for avi in range(0, self.gwvw//self.atomicW):
          dataV = self.ss.elementData[elementIdx] + int(avi*self.ss.cfg.numVgprsPerDataPerVI)
          atomicDestVgpr = dataV if self.kernel["BufferStore"] else dataV+2
          if self.kernel["ProblemType"]["DestDataType"].isDouble():
            module.add(VCmpNeU64(sgpr(mask, self.laneSGPRC), vgpr(atomicDestVgpr,2), \
                vgpr(dataV+2,2), comment="c read during atomic != c read during prior load"))
          else:
            module.add(VCmpNeU32(sgpr(mask, self.laneSGPRC), vgpr(atomicDestVgpr), \
                vgpr(dataV+1), comment="c read during atomic != c read during prior load"))

    # or masks together to check early exit
    module.addComment1("or masks to check for exit")
    module.add(self.getEdgeMovInstType()(sgpr(self.tmpS01, self.laneSGPRC), 0, "empty mask"))
    for elementIdx in range(0, len(self.batchElements)):
      mask = self.ss.elementMask[elementIdx]
      module.add(self.getEdgeOrInstType()(sgpr(self.tmpS01, self.laneSGPRC), sgpr(mask, self.laneSGPRC), sgpr(self.tmpS01, self.laneSGPRC), "or to add threads"))
    module.add(self.getSOrSaveExecType()(sgpr(self.tmpS23,self.laneSGPRC), sgpr(self.tmpS01,self.laneSGPRC), "apply combined mask"))
    module.add(SCBranchExecZ(labelAfterAtomicLoop.getLabelName(), "if exec is zero skip loop"))

    # begin atomic loop
    module.addComment1("atomic CAS loop")
    module.add(label)

    module.addComment1("apply updated masks and issue writes again")
    for elementIdx in range(0, len(self.batchElements)):
      addrCalc = self.ss.elementAddr[elementIdx]
      addr = addrCalc.addrDVgpr
      mask = self.ss.elementMask[elementIdx]
      vgprCnt = 2 if self.kernel["ProblemType"]["DestDataType"].isDouble() else 1   # number of registers for f32/f64
      bpm = self.parentWriter.states.bpeCexternal * self.atomicW
      vgprIdx = 1*(bpm//4)   # index register

      for avi in range(0, self.gwvw//self.atomicW):
        dataV = self.ss.elementData[elementIdx] + int(avi*self.ss.cfg.numVgprsPerDataPerVI)
        atomicDestVgpr = dataV if self.kernel["BufferStore"] else dataV+2
        sumIdxV = self.ss.elementSumIdx[elementIdx] + avi
        if self.kernel["ProblemType"]["DestDataType"].numRegisters() < 1 and not self.kernel["_GlobalAccumulation"]:
          sumIdxV //= 2
        if self.kernel["ProblemType"]["DestDataType"].isDouble():
          sumIdxV =  sumIdxV * 2
        newSumIdxV = sumIdxV - self.parentWriter.states.c.startVgprValu

        # apply mask for element
        module.add(self.getEdgeMovInstType()(EXEC(), sgpr(mask,self.laneSGPRC), "must try again"))
        if self.kernel["ProblemType"]["DestDataType"].isDouble():
          #64-bit C val move by 2 32-bit instructions
          module.add(VMovB32(vgpr(dataV+2), vgpr(atomicDestVgpr), "dataV+2 = tmp (new original C)" ))
          module.add(VMovB32(vgpr(dataV+3), vgpr(atomicDestVgpr+1), "dataV+3 = tmp (new original C)" ))
        else:
          module.add(VMovB32(dst=vgpr(dataV+1), src=vgpr(atomicDestVgpr), comment="dataV+1 = tmp (new original C)" ))
        module.add(self._chooseAddForAtomic(self.kernel, \
                        vgpr(dataV+0,vgprCnt), vgpr(dataV+1*vgprIdx,vgprCnt), vgpr("ValuC+%u"%newSumIdxV,vgprCnt), \
                        "newC = rC + originalC"))
        if self.parentWriter.do["GlobalWrite"]:
          if self.kernel["BufferStore"]:
            # Using no-ret version here?
            # cmpswap_x2 for DGEMM
            if self.kernel["ProblemType"]["DestDataType"].isDouble():
              module.add(BufferAtomicCmpswapB64(vgpr(dataV,4), \
                          vgpr(addr,1), \
                          sgpr("SrdD", 4), \
                          0,
                          MUBUFModifiers(offen=True, offset12=addrCalc.globalOffset, glc=True,),
                          "try again"))
            else:
              module.add(BufferAtomicCmpswapB32(
                          vgpr(dataV,2), \
                          vgpr(addr,1), \
                          sgpr("SrdD", 4), \
                          0,
                          MUBUFModifiers(offen=True, offset12=addrCalc.globalOffset, glc=True),
                          "try again"))
          else:
            module.add(FlatAtomicCmpswapB32(vgpr(atomicDestVgpr), \
                                            vgpr(addr,2), \
                                            vgpr(dataV,2), \
                                            FLATModifiers(glc=True), \
                                            "try again"))

    # wait for batched write
    module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="wait for atomic writes"))
    # check batched write success
    module.addComment1("apply masks and check for success")
    for elementIdx in range(0, len(self.batchElements)):
      data = self.ss.elementData[elementIdx]
      mask = self.ss.elementMask[elementIdx]
      for avi in range(0, self.gwvw//self.atomicW):
        dataV = self.ss.elementData[elementIdx] + int(avi*self.ss.cfg.numVgprsPerDataPerVI)
        atomicDestVgpr = dataV if self.kernel["BufferStore"] else dataV+2

        # apply mask for element
        module.add(self.getEdgeMovInstType()(EXEC(), sgpr(mask,self.laneSGPRC), "must try again"))

        # compare success
        if self.kernel["ProblemType"]["DestDataType"].isDouble():
          module.add(VCmpNeU64(sgpr(self.tmpS01,self.laneSGPRC), vgpr(data+2,2), vgpr(atomicDestVgpr,2), \
              comment="c read during atomic != c read during prior load"))
        else:
          module.add(VCmpNeU32(sgpr(self.tmpS01,self.laneSGPRC), vgpr(data+1), vgpr(atomicDestVgpr), \
              comment="c read during atomic == c read during prior load"))
        # update element mask
        module.add(self.getEdgeAndInstType()(sgpr(mask,self.laneSGPRC), sgpr(self.tmpS01,self.laneSGPRC), sgpr(mask,self.laneSGPRC), "inBounds & must try again"))

    # or masks together
    module.addComment1("or masks to check for exit")
    module.add(self.getEdgeMovInstType()(sgpr(self.tmpS01,self.laneSGPRC), 0, "empty mask"))
    for elementIdx in range(0, len(self.batchElements)):
      mask = self.ss.elementMask[elementIdx]
      module.add(self.getEdgeOrInstType()(sgpr(self.tmpS01,self.laneSGPRC), sgpr(mask,self.laneSGPRC), sgpr(self.tmpS01,self.laneSGPRC), "or to add threads"))

    # apply combined masks and exit
    module.add(self.getSOrSaveExecType()(sgpr(self.tmpS23, self.laneSGPRC), sgpr(self.tmpS01,self.laneSGPRC), "apply combined mask"))
    module.add(SCBranchExecNZ(label.getLabelName(), "try again if not complete"))
    module.add(labelAfterAtomicLoop)
    module.add(self.getEdgeMovInstType()(EXEC(), -1, "full mask -> exec"))

  def _checkAtomicPreconditions(self) -> bool:
    if self.atomic:
      # all kinds of code relies on this assumption:
      if self.atomicW > self.gwvw:
        return False

      if (self.kernel["ProblemType"]["DataType"].isHalf() or self.kernel["ProblemType"]["DataType"].isBFloat16()) \
        and not self.kernel["_GlobalAccumulation"]:
        return self.atomicW >= 2
    return True

  def _storeSyncOpt(self, module: Module):
    module.add(SSleep(self.kernel["StoreSyncOpt"] - 1, "optimization: sync and wait"))
    module.add(SBarrier())

  def _applyAlpha(self, kernel, gwvw, elementSumIdx, elementIdx, tmpS01, usePK=False):
    module = Module("applyAlpha")

    if kernel["_GlobalAccumulation"] == 'MultipleBuffer':
      return module

    if self.parentWriter.do["ApplyAlpha"]:
      for vi in range(0, gwvw):
        sumIdxV = elementSumIdx[elementIdx] + vi

        if kernel["ProblemType"]["ComputeDataType"].isHalf() and not kernel["ProblemType"]["HighPrecisionAccumulate"]:
          # (h,h,h,h,h,h), internal alpha is f16 (2-16bits)
          if sumIdxV%2:
            newSumIdx = sumIdxV // 2 - self.parentWriter.states.c.startVgprValu
            module.add(VMulPKF16(dst=vgpr("ValuC+%u"%(newSumIdx)), src0=sgpr("Alpha"), src1=vgpr("ValuC+%u"%(newSumIdx)), comment="*= alpha sumIdx=%u vi=%u"%(elementSumIdx[elementIdx], vi)))

        # Int8 (TODO- Int8x4 not checked, but should be OK)
        elif kernel["ProblemType"]["ComputeDataType"].isInt32():
          newSumIdx = sumIdxV - self.parentWriter.states.c.startVgprValu
          # below assume we use v_mul_lo_u32. Could also use v_mul_i32_i24.
          # module.add(VMulI32I24(dst=vgpr("ValuC+%u"%newSumIdx), src0=sgpr("Alpha"), src1=vgpr("ValuC+%u"%newSumIdx), comment="*= alpha" )_
          module.add(VMulLOU32(dst=vgpr("ValuC+%u"%newSumIdx), src0=sgpr("Alpha"), src1=vgpr("ValuC+%u"%newSumIdx), comment="*= alpha" ))
          if usePK:
            module.add(VMulLOU32(dst=vgpr("ValuC+%u"%(newSumIdx+1)), src0=sgpr("Alpha"), src1=vgpr("ValuC+%u"%(newSumIdx+1)), comment="*= alpha" ))
          if self.parentWriter.db["ForceExpectedValue"]:
            module.add(VMovB32(dst=vgpr("ValuC+%u"%newSumIdx), src=self.parentWriter.db["ValueCExpectedValue"], comment="force expected value" ))
          if self.parentWriter.db["CheckValueC"]:
            module.add(SMovB32(dst=sgpr(tmpS01), src=self.parentWriter.db["ValueCExpectedValue"], comment="Move expected value"))
            module.add(self.parentWriter.getCmpAssert(self.parentWriter.asmAssert.eq, vgpr("ValuC+%u"%newSumIdx), sgpr(tmpS01)))

        # sgemm, HPA-bfgemm(b,b,b,b,s,s), and HPA-hgemm(h,h,h,h,s,s)
        # (h,h,h,h,h,h) + HPA (will be converted to (h,h,h,h,s,s)), internal alpha is single
        elif kernel["ProblemType"]["ComputeDataType"].isSingle() or (kernel["ProblemType"]["ComputeDataType"].isHalf() and kernel["ProblemType"]["HighPrecisionAccumulate"]):

          if kernel["ProblemType"]["DataType"].isInt8() and kernel["ProblemType"]["HighPrecisionAccumulate"]:
            if usePK or gwvw > 1:
              if vi % 2 == 0:
                module.add(VCvtI32toF32(dst=vgpr("ValuC+%u"%sumIdxV), src=vgpr("ValuC+%u"%sumIdxV), comment="convert to fp32" ))
                module.add(VCvtI32toF32(dst=vgpr("ValuC+%u"%(sumIdxV+1)), src=vgpr("ValuC+%u"%(sumIdxV+1)), comment="convert to fp32" ))
            else:
              module.add(VCvtI32toF32(dst=vgpr("ValuC+%u"%sumIdxV), src=vgpr("ValuC+%u"%sumIdxV), comment="convert to fp32" ))

          newSumIdx = sumIdxV - self.parentWriter.states.c.startVgprValu
          # Use pk if possible
          if usePK or gwvw > 1:
            if newSumIdx % 2 == 0:
              module.add(VMulPKF32(dst=vgpr("ValuC+%u"%newSumIdx, 2), src0=sgpr("Alpha",2), src1=vgpr("ValuC+%u"%newSumIdx,2), vop3=VOP3PModifiers(op_sel_hi=[0,1,1]), comment="*= alpha (pk)"))
          else:
            module.add(VMulF32(dst=vgpr("ValuC+%u"%newSumIdx), src0=sgpr("Alpha"), src1=vgpr("ValuC+%u"%newSumIdx), comment="*= alpha" ))
          if self.parentWriter.db["ForceExpectedValue"]:
            module.add(VMovB32(dst=vgpr("ValuC+%u"%newSumIdx), src=self.parentWriter.db["ValueCExpectedValue"], comment="force expected value" ))
          if self.parentWriter.db["ForceVSerial"]:
            module.add(VMovB32(dst=vgpr("ValuC+%u"%newSumIdx), src=vgpr("Serial"), comment="force expected value to serial" ))
          if self.parentWriter.db["CheckValueC"]:
            module.add(SMovB32(dst=sgpr(tmpS01), src=self.parentWriter.db["ValueCExpectedValue"], comment="Move expected value"))
            module.add(self.parentWriter.getCmpAssert(self.parentWriter.asmAssert.eq, vgpr("ValuC+%u"%newSumIdx), sgpr(tmpS01)))

        # dgemm
        elif kernel["ProblemType"]["ComputeDataType"].isDouble():
          newSumIdx = sumIdxV * 2 - self.parentWriter.states.c.startVgprValu
          module.add(VMulF64(dst=vgpr("ValuC+%u"%(newSumIdx),2), src0=sgpr("Alpha",2), src1=vgpr("ValuC+%u"%(newSumIdx),2), comment="*= alpha"))
          if usePK:
            module.add(VMulF64(dst=vgpr("ValuC+%u"%(newSumIdx+2),2), src0=sgpr("Alpha",2), src1=vgpr("ValuC+%u"%(newSumIdx+2),2), comment="*= alpha"))

        # single precision complex
        elif kernel["ProblemType"]["ComputeDataType"].isSingleComplex():
          newSumIdx = sumIdxV * 2 - self.parentWriter.states.c.startVgprValu
          tmpVgpr = self.parentWriter.vgprPool.checkOut(1, tag="_applyAlpha_tmpVgpr")
          module.add(VMovB32(dst=vgpr(tmpVgpr), src=vgpr("ValuC+%u"%(newSumIdx)), comment="store Cr"))
          module.add(VMulF32(dst=vgpr("ValuC+%u"%(newSumIdx)), src0=sgpr("Alpha"), src1=vgpr("ValuC+%u"%(newSumIdx)), comment="*= alpha ( Cr = Ar * Cr)"))
          module.add(VMacF32(dst=vgpr("ValuC+%u"%(newSumIdx)), src0=(sgpr("Alpha+1").getMinus()), src1=vgpr("ValuC+%u"%(newSumIdx+1)), comment="*= alpha ( Cr += -Ai * Ci )"))
          module.add(VMulF32(dst=vgpr("ValuC+%u"%(newSumIdx+1)), src0=sgpr("Alpha"), src1=vgpr("ValuC+%u"%(newSumIdx+1)), comment="*= alpha ( Ci = Ar * Ci)"))
          module.add(VMacF32(dst=vgpr("ValuC+%u"%(newSumIdx+1)), src0=sgpr("Alpha+1"), src1=vgpr(tmpVgpr), comment="*= alpha ( Ci += Ai * Cr_backup )"))
          if usePK:
            newSumIdx2 = newSumIdx + 2
            module.add(VMovB32(dst=vgpr(tmpVgpr), src=vgpr("ValuC+%u"%(newSumIdx2)), comment="store Cr"))
            module.add(VMulF32(dst=vgpr("ValuC+%u"%(newSumIdx2)), src0=sgpr("Alpha"), src1=vgpr("ValuC+%u"%(newSumIdx2)), comment="*= alpha ( Cr = Ar * Cr)"))
            module.add(VMacF32(dst=vgpr("ValuC+%u"%(newSumIdx2)), src0=(sgpr("Alpha+1").getMinus()), src1=vgpr("ValuC+%u"%(newSumIdx2+1)), comment="*= alpha ( Cr += -Ai * Ci )"))
            module.add(VMulF32(dst=vgpr("ValuC+%u"%(newSumIdx2+1)), src0=sgpr("Alpha"), src1=vgpr("ValuC+%u"%(newSumIdx2+1)), comment="*= alpha ( Ci = Ar * Ci)"))
            module.add(VMacF32(dst=vgpr("ValuC+%u"%(newSumIdx2+1)), src0=sgpr("Alpha+1"), src1=vgpr(tmpVgpr), comment="*= alpha ( Ci += Ai * Cr_backup )"))
          self.parentWriter.vgprPool.checkIn(tmpVgpr)

        # double precision complex
        elif kernel["ProblemType"]["ComputeDataType"].isDoubleComplex():
          newSumIdx = sumIdxV * 4 - self.parentWriter.states.c.startVgprValu
          vtmp1 = self.parentWriter.vgprPool.checkOutAligned(2, 2, tag="_applyAlpha_vtmp1")
          vtmp2 = self.parentWriter.vgprPool.checkOutAligned(2, 2, tag="_applyAlpha_vtmp2")
          # tmp1 = a.real * b.real (t1 = Ar*Cr)
          module.add(VMulF64(dst=vgpr(vtmp1,2), src0=sgpr("Alpha+0",2), src1=vgpr("ValuC+%u"%(newSumIdx+0),2)))
          # tmp2 = a.imag * b.real (t2 = Ai*Cr)
          module.add(VMulF64(dst=vgpr(vtmp2,2), src0=sgpr("Alpha+2",2), src1=vgpr("ValuC+%u"%(newSumIdx+0),2)))
          # c.real = a.real * b.real - a.imag * b.imag = tmp1 - a.imag * b.imag (Cr = -Ai*Ci + t1).
          module.add(VFmaF64(dst=vgpr("ValuC+%u"%(newSumIdx+0),2), src0=sgpr("Alpha+2",2), src1=vgpr("ValuC+%u"%(newSumIdx+2),2).getMinus(), src2=vgpr(vtmp1,2)))
          # c.imag = a.real * b.imag + a.imag * b.real = a.real * b.imag + tmp2 (Ci = Ar*Ci + t2).
          module.add(VFmaF64(dst=vgpr("ValuC+%u"%(newSumIdx+2),2), src0=sgpr("Alpha+0",2), src1=vgpr("ValuC+%u"%(newSumIdx+2),2), src2=vgpr(vtmp2,2)))
          if usePK:
            newSumIdx2 = newSumIdx + 4
            # tmp1 = a.real * b.real
            module.add(VMulF64(dst=vgpr(vtmp1,2), src0=sgpr("Alpha+0",2), src1=vgpr("ValuC+%u"%(newSumIdx2+0),2)))
            # tmp2 = a.imag * b.real
            module.add(VMulF64(dst=vgpr(vtmp2,2), src0=sgpr("Alpha+2",2), src1=vgpr("ValuC+%u"%(newSumIdx2+0),2)))
            # c.real = a.real * b.real - a.imag * b.imag = tmp1 - a.imag * b.imag
            module.add(VFmaF64(dst=vgpr("ValuC+%u"%(newSumIdx2+0),2), src0=sgpr("Alpha+2",2), src1=vgpr("ValuC+%u"%(newSumIdx2+2),2).getMinus(), src2=vgpr(vtmp1,2)))
            # c.imag = a.real * b.imag + a.imag * b.real = a.real * b.imag + tmp2
            module.add(VFmaF64(dst=vgpr("ValuC+%u"%(newSumIdx2+2),2), src0=sgpr("Alpha+0",2), src1=vgpr("ValuC+%u"%(newSumIdx2+2),2), src2=vgpr(vtmp2,2)))
          self.parentWriter.vgprPool.checkIn(vtmp1)
          self.parentWriter.vgprPool.checkIn(vtmp2)
    return module

  def _addSumAlphaWithCBeta(self, kernel, ss, gwvw, elementIdx, vc0, tmpVgpr, cvtVgprStruct):
    module = Module("addSumAlphaWithCBeta #elementIdx%u, vc0 %u"%(elementIdx, vc0))
    for vi in range(0, gwvw):
      dataV = ss.elementData[elementIdx] + int(vi*ss.cfg.numVgprsPerDataPerVI)
      sumIdxV = ss.elementSumIdx[elementIdx] + vi
      if kernel["ProblemType"]["DestDataType"].isHalf():
        if not kernel["ProblemType"]["HighPrecisionAccumulate"]:
          if self.parentWriter.states.asmCaps["HasWMMA_V1"] and kernel["EnableMatrixInstruction"]:
            dataV = ss.elementData[elementIdx] + int(vi / 2 * ss.cfg.numVgprsPerDataPerVI)
            if (vi % 2) == 0:
              module.add(VMulPKF16(dst=vgpr(dataV), src0=sgpr("Beta"), src1=vgpr(dataV+0), \
                    comment="%s = C*beta ei=%u vi=%u"%(vgpr(dataV),elementIdx, vi)))
            else:
              module.add(VLShiftRightB32(dst=vgpr(dataV), shiftHex=16, src=vgpr(dataV), \
                    comment="shift 16bit to get next half of packed ValueC"))
            # dataV+0 = new c = old c*beta + rC
            module.add(VAddPKF16(dst=vgpr("ValuC+%u"%(sumIdxV)), src0=vgpr(dataV), src1=vgpr("ValuC+%u"%(sumIdxV)), \
                comment="sum*alpha + C*beta"))
          elif sumIdxV%2==0 or (not ss.cfg.halfDataRegPerVI and gwvw==1):
            newSumIdxV = sumIdxV // 2 - self.parentWriter.states.c.startVgprValu
            # dataV+0 = new c = old c*beta
            module.add(VMulPKF16(dst=vgpr(dataV), src0=sgpr("Beta"), src1=vgpr(dataV+0), \
                comment="%s = C*beta ei=%u vi=%u"%(vgpr(dataV),elementIdx, vi)))
            # dataV+0 = new c = old c*beta + rC
            module.add(VAddPKF16(dst=vgpr("ValuC+%u"%(newSumIdxV)), src0=vgpr(dataV), src1=vgpr("ValuC+%u"%(newSumIdxV)), \
                comment="sum*alpha + C*beta"))
          else:
            pass # add will have been done previously
        else: # HPA
          newSumIdxV = sumIdxV - self.parentWriter.states.c.startVgprValu
          # dataV+0 = new c = old c*beta + rC
          # src0 = beta = f32 = opsel 00
          # src1 = dataV = f16.lo = opsel 10 or 11 depending on even/odd
          # src2 = sumIdxV = f32 = opsel 00
          dataCExternal = ss.elementData[elementIdx] + vi//2
          hi16 = (vi + gwvw*vc0) % 2
          module.add(self.parentWriter.states.mixinst(dst=vgpr("ValuC+%u"%newSumIdxV), src0=sgpr("Beta"), \
              src1=vgpr(dataCExternal), src2=vgpr("ValuC+%u"%newSumIdxV), \
              vop3=VOP3PModifiers(op_sel=[0,hi16,0], op_sel_hi=[0,1,0]),
              comment="//C*=beta"))

      elif kernel["ProblemType"]["DestDataType"].isBFloat16():
        if kernel["ProblemType"]["HighPrecisionAccumulate"]:
          # dataV+0 = new c = old c*beta + rC
          # src0 = beta = f32 = opsel 00
          # src1 = dataV = f16.lo = opsel 10 or 11 depending on even/odd
          # src2 = sumIdxV = f32 = opsel 00
          dataCExternal = ss.elementData[elementIdx] + vi//2
          module.add(VCvtBF16toFP32(dst=vgpr(tmpVgpr), src=vgpr(dataCExternal), vgprMask=vgpr(cvtVgprStruct.vgprBf16Mask), vi=(vi)))
          newSumIdxV = sumIdxV - self.parentWriter.states.c.startVgprValu
          module.add(VMacF32(dst=vgpr("ValuC+%u"%newSumIdxV), src0=vgpr(tmpVgpr), src1=sgpr("Beta"), \
              comment="finalSum = sum*alpha + C*beta"))
      elif kernel["ProblemType"]["DestDataType"].isSingle():
        newSumIdxV = sumIdxV - self.parentWriter.states.c.startVgprValu
        module.add(VMacF32(dst=vgpr("ValuC+%u"%newSumIdxV), src0=vgpr(dataV+0), src1=sgpr("Beta"), \
            comment="finalSum = sum*alpha + C*beta"))

      elif kernel["ProblemType"]["DestDataType"].isInt8():
        if kernel["ProblemType"]["HighPrecisionAccumulate"]:
          dataCExternal   = ss.elementData[elementIdx] + (vi // 4)
          byteIdx       = vi %  4
          if (vi%4) != 3:
            module.add(VMovB32(dst=vgpr(tmpVgpr+1), src=hex(byteIdx * 8), comment="value = %u"%(byteIdx * 8)))
            module.add(VBfeI32(dst=vgpr(tmpVgpr), src0=vgpr(dataCExternal), src1=vgpr(tmpVgpr+1), src2=8, comment="int8 to int32"))
          else:
            module.add(VAShiftRightI32(dst=vgpr(tmpVgpr), shiftHex=24, src=vgpr(dataCExternal), comment="int8 to int32"))

          newSumIdxV = sumIdxV - self.parentWriter.states.c.startVgprValu
          if kernel["ProblemType"]["ComputeDataType"].isSingle():
            module.add(VCvtI32toF32(dst=vgpr(tmpVgpr), src=vgpr(tmpVgpr), comment="convert to fp32" ))
            module.add(VMacF32(dst=vgpr("ValuC+%u"%newSumIdxV), src0=vgpr(tmpVgpr), src1=sgpr("Beta"), \
                               comment="finalSum = sum*alpha + C*beta"))
          else:
            module.add(VMulLOU32(dst=vgpr(tmpVgpr), src0=sgpr("Beta"), src1=vgpr(tmpVgpr), comment="C = C*beta"))
            module.add(VAddU32(dst=vgpr("ValuC+%u"%newSumIdxV), src0=vgpr(tmpVgpr), src1=vgpr("ValuC+%u"%newSumIdxV), comment="finalSum = sum*alpha + C*beta"))

      elif kernel["ProblemType"]["DestDataType"].isInt32():
        newSumIdxV = sumIdxV - self.parentWriter.states.c.startVgprValu
        if kernel["ProblemType"]["ComputeDataType"].isSingle():
          module.add(VCvtI32toF32(dst=vgpr(dataV+0), src=vgpr(dataV+0), comment="convert to fp32" ))
          module.add(VMacF32(dst=vgpr("ValuC+%u"%newSumIdxV), src0=vgpr(dataV+0), src1=sgpr("Beta"), comment="finalSum = sum*alpha + C*beta"))
        else:
          # assume we will need to replace v_mac_f32 with v_add_u32 and s_mul_lo_i32
          # v_mad_i32_i24
          # module.add(VMadI32I24(dst=vgpr("ValuC+%u"%sumIdxV), src0=vgpr(dataV+0), src1=sgpr("Beta"), src2=vgpr("ValuC+%u"%sumIdxV), \
          #     comment="finalSum = sum*alpha + C*beta"))
          module.add(VMulLOU32(dst=vgpr(dataV+0), src0=sgpr("Beta"), src1=vgpr(dataV+0), comment="C = C*beta"))
          module.add(VAddU32(dst=vgpr("ValuC+%u"%newSumIdxV), src0=vgpr(dataV+0), src1=vgpr("ValuC+%u"%newSumIdxV), comment="finalSum = sum*alpha + C*beta"))

      elif kernel["ProblemType"]["DestDataType"].isDouble():
        newSumIdxV = sumIdxV * 2 - self.parentWriter.states.c.startVgprValu
        # dataV+0 = new c = old c*beta
        module.add(VFmaF64(dst=vgpr("ValuC+%u"%(newSumIdxV),2), src0=vgpr(dataV+0,2), src1=sgpr("Beta",2), src2=vgpr("ValuC+%u"%(newSumIdxV),2), \
            comment="finalSum = sum*alpha + C*beta"))

      # single precision complex
      elif kernel["ProblemType"]["DestDataType"].isSingleComplex():
        newSumIdxV = sumIdxV * 2 - self.parentWriter.states.c.startVgprValu
        module.add(VMacF32(dst=vgpr("ValuC+%u"%(newSumIdxV)), src0=vgpr(dataV+0), src1=sgpr("Beta"), comment="finalSum Cr += old Cr * Br"))
        module.add(VMacF32(dst=vgpr("ValuC+%u"%(newSumIdxV)), src0=vgpr(dataV+1), src1=sgpr("Beta+1").getMinus(), comment="finalSum Cr += old Ci * -Bi"))
        module.add(VMacF32(dst=vgpr("ValuC+%u"%(newSumIdxV+1)), src0=vgpr(dataV+1), src1=sgpr("Beta"), comment="finalSum Ci += old Ci * Br"))
        module.add(VMacF32(dst=vgpr("ValuC+%u"%(newSumIdxV+1)), src0=vgpr(dataV+0), src1=sgpr("Beta+1"), comment="finalSum Ci += old Cr * Bi"))

      # double precision complex
      elif kernel["ProblemType"]["DestDataType"].isDoubleComplex():
        newSumIdxV = sumIdxV * 4 - self.parentWriter.states.c.startVgprValu
        module.add(VFmaF64(dst=vgpr("ValuC+%u"%(newSumIdxV+0),2), src0=vgpr(dataV+0,2), src1=sgpr("Beta+0",2), src2=vgpr("ValuC+%u"%(newSumIdxV+0),2), comment="c.real += a.real * b.real"))
        module.add(VFmaF64(dst=vgpr("ValuC+%u"%(newSumIdxV+0),2), src0=vgpr(dataV+2,2), src1=sgpr("Beta+2",2).getMinus(), src2=vgpr("ValuC+%u"%(newSumIdxV+0),2), comment="c.real -= a.imag * b.imag"))
        module.add(VFmaF64(dst=vgpr("ValuC+%u"%(newSumIdxV+2),2), src0=vgpr(dataV+0,2), src1=sgpr("Beta+2",2), src2=vgpr("ValuC+%u"%(newSumIdxV+2),2), comment="c.imag += a.real * b.imag"))
        module.add(VFmaF64(dst=vgpr("ValuC+%u"%(newSumIdxV+2),2), src0=vgpr(dataV+2,2), src1=sgpr("Beta+0",2), src2=vgpr("ValuC+%u"%(newSumIdxV+2),2), comment="c.imag += a.imag * b.real"))

      # float8 precision
      elif kernel["ProblemType"]["DestDataType"].isAnyFloat8():
        if kernel["ProblemType"]["HighPrecisionAccumulate"]:
          dataCExternal   = ss.elementData[elementIdx] + (vi // 4)
          newSumIdxV = sumIdxV - self.parentWriter.states.c.startVgprValu
          # Generate single f32 code if edge is detected.
          isPK = False
          if ((vi + 1) == self.gwvw) and ((self.gwvw % 2) == 1):
            if self.parentWriter.states.archCaps["VOP3ByteSel"]:
              sb = 0 if self.gwvw == 1 else 1
              if not self.amdClangVersion.major >= 19:
                module.add(VCvtFP8toF32(dst=vgpr(tmpVgpr), src=vgpr(dataCExternal), vop3=VOP3PModifiers(op_sel=[0,sb])))
              else:
                module.add(VCvtFP8toF32(dst=vgpr(tmpVgpr), src=vgpr(dataCExternal), vop3=VOP3PModifiers(byte_sel=[sb])))
            else:
              sb = SelectBit.BYTE_0 if self.gwvw == 1 else SelectBit.BYTE_2
              module.add(VCvtFP8toF32(dst=vgpr(tmpVgpr), src=vgpr(dataCExternal), sdwa=SDWAModifiers(src0_sel=sb)))
          # Original packed route
          elif vi%2 == 1:
            continue
          else:
            isPK = True
            module.add(ECvtPkFP8toF32(dst=vgpr(tmpVgpr, 2), src=vgpr(dataCExternal), sel=HighBitSel.LOW if vi%4 == 0 else HighBitSel.HIGH))
          module.add(SNop(waitState=0))
          if kernel["ProblemType"]["ComputeDataType"].isSingle():
            module.add(VMacF32(dst=vgpr("ValuC+%u"%newSumIdxV), src0=vgpr(tmpVgpr), src1=sgpr("Beta"), comment="finalSum = sum*alpha + C*beta"))
            if isPK:
              module.add(VMacF32(dst=vgpr("ValuC+%u"%(newSumIdxV+1)), src0=vgpr(tmpVgpr+1), src1=sgpr("Beta"), comment="finalSum = sum*alpha + C*beta (PK)"))
      # bfloat8 precision
      elif kernel["ProblemType"]["DestDataType"].isAnyBFloat8():
        if kernel["ProblemType"]["HighPrecisionAccumulate"]:
          dataCExternal   = ss.elementData[elementIdx] + (vi // 4)
          newSumIdxV = sumIdxV - self.parentWriter.states.c.startVgprValu
          # Generate single f32 code if edge is detected.
          isPK = False
          if ((vi + 1) == self.gwvw) and ((self.gwvw % 2) == 1):
            if self.parentWriter.states.archCaps["VOP3ByteSel"]:
              sb = 0 if self.gwvw == 1 else 1
              if not self.amdClangVersion.major >= 19:
                module.add(VCvtBF8toF32(dst=vgpr(tmpVgpr), src=vgpr(dataCExternal), vop3=VOP3PModifiers(op_sel=[0,sb])))
              else:
                module.add(VCvtBF8toF32(dst=vgpr(tmpVgpr), src=vgpr(dataCExternal), vop3=VOP3PModifiers(byte_sel=[sb])))
            else:
              sb = SelectBit.BYTE_0 if self.gwvw == 1 else SelectBit.BYTE_2
              module.add(VCvtBF8toF32(dst=vgpr(tmpVgpr), src=vgpr(dataCExternal), sdwa=SDWAModifiers(src0_sel=sb)))
          # Original packed route
          elif vi%2 == 1:
            continue
          else:
            isPK = True
            module.add(ECvtPkBF8toF32(dst=vgpr(tmpVgpr, 2), src=vgpr(dataCExternal), sel=HighBitSel.LOW if vi%4 == 0 else HighBitSel.HIGH))
          module.add(SNop(waitState=0))
          if kernel["ProblemType"]["ComputeDataType"].isSingle():
            module.add(VMacF32(dst=vgpr("ValuC+%u"%newSumIdxV), src0=vgpr(tmpVgpr), src1=sgpr("Beta"), comment="finalSum = sum*alpha + C*beta"))
            if isPK:
              module.add(VMacF32(dst=vgpr("ValuC+%u"%(newSumIdxV+1)), src0=vgpr(tmpVgpr+1), src1=sgpr("Beta"), comment="finalSum = sum*alpha + C*beta (PK)"))
    return module

def copyData(computeDataType, elementSumIdx, gwvw, vgprStart, direction=0):
  module = Module("Copy Data")
  vi = 0
  while vi < gwvw:
    sumIdxV = elementSumIdx + vi
    if computeDataType.isHalf() or computeDataType.isBFloat16():
      if (sumIdxV % 2 != 0):
        vi += 1
        continue
      vgprIdx = elementSumIdx + vi // 2
      if (vi + 1 < gwvw) and ((vgprStart + (vi // 2)) % 2 == 0) and (vgprIdx % 2 == 0):
        module.add(VMovB64(dst=vgpr(vgprStart + (vi // 2), 2), src=vgpr(vgprIdx, 2)))
        vi += 2
      else:
        module.add(VMovB32(dst=vgpr(vgprStart + (vi // 2)), src=vgpr(vgprIdx)))
        vi += 1
    elif computeDataType.isSingle() or computeDataType.isInt32():
      vgprIdx = sumIdxV
      if (vi + 1 < gwvw) and ((vgprStart + vi) % 2 == 0) and (vgprIdx % 2 == 0):
        module.add(VMovB64(dst=vgpr(vgprStart + vi, 2), src=vgpr(vgprIdx, 2)))
        vi += 2
      else:
        module.add(VMovB32(dst=vgpr(vgprStart + vi), src=vgpr(vgprIdx)))
        vi += 1
    elif computeDataType.isDouble():
      vgprIdx = elementSumIdx + vi * 2
      module.add(VMovB64(dst=vgpr(vgprStart + vi * 2, 2), src=vgpr(vgprIdx, 2)))
      vi += 1
    else:
      assert 0

  if direction == 1:
    for i in module.items():
      srcs = i.srcs
      tmp = srcs[0]
      srcs[0] = i.dst
      i.dst = tmp
      i.srcs = srcs
  return module

def convertData(gwvw, elementSumIdx, cvtType: CvtType, roundType: RoundType = RoundType.ROUND_UP, inputPrefix="", prefixOffset=0):
  module = Module("ConvertData")
  for vi in range(0, gwvw):
    sumIdxV = elementSumIdx + vi
    formatVgpr = formatting(sumIdxV, inputPrefix, prefixOffset)
    if cvtType == CvtType.CVT_F32_to_I32:
        if roundType == RoundType.ROUND_TO_NEAREST_EVEN:
          module.add(VRndneF32(dst=vgpr(formatVgpr), src=vgpr(formatVgpr), comment=" round to even"))
        module.add(VCvtF32toI32(dst=vgpr(formatVgpr), src=vgpr(formatVgpr), comment=" convert fp32 to i32"))
    elif cvtType == CvtType.CVT_I32_to_F32:
        module.add(VCvtI32toF32(dst=vgpr(formatVgpr), src=vgpr(formatVgpr), comment=" convert to fp32"))
    else:
      #TODO add other convert types here.
      assert 0
  return module