[llvm] [AMDGPU] Pipeline DS_READ_B128 with F16 MFMA on gfx950 (PR #212330)

[Juan Manuel Martinez Caamaño via llvm-commits]

================
@@ -240,6 +288,553 @@ static bool canUsePressureDiffs(const SUnit &SU) {
   return true;
 }
 
+static bool isTargetDSReadOpcode(unsigned Opcode) {
+  switch (Opcode) {
+  case AMDGPU::DS_READ_B128:
+  case AMDGPU::DS_READ_B128_gfx9:
+    return true;
+  default:
+    return false;
+  }
+}
+
+static bool isTargetMFMAOpcode(unsigned Opcode) {
+  switch (Opcode) {
+  case AMDGPU::V_MFMA_F32_32X32X16_F16_mac_vgprcd_e64:
+    return true;
+  default:
+    return false;
+  }
+}
+
+template <typename PredicateT>
+static bool hasBundledMI(const MachineInstr *MI, PredicateT Predicate) {
+  if (!MI)
+    return false;
+
+  if (Predicate(*MI))
+    return true;
+
+  if (!MI->isBundle())
+    return false;
+
+  MachineBasicBlock::const_instr_iterator BundleI = MI->getIterator();
+  for (++BundleI;
+       BundleI != MI->getParent()->instr_end() && BundleI->isBundledWithPred();
+       ++BundleI) {
+    if (Predicate(*BundleI))
+      return true;
+  }
+
+  return false;
+}
+
+static bool isDSReadLike(const SUnit *SU) {
+  if (!SU)
+    return false;
+
+  return hasBundledMI(SU->getInstr(), [](const MachineInstr &MI) {
+    return isTargetDSReadOpcode(MI.getOpcode());
+  });
+}
+
+static bool isMFMALike(const SUnit *SU) {
+  if (!SU)
+    return false;
+
+  return hasBundledMI(SU->getInstr(), [](const MachineInstr &MI) {
+    return isTargetMFMAOpcode(MI.getOpcode());
+  });
+}
+
+static bool hasMFMAFragmentPipeline(const ScheduleDAGMI &DAG) {
+  for (const SUnit &SU : DAG.SUnits) {
+    if (!isDSReadLike(&SU))
+      continue;
+    for (const SDep &Succ : SU.Succs)
+      if (Succ.getKind() == SDep::Data && isMFMALike(Succ.getSUnit()))
+        return true;
+  }
+  return false;
+}
+
+static bool isSafeMFMAFragmentRampUpFiller(const SUnit *SU) {
+  if (!SU || isDSReadLike(SU) || isMFMALike(SU))
+    return false;
+
+  // This tune does not model buffered-load maturity windows or available
+  // in-flight memory capacity. Keep memory operations, calls, barriers, and
+  // other side effects under the normal scheduler policy.
+  return !hasBundledMI(SU->getInstr(), [](const MachineInstr &MI) {
+    return MI.mayLoadOrStore() || MI.isCall() || MI.isTerminator() ||
+           MI.isBarrier() || MI.hasUnmodeledSideEffects();
+  });
+}
+
+static unsigned getBoundaryReadyCycle(const SchedBoundary &Zone,
+                                      const SUnit *SU) {
+  return Zone.isTop() ? SU->TopReadyCycle : SU->BotReadyCycle;
+}
+
+static unsigned getReadyStall(const SchedBoundary &Zone, const SUnit *SU) {
+  unsigned ReadyCycle = getBoundaryReadyCycle(Zone, SU);
+  unsigned CurrCycle = Zone.getCurrCycle();
+  return ReadyCycle > CurrCycle ? ReadyCycle - CurrCycle : 0;
+}
+
+static unsigned getIssueStall(SchedBoundary &Zone, SUnit *SU) {
+  if (!SU || !SU->hasReservedResource || !Zone.SchedModel ||
+      !Zone.SchedModel->hasInstrSchedModel())
+    return 0;
+
+  const MCSchedClassDesc *SC = Zone.DAG->getSchedClass(SU);
+  unsigned Stall = 0;
+  for (const MCWriteProcResEntry &PE :
+       make_range(Zone.SchedModel->getWriteProcResBegin(SC),
+                  Zone.SchedModel->getWriteProcResEnd(SC))) {
+    unsigned NextCycle;
+    unsigned InstanceIdx;
+    std::tie(NextCycle, InstanceIdx) = Zone.getNextResourceCycle(
+        SC, PE.ProcResourceIdx, PE.ReleaseAtCycle, PE.AcquireAtCycle);
+    (void)InstanceIdx;
+    if (NextCycle > Zone.getCurrCycle())
+      Stall = std::max(Stall, NextCycle - Zone.getCurrCycle());
+  }
+
+  return Stall;
+}
+
+static unsigned getEffectiveStall(SchedBoundary &Zone, SUnit *SU) {
+  return std::max(getReadyStall(Zone, SU), getIssueStall(Zone, SU));
+}
+
+static bool consumesThisDSRead(const SUnit *SU, const SUnit *Pred) {
+  if (!isMFMALike(SU) || !isDSReadLike(Pred))
+    return false;
+
+  for (const SDep &PredDep : SU->Preds) {
+    if (PredDep.getSUnit() == Pred)
+      return true;
+  }
+
+  return false;
+}
+
+static bool hasScheduledMFMASuccessor(const SUnit *SU) {
+  if (!isDSReadLike(SU))
+    return false;
+
+  for (const SDep &SuccDep : SU->Succs) {
+    const SUnit *Succ = SuccDep.getSUnit();
+    if (isMFMALike(Succ) && Succ->isScheduled)
+      return true;
+  }
+
+  return false;
+}
+
+enum class PipeKind {
+  // Not part of the DS_READ/MFMA fragment pipeline.
+  None,
+
+  // A DS_READ that may produce an operand fragment for an MFMA.
+  FragmentProducer,
+
+  // An MFMA with no remaining scheduler-model dependency stall.
+  ReadyMFMA,
+
+  // An MFMA whose scheduler-model dependency-ready cycle is still ahead of
+  // the current scheduling boundary.
+  PendingMFMA,
+};
+
+static PipeKind classifyPipeKind(SchedBoundary &Zone, SUnit *SU) {
+  if (isDSReadLike(SU))
+    return PipeKind::FragmentProducer;
+
+  if (isMFMALike(SU))
+    return getReadyStall(Zone, SU) ? PipeKind::PendingMFMA
+                                   : PipeKind::ReadyMFMA;
+
+  return PipeKind::None;
+}
+
+static unsigned getNumDSReadsToUnlockClosestMFMA(const SUnit *SU) {
+  // For every unscheduled MFMA successor, count its other unscheduled DS_READ
+  // predecessors. Return the minimum count, excluding SU itself, to prefer the
+  // DS_READ that can unlock any MFMA with the fewest additional reads.
+  if (!isDSReadLike(SU))
+    return std::numeric_limits<unsigned>::max();
+
+  unsigned MinRemaining = std::numeric_limits<unsigned>::max();
+  for (const SDep &SuccDep : SU->Succs) {
+    const SUnit *Succ = SuccDep.getSUnit();
+    if (!isMFMALike(Succ) || Succ->isScheduled)
+      continue;
+
+    unsigned Remaining = 0;
+    for (const SDep &PredDep : Succ->Preds) {
+      const SUnit *Pred = PredDep.getSUnit();
+      if (!isDSReadLike(Pred) || Pred == SU || Pred->isScheduled)
+        continue;
+      ++Remaining;
+    }
----------------
jmmartinez wrote:

Isn't this loop `countUnscheduledDSReadPredsToMFMASuccessorAfter` ?

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