[StructurizeCFG] Annotate branches that were treated as uniform

[nhaehnle]

Summary

This fully solves the problem where the StructurizeCFG pass does not
consider the same branches as uniform as the SIAnnotateControlFlow pass.
The patch in D19013 helps with this problem, but is not sufficient
(and, interestingly, causes a "regression" with one of the existing
test cases).

No tests included here, because tests in D19013 already cover this.

[• tstellarAMD]

LGTM.

llvm/trunk/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp

return !checkType(MemVal, AMDGPUAS::LOCAL_ADDRESS) &&
     !checkType(MemVal, AMDGPUAS::REGION_ADDRESS) &&
     !checkType(MemVal, AMDGPUAS::CONSTANT_ADDRESS) &&
     !checkType(MemVal, AMDGPUAS::PARAM_D_ADDRESS) &&
     !checkType(MemVal, AMDGPUAS::PARAM_I_ADDRESS);
 }
 
 bool AMDGPUDAGToDAGISel::isUniformBr(const SDNode *N) const {
   const BasicBlock *BB = FuncInfo->MBB->getBasicBlock();
-  return BB->getTerminator()->getMetadata("amdgpu.uniform");
+  const Instruction *Term = BB->getTerminator();
+  return Term->getMetadata("amdgpu.uniform") ||
+         Term->getMetadata("structurizecfg.uniform");
 }
 
 const char *AMDGPUDAGToDAGISel::getPassName() const {
   return "AMDGPU DAG->DAG Pattern Instruction Selection";
 }
 
 //===----------------------------------------------------------------------===//
 // Complex Patterns

llvm/trunk/lib/Target/AMDGPU/SIAnnotateControlFlow.cpp

class SIAnnotateControlFlow : public FunctionPass {
   Constant *Loop;
   Constant *EndCf;
 
   DominatorTree *DT;
   StackVector Stack;
 
   LoopInfo *LI;
 
+  bool isUniform(BranchInst *T);
+
   bool isTopOfStack(BasicBlock *BB);
 
   Value *popSaved();
 
   void push(BasicBlock *BB, Value *Saved);
 
   bool isElse(PHINode *Phi);
 
Loop = M.getOrInsertFunction(
     LoopIntrinsic, Boolean, Int64, (Type *)nullptr);
 
   EndCf = M.getOrInsertFunction(
     EndCfIntrinsic, Void, Int64, (Type *)nullptr);
 
   return false;
 }
 
+/// \brief Is the branch condition uniform or did the StructurizeCFG pass
+/// consider it as such?
+bool SIAnnotateControlFlow::isUniform(BranchInst *T) {
+  return DA->isUniform(T->getCondition()) ||
+         T->getMetadata("structurizecfg.uniform") != nullptr;
+}
+
 /// \brief Is BB the last block saved on the stack ?
 bool SIAnnotateControlFlow::isTopOfStack(BasicBlock *BB) {
   return !Stack.empty() && Stack.back().first == BB;
 }
 
 /// \brief Pop the last saved value from the control flow stack
 Value *SIAnnotateControlFlow::popSaved() {
   return Stack.pop_back_val().second;

 // \brief Erase "Phi" if it is not used any more
 void SIAnnotateControlFlow::eraseIfUnused(PHINode *Phi) {
   if (!Phi->hasNUsesOrMore(1))
     Phi->eraseFromParent();
 }
 
 /// \brief Open a new "If" block
 void SIAnnotateControlFlow::openIf(BranchInst *Term) {
-  if (DA->isUniform(Term->getCondition())) {
+  if (isUniform(Term)) {
     return;
   }
   Value *Ret = CallInst::Create(If, Term->getCondition(), "", Term);
   Term->setCondition(ExtractValueInst::Create(Ret, 0, "", Term));
   push(Term->getSuccessor(1), ExtractValueInst::Create(Ret, 1, "", Term));
 }
 
 /// \brief Close the last "If" block and open a new "Else" block
 void SIAnnotateControlFlow::insertElse(BranchInst *Term) {
-  if (DA->isUniform(Term->getCondition())) {
+  if (isUniform(Term)) {
     return;
   }
   Value *Ret = CallInst::Create(Else, popSaved(), "", Term);
   Term->setCondition(ExtractValueInst::Create(Ret, 0, "", Term));
   push(Term->getSuccessor(1), ExtractValueInst::Create(Ret, 1, "", Term));
 }
 
 /// \brief Recursively handle the condition leading to a loop
Value *SIAnnotateControlFlow::handleLoopCondition(Value *Cond, PHINode *Broken,
   } else {
     llvm_unreachable("Unhandled loop condition!");
   }
   return 0;
 }
 
 /// \brief Handle a back edge (loop)
 void SIAnnotateControlFlow::handleLoop(BranchInst *Term) {
-  if (DA->isUniform(Term->getCondition())) {
+  if (isUniform(Term)) {
     return;
   }
 
   BasicBlock *BB = Term->getParent();
   llvm::Loop *L = LI->getLoopFor(BB);
   BasicBlock *Target = Term->getSuccessor(1);
   PHINode *Broken = PHINode::Create(Int64, 0, "", &Target->front());
 

llvm/trunk/lib/Transforms/Scalar/StructurizeCFG.cpp

bool StructurizeCFG::runOnRegion(Region *R, RGPassManager &RGM) {
   if (R->isTopLevelRegion())
     return false;
 
   if (SkipUniformRegions) {
     DA = &getAnalysis<DivergenceAnalysis>();
     // TODO: We could probably be smarter here with how we handle sub-regions.
     if (hasOnlyUniformBranches(R)) {
       DEBUG(dbgs() << "Skipping region with uniform control flow: " << *R << '\n');
+
+      // Mark all direct child block terminators as having been treated as
+      // uniform. To account for a possible future in which non-uniform
+      // sub-regions are treated more cleverly, indirect children are not
+      // marked as uniform.
+      MDNode *MD = MDNode::get(R->getEntry()->getParent()->getContext(), {});
+      Region::element_iterator E = R->element_end();
+      for (Region::element_iterator I = R->element_begin(); I != E; ++I) {
+        if (I->isSubRegion())
+          continue;
+
+        if (Instruction *Term = I->getEntry()->getTerminator())
+          Term->setMetadata("structurizecfg.uniform", MD);
+      }
+
       return false;
     }
   }
 
   Func = R->getEntry()->getParent();
   ParentRegion = R;
 
   DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();