perf(hipblaslt): Split subtile cluster barrier signal/wait (gfx1250) (ROCm#8805)

ISSUE ID : AIHPBLAS-3984

### Motivation

The subtile cluster barrier exposed its cross-CU latency as a stall, and
loop-control / cluster branches were issued without regard to the
surrounding WMMAs, costing branch latency in the mainloop on gfx1250.

### Technical Details

Added features:

- Split the cluster barrier into separate signal and wait halves 
- Wave-0 election replacing the prior isfirst-barrier scheme; only wave
0 branches into the cluster signal.
- Post-schedule splicing (insertClusterBarrier) — the signal is spliced
right after the mainloop's existing workgroup barrier (reusing that sync
rather than emitting a second one);
- Branch-after-WMMA placement to hide branching latency:
- 16-byte alignment on branch targets (LoopBeginL, skipCBPreSignal).
- gfx1250 gating via HasWmmaArbStallBit, with an assert that
ClusterBarrier is gfx1250-only.

### Test Result

- test_SubtileBasedLogicalScheduler.py — 131 passed.
- The three subtile bf16 gfx1250 gemm common tests pass:
  - Tensile/Tests/common/gemm/gfx12/subtile_bf16_gfx1250.yaml
  - Tensile/Tests/common/gemm/gfx12/subtile_bf16_gfx1250_dp.yaml
  - Tensile/Tests/common/gemm/gfx12/subtile_bf16_gfx1250_cluster.yaml

Author: sebvince

projects/hipblaslt/tensilelite/Tensile/Components/Subtile/ClusterBarrier.py

@@ -3,33 +3,130 @@
 
 """Cluster-scope barrier handshake for the subtile mainloop.
 
-Subtile-specific equivalent of the StinkyTofu InsertClusterBarrierPass (which
-does not run at OptLevel 0 / ScheduleIterAlg=3). Free functions take the writer
-explicitly to keep cluster logic in Subtile/ rather than KernelWriterAssembly.
+The handshake is split into a *signal* half and a *wait* half so the wait can be
+moved away from the signal, hiding the cluster barrier's cross-CU latency behind
+the WMMAs that issue in between instead of exposing it as a stall.
 """
 
 from __future__ import annotations
 
-from rocisa.code import Label, Module, TextBlock
-from rocisa.instruction import SBarrier, SBarrierSignalIsFirst, SCBranchSCC0
+from rocisa.code import Label, Module
+from rocisa.container import sgpr
+from rocisa.instruction import (SBarrier, SCBranchSCC0, SCmpEQU32,
+                                MFMAInstruction, MXMFMAInstruction)
 
+_isWgBarrier = lambda x: isinstance(x, SBarrier) and "s_barrier_wait -1" in str(x)
 
-def subtileClusterBarrier(writer, kernel, label="") -> Module:
-    # Balanced workgroup barrier + cluster signal/wait -3 per occurrence.
-    mod = Module("subtile_cluster_barrier")
-    # Workgroup barrier via isfirst: the first wave to arrive gets SCC=1 and so
-    # is the single wave that signals the cluster barrier (one arrival per WG).
-    skipPreSignal = Label(writer.labels.getUniqueNamePrefix("skipCBPreSignal"), "")
-    mod.add(SBarrierSignalIsFirst(False, "workgroup barrier signal (isfirst)"))
-    mod.add(SBarrier(True, True, False, "workgroup barrier wait"))
-    mod.add(SCBranchSCC0(skipPreSignal.getLabelName(), "only the first-arriving wave signals the cluster"))
+
+def _findNextMFMA(items, start):
+    """Index of the first MFMA at/after ``start``, or ``None`` if none follows."""
+    for j in range(start, len(items)):
+        if isinstance(items[j], (MFMAInstruction, MXMFMAInstruction)):
+            return j
+    return None
+
+
+def subtileClusterBarrierSignal(writer, kernel) -> Module:
+    """Wave-0-only cluster_barrier signal.
+
+    Wave 0 alone issues the cluster_barrier signal; all other waves branch over
+    it. Ends at the ``skipPreSignal`` label so all waves fall through to whatever
+    work follows; the matching wait is emitted later by ``subtileClusterBarrierWait``.
+    """
+    mod = Module("subtile_cluster_barrier_signal")
+    skipPreSignal = Label(writer.labels.getUniqueNamePrefix("skipCBPreSignal"), "", 16)
+    # Elect wave 0 to issue the single cluster_barrier signal.
+    mod.add(SCmpEQU32(sgpr("WaveIdx"), 0, "wave 0?"))
+    mod.add(SCBranchSCC0(skipPreSignal.getLabelName(), "only wave 0 signals the cluster"))
     mod.add(SBarrier(True, False, True, "cluster_barrier signal"))
     mod.add(skipPreSignal)
