[AMDGPU] Alternative control flow lowering. Handling uniform if and lop inside the divergent CF in SIAnnotateControlFlow

Author: alex-t

llvm/lib/Target/AMDGPU/SIAnnotateControlFlow.cpp

@@ -15,6 +15,7 @@
 #include "GCNSubtarget.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/UniformityAnalysis.h"
+#include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
@@ -86,7 +87,7 @@ class SIAnnotateControlFlow : public FunctionPass {
 
   bool handleLoop(BranchInst *Term);
 
-  bool insertWaveReconverge(BasicBlock *BB);
+  bool tryWaveReconverge(BasicBlock *BB);
 
 public:
   static char ID;
@@ -203,8 +204,6 @@ bool SIAnnotateControlFlow::eraseIfUnused(PHINode *Phi) {
 
 /// Open a new "If" block
 bool SIAnnotateControlFlow::openIf(BranchInst *Term) {
-  if (isUniform(Term))
-    return false;
 
   IRBuilder<> IRB(Term);
   Value *IfCall = IRB.CreateCall(If, {Term->getCondition()});
@@ -305,20 +304,44 @@ bool SIAnnotateControlFlow::handleLoop(BranchInst *Term) {
 }
 
 /// Close the last opened control flow
-bool SIAnnotateControlFlow::insertWaveReconverge(BasicBlock *BB) {
-      assert(succ_empty(BB) || succ_size(BB) == 1);
-      
-      if (succ_empty(BB))
-        return false;
+bool SIAnnotateControlFlow::tryWaveReconverge(BasicBlock *BB) {
+
+  if (succ_empty(BB))
+    return false;
 
-      BasicBlock *SingleSucc = *succ_begin(BB);
-      BranchInst *Term = dyn_cast<BranchInst>(BB->getTerminator());
-      BasicBlock::iterator InsPt = Term ? BasicBlock::iterator(Term) : BB->end();
+  BranchInst *Term = dyn_cast<BranchInst>(BB->getTerminator());
+  if (Term->getNumSuccessors() == 1) {
+    // The current BBs single successor is a top of the stack. We need to
+    // reconverge over thaqt path.
+    BasicBlock *SingleSucc = *succ_begin(BB);
+    BasicBlock::iterator InsPt = Term ? BasicBlock::iterator(Term) : BB->end();
 
-      if (isTopOfStack(SingleSucc)) {
-        Value *Exec = Stack.back().second;
-        IRBuilder<>(BB, InsPt).CreateCall(WaveReconverge, {Exec});
+    if (isTopOfStack(SingleSucc)) {
+      Value *Exec = Stack.back().second;
+      IRBuilder<>(BB, InsPt).CreateCall(WaveReconverge, {Exec});
+    }
+  } else {
+    // We have a uniform conditional branch terminating the block.
+    // THis block may be the last in the Then path of the enclosing divergent
+    // IF.
+    if (!isUniform(Term))
+      // Divergent loop is going to be further processed in another place
+      return false;
+
+    for (auto Succ : Term->successors()) {
+      if (isTopOfStack(Succ)) {
+        // Just split to make a room for further WAVE_RECONVERGE insertion
+        SmallVector<BasicBlock*, 2> Preds;
+        for (auto P : predecessors(Succ)) {
+          if (DT->dominates(BB, P))
+            Preds.push_back(P);
+        }
+        DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
+        SplitBlockPredecessors(Succ, Preds, ".reconverge", &DTU, LI,
+                                            nullptr, false);
       }
+    }
+  }
 
   return true;
 }
@@ -342,8 +365,8 @@ bool SIAnnotateControlFlow::runOnFunction(Function &F) {
     if (!Term || Term->isUnconditional()) {
       if (isTopOfStack(BB))
         Stack.pop_back();
-      
-      insertWaveReconverge(BB);
+
+      Changed |= tryWaveReconverge(BB);
 
       continue;
     }
@@ -352,6 +375,10 @@ bool SIAnnotateControlFlow::runOnFunction(Function &F) {
       if (isTopOfStack(BB))
         Stack.pop_back();
 
+      // Let's take care of uniform loop latch that may be closing the Then
+      // path of the enclosing divergent branch.
+      Changed |= tryWaveReconverge(BB);
+
       if (DT->dominates(Term->getSuccessor(1), BB))
         Changed |= handleLoop(Term);
       continue;
@@ -368,7 +395,12 @@ bool SIAnnotateControlFlow::runOnFunction(Function &F) {
       Stack.pop_back();
     }
 
-    Changed |= openIf(Term);
+    if (isUniform(Term))
+      // Uniform conditional branch may be in the block that closes the Then
+      // path of the divergent conditional branch.
+      Changed |= tryWaveReconverge(BB);
+    else
+      Changed |= openIf(Term);
   }
 
   if (!Stack.empty()) {

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -15741,9 +15741,8 @@ void SITargetLowering::finalizeLowering(MachineFunction &MF) const {
   }
 
   // ISel inserts copy to regs for the successor PHIs
-  // at the BB end. We need to move the SI_WAVE_RECONVERGE right before the branch.
-  // Even we don't have to move SI_WAVE_RECONVERGE we need to take care of the
-  // S_CBRANCH_SCC0/1 as SI_WAVE_RECONVERGE overwrites SCC
+  // at the BB end. We need to move the SI_WAVE_RECONVERGE right before the
+  // branch.
   for (auto &MBB : MF) {
     for (auto &MI : MBB) {
       if (MI.getOpcode() == AMDGPU::SI_WAVE_RECONVERGE) {
@@ -15755,66 +15754,8 @@ void SITargetLowering::finalizeLowering(MachineFunction &MF) const {
           Next++;
         }
 
-        // Lets take care of SCC users as SI_WAVE_RECONVERGE defines SCC
-        bool NeedPreserveSCC =
-            Next != MBB.end() && Next->readsRegister(AMDGPU::SCC);
-        MachineBasicBlock::iterator SCCDefUse(Next);
-        // This loop will be never taken as we always have S_CBRANCH_SCC1/0 at
-        // the end of the block.
-        while (!NeedPreserveSCC && SCCDefUse != MBB.end()) {
-          if (SCCDefUse->definesRegister(AMDGPU::SCC))
-            // This should never happen - SCC def after the branch reading SCC
-            break;
-          if (SCCDefUse->readsRegister(AMDGPU::SCC)) {
-            NeedPreserveSCC = true;
-            break;
-          }
-          SCCDefUse++;
-        }
-        if (NeedPreserveSCC) {
-          MachineBasicBlock::reverse_iterator BackSeeker(Next);
-          while (BackSeeker != MBB.rend()) {
-            if (BackSeeker != MI && BackSeeker->definesRegister(AMDGPU::SCC))
-              break;
-            BackSeeker++;
-          }
-          // we need this to makes some artificial MIR tests happy
-          bool NeedSetSCCUndef = false;
-          if (BackSeeker == MBB.rend()) {
-            // We have reached the begin of the block but haven't seen the SCC
-            // def Given that the MIR is correct, we either have SCC live in
-            // or SCCUser SCC operand is undef. In fact, we don't need to emit
-            // the instructions that preserve thje SCC if the use is Undef. We
-            // do this just because the MIR looks weird otherwise.
-            MachineOperand *SCCUseOp =
-                SCCDefUse->findRegisterUseOperand(AMDGPU::SCC, false, TRI);
-            assert(SCCUseOp);
-            bool IsSCCLiveIn = MBB.isLiveIn(AMDGPU::SCC);
-            bool IsUseUndef = SCCUseOp->isUndef();
-            NeedSetSCCUndef = (!IsSCCLiveIn && IsUseUndef);
-          }
-          MachineBasicBlock::iterator InsPt(BackSeeker);
-          Register SavedSCC =
-              MRI.createVirtualRegister(&AMDGPU::SReg_32RegClass);
-          MachineInstr *SaveSCC =
-              BuildMI(MBB, InsPt, InsPt->getDebugLoc(),
-                      TII->get(AMDGPU::S_CSELECT_B32), SavedSCC)
-                  .addImm(1)
-                  .addImm(0);
-          if (NeedSetSCCUndef) {
-
-            MachineOperand *SCCOp =
-                SaveSCC->findRegisterUseOperand(AMDGPU::SCC, false, TRI);
-            if (SCCOp)
-              SCCOp->setIsUndef();
-          }
-          Register Tmp =
-              MRI.createVirtualRegister(&AMDGPU::SReg_32RegClass);
-          Next = BuildMI(MBB, Next, Next->getDebugLoc(),
-                         TII->get(AMDGPU::S_AND_B32_term), Tmp)
-                     .addReg(SavedSCC)
-                     .addImm(1);
-        }
+        assert((Next == MBB.end() || !Next->readsRegister(AMDGPU::SCC)) &&
+               "Malformed CFG detected!\n");
 
         if (NeedToMove) {
           MBB.splice(Next, &MBB, &MI);