[MachineScheduler] Optional scheduling of single-MI regions (llvm#129704)

lucas-rami · jrbyrnes · commit e0b64882dc0f · 2025-05-01T08:52:01.000-07:00
Following 15e295d the machine scheduler no longer filters-out single-MI regions when emitting regions to schedule. While this has no functional impact at the moment, it generally has a negative compile-time impact (see llvm#128739). Since all targets but AMDGPU do not care for this behavior, this introduces an off-by-default flag to `ScheduleDAGInstrs` to control whether such regions are going to be scheduled, effectively reverting 15e295d for all targets but AMDGPU (currently the only target enabling this flag). Change-Id: Ib38f9b7e8d2bb1073cb43ed7e58eaf251ffdce48
diff --git a/llvm/include/llvm/CodeGen/ScheduleDAGInstrs.h b/llvm/include/llvm/CodeGen/ScheduleDAGInstrs.h
@@ -124,6 +124,9 @@ namespace llvm {
     /// rescheduling).
     bool RemoveKillFlags;
 
+    /// True if regions with a single MI should be scheduled.
+    bool ScheduleSingleMIRegions = false;
+
     /// The standard DAG builder does not normally include terminators as DAG
     /// nodes because it does not create the necessary dependencies to prevent
     /// reordering. A specialized scheduler can override
@@ -288,6 +291,11 @@ namespace llvm {
       return Topo.IsReachable(SU, TargetSU);
     }
 
+    /// Whether regions with a single MI should be scheduled.
+    bool shouldScheduleSingleMIRegions() const {
+      return ScheduleSingleMIRegions;
+    }
+
     /// Returns an iterator to the top of the current scheduling region.
     MachineBasicBlock::iterator begin() const { return RegionBegin; }
 
diff --git a/llvm/lib/CodeGen/MachineScheduler.cpp b/llvm/lib/CodeGen/MachineScheduler.cpp
@@ -618,6 +618,7 @@ void MachineSchedulerBase::scheduleRegions(ScheduleDAGInstrs &Scheduler,
 
     MBBRegionsVector MBBRegions;
     getSchedRegions(&*MBB, MBBRegions, Scheduler.doMBBSchedRegionsTopDown());
+    bool ScheduleSingleMI = Scheduler.shouldScheduleSingleMIRegions();
     for (const SchedRegion &R : MBBRegions) {
       MachineBasicBlock::iterator I = R.RegionBegin;
       MachineBasicBlock::iterator RegionEnd = R.RegionEnd;
@@ -627,8 +628,9 @@ void MachineSchedulerBase::scheduleRegions(ScheduleDAGInstrs &Scheduler,
       // it. Perhaps it still needs to be bundled.
       Scheduler.enterRegion(&*MBB, I, RegionEnd, NumRegionInstrs);
 
-      // Skip empty scheduling regions (0 or 1 schedulable instructions).
-      if (I == RegionEnd || I == std::prev(RegionEnd)) {
+      // Skip empty scheduling regions and, conditionally, regions with a single
+      // MI.
+      if (I == RegionEnd || (!ScheduleSingleMI && I == std::prev(RegionEnd))) {
         // Close the current region. Bundle the terminator if needed.
         // This invalidates 'RegionEnd' and 'I'.
         Scheduler.exitRegion();
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
@@ -765,6 +765,10 @@ GCNScheduleDAGMILive::GCNScheduleDAGMILive(
       StartingOccupancy(MFI.getOccupancy()), MinOccupancy(StartingOccupancy),
       RegionLiveOuts(this, /*IsLiveOut=*/true) {
 
+  // We want regions with a single MI to be scheduled so that we can reason
+  // about them correctly during scheduling stages that move MIs between regions
+  // (e.g., rematerialization).
+  ScheduleSingleMIRegions = true;
   LLVM_DEBUG(dbgs() << "Starting occupancy is " << StartingOccupancy << ".\n");
   if (RelaxedOcc) {
     MinOccupancy = std::min(MFI.getMinAllowedOccupancy(), StartingOccupancy);
@@ -1098,6 +1102,13 @@ bool PreRARematStage::initGCNSchedStage() {
       DAG.Regions.size() == 1 || !canIncreaseOccupancyOrReduceSpill())
     return false;
 
+  // Map all MIs to their parent region.
+  for (unsigned I = 0, E = DAG.Regions.size(); I < E; ++I) {
+    RegionBoundaries Region = DAG.Regions[I];
+    for (auto MI = Region.first; MI != Region.second; ++MI)
+      MIRegion.insert({&*MI, I});
+  }
+
   // Rematerialize identified instructions and update scheduler's state.
   rematerialize();
   if (GCNTrackers)
@@ -1831,7 +1842,7 @@ bool PreRARematStage::canIncreaseOccupancyOrReduceSpill() {
   // Cache set of registers that are going to be rematerialized.
   DenseSet<unsigned> RematRegs;
 
-  // identify rematerializable instructions in the function.
+  // Identify rematerializable instructions in the function.
   for (unsigned I = 0, E = DAG.Regions.size(); I != E; ++I) {
     auto Region = DAG.Regions[I];
     for (auto MI = Region.first; MI != Region.second; ++MI) {
@@ -1941,6 +1952,8 @@ void PreRARematStage::rematerialize() {
     MachineBasicBlock::iterator InsertPos = Remat.InsertPos;
     Register Reg = DefMI->getOperand(0).getReg();
     unsigned SubReg = DefMI->getOperand(0).getSubReg();
+    unsigned DefRegion = MIRegion.at(DefMI);
+
     // Rematerialize DefMI to its use block.
     TII->reMaterialize(*InsertPos->getParent(), InsertPos, Reg, SubReg, *DefMI,
                        *DAG.TRI);
@@ -1952,8 +1965,12 @@ void PreRARematStage::rematerialize() {
     // Update region boundaries in regions we sinked from (remove defining MI)
     // and to (insert MI rematerialized in use block). Only then we can erase
     // the original MI.
-    DAG.updateRegionBoundaries(DAG.Regions, DefMI, nullptr);
-    DAG.updateRegionBoundaries(DAG.Regions, InsertPos, Remat.RematMI);
+    DAG.updateRegionBoundaries(DAG.Regions[DefRegion], DefMI, nullptr);
+    auto UseRegion = MIRegion.find(Remat.UseMI);
+    if (UseRegion != MIRegion.end()) {
+      DAG.updateRegionBoundaries(DAG.Regions[UseRegion->second], InsertPos,
+                                 Remat.RematMI);
+    }
     DefMI->eraseFromParent();
     DAG.LIS->RemoveMachineInstrFromMaps(*DefMI);
 
@@ -2008,8 +2025,8 @@ void PreRARematStage::rematerialize() {
     }
     // RP in the region from which the instruction was rematerialized may or may
     // not decrease.
-    ImpactedRegions.insert({Remat.DefRegion, DAG.Pressure[Remat.DefRegion]});
-    RecomputeRP.insert(Remat.DefRegion);
+    ImpactedRegions.insert({DefRegion, DAG.Pressure[DefRegion]});
+    RecomputeRP.insert(DefRegion);
 
     // Update the register's live interval manually. This is aligned with the
     // instruction collection phase in that it assumes a single def and use
@@ -2150,7 +2167,8 @@ void PreRARematStage::finalizeGCNSchedStage() {
   // Rollback the rematerializations.
   for (auto &Remat : Remats) {
     MachineInstr &RematMI = *Remat.RematMI;
-    MachineBasicBlock *MBB = getRegionMBB(MF, DAG.Regions[Remat.DefRegion]);
+    unsigned DefRegion = MIRegion.at(DefMI);
+    MachineBasicBlock *MBB = getRegionMBB(MF, DAG.Regions[DefRegion]);
     MachineBasicBlock::iterator InsertPos(MBB->end());
 
     Register Reg = RematMI.getOperand(0).getReg();
@@ -2164,6 +2182,13 @@ void PreRARematStage::finalizeGCNSchedStage() {
     NewMI->getOperand(0).setSubReg(SubReg);
     DAG.LIS->InsertMachineInstrInMaps(*NewMI);
 
+    auto UseRegion = MIRegion.find(Remat.UseMI);
+    if (UseRegion != MIRegion.end()) {
+      DAG.updateRegionBoundaries(DAG.Regions[UseRegion->second], RematMI,
+                                 nullptr);
+    }
+    DAG.updateRegionBoundaries(DAG.Regions[DefRegion], InsertPos, NewMI);
+
     // Erase rematerialized MI.
     DAG.updateRegionBoundaries(DAG.Regions, RematMI, nullptr);
     RematMI.eraseFromParent();
@@ -2187,49 +2212,25 @@ void PreRARematStage::finalizeGCNSchedStage() {
 }
 
 void GCNScheduleDAGMILive::updateRegionBoundaries(
-    SmallVectorImpl<RegionBoundaries> &RegionBoundaries,
-    MachineBasicBlock::iterator MI, MachineInstr *NewMI) {
-  unsigned I = 0, E = RegionBoundaries.size();
-  // Search for first region of the block where MI is located. We may encounter
-  // an empty region if all instructions from an initially non-empty region were
-  // removed.
-  while (I != E && RegionBoundaries[I].first != RegionBoundaries[I].second &&
-         MI->getParent() != RegionBoundaries[I].first->getParent())
-    ++I;
-
-  for (; I != E; ++I) {
-    auto &Bounds = RegionBoundaries[I];
-    assert(MI != Bounds.second && "cannot insert at region end");
-    assert(!NewMI || NewMI != Bounds.second && "cannot remove at region end");
-
-    // We may encounter an empty region if all of the region' instructions were
-    // previously removed.
-    if (Bounds.first == Bounds.second) {
-      if (MI->getParent()->end() != Bounds.second)
-        return;
-      continue;
-    }
-    if (MI->getParent() != Bounds.first->getParent())
-      return;
-
-    // We only care for modifications at the beginning of the region since the
-    // upper region boundary is exclusive.
-    if (MI != Bounds.first)
-      continue;
-    if (!NewMI) {
-      // This is an MI removal, which may leave the region empty; in such cases
-      // set both boundaries to the removed instruction's MBB's end.
-      MachineBasicBlock::iterator NextMI = std::next(MI);
-      if (NextMI != Bounds.second)
-        Bounds.first = NextMI;
-      else
-        Bounds.first = Bounds.second;
-    } else {
-      // This is an MI insertion at the beggining of the region.
-      Bounds.first = NewMI;
-    }
+    RegionBoundaries &RegionBounds, MachineBasicBlock::iterator MI,
+    MachineInstr *NewMI) {
+  assert(!NewMI ||
+         NewMI != RegionBounds.second && "cannot remove at region end");
+
+  if (RegionBounds.first == RegionBounds.second) {
+    assert(NewMI && "cannot remove from an empty region");
+    RegionBounds.first = NewMI;
     return;
   }
+
+  // We only care for modifications at the beginning of a non-empty region since
+  // the upper region boundary is exclusive.
+  if (MI != RegionBounds.first)
+    return;
+  if (!NewMI)
+    RegionBounds.first = std::next(MI); // Removal
+  else
+    RegionBounds.first = NewMI; // Insertion
 }
 
 static bool hasIGLPInstrs(ScheduleDAGInstrs *DAG) {
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
@@ -295,11 +295,10 @@ class GCNScheduleDAGMILive final : public ScheduleDAGMILive {
   /// If necessary, updates a region's boundaries following insertion ( \p NewMI
   /// != nullptr) or removal ( \p NewMI == nullptr) of a \p MI in the region.
   /// For an MI removal, this must be called before the MI is actually erased
-  /// from its parent MBB. If a region is left empty by a removal, both
-  /// boundaries are set to the last removed MI's MBB's end.
-  void
-  updateRegionBoundaries(SmallVectorImpl<RegionBoundaries> &RegionBoundaries,
-                         MachineBasicBlock::iterator MI, MachineInstr *NewMI);
+  /// from its parent MBB.
+  void updateRegionBoundaries(RegionBoundaries &RegionBounds,
+                              MachineBasicBlock::iterator MI,
+                              MachineInstr *NewMI);
 
   void runSchedStages();
 
@@ -463,8 +462,6 @@ class PreRARematStage : public GCNSchedStage {
     /// Set of regions in which the rematerializable instruction's defined
     /// register is a live-in.
     SmallDenseSet<unsigned, 4> LiveInRegions;
-    /// Region containing the rematerializable instruction.
-    unsigned DefRegion;
 
     /// The position at which to insert the remat.
     MachineBasicBlock::iterator InsertPos;
@@ -611,6 +608,12 @@ class PreRARematStage : public GCNSchedStage {
 
   RematInstructions Remats;
 
+    RematInstruction(MachineInstr *UseMI) : UseMI(UseMI) {}
+  };
+
+  /// Maps all MIs to their parent region. MI terminators are considered to be
+  /// outside the region they delimitate, and as such are not stored in the map.
+  DenseMap<MachineInstr *, unsigned> MIRegion;
   /// Collect regions whose live-ins or register pressure will change due to
   /// rematerializations.
   DenseMap<unsigned, GCNRegPressure> ImpactedRegions;
diff --git a/llvm/test/CodeGen/ARM/misched-branch-targets.mir b/llvm/test/CodeGen/ARM/misched-branch-targets.mir
@@ -1,5 +1,7 @@
 # RUN: llc -o - -run-pass=machine-scheduler -misched=shuffle %s | FileCheck %s
+# RUN: llc -o - -passes=machine-scheduler -misched=shuffle %s | FileCheck %s
 # RUN: llc -o - -run-pass=postmisched %s | FileCheck %s
+# RUN: llc -o - -passes=postmisched %s | FileCheck %s
 
 # REQUIRES: asserts
 # -misched=shuffle is only available with assertions enabled
@@ -145,7 +147,7 @@ body:             |
 
 # CHECK-LABEL: name:            foo_setjmp
 # CHECK:       body:
-# CHECK:         tBL 14 /* CC::al */, $noreg, @setjmp, csr_aapcs, implicit-def dead $lr, implicit $sp, implicit killed $r0, implicit-def $sp, implicit-def $r0
+# CHECK: tBL 14 /* CC::al */, $noreg, @setjmp, csr_aapcs, implicit-def dead $lr, implicit $sp, implicit killed $r0, implicit-def $sp, implicit-def $r0
 # CHECK-NEXT:    t2BTI
 
 ---
