[NewPM][CodeGen] Port StackColoring to NPM (#111812)

13 files changed, 75 insertions, 24 deletions

diff --git a/llvm/include/llvm/CodeGen/Passes.h b/llvm/include/llvm/CodeGen/Passes.h
index ddb2012..99421bd 100644
--- a/llvm/include/llvm/CodeGen/Passes.h
+++ b/llvm/include/llvm/CodeGen/Passes.h
@@ -285,7 +285,7 @@ namespace llvm {
 
   /// StackSlotColoring - This pass performs stack coloring and merging.
   /// It merges disjoint allocas to reduce the stack size.
-  extern char &StackColoringID;
+  extern char &StackColoringLegacyID;
 
   /// StackFramePrinter - This pass prints the stack frame layout and variable
   /// mappings.
diff --git a/llvm/include/llvm/CodeGen/StackColoring.h b/llvm/include/llvm/CodeGen/StackColoring.h
new file mode 100644
index 0000000..65f8916
--- /dev/null
+++ b/llvm/include/llvm/CodeGen/StackColoring.h
@@ -0,0 +1,24 @@
+//===- llvm/CodeGen/StackColoring.h -----------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_STACKCOLORINGPASS_H
+#define LLVM_CODEGEN_STACKCOLORINGPASS_H
+
+#include "llvm/CodeGen/MachinePassManager.h"
+
+namespace llvm {
+
+class StackColoringPass : public PassInfoMixin<StackColoringPass> {
+public:
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_STACKCOLORINGPASS_H
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index 4352099..6a75dc0 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/llvm/include/llvm/InitializePasses.h
@@ -288,7 +288,7 @@ void initializeSjLjEHPreparePass(PassRegistry &);
 void initializeSlotIndexesWrapperPassPass(PassRegistry &);
 void initializeSpeculativeExecutionLegacyPassPass(PassRegistry &);
 void initializeSpillPlacementPass(PassRegistry &);
-void initializeStackColoringPass(PassRegistry &);
+void initializeStackColoringLegacyPass(PassRegistry &);
 void initializeStackFrameLayoutAnalysisPassPass(PassRegistry &);
 void initializeStackMapLivenessPass(PassRegistry &);
 void initializeStackProtectorPass(PassRegistry &);
diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index 13bc4700d..0d45df0 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -56,6 +56,7 @@
 #include "llvm/CodeGen/SelectOptimize.h"
 #include "llvm/CodeGen/ShadowStackGCLowering.h"
 #include "llvm/CodeGen/SjLjEHPrepare.h"
+#include "llvm/CodeGen/StackColoring.h"
 #include "llvm/CodeGen/StackProtector.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TwoAddressInstructionPass.h"
diff --git a/llvm/include/llvm/Passes/MachinePassRegistry.def b/llvm/include/llvm/Passes/MachinePassRegistry.def
index 6ae80a4..2aa5f4f 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -152,6 +152,7 @@ MACHINE_FUNCTION_PASS("print<machine-post-dom-tree>",
 MACHINE_FUNCTION_PASS("print<slot-indexes>", SlotIndexesPrinterPass(dbgs()))
 MACHINE_FUNCTION_PASS("require-all-machine-function-properties",
                       RequireAllMachineFunctionPropertiesPass())
+MACHINE_FUNCTION_PASS("stack-coloring", StackColoringPass())
 MACHINE_FUNCTION_PASS("trigger-verifier-error", TriggerVerifierErrorPass())
 MACHINE_FUNCTION_PASS("two-address-instruction", TwoAddressInstructionPass())
 MACHINE_FUNCTION_PASS("verify", MachineVerifierPass())
@@ -255,7 +256,6 @@ DUMMY_MACHINE_FUNCTION_PASS("rename-independent-subregs", RenameIndependentSubre
 DUMMY_MACHINE_FUNCTION_PASS("reset-machine-function", ResetMachineFunctionPass)
 DUMMY_MACHINE_FUNCTION_PASS("shrink-wrap", ShrinkWrapPass)
 DUMMY_MACHINE_FUNCTION_PASS("simple-register-coalescing", RegisterCoalescerPass)
-DUMMY_MACHINE_FUNCTION_PASS("stack-coloring", StackColoringPass)
 DUMMY_MACHINE_FUNCTION_PASS("stack-frame-layout", StackFrameLayoutAnalysisPass)
 DUMMY_MACHINE_FUNCTION_PASS("stack-slot-coloring", StackSlotColoringPass)
 DUMMY_MACHINE_FUNCTION_PASS("stackmap-liveness", StackMapLivenessPass)
diff --git a/llvm/lib/CodeGen/CodeGen.cpp b/llvm/lib/CodeGen/CodeGen.cpp
index 16b8d45..48cc21ee 100644
--- a/llvm/lib/CodeGen/CodeGen.cpp
+++ b/llvm/lib/CodeGen/CodeGen.cpp
@@ -125,7 +125,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeShrinkWrapPass(Registry);
   initializeSjLjEHPreparePass(Registry);
   initializeSlotIndexesWrapperPassPass(Registry);
-  initializeStackColoringPass(Registry);
+  initializeStackColoringLegacyPass(Registry);
   initializeStackFrameLayoutAnalysisPassPass(Registry);
   initializeStackMapLivenessPass(Registry);
   initializeStackProtectorPass(Registry);
diff --git a/llvm/lib/CodeGen/StackColoring.cpp b/llvm/lib/CodeGen/StackColoring.cpp
index 341ec62..0be31d5 100644
--- a/llvm/lib/CodeGen/StackColoring.cpp
+++ b/llvm/lib/CodeGen/StackColoring.cpp
@@ -20,6 +20,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/CodeGen/StackColoring.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DepthFirstIterator.h"
@@ -376,7 +377,7 @@ namespace {
 
 /// StackColoring - A machine pass for merging disjoint stack allocations,
 /// marked by the LIFETIME_START and LIFETIME_END pseudo instructions.
-class StackColoring : public MachineFunctionPass {
+class StackColoring {
   MachineFrameInfo *MFI = nullptr;
   MachineFunction *MF = nullptr;
 
@@ -436,14 +437,8 @@ class StackColoring : public MachineFunctionPass {
   unsigned NumIterations;
 
 public:
-  static char ID;
-
-  StackColoring() : MachineFunctionPass(ID) {
-    initializeStackColoringPass(*PassRegistry::getPassRegistry());
-  }
-
-  void getAnalysisUsage(AnalysisUsage &AU) const override;
-  bool runOnMachineFunction(MachineFunction &Func) override;
+  StackColoring(SlotIndexes *Indexes) : Indexes(Indexes) {}
+  bool run(MachineFunction &Func);
 
 private:
   /// Used in collectMarkers
@@ -509,19 +504,29 @@ private:
   void expungeSlotMap(DenseMap<int, int> &SlotRemap, unsigned NumSlots);
 };
 
+class StackColoringLegacy : public MachineFunctionPass {
+public:
+  static char ID;
+
+  StackColoringLegacy() : MachineFunctionPass(ID) {}
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+  bool runOnMachineFunction(MachineFunction &Func) override;
+};
+
 } // end anonymous namespace
 
-char StackColoring::ID = 0;
+char StackColoringLegacy::ID = 0;
 
-char &llvm::StackColoringID = StackColoring::ID;
+char &llvm::StackColoringLegacyID = StackColoringLegacy::ID;
 
-INITIALIZE_PASS_BEGIN(StackColoring, DEBUG_TYPE,
+INITIALIZE_PASS_BEGIN(StackColoringLegacy, DEBUG_TYPE,
                       "Merge disjoint stack slots", false, false)
 INITIALIZE_PASS_DEPENDENCY(SlotIndexesWrapperPass)
-INITIALIZE_PASS_END(StackColoring, DEBUG_TYPE,
+INITIALIZE_PASS_END(StackColoringLegacy, DEBUG_TYPE,
                     "Merge disjoint stack slots", false, false)
 
-void StackColoring::getAnalysisUsage(AnalysisUsage &AU) const {
+void StackColoringLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<SlotIndexesWrapperPass>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
@@ -1178,12 +1183,30 @@ void StackColoring::expungeSlotMap(DenseMap<int, int> &SlotRemap,
   }
 }
 
-bool StackColoring::runOnMachineFunction(MachineFunction &Func) {
+bool StackColoringLegacy::runOnMachineFunction(MachineFunction &MF) {
+  if (skipFunction(MF.getFunction()))
+    return false;
+
+  StackColoring SC(&getAnalysis<SlotIndexesWrapperPass>().getSI());
+  return SC.run(MF);
+}
+
+PreservedAnalyses StackColoringPass::run(MachineFunction &MF,
+                                         MachineFunctionAnalysisManager &MFAM) {
+  if (MF.getFunction().hasOptNone())
+    return PreservedAnalyses::all();
+
+  StackColoring SC(&MFAM.getResult<SlotIndexesAnalysis>(MF));
+  if (SC.run(MF))
+    return PreservedAnalyses::none();
+  return PreservedAnalyses::all();
+}
+
+bool StackColoring::run(MachineFunction &Func) {
   LLVM_DEBUG(dbgs() << "********** Stack Coloring **********\n"
                     << "********** Function: " << Func.getName() << '\n');
   MF = &Func;
   MFI = &MF->getFrameInfo();
-  Indexes = &getAnalysis<SlotIndexesWrapperPass>().getSI();
   BlockLiveness.clear();
   BasicBlocks.clear();
   BasicBlockNumbering.clear();
@@ -1220,8 +1243,7 @@ bool StackColoring::runOnMachineFunction(MachineFunction &Func) {
 
   // Don't continue because there are not enough lifetime markers, or the
   // stack is too small, or we are told not to optimize the slots.
-  if (NumMarkers < 2 || TotalSize < 16 || DisableColoring ||
-      skipFunction(Func.getFunction())) {
+  if (NumMarkers < 2 || TotalSize < 16 || DisableColoring) {
     LLVM_DEBUG(dbgs() << "Will not try to merge slots.\n");
     return removeAllMarkers();
   }
diff --git a/llvm/lib/CodeGen/TargetPassConfig.cpp b/llvm/lib/CodeGen/TargetPassConfig.cpp
index 11a7752..cf9d63d 100644
--- a/llvm/lib/CodeGen/TargetPassConfig.cpp
+++ b/llvm/lib/CodeGen/TargetPassConfig.cpp
@@ -1287,7 +1287,7 @@ void TargetPassConfig::addMachineSSAOptimization() {
 
   // This pass merges large allocas. StackSlotColoring is a different pass
   // which merges spill slots.
-  addPass(&StackColoringID);
+  addPass(&StackColoringLegacyID);
 
   // If the target requests it, assign local variables to stack slots relative
   // to one another and simplify frame index references where possible.
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index f47f8ca..36c0cea 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -118,6 +118,7 @@
 #include "llvm/CodeGen/ShadowStackGCLowering.h"
 #include "llvm/CodeGen/SjLjEHPrepare.h"
 #include "llvm/CodeGen/SlotIndexes.h"
+#include "llvm/CodeGen/StackColoring.h"
 #include "llvm/CodeGen/StackProtector.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TwoAddressInstructionPass.h"
diff --git a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
index 2eb8b17..7d04cf3 100644
--- a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
@@ -445,7 +445,7 @@ void NVPTXPassConfig::addMachineSSAOptimization() {
 
   // This pass merges large allocas. StackSlotColoring is a different pass
   // which merges spill slots.
-  addPass(&StackColoringID);
+  addPass(&StackColoringLegacyID);
 
   // If the target requests it, assign local variables to stack slots relative
   // to one another and simplify frame index references where possible.
diff --git a/llvm/test/CodeGen/X86/PR37310.mir b/llvm/test/CodeGen/X86/PR37310.mir
index fa0368a..6c68f79 100644
--- a/llvm/test/CodeGen/X86/PR37310.mir
+++ b/llvm/test/CodeGen/X86/PR37310.mir
@@ -1,4 +1,5 @@
 # RUN: llc -mtriple=x86_64-unknown-linux-gnu -mcpu=corei7 -no-stack-coloring=false -run-pass stack-coloring -o - %s
+# RUN: llc -mtriple=x86_64-unknown-linux-gnu -mcpu=corei7 -no-stack-coloring=false -passes stack-coloring -o - %s
 
 # Test to insure that the liveness analysis in the StackColoring
 # pass gracefully handles statically unreachable blocks. See PR 37310.
diff --git a/llvm/test/CodeGen/X86/StackColoring-dbg-invariance.mir b/llvm/test/CodeGen/X86/StackColoring-dbg-invariance.mir
index 9391430..348a290 100644
--- a/llvm/test/CodeGen/X86/StackColoring-dbg-invariance.mir
+++ b/llvm/test/CodeGen/X86/StackColoring-dbg-invariance.mir
@@ -1,5 +1,6 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -mcpu=corei7 -no-stack-coloring=false -run-pass=stack-coloring -o - %s | FileCheck %s
+# RUN: llc -mcpu=corei7 -no-stack-coloring=false -passes=stack-coloring -o - %s | FileCheck %s
 
 # Difference between test_1 and test_2 is that there is a DBG_VALUE in test_1.
 # If transformation is debug invariant the resulting LEA instruction should be
diff --git a/llvm/test/CodeGen/X86/pr48064.mir b/llvm/test/CodeGen/X86/pr48064.mir
index a6c9271..9712a3ca 100644
--- a/llvm/test/CodeGen/X86/pr48064.mir
+++ b/llvm/test/CodeGen/X86/pr48064.mir
@@ -1,4 +1,5 @@
 # RUN: llc -mtriple="i386-pc-windows-msvc" -run-pass=stack-coloring %s -o - | FileCheck %s
+# RUN: llc -mtriple="i386-pc-windows-msvc" -passes=stack-coloring %s -o - | FileCheck %s
 
 # There is a problem with the exception handler, we found in windows, when set
 # LifetimeStartOnFirstUse=true for stack-coloring in default. Take the following