[CodeGen][NewPM] Port RegisterUsageInfo to NPM

[optimisan]

And add to the codegen pipeline if ipra is enabled with a RequireAnalysisPass since this is a module pass.

[optimisan]

• [CodeGen][NewPM] Port RegUsageInfoPropagation pass to NPM #114010
• [CodeGen][NewPM] Port RegUsageInfoCollector pass to NPM #113874
• [CodeGen][NewPM] Port RegisterUsageInfo to NPM #113873 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[llvmbot]

@llvm/pr-subscribers-llvm-regalloc

Author: Akshat Oke (optimisan)

Changes

And add to the codegen pipeline if ipra is enabled with a RequireAnalysisPass since this is a module pass.

---

Full diff: https://github.com/llvm/llvm-project/pull/113873.diff

9 Files Affected:

• (modified) llvm/include/llvm/CodeGen/RegisterUsageInfo.h (+54-11)
• (modified) llvm/include/llvm/InitializePasses.h (+1-1)
• (modified) llvm/include/llvm/Passes/CodeGenPassBuilder.h (+4-2)
• (modified) llvm/include/llvm/Passes/MachinePassRegistry.def (+1)
• (modified) llvm/lib/CodeGen/RegUsageInfoCollector.cpp (+4-3)
• (modified) llvm/lib/CodeGen/RegUsageInfoPropagate.cpp (+4-3)
• (modified) llvm/lib/CodeGen/RegisterUsageInfo.cpp (+27-2)
• (modified) llvm/lib/Passes/PassBuilder.cpp (+1)
• (modified) llvm/lib/Passes/PassRegistry.def (+2)

diff --git a/llvm/include/llvm/CodeGen/RegisterUsageInfo.h b/llvm/include/llvm/CodeGen/RegisterUsageInfo.h
index aa1f5ef8110b0c..3073b62f37be7e 100644
--- 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 LLVMTargetMachine;
 
-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 LLVMTargetMachine &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 LLVMTargetMachine *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 PhysicalRegisterUsageInfoAnalysis
+    : public AnalysisInfoMixin<PhysicalRegisterUsageInfoAnalysis> {
+  friend AnalysisInfoMixin<PhysicalRegisterUsageInfoAnalysis>;
+  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
index 26f5d63553c5a8..f6f6797ec9f87c 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/llvm/include/llvm/InitializePasses.h
@@ -228,7 +228,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
index ad80c661147d6f..e5de62935a8e48 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -53,6 +53,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"
@@ -900,9 +901,10 @@ Error CodeGenPassBuilder<Derived, TargetMachineT>::addMachinePasses(
     addPass(LocalStackSlotAllocationPass());
   }
 
-  if (TM.Options.EnableIPRA)
+  if (TM.Options.EnableIPRA) {
+    addPass(RequireAnalysisPass<PhysicalRegisterUsageInfoAnalysis, 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
index 4f32a917738c13..183a777a93b9fa 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -29,6 +29,7 @@ MODULE_PASS("jmc-instrumenter", JMCInstrumenterPass())
 MODULE_PASS("lower-emutls", LowerEmuTLSPass())
 MODULE_PASS("pre-isel-intrinsic-lowering", PreISelIntrinsicLoweringPass())
 MODULE_PASS("shadow-stack-gc-lowering", ShadowStackGCLoweringPass())
+MODULE_PASS("print<regusage>", PhysicalRegisterUsageInfoPrinterPass(dbgs()))
 #undef MODULE_PASS
 
 #ifndef FUNCTION_ANALYSIS
diff --git a/llvm/lib/CodeGen/RegUsageInfoCollector.cpp b/llvm/lib/CodeGen/RegUsageInfoCollector.cpp
index ca5e0b428c4772..fc7664e2ba2357 100644
--- 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
index d356962e0d78a2..5ffe6acc83d601 100644
--- 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
index 51bac3fc0a233a..eae8d51617c26d 100644
--- 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 LLVMTargetMachine &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<PhysicalRegisterUsageInfoAnalysis>();
+  return !PAC.preservedWhenStateless();
+}
+
+AnalysisKey PhysicalRegisterUsageInfoAnalysis::Key;
+PhysicalRegisterUsageInfo
+PhysicalRegisterUsageInfoAnalysis::run(Module &M, ModuleAnalysisManager &) {
+  PhysicalRegisterUsageInfo PRUI;
+  PRUI.doInitialization(M);
+  return PRUI;
+}
+
+PreservedAnalyses
+PhysicalRegisterUsageInfoPrinterPass::run(Module &M,
+                                          ModuleAnalysisManager &AM) {
+  auto *PRUI = &AM.getResult<PhysicalRegisterUsageInfoAnalysis>(M);
+  PRUI->print(OS, &M);
+  return PreservedAnalyses::all();
+}
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index d1f75dfb5350a0..3d96b887afcf5a 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -118,6 +118,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
index 017ae311c55eb4..5e4610e3a717bd 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -30,6 +30,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-info", PhysicalRegisterUsageInfoAnalysis())
 MODULE_ANALYSIS("stack-safety", StackSafetyGlobalAnalysis())
 MODULE_ANALYSIS("verify", VerifierAnalysis())
 
@@ -127,6 +128,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<regusage>", PhysicalRegisterUsageInfoPrinterPass(dbgs()))
 MODULE_PASS("pseudo-probe", SampleProfileProbePass(TM))
 MODULE_PASS("pseudo-probe-update", PseudoProbeUpdatePass())
 MODULE_PASS("recompute-globalsaa", RecomputeGlobalsAAPass())

[arsenm]

Eventually we should get rid of this and just directly update the regmask used directly in the MIR

[optimisan]

I don't know much about regmasks, but do you mean forego setting it in the call instruction MO?

[arsenm]

No. I mean this analysis largely serves just as a place to hold allocated regmasks. There is a parallel mechanism for storing custom register masks in MachineRegisterInfo already. We need to consolidate these, and improve the serialization for them (i.e. introduce some kind of printed MachineFunction section for contained regmasks with an ID that the instructions can refer to them by inline, otherwise we'll have hundreds of lines printed for every call instruction)

[MatzeB]

From what I remember this analysis is basically equivalent to MachineRegisterInfo::UsedPhysRegMask however I think for interprocedural regalloc we did not want to keep the full MachineRegisterInfo around for every function in the module, and that's why this simpler analysis pass exists? (and I think this analysis is not used anyway when IPRA is not enabled)

[arsenm]

The problem is then you can't properly serialize the MIR. If your testcase depends on the IPRA state, -start-before to run the RA pipeline does not work

[paperchalice]

Just some nitpick comments.

[arsenm]

No. I mean this analysis largely serves just as a place to hold allocated regmasks. There is a parallel mechanism for storing custom register masks in MachineRegisterInfo already. We need to consolidate these, and improve the serialization for them (i.e. introduce some kind of printed MachineFunction section for contained regmasks with an ID that the instructions can refer to them by inline, otherwise we'll have hundreds of lines printed for every call instruction)

[optimisan]

Merge activity

• Nov 14, 8:37 AM EST: A user started a stack merge that includes this pull request via Graphite.
• Nov 14, 8:40 AM EST: Graphite couldn't merge this PR because it had conflicts with the trunk branch.
• Nov 14, 8:58 AM EST: A user started a stack merge that includes this pull request via Graphite.
• Nov 15, 12:17 AM EST: A user started a stack merge that includes this pull request via Graphite.
• Nov 15, 12:19 AM EST: A user merged this pull request with Graphite.