+    return mod
+
+
+def subtileClusterBarrierWait(writer, kernel) -> Module:
+    """The all-waves cluster_barrier wait that closes the handshake."""
+    mod = Module("subtile_cluster_barrier_wait")
     mod.add(SBarrier(True, True, True, "cluster_barrier wait"))
-    # SQTT trace marker tagging which loop section this barrier belongs to.
-    base = {"PRELOOP": 0x00, "MAINLOOP": 0x10, "NGLL": 0x20,
-            "NLL": 0x30, "TAILLOOP": 0x40}
-    kind, _, suffix = label.partition("_C")
-    loopId = (base.get(kind, 0xf0) + (int(suffix) if suffix.isdigit() else 0)) & 0xff
-    mod.add(TextBlock(f"s_ttracedata_imm {0xc100 | loopId:#06x}\n"))
     return mod
+
+
+def insertClusterBarrier(module, writer, kernel):
+    """Splice the cluster-scope barrier handshake into the post-schedule order.
+
+    No-op unless ``ClusterBarrier`` is enabled. The signal is spliced in right
+    after the mainloop's existing workgroup barrier (reusing that sync instead of
+    emitting a second one); the wait is appended at the end of the section, so the
+    barrier's cross-CU latency overlaps the whole macro tile's WMMAs before the
+    handshake is closed.
+
+    If no workgroup barrier is found in this section, the signal is prepended at
+    the start so the handshake is still opened (correctness over reuse).
+
+    Returns a rebuilt Module; the input is left untouched.
+    """
+    if not kernel.get("ClusterBarrier"):
+        return module
+
+    signalItems = subtileClusterBarrierSignal(writer, kernel).flatitems()
+    waitItems = subtileClusterBarrierWait(writer, kernel).flatitems()
+
+    # ClusterBarrier is only supported on gfx1250.
+    assert writer.states.asmCaps.get("HasClusterBarrier", False), \
+        "ClusterBarrier requires the HasClusterBarrier asm capability"
+
+    # Place the wave-0-election branch right after a WMMA to hide branching
+    # latency: keep s_cmp before the next scheduled MFMA and emit the branch
+    # after it.
+
+    items = module.flatitems()
+    result = Module(module.name)
+    done = False
+    skip = set()
+    for i, inst in enumerate(items):
+        if i in skip:
+            continue
+        result.add(inst)
+        if not done and _isWgBarrier(inst):
+            done = True
+            mfmaIdx = _findNextMFMA(items, i + 1)
+            if mfmaIdx is None:
+                # No following MFMA to pin the branch to: emit the block intact
+                # (best-effort).
+                for s in signalItems:
+                    result.add(s)
+            else:
+                # Split the signal block at the wave-0 election branch. The
+                # block is authored with exactly one conditional branch; assert
+                # it so a future change that adds another fails loudly here.
+                brIdxs = [k for k, s in enumerate(signalItems)
+                          if isinstance(s, SCBranchSCC0)]
+                assert len(brIdxs) == 1, \
+                    "signal block must contain exactly one wave-0 election branch"
+                brIdx = brIdxs[0]
+                pre, post = signalItems[:brIdx], signalItems[brIdx:]
+                # Everything up to the MFMA (incl. its s_set_vgpr_msb primer)
+                # keeps its order, then s_cmp, the MFMA, and the branch. SCC
+                # survives the MFMA and vgpr-msb is a persistent mode, so the
+                # intervening compare disturbs neither.
+                for k in range(i + 1, mfmaIdx):
+                    result.add(items[k])
+                    skip.add(k)
+                for s in pre:
+                    result.add(s)
+                result.add(items[mfmaIdx])
+                skip.add(mfmaIdx)
+                for s in post:
+                    result.add(s)
+    if not done:  # no workgroup barrier: open the handshake at the start
+        head = Module(module.name)
+        for s in signalItems:
+            head.add(s)
+        for inst in result.flatitems():
+            head.add(inst)
+        result = head
+
+    # Wait: append at the end of the section so cluster latency hides behind the
+    # whole macro tile's WMMAs before the handshake is closed.
+    for w in waitItems:
+        result.add(w)
+    return result

