[CodeGen][NewPM] Port RegisterUsageInfo to NPM (#113873)

optimisan · web-flow · commit 7b54976d11a5 · 2024-11-15T10:49:00.000+05:30
And add to the codegen pipeline if ipra is enabled with a `RequireAnalysisPass` since this is a module pass.
diff --git a/llvm/include/llvm/CodeGen/RegisterUsageInfo.h b/llvm/include/llvm/CodeGen/RegisterUsageInfo.h
@@ -20,6 +20,7 @@
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/PassRegistry.h"
@@ -31,21 +32,14 @@ namespace llvm {
 class Function;
 class TargetMachine;
 
-class PhysicalRegisterUsageInfo : public ImmutablePass {
+class PhysicalRegisterUsageInfo {
 public:
-  static char ID;
-
-  PhysicalRegisterUsageInfo() : ImmutablePass(ID) {
-    PassRegistry &Registry = *PassRegistry::getPassRegistry();
-    initializePhysicalRegisterUsageInfoPass(Registry);
-  }
-
   /// Set TargetMachine which is used to print analysis.
   void setTargetMachine(const TargetMachine &TM);
 
-  bool doInitialization(Module &M) override;
+  bool doInitialization(Module &M);
 
-  bool doFinalization(Module &M) override;
+  bool doFinalization(Module &M);
 
   /// To store RegMask for given Function *.
   void storeUpdateRegUsageInfo(const Function &FP,
@@ -55,7 +49,10 @@ class PhysicalRegisterUsageInfo : public ImmutablePass {
   /// array if function is not known.
   ArrayRef<uint32_t> getRegUsageInfo(const Function &FP);
 
-  void print(raw_ostream &OS, const Module *M = nullptr) const override;
+  void print(raw_ostream &OS, const Module *M = nullptr) const;
+
+  bool invalidate(Module &M, const PreservedAnalyses &PA,
+                  ModuleAnalysisManager::Invalidator &Inv);
 
 private:
   /// A Dense map from Function * to RegMask.
@@ -66,6 +63,52 @@ class PhysicalRegisterUsageInfo : public ImmutablePass {
   const TargetMachine *TM = nullptr;
 };
 
+class PhysicalRegisterUsageInfoWrapperLegacy : public ImmutablePass {
+  std::unique_ptr<PhysicalRegisterUsageInfo> PRUI;
+
+public:
+  static char ID;
+  PhysicalRegisterUsageInfoWrapperLegacy() : ImmutablePass(ID) {
+    initializePhysicalRegisterUsageInfoWrapperLegacyPass(
+        *PassRegistry::getPassRegistry());
+  }
+
+  PhysicalRegisterUsageInfo &getPRUI() { return *PRUI; }
+  const PhysicalRegisterUsageInfo &getPRUI() const { return *PRUI; }
+
+  bool doInitialization(Module &M) override {
+    PRUI.reset(new PhysicalRegisterUsageInfo());
+    return PRUI->doInitialization(M);
+  }
+
+  bool doFinalization(Module &M) override { return PRUI->doFinalization(M); }
+
+  void print(raw_ostream &OS, const Module *M = nullptr) const override {
+    PRUI->print(OS, M);
+  }
+};
+
+class PhysicalRegisterUsageAnalysis
+    : public AnalysisInfoMixin<PhysicalRegisterUsageAnalysis> {
+  friend AnalysisInfoMixin<PhysicalRegisterUsageAnalysis>;
+  static AnalysisKey Key;
+
+public:
+  using Result = PhysicalRegisterUsageInfo;
+
+  PhysicalRegisterUsageInfo run(Module &M, ModuleAnalysisManager &);
+};
+
+class PhysicalRegisterUsageInfoPrinterPass
+    : public PassInfoMixin<PhysicalRegisterUsageInfoPrinterPass> {
+  raw_ostream &OS;
+
+public:
+  explicit PhysicalRegisterUsageInfoPrinterPass(raw_ostream &OS) : OS(OS) {}
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+  static bool isRequired() { return true; }
+};
+
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_REGISTERUSAGEINFO_H
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
@@ -229,7 +229,7 @@ void initializePartiallyInlineLibCallsLegacyPassPass(PassRegistry &);
 void initializePatchableFunctionPass(PassRegistry &);
 void initializePeepholeOptimizerPass(PassRegistry &);
 void initializePhiValuesWrapperPassPass(PassRegistry &);
-void initializePhysicalRegisterUsageInfoPass(PassRegistry &);
+void initializePhysicalRegisterUsageInfoWrapperLegacyPass(PassRegistry &);
 void initializePlaceBackedgeSafepointsLegacyPassPass(PassRegistry &);
 void initializePostDomOnlyPrinterWrapperPassPass(PassRegistry &);
 void initializePostDomOnlyViewerWrapperPassPass(PassRegistry &);
diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -54,6 +54,7 @@
 #include "llvm/CodeGen/PHIElimination.h"
 #include "llvm/CodeGen/PreISelIntrinsicLowering.h"
 #include "llvm/CodeGen/RegAllocFast.h"
+#include "llvm/CodeGen/RegisterUsageInfo.h"
 #include "llvm/CodeGen/ReplaceWithVeclib.h"
 #include "llvm/CodeGen/SafeStack.h"
 #include "llvm/CodeGen/SelectOptimize.h"
@@ -906,9 +907,10 @@ Error CodeGenPassBuilder<Derived, TargetMachineT>::addMachinePasses(
     addPass(LocalStackSlotAllocationPass());
   }
 
-  if (TM.Options.EnableIPRA)
+  if (TM.Options.EnableIPRA) {
+    addPass(RequireAnalysisPass<PhysicalRegisterUsageAnalysis, Module>());
     addPass(RegUsageInfoPropagationPass());
-
+  }
   // Run pre-ra passes.
   derived().addPreRegAlloc(addPass);
 
diff --git a/llvm/include/llvm/Passes/MachinePassRegistry.def b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -28,6 +28,7 @@ MODULE_PASS("global-merge", GlobalMergePass(TM, GlobalMergeOptions()))
 MODULE_PASS("jmc-instrumenter", JMCInstrumenterPass())
 MODULE_PASS("lower-emutls", LowerEmuTLSPass())
 MODULE_PASS("pre-isel-intrinsic-lowering", PreISelIntrinsicLoweringPass())
+MODULE_PASS("print<regusage>", PhysicalRegisterUsageInfoPrinterPass(dbgs()))
 MODULE_PASS("shadow-stack-gc-lowering", ShadowStackGCLoweringPass())
 MODULE_PASS("global-merge-func", GlobalMergeFuncPass())
 #undef MODULE_PASS
diff --git a/llvm/lib/CodeGen/RegUsageInfoCollector.cpp b/llvm/lib/CodeGen/RegUsageInfoCollector.cpp
@@ -48,7 +48,7 @@ class RegUsageInfoCollector : public MachineFunctionPass {
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.addRequired<PhysicalRegisterUsageInfo>();
+    AU.addRequired<PhysicalRegisterUsageInfoWrapperLegacy>();
     AU.setPreservesAll();
     MachineFunctionPass::getAnalysisUsage(AU);
   }
@@ -68,7 +68,7 @@ char RegUsageInfoCollector::ID = 0;
 
 INITIALIZE_PASS_BEGIN(RegUsageInfoCollector, "RegUsageInfoCollector",
                       "Register Usage Information Collector", false, false)
-INITIALIZE_PASS_DEPENDENCY(PhysicalRegisterUsageInfo)
+INITIALIZE_PASS_DEPENDENCY(PhysicalRegisterUsageInfoWrapperLegacy)
 INITIALIZE_PASS_END(RegUsageInfoCollector, "RegUsageInfoCollector",
                     "Register Usage Information Collector", false, false)
 
@@ -129,7 +129,8 @@ bool RegUsageInfoCollector::runOnMachineFunction(MachineFunction &MF) {
 
   const Function &F = MF.getFunction();
 
-  PhysicalRegisterUsageInfo &PRUI = getAnalysis<PhysicalRegisterUsageInfo>();
+  PhysicalRegisterUsageInfo &PRUI =
+      getAnalysis<PhysicalRegisterUsageInfoWrapperLegacy>().getPRUI();
   PRUI.setTargetMachine(TM);
 
   LLVM_DEBUG(dbgs() << "Clobbered Registers: ");
diff --git a/llvm/lib/CodeGen/RegUsageInfoPropagate.cpp b/llvm/lib/CodeGen/RegUsageInfoPropagate.cpp
@@ -50,7 +50,7 @@ class RegUsageInfoPropagation : public MachineFunctionPass {
   bool runOnMachineFunction(MachineFunction &MF) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.addRequired<PhysicalRegisterUsageInfo>();
+    AU.addRequired<PhysicalRegisterUsageInfoWrapperLegacy>();
     AU.setPreservesAll();
     MachineFunctionPass::getAnalysisUsage(AU);
   }
@@ -75,7 +75,7 @@ class RegUsageInfoPropagation : public MachineFunctionPass {
 
 INITIALIZE_PASS_BEGIN(RegUsageInfoPropagation, "reg-usage-propagation",
                       RUIP_NAME, false, false)
-INITIALIZE_PASS_DEPENDENCY(PhysicalRegisterUsageInfo)
+INITIALIZE_PASS_DEPENDENCY(PhysicalRegisterUsageInfoWrapperLegacy)
 INITIALIZE_PASS_END(RegUsageInfoPropagation, "reg-usage-propagation",
                     RUIP_NAME, false, false)
 
@@ -97,7 +97,8 @@ static const Function *findCalledFunction(const Module &M,
 
 bool RegUsageInfoPropagation::runOnMachineFunction(MachineFunction &MF) {
   const Module &M = *MF.getFunction().getParent();
-  PhysicalRegisterUsageInfo *PRUI = &getAnalysis<PhysicalRegisterUsageInfo>();
+  PhysicalRegisterUsageInfo *PRUI =
+      &getAnalysis<PhysicalRegisterUsageInfoWrapperLegacy>().getPRUI();
 
   LLVM_DEBUG(dbgs() << " ++++++++++++++++++++ " << getPassName()
                     << " ++++++++++++++++++++  \n");
diff --git a/llvm/lib/CodeGen/RegisterUsageInfo.cpp b/llvm/lib/CodeGen/RegisterUsageInfo.cpp
@@ -16,8 +16,10 @@
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/IR/Analysis.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/raw_ostream.h"
@@ -32,10 +34,10 @@ static cl::opt<bool> DumpRegUsage(
     "print-regusage", cl::init(false), cl::Hidden,
     cl::desc("print register usage details collected for analysis."));
 
-INITIALIZE_PASS(PhysicalRegisterUsageInfo, "reg-usage-info",
+INITIALIZE_PASS(PhysicalRegisterUsageInfoWrapperLegacy, "reg-usage-info",
                 "Register Usage Information Storage", false, true)
 
-char PhysicalRegisterUsageInfo::ID = 0;
+char PhysicalRegisterUsageInfoWrapperLegacy::ID = 0;
 
 void PhysicalRegisterUsageInfo::setTargetMachine(const TargetMachine &TM) {
   this->TM = &TM;
@@ -97,3 +99,26 @@ void PhysicalRegisterUsageInfo::print(raw_ostream &OS, const Module *M) const {
     OS << "\n";
   }
 }
+
+bool PhysicalRegisterUsageInfo::invalidate(
+    Module &M, const PreservedAnalyses &PA,
+    ModuleAnalysisManager::Invalidator &) {
+  auto PAC = PA.getChecker<PhysicalRegisterUsageAnalysis>();
+  return !PAC.preservedWhenStateless();
+}
+
+AnalysisKey PhysicalRegisterUsageAnalysis::Key;
+PhysicalRegisterUsageInfo
+PhysicalRegisterUsageAnalysis::run(Module &M, ModuleAnalysisManager &) {
+  PhysicalRegisterUsageInfo PRUI;
+  PRUI.doInitialization(M);
+  return PRUI;
+}
+
+PreservedAnalyses
+PhysicalRegisterUsageInfoPrinterPass::run(Module &M,
+                                          ModuleAnalysisManager &AM) {
+  auto *PRUI = &AM.getResult<PhysicalRegisterUsageAnalysis>(M);
+  PRUI->print(OS, &M);
+  return PreservedAnalyses::all();
+}
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
@@ -120,6 +120,7 @@
 #include "llvm/CodeGen/PHIElimination.h"
 #include "llvm/CodeGen/PreISelIntrinsicLowering.h"
 #include "llvm/CodeGen/RegAllocFast.h"
+#include "llvm/CodeGen/RegisterUsageInfo.h"
 #include "llvm/CodeGen/SafeStack.h"
 #include "llvm/CodeGen/SelectOptimize.h"
 #include "llvm/CodeGen/ShadowStackGCLowering.h"
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
@@ -31,6 +31,7 @@ MODULE_ANALYSIS("module-summary", ModuleSummaryIndexAnalysis())
 MODULE_ANALYSIS("no-op-module", NoOpModuleAnalysis())
 MODULE_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC))
 MODULE_ANALYSIS("profile-summary", ProfileSummaryAnalysis())
+MODULE_ANALYSIS("reg-usage", PhysicalRegisterUsageAnalysis())
 MODULE_ANALYSIS("stack-safety", StackSafetyGlobalAnalysis())
 MODULE_ANALYSIS("verify", VerifierAnalysis())
 
@@ -129,6 +130,7 @@ MODULE_PASS("print<dxil-metadata>", DXILMetadataAnalysisPrinterPass(dbgs()))
 MODULE_PASS("print<dxil-resource>", DXILResourcePrinterPass(dbgs()))
 MODULE_PASS("print<inline-advisor>", InlineAdvisorAnalysisPrinterPass(dbgs()))
 MODULE_PASS("print<module-debuginfo>", ModuleDebugInfoPrinterPass(dbgs()))
+MODULE_PASS("print<reg-usage>", PhysicalRegisterUsageInfoPrinterPass(dbgs()))
 MODULE_PASS("pseudo-probe", SampleProfileProbePass(TM))
 MODULE_PASS("pseudo-probe-update", PseudoProbeUpdatePass())
 MODULE_PASS("recompute-globalsaa", RecomputeGlobalsAAPass())
