[AMDGPU] Port AMDGPURewriteUndefForPHI to new pass manager (#66008)

This patch ports the AMDGPURewriteUndefForPHI pass to the new pass
manager. With this, the pass is supported under both the legacy and the
new pass managers.

---------

Co-authored-by: Jun Wang <[email protected]>

Author: jwanggit86

llvm/lib/Target/AMDGPU/AMDGPU.h

@@ -287,9 +287,19 @@ extern char &AMDGPURemoveIncompatibleFunctionsID;
 void initializeAMDGPULateCodeGenPreparePass(PassRegistry &);
 extern char &AMDGPULateCodeGenPrepareID;
 
-FunctionPass *createAMDGPURewriteUndefForPHIPass();
-void initializeAMDGPURewriteUndefForPHIPass(PassRegistry &);
-extern char &AMDGPURewriteUndefForPHIPassID;
+FunctionPass *createAMDGPURewriteUndefForPHILegacyPass();
+void initializeAMDGPURewriteUndefForPHILegacyPass(PassRegistry &);
+extern char &AMDGPURewriteUndefForPHILegacyPassID;
+
+class AMDGPURewriteUndefForPHIPass
+    : public PassInfoMixin<AMDGPURewriteUndefForPHIPass> {
+private:
+  TargetMachine &TM;
+
+public:
+  AMDGPURewriteUndefForPHIPass(TargetMachine &TM) : TM(TM){};
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
 
 void initializeSIAnnotateControlFlowPass(PassRegistry&);
 extern char &SIAnnotateControlFlowPassID;

llvm/lib/Target/AMDGPU/AMDGPURewriteUndefForPHI.cpp

@@ -69,11 +69,11 @@ using namespace llvm;
 
 namespace {
 
-class AMDGPURewriteUndefForPHI : public FunctionPass {
+class AMDGPURewriteUndefForPHILegacy : public FunctionPass {
 public:
   static char ID;
-  AMDGPURewriteUndefForPHI() : FunctionPass(ID) {
-    initializeAMDGPURewriteUndefForPHIPass(*PassRegistry::getPassRegistry());
+  AMDGPURewriteUndefForPHILegacy() : FunctionPass(ID) {
+    initializeAMDGPURewriteUndefForPHILegacyPass(*PassRegistry::getPassRegistry());
   }
   bool runOnFunction(Function &F) override;
   StringRef getPassName() const override {
@@ -91,13 +91,13 @@ class AMDGPURewriteUndefForPHI : public FunctionPass {
 };
 
 } // end anonymous namespace
-char AMDGPURewriteUndefForPHI::ID = 0;
+char AMDGPURewriteUndefForPHILegacy::ID = 0;
 
-INITIALIZE_PASS_BEGIN(AMDGPURewriteUndefForPHI, DEBUG_TYPE,
+INITIALIZE_PASS_BEGIN(AMDGPURewriteUndefForPHILegacy, DEBUG_TYPE,
                       "Rewrite undef for PHI", false, false)
 INITIALIZE_PASS_DEPENDENCY(UniformityInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
-INITIALIZE_PASS_END(AMDGPURewriteUndefForPHI, DEBUG_TYPE,
+INITIALIZE_PASS_END(AMDGPURewriteUndefForPHILegacy, DEBUG_TYPE,
                     "Rewrite undef for PHI", false, false)
 
 bool rewritePHIs(Function &F, UniformityInfo &UA, DominatorTree *DT) {
@@ -170,13 +170,27 @@ bool rewritePHIs(Function &F, UniformityInfo &UA, DominatorTree *DT) {
   return Changed;
 }
 
-bool AMDGPURewriteUndefForPHI::runOnFunction(Function &F) {
+bool AMDGPURewriteUndefForPHILegacy::runOnFunction(Function &F) {
   UniformityInfo &UA =
       getAnalysis<UniformityInfoWrapperPass>().getUniformityInfo();
   DominatorTree *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   return rewritePHIs(F, UA, DT);
 }
 
-FunctionPass *llvm::createAMDGPURewriteUndefForPHIPass() {
-  return new AMDGPURewriteUndefForPHI();
+PreservedAnalyses
+AMDGPURewriteUndefForPHIPass::run(Function &F, FunctionAnalysisManager &AM) {
+  UniformityInfo &UA = AM.getResult<UniformityInfoAnalysis>(F);
+  DominatorTree *DT = &AM.getResult<DominatorTreeAnalysis>(F);
+  bool Changed = rewritePHIs(F, UA, DT);
+  if (Changed) {
+    PreservedAnalyses PA;
+    PA.preserveSet<CFGAnalyses>();
+    return PA;
+  }
+
+  return PreservedAnalyses::all();
+}
+
+FunctionPass *llvm::createAMDGPURewriteUndefForPHILegacyPass() {
+  return new AMDGPURewriteUndefForPHILegacy();
 }

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

@@ -390,7 +390,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAMDGPUTarget() {
   initializeAMDGPURemoveIncompatibleFunctionsPass(*PR);
   initializeAMDGPULowerModuleLDSLegacyPass(*PR);
   initializeAMDGPURewriteOutArgumentsPass(*PR);
-  initializeAMDGPURewriteUndefForPHIPass(*PR);
+  initializeAMDGPURewriteUndefForPHILegacyPass(*PR);
   initializeAMDGPUUnifyMetadataPass(*PR);
   initializeSIAnnotateControlFlowPass(*PR);
   initializeAMDGPUInsertDelayAluPass(*PR);
@@ -663,6 +663,10 @@ void AMDGPUTargetMachine::registerPassBuilderCallbacks(PassBuilder &PB) {
           PM.addPass(AMDGPULowerKernelArgumentsPass(*this));
           return true;
         }
+        if (PassName == "amdgpu-rewrite-undef-for-phi") {
+          PM.addPass(AMDGPURewriteUndefForPHIPass(*this));
+          return true;
+        }
         return false;
       });
 
@@ -1150,7 +1154,7 @@ bool GCNPassConfig::addPreISel() {
     // TODO: Move this right after structurizeCFG to avoid extra divergence
     // analysis. This depends on stopping SIAnnotateControlFlow from making
     // control flow modifications.
-    addPass(createAMDGPURewriteUndefForPHIPass());
+    addPass(createAMDGPURewriteUndefForPHILegacyPass());
   }
   addPass(createLCSSAPass());
 

llvm/test/CodeGen/AMDGPU/rewrite-undef-for-phi.ll

@@ -1,5 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -mtriple=amdgcn-- -S -amdgpu-rewrite-undef-for-phi %s | FileCheck -check-prefix=OPT %s
+; RUN: opt -mtriple=amdgcn-- -S -passes=amdgpu-rewrite-undef-for-phi %s | FileCheck -check-prefix=OPT %s
 
 define amdgpu_ps float @basic(float inreg %c, i32 %x) #0 {
 ; OPT-LABEL: @basic(