[AMDGPU] Fall back to SavedMutations when not applying IGLP

jrbyrnes · jrbyrnes · commit 06bf083312b2 · 2024-02-19T11:58:27.000-08:00
Change-Id: I4ef67afa5e68714231a01afb1e9f58dbddd07368
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUIGroupLP.cpp b/llvm/lib/Target/AMDGPU/AMDGPUIGroupLP.cpp
@@ -849,7 +849,7 @@ class IGLPStrategy {
 
 public:
   /// Add SchedGroups to \p SyncedSchedGroups to implement this Strategy.
-  virtual void applyIGLPStrategy(
+  virtual bool applyIGLPStrategy(
       DenseMap<int, SUnitsToCandidateSGsMap> &SyncedInstrs,
       DenseMap<int, SmallVector<SchedGroup, 4>> &SyncedSchedGroups,
       IGLPPhase Phase) = 0;
@@ -868,7 +868,7 @@ class IGLPStrategy {
 class MFMASmallGemmOpt final : public IGLPStrategy {
 private:
 public:
-  void applyIGLPStrategy(
+  bool applyIGLPStrategy(
       DenseMap<int, SUnitsToCandidateSGsMap> &SyncedInstrs,
       DenseMap<int, SmallVector<SchedGroup, 4>> &SyncedSchedGroups,
       IGLPPhase Phase) override;
@@ -881,7 +881,7 @@ class MFMASmallGemmOpt final : public IGLPStrategy {
   }
 };
 
-void MFMASmallGemmOpt::applyIGLPStrategy(
+bool MFMASmallGemmOpt::applyIGLPStrategy(
     DenseMap<int, SUnitsToCandidateSGsMap> &SyncedInstrs,
     DenseMap<int, SmallVector<SchedGroup, 4>> &SyncedSchedGroups,
     IGLPPhase Phase) {
@@ -902,6 +902,8 @@ void MFMASmallGemmOpt::applyIGLPStrategy(
         SchedGroupMask::MFMA, 1, PipelineSyncID, DAG, TII);
     SG->initSchedGroup(SyncedInstrs[SG->getSyncID()]);
   }
+
+  return true;
 }
 
 class MFMASmallGemmSingleWaveOpt final : public IGLPStrategy {
@@ -1098,7 +1100,7 @@ class MFMASmallGemmSingleWaveOpt final : public IGLPStrategy {
   };
 
 public:
-  void applyIGLPStrategy(
+  bool applyIGLPStrategy(
       DenseMap<int, SUnitsToCandidateSGsMap> &SyncedInstrs,
       DenseMap<int, SmallVector<SchedGroup, 4>> &SyncedSchedGroups,
       IGLPPhase Phase) override;
@@ -1115,7 +1117,7 @@ static unsigned DSWCount = 0;
 static unsigned DSWWithPermCount = 0;
 static unsigned DSWWithSharedVMEMCount = 0;
 
-void MFMASmallGemmSingleWaveOpt::applyIGLPStrategy(
+bool MFMASmallGemmSingleWaveOpt::applyIGLPStrategy(
     DenseMap<int, SUnitsToCandidateSGsMap> &SyncedInstrs,
     DenseMap<int, SmallVector<SchedGroup, 4>> &SyncedSchedGroups,
     IGLPPhase Phase) {
@@ -1355,6 +1357,8 @@ void MFMASmallGemmSingleWaveOpt::applyIGLPStrategy(
         SchedGroupMask::MFMA, 1, PipelineSyncID, DAG, TII);
     SG->initSchedGroup(SyncedInstrs[SG->getSyncID()]);
   }
+
+  return true;
 }
 
 static std::unique_ptr<IGLPStrategy>
@@ -1376,6 +1380,8 @@ class IGroupLPDAGMutation : public ScheduleDAGMutation {
 
   ScheduleDAGMI *DAG;
 
+  std::vector<std::unique_ptr<ScheduleDAGMutation>> *SavedMutations;
+
   // Organize lists of SchedGroups by their SyncID. SchedGroups /
   // SCHED_GROUP_BARRIERs with different SyncIDs will have no edges added
   // between then.
@@ -1402,7 +1408,7 @@ class IGroupLPDAGMutation : public ScheduleDAGMutation {
   void initSchedGroupBarrierPipelineStage(
       std::vector<SUnit>::reverse_iterator RIter);
 
-  void initIGLPOpt(SUnit &SU);
+  bool initIGLPOpt(SUnit &SU);
 
 public:
   void apply(ScheduleDAGInstrs *DAGInstrs) override;
@@ -1418,7 +1424,10 @@ class IGroupLPDAGMutation : public ScheduleDAGMutation {
   IGLPPhase Phase = IGLPPhase::Initial;
 
   IGroupLPDAGMutation() = default;
-  IGroupLPDAGMutation(IGLPPhase Phase) : Phase(Phase) {}
+  IGroupLPDAGMutation(
+      IGLPPhase Phase,
+      std::vector<std::unique_ptr<ScheduleDAGMutation>> *SavedMutations)
+      : SavedMutations(SavedMutations), Phase(Phase) {}
 };
 
 unsigned SchedGroup::NumSchedGroups = 0;
@@ -1609,31 +1618,41 @@ void IGroupLPDAGMutation::apply(ScheduleDAGInstrs *DAGInstrs) {
   DAG = static_cast<ScheduleDAGMI *>(DAGInstrs);
   SyncedSchedGroups.clear();
   SyncedInstrs.clear();
-  bool foundSB = false;
-  bool foundIGLP = false;
+  bool FoundSB = false;
+  bool FoundIGLP = false;
+  bool ShouldApplyIGLP = false;
   for (auto R = DAG->SUnits.rbegin(), E = DAG->SUnits.rend(); R != E; ++R) {
     unsigned Opc = R->getInstr()->getOpcode();
     // SCHED_[GROUP_]BARRIER and IGLP are mutually exclusive.
     if (Opc == AMDGPU::SCHED_BARRIER) {
       addSchedBarrierEdges(*R);
-      foundSB = true;
+      FoundSB = true;
     } else if (Opc == AMDGPU::SCHED_GROUP_BARRIER) {
       initSchedGroupBarrierPipelineStage(R);
-      foundSB = true;
+      FoundSB = true;
     } else if (Opc == AMDGPU::IGLP_OPT) {
       resetEdges(*R, DAG);
-      if (!foundSB && !foundIGLP)
-        initIGLPOpt(*R);
-      foundIGLP = true;
+      if (!FoundSB && !FoundIGLP) {
+        FoundIGLP = true;
+        ShouldApplyIGLP = initIGLPOpt(*R);
+      }
     }
   }
 
-  if (foundSB || foundIGLP) {
+  if (FoundSB || (FoundIGLP && ShouldApplyIGLP)) {
     PipelineSolver PS(SyncedSchedGroups, SyncedInstrs, DAG, IsBottomUp);
     // PipelineSolver performs the mutation by adding the edges it
     // determined as the best
     PS.solve();
+    return;
   }
+
+  if (!SavedMutations)
+    return;
+
+  // We did not apply a mutation, fall back to SavedMutations
+  for (auto &m : *SavedMutations)
+    m->apply(DAG);
 }
 
 void IGroupLPDAGMutation::addSchedBarrierEdges(SUnit &SchedBarrier) {
@@ -1712,14 +1731,15 @@ void IGroupLPDAGMutation::initSchedGroupBarrierPipelineStage(
   SG.initSchedGroup(RIter, SyncedInstrs[SG.getSyncID()]);
 }
 
-void IGroupLPDAGMutation::initIGLPOpt(SUnit &SU) {
+bool IGroupLPDAGMutation::initIGLPOpt(SUnit &SU) {
   IGLPStrategyID StrategyID =
       (IGLPStrategyID)SU.getInstr()->getOperand(0).getImm();
   auto S = createIGLPStrategy(StrategyID, DAG, TII);
-  if (S->shouldApplyStrategy(DAG)) {
-    IsBottomUp = S->IsBottomUp;
-    S->applyIGLPStrategy(SyncedInstrs, SyncedSchedGroups, Phase);
-  }
+  if (!S->shouldApplyStrategy(DAG))
+    return false;
+
+  IsBottomUp = S->IsBottomUp;
+  return S->applyIGLPStrategy(SyncedInstrs, SyncedSchedGroups, Phase);
 }
 
 } // namespace
@@ -1731,9 +1751,10 @@ namespace llvm {
 /// same scheduling region (e.g. pre and post-RA scheduling / multiple
 /// scheduling "phases"), we can reenter this mutation framework more than once
 /// for a given region.
-std::unique_ptr<ScheduleDAGMutation>
-createIGroupLPDAGMutation(IGLPPhase Phase) {
-  return std::make_unique<IGroupLPDAGMutation>(Phase);
+std::unique_ptr<ScheduleDAGMutation> createIGroupLPDAGMutation(
+    IGLPPhase Phase,
+    std::vector<std::unique_ptr<ScheduleDAGMutation>> *SavedMutations) {
+  return std::make_unique<IGroupLPDAGMutation>(Phase, SavedMutations);
 }
 
 } // end namespace llvm
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUIGroupLP.h b/llvm/lib/Target/AMDGPU/AMDGPUIGroupLP.h
@@ -11,14 +11,17 @@
 
 #include "llvm/CodeGen/ScheduleDAGMutation.h"
 #include <memory>
+#include <vector>
 
 namespace llvm {
 
 // Components of the mask that determines which instruction types may be may be
 // classified into a SchedGroup.
 enum class IGLPPhase { Initial = 0u, PreRAReentry = 1u << 0, PostRA = 1u << 1 };
 
-std::unique_ptr<ScheduleDAGMutation> createIGroupLPDAGMutation(IGLPPhase Phase);
+std::unique_ptr<ScheduleDAGMutation> createIGroupLPDAGMutation(
+    IGLPPhase Phase,
+    std::vector<std::unique_ptr<ScheduleDAGMutation>> *SavedMutations);
 
 } // namespace llvm
 
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -461,7 +461,7 @@ createGCNMaxOccupancyMachineScheduler(MachineSchedContext *C) {
   DAG->addMutation(createLoadClusterDAGMutation(DAG->TII, DAG->TRI));
   if (ST.shouldClusterStores())
     DAG->addMutation(createStoreClusterDAGMutation(DAG->TII, DAG->TRI));
-  DAG->addMutation(createIGroupLPDAGMutation(IGLPPhase::Initial));
+  DAG->addMutation(createIGroupLPDAGMutation(IGLPPhase::Initial, nullptr));
   DAG->addMutation(createAMDGPUMacroFusionDAGMutation());
   DAG->addMutation(createAMDGPUExportClusteringDAGMutation());
   return DAG;
@@ -471,7 +471,7 @@ static ScheduleDAGInstrs *
 createGCNMaxILPMachineScheduler(MachineSchedContext *C) {
   ScheduleDAGMILive *DAG =
       new GCNScheduleDAGMILive(C, std::make_unique<GCNMaxILPSchedStrategy>(C));
-  DAG->addMutation(createIGroupLPDAGMutation(IGLPPhase::Initial));
+  DAG->addMutation(createIGroupLPDAGMutation(IGLPPhase::Initial, nullptr));
   return DAG;
 }
 
@@ -934,7 +934,7 @@ class GCNPassConfig final : public AMDGPUPassConfig {
     if (ST.shouldClusterStores())
       DAG->addMutation(createStoreClusterDAGMutation(DAG->TII, DAG->TRI));
     DAG->addMutation(ST.createFillMFMAShadowMutation(DAG->TII));
-    DAG->addMutation(createIGroupLPDAGMutation(IGLPPhase::PostRA));
+    DAG->addMutation(createIGroupLPDAGMutation(IGLPPhase::PostRA, nullptr));
     if (isPassEnabled(EnableVOPD, CodeGenOptLevel::Less))
       DAG->addMutation(createVOPDPairingMutation());
     return DAG;
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
@@ -713,7 +713,7 @@ bool UnclusteredHighRPStage::initGCNSchedStage() {
     return false;
 
   SavedMutations.swap(DAG.Mutations);
-  DAG.addMutation(createIGroupLPDAGMutation(IGLPPhase::PreRAReentry));
+  DAG.addMutation(createIGroupLPDAGMutation(IGLPPhase::PreRAReentry, nullptr));
 
   InitialOccupancy = DAG.MinOccupancy;
   // Aggressivly try to reduce register pressure in the unclustered high RP
@@ -856,7 +856,8 @@ bool GCNSchedStage::initGCNRegion() {
     bool IsInitialStage = StageID == GCNSchedStageID::OccInitialSchedule ||
                           StageID == GCNSchedStageID::ILPInitialSchedule;
     DAG.addMutation(createIGroupLPDAGMutation(
-        IsInitialStage ? IGLPPhase::Initial : IGLPPhase::PreRAReentry));
+        IsInitialStage ? IGLPPhase::Initial : IGLPPhase::PreRAReentry,
+        &SavedMutations));
   }
 
   return true;
@@ -1570,7 +1571,8 @@ void GCNPostScheduleDAGMILive::schedule() {
   if (HasIGLPInstrs) {
     SavedMutations.clear();
     SavedMutations.swap(Mutations);
-    addMutation(createIGroupLPDAGMutation(/*IsReentry=*/IGLPPhase::PostRA));
+    addMutation(createIGroupLPDAGMutation(/*IsReentry=*/IGLPPhase::PostRA,
+                                          &SavedMutations));
   }
 
   ScheduleDAGMI::schedule();
