[llvm] [AMDGPU] Hoist WMMA coexecution hazard V_NOPs from loops to preheaders (PR #176895)

[Dark Steve via llvm-commits]

================
@@ -2190,6 +2159,161 @@ int GCNHazardRecognizer::checkWMMACoexecutionHazards(MachineInstr *MI) {
   return WaitStatesNeeded;
 }
 
+void GCNHazardRecognizer::insertVnopsBeforeTerminator(MachineBasicBlock *MBB,
+                                                      int Count) {
+  MachineBasicBlock::iterator InsertPt = MBB->getFirstTerminator();
+  const DebugLoc &DL =
+      InsertPt != MBB->end() ? InsertPt->getDebugLoc() : DebugLoc();
+
+  for (int i = 0; i < Count; ++i) {
+    BuildMI(*MBB, InsertPt, DL, TII.get(AMDGPU::V_NOP_e32));
+  }
+}
+
+bool GCNHazardRecognizer::hasWMMAToWMMARegOverlap(
+    const MachineInstr &WMMA, const MachineInstr &MI) const {
+  Register D0 = TII.getNamedOperand(WMMA, AMDGPU::OpName::vdst)->getReg();
+  Register A1 = TII.getNamedOperand(MI, AMDGPU::OpName::src0)->getReg();
+  Register B1 = TII.getNamedOperand(MI, AMDGPU::OpName::src1)->getReg();
+
+  // WMMA0 writes (D0), WMMA1 reads (A1/B1/Idx1).
+  if (TRI.regsOverlap(D0, A1) || TRI.regsOverlap(D0, B1))
+    return true;
+
+  if (SIInstrInfo::isSWMMAC(MI)) {
+    Register Idx1 = TII.getNamedOperand(MI, AMDGPU::OpName::src2)->getReg();
+    if (TRI.regsOverlap(D0, Idx1))
+      return true;
+  }
+  return false;
+}
+
+bool GCNHazardRecognizer::hasWMMAToVALURegOverlap(
+    const MachineInstr &WMMA, const MachineInstr &MI) const {
+  // WMMA writes, VALU reads.
+  Register D0 = TII.getNamedOperand(WMMA, AMDGPU::OpName::vdst)->getReg();
+  for (const MachineOperand &ValuUse : MI.explicit_uses()) {
+    if (ValuUse.isReg() && TRI.regsOverlap(D0, ValuUse.getReg()))
+      return true;
+  }
+
+  auto *ValuDst = TII.getNamedOperand(MI, AMDGPU::OpName::vdst);
+  if (!ValuDst || !ValuDst->isReg())
+    return false;
+  Register D1 = ValuDst->getReg();
+
+  // WMMA writes, VALU writes.
+  if (TRI.regsOverlap(D0, D1))
+    return true;
+
+  // WMMA reads, VALU writes.
+  Register A0 = TII.getNamedOperand(WMMA, AMDGPU::OpName::src0)->getReg();
+  Register B0 = TII.getNamedOperand(WMMA, AMDGPU::OpName::src1)->getReg();
+  if (TRI.regsOverlap(A0, D1) || TRI.regsOverlap(B0, D1))
+    return true;
+
+  if (SIInstrInfo::isSWMMAC(WMMA)) {
+    Register Idx0 = TII.getNamedOperand(WMMA, AMDGPU::OpName::src2)->getReg();
+    if (TRI.regsOverlap(D1, Idx0))
+      return true;
+  }
+  return false;
+}
+
+bool GCNHazardRecognizer::isCoexecutionHazardFor(const MachineInstr &I,
+                                                 const MachineInstr &MI) const {
+  if (!TII.isXDLWMMA(I))
+    return false;
+
+  // Dispatch based on MI type
+  if (TII.isXDLWMMA(MI))
+    return hasWMMAToWMMARegOverlap(I, MI);
+  else if (isCoexecutableVALUInst(MI))
+    return hasWMMAToVALURegOverlap(I, MI);
+
+  return false;
+}
+
+bool GCNHazardRecognizer::hasWMMAHazardInLoop(MachineLoop *L, MachineInstr *MI,
+                                              bool IncludeSubloops) {
+  // Scan loop for any WMMA that hazards MI.
+  // TODO: Avoid full loop scan when WMMA is beyond VALU distance.
+  for (MachineBasicBlock *MBB : L->getBlocks()) {
+    if (!IncludeSubloops && MLI->getLoopFor(MBB) != L)
+      continue;
+    for (MachineInstr &I : *MBB) {
+      if (&I == MI)
+        continue;
+      if (isCoexecutionHazardFor(I, *MI))
+        return true;
+    }
+  }
+  return false;
+}
+
+void GCNHazardRecognizer::ensureLoopInfoAvailable() {
+  // Lazily compute MDT and MLI only when needed
+  if (MLI)
+    return;
+
+  OwnedMDT =
+      std::make_unique<MachineDominatorTree>(const_cast<MachineFunction &>(MF));
+  OwnedMLI = std::make_unique<MachineLoopInfo>();
+  OwnedMLI->analyze(*OwnedMDT);
+
+  MLI = OwnedMLI.get();
+}
+
+bool GCNHazardRecognizer::tryHoistWMMAVnopsFromLoop(MachineInstr *MI,
+                                                    int WaitStatesNeeded) {
+  ensureLoopInfoAvailable();
+
+  MachineLoop *L = MLI->getLoopFor(MI->getParent());
+  if (!L) {
+    ++NumWMMAHoistingBailed;
+    return false;
+  }
+
+  // If innermost loop has WMMA hazard, we can't hoist at all
+  if (hasWMMAHazardInLoop(L, MI)) {
+    ++NumWMMAHoistingBailed;
+    return false;
+  }
+
+  // Find outermost loop with no internal hazard
+  MachineLoop *TargetLoop = L;
+  while (MachineLoop *Parent = TargetLoop->getParentLoop()) {
+    if (hasWMMAHazardInLoop(Parent, MI, false))
+      break;             // Parent has hazard in its own blocks, stop here
+    TargetLoop = Parent; // Safe to hoist further out
+  }
+
+  // Need valid preheader to insert V_NOPs
+  MachineBasicBlock *Preheader = TargetLoop->getLoopPreheader();
+  if (!Preheader) {
+    ++NumWMMAHoistingBailed;
+    return false;
+  }
+
+  LLVM_DEBUG(dbgs() << "WMMA V_NOP Hoisting: Moving " << WaitStatesNeeded
+                    << " V_NOPs from loop to " << Preheader->getName() << "\n");
+
+  insertVnopsBeforeTerminator(Preheader, WaitStatesNeeded);
+  NumWMMANopsHoisted += WaitStatesNeeded;
+  return true;
+}
+
+bool GCNHazardRecognizer::fixWMMACoexecutionHazards(MachineInstr *MI) {
----------------
PrasoonMishra wrote:

This is a real hazard (SPG 4.6.12.1. "Requirements for WMMA data hazards")

https://github.com/llvm/llvm-project/pull/176895