[MachineScheduler] Add support for scheduling while in SSA

llvm/include/llvm/CodeGen/MachineScheduler.h

@@ -103,6 +103,7 @@ namespace impl_detail {
 // FIXME: Remove these declarations once RegisterClassInfo is queryable as an
 // analysis.
 class MachineSchedulerImpl;
+class SSAMachineSchedulerImpl;
 class PostMachineSchedulerImpl;
 } // namespace impl_detail
 
@@ -1464,6 +1465,20 @@ class MachineSchedulerPass : public PassInfoMixin<MachineSchedulerPass> {
                                  MachineFunctionAnalysisManager &MFAM);
 };
 
+class SSAMachineSchedulerPass : public PassInfoMixin<SSAMachineSchedulerPass> {
+  // FIXME: Remove this member once RegisterClassInfo is queryable as an
+  // analysis.
+  std::unique_ptr<impl_detail::SSAMachineSchedulerImpl> Impl;
+  const TargetMachine *TM;
+
+public:
+  SSAMachineSchedulerPass(const TargetMachine *TM);
+  SSAMachineSchedulerPass(SSAMachineSchedulerPass &&Other);
+  ~SSAMachineSchedulerPass();
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+};
+
 class PostMachineSchedulerPass
     : public PassInfoMixin<PostMachineSchedulerPass> {
   // FIXME: Remove this member once RegisterClassInfo is queryable as an

llvm/include/llvm/CodeGen/Passes.h

@@ -165,6 +165,9 @@ LLVM_ABI extern char &MachineSchedulerID;
 /// PostMachineScheduler - This pass schedules machine instructions postRA.
 LLVM_ABI extern char &PostMachineSchedulerID;
 
+/// SSAMachineScheduler - This pass schedules machine instructions in SSA.
+LLVM_ABI extern char &SSAMachineSchedulerID;
+
 /// SpillPlacement analysis. Suggest optimal placement of spill code between
 /// basic blocks.
 LLVM_ABI extern char &SpillPlacementID;

llvm/include/llvm/CodeGen/TargetPassConfig.h

@@ -135,6 +135,10 @@ class LLVM_ABI TargetPassConfig : public ImmutablePass {
   /// replace a copy.
   bool EnableSinkAndFold = false;
 
+  /// Enable insertion of SSAMachineScheduler pass, this triggers early
+  /// computation of live intervals.
+  bool EnableSSAMachineScheduler = false;
+
   /// Require processing of functions such that callees are generated before
   /// callers.
   bool RequireCodeGenSCCOrder = false;
@@ -205,6 +209,13 @@ class LLVM_ABI TargetPassConfig : public ImmutablePass {
     setOpt(RequireCodeGenSCCOrder, Enable);
   }
 
+  bool getEnableSSAMachineScheduler() const {
+    return EnableSSAMachineScheduler;
+  }
+  void setEnableSSAMachineScheduler(bool Enable) {
+    setOpt(EnableSSAMachineScheduler, Enable);
+  }
+
   /// Allow the target to override a specific pass without overriding the pass
   /// pipeline. When passes are added to the standard pipeline at the
   /// point where StandardID is expected, add TargetID in its place.

llvm/include/llvm/CodeGen/TargetSubtargetInfo.h

@@ -220,6 +220,10 @@ class LLVM_ABI TargetSubtargetInfo : public MCSubtargetInfo {
   /// allocation.
   virtual bool enablePostRAMachineScheduler() const;
 
+  /// True if the subtarget should run a machine scheduler before PHI
+  /// elimination.
+  virtual bool enableSSAMachineScheduler() const;
+
   /// True if the subtarget should run the atomic expansion pass.
   virtual bool enableAtomicExpand() const;
 

llvm/include/llvm/InitializePasses.h

@@ -289,6 +289,7 @@ LLVM_ABI void initializeReplaceWithVeclibLegacyPass(PassRegistry &);
 LLVM_ABI void initializeResetMachineFunctionPass(PassRegistry &);
 LLVM_ABI void initializeSCEVAAWrapperPassPass(PassRegistry &);
 LLVM_ABI void initializeSROALegacyPassPass(PassRegistry &);
+LLVM_ABI void initializeSSAMachineSchedulerPass(PassRegistry &);
 LLVM_ABI void initializeSafeStackLegacyPassPass(PassRegistry &);
 LLVM_ABI void initializeSafepointIRVerifierPass(PassRegistry &);
 LLVM_ABI void initializeSelectOptimizePass(PassRegistry &);

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -122,6 +122,7 @@ MACHINE_FUNCTION_PASS("machine-cse", MachineCSEPass())
 MACHINE_FUNCTION_PASS("machine-latecleanup", MachineLateInstrsCleanupPass())
 MACHINE_FUNCTION_PASS("machine-sanmd", MachineSanitizerBinaryMetadataPass())
 MACHINE_FUNCTION_PASS("machine-scheduler", MachineSchedulerPass(TM))
+MACHINE_FUNCTION_PASS("ssa-machine-scheduler", SSAMachineSchedulerPass(TM))
 MACHINE_FUNCTION_PASS("machinelicm", MachineLICMPass())
 MACHINE_FUNCTION_PASS("no-op-machine-function", NoOpMachineFunctionPass())
 MACHINE_FUNCTION_PASS("opt-phis", OptimizePHIsPass())

llvm/lib/CodeGen/CodeGen.cpp

@@ -120,6 +120,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeRemoveLoadsIntoFakeUsesLegacyPass(Registry);
   initializeRemoveRedundantDebugValuesLegacyPass(Registry);
   initializeRenameIndependentSubregsLegacyPass(Registry);
+  initializeSSAMachineSchedulerPass(Registry);
   initializeSafeStackLegacyPassPass(Registry);
   initializeSelectOptimizePass(Registry);
   initializeShadowStackGCLoweringPass(Registry);

llvm/lib/CodeGen/MachineScheduler.cpp

@@ -350,6 +350,33 @@ class MachineSchedulerImpl : public MachineSchedulerBase {
   ScheduleDAGInstrs *createMachineScheduler();
 };
 
+/// Impl class for SSAMachineScheduler.
+class SSAMachineSchedulerImpl : public MachineSchedulerBase {
+  // These are only for using MF.verify()
+  // remove when verify supports passing in all analyses
+  MachineFunctionPass *P = nullptr;
+  MachineFunctionAnalysisManager *MFAM = nullptr;
+
+public:
+  struct RequiredAnalyses {
+    MachineLoopInfo &MLI;
+    MachineDominatorTree &MDT;
+    AAResults &AA;
+    LiveIntervals &LIS;
+  };
+
+  SSAMachineSchedulerImpl() {}
+  // Migration only
+  void setLegacyPass(MachineFunctionPass *P) { this->P = P; }
+  void setMFAM(MachineFunctionAnalysisManager *MFAM) { this->MFAM = MFAM; }
+
+  bool run(MachineFunction &MF, const TargetMachine &TM,
+           const RequiredAnalyses &Analyses);
+
+protected:
+  ScheduleDAGInstrs *createMachineScheduler();
+};
+
 /// Impl class for PostMachineScheduler.
 class PostMachineSchedulerImpl : public MachineSchedulerBase {
   // These are only for using MF.verify()
@@ -380,6 +407,7 @@ class PostMachineSchedulerImpl : public MachineSchedulerBase {
 using impl_detail::MachineSchedulerBase;
 using impl_detail::MachineSchedulerImpl;
 using impl_detail::PostMachineSchedulerImpl;
+using impl_detail::SSAMachineSchedulerImpl;
 
 namespace {
 /// MachineScheduler runs after coalescing and before register allocation.
@@ -394,6 +422,18 @@ class MachineSchedulerLegacy : public MachineFunctionPass {
   static char ID; // Class identification, replacement for typeinfo
 };
 
+/// SSAMachineScheduler runs before PHI elimination.
+class SSAMachineScheduler : public MachineFunctionPass {
+  SSAMachineSchedulerImpl Impl;
+
+public:
+  SSAMachineScheduler();
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+  bool runOnMachineFunction(MachineFunction &) override;
+
+  static char ID; // Class identification, replacement for typeinfo
+};
+
 /// PostMachineScheduler runs after shortly before code emission.
 class PostMachineSchedulerLegacy : public MachineFunctionPass {
   PostMachineSchedulerImpl Impl;
@@ -439,6 +479,35 @@ void MachineSchedulerLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
+char SSAMachineScheduler::ID = 0;
+
+char &llvm::SSAMachineSchedulerID = SSAMachineScheduler::ID;
+
+INITIALIZE_PASS_BEGIN(SSAMachineScheduler, "ssa-machine-scheduler",
+                      "SSA Machine Instruction Scheduler", false, false)
+INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(MachineDominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(SlotIndexesWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(LiveIntervalsWrapperPass)
+INITIALIZE_PASS_END(SSAMachineScheduler, "ssa-machine-scheduler",
+                    "SSA Machine Instruction Scheduler", false, false)
+
+SSAMachineScheduler::SSAMachineScheduler() : MachineFunctionPass(ID) {
+  initializeSSAMachineSchedulerPass(*PassRegistry::getPassRegistry());
+}
+
+void SSAMachineScheduler::getAnalysisUsage(AnalysisUsage &AU) const {
+  AU.setPreservesCFG();
+  AU.addRequired<MachineDominatorTreeWrapperPass>();
+  AU.addRequired<MachineLoopInfoWrapperPass>();
+  AU.addRequired<AAResultsWrapperPass>();
+  AU.addRequired<TargetPassConfig>();
+  AU.addRequired<SlotIndexesWrapperPass>();
+  AU.addRequired<LiveIntervalsWrapperPass>();
+  MachineFunctionPass::getAnalysisUsage(AU);
+}
+
 char PostMachineSchedulerLegacy::ID = 0;
 
 char &llvm::PostMachineSchedulerID = PostMachineSchedulerLegacy::ID;
@@ -490,6 +559,11 @@ static cl::opt<bool> EnableMachineSched(
     cl::desc("Enable the machine instruction scheduling pass."), cl::init(true),
     cl::Hidden);
 
+static cl::opt<bool> EnableSSAMachineSched(
+    "enable-ssa-misched",
+    cl::desc("Enable the machine instruction scheduling pass in SSA."),
+    cl::init(false), cl::Hidden);
+
 static cl::opt<bool> EnablePostRAMachineSched(
     "enable-post-misched",
     cl::desc("Enable the post-ra machine instruction scheduling pass."),
@@ -586,6 +660,53 @@ bool MachineSchedulerImpl::run(MachineFunction &Func, const TargetMachine &TM,
   return true;
 }
 
+/// Instantiate a ScheduleDAGInstrs that will be owned by the caller.
+ScheduleDAGInstrs *SSAMachineSchedulerImpl::createMachineScheduler() {
+  // Get the default scheduler set by the target for this function.
+  ScheduleDAGInstrs *Scheduler = TM->createMachineScheduler(this);
+  if (Scheduler)
+    return Scheduler;
+
+  // Default to GenericScheduler.
+  return createSchedLive(this);
+}
+
+bool SSAMachineSchedulerImpl::run(MachineFunction &Func,
+                                  const TargetMachine &TM,
+                                  const RequiredAnalyses &Analyses) {
+  MF = &Func;
+  MLI = &Analyses.MLI;
+  MDT = &Analyses.MDT;
+  this->TM = &TM;
+  AA = &Analyses.AA;
+  LIS = &Analyses.LIS;
+
+  if (VerifyScheduling) {
+    LLVM_DEBUG(LIS->dump());
+    const char *MSchedBanner = "Before machine scheduling.";
+    if (P)
+      MF->verify(P, MSchedBanner, &errs());
+    else
+      MF->verify(*MFAM, MSchedBanner, &errs());
+  }
+  RegClassInfo->runOnMachineFunction(*MF);
+
+  // Instantiate the selected scheduler for this target, function, and
+  // optimization level.
+  std::unique_ptr<ScheduleDAGInstrs> Scheduler(createMachineScheduler());
+  scheduleRegions(*Scheduler, false);
+
+  LLVM_DEBUG(LIS->dump());
+  if (VerifyScheduling) {
+    const char *MSchedBanner = "After machine scheduling.";
+    if (P)
+      MF->verify(P, MSchedBanner, &errs());
+    else
+      MF->verify(*MFAM, MSchedBanner, &errs());
+  }
+  return true;
+}
+
 /// Instantiate a ScheduleDAGInstrs for PostRA scheduling that will be owned by
 /// the caller. We don't have a command line option to override the postRA
 /// scheduler. The Target must configure it.
@@ -668,12 +789,38 @@ bool MachineSchedulerLegacy::runOnMachineFunction(MachineFunction &MF) {
   return Impl.run(MF, TM, {MLI, MDT, AA, LIS});
 }
 
+bool SSAMachineScheduler::runOnMachineFunction(MachineFunction &MF) {
+  if (skipFunction(MF.getFunction()))
+    return false;
+
+  if (EnableSSAMachineSched.getNumOccurrences()) {
+    if (!EnableSSAMachineSched)
+      return false;
+  } else if (!MF.getSubtarget().enableSSAMachineScheduler()) {
+    return false;
+  }
+
+  auto &MLI = getAnalysis<MachineLoopInfoWrapperPass>().getLI();
+  auto &MDT = getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
+  auto &TM = getAnalysis<TargetPassConfig>().getTM<TargetMachine>();
+  auto &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
+  auto &LIS = getAnalysis<LiveIntervalsWrapperPass>().getLIS();
+  Impl.setLegacyPass(this);
+  return Impl.run(MF, TM, {MLI, MDT, AA, LIS});
+}
+
 MachineSchedulerPass::MachineSchedulerPass(const TargetMachine *TM)
     : Impl(std::make_unique<MachineSchedulerImpl>()), TM(TM) {}
 MachineSchedulerPass::~MachineSchedulerPass() = default;
 MachineSchedulerPass::MachineSchedulerPass(MachineSchedulerPass &&Other) =
     default;
 
+SSAMachineSchedulerPass::SSAMachineSchedulerPass(const TargetMachine *TM)
+    : Impl(std::make_unique<SSAMachineSchedulerImpl>()), TM(TM) {}
+SSAMachineSchedulerPass::SSAMachineSchedulerPass(
+    SSAMachineSchedulerPass &&Other) = default;
+SSAMachineSchedulerPass::~SSAMachineSchedulerPass() = default;
+
 PostMachineSchedulerPass::PostMachineSchedulerPass(const TargetMachine *TM)
     : Impl(std::make_unique<PostMachineSchedulerImpl>()), TM(TM) {}
 PostMachineSchedulerPass::PostMachineSchedulerPass(
@@ -708,6 +855,33 @@ MachineSchedulerPass::run(MachineFunction &MF,
       .preserve<LiveIntervalsAnalysis>();
 }
 
+PreservedAnalyses
+SSAMachineSchedulerPass::run(MachineFunction &MF,
+                             MachineFunctionAnalysisManager &MFAM) {
+  if (EnableSSAMachineSched.getNumOccurrences()) {
+    if (!EnableSSAMachineSched)
+      return PreservedAnalyses::all();
+  } else if (!MF.getSubtarget().enableSSAMachineScheduler()) {
+    LLVM_DEBUG(dbgs() << "Subtarget disables ssa-MI-sched.\n");
+    return PreservedAnalyses::all();
+  }
+
+  auto &MLI = MFAM.getResult<MachineLoopAnalysis>(MF);
+  auto &MDT = MFAM.getResult<MachineDominatorTreeAnalysis>(MF);
+  auto &FAM = MFAM.getResult<FunctionAnalysisManagerMachineFunctionProxy>(MF)
+                  .getManager();
+  auto &AA = FAM.getResult<AAManager>(MF.getFunction());
+  auto &LIS = MFAM.getResult<LiveIntervalsAnalysis>(MF);
+  Impl->setMFAM(&MFAM);
+  bool Changed = Impl->run(MF, *TM, {MLI, MDT, AA, LIS});
+  if (!Changed)
+    return PreservedAnalyses::all();
+
+  PreservedAnalyses PA = getMachineFunctionPassPreservedAnalyses();
+  PA.preserveSet<CFGAnalyses>();
+  return PA;
+}
+
 bool PostMachineSchedulerLegacy::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
@@ -764,11 +938,10 @@ PostMachineSchedulerPass::run(MachineFunction &MF,
 /// the boundary, but there would be no benefit to postRA scheduling across
 /// calls this late anyway.
 static bool isSchedBoundary(MachineBasicBlock::iterator MI,
-                            MachineBasicBlock *MBB,
-                            MachineFunction *MF,
+                            MachineBasicBlock *MBB, MachineFunction *MF,
                             const TargetInstrInfo *TII) {
   return MI->isCall() || TII->isSchedulingBoundary(*MI, MBB, *MF) ||
-         MI->isFakeUse();
+         MI->isFakeUse() || MI->isPHI();
 }
 
 using MBBRegionsVector = SmallVector<SchedRegion, 16>;

llvm/lib/CodeGen/TargetPassConfig.cpp

@@ -1479,6 +1479,12 @@ void TargetPassConfig::addOptimizedRegAlloc() {
   addPass(&UnreachableMachineBlockElimID);
   addPass(&LiveVariablesID);
 
+  // Run SSA machine scheduler runs just before PHI elimination.
+  if (EnableSSAMachineScheduler) {
+    addPass(&LiveIntervalsID);
+    addPass(&SSAMachineSchedulerID);
+  }
+
   // Edge splitting is smarter with machine loop info.
   addPass(&MachineLoopInfoID);
   addPass(&PHIEliminationID);

llvm/lib/CodeGen/TargetSubtargetInfo.cpp

@@ -54,6 +54,8 @@ bool TargetSubtargetInfo::enablePostRAMachineScheduler() const {
   return enableMachineScheduler() && enablePostRAScheduler();
 }
 
+bool TargetSubtargetInfo::enableSSAMachineScheduler() const { return false; }
+
 bool TargetSubtargetInfo::useAA() const {
   return false;
 }