projects/hipblaslt/tensilelite/Tensile/Components/Subtile/LogicalScheduler.py

@@ -62,6 +62,62 @@ def _checkin_tile(tile):
     tile.regList.pool.checkIn(tile.regList.indices[0])
 
 
+def _hoistSetupBeforeLastMFMA(loopBody, setupInsts):
+    """Move loop-control setup (dec + compare) above the body's last MFMA.
+
+    Place the loop-control branch immediately after the last MFMA to hide
+    branching latency behind it. The decrement+compare that feed the branch's
+    SCC are normally emitted right after the body (MFMA; s_sub; s_cmp;
+    s_cbranch); hoisting them above the final MFMA produces (s_sub; s_cmp;
+    MFMA; s_cbranch) so the branch directly follows the MFMA.
+
+    The hoist is dependency-safe: SCC set by s_cmp survives the MFMA (which
+    never touches SCC), and the MFMA does not read LoopCounterL.
+
+    Falls back to appending the setup when the body has no MFMA (e.g. a skip
+    branch with no preceding compute), leaving behavior unchanged there.
+    Returns a rebuilt Module; the input is left untouched.
+    """
+    from rocisa.instruction import MFMAInstruction, MXMFMAInstruction
+
+    items = loopBody.flatitems()
+    lastMfma = -1
+    for idx, it in enumerate(items):
+        if isinstance(it, (MFMAInstruction, MXMFMAInstruction)):
+            lastMfma = idx
+
+    rebuilt = Module(loopBody.name)
+    if lastMfma < 0:
+        for it in items:
+            rebuilt.add(it)
+        for s in setupInsts:
+            rebuilt.add(s)
+        return rebuilt
+
+    for idx, it in enumerate(items):
+        if idx == lastMfma:
+            for s in setupInsts:
+                rebuilt.add(s)
+        rebuilt.add(it)
+    return rebuilt
+
+
+def _emitBodySetupBranch(module, loopBody, setupInsts, branch, hoist):
+    """Append a loop body, its dec+compare setup, and the trailing control branch.
+
+    When ``hoist`` is set, the setup is moved above the body's last MFMA (via
+    _hoistSetupBeforeLastMFMA) so the branch lands right after that MFMA to hide
+    branching latency. Otherwise the setup and branch simply follow the body.
+    """
+    if hoist:
+        module.add(_hoistSetupBeforeLastMFMA(loopBody, setupInsts))
+    else:
+        module.add(loopBody)
+        for s in setupInsts:
+            module.add(s)
+    module.add(branch)
+
+
 class Pass(IntEnum):
     """Scheduler passes in dependency order.
 
@@ -3315,10 +3371,6 @@ def _emitLoop(self, writer, kernel, label, emitted_3d, schedule=True):
 
         module = Module(label)
         module.addComment0(f"{label} start")
-        if kernel.get("ClusterBarrier"):
-            from Tensile.Components.Subtile.ClusterBarrier import subtileClusterBarrier
-            for inst in subtileClusterBarrier(writer, kernel, label=label).flatitems():
-                module.add(inst)
         use_pap_preloop_skip = (
             label == "PRELOOP"
             and kernel.get("UseSubtileImpl")
@@ -3372,6 +3424,11 @@ def _emitLoop(self, writer, kernel, label, emitted_3d, schedule=True):
         # them (swap hoisted ahead, dscnt drain between), so no WAR can form.
         if self.config.pgr == 0 and label.startswith("MAINLOOP"):
             module = insertLRSwapWarWaitAlu(module, writer, kernel)
+        # Cluster barrier: splice both halves against the final post-schedule order.
+        # Signal goes right after the mainloop's existing workgroup barrier (reusing
+        # that sync); the wait is appended at the end to hide its cross-CU latency.
+        from Tensile.Components.Subtile.ClusterBarrier import insertClusterBarrier
+        module = insertClusterBarrier(module, writer, kernel)
         return module
 
     def _emit_pgr2_tail_lw_align(self, kernel):
@@ -3433,6 +3490,10 @@ def emitMainAndExitLoops(self, writer, kernel, tensorParametersA=None, tensorPar
         module = Module("MainAndExitLoops")
         uf = self.unroll_factor
 
+        # gfx1250 requires a loop-back/exit s_cbranch to be alone/first after an
+        # MFMA, so hoist the per-copy dec+compare above the body's last MFMA.
+        isGfx1250 = writer.states.archCaps.get("HasWmmaArbStallBit", False)
+
         def make_subtile_pap_module(skip_barrier=False):
             if (kernel.get("UseSubtileImpl")
                 and kernel.get("PrefetchAcrossPersistent")
@@ -3502,7 +3563,7 @@ def insert_after_drain(em_list):
 
         # ── Mainloop ──
         module.addComment0("MAINLOOP")
-        loopBegin = Label("LoopBeginL", "")
+        loopBegin = Label("LoopBeginL", "", alignment=16)
 
         exitValue = self.config.pgr
 
@@ -3524,21 +3585,25 @@ def insert_after_drain(em_list):
                 module.add(_SMovB32(dst=mgpr(0), src=sgpr("LoopCounterL"),
                                     comment="trace: M0 = LoopCounterL"))
                 module.add(_STtraceData(comment="trace: emit M0 to SQTT"))
-            module.add(self._emitLoop(writer, kernel, f"MAINLOOP_C{ui}",
-                                      self._emitted_per_unroll[ui]))
-            module.add(SSubU32(dst=sgpr("LoopCounterL"),
-                               src0=sgpr("LoopCounterL"), src1=1,
-                               comment=f"dec counterL (copy {ui})"))
-            module.add(SCmpEQU32(src0=sgpr("LoopCounterL"), src1=exitValue,
-                                 comment=f"counterL == {exitValue}? (copy {ui} exit)"))
+            loopBody = self._emitLoop(writer, kernel, f"MAINLOOP_C{ui}",
+                                      self._emitted_per_unroll[ui])
+            # dec + compare that produce the SCC the loop-control branch reads.
+            setupInsts = [
+                SSubU32(dst=sgpr("LoopCounterL"),
+                        src0=sgpr("LoopCounterL"), src1=1,
+                        comment=f"dec counterL (copy {ui})"),
+                SCmpEQU32(src0=sgpr("LoopCounterL"), src1=exitValue,
+                          comment=f"counterL == {exitValue}? (copy {ui} exit)"),
+            ]
             if ui < uf - 1:
-                module.add(SCBranchSCC1(
+                branch = SCBranchSCC1(
                     labelName=exitLabels[ui].getLabelName(),
-                    comment=f"copy {ui} exit → NGLL_C{ui}"))
+                    comment=f"copy {ui} exit → NGLL_C{ui}")
             else:
-                module.add(SCBranchSCC0(
+                branch = SCBranchSCC0(
                     labelName=loopBegin.getLabelName(),
-                    comment="restart mainloop"))
+                    comment="restart mainloop")
+            _emitBodySetupBranch(module, loopBody, setupInsts, branch, hoist=isGfx1250)
 
         # ── NGLL + NLL exit paths ──
         hasNGLL = self.config.pgr >= 2