xref: /freebsd/contrib/llvm-project/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp (revision fe75646a0234a261c0013bf1840fdac4acaf0cec)
1 //===- DFAJumpThreading.cpp - Threads a switch statement inside a loop ----===//
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 // Transform each threading path to effectively jump thread the DFA. For
10 // example, the CFG below could be transformed as follows, where the cloned
11 // blocks unconditionally branch to the next correct case based on what is
12 // identified in the analysis.
13 //
14 //          sw.bb                        sw.bb
15 //        /   |   \                    /   |   \
16 //   case1  case2  case3          case1  case2  case3
17 //        \   |   /                 |      |      |
18 //       determinator            det.2   det.3  det.1
19 //        br sw.bb                /        |        \
20 //                          sw.bb.2     sw.bb.3     sw.bb.1
21 //                           br case2    br case3    br case1§
22 //
23 // Definitions and Terminology:
24 //
25 // * Threading path:
26 //   a list of basic blocks, the exit state, and the block that determines
27 //   the next state, for which the following notation will be used:
28 //   < path of BBs that form a cycle > [ state, determinator ]
29 //
30 // * Predictable switch:
31 //   The switch variable is always a known constant so that all conditional
32 //   jumps based on switch variable can be converted to unconditional jump.
33 //
34 // * Determinator:
35 //   The basic block that determines the next state of the DFA.
36 //
37 // Representing the optimization in C-like pseudocode: the code pattern on the
38 // left could functionally be transformed to the right pattern if the switch
39 // condition is predictable.
40 //
41 //  X = A                       goto A
42 //  for (...)                   A:
43 //    switch (X)                  ...
44 //      case A                    goto B
45 //        X = B                 B:
46 //      case B                    ...
47 //        X = C                   goto C
48 //
49 // The pass first checks that switch variable X is decided by the control flow
50 // path taken in the loop; for example, in case B, the next value of X is
51 // decided to be C. It then enumerates through all paths in the loop and labels
52 // the basic blocks where the next state is decided.
53 //
54 // Using this information it creates new paths that unconditionally branch to
55 // the next case. This involves cloning code, so it only gets triggered if the
56 // amount of code duplicated is below a threshold.
57 //
58 //===----------------------------------------------------------------------===//
59 
60 #include "llvm/Transforms/Scalar/DFAJumpThreading.h"
61 #include "llvm/ADT/APInt.h"
62 #include "llvm/ADT/DenseMap.h"
63 #include "llvm/ADT/SmallSet.h"
64 #include "llvm/ADT/Statistic.h"
65 #include "llvm/Analysis/AssumptionCache.h"
66 #include "llvm/Analysis/CodeMetrics.h"
67 #include "llvm/Analysis/DomTreeUpdater.h"
68 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
69 #include "llvm/Analysis/TargetTransformInfo.h"
70 #include "llvm/IR/CFG.h"
71 #include "llvm/IR/Constants.h"
72 #include "llvm/IR/IntrinsicInst.h"
73 #include "llvm/Support/CommandLine.h"
74 #include "llvm/Support/Debug.h"
75 #include "llvm/Transforms/Utils/Cloning.h"
76 #include "llvm/Transforms/Utils/SSAUpdaterBulk.h"
77 #include "llvm/Transforms/Utils/ValueMapper.h"
78 #include <algorithm>
79 #include <deque>
80 
81 #ifdef EXPENSIVE_CHECKS
82 #include "llvm/IR/Verifier.h"
83 #endif
84 
85 using namespace llvm;
86 
87 #define DEBUG_TYPE "dfa-jump-threading"
88 
89 STATISTIC(NumTransforms, "Number of transformations done");
90 STATISTIC(NumCloned, "Number of blocks cloned");
91 STATISTIC(NumPaths, "Number of individual paths threaded");
92 
93 static cl::opt<bool>
94     ClViewCfgBefore("dfa-jump-view-cfg-before",
95                     cl::desc("View the CFG before DFA Jump Threading"),
96                     cl::Hidden, cl::init(false));
97 
98 static cl::opt<unsigned> MaxPathLength(
99     "dfa-max-path-length",
100     cl::desc("Max number of blocks searched to find a threading path"),
101     cl::Hidden, cl::init(20));
102 
103 static cl::opt<unsigned> MaxNumPaths(
104     "dfa-max-num-paths",
105     cl::desc("Max number of paths enumerated around a switch"),
106     cl::Hidden, cl::init(200));
107 
108 static cl::opt<unsigned>
109     CostThreshold("dfa-cost-threshold",
110                   cl::desc("Maximum cost accepted for the transformation"),
111                   cl::Hidden, cl::init(50));
112 
113 namespace {
114 
115 class SelectInstToUnfold {
116   SelectInst *SI;
117   PHINode *SIUse;
118 
119 public:
120   SelectInstToUnfold(SelectInst *SI, PHINode *SIUse) : SI(SI), SIUse(SIUse) {}
121 
122   SelectInst *getInst() { return SI; }
123   PHINode *getUse() { return SIUse; }
124 
125   explicit operator bool() const { return SI && SIUse; }
126 };
127 
128 void unfold(DomTreeUpdater *DTU, SelectInstToUnfold SIToUnfold,
129             std::vector<SelectInstToUnfold> *NewSIsToUnfold,
130             std::vector<BasicBlock *> *NewBBs);
131 
132 class DFAJumpThreading {
133 public:
134   DFAJumpThreading(AssumptionCache *AC, DominatorTree *DT,
135                    TargetTransformInfo *TTI, OptimizationRemarkEmitter *ORE)
136       : AC(AC), DT(DT), TTI(TTI), ORE(ORE) {}
137 
138   bool run(Function &F);
139 
140 private:
141   void
142   unfoldSelectInstrs(DominatorTree *DT,
143                      const SmallVector<SelectInstToUnfold, 4> &SelectInsts) {
144     DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
145     SmallVector<SelectInstToUnfold, 4> Stack;
146     for (SelectInstToUnfold SIToUnfold : SelectInsts)
147       Stack.push_back(SIToUnfold);
148 
149     while (!Stack.empty()) {
150       SelectInstToUnfold SIToUnfold = Stack.pop_back_val();
151 
152       std::vector<SelectInstToUnfold> NewSIsToUnfold;
153       std::vector<BasicBlock *> NewBBs;
154       unfold(&DTU, SIToUnfold, &NewSIsToUnfold, &NewBBs);
155 
156       // Put newly discovered select instructions into the work list.
157       for (const SelectInstToUnfold &NewSIToUnfold : NewSIsToUnfold)
158         Stack.push_back(NewSIToUnfold);
159     }
160   }
161 
162   AssumptionCache *AC;
163   DominatorTree *DT;
164   TargetTransformInfo *TTI;
165   OptimizationRemarkEmitter *ORE;
166 };
167 
168 } // end anonymous namespace
169 
170 namespace {
171 
172 /// Create a new basic block and sink \p SIToSink into it.
173 void createBasicBlockAndSinkSelectInst(
174     DomTreeUpdater *DTU, SelectInst *SI, PHINode *SIUse, SelectInst *SIToSink,
175     BasicBlock *EndBlock, StringRef NewBBName, BasicBlock **NewBlock,
176     BranchInst **NewBranch, std::vector<SelectInstToUnfold> *NewSIsToUnfold,
177     std::vector<BasicBlock *> *NewBBs) {
178   assert(SIToSink->hasOneUse());
179   assert(NewBlock);
180   assert(NewBranch);
181   *NewBlock = BasicBlock::Create(SI->getContext(), NewBBName,
182                                  EndBlock->getParent(), EndBlock);
183   NewBBs->push_back(*NewBlock);
184   *NewBranch = BranchInst::Create(EndBlock, *NewBlock);
185   SIToSink->moveBefore(*NewBranch);
186   NewSIsToUnfold->push_back(SelectInstToUnfold(SIToSink, SIUse));
187   DTU->applyUpdates({{DominatorTree::Insert, *NewBlock, EndBlock}});
188 }
189 
190 /// Unfold the select instruction held in \p SIToUnfold by replacing it with
191 /// control flow.
192 ///
193 /// Put newly discovered select instructions into \p NewSIsToUnfold. Put newly
194 /// created basic blocks into \p NewBBs.
195 ///
196 /// TODO: merge it with CodeGenPrepare::optimizeSelectInst() if possible.
197 void unfold(DomTreeUpdater *DTU, SelectInstToUnfold SIToUnfold,
198             std::vector<SelectInstToUnfold> *NewSIsToUnfold,
199             std::vector<BasicBlock *> *NewBBs) {
200   SelectInst *SI = SIToUnfold.getInst();
201   PHINode *SIUse = SIToUnfold.getUse();
202   BasicBlock *StartBlock = SI->getParent();
203   BasicBlock *EndBlock = SIUse->getParent();
204   BranchInst *StartBlockTerm =
205       dyn_cast<BranchInst>(StartBlock->getTerminator());
206 
207   assert(StartBlockTerm && StartBlockTerm->isUnconditional());
208   assert(SI->hasOneUse());
209 
210   // These are the new basic blocks for the conditional branch.
211   // At least one will become an actual new basic block.
212   BasicBlock *TrueBlock = nullptr;
213   BasicBlock *FalseBlock = nullptr;
214   BranchInst *TrueBranch = nullptr;
215   BranchInst *FalseBranch = nullptr;
216 
217   // Sink select instructions to be able to unfold them later.
218   if (SelectInst *SIOp = dyn_cast<SelectInst>(SI->getTrueValue())) {
219     createBasicBlockAndSinkSelectInst(DTU, SI, SIUse, SIOp, EndBlock,
220                                       "si.unfold.true", &TrueBlock, &TrueBranch,
221                                       NewSIsToUnfold, NewBBs);
222   }
223   if (SelectInst *SIOp = dyn_cast<SelectInst>(SI->getFalseValue())) {
224     createBasicBlockAndSinkSelectInst(DTU, SI, SIUse, SIOp, EndBlock,
225                                       "si.unfold.false", &FalseBlock,
226                                       &FalseBranch, NewSIsToUnfold, NewBBs);
227   }
228 
229   // If there was nothing to sink, then arbitrarily choose the 'false' side
230   // for a new input value to the PHI.
231   if (!TrueBlock && !FalseBlock) {
232     FalseBlock = BasicBlock::Create(SI->getContext(), "si.unfold.false",
233                                     EndBlock->getParent(), EndBlock);
234     NewBBs->push_back(FalseBlock);
235     BranchInst::Create(EndBlock, FalseBlock);
236     DTU->applyUpdates({{DominatorTree::Insert, FalseBlock, EndBlock}});
237   }
238 
239   // Insert the real conditional branch based on the original condition.
240   // If we did not create a new block for one of the 'true' or 'false' paths
241   // of the condition, it means that side of the branch goes to the end block
242   // directly and the path originates from the start block from the point of
243   // view of the new PHI.
244   BasicBlock *TT = EndBlock;
245   BasicBlock *FT = EndBlock;
246   if (TrueBlock && FalseBlock) {
247     // A diamond.
248     TT = TrueBlock;
249     FT = FalseBlock;
250 
251     // Update the phi node of SI.
252     SIUse->removeIncomingValue(StartBlock, /* DeletePHIIfEmpty = */ false);
253     SIUse->addIncoming(SI->getTrueValue(), TrueBlock);
254     SIUse->addIncoming(SI->getFalseValue(), FalseBlock);
255 
256     // Update any other PHI nodes in EndBlock.
257     for (PHINode &Phi : EndBlock->phis()) {
258       if (&Phi != SIUse) {
259         Phi.addIncoming(Phi.getIncomingValueForBlock(StartBlock), TrueBlock);
260         Phi.addIncoming(Phi.getIncomingValueForBlock(StartBlock), FalseBlock);
261       }
262     }
263   } else {
264     BasicBlock *NewBlock = nullptr;
265     Value *SIOp1 = SI->getTrueValue();
266     Value *SIOp2 = SI->getFalseValue();
267 
268     // A triangle pointing right.
269     if (!TrueBlock) {
270       NewBlock = FalseBlock;
271       FT = FalseBlock;
272     }
273     // A triangle pointing left.
274     else {
275       NewBlock = TrueBlock;
276       TT = TrueBlock;
277       std::swap(SIOp1, SIOp2);
278     }
279 
280     // Update the phi node of SI.
281     for (unsigned Idx = 0; Idx < SIUse->getNumIncomingValues(); ++Idx) {
282       if (SIUse->getIncomingBlock(Idx) == StartBlock)
283         SIUse->setIncomingValue(Idx, SIOp1);
284     }
285     SIUse->addIncoming(SIOp2, NewBlock);
286 
287     // Update any other PHI nodes in EndBlock.
288     for (auto II = EndBlock->begin(); PHINode *Phi = dyn_cast<PHINode>(II);
289          ++II) {
290       if (Phi != SIUse)
291         Phi->addIncoming(Phi->getIncomingValueForBlock(StartBlock), NewBlock);
292     }
293   }
294   StartBlockTerm->eraseFromParent();
295   BranchInst::Create(TT, FT, SI->getCondition(), StartBlock);
296   DTU->applyUpdates({{DominatorTree::Insert, StartBlock, TT},
297                      {DominatorTree::Insert, StartBlock, FT}});
298 
299   // The select is now dead.
300   SI->eraseFromParent();
301 }
302 
303 struct ClonedBlock {
304   BasicBlock *BB;
305   uint64_t State; ///< \p State corresponds to the next value of a switch stmnt.
306 };
307 
308 typedef std::deque<BasicBlock *> PathType;
309 typedef std::vector<PathType> PathsType;
310 typedef SmallPtrSet<const BasicBlock *, 8> VisitedBlocks;
311 typedef std::vector<ClonedBlock> CloneList;
312 
313 // This data structure keeps track of all blocks that have been cloned.  If two
314 // different ThreadingPaths clone the same block for a certain state it should
315 // be reused, and it can be looked up in this map.
316 typedef DenseMap<BasicBlock *, CloneList> DuplicateBlockMap;
317 
318 // This map keeps track of all the new definitions for an instruction. This
319 // information is needed when restoring SSA form after cloning blocks.
320 typedef MapVector<Instruction *, std::vector<Instruction *>> DefMap;
321 
322 inline raw_ostream &operator<<(raw_ostream &OS, const PathType &Path) {
323   OS << "< ";
324   for (const BasicBlock *BB : Path) {
325     std::string BBName;
326     if (BB->hasName())
327       raw_string_ostream(BBName) << BB->getName();
328     else
329       raw_string_ostream(BBName) << BB;
330     OS << BBName << " ";
331   }
332   OS << ">";
333   return OS;
334 }
335 
336 /// ThreadingPath is a path in the control flow of a loop that can be threaded
337 /// by cloning necessary basic blocks and replacing conditional branches with
338 /// unconditional ones. A threading path includes a list of basic blocks, the
339 /// exit state, and the block that determines the next state.
340 struct ThreadingPath {
341   /// Exit value is DFA's exit state for the given path.
342   uint64_t getExitValue() const { return ExitVal; }
343   void setExitValue(const ConstantInt *V) {
344     ExitVal = V->getZExtValue();
345     IsExitValSet = true;
346   }
347   bool isExitValueSet() const { return IsExitValSet; }
348 
349   /// Determinator is the basic block that determines the next state of the DFA.
350   const BasicBlock *getDeterminatorBB() const { return DBB; }
351   void setDeterminator(const BasicBlock *BB) { DBB = BB; }
352 
353   /// Path is a list of basic blocks.
354   const PathType &getPath() const { return Path; }
355   void setPath(const PathType &NewPath) { Path = NewPath; }
356 
357   void print(raw_ostream &OS) const {
358     OS << Path << " [ " << ExitVal << ", " << DBB->getName() << " ]";
359   }
360 
361 private:
362   PathType Path;
363   uint64_t ExitVal;
364   const BasicBlock *DBB = nullptr;
365   bool IsExitValSet = false;
366 };
367 
368 #ifndef NDEBUG
369 inline raw_ostream &operator<<(raw_ostream &OS, const ThreadingPath &TPath) {
370   TPath.print(OS);
371   return OS;
372 }
373 #endif
374 
375 struct MainSwitch {
376   MainSwitch(SwitchInst *SI, OptimizationRemarkEmitter *ORE) {
377     if (isCandidate(SI)) {
378       Instr = SI;
379     } else {
380       ORE->emit([&]() {
381         return OptimizationRemarkMissed(DEBUG_TYPE, "SwitchNotPredictable", SI)
382                << "Switch instruction is not predictable.";
383       });
384     }
385   }
386 
387   virtual ~MainSwitch() = default;
388 
389   SwitchInst *getInstr() const { return Instr; }
390   const SmallVector<SelectInstToUnfold, 4> getSelectInsts() {
391     return SelectInsts;
392   }
393 
394 private:
395   /// Do a use-def chain traversal starting from the switch condition to see if
396   /// \p SI is a potential condidate.
397   ///
398   /// Also, collect select instructions to unfold.
399   bool isCandidate(const SwitchInst *SI) {
400     std::deque<Value *> Q;
401     SmallSet<Value *, 16> SeenValues;
402     SelectInsts.clear();
403 
404     Value *SICond = SI->getCondition();
405     LLVM_DEBUG(dbgs() << "\tSICond: " << *SICond << "\n");
406     if (!isa<PHINode>(SICond))
407       return false;
408 
409     addToQueue(SICond, Q, SeenValues);
410 
411     while (!Q.empty()) {
412       Value *Current = Q.front();
413       Q.pop_front();
414 
415       if (auto *Phi = dyn_cast<PHINode>(Current)) {
416         for (Value *Incoming : Phi->incoming_values()) {
417           addToQueue(Incoming, Q, SeenValues);
418         }
419         LLVM_DEBUG(dbgs() << "\tphi: " << *Phi << "\n");
420       } else if (SelectInst *SelI = dyn_cast<SelectInst>(Current)) {
421         if (!isValidSelectInst(SelI))
422           return false;
423         addToQueue(SelI->getTrueValue(), Q, SeenValues);
424         addToQueue(SelI->getFalseValue(), Q, SeenValues);
425         LLVM_DEBUG(dbgs() << "\tselect: " << *SelI << "\n");
426         if (auto *SelIUse = dyn_cast<PHINode>(SelI->user_back()))
427           SelectInsts.push_back(SelectInstToUnfold(SelI, SelIUse));
428       } else if (isa<Constant>(Current)) {
429         LLVM_DEBUG(dbgs() << "\tconst: " << *Current << "\n");
430         continue;
431       } else {
432         LLVM_DEBUG(dbgs() << "\tother: " << *Current << "\n");
433         // Allow unpredictable values. The hope is that those will be the
434         // initial switch values that can be ignored (they will hit the
435         // unthreaded switch) but this assumption will get checked later after
436         // paths have been enumerated (in function getStateDefMap).
437         continue;
438       }
439     }
440 
441     return true;
442   }
443 
444   void addToQueue(Value *Val, std::deque<Value *> &Q,
445                   SmallSet<Value *, 16> &SeenValues) {
446     if (SeenValues.contains(Val))
447       return;
448     Q.push_back(Val);
449     SeenValues.insert(Val);
450   }
451 
452   bool isValidSelectInst(SelectInst *SI) {
453     if (!SI->hasOneUse())
454       return false;
455 
456     Instruction *SIUse = dyn_cast<Instruction>(SI->user_back());
457     // The use of the select inst should be either a phi or another select.
458     if (!SIUse && !(isa<PHINode>(SIUse) || isa<SelectInst>(SIUse)))
459       return false;
460 
461     BasicBlock *SIBB = SI->getParent();
462 
463     // Currently, we can only expand select instructions in basic blocks with
464     // one successor.
465     BranchInst *SITerm = dyn_cast<BranchInst>(SIBB->getTerminator());
466     if (!SITerm || !SITerm->isUnconditional())
467       return false;
468 
469     if (isa<PHINode>(SIUse) &&
470         SIBB->getSingleSuccessor() != cast<Instruction>(SIUse)->getParent())
471       return false;
472 
473     // If select will not be sunk during unfolding, and it is in the same basic
474     // block as another state defining select, then cannot unfold both.
475     for (SelectInstToUnfold SIToUnfold : SelectInsts) {
476       SelectInst *PrevSI = SIToUnfold.getInst();
477       if (PrevSI->getTrueValue() != SI && PrevSI->getFalseValue() != SI &&
478           PrevSI->getParent() == SI->getParent())
479         return false;
480     }
481 
482     return true;
483   }
484 
485   SwitchInst *Instr = nullptr;
486   SmallVector<SelectInstToUnfold, 4> SelectInsts;
487 };
488 
489 struct AllSwitchPaths {
490   AllSwitchPaths(const MainSwitch *MSwitch, OptimizationRemarkEmitter *ORE)
491       : Switch(MSwitch->getInstr()), SwitchBlock(Switch->getParent()),
492         ORE(ORE) {}
493 
494   std::vector<ThreadingPath> &getThreadingPaths() { return TPaths; }
495   unsigned getNumThreadingPaths() { return TPaths.size(); }
496   SwitchInst *getSwitchInst() { return Switch; }
497   BasicBlock *getSwitchBlock() { return SwitchBlock; }
498 
499   void run() {
500     VisitedBlocks Visited;
501     PathsType LoopPaths = paths(SwitchBlock, Visited, /* PathDepth = */ 1);
502     StateDefMap StateDef = getStateDefMap(LoopPaths);
503 
504     if (StateDef.empty()) {
505       ORE->emit([&]() {
506         return OptimizationRemarkMissed(DEBUG_TYPE, "SwitchNotPredictable",
507                                         Switch)
508                << "Switch instruction is not predictable.";
509       });
510       return;
511     }
512 
513     for (PathType Path : LoopPaths) {
514       ThreadingPath TPath;
515 
516       const BasicBlock *PrevBB = Path.back();
517       for (const BasicBlock *BB : Path) {
518         if (StateDef.count(BB) != 0) {
519           const PHINode *Phi = dyn_cast<PHINode>(StateDef[BB]);
520           assert(Phi && "Expected a state-defining instr to be a phi node.");
521 
522           const Value *V = Phi->getIncomingValueForBlock(PrevBB);
523           if (const ConstantInt *C = dyn_cast<const ConstantInt>(V)) {
524             TPath.setExitValue(C);
525             TPath.setDeterminator(BB);
526             TPath.setPath(Path);
527           }
528         }
529 
530         // Switch block is the determinator, this is the final exit value.
531         if (TPath.isExitValueSet() && BB == Path.front())
532           break;
533 
534         PrevBB = BB;
535       }
536 
537       if (TPath.isExitValueSet() && isSupported(TPath))
538         TPaths.push_back(TPath);
539     }
540   }
541 
542 private:
543   // Value: an instruction that defines a switch state;
544   // Key: the parent basic block of that instruction.
545   typedef DenseMap<const BasicBlock *, const PHINode *> StateDefMap;
546 
547   PathsType paths(BasicBlock *BB, VisitedBlocks &Visited,
548                   unsigned PathDepth) const {
549     PathsType Res;
550 
551     // Stop exploring paths after visiting MaxPathLength blocks
552     if (PathDepth > MaxPathLength) {
553       ORE->emit([&]() {
554         return OptimizationRemarkAnalysis(DEBUG_TYPE, "MaxPathLengthReached",
555                                           Switch)
556                << "Exploration stopped after visiting MaxPathLength="
557                << ore::NV("MaxPathLength", MaxPathLength) << " blocks.";
558       });
559       return Res;
560     }
561 
562     Visited.insert(BB);
563 
564     // Some blocks have multiple edges to the same successor, and this set
565     // is used to prevent a duplicate path from being generated
566     SmallSet<BasicBlock *, 4> Successors;
567     for (BasicBlock *Succ : successors(BB)) {
568       if (!Successors.insert(Succ).second)
569         continue;
570 
571       // Found a cycle through the SwitchBlock
572       if (Succ == SwitchBlock) {
573         Res.push_back({BB});
574         continue;
575       }
576 
577       // We have encountered a cycle, do not get caught in it
578       if (Visited.contains(Succ))
579         continue;
580 
581       PathsType SuccPaths = paths(Succ, Visited, PathDepth + 1);
582       for (const PathType &Path : SuccPaths) {
583         PathType NewPath(Path);
584         NewPath.push_front(BB);
585         Res.push_back(NewPath);
586         if (Res.size() >= MaxNumPaths) {
587           return Res;
588         }
589       }
590     }
591     // This block could now be visited again from a different predecessor. Note
592     // that this will result in exponential runtime. Subpaths could possibly be
593     // cached but it takes a lot of memory to store them.
594     Visited.erase(BB);
595     return Res;
596   }
597 
598   /// Walk the use-def chain and collect all the state-defining instructions.
599   ///
600   /// Return an empty map if unpredictable values encountered inside the basic
601   /// blocks of \p LoopPaths.
602   StateDefMap getStateDefMap(const PathsType &LoopPaths) const {
603     StateDefMap Res;
604 
605     // Basic blocks belonging to any of the loops around the switch statement.
606     SmallPtrSet<BasicBlock *, 16> LoopBBs;
607     for (const PathType &Path : LoopPaths) {
608       for (BasicBlock *BB : Path)
609         LoopBBs.insert(BB);
610     }
611 
612     Value *FirstDef = Switch->getOperand(0);
613 
614     assert(isa<PHINode>(FirstDef) && "The first definition must be a phi.");
615 
616     SmallVector<PHINode *, 8> Stack;
617     Stack.push_back(dyn_cast<PHINode>(FirstDef));
618     SmallSet<Value *, 16> SeenValues;
619 
620     while (!Stack.empty()) {
621       PHINode *CurPhi = Stack.pop_back_val();
622 
623       Res[CurPhi->getParent()] = CurPhi;
624       SeenValues.insert(CurPhi);
625 
626       for (BasicBlock *IncomingBB : CurPhi->blocks()) {
627         Value *Incoming = CurPhi->getIncomingValueForBlock(IncomingBB);
628         bool IsOutsideLoops = LoopBBs.count(IncomingBB) == 0;
629         if (Incoming == FirstDef || isa<ConstantInt>(Incoming) ||
630             SeenValues.contains(Incoming) || IsOutsideLoops) {
631           continue;
632         }
633 
634         // Any unpredictable value inside the loops means we must bail out.
635         if (!isa<PHINode>(Incoming))
636           return StateDefMap();
637 
638         Stack.push_back(cast<PHINode>(Incoming));
639       }
640     }
641 
642     return Res;
643   }
644 
645   /// The determinator BB should precede the switch-defining BB.
646   ///
647   /// Otherwise, it is possible that the state defined in the determinator block
648   /// defines the state for the next iteration of the loop, rather than for the
649   /// current one.
650   ///
651   /// Currently supported paths:
652   /// \code
653   /// < switch bb1 determ def > [ 42, determ ]
654   /// < switch_and_def bb1 determ > [ 42, determ ]
655   /// < switch_and_def_and_determ bb1 > [ 42, switch_and_def_and_determ ]
656   /// \endcode
657   ///
658   /// Unsupported paths:
659   /// \code
660   /// < switch bb1 def determ > [ 43, determ ]
661   /// < switch_and_determ bb1 def > [ 43, switch_and_determ ]
662   /// \endcode
663   bool isSupported(const ThreadingPath &TPath) {
664     Instruction *SwitchCondI = dyn_cast<Instruction>(Switch->getCondition());
665     assert(SwitchCondI);
666     if (!SwitchCondI)
667       return false;
668 
669     const BasicBlock *SwitchCondDefBB = SwitchCondI->getParent();
670     const BasicBlock *SwitchCondUseBB = Switch->getParent();
671     const BasicBlock *DeterminatorBB = TPath.getDeterminatorBB();
672 
673     assert(
674         SwitchCondUseBB == TPath.getPath().front() &&
675         "The first BB in a threading path should have the switch instruction");
676     if (SwitchCondUseBB != TPath.getPath().front())
677       return false;
678 
679     // Make DeterminatorBB the first element in Path.
680     PathType Path = TPath.getPath();
681     auto ItDet = llvm::find(Path, DeterminatorBB);
682     std::rotate(Path.begin(), ItDet, Path.end());
683 
684     bool IsDetBBSeen = false;
685     bool IsDefBBSeen = false;
686     bool IsUseBBSeen = false;
687     for (BasicBlock *BB : Path) {
688       if (BB == DeterminatorBB)
689         IsDetBBSeen = true;
690       if (BB == SwitchCondDefBB)
691         IsDefBBSeen = true;
692       if (BB == SwitchCondUseBB)
693         IsUseBBSeen = true;
694       if (IsDetBBSeen && IsUseBBSeen && !IsDefBBSeen)
695         return false;
696     }
697 
698     return true;
699   }
700 
701   SwitchInst *Switch;
702   BasicBlock *SwitchBlock;
703   OptimizationRemarkEmitter *ORE;
704   std::vector<ThreadingPath> TPaths;
705 };
706 
707 struct TransformDFA {
708   TransformDFA(AllSwitchPaths *SwitchPaths, DominatorTree *DT,
709                AssumptionCache *AC, TargetTransformInfo *TTI,
710                OptimizationRemarkEmitter *ORE,
711                SmallPtrSet<const Value *, 32> EphValues)
712       : SwitchPaths(SwitchPaths), DT(DT), AC(AC), TTI(TTI), ORE(ORE),
713         EphValues(EphValues) {}
714 
715   void run() {
716     if (isLegalAndProfitableToTransform()) {
717       createAllExitPaths();
718       NumTransforms++;
719     }
720   }
721 
722 private:
723   /// This function performs both a legality check and profitability check at
724   /// the same time since it is convenient to do so. It iterates through all
725   /// blocks that will be cloned, and keeps track of the duplication cost. It
726   /// also returns false if it is illegal to clone some required block.
727   bool isLegalAndProfitableToTransform() {
728     CodeMetrics Metrics;
729     SwitchInst *Switch = SwitchPaths->getSwitchInst();
730 
731     // Note that DuplicateBlockMap is not being used as intended here. It is
732     // just being used to ensure (BB, State) pairs are only counted once.
733     DuplicateBlockMap DuplicateMap;
734 
735     for (ThreadingPath &TPath : SwitchPaths->getThreadingPaths()) {
736       PathType PathBBs = TPath.getPath();
737       uint64_t NextState = TPath.getExitValue();
738       const BasicBlock *Determinator = TPath.getDeterminatorBB();
739 
740       // Update Metrics for the Switch block, this is always cloned
741       BasicBlock *BB = SwitchPaths->getSwitchBlock();
742       BasicBlock *VisitedBB = getClonedBB(BB, NextState, DuplicateMap);
743       if (!VisitedBB) {
744         Metrics.analyzeBasicBlock(BB, *TTI, EphValues);
745         DuplicateMap[BB].push_back({BB, NextState});
746       }
747 
748       // If the Switch block is the Determinator, then we can continue since
749       // this is the only block that is cloned and we already counted for it.
750       if (PathBBs.front() == Determinator)
751         continue;
752 
753       // Otherwise update Metrics for all blocks that will be cloned. If any
754       // block is already cloned and would be reused, don't double count it.
755       auto DetIt = llvm::find(PathBBs, Determinator);
756       for (auto BBIt = DetIt; BBIt != PathBBs.end(); BBIt++) {
757         BB = *BBIt;
758         VisitedBB = getClonedBB(BB, NextState, DuplicateMap);
759         if (VisitedBB)
760           continue;
761         Metrics.analyzeBasicBlock(BB, *TTI, EphValues);
762         DuplicateMap[BB].push_back({BB, NextState});
763       }
764 
765       if (Metrics.notDuplicatable) {
766         LLVM_DEBUG(dbgs() << "DFA Jump Threading: Not jump threading, contains "
767                           << "non-duplicatable instructions.\n");
768         ORE->emit([&]() {
769           return OptimizationRemarkMissed(DEBUG_TYPE, "NonDuplicatableInst",
770                                           Switch)
771                  << "Contains non-duplicatable instructions.";
772         });
773         return false;
774       }
775 
776       if (Metrics.convergent) {
777         LLVM_DEBUG(dbgs() << "DFA Jump Threading: Not jump threading, contains "
778                           << "convergent instructions.\n");
779         ORE->emit([&]() {
780           return OptimizationRemarkMissed(DEBUG_TYPE, "ConvergentInst", Switch)
781                  << "Contains convergent instructions.";
782         });
783         return false;
784       }
785 
786       if (!Metrics.NumInsts.isValid()) {
787         LLVM_DEBUG(dbgs() << "DFA Jump Threading: Not jump threading, contains "
788                           << "instructions with invalid cost.\n");
789         ORE->emit([&]() {
790           return OptimizationRemarkMissed(DEBUG_TYPE, "ConvergentInst", Switch)
791                  << "Contains instructions with invalid cost.";
792         });
793         return false;
794       }
795     }
796 
797     InstructionCost DuplicationCost = 0;
798 
799     unsigned JumpTableSize = 0;
800     TTI->getEstimatedNumberOfCaseClusters(*Switch, JumpTableSize, nullptr,
801                                           nullptr);
802     if (JumpTableSize == 0) {
803       // Factor in the number of conditional branches reduced from jump
804       // threading. Assume that lowering the switch block is implemented by
805       // using binary search, hence the LogBase2().
806       unsigned CondBranches =
807           APInt(32, Switch->getNumSuccessors()).ceilLogBase2();
808       DuplicationCost = Metrics.NumInsts / CondBranches;
809     } else {
810       // Compared with jump tables, the DFA optimizer removes an indirect branch
811       // on each loop iteration, thus making branch prediction more precise. The
812       // more branch targets there are, the more likely it is for the branch
813       // predictor to make a mistake, and the more benefit there is in the DFA
814       // optimizer. Thus, the more branch targets there are, the lower is the
815       // cost of the DFA opt.
816       DuplicationCost = Metrics.NumInsts / JumpTableSize;
817     }
818 
819     LLVM_DEBUG(dbgs() << "\nDFA Jump Threading: Cost to jump thread block "
820                       << SwitchPaths->getSwitchBlock()->getName()
821                       << " is: " << DuplicationCost << "\n\n");
822 
823     if (DuplicationCost > CostThreshold) {
824       LLVM_DEBUG(dbgs() << "Not jump threading, duplication cost exceeds the "
825                         << "cost threshold.\n");
826       ORE->emit([&]() {
827         return OptimizationRemarkMissed(DEBUG_TYPE, "NotProfitable", Switch)
828                << "Duplication cost exceeds the cost threshold (cost="
829                << ore::NV("Cost", DuplicationCost)
830                << ", threshold=" << ore::NV("Threshold", CostThreshold) << ").";
831       });
832       return false;
833     }
834 
835     ORE->emit([&]() {
836       return OptimizationRemark(DEBUG_TYPE, "JumpThreaded", Switch)
837              << "Switch statement jump-threaded.";
838     });
839 
840     return true;
841   }
842 
843   /// Transform each threading path to effectively jump thread the DFA.
844   void createAllExitPaths() {
845     DomTreeUpdater DTU(*DT, DomTreeUpdater::UpdateStrategy::Eager);
846 
847     // Move the switch block to the end of the path, since it will be duplicated
848     BasicBlock *SwitchBlock = SwitchPaths->getSwitchBlock();
849     for (ThreadingPath &TPath : SwitchPaths->getThreadingPaths()) {
850       LLVM_DEBUG(dbgs() << TPath << "\n");
851       PathType NewPath(TPath.getPath());
852       NewPath.push_back(SwitchBlock);
853       TPath.setPath(NewPath);
854     }
855 
856     // Transform the ThreadingPaths and keep track of the cloned values
857     DuplicateBlockMap DuplicateMap;
858     DefMap NewDefs;
859 
860     SmallSet<BasicBlock *, 16> BlocksToClean;
861     for (BasicBlock *BB : successors(SwitchBlock))
862       BlocksToClean.insert(BB);
863 
864     for (ThreadingPath &TPath : SwitchPaths->getThreadingPaths()) {
865       createExitPath(NewDefs, TPath, DuplicateMap, BlocksToClean, &DTU);
866       NumPaths++;
867     }
868 
869     // After all paths are cloned, now update the last successor of the cloned
870     // path so it skips over the switch statement
871     for (ThreadingPath &TPath : SwitchPaths->getThreadingPaths())
872       updateLastSuccessor(TPath, DuplicateMap, &DTU);
873 
874     // For each instruction that was cloned and used outside, update its uses
875     updateSSA(NewDefs);
876 
877     // Clean PHI Nodes for the newly created blocks
878     for (BasicBlock *BB : BlocksToClean)
879       cleanPhiNodes(BB);
880   }
881 
882   /// For a specific ThreadingPath \p Path, create an exit path starting from
883   /// the determinator block.
884   ///
885   /// To remember the correct destination, we have to duplicate blocks
886   /// corresponding to each state. Also update the terminating instruction of
887   /// the predecessors, and phis in the successor blocks.
888   void createExitPath(DefMap &NewDefs, ThreadingPath &Path,
889                       DuplicateBlockMap &DuplicateMap,
890                       SmallSet<BasicBlock *, 16> &BlocksToClean,
891                       DomTreeUpdater *DTU) {
892     uint64_t NextState = Path.getExitValue();
893     const BasicBlock *Determinator = Path.getDeterminatorBB();
894     PathType PathBBs = Path.getPath();
895 
896     // Don't select the placeholder block in front
897     if (PathBBs.front() == Determinator)
898       PathBBs.pop_front();
899 
900     auto DetIt = llvm::find(PathBBs, Determinator);
901     auto Prev = std::prev(DetIt);
902     BasicBlock *PrevBB = *Prev;
903     for (auto BBIt = DetIt; BBIt != PathBBs.end(); BBIt++) {
904       BasicBlock *BB = *BBIt;
905       BlocksToClean.insert(BB);
906 
907       // We already cloned BB for this NextState, now just update the branch
908       // and continue.
909       BasicBlock *NextBB = getClonedBB(BB, NextState, DuplicateMap);
910       if (NextBB) {
911         updatePredecessor(PrevBB, BB, NextBB, DTU);
912         PrevBB = NextBB;
913         continue;
914       }
915 
916       // Clone the BB and update the successor of Prev to jump to the new block
917       BasicBlock *NewBB = cloneBlockAndUpdatePredecessor(
918           BB, PrevBB, NextState, DuplicateMap, NewDefs, DTU);
919       DuplicateMap[BB].push_back({NewBB, NextState});
920       BlocksToClean.insert(NewBB);
921       PrevBB = NewBB;
922     }
923   }
924 
925   /// Restore SSA form after cloning blocks.
926   ///
927   /// Each cloned block creates new defs for a variable, and the uses need to be
928   /// updated to reflect this. The uses may be replaced with a cloned value, or
929   /// some derived phi instruction. Note that all uses of a value defined in the
930   /// same block were already remapped when cloning the block.
931   void updateSSA(DefMap &NewDefs) {
932     SSAUpdaterBulk SSAUpdate;
933     SmallVector<Use *, 16> UsesToRename;
934 
935     for (const auto &KV : NewDefs) {
936       Instruction *I = KV.first;
937       BasicBlock *BB = I->getParent();
938       std::vector<Instruction *> Cloned = KV.second;
939 
940       // Scan all uses of this instruction to see if it is used outside of its
941       // block, and if so, record them in UsesToRename.
942       for (Use &U : I->uses()) {
943         Instruction *User = cast<Instruction>(U.getUser());
944         if (PHINode *UserPN = dyn_cast<PHINode>(User)) {
945           if (UserPN->getIncomingBlock(U) == BB)
946             continue;
947         } else if (User->getParent() == BB) {
948           continue;
949         }
950 
951         UsesToRename.push_back(&U);
952       }
953 
954       // If there are no uses outside the block, we're done with this
955       // instruction.
956       if (UsesToRename.empty())
957         continue;
958       LLVM_DEBUG(dbgs() << "DFA-JT: Renaming non-local uses of: " << *I
959                         << "\n");
960 
961       // We found a use of I outside of BB.  Rename all uses of I that are
962       // outside its block to be uses of the appropriate PHI node etc.  See
963       // ValuesInBlocks with the values we know.
964       unsigned VarNum = SSAUpdate.AddVariable(I->getName(), I->getType());
965       SSAUpdate.AddAvailableValue(VarNum, BB, I);
966       for (Instruction *New : Cloned)
967         SSAUpdate.AddAvailableValue(VarNum, New->getParent(), New);
968 
969       while (!UsesToRename.empty())
970         SSAUpdate.AddUse(VarNum, UsesToRename.pop_back_val());
971 
972       LLVM_DEBUG(dbgs() << "\n");
973     }
974     // SSAUpdater handles phi placement and renaming uses with the appropriate
975     // value.
976     SSAUpdate.RewriteAllUses(DT);
977   }
978 
979   /// Clones a basic block, and adds it to the CFG.
980   ///
981   /// This function also includes updating phi nodes in the successors of the
982   /// BB, and remapping uses that were defined locally in the cloned BB.
983   BasicBlock *cloneBlockAndUpdatePredecessor(BasicBlock *BB, BasicBlock *PrevBB,
984                                              uint64_t NextState,
985                                              DuplicateBlockMap &DuplicateMap,
986                                              DefMap &NewDefs,
987                                              DomTreeUpdater *DTU) {
988     ValueToValueMapTy VMap;
989     BasicBlock *NewBB = CloneBasicBlock(
990         BB, VMap, ".jt" + std::to_string(NextState), BB->getParent());
991     NewBB->moveAfter(BB);
992     NumCloned++;
993 
994     for (Instruction &I : *NewBB) {
995       // Do not remap operands of PHINode in case a definition in BB is an
996       // incoming value to a phi in the same block. This incoming value will
997       // be renamed later while restoring SSA.
998       if (isa<PHINode>(&I))
999         continue;
1000       RemapInstruction(&I, VMap,
1001                        RF_IgnoreMissingLocals | RF_NoModuleLevelChanges);
1002       if (AssumeInst *II = dyn_cast<AssumeInst>(&I))
1003         AC->registerAssumption(II);
1004     }
1005 
1006     updateSuccessorPhis(BB, NewBB, NextState, VMap, DuplicateMap);
1007     updatePredecessor(PrevBB, BB, NewBB, DTU);
1008     updateDefMap(NewDefs, VMap);
1009 
1010     // Add all successors to the DominatorTree
1011     SmallPtrSet<BasicBlock *, 4> SuccSet;
1012     for (auto *SuccBB : successors(NewBB)) {
1013       if (SuccSet.insert(SuccBB).second)
1014         DTU->applyUpdates({{DominatorTree::Insert, NewBB, SuccBB}});
1015     }
1016     SuccSet.clear();
1017     return NewBB;
1018   }
1019 
1020   /// Update the phi nodes in BB's successors.
1021   ///
1022   /// This means creating a new incoming value from NewBB with the new
1023   /// instruction wherever there is an incoming value from BB.
1024   void updateSuccessorPhis(BasicBlock *BB, BasicBlock *ClonedBB,
1025                            uint64_t NextState, ValueToValueMapTy &VMap,
1026                            DuplicateBlockMap &DuplicateMap) {
1027     std::vector<BasicBlock *> BlocksToUpdate;
1028 
1029     // If BB is the last block in the path, we can simply update the one case
1030     // successor that will be reached.
1031     if (BB == SwitchPaths->getSwitchBlock()) {
1032       SwitchInst *Switch = SwitchPaths->getSwitchInst();
1033       BasicBlock *NextCase = getNextCaseSuccessor(Switch, NextState);
1034       BlocksToUpdate.push_back(NextCase);
1035       BasicBlock *ClonedSucc = getClonedBB(NextCase, NextState, DuplicateMap);
1036       if (ClonedSucc)
1037         BlocksToUpdate.push_back(ClonedSucc);
1038     }
1039     // Otherwise update phis in all successors.
1040     else {
1041       for (BasicBlock *Succ : successors(BB)) {
1042         BlocksToUpdate.push_back(Succ);
1043 
1044         // Check if a successor has already been cloned for the particular exit
1045         // value. In this case if a successor was already cloned, the phi nodes
1046         // in the cloned block should be updated directly.
1047         BasicBlock *ClonedSucc = getClonedBB(Succ, NextState, DuplicateMap);
1048         if (ClonedSucc)
1049           BlocksToUpdate.push_back(ClonedSucc);
1050       }
1051     }
1052 
1053     // If there is a phi with an incoming value from BB, create a new incoming
1054     // value for the new predecessor ClonedBB. The value will either be the same
1055     // value from BB or a cloned value.
1056     for (BasicBlock *Succ : BlocksToUpdate) {
1057       for (auto II = Succ->begin(); PHINode *Phi = dyn_cast<PHINode>(II);
1058            ++II) {
1059         Value *Incoming = Phi->getIncomingValueForBlock(BB);
1060         if (Incoming) {
1061           if (isa<Constant>(Incoming)) {
1062             Phi->addIncoming(Incoming, ClonedBB);
1063             continue;
1064           }
1065           Value *ClonedVal = VMap[Incoming];
1066           if (ClonedVal)
1067             Phi->addIncoming(ClonedVal, ClonedBB);
1068           else
1069             Phi->addIncoming(Incoming, ClonedBB);
1070         }
1071       }
1072     }
1073   }
1074 
1075   /// Sets the successor of PrevBB to be NewBB instead of OldBB. Note that all
1076   /// other successors are kept as well.
1077   void updatePredecessor(BasicBlock *PrevBB, BasicBlock *OldBB,
1078                          BasicBlock *NewBB, DomTreeUpdater *DTU) {
1079     // When a path is reused, there is a chance that predecessors were already
1080     // updated before. Check if the predecessor needs to be updated first.
1081     if (!isPredecessor(OldBB, PrevBB))
1082       return;
1083 
1084     Instruction *PrevTerm = PrevBB->getTerminator();
1085     for (unsigned Idx = 0; Idx < PrevTerm->getNumSuccessors(); Idx++) {
1086       if (PrevTerm->getSuccessor(Idx) == OldBB) {
1087         OldBB->removePredecessor(PrevBB, /* KeepOneInputPHIs = */ true);
1088         PrevTerm->setSuccessor(Idx, NewBB);
1089       }
1090     }
1091     DTU->applyUpdates({{DominatorTree::Delete, PrevBB, OldBB},
1092                        {DominatorTree::Insert, PrevBB, NewBB}});
1093   }
1094 
1095   /// Add new value mappings to the DefMap to keep track of all new definitions
1096   /// for a particular instruction. These will be used while updating SSA form.
1097   void updateDefMap(DefMap &NewDefs, ValueToValueMapTy &VMap) {
1098     SmallVector<std::pair<Instruction *, Instruction *>> NewDefsVector;
1099     NewDefsVector.reserve(VMap.size());
1100 
1101     for (auto Entry : VMap) {
1102       Instruction *Inst =
1103           dyn_cast<Instruction>(const_cast<Value *>(Entry.first));
1104       if (!Inst || !Entry.second || isa<BranchInst>(Inst) ||
1105           isa<SwitchInst>(Inst)) {
1106         continue;
1107       }
1108 
1109       Instruction *Cloned = dyn_cast<Instruction>(Entry.second);
1110       if (!Cloned)
1111         continue;
1112 
1113       NewDefsVector.push_back({Inst, Cloned});
1114     }
1115 
1116     // Sort the defs to get deterministic insertion order into NewDefs.
1117     sort(NewDefsVector, [](const auto &LHS, const auto &RHS) {
1118       if (LHS.first == RHS.first)
1119         return LHS.second->comesBefore(RHS.second);
1120       return LHS.first->comesBefore(RHS.first);
1121     });
1122 
1123     for (const auto &KV : NewDefsVector)
1124       NewDefs[KV.first].push_back(KV.second);
1125   }
1126 
1127   /// Update the last branch of a particular cloned path to point to the correct
1128   /// case successor.
1129   ///
1130   /// Note that this is an optional step and would have been done in later
1131   /// optimizations, but it makes the CFG significantly easier to work with.
1132   void updateLastSuccessor(ThreadingPath &TPath,
1133                            DuplicateBlockMap &DuplicateMap,
1134                            DomTreeUpdater *DTU) {
1135     uint64_t NextState = TPath.getExitValue();
1136     BasicBlock *BB = TPath.getPath().back();
1137     BasicBlock *LastBlock = getClonedBB(BB, NextState, DuplicateMap);
1138 
1139     // Note multiple paths can end at the same block so check that it is not
1140     // updated yet
1141     if (!isa<SwitchInst>(LastBlock->getTerminator()))
1142       return;
1143     SwitchInst *Switch = cast<SwitchInst>(LastBlock->getTerminator());
1144     BasicBlock *NextCase = getNextCaseSuccessor(Switch, NextState);
1145 
1146     std::vector<DominatorTree::UpdateType> DTUpdates;
1147     SmallPtrSet<BasicBlock *, 4> SuccSet;
1148     for (BasicBlock *Succ : successors(LastBlock)) {
1149       if (Succ != NextCase && SuccSet.insert(Succ).second)
1150         DTUpdates.push_back({DominatorTree::Delete, LastBlock, Succ});
1151     }
1152 
1153     Switch->eraseFromParent();
1154     BranchInst::Create(NextCase, LastBlock);
1155 
1156     DTU->applyUpdates(DTUpdates);
1157   }
1158 
1159   /// After cloning blocks, some of the phi nodes have extra incoming values
1160   /// that are no longer used. This function removes them.
1161   void cleanPhiNodes(BasicBlock *BB) {
1162     // If BB is no longer reachable, remove any remaining phi nodes
1163     if (pred_empty(BB)) {
1164       std::vector<PHINode *> PhiToRemove;
1165       for (auto II = BB->begin(); PHINode *Phi = dyn_cast<PHINode>(II); ++II) {
1166         PhiToRemove.push_back(Phi);
1167       }
1168       for (PHINode *PN : PhiToRemove) {
1169         PN->replaceAllUsesWith(PoisonValue::get(PN->getType()));
1170         PN->eraseFromParent();
1171       }
1172       return;
1173     }
1174 
1175     // Remove any incoming values that come from an invalid predecessor
1176     for (auto II = BB->begin(); PHINode *Phi = dyn_cast<PHINode>(II); ++II) {
1177       std::vector<BasicBlock *> BlocksToRemove;
1178       for (BasicBlock *IncomingBB : Phi->blocks()) {
1179         if (!isPredecessor(BB, IncomingBB))
1180           BlocksToRemove.push_back(IncomingBB);
1181       }
1182       for (BasicBlock *BB : BlocksToRemove)
1183         Phi->removeIncomingValue(BB);
1184     }
1185   }
1186 
1187   /// Checks if BB was already cloned for a particular next state value. If it
1188   /// was then it returns this cloned block, and otherwise null.
1189   BasicBlock *getClonedBB(BasicBlock *BB, uint64_t NextState,
1190                           DuplicateBlockMap &DuplicateMap) {
1191     CloneList ClonedBBs = DuplicateMap[BB];
1192 
1193     // Find an entry in the CloneList with this NextState. If it exists then
1194     // return the corresponding BB
1195     auto It = llvm::find_if(ClonedBBs, [NextState](const ClonedBlock &C) {
1196       return C.State == NextState;
1197     });
1198     return It != ClonedBBs.end() ? (*It).BB : nullptr;
1199   }
1200 
1201   /// Helper to get the successor corresponding to a particular case value for
1202   /// a switch statement.
1203   BasicBlock *getNextCaseSuccessor(SwitchInst *Switch, uint64_t NextState) {
1204     BasicBlock *NextCase = nullptr;
1205     for (auto Case : Switch->cases()) {
1206       if (Case.getCaseValue()->getZExtValue() == NextState) {
1207         NextCase = Case.getCaseSuccessor();
1208         break;
1209       }
1210     }
1211     if (!NextCase)
1212       NextCase = Switch->getDefaultDest();
1213     return NextCase;
1214   }
1215 
1216   /// Returns true if IncomingBB is a predecessor of BB.
1217   bool isPredecessor(BasicBlock *BB, BasicBlock *IncomingBB) {
1218     return llvm::is_contained(predecessors(BB), IncomingBB);
1219   }
1220 
1221   AllSwitchPaths *SwitchPaths;
1222   DominatorTree *DT;
1223   AssumptionCache *AC;
1224   TargetTransformInfo *TTI;
1225   OptimizationRemarkEmitter *ORE;
1226   SmallPtrSet<const Value *, 32> EphValues;
1227   std::vector<ThreadingPath> TPaths;
1228 };
1229 
1230 bool DFAJumpThreading::run(Function &F) {
1231   LLVM_DEBUG(dbgs() << "\nDFA Jump threading: " << F.getName() << "\n");
1232 
1233   if (F.hasOptSize()) {
1234     LLVM_DEBUG(dbgs() << "Skipping due to the 'minsize' attribute\n");
1235     return false;
1236   }
1237 
1238   if (ClViewCfgBefore)
1239     F.viewCFG();
1240 
1241   SmallVector<AllSwitchPaths, 2> ThreadableLoops;
1242   bool MadeChanges = false;
1243 
1244   for (BasicBlock &BB : F) {
1245     auto *SI = dyn_cast<SwitchInst>(BB.getTerminator());
1246     if (!SI)
1247       continue;
1248 
1249     LLVM_DEBUG(dbgs() << "\nCheck if SwitchInst in BB " << BB.getName()
1250                       << " is a candidate\n");
1251     MainSwitch Switch(SI, ORE);
1252 
1253     if (!Switch.getInstr())
1254       continue;
1255 
1256     LLVM_DEBUG(dbgs() << "\nSwitchInst in BB " << BB.getName() << " is a "
1257                       << "candidate for jump threading\n");
1258     LLVM_DEBUG(SI->dump());
1259 
1260     unfoldSelectInstrs(DT, Switch.getSelectInsts());
1261     if (!Switch.getSelectInsts().empty())
1262       MadeChanges = true;
1263 
1264     AllSwitchPaths SwitchPaths(&Switch, ORE);
1265     SwitchPaths.run();
1266 
1267     if (SwitchPaths.getNumThreadingPaths() > 0) {
1268       ThreadableLoops.push_back(SwitchPaths);
1269 
1270       // For the time being limit this optimization to occurring once in a
1271       // function since it can change the CFG significantly. This is not a
1272       // strict requirement but it can cause buggy behavior if there is an
1273       // overlap of blocks in different opportunities. There is a lot of room to
1274       // experiment with catching more opportunities here.
1275       break;
1276     }
1277   }
1278 
1279   SmallPtrSet<const Value *, 32> EphValues;
1280   if (ThreadableLoops.size() > 0)
1281     CodeMetrics::collectEphemeralValues(&F, AC, EphValues);
1282 
1283   for (AllSwitchPaths SwitchPaths : ThreadableLoops) {
1284     TransformDFA Transform(&SwitchPaths, DT, AC, TTI, ORE, EphValues);
1285     Transform.run();
1286     MadeChanges = true;
1287   }
1288 
1289 #ifdef EXPENSIVE_CHECKS
1290   assert(DT->verify(DominatorTree::VerificationLevel::Full));
1291   verifyFunction(F, &dbgs());
1292 #endif
1293 
1294   return MadeChanges;
1295 }
1296 
1297 } // end anonymous namespace
1298 
1299 /// Integrate with the new Pass Manager
1300 PreservedAnalyses DFAJumpThreadingPass::run(Function &F,
1301                                             FunctionAnalysisManager &AM) {
1302   AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(F);
1303   DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(F);
1304   TargetTransformInfo &TTI = AM.getResult<TargetIRAnalysis>(F);
1305   OptimizationRemarkEmitter ORE(&F);
1306 
1307   if (!DFAJumpThreading(&AC, &DT, &TTI, &ORE).run(F))
1308     return PreservedAnalyses::all();
1309 
1310   PreservedAnalyses PA;
1311   PA.preserve<DominatorTreeAnalysis>();
1312   return PA;
1313 }
1314