1//===- StructurizeCFG.cpp -------------------------------------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#include "llvm/ADT/DenseMap.h"
10#include "llvm/ADT/MapVector.h"
11#include "llvm/ADT/PostOrderIterator.h"
12#include "llvm/ADT/STLExtras.h"
13#include "llvm/ADT/SmallPtrSet.h"
14#include "llvm/ADT/SmallVector.h"
15#include "llvm/Analysis/InstructionSimplify.h"
16#include "llvm/Analysis/LegacyDivergenceAnalysis.h"
17#include "llvm/Analysis/LoopInfo.h"
18#include "llvm/Analysis/RegionInfo.h"
19#include "llvm/Analysis/RegionIterator.h"
20#include "llvm/Analysis/RegionPass.h"
21#include "llvm/IR/Argument.h"
22#include "llvm/IR/BasicBlock.h"
23#include "llvm/IR/CFG.h"
24#include "llvm/IR/Constant.h"
25#include "llvm/IR/Constants.h"
26#include "llvm/IR/Dominators.h"
27#include "llvm/IR/Function.h"
28#include "llvm/IR/InstrTypes.h"
29#include "llvm/IR/Instruction.h"
30#include "llvm/IR/Instructions.h"
31#include "llvm/IR/Metadata.h"
32#include "llvm/IR/PatternMatch.h"
33#include "llvm/IR/Type.h"
34#include "llvm/IR/Use.h"
35#include "llvm/IR/User.h"
36#include "llvm/IR/Value.h"
37#include "llvm/InitializePasses.h"
38#include "llvm/Pass.h"
39#include "llvm/Support/Casting.h"
40#include "llvm/Support/CommandLine.h"
41#include "llvm/Support/Debug.h"
42#include "llvm/Support/ErrorHandling.h"
43#include "llvm/Support/raw_ostream.h"
44#include "llvm/Transforms/Scalar.h"
45#include "llvm/Transforms/Utils.h"
46#include "llvm/Transforms/Utils/SSAUpdater.h"
47#include <algorithm>
48#include <cassert>
49#include <utility>
50
51using namespace llvm;
52using namespace llvm::PatternMatch;
53
54#define DEBUG_TYPE "structurizecfg"
55
56// The name for newly created blocks.
57static const char *const FlowBlockName = "Flow";
58
59namespace {
60
61static cl::opt<bool> ForceSkipUniformRegions(
62  "structurizecfg-skip-uniform-regions",
63  cl::Hidden,
64  cl::desc("Force whether the StructurizeCFG pass skips uniform regions"),
65  cl::init(false));
66
67static cl::opt<bool>
68    RelaxedUniformRegions("structurizecfg-relaxed-uniform-regions", cl::Hidden,
69                          cl::desc("Allow relaxed uniform region checks"),
70                          cl::init(true));
71
72// Definition of the complex types used in this pass.
73
74using BBValuePair = std::pair<BasicBlock *, Value *>;
75
76using RNVector = SmallVector<RegionNode *, 8>;
77using BBVector = SmallVector<BasicBlock *, 8>;
78using BranchVector = SmallVector<BranchInst *, 8>;
79using BBValueVector = SmallVector<BBValuePair, 2>;
80
81using BBSet = SmallPtrSet<BasicBlock *, 8>;
82
83using PhiMap = MapVector<PHINode *, BBValueVector>;
84using BB2BBVecMap = MapVector<BasicBlock *, BBVector>;
85
86using BBPhiMap = DenseMap<BasicBlock *, PhiMap>;
87using BBPredicates = DenseMap<BasicBlock *, Value *>;
88using PredMap = DenseMap<BasicBlock *, BBPredicates>;
89using BB2BBMap = DenseMap<BasicBlock *, BasicBlock *>;
90
91/// Finds the nearest common dominator of a set of BasicBlocks.
92///
93/// For every BB you add to the set, you can specify whether we "remember" the
94/// block.  When you get the common dominator, you can also ask whether it's one
95/// of the blocks we remembered.
96class NearestCommonDominator {
97  DominatorTree *DT;
98  BasicBlock *Result = nullptr;
99  bool ResultIsRemembered = false;
100
101  /// Add BB to the resulting dominator.
102  void addBlock(BasicBlock *BB, bool Remember) {
103    if (!Result) {
104      Result = BB;
105      ResultIsRemembered = Remember;
106      return;
107    }
108
109    BasicBlock *NewResult = DT->findNearestCommonDominator(Result, BB);
110    if (NewResult != Result)
111      ResultIsRemembered = false;
112    if (NewResult == BB)
113      ResultIsRemembered |= Remember;
114    Result = NewResult;
115  }
116
117public:
118  explicit NearestCommonDominator(DominatorTree *DomTree) : DT(DomTree) {}
119
120  void addBlock(BasicBlock *BB) {
121    addBlock(BB, /* Remember = */ false);
122  }
123
124  void addAndRememberBlock(BasicBlock *BB) {
125    addBlock(BB, /* Remember = */ true);
126  }
127
128  /// Get the nearest common dominator of all the BBs added via addBlock() and
129  /// addAndRememberBlock().
130  BasicBlock *result() { return Result; }
131
132  /// Is the BB returned by getResult() one of the blocks we added to the set
133  /// with addAndRememberBlock()?
134  bool resultIsRememberedBlock() { return ResultIsRemembered; }
135};
136
137/// Transforms the control flow graph on one single entry/exit region
138/// at a time.
139///
140/// After the transform all "If"/"Then"/"Else" style control flow looks like
141/// this:
142///
143/// \verbatim
144/// 1
145/// ||
146/// | |
147/// 2 |
148/// | /
149/// |/
150/// 3
151/// ||   Where:
152/// | |  1 = "If" block, calculates the condition
153/// 4 |  2 = "Then" subregion, runs if the condition is true
154/// | /  3 = "Flow" blocks, newly inserted flow blocks, rejoins the flow
155/// |/   4 = "Else" optional subregion, runs if the condition is false
156/// 5    5 = "End" block, also rejoins the control flow
157/// \endverbatim
158///
159/// Control flow is expressed as a branch where the true exit goes into the
160/// "Then"/"Else" region, while the false exit skips the region
161/// The condition for the optional "Else" region is expressed as a PHI node.
162/// The incoming values of the PHI node are true for the "If" edge and false
163/// for the "Then" edge.
164///
165/// Additionally to that even complicated loops look like this:
166///
167/// \verbatim
168/// 1
169/// ||
170/// | |
171/// 2 ^  Where:
172/// | /  1 = "Entry" block
173/// |/   2 = "Loop" optional subregion, with all exits at "Flow" block
174/// 3    3 = "Flow" block, with back edge to entry block
175/// |
176/// \endverbatim
177///
178/// The back edge of the "Flow" block is always on the false side of the branch
179/// while the true side continues the general flow. So the loop condition
180/// consist of a network of PHI nodes where the true incoming values expresses
181/// breaks and the false values expresses continue states.
182class StructurizeCFG : public RegionPass {
183  bool SkipUniformRegions;
184
185  Type *Boolean;
186  ConstantInt *BoolTrue;
187  ConstantInt *BoolFalse;
188  UndefValue *BoolUndef;
189
190  Function *Func;
191  Region *ParentRegion;
192
193  LegacyDivergenceAnalysis *DA;
194  DominatorTree *DT;
195  LoopInfo *LI;
196
197  SmallVector<RegionNode *, 8> Order;
198  BBSet Visited;
199
200  BBPhiMap DeletedPhis;
201  BB2BBVecMap AddedPhis;
202
203  PredMap Predicates;
204  BranchVector Conditions;
205
206  BB2BBMap Loops;
207  PredMap LoopPreds;
208  BranchVector LoopConds;
209
210  RegionNode *PrevNode;
211
212  void orderNodes();
213
214  Loop *getAdjustedLoop(RegionNode *RN);
215  unsigned getAdjustedLoopDepth(RegionNode *RN);
216
217  void analyzeLoops(RegionNode *N);
218
219  Value *invert(Value *Condition);
220
221  Value *buildCondition(BranchInst *Term, unsigned Idx, bool Invert);
222
223  void gatherPredicates(RegionNode *N);
224
225  void collectInfos();
226
227  void insertConditions(bool Loops);
228
229  void delPhiValues(BasicBlock *From, BasicBlock *To);
230
231  void addPhiValues(BasicBlock *From, BasicBlock *To);
232
233  void setPhiValues();
234
235  void killTerminator(BasicBlock *BB);
236
237  void changeExit(RegionNode *Node, BasicBlock *NewExit,
238                  bool IncludeDominator);
239
240  BasicBlock *getNextFlow(BasicBlock *Dominator);
241
242  BasicBlock *needPrefix(bool NeedEmpty);
243
244  BasicBlock *needPostfix(BasicBlock *Flow, bool ExitUseAllowed);
245
246  void setPrevNode(BasicBlock *BB);
247
248  bool dominatesPredicates(BasicBlock *BB, RegionNode *Node);
249
250  bool isPredictableTrue(RegionNode *Node);
251
252  void wireFlow(bool ExitUseAllowed, BasicBlock *LoopEnd);
253
254  void handleLoops(bool ExitUseAllowed, BasicBlock *LoopEnd);
255
256  void createFlow();
257
258  void rebuildSSA();
259
260public:
261  static char ID;
262
263  explicit StructurizeCFG(bool SkipUniformRegions_ = false)
264      : RegionPass(ID),
265        SkipUniformRegions(SkipUniformRegions_) {
266    if (ForceSkipUniformRegions.getNumOccurrences())
267      SkipUniformRegions = ForceSkipUniformRegions.getValue();
268    initializeStructurizeCFGPass(*PassRegistry::getPassRegistry());
269  }
270
271  bool doInitialization(Region *R, RGPassManager &RGM) override;
272
273  bool runOnRegion(Region *R, RGPassManager &RGM) override;
274
275  StringRef getPassName() const override { return "Structurize control flow"; }
276
277  void getAnalysisUsage(AnalysisUsage &AU) const override {
278    if (SkipUniformRegions)
279      AU.addRequired<LegacyDivergenceAnalysis>();
280    AU.addRequiredID(LowerSwitchID);
281    AU.addRequired<DominatorTreeWrapperPass>();
282    AU.addRequired<LoopInfoWrapperPass>();
283
284    AU.addPreserved<DominatorTreeWrapperPass>();
285    RegionPass::getAnalysisUsage(AU);
286  }
287};
288
289} // end anonymous namespace
290
291char StructurizeCFG::ID = 0;
292
293INITIALIZE_PASS_BEGIN(StructurizeCFG, "structurizecfg", "Structurize the CFG",
294                      false, false)
295INITIALIZE_PASS_DEPENDENCY(LegacyDivergenceAnalysis)
296INITIALIZE_PASS_DEPENDENCY(LowerSwitch)
297INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
298INITIALIZE_PASS_DEPENDENCY(RegionInfoPass)
299INITIALIZE_PASS_END(StructurizeCFG, "structurizecfg", "Structurize the CFG",
300                    false, false)
301
302/// Initialize the types and constants used in the pass
303bool StructurizeCFG::doInitialization(Region *R, RGPassManager &RGM) {
304  LLVMContext &Context = R->getEntry()->getContext();
305
306  Boolean = Type::getInt1Ty(Context);
307  BoolTrue = ConstantInt::getTrue(Context);
308  BoolFalse = ConstantInt::getFalse(Context);
309  BoolUndef = UndefValue::get(Boolean);
310
311  return false;
312}
313
314/// Use the exit block to determine the loop if RN is a SubRegion.
315Loop *StructurizeCFG::getAdjustedLoop(RegionNode *RN) {
316  if (RN->isSubRegion()) {
317    Region *SubRegion = RN->getNodeAs<Region>();
318    return LI->getLoopFor(SubRegion->getExit());
319  }
320
321  return LI->getLoopFor(RN->getEntry());
322}
323
324/// Use the exit block to determine the loop depth if RN is a SubRegion.
325unsigned StructurizeCFG::getAdjustedLoopDepth(RegionNode *RN) {
326  if (RN->isSubRegion()) {
327    Region *SubR = RN->getNodeAs<Region>();
328    return LI->getLoopDepth(SubR->getExit());
329  }
330
331  return LI->getLoopDepth(RN->getEntry());
332}
333
334/// Build up the general order of nodes
335void StructurizeCFG::orderNodes() {
336  ReversePostOrderTraversal<Region*> RPOT(ParentRegion);
337  SmallDenseMap<Loop*, unsigned, 8> LoopBlocks;
338
339  // The reverse post-order traversal of the list gives us an ordering close
340  // to what we want.  The only problem with it is that sometimes backedges
341  // for outer loops will be visited before backedges for inner loops.
342  for (RegionNode *RN : RPOT) {
343    Loop *Loop = getAdjustedLoop(RN);
344    ++LoopBlocks[Loop];
345  }
346
347  unsigned CurrentLoopDepth = 0;
348  Loop *CurrentLoop = nullptr;
349  for (auto I = RPOT.begin(), E = RPOT.end(); I != E; ++I) {
350    RegionNode *RN = cast<RegionNode>(*I);
351    unsigned LoopDepth = getAdjustedLoopDepth(RN);
352
353    if (is_contained(Order, *I))
354      continue;
355
356    if (LoopDepth < CurrentLoopDepth) {
357      // Make sure we have visited all blocks in this loop before moving back to
358      // the outer loop.
359
360      auto LoopI = I;
361      while (unsigned &BlockCount = LoopBlocks[CurrentLoop]) {
362        LoopI++;
363        if (getAdjustedLoop(cast<RegionNode>(*LoopI)) == CurrentLoop) {
364          --BlockCount;
365          Order.push_back(*LoopI);
366        }
367      }
368    }
369
370    CurrentLoop = getAdjustedLoop(RN);
371    if (CurrentLoop)
372      LoopBlocks[CurrentLoop]--;
373
374    CurrentLoopDepth = LoopDepth;
375    Order.push_back(*I);
376  }
377
378  // This pass originally used a post-order traversal and then operated on
379  // the list in reverse. Now that we are using a reverse post-order traversal
380  // rather than re-working the whole pass to operate on the list in order,
381  // we just reverse the list and continue to operate on it in reverse.
382  std::reverse(Order.begin(), Order.end());
383}
384
385/// Determine the end of the loops
386void StructurizeCFG::analyzeLoops(RegionNode *N) {
387  if (N->isSubRegion()) {
388    // Test for exit as back edge
389    BasicBlock *Exit = N->getNodeAs<Region>()->getExit();
390    if (Visited.count(Exit))
391      Loops[Exit] = N->getEntry();
392
393  } else {
394    // Test for successors as back edge
395    BasicBlock *BB = N->getNodeAs<BasicBlock>();
396    BranchInst *Term = cast<BranchInst>(BB->getTerminator());
397
398    for (BasicBlock *Succ : Term->successors())
399      if (Visited.count(Succ))
400        Loops[Succ] = BB;
401  }
402}
403
404/// Invert the given condition
405Value *StructurizeCFG::invert(Value *Condition) {
406  // First: Check if it's a constant
407  if (Constant *C = dyn_cast<Constant>(Condition))
408    return ConstantExpr::getNot(C);
409
410  // Second: If the condition is already inverted, return the original value
411  Value *NotCondition;
412  if (match(Condition, m_Not(m_Value(NotCondition))))
413    return NotCondition;
414
415  if (Instruction *Inst = dyn_cast<Instruction>(Condition)) {
416    // Third: Check all the users for an invert
417    BasicBlock *Parent = Inst->getParent();
418    for (User *U : Condition->users())
419      if (Instruction *I = dyn_cast<Instruction>(U))
420        if (I->getParent() == Parent && match(I, m_Not(m_Specific(Condition))))
421          return I;
422
423    // Last option: Create a new instruction
424    return BinaryOperator::CreateNot(Condition, "", Parent->getTerminator());
425  }
426
427  if (Argument *Arg = dyn_cast<Argument>(Condition)) {
428    BasicBlock &EntryBlock = Arg->getParent()->getEntryBlock();
429    return BinaryOperator::CreateNot(Condition,
430                                     Arg->getName() + ".inv",
431                                     EntryBlock.getTerminator());
432  }
433
434  llvm_unreachable("Unhandled condition to invert");
435}
436
437/// Build the condition for one edge
438Value *StructurizeCFG::buildCondition(BranchInst *Term, unsigned Idx,
439                                      bool Invert) {
440  Value *Cond = Invert ? BoolFalse : BoolTrue;
441  if (Term->isConditional()) {
442    Cond = Term->getCondition();
443
444    if (Idx != (unsigned)Invert)
445      Cond = invert(Cond);
446  }
447  return Cond;
448}
449
450/// Analyze the predecessors of each block and build up predicates
451void StructurizeCFG::gatherPredicates(RegionNode *N) {
452  RegionInfo *RI = ParentRegion->getRegionInfo();
453  BasicBlock *BB = N->getEntry();
454  BBPredicates &Pred = Predicates[BB];
455  BBPredicates &LPred = LoopPreds[BB];
456
457  for (BasicBlock *P : predecessors(BB)) {
458    // Ignore it if it's a branch from outside into our region entry
459    if (!ParentRegion->contains(P))
460      continue;
461
462    Region *R = RI->getRegionFor(P);
463    if (R == ParentRegion) {
464      // It's a top level block in our region
465      BranchInst *Term = cast<BranchInst>(P->getTerminator());
466      for (unsigned i = 0, e = Term->getNumSuccessors(); i != e; ++i) {
467        BasicBlock *Succ = Term->getSuccessor(i);
468        if (Succ != BB)
469          continue;
470
471        if (Visited.count(P)) {
472          // Normal forward edge
473          if (Term->isConditional()) {
474            // Try to treat it like an ELSE block
475            BasicBlock *Other = Term->getSuccessor(!i);
476            if (Visited.count(Other) && !Loops.count(Other) &&
477                !Pred.count(Other) && !Pred.count(P)) {
478
479              Pred[Other] = BoolFalse;
480              Pred[P] = BoolTrue;
481              continue;
482            }
483          }
484          Pred[P] = buildCondition(Term, i, false);
485        } else {
486          // Back edge
487          LPred[P] = buildCondition(Term, i, true);
488        }
489      }
490    } else {
491      // It's an exit from a sub region
492      while (R->getParent() != ParentRegion)
493        R = R->getParent();
494
495      // Edge from inside a subregion to its entry, ignore it
496      if (*R == *N)
497        continue;
498
499      BasicBlock *Entry = R->getEntry();
500      if (Visited.count(Entry))
501        Pred[Entry] = BoolTrue;
502      else
503        LPred[Entry] = BoolFalse;
504    }
505  }
506}
507
508/// Collect various loop and predicate infos
509void StructurizeCFG::collectInfos() {
510  // Reset predicate
511  Predicates.clear();
512
513  // and loop infos
514  Loops.clear();
515  LoopPreds.clear();
516
517  // Reset the visited nodes
518  Visited.clear();
519
520  for (RegionNode *RN : reverse(Order)) {
521    LLVM_DEBUG(dbgs() << "Visiting: "
522                      << (RN->isSubRegion() ? "SubRegion with entry: " : "")
523                      << RN->getEntry()->getName() << " Loop Depth: "
524                      << LI->getLoopDepth(RN->getEntry()) << "\n");
525
526    // Analyze all the conditions leading to a node
527    gatherPredicates(RN);
528
529    // Remember that we've seen this node
530    Visited.insert(RN->getEntry());
531
532    // Find the last back edges
533    analyzeLoops(RN);
534  }
535}
536
537/// Insert the missing branch conditions
538void StructurizeCFG::insertConditions(bool Loops) {
539  BranchVector &Conds = Loops ? LoopConds : Conditions;
540  Value *Default = Loops ? BoolTrue : BoolFalse;
541  SSAUpdater PhiInserter;
542
543  for (BranchInst *Term : Conds) {
544    assert(Term->isConditional());
545
546    BasicBlock *Parent = Term->getParent();
547    BasicBlock *SuccTrue = Term->getSuccessor(0);
548    BasicBlock *SuccFalse = Term->getSuccessor(1);
549
550    PhiInserter.Initialize(Boolean, "");
551    PhiInserter.AddAvailableValue(&Func->getEntryBlock(), Default);
552    PhiInserter.AddAvailableValue(Loops ? SuccFalse : Parent, Default);
553
554    BBPredicates &Preds = Loops ? LoopPreds[SuccFalse] : Predicates[SuccTrue];
555
556    NearestCommonDominator Dominator(DT);
557    Dominator.addBlock(Parent);
558
559    Value *ParentValue = nullptr;
560    for (std::pair<BasicBlock *, Value *> BBAndPred : Preds) {
561      BasicBlock *BB = BBAndPred.first;
562      Value *Pred = BBAndPred.second;
563
564      if (BB == Parent) {
565        ParentValue = Pred;
566        break;
567      }
568      PhiInserter.AddAvailableValue(BB, Pred);
569      Dominator.addAndRememberBlock(BB);
570    }
571
572    if (ParentValue) {
573      Term->setCondition(ParentValue);
574    } else {
575      if (!Dominator.resultIsRememberedBlock())
576        PhiInserter.AddAvailableValue(Dominator.result(), Default);
577
578      Term->setCondition(PhiInserter.GetValueInMiddleOfBlock(Parent));
579    }
580  }
581}
582
583/// Remove all PHI values coming from "From" into "To" and remember
584/// them in DeletedPhis
585void StructurizeCFG::delPhiValues(BasicBlock *From, BasicBlock *To) {
586  PhiMap &Map = DeletedPhis[To];
587  for (PHINode &Phi : To->phis()) {
588    while (Phi.getBasicBlockIndex(From) != -1) {
589      Value *Deleted = Phi.removeIncomingValue(From, false);
590      Map[&Phi].push_back(std::make_pair(From, Deleted));
591    }
592  }
593}
594
595/// Add a dummy PHI value as soon as we knew the new predecessor
596void StructurizeCFG::addPhiValues(BasicBlock *From, BasicBlock *To) {
597  for (PHINode &Phi : To->phis()) {
598    Value *Undef = UndefValue::get(Phi.getType());
599    Phi.addIncoming(Undef, From);
600  }
601  AddedPhis[To].push_back(From);
602}
603
604/// Add the real PHI value as soon as everything is set up
605void StructurizeCFG::setPhiValues() {
606  SmallVector<PHINode *, 8> InsertedPhis;
607  SSAUpdater Updater(&InsertedPhis);
608  for (const auto &AddedPhi : AddedPhis) {
609    BasicBlock *To = AddedPhi.first;
610    const BBVector &From = AddedPhi.second;
611
612    if (!DeletedPhis.count(To))
613      continue;
614
615    PhiMap &Map = DeletedPhis[To];
616    for (const auto &PI : Map) {
617      PHINode *Phi = PI.first;
618      Value *Undef = UndefValue::get(Phi->getType());
619      Updater.Initialize(Phi->getType(), "");
620      Updater.AddAvailableValue(&Func->getEntryBlock(), Undef);
621      Updater.AddAvailableValue(To, Undef);
622
623      NearestCommonDominator Dominator(DT);
624      Dominator.addBlock(To);
625      for (const auto &VI : PI.second) {
626        Updater.AddAvailableValue(VI.first, VI.second);
627        Dominator.addAndRememberBlock(VI.first);
628      }
629
630      if (!Dominator.resultIsRememberedBlock())
631        Updater.AddAvailableValue(Dominator.result(), Undef);
632
633      for (BasicBlock *FI : From)
634        Phi->setIncomingValueForBlock(FI, Updater.GetValueAtEndOfBlock(FI));
635    }
636
637    DeletedPhis.erase(To);
638  }
639  assert(DeletedPhis.empty());
640
641  // Simplify any phis inserted by the SSAUpdater if possible
642  bool Changed;
643  do {
644    Changed = false;
645
646    SimplifyQuery Q(Func->getParent()->getDataLayout());
647    Q.DT = DT;
648    for (size_t i = 0; i < InsertedPhis.size(); ++i) {
649      PHINode *Phi = InsertedPhis[i];
650      if (Value *V = SimplifyInstruction(Phi, Q)) {
651        Phi->replaceAllUsesWith(V);
652        Phi->eraseFromParent();
653        InsertedPhis[i] = InsertedPhis.back();
654        InsertedPhis.pop_back();
655        i--;
656        Changed = true;
657      }
658    }
659  } while (Changed);
660}
661
662/// Remove phi values from all successors and then remove the terminator.
663void StructurizeCFG::killTerminator(BasicBlock *BB) {
664  Instruction *Term = BB->getTerminator();
665  if (!Term)
666    return;
667
668  for (succ_iterator SI = succ_begin(BB), SE = succ_end(BB);
669       SI != SE; ++SI)
670    delPhiValues(BB, *SI);
671
672  if (DA)
673    DA->removeValue(Term);
674  Term->eraseFromParent();
675}
676
677/// Let node exit(s) point to NewExit
678void StructurizeCFG::changeExit(RegionNode *Node, BasicBlock *NewExit,
679                                bool IncludeDominator) {
680  if (Node->isSubRegion()) {
681    Region *SubRegion = Node->getNodeAs<Region>();
682    BasicBlock *OldExit = SubRegion->getExit();
683    BasicBlock *Dominator = nullptr;
684
685    // Find all the edges from the sub region to the exit
686    for (auto BBI = pred_begin(OldExit), E = pred_end(OldExit); BBI != E;) {
687      // Incrememt BBI before mucking with BB's terminator.
688      BasicBlock *BB = *BBI++;
689
690      if (!SubRegion->contains(BB))
691        continue;
692
693      // Modify the edges to point to the new exit
694      delPhiValues(BB, OldExit);
695      BB->getTerminator()->replaceUsesOfWith(OldExit, NewExit);
696      addPhiValues(BB, NewExit);
697
698      // Find the new dominator (if requested)
699      if (IncludeDominator) {
700        if (!Dominator)
701          Dominator = BB;
702        else
703          Dominator = DT->findNearestCommonDominator(Dominator, BB);
704      }
705    }
706
707    // Change the dominator (if requested)
708    if (Dominator)
709      DT->changeImmediateDominator(NewExit, Dominator);
710
711    // Update the region info
712    SubRegion->replaceExit(NewExit);
713  } else {
714    BasicBlock *BB = Node->getNodeAs<BasicBlock>();
715    killTerminator(BB);
716    BranchInst::Create(NewExit, BB);
717    addPhiValues(BB, NewExit);
718    if (IncludeDominator)
719      DT->changeImmediateDominator(NewExit, BB);
720  }
721}
722
723/// Create a new flow node and update dominator tree and region info
724BasicBlock *StructurizeCFG::getNextFlow(BasicBlock *Dominator) {
725  LLVMContext &Context = Func->getContext();
726  BasicBlock *Insert = Order.empty() ? ParentRegion->getExit() :
727                       Order.back()->getEntry();
728  BasicBlock *Flow = BasicBlock::Create(Context, FlowBlockName,
729                                        Func, Insert);
730  DT->addNewBlock(Flow, Dominator);
731  ParentRegion->getRegionInfo()->setRegionFor(Flow, ParentRegion);
732  return Flow;
733}
734
735/// Create a new or reuse the previous node as flow node
736BasicBlock *StructurizeCFG::needPrefix(bool NeedEmpty) {
737  BasicBlock *Entry = PrevNode->getEntry();
738
739  if (!PrevNode->isSubRegion()) {
740    killTerminator(Entry);
741    if (!NeedEmpty || Entry->getFirstInsertionPt() == Entry->end())
742      return Entry;
743  }
744
745  // create a new flow node
746  BasicBlock *Flow = getNextFlow(Entry);
747
748  // and wire it up
749  changeExit(PrevNode, Flow, true);
750  PrevNode = ParentRegion->getBBNode(Flow);
751  return Flow;
752}
753
754/// Returns the region exit if possible, otherwise just a new flow node
755BasicBlock *StructurizeCFG::needPostfix(BasicBlock *Flow,
756                                        bool ExitUseAllowed) {
757  if (!Order.empty() || !ExitUseAllowed)
758    return getNextFlow(Flow);
759
760  BasicBlock *Exit = ParentRegion->getExit();
761  DT->changeImmediateDominator(Exit, Flow);
762  addPhiValues(Flow, Exit);
763  return Exit;
764}
765
766/// Set the previous node
767void StructurizeCFG::setPrevNode(BasicBlock *BB) {
768  PrevNode = ParentRegion->contains(BB) ? ParentRegion->getBBNode(BB)
769                                        : nullptr;
770}
771
772/// Does BB dominate all the predicates of Node?
773bool StructurizeCFG::dominatesPredicates(BasicBlock *BB, RegionNode *Node) {
774  BBPredicates &Preds = Predicates[Node->getEntry()];
775  return llvm::all_of(Preds, [&](std::pair<BasicBlock *, Value *> Pred) {
776    return DT->dominates(BB, Pred.first);
777  });
778}
779
780/// Can we predict that this node will always be called?
781bool StructurizeCFG::isPredictableTrue(RegionNode *Node) {
782  BBPredicates &Preds = Predicates[Node->getEntry()];
783  bool Dominated = false;
784
785  // Regionentry is always true
786  if (!PrevNode)
787    return true;
788
789  for (std::pair<BasicBlock*, Value*> Pred : Preds) {
790    BasicBlock *BB = Pred.first;
791    Value *V = Pred.second;
792
793    if (V != BoolTrue)
794      return false;
795
796    if (!Dominated && DT->dominates(BB, PrevNode->getEntry()))
797      Dominated = true;
798  }
799
800  // TODO: The dominator check is too strict
801  return Dominated;
802}
803
804/// Take one node from the order vector and wire it up
805void StructurizeCFG::wireFlow(bool ExitUseAllowed,
806                              BasicBlock *LoopEnd) {
807  RegionNode *Node = Order.pop_back_val();
808  Visited.insert(Node->getEntry());
809
810  if (isPredictableTrue(Node)) {
811    // Just a linear flow
812    if (PrevNode) {
813      changeExit(PrevNode, Node->getEntry(), true);
814    }
815    PrevNode = Node;
816  } else {
817    // Insert extra prefix node (or reuse last one)
818    BasicBlock *Flow = needPrefix(false);
819
820    // Insert extra postfix node (or use exit instead)
821    BasicBlock *Entry = Node->getEntry();
822    BasicBlock *Next = needPostfix(Flow, ExitUseAllowed);
823
824    // let it point to entry and next block
825    Conditions.push_back(BranchInst::Create(Entry, Next, BoolUndef, Flow));
826    addPhiValues(Flow, Entry);
827    DT->changeImmediateDominator(Entry, Flow);
828
829    PrevNode = Node;
830    while (!Order.empty() && !Visited.count(LoopEnd) &&
831           dominatesPredicates(Entry, Order.back())) {
832      handleLoops(false, LoopEnd);
833    }
834
835    changeExit(PrevNode, Next, false);
836    setPrevNode(Next);
837  }
838}
839
840void StructurizeCFG::handleLoops(bool ExitUseAllowed,
841                                 BasicBlock *LoopEnd) {
842  RegionNode *Node = Order.back();
843  BasicBlock *LoopStart = Node->getEntry();
844
845  if (!Loops.count(LoopStart)) {
846    wireFlow(ExitUseAllowed, LoopEnd);
847    return;
848  }
849
850  if (!isPredictableTrue(Node))
851    LoopStart = needPrefix(true);
852
853  LoopEnd = Loops[Node->getEntry()];
854  wireFlow(false, LoopEnd);
855  while (!Visited.count(LoopEnd)) {
856    handleLoops(false, LoopEnd);
857  }
858
859  // If the start of the loop is the entry block, we can't branch to it so
860  // insert a new dummy entry block.
861  Function *LoopFunc = LoopStart->getParent();
862  if (LoopStart == &LoopFunc->getEntryBlock()) {
863    LoopStart->setName("entry.orig");
864
865    BasicBlock *NewEntry =
866      BasicBlock::Create(LoopStart->getContext(),
867                         "entry",
868                         LoopFunc,
869                         LoopStart);
870    BranchInst::Create(LoopStart, NewEntry);
871    DT->setNewRoot(NewEntry);
872  }
873
874  // Create an extra loop end node
875  LoopEnd = needPrefix(false);
876  BasicBlock *Next = needPostfix(LoopEnd, ExitUseAllowed);
877  LoopConds.push_back(BranchInst::Create(Next, LoopStart,
878                                         BoolUndef, LoopEnd));
879  addPhiValues(LoopEnd, LoopStart);
880  setPrevNode(Next);
881}
882
883/// After this function control flow looks like it should be, but
884/// branches and PHI nodes only have undefined conditions.
885void StructurizeCFG::createFlow() {
886  BasicBlock *Exit = ParentRegion->getExit();
887  bool EntryDominatesExit = DT->dominates(ParentRegion->getEntry(), Exit);
888
889  DeletedPhis.clear();
890  AddedPhis.clear();
891  Conditions.clear();
892  LoopConds.clear();
893
894  PrevNode = nullptr;
895  Visited.clear();
896
897  while (!Order.empty()) {
898    handleLoops(EntryDominatesExit, nullptr);
899  }
900
901  if (PrevNode)
902    changeExit(PrevNode, Exit, EntryDominatesExit);
903  else
904    assert(EntryDominatesExit);
905}
906
907/// Handle a rare case where the disintegrated nodes instructions
908/// no longer dominate all their uses. Not sure if this is really necessary
909void StructurizeCFG::rebuildSSA() {
910  SSAUpdater Updater;
911  for (BasicBlock *BB : ParentRegion->blocks())
912    for (Instruction &I : *BB) {
913      bool Initialized = false;
914      // We may modify the use list as we iterate over it, so be careful to
915      // compute the next element in the use list at the top of the loop.
916      for (auto UI = I.use_begin(), E = I.use_end(); UI != E;) {
917        Use &U = *UI++;
918        Instruction *User = cast<Instruction>(U.getUser());
919        if (User->getParent() == BB) {
920          continue;
921        } else if (PHINode *UserPN = dyn_cast<PHINode>(User)) {
922          if (UserPN->getIncomingBlock(U) == BB)
923            continue;
924        }
925
926        if (DT->dominates(&I, User))
927          continue;
928
929        if (!Initialized) {
930          Value *Undef = UndefValue::get(I.getType());
931          Updater.Initialize(I.getType(), "");
932          Updater.AddAvailableValue(&Func->getEntryBlock(), Undef);
933          Updater.AddAvailableValue(BB, &I);
934          Initialized = true;
935        }
936        Updater.RewriteUseAfterInsertions(U);
937      }
938    }
939}
940
941static bool hasOnlyUniformBranches(Region *R, unsigned UniformMDKindID,
942                                   const LegacyDivergenceAnalysis &DA) {
943  // Bool for if all sub-regions are uniform.
944  bool SubRegionsAreUniform = true;
945  // Count of how many direct children are conditional.
946  unsigned ConditionalDirectChildren = 0;
947
948  for (auto E : R->elements()) {
949    if (!E->isSubRegion()) {
950      auto Br = dyn_cast<BranchInst>(E->getEntry()->getTerminator());
951      if (!Br || !Br->isConditional())
952        continue;
953
954      if (!DA.isUniform(Br))
955        return false;
956
957      // One of our direct children is conditional.
958      ConditionalDirectChildren++;
959
960      LLVM_DEBUG(dbgs() << "BB: " << Br->getParent()->getName()
961                        << " has uniform terminator\n");
962    } else {
963      // Explicitly refuse to treat regions as uniform if they have non-uniform
964      // subregions. We cannot rely on DivergenceAnalysis for branches in
965      // subregions because those branches may have been removed and re-created,
966      // so we look for our metadata instead.
967      //
968      // Warning: It would be nice to treat regions as uniform based only on
969      // their direct child basic blocks' terminators, regardless of whether
970      // subregions are uniform or not. However, this requires a very careful
971      // look at SIAnnotateControlFlow to make sure nothing breaks there.
972      for (auto BB : E->getNodeAs<Region>()->blocks()) {
973        auto Br = dyn_cast<BranchInst>(BB->getTerminator());
974        if (!Br || !Br->isConditional())
975          continue;
976
977        if (!Br->getMetadata(UniformMDKindID)) {
978          // Early exit if we cannot have relaxed uniform regions.
979          if (!RelaxedUniformRegions)
980            return false;
981
982          SubRegionsAreUniform = false;
983          break;
984        }
985      }
986    }
987  }
988
989  // Our region is uniform if:
990  // 1. All conditional branches that are direct children are uniform (checked
991  // above).
992  // 2. And either:
993  //   a. All sub-regions are uniform.
994  //   b. There is one or less conditional branches among the direct children.
995  return SubRegionsAreUniform || (ConditionalDirectChildren <= 1);
996}
997
998/// Run the transformation for each region found
999bool StructurizeCFG::runOnRegion(Region *R, RGPassManager &RGM) {
1000  if (R->isTopLevelRegion())
1001    return false;
1002
1003  DA = nullptr;
1004
1005  if (SkipUniformRegions) {
1006    // TODO: We could probably be smarter here with how we handle sub-regions.
1007    // We currently rely on the fact that metadata is set by earlier invocations
1008    // of the pass on sub-regions, and that this metadata doesn't get lost --
1009    // but we shouldn't rely on metadata for correctness!
1010    unsigned UniformMDKindID =
1011        R->getEntry()->getContext().getMDKindID("structurizecfg.uniform");
1012    DA = &getAnalysis<LegacyDivergenceAnalysis>();
1013
1014    if (hasOnlyUniformBranches(R, UniformMDKindID, *DA)) {
1015      LLVM_DEBUG(dbgs() << "Skipping region with uniform control flow: " << *R
1016                        << '\n');
1017
1018      // Mark all direct child block terminators as having been treated as
1019      // uniform. To account for a possible future in which non-uniform
1020      // sub-regions are treated more cleverly, indirect children are not
1021      // marked as uniform.
1022      MDNode *MD = MDNode::get(R->getEntry()->getParent()->getContext(), {});
1023      for (RegionNode *E : R->elements()) {
1024        if (E->isSubRegion())
1025          continue;
1026
1027        if (Instruction *Term = E->getEntry()->getTerminator())
1028          Term->setMetadata(UniformMDKindID, MD);
1029      }
1030
1031      return false;
1032    }
1033  }
1034
1035  Func = R->getEntry()->getParent();
1036  ParentRegion = R;
1037
1038  DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
1039  LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
1040
1041  orderNodes();
1042  collectInfos();
1043  createFlow();
1044  insertConditions(false);
1045  insertConditions(true);
1046  setPhiValues();
1047  rebuildSSA();
1048
1049  // Cleanup
1050  Order.clear();
1051  Visited.clear();
1052  DeletedPhis.clear();
1053  AddedPhis.clear();
1054  Predicates.clear();
1055  Conditions.clear();
1056  Loops.clear();
1057  LoopPreds.clear();
1058  LoopConds.clear();
1059
1060  return true;
1061}
1062
1063Pass *llvm::createStructurizeCFGPass(bool SkipUniformRegions) {
1064  return new StructurizeCFG(SkipUniformRegions);
1065}
1066