[CodeGen][NPM] Port VirtRegRewriter to NPM (#130564)

llvm/include/llvm/CodeGen/VirtRegMap.h

@@ -235,6 +235,21 @@ public:
                         MachineFunctionAnalysisManager &MFAM);
   static bool isRequired() { return true; }
 };
+
+class VirtRegRewriterPass : public PassInfoMixin<VirtRegRewriterPass> {
+  bool ClearVirtRegs = true;
+
+public:
+  VirtRegRewriterPass(bool ClearVirtRegs = true)
+      : ClearVirtRegs(ClearVirtRegs) {}
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+
+  static bool isRequired() { return true; }
+
+  void printPipeline(raw_ostream &OS, function_ref<StringRef(StringRef)>) const;
+};
+
 } // end llvm namespace
 
 #endif // LLVM_CODEGEN_VIRTREGMAP_H

llvm/include/llvm/InitializePasses.h

@@ -319,7 +319,7 @@ void initializeUnreachableBlockElimLegacyPassPass(PassRegistry &);
 void initializeUnreachableMachineBlockElimLegacyPass(PassRegistry &);
 void initializeVerifierLegacyPassPass(PassRegistry &);
 void initializeVirtRegMapWrapperLegacyPass(PassRegistry &);
-void initializeVirtRegRewriterPass(PassRegistry &);
+void initializeVirtRegRewriterLegacyPass(PassRegistry &);
 void initializeWasmEHPreparePass(PassRegistry &);
 void initializeWinEHPreparePass(PassRegistry &);
 void initializeWriteBitcodePassPass(PassRegistry &);

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -94,6 +94,7 @@
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TwoAddressInstructionPass.h"
 #include "llvm/CodeGen/UnreachableBlockElim.h"
+#include "llvm/CodeGen/VirtRegMap.h"
 #include "llvm/CodeGen/WasmEHPrepare.h"
 #include "llvm/CodeGen/WinEHPrepare.h"
 #include "llvm/CodeGen/XRayInstrumentation.h"

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -262,6 +262,12 @@ MACHINE_FUNCTION_PASS_WITH_PARAMS(
       return parseRegAllocGreedyFilterFunc(*PB, Params);
     }, "reg-filter"
 )
+
+MACHINE_FUNCTION_PASS_WITH_PARAMS(
+    "virt-reg-rewriter", "VirtRegRewriterPass",
+    [](bool ClearVirtRegs) { return VirtRegRewriterPass(ClearVirtRegs); },
+    parseVirtRegRewriterPassOptions, "no-clear-vregs;clear-vregs")
+
 #undef MACHINE_FUNCTION_PASS_WITH_PARAMS
 
 // After a pass is converted to new pass manager, its entry should be moved from
@@ -320,5 +326,4 @@ DUMMY_MACHINE_FUNCTION_PASS("regbankselect", RegBankSelectPass)
 DUMMY_MACHINE_FUNCTION_PASS("reset-machine-function", ResetMachineFunctionPass)
 DUMMY_MACHINE_FUNCTION_PASS("stackmap-liveness", StackMapLivenessPass)
 DUMMY_MACHINE_FUNCTION_PASS("unpack-mi-bundles", UnpackMachineBundlesPass)
-DUMMY_MACHINE_FUNCTION_PASS("virtregrewriter", VirtRegRewriterPass)
 #undef DUMMY_MACHINE_FUNCTION_PASS

llvm/lib/CodeGen/CodeGen.cpp

@@ -142,7 +142,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeUnreachableBlockElimLegacyPassPass(Registry);
   initializeUnreachableMachineBlockElimLegacyPass(Registry);
   initializeVirtRegMapWrapperLegacyPass(Registry);
-  initializeVirtRegRewriterPass(Registry);
+  initializeVirtRegRewriterLegacyPass(Registry);
   initializeWasmEHPreparePass(Registry);
   initializeWinEHPreparePass(Registry);
   initializeXRayInstrumentationLegacyPass(Registry);

llvm/lib/CodeGen/VirtRegMap.cpp

