[AMDGPU][NewPM] Port GCNNSAReassign pass to new pass manager (#125034)

tests to be added while porting virtregrewrite and greedy regalloc

6 files changed, 93 insertions, 39 deletions

diff --git a/llvm/lib/Target/AMDGPU/AMDGPU.h b/llvm/lib/Target/AMDGPU/AMDGPU.h
index 4a0e5ef..42392e2 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPU.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPU.h
@@ -448,7 +448,7 @@ ModulePass *createAMDGPUOpenCLEnqueuedBlockLoweringLegacyPass();
 void initializeAMDGPUOpenCLEnqueuedBlockLoweringLegacyPass(PassRegistry &);
 extern char &AMDGPUOpenCLEnqueuedBlockLoweringLegacyID;
 
-void initializeGCNNSAReassignPass(PassRegistry &);
+void initializeGCNNSAReassignLegacyPass(PassRegistry &);
 extern char &GCNNSAReassignID;
 
 void initializeGCNPreRALongBranchRegLegacyPass(PassRegistry &);
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUPassRegistry.def b/llvm/lib/Target/AMDGPU/AMDGPUPassRegistry.def
index 225f847..fd1341e 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUPassRegistry.def
+++ b/llvm/lib/Target/AMDGPU/AMDGPUPassRegistry.def
@@ -100,6 +100,7 @@ MACHINE_FUNCTION_PASS("amdgpu-isel", AMDGPUISelDAGToDAGPass(*this))
 MACHINE_FUNCTION_PASS("amdgpu-pre-ra-long-branch-reg", GCNPreRALongBranchRegPass())
 MACHINE_FUNCTION_PASS("amdgpu-rewrite-partial-reg-uses", GCNRewritePartialRegUsesPass())
 MACHINE_FUNCTION_PASS("amdgpu-pre-ra-optimizations", GCNPreRAOptimizationsPass())
+MACHINE_FUNCTION_PASS("amdgpu-nsa-reassign", GCNNSAReassignPass())
 MACHINE_FUNCTION_PASS("gcn-dpp-combine", GCNDPPCombinePass())
 MACHINE_FUNCTION_PASS("si-fix-sgpr-copies", SIFixSGPRCopiesPass())
 MACHINE_FUNCTION_PASS("si-fix-vgpr-copies", SIFixVGPRCopiesPass())
@@ -120,7 +121,6 @@ MACHINE_FUNCTION_PASS("si-wqm", SIWholeQuadModePass())
 
 #define DUMMY_MACHINE_FUNCTION_PASS(NAME, CREATE_PASS)
 DUMMY_MACHINE_FUNCTION_PASS("amdgpu-insert-delay-alu", AMDGPUInsertDelayAluPass())
-DUMMY_MACHINE_FUNCTION_PASS("amdgpu-nsa-reassign", GCNNSAReassignPass())
 DUMMY_MACHINE_FUNCTION_PASS("amdgpu-pre-ra-optimizations", GCNPreRAOptimizationsPass())
 DUMMY_MACHINE_FUNCTION_PASS("amdgpu-rewrite-partial-reg-uses", GCNRewritePartialRegUsesPass())
 DUMMY_MACHINE_FUNCTION_PASS("amdgpu-set-wave-priority", AMDGPUSetWavePriorityPass())
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index 92ab106..7c9377e 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -32,6 +32,7 @@
 #include "AMDGPUWaitSGPRHazards.h"
 #include "GCNDPPCombine.h"
 #include "GCNIterativeScheduler.h"
+#include "GCNNSAReassign.h"
 #include "GCNPreRALongBranchReg.h"
 #include "GCNPreRAOptimizations.h"
 #include "GCNRewritePartialRegUses.h"
@@ -550,7 +551,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAMDGPUTarget() {
   initializeAMDGPUImageIntrinsicOptimizerPass(*PR);
   initializeAMDGPUPrintfRuntimeBindingPass(*PR);
   initializeAMDGPUResourceUsageAnalysisPass(*PR);
-  initializeGCNNSAReassignPass(*PR);
+  initializeGCNNSAReassignLegacyPass(*PR);
   initializeGCNPreRAOptimizationsLegacyPass(*PR);
   initializeGCNPreRALongBranchRegLegacyPass(*PR);
   initializeGCNRewritePartialRegUsesLegacyPass(*PR);
@@ -2112,6 +2113,12 @@ Error AMDGPUCodeGenPassBuilder::addInstSelector(AddMachinePass &addPass) const {
   return Error::success();
 }
 
+void AMDGPUCodeGenPassBuilder::addPreRewrite(AddMachinePass &addPass) const {
+  if (EnableRegReassign) {
+    addPass(GCNNSAReassignPass());
+  }
+}
+
 void AMDGPUCodeGenPassBuilder::addMachineSSAOptimization(
     AddMachinePass &addPass) const {
   Base::addMachineSSAOptimization(addPass);
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h
index 1455494..eb5a9ca 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h
@@ -177,6 +177,7 @@ public:
   void addILPOpts(AddMachinePass &) const;
   void addAsmPrinter(AddMachinePass &, CreateMCStreamer) const;
   Error addInstSelector(AddMachinePass &) const;
+  void addPreRewrite(AddMachinePass &) const;
   void addMachineSSAOptimization(AddMachinePass &) const;
   void addPostRegAlloc(AddMachinePass &) const;
 
diff --git a/llvm/lib/Target/AMDGPU/GCNNSAReassign.cpp b/llvm/lib/Target/AMDGPU/GCNNSAReassign.cpp
index 85e79aa..13eb0ca 100644
--- a/llvm/lib/Target/AMDGPU/GCNNSAReassign.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNNSAReassign.cpp
@@ -13,6 +13,7 @@
 ///
 //===----------------------------------------------------------------------===//
 
+#include "GCNNSAReassign.h"
 #include "AMDGPU.h"
 #include "GCNSubtarget.h"
 #include "SIMachineFunctionInfo.h"
@@ -34,26 +35,12 @@ STATISTIC(NumNSAConverted,
           "Number of NSA instructions changed to sequential");
 
 namespace {
-
-class GCNNSAReassign : public MachineFunctionPass {
+class GCNNSAReassignImpl {
 public:
-  static char ID;
-
-  GCNNSAReassign() : MachineFunctionPass(ID) {
-    initializeGCNNSAReassignPass(*PassRegistry::getPassRegistry());
-  }
-
-  bool runOnMachineFunction(MachineFunction &MF) override;
+  GCNNSAReassignImpl(VirtRegMap *VM, LiveRegMatrix *LM, LiveIntervals *LS)
+      : VRM(VM), LRM(LM), LIS(LS) {}
 
-  StringRef getPassName() const override { return "GCN NSA Reassign"; }
-
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.addRequired<LiveIntervalsWrapperPass>();
-    AU.addRequired<VirtRegMapWrapperLegacy>();
-    AU.addRequired<LiveRegMatrixWrapperLegacy>();
-    AU.setPreservesAll();
-    MachineFunctionPass::getAnalysisUsage(AU);
-  }
+  bool run(MachineFunction &MF);
 
 private:
   using NSA_Status = enum {
@@ -90,24 +77,43 @@ private:
   bool scavengeRegs(SmallVectorImpl<LiveInterval *> &Intervals) const;
 };
 
+class GCNNSAReassignLegacy : public MachineFunctionPass {
+public:
+  static char ID;
+
+  GCNNSAReassignLegacy() : MachineFunctionPass(ID) {
+    initializeGCNNSAReassignLegacyPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+  StringRef getPassName() const override { return "GCN NSA Reassign"; };
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<LiveIntervalsWrapperPass>();
+    AU.addRequired<VirtRegMapWrapperLegacy>();
+    AU.addRequired<LiveRegMatrixWrapperLegacy>();
+    AU.setPreservesAll();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+};
+
 } // End anonymous namespace.
 
-INITIALIZE_PASS_BEGIN(GCNNSAReassign, DEBUG_TYPE, "GCN NSA Reassign",
+INITIALIZE_PASS_BEGIN(GCNNSAReassignLegacy, DEBUG_TYPE, "GCN NSA Reassign",
                       false, false)
 INITIALIZE_PASS_DEPENDENCY(LiveIntervalsWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(VirtRegMapWrapperLegacy)
 INITIALIZE_PASS_DEPENDENCY(LiveRegMatrixWrapperLegacy)
-INITIALIZE_PASS_END(GCNNSAReassign, DEBUG_TYPE, "GCN NSA Reassign",
-                    false, false)
-
+INITIALIZE_PASS_END(GCNNSAReassignLegacy, DEBUG_TYPE, "GCN NSA Reassign", false,
+                    false)
 
-char GCNNSAReassign::ID = 0;
+char GCNNSAReassignLegacy::ID = 0;
 
-char &llvm::GCNNSAReassignID = GCNNSAReassign::ID;
+char &llvm::GCNNSAReassignID = GCNNSAReassignLegacy::ID;
 
-bool
-GCNNSAReassign::tryAssignRegisters(SmallVectorImpl<LiveInterval *> &Intervals,
-                                   unsigned StartReg) const {
+bool GCNNSAReassignImpl::tryAssignRegisters(
+    SmallVectorImpl<LiveInterval *> &Intervals, unsigned StartReg) const {
   unsigned NumRegs = Intervals.size();
 
   for (unsigned N = 0; N < NumRegs; ++N)
@@ -124,7 +130,7 @@ GCNNSAReassign::tryAssignRegisters(SmallVectorImpl<LiveInterval *> &Intervals,
   return true;
 }
 
-bool GCNNSAReassign::canAssign(unsigned StartReg, unsigned NumRegs) const {
+bool GCNNSAReassignImpl::canAssign(unsigned StartReg, unsigned NumRegs) const {
   for (unsigned N = 0; N < NumRegs; ++N) {
     unsigned Reg = StartReg + N;
     if (!MRI->isAllocatable(Reg))
@@ -139,8 +145,8 @@ bool GCNNSAReassign::canAssign(unsigned StartReg, unsigned NumRegs) const {
   return true;
 }
 
-bool
-GCNNSAReassign::scavengeRegs(SmallVectorImpl<LiveInterval *> &Intervals) const {
+bool GCNNSAReassignImpl::scavengeRegs(
+    SmallVectorImpl<LiveInterval *> &Intervals) const {
   unsigned NumRegs = Intervals.size();
 
   if (NumRegs > MaxNumVGPRs)
@@ -158,8 +164,8 @@ GCNNSAReassign::scavengeRegs(SmallVectorImpl<LiveInterval *> &Intervals) const {
   return false;
 }
 
-GCNNSAReassign::NSA_Status
-GCNNSAReassign::CheckNSA(const MachineInstr &MI, bool Fast) const {
+GCNNSAReassignImpl::NSA_Status
+GCNNSAReassignImpl::CheckNSA(const MachineInstr &MI, bool Fast) const {
   const AMDGPU::MIMGInfo *Info = AMDGPU::getMIMGInfo(MI.getOpcode());
   if (!Info)
     return NSA_Status::NOT_NSA;
@@ -235,16 +241,13 @@ GCNNSAReassign::CheckNSA(const MachineInstr &MI, bool Fast) const {
   return NSA ? NSA_Status::NON_CONTIGUOUS : NSA_Status::CONTIGUOUS;
 }
 
-bool GCNNSAReassign::runOnMachineFunction(MachineFunction &MF) {
+bool GCNNSAReassignImpl::run(MachineFunction &MF) {
   ST = &MF.getSubtarget<GCNSubtarget>();
   if (!ST->hasNSAEncoding() || !ST->hasNonNSAEncoding())
     return false;
 
   MRI = &MF.getRegInfo();
   TRI = ST->getRegisterInfo();
-  VRM = &getAnalysis<VirtRegMapWrapperLegacy>().getVRM();
-  LRM = &getAnalysis<LiveRegMatrixWrapperLegacy>().getLRM();
-  LIS = &getAnalysis<LiveIntervalsWrapperPass>().getLIS();
 
   const SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
   MaxNumVGPRs = ST->getMaxNumVGPRs(MF);
@@ -367,3 +370,24 @@ bool GCNNSAReassign::runOnMachineFunction(MachineFunction &MF) {
 
   return Changed;
 }
+
+bool GCNNSAReassignLegacy::runOnMachineFunction(MachineFunction &MF) {
+  auto *VRM = &getAnalysis<VirtRegMapWrapperLegacy>().getVRM();
+  auto *LRM = &getAnalysis<LiveRegMatrixWrapperLegacy>().getLRM();
+  auto *LIS = &getAnalysis<LiveIntervalsWrapperPass>().getLIS();
+
+  GCNNSAReassignImpl Impl(VRM, LRM, LIS);
+  return Impl.run(MF);
+}
+
+PreservedAnalyses
+GCNNSAReassignPass::run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM) {
+  auto &VRM = MFAM.getResult<VirtRegMapAnalysis>(MF);
+  auto &LRM = MFAM.getResult<LiveRegMatrixAnalysis>(MF);
+  auto &LIS = MFAM.getResult<LiveIntervalsAnalysis>(MF);
+
+  GCNNSAReassignImpl Impl(&VRM, &LRM, &LIS);
+  Impl.run(MF);
+  return PreservedAnalyses::all();
+}
diff --git a/llvm/lib/Target/AMDGPU/GCNNSAReassign.h b/llvm/lib/Target/AMDGPU/GCNNSAReassign.h
new file mode 100644
index 0000000..97a72e7
--- /dev/null
+++ b/llvm/lib/Target/AMDGPU/GCNNSAReassign.h
@@ -0,0 +1,22 @@
+//===- GCNNSAReassign.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_LIB_TARGET_AMDGPU_GCNNSAREASSIGN_H
+#define LLVM_LIB_TARGET_AMDGPU_GCNNSAREASSIGN_H
+
+#include "llvm/CodeGen/MachinePassManager.h"
+
+namespace llvm {
+class GCNNSAReassignPass : public PassInfoMixin<GCNNSAReassignPass> {
+public:
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+};
+} // namespace llvm
+
+#endif // LLVM_LIB_TARGET_AMDGPU_GCNNSAREASSIGN_H