[AMDGPU] Use UniformityAnalysis in RewriteUndefsForPHI

[foad]

LGTM, thanks!

[sameerds]

Please make sure this passes internal builds before submitting! We have faced issues with similar recent changes.

[Pierre-vh]

In D145359#4170725, @sameerds wrote:

Please make sure this passes internal builds before submitting! We have faced issues with similar recent changes.

I first checked that it passed internal builds before submitting both of my patches - seems like there was no issue.

llvm/lib/Target/AMDGPU/AMDGPURewriteUndefForPHI.cpp

 //===- AMDGPURewriteUndefForPHI.cpp ---------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 // This file implements the idea to rewrite undef incoming operand for certain
 // PHIs in structurized CFG. This pass only works on IR that has gone through
 // StructurizedCFG pass, and this pass has some additional limitation that make
 // it can only run after SIAnnotateControlFlow.
 //
-// To achieve optimal code generation for AMDGPU, we assume that divergence
+// To achieve optimal code generation for AMDGPU, we assume that uniformity
 // analysis reports the PHI in join block of divergent branch as uniform if
 // it has one unique uniform value plus additional undefined/poisoned incoming
 // value. That is to say the later compiler pipeline will ensure such PHI always
 // return uniform value and ensure it work correctly. Let's take a look at two
 // typical patterns in structured CFG that need to be taken care: (In both
 // patterns, block %if terminate with divergent branch.)
 //
 // Pattern A: Block with undefined incoming value dominates defined predecessor

 //     | \
 //     | %then
 //     | /
 // -> %header: %phi = phi [%uniform, %if], [%undef, %then], [%uniform2, %header]
 // |   |
 // \---
 
 #include "AMDGPU.h"
-#include "llvm/Analysis/LegacyDivergenceAnalysis.h"
+#include "llvm/Analysis/UniformityAnalysis.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/InitializePasses.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "amdgpu-rewrite-undef-for-phi"
 
 namespace {
 
 class AMDGPURewriteUndefForPHI : public FunctionPass {
 public:
   static char ID;
   AMDGPURewriteUndefForPHI() : FunctionPass(ID) {
     initializeAMDGPURewriteUndefForPHIPass(*PassRegistry::getPassRegistry());
   }
   bool runOnFunction(Function &F) override;
   StringRef getPassName() const override {
     return "AMDGPU Rewrite Undef for PHI";
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.addRequired<LegacyDivergenceAnalysis>();
+    AU.addRequired<UniformityInfoWrapperPass>();
     AU.addRequired<DominatorTreeWrapperPass>();
 
     AU.addPreserved<DominatorTreeWrapperPass>();
-    AU.addPreserved<LegacyDivergenceAnalysis>();
+    AU.addPreserved<UniformityInfoWrapperPass>();
     AU.setPreservesCFG();
   }
 };
 
 } // end anonymous namespace
 char AMDGPURewriteUndefForPHI::ID = 0;
 
 INITIALIZE_PASS_BEGIN(AMDGPURewriteUndefForPHI, DEBUG_TYPE,
                       "Rewrite undef for PHI", false, false)
-INITIALIZE_PASS_DEPENDENCY(LegacyDivergenceAnalysis)
+INITIALIZE_PASS_DEPENDENCY(UniformityInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_END(AMDGPURewriteUndefForPHI, DEBUG_TYPE,
                     "Rewrite undef for PHI", false, false)
 
-bool rewritePHIs(Function &F, LegacyDivergenceAnalysis *DA, DominatorTree *DT) {
+bool rewritePHIs(Function &F, UniformityInfo &UA, DominatorTree *DT) {
   bool Changed = false;
   SmallVector<PHINode *> ToBeDeleted;
   for (auto &BB : F) {
     for (auto &PHI : BB.phis()) {
-      if (DA->isDivergent(&PHI))
+      if (UA.isDivergent(&PHI))
         continue;
 
       // The unique incoming value except undef/poison for the PHI node.
       Value *UniqueDefinedIncoming = nullptr;
       // The divergent block with defined incoming value that dominates all
       // other block with the same incoming value.
       BasicBlock *DominateBB = nullptr;
       // Predecessors with undefined incoming value (excluding loop backedge).
for (auto &PHI : BB.phis()) {
           break;
         }
       }
       // We only need to replace the undef for the PHI which is merging
       // defined/undefined values from divergent threads.
       // TODO: We should still be able to replace undef value if the unique
       // value is a Constant.
       if (!UniqueDefinedIncoming || Undefs.empty() ||
-          !DA->isDivergent(DominateBB->getTerminator()))
+          !UA.isDivergent(DominateBB->getTerminator()))
         continue;
 
       // We only replace the undef when DominateBB truly dominates all the
       // other predecessors with undefined incoming value. Make sure DominateBB
       // dominates BB so that UniqueDefinedIncoming is available in BB and
       // afterwards.
       if (DT->dominates(DominateBB, &BB) && all_of(Undefs, [&](BasicBlock *UD) {
             return DT->dominates(DominateBB, UD);
           })) {
         PHI.replaceAllUsesWith(UniqueDefinedIncoming);
         ToBeDeleted.push_back(&PHI);
         Changed = true;
       }
     }
   }
 
   for (auto *PHI : ToBeDeleted)
     PHI->eraseFromParent();
 
   return Changed;
 }
 
 bool AMDGPURewriteUndefForPHI::runOnFunction(Function &F) {
-  LegacyDivergenceAnalysis *DA = &getAnalysis<LegacyDivergenceAnalysis>();
+  UniformityInfo &UA =
+      getAnalysis<UniformityInfoWrapperPass>().getUniformityInfo();
   DominatorTree *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
-  return rewritePHIs(F, DA, DT);
+  return rewritePHIs(F, UA, DT);
 }
 
 FunctionPass *llvm::createAMDGPURewriteUndefForPHIPass() {
   return new AMDGPURewriteUndefForPHI();
 }

llvm/test/CodeGen/AMDGPU/llc-pipeline.ll

 ; GCN-O0-NEXT:        Cycle Info Analysis
 ; GCN-O0-NEXT:        Uniformity Analysis
 ; GCN-O0-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O0-NEXT:        Function Alias Analysis Results
 ; GCN-O0-NEXT:        Memory SSA
 ; GCN-O0-NEXT:        AMDGPU Annotate Uniform Values
 ; GCN-O0-NEXT:        Natural Loop Information
 ; GCN-O0-NEXT:        SI annotate control flow
-; GCN-O0-NEXT:        Post-Dominator Tree Construction
-; GCN-O0-NEXT:        Legacy Divergence Analysis
+; GCN-O0-NEXT:        Cycle Info Analysis
+; GCN-O0-NEXT:        Uniformity Analysis
 ; GCN-O0-NEXT:        AMDGPU Rewrite Undef for PHI
 ; GCN-O0-NEXT:        LCSSA Verifier
 ; GCN-O0-NEXT:        Loop-Closed SSA Form Pass
 ; GCN-O0-NEXT:      DummyCGSCCPass
 ; GCN-O0-NEXT:      FunctionPass Manager
 ; GCN-O0-NEXT:        Prepare callbr
 ; GCN-O0-NEXT:        Safe Stack instrumentation pass
 ; GCN-O0-NEXT:        Insert stack protectors

 ; GCN-O1-NEXT:        Cycle Info Analysis
 ; GCN-O1-NEXT:        Uniformity Analysis
 ; GCN-O1-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O1-NEXT:        Function Alias Analysis Results
 ; GCN-O1-NEXT:        Memory SSA
 ; GCN-O1-NEXT:        AMDGPU Annotate Uniform Values
 ; GCN-O1-NEXT:        Natural Loop Information
 ; GCN-O1-NEXT:        SI annotate control flow
-; GCN-O1-NEXT:        Post-Dominator Tree Construction
-; GCN-O1-NEXT:        Legacy Divergence Analysis
+; GCN-O1-NEXT:        Cycle Info Analysis
+; GCN-O1-NEXT:        Uniformity Analysis
 ; GCN-O1-NEXT:        AMDGPU Rewrite Undef for PHI
 ; GCN-O1-NEXT:        LCSSA Verifier
 ; GCN-O1-NEXT:        Loop-Closed SSA Form Pass
 ; GCN-O1-NEXT:      DummyCGSCCPass
 ; GCN-O1-NEXT:      FunctionPass Manager
 ; GCN-O1-NEXT:        Prepare callbr
 ; GCN-O1-NEXT:        Safe Stack instrumentation pass
 ; GCN-O1-NEXT:        Insert stack protectors

 ; GCN-O1-OPTS-NEXT:        Cycle Info Analysis
 ; GCN-O1-OPTS-NEXT:        Uniformity Analysis
 ; GCN-O1-OPTS-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O1-OPTS-NEXT:        Function Alias Analysis Results
 ; GCN-O1-OPTS-NEXT:        Memory SSA
 ; GCN-O1-OPTS-NEXT:        AMDGPU Annotate Uniform Values
 ; GCN-O1-OPTS-NEXT:        Natural Loop Information
 ; GCN-O1-OPTS-NEXT:        SI annotate control flow
-; GCN-O1-OPTS-NEXT:        Post-Dominator Tree Construction
-; GCN-O1-OPTS-NEXT:        Legacy Divergence Analysis
+; GCN-O1-OPTS-NEXT:        Cycle Info Analysis
+; GCN-O1-OPTS-NEXT:        Uniformity Analysis
 ; GCN-O1-OPTS-NEXT:        AMDGPU Rewrite Undef for PHI
 ; GCN-O1-OPTS-NEXT:        LCSSA Verifier
 ; GCN-O1-OPTS-NEXT:        Loop-Closed SSA Form Pass
 ; GCN-O1-OPTS-NEXT:      DummyCGSCCPass
 ; GCN-O1-OPTS-NEXT:      FunctionPass Manager
 ; GCN-O1-OPTS-NEXT:        Prepare callbr
 ; GCN-O1-OPTS-NEXT:        Safe Stack instrumentation pass
 ; GCN-O1-OPTS-NEXT:        Insert stack protectors

 ; GCN-O2-NEXT:        Cycle Info Analysis
 ; GCN-O2-NEXT:        Uniformity Analysis
 ; GCN-O2-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O2-NEXT:        Function Alias Analysis Results
 ; GCN-O2-NEXT:        Memory SSA
 ; GCN-O2-NEXT:        AMDGPU Annotate Uniform Values
 ; GCN-O2-NEXT:        Natural Loop Information
 ; GCN-O2-NEXT:        SI annotate control flow
-; GCN-O2-NEXT:        Post-Dominator Tree Construction
-; GCN-O2-NEXT:        Legacy Divergence Analysis
+; GCN-O2-NEXT:        Cycle Info Analysis
+; GCN-O2-NEXT:        Uniformity Analysis
 ; GCN-O2-NEXT:        AMDGPU Rewrite Undef for PHI
 ; GCN-O2-NEXT:        LCSSA Verifier
 ; GCN-O2-NEXT:        Loop-Closed SSA Form Pass
 ; GCN-O2-NEXT:      Analysis if a function is memory bound
 ; GCN-O2-NEXT:      DummyCGSCCPass
 ; GCN-O2-NEXT:      FunctionPass Manager
 ; GCN-O2-NEXT:        Prepare callbr
 ; GCN-O2-NEXT:        Safe Stack instrumentation pass

 ; GCN-O3-NEXT:        Cycle Info Analysis
 ; GCN-O3-NEXT:        Uniformity Analysis
 ; GCN-O3-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O3-NEXT:        Function Alias Analysis Results
 ; GCN-O3-NEXT:        Memory SSA
 ; GCN-O3-NEXT:        AMDGPU Annotate Uniform Values
 ; GCN-O3-NEXT:        Natural Loop Information
 ; GCN-O3-NEXT:        SI annotate control flow
-; GCN-O3-NEXT:        Post-Dominator Tree Construction
-; GCN-O3-NEXT:        Legacy Divergence Analysis
+; GCN-O3-NEXT:        Cycle Info Analysis
+; GCN-O3-NEXT:        Uniformity Analysis
 ; GCN-O3-NEXT:        AMDGPU Rewrite Undef for PHI
 ; GCN-O3-NEXT:        LCSSA Verifier
 ; GCN-O3-NEXT:        Loop-Closed SSA Form Pass
 ; GCN-O3-NEXT:      Analysis if a function is memory bound
 ; GCN-O3-NEXT:      DummyCGSCCPass
 ; GCN-O3-NEXT:      FunctionPass Manager
 ; GCN-O3-NEXT:        Prepare callbr
 ; GCN-O3-NEXT:        Safe Stack instrumentation pass