@@ -197,7 +197,7 @@ VirtRegMap VirtRegMapAnalysis::run(MachineFunction &MF,
 //
 namespace {
 
-class VirtRegRewriter : public MachineFunctionPass {
+class VirtRegRewriter {
   MachineFunction *MF = nullptr;
   const TargetRegisterInfo *TRI = nullptr;
   const TargetInstrInfo *TII = nullptr;
@@ -223,9 +223,21 @@ class VirtRegRewriter : public MachineFunctionPass {
 
 public:
   static char ID;
-  VirtRegRewriter(bool ClearVirtRegs_ = true) :
-    MachineFunctionPass(ID),
-    ClearVirtRegs(ClearVirtRegs_) {}
+  VirtRegRewriter(bool ClearVirtRegs, SlotIndexes *Indexes, LiveIntervals *LIS,
+                  LiveRegMatrix *LRM, VirtRegMap *VRM,
+                  LiveDebugVariables *DebugVars)
+      : Indexes(Indexes), LIS(LIS), LRM(LRM), VRM(VRM), DebugVars(DebugVars),
+        ClearVirtRegs(ClearVirtRegs) {}
+
+  bool run(MachineFunction &);
+};
+
+class VirtRegRewriterLegacy : public MachineFunctionPass {
+public:
+  static char ID;
+  bool ClearVirtRegs;
+  VirtRegRewriterLegacy(bool ClearVirtRegs = true)
+      : MachineFunctionPass(ID), ClearVirtRegs(ClearVirtRegs) {}
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
@@ -243,11 +255,11 @@ public:
 
 } // end anonymous namespace
 
-char VirtRegRewriter::ID = 0;
+char VirtRegRewriterLegacy::ID = 0;
 
-char &llvm::VirtRegRewriterID = VirtRegRewriter::ID;
+char &llvm::VirtRegRewriterID = VirtRegRewriterLegacy::ID;
 
-INITIALIZE_PASS_BEGIN(VirtRegRewriter, "virtregrewriter",
+INITIALIZE_PASS_BEGIN(VirtRegRewriterLegacy, "virtregrewriter",
                       "Virtual Register Rewriter", false, false)
 INITIALIZE_PASS_DEPENDENCY(SlotIndexesWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LiveIntervalsWrapperPass)
@@ -255,10 +267,10 @@ INITIALIZE_PASS_DEPENDENCY(LiveDebugVariablesWrapperLegacy)
 INITIALIZE_PASS_DEPENDENCY(LiveRegMatrixWrapperLegacy)
 INITIALIZE_PASS_DEPENDENCY(LiveStacksWrapperLegacy)
 INITIALIZE_PASS_DEPENDENCY(VirtRegMapWrapperLegacy)
-INITIALIZE_PASS_END(VirtRegRewriter, "virtregrewriter",
+INITIALIZE_PASS_END(VirtRegRewriterLegacy, "virtregrewriter",
                     "Virtual Register Rewriter", false, false)
 
-void VirtRegRewriter::getAnalysisUsage(AnalysisUsage &AU) const {
+void VirtRegRewriterLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesCFG();
   AU.addRequired<LiveIntervalsWrapperPass>();
   AU.addPreserved<LiveIntervalsWrapperPass>();
@@ -276,16 +288,50 @@ void VirtRegRewriter::getAnalysisUsage(AnalysisUsage &AU) const {
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
-bool VirtRegRewriter::runOnMachineFunction(MachineFunction &fn) {
+bool VirtRegRewriterLegacy::runOnMachineFunction(MachineFunction &MF) {
+  VirtRegMap &VRM = getAnalysis<VirtRegMapWrapperLegacy>().getVRM();
+  LiveIntervals &LIS = getAnalysis<LiveIntervalsWrapperPass>().getLIS();
+  LiveRegMatrix &LRM = getAnalysis<LiveRegMatrixWrapperLegacy>().getLRM();
+  SlotIndexes &Indexes = getAnalysis<SlotIndexesWrapperPass>().getSI();
+  LiveDebugVariables &DebugVars =
+      getAnalysis<LiveDebugVariablesWrapperLegacy>().getLDV();
+
+  VirtRegRewriter R(ClearVirtRegs, &Indexes, &LIS, &LRM, &VRM, &DebugVars);
+  return R.run(MF);
+}
+
+PreservedAnalyses
+VirtRegRewriterPass::run(MachineFunction &MF,
+                         MachineFunctionAnalysisManager &MFAM) {
+  VirtRegMap &VRM = MFAM.getResult<VirtRegMapAnalysis>(MF);
+  LiveIntervals &LIS = MFAM.getResult<LiveIntervalsAnalysis>(MF);
+  LiveRegMatrix &LRM = MFAM.getResult<LiveRegMatrixAnalysis>(MF);
+  SlotIndexes &Indexes = MFAM.getResult<SlotIndexesAnalysis>(MF);
+  LiveDebugVariables &DebugVars =
+      MFAM.getResult<LiveDebugVariablesAnalysis>(MF);
+
+  VirtRegRewriter R(ClearVirtRegs, &Indexes, &LIS, &LRM, &VRM, &DebugVars);
+  if (!R.run(MF))
+    return PreservedAnalyses::all();
+
+  auto PA = getMachineFunctionPassPreservedAnalyses();
+  PA.preserveSet<CFGAnalyses>();
+  PA.preserve<LiveIntervalsAnalysis>();
+  PA.preserve<SlotIndexesAnalysis>();
+  PA.preserve<LiveStacksAnalysis>();
+  // LiveDebugVariables is preserved by default, so clear it
+  // if this VRegRewriter is the last one in the pipeline.
+  if (ClearVirtRegs)
+    PA.abandon<LiveDebugVariablesAnalysis>();
+  return PA;
+}
+
+bool VirtRegRewriter::run(MachineFunction &fn) {
   MF = &fn;
   TRI = MF->getSubtarget().getRegisterInfo();
   TII = MF->getSubtarget().getInstrInfo();
   MRI = &MF->getRegInfo();
-  Indexes = &getAnalysis<SlotIndexesWrapperPass>().getSI();
-  LIS = &getAnalysis<LiveIntervalsWrapperPass>().getLIS();
-  LRM = &getAnalysis<LiveRegMatrixWrapperLegacy>().getLRM();
-  VRM = &getAnalysis<VirtRegMapWrapperLegacy>().getVRM();
-  DebugVars = &getAnalysis<LiveDebugVariablesWrapperLegacy>().getLDV();
+
   LLVM_DEBUG(dbgs() << "********** REWRITE VIRTUAL REGISTERS **********\n"
                     << "********** Function: " << MF->getName() << '\n');
   LLVM_DEBUG(VRM->dump());
@@ -726,6 +772,13 @@ void VirtRegRewriter::rewrite() {
   RewriteRegs.clear();
 }
 
-FunctionPass *llvm::createVirtRegRewriter(bool ClearVirtRegs) {
-  return new VirtRegRewriter(ClearVirtRegs);
+void VirtRegRewriterPass::printPipeline(
+    raw_ostream &OS, function_ref<StringRef(StringRef)>) const {
+  OS << "virt-reg-rewriter";
+  if (!ClearVirtRegs)
+    OS << "<no-clear-vregs>";
+}
+
+FunctionPass *llvm::createVirtRegRewriter(bool ClearVirtRegs) {
+  return new VirtRegRewriterLegacy(ClearVirtRegs);
 }

llvm/lib/Passes/PassBuilder.cpp

@@ -1505,6 +1505,19 @@ Expected<bool> parseMachineBlockPlacementPassOptions(StringRef Params) {
           inconvertibleErrorCode());
   }
   return AllowTailMerge;
+};
+
+Expected<bool> parseVirtRegRewriterPassOptions(StringRef Params) {
+  bool ClearVirtRegs = true;
+  if (!Params.empty()) {
+    ClearVirtRegs = !Params.consume_front("no-");
+    if (Params != "clear-vregs")
+      return make_error<StringError>(
+          formatv("invalid VirtRegRewriter pass parameter '{0}' ", Params)
+              .str(),
+          inconvertibleErrorCode());
+  }
+  return ClearVirtRegs;
 }
 
 } // namespace

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

@@ -2160,6 +2160,52 @@ void AMDGPUCodeGenPassBuilder::addMachineSSAOptimization(
   addPass(SIShrinkInstructionsPass());
 }
 
+
+
+Error AMDGPUCodeGenPassBuilder::addRegAssignmentOptimized(
+    AddMachinePass &addPass) const {
+  // TODO: Check --regalloc-npm option
+
+  addPass(GCNPreRALongBranchRegPass());
+
+  addPass(RAGreedyPass({onlyAllocateSGPRs, "sgpr"}));
+
+  // Commit allocated register changes. This is mostly necessary because too
+  // many things rely on the use lists of the physical registers, such as the
+  // verifier. This is only necessary with allocators which use LiveIntervals,
+  // since FastRegAlloc does the replacements itself.
+  addPass(VirtRegRewriterPass(false));
+
+  // At this point, the sgpr-regalloc has been done and it is good to have the
+  // stack slot coloring to try to optimize the SGPR spill stack indices before
+  // attempting the custom SGPR spill lowering.
+  addPass(StackSlotColoringPass());
+
+  // Equivalent of PEI for SGPRs.
+  addPass(SILowerSGPRSpillsPass());
+
+  // To Allocate wwm registers used in whole quad mode operations (for shaders).
+  addPass(SIPreAllocateWWMRegsPass());
+
+  // For allocating other wwm register operands.
+  // addRegAlloc<RAGreedyPass>(addPass, RegAllocPhase::WWM);
+  addPass(RAGreedyPass({onlyAllocateWWMRegs, "wwm"}));
+  addPass(SILowerWWMCopiesPass());
+  addPass(VirtRegRewriterPass(false));
+  addPass(AMDGPUReserveWWMRegsPass());
+
+  // For allocating per-thread VGPRs.
+  // addRegAlloc<RAGreedyPass>(addPass, RegAllocPhase::VGPR);
+  addPass(RAGreedyPass({onlyAllocateVGPRs, "vgpr"}));
+
+
+  addPreRewrite(addPass);
+  addPass(VirtRegRewriterPass(true));
+
+  // TODO: addPass(AMDGPUMarkLastScratchLoadPass());
+  return Error::success();
+}
+
 void AMDGPUCodeGenPassBuilder::addPostRegAlloc(AddMachinePass &addPass) const {
   addPass(SIFixVGPRCopiesPass());
   if (TM.getOptLevel() > CodeGenOptLevel::None)

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h

@@ -181,6 +181,7 @@ public:
   void addMachineSSAOptimization(AddMachinePass &) const;
   void addPostRegAlloc(AddMachinePass &) const;
   void addPreEmitPass(AddMachinePass &) const;
+  Error addRegAssignmentOptimized(AddMachinePass &) const;
 
   /// Check if a pass is enabled given \p Opt option. The option always
   /// overrides defaults if explicitly used. Otherwise its default will be used

llvm/test/CodeGen/AMDGPU/alloc-aligned-tuples-gfx90a.mir

@@ -1,4 +1,6 @@
 # RUN: llc -mtriple=amdgcn -mcpu=gfx90a -start-before=greedy,0 -stop-after=virtregrewriter,2 -verify-machineinstrs -o - %s | FileCheck --check-prefixes=GCN,GFX90A %s
+
+# RUN: llc -enable-new-pm -mtriple=amdgcn -mcpu=gfx90a -passes='greedy<sgpr>,virt-reg-rewriter<no-clear-vregs>,stack-slot-coloring,si-lower-sgpr-spills,si-pre-allocate-wwm-regs,greedy<wwm>,si-lower-wwm-copies,virt-reg-rewriter<no-clear-vregs>,amdgpu-reserve-wwm-regs,greedy<vgpr>,amdgpu-nsa-reassign,virt-reg-rewriter' -o - %s | FileCheck --check-prefixes=GCN,GFX90A %s
 # Using the unaligned vector tuples are OK as long as they aren't used
 # in a real instruction.
 

llvm/test/CodeGen/AMDGPU/fold-restore-undef-use.mir

@@ -1,6 +1,8 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx90a -stress-regalloc=4 -verify-regalloc -start-before=greedy,0 -stop-after=virtregrewriter,0 %s -o - | FileCheck %s
 
+# RUN: llc -enable-new-pm -mtriple=amdgcn-amd-amdhsa -mcpu=gfx90a -stress-regalloc=4 -verify-regalloc -passes="greedy<sgpr>,virt-reg-rewriter<no-clear-vregs>" %s -o - | FileCheck %s
+
 # Check that we don't generate *** Bad machine code: Instruction loads
 # from dead spill slot ***
 

llvm/test/CodeGen/AMDGPU/greedy-remark-crash-unassigned-reg.mir

@@ -2,6 +2,10 @@
 # RUN:   -start-before=greedy,0 -stop-after=virtregrewriter,0 -pass-remarks='.*' -pass-remarks-output=%t.yaml -o /dev/null %s
 # RUN: FileCheck %s < %t.yaml
 
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 \
+# RUN:   -passes='greedy<sgpr>,virt-reg-rewriter<no-clear-vregs>' -pass-remarks='.*' -pass-remarks-output=%t.yaml -o /dev/null %s
+# RUN: FileCheck %s < %t.yaml
+
 # CHECK: Name:            SpillReloadCopies
 # CHECK-NEXT: Function:        func
 # CHECK-NEXT: Args:

llvm/test/CodeGen/X86/pr30821.mir

@@ -1,4 +1,5 @@
 # RUN: llc -x mir < %s -run-pass=greedy,virtregrewriter,stack-slot-coloring | FileCheck %s
+# RUN: llc -x mir < %s -passes=greedy,virt-reg-rewriter,stack-slot-coloring | FileCheck %s
 --- |
   target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
   target triple = "x86_64-unknown-linux-gnu"