llvm · lucas-rami · Mar 4, 2025 · Feb 25, 2025 · Mar 4, 2025
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
@@ -769,6 +769,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;
@@ -778,11 +779,9 @@ void MachineSchedulerBase::scheduleRegions(ScheduleDAGInstrs &Scheduler,
      // it. Perhaps it still needs to be bundled.
      Scheduler.enterRegion(&*MBB, I, RegionEnd, NumRegionInstrs);

-      // Skip empty scheduling regions but include single-MI regions; we want
-      // those to be scheduled so that backends which move MIs across regions
-      // during scheduling can reason about and schedule those regions
-      // correctly.
-      if (I == 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
@@ -760,6 +760,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);

diff --git a/llvm/test/CodeGen/ARM/misched-branch-targets.mir b/llvm/test/CodeGen/ARM/misched-branch-targets.mir
@@ -1,7 +1,7 @@
-# RUN: llc -o - -run-pass=machine-scheduler -misched=shuffle %s | FileCheck -check-prefixes=CHECK,CHECK-MISCHED %s
-# RUN: llc -o - -passes=machine-scheduler -misched=shuffle %s | FileCheck -check-prefixes=CHECK,CHECK-MISCHED %s
-# RUN: llc -o - -run-pass=postmisched %s | FileCheck -check-prefixes=CHECK,CHECK-POSTMISCHED %s
-# RUN: llc -o - -passes=postmisched %s | FileCheck -check-prefixes=CHECK,CHECK-POSTMISCHED %s
+# 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
@@ -147,8 +147,7 @@ body:             |

 # CHECK-LABEL: name:            foo_setjmp
 # CHECK:       body:
-# CHECK-MISCHED:         tBL 14 /* CC::al */, $noreg, @setjmp, csr_aapcs, implicit-def dead $lr, implicit $sp, implicit $r0, implicit-def $sp, implicit-def $r0
-# CHECK-POSTMISCHED:     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

 ---

diff --git a/llvm/test/CodeGen/X86/fake-use-scheduler.mir b/llvm/test/CodeGen/X86/fake-use-scheduler.mir
@@ -9,12 +9,6 @@
 #
 # CHECK: ********** MI Scheduling **********
 # CHECK-NEXT: foo:%bb.0 entry
-# CHECK-NEXT:   From: $rax = COPY %5:gr64
-# CHECK-NEXT:     To: RET 0, killed $rax
-# CHECK-NEXT:  RegionInstrs: 1
-#
-# CHECK: ********** MI Scheduling **********
-# CHECK-NEXT: foo:%bb.0 entry
 # CHECK-NEXT:   From: %0:gr64 = COPY $rdi
 # CHECK-NEXT:     To: FAKE_USE %5:gr64
 # CHECK-NEXT:  RegionInstrs: 7