PredicateInfo: Support switch statements

[• dberlin]

Summary

Depends on D29606 and D29682

Makes us pass GVN's edge.ll (we also will pass a few other testcases
they just need cleaning up).

Thoughts on the Predicate* hiearchy of classes especially welcome :)
(it's not clear to me how best to organize it, and currently, the getBlock* seems ... uglier than maybe wasting a field somewhere or something).

[• dberlin]

lib/Transforms/Scalar/NewGVN.cpp

862

sigh, clang format, i'll fix

876

ditto

880

ditto

[• dberlin]

• Update new tests

[• dberlin]

lib/Transforms/Scalar/NewGVN.cpp

844

See comment below about branch having two parents.
The switch code is the only real new code in this function,

1063–1064

Just a friendly note: This branch has two parents (the critical edge patch and the initial newgvn patch).
So this diff shows a lot more than is actually here.
In particular, this function is not really being added here.

[davide]

Once the basic predicate support goes in, can you please rebase this patch so I can review it more easily? Thanks!

[• dberlin]

Yup, as soon as those two are committed i'll rebase and ping you.

[• dberlin]

Rebase on top of master and clean up.

[davide]

lib/Transforms/Scalar/NewGVN.cpp

1094

Spurious newline

lib/Transforms/Utils/PredicateInfo.cpp

56–59

assert(isa<>) + return static_cast<> ?

65–67

Also here and in the other places.

425–433

range for?

[• dberlin]

• Update for review comments

Note: This adds two pair conversion function for BasicBlockEdge in
dominators to make life easier. We can't just use basicblockedge
ourselves, because we need non-const basic blocks in a few places.

[davide]

LGTM. Thanks.

include/llvm/Transforms/Utils/PredicateInfo.h

 class DominatorTree;
 class Function;
 class Instruction;
 class MemoryAccess;
 class LLVMContext;
 class raw_ostream;
 class OrderedBasicBlock;
 
-enum PredicateType { PT_Branch, PT_Assume };
+enum PredicateType { PT_Branch, PT_Assume, PT_Switch };
 
 // Base class for all predicate information we provide.
 // All of our predicate information has at least a comparison.
 class PredicateBase : public ilist_node<PredicateBase> {
 public:
   PredicateType Type;
   // The original operand before we renamed it.
   // This can be use by passes, when destroying predicateinfo, to know
   // whether they can just drop the intrinsic, or have to merge metadata.
   Value *OriginalOp;
-  CmpInst *Comparison;
   PredicateBase(const PredicateBase &) = delete;
   PredicateBase &operator=(const PredicateBase &) = delete;
   PredicateBase() = delete;
 
 protected:
-  PredicateBase(PredicateType PT, Value *Op, CmpInst *Comparison)
-      : Type(PT), OriginalOp(Op), Comparison(Comparison) {}
+  PredicateBase(PredicateType PT, Value *Op) : Type(PT), OriginalOp(Op) {}
+};
+
+class PredicateWithComparison : public PredicateBase {
+public:
+  CmpInst *Comparison;
+  static inline bool classof(const PredicateBase *PB) {
+    return PB->Type == PT_Assume || PB->Type == PT_Branch;
+  }
+
+protected:
+  PredicateWithComparison(PredicateType PT, Value *Op, CmpInst *Comparison)
+      : PredicateBase(PT, Op), Comparison(Comparison) {}
 };
 
 // Provides predicate information for assumes.  Since assumes are always true,
 // we simply provide the assume instruction, so you can tell your relative
 // position to it.
-class PredicateAssume : public PredicateBase {
+class PredicateAssume : public PredicateWithComparison {
 public:
   IntrinsicInst *AssumeInst;
   PredicateAssume(Value *Op, IntrinsicInst *AssumeInst, CmpInst *Comparison)
-      : PredicateBase(PT_Assume, Op, Comparison), AssumeInst(AssumeInst) {}
+      : PredicateWithComparison(PT_Assume, Op, Comparison),
+        AssumeInst(AssumeInst) {}
   PredicateAssume() = delete;
   static inline bool classof(const PredicateBase *PB) {
     return PB->Type == PT_Assume;
   }
 };
 
 // Provides predicate information for branches.
-class PredicateBranch : public PredicateBase {
+class PredicateBranch : public PredicateWithComparison {
 public:
   // This is the block that is conditional upon the comparison.
   BasicBlock *BranchBB;
   // This is one of the true/false successors of BranchBB.
   BasicBlock *SplitBB;
   // If true, SplitBB is the true successor, otherwise it's the false successor.
   bool TrueEdge;
   PredicateBranch(Value *Op, BasicBlock *BranchBB, BasicBlock *SplitBB,
                   CmpInst *Comparison, bool TakenEdge)
-      : PredicateBase(PT_Branch, Op, Comparison), BranchBB(BranchBB),
+      : PredicateWithComparison(PT_Branch, Op, Comparison), BranchBB(BranchBB),
         SplitBB(SplitBB), TrueEdge(TakenEdge) {}
   PredicateBranch() = delete;
   static inline bool classof(const PredicateBase *PB) {
     return PB->Type == PT_Branch;
   }
 };
 
+class PredicateSwitch : public PredicateBase {
+public:
+  // This is the block of the switch.
+  BasicBlock *SwitchBB;
+  // This is the target block of the case statement.
+  BasicBlock *TargetBB;
+  // This is the case value for the switch to that block.
+  Value *CaseValue;
+  // This is the switch instruction.
+  SwitchInst *Switch;
+  PredicateSwitch(Value *Op, BasicBlock *SwitchBB, BasicBlock *TargetBB,
+                  Value *CaseValue, SwitchInst *SI)
+      : PredicateBase(PT_Switch, Op), SwitchBB(SwitchBB), TargetBB(TargetBB),
+        CaseValue(CaseValue), Switch(SI) {}
+  static inline bool classof(const PredicateBase *PB) {
+    return PB->Type == PT_Switch;
+  }
+};
+
 // This name is used in a few places, so kick it into their own namespace
 namespace PredicateInfoClasses {
 struct ValueDFS;
 }
 
 /// \brief Encapsulates PredicateInfo, including all data associated with memory
 /// accesses.
 class PredicateInfo {
protected:
   // Used by PredicateInfo annotater, dumpers, and wrapper pass.
   friend class PredicateInfoAnnotatedWriter;
   friend class PredicateInfoPrinterLegacyPass;
 
 private:
   void buildPredicateInfo();
   void processAssume(IntrinsicInst *, BasicBlock *, SmallPtrSetImpl<Value *> &);
   void processBranch(BranchInst *, BasicBlock *, SmallPtrSetImpl<Value *> &);
+  void processSwitch(SwitchInst *, BasicBlock *, SmallPtrSetImpl<Value *> &);
   void renameUses(SmallPtrSetImpl<Value *> &);
   using ValueDFS = PredicateInfoClasses::ValueDFS;
   typedef SmallVectorImpl<ValueDFS> ValueDFSStack;
   void convertUsesToDFSOrdered(Value *, SmallVectorImpl<ValueDFS> &);
   Value *materializeStack(unsigned int &, ValueDFSStack &, Value *);
   bool stackIsInScope(const ValueDFSStack &, const ValueDFS &) const;
   void popStackUntilDFSScope(ValueDFSStack &, const ValueDFS &);
   ValueInfo &getOrCreateValueInfo(Value *);
private:
   // This gives the index into the ValueInfos array for a given Value.  Because
   // 0 is not a valid Value Info index, you can use DenseMap::lookup and tell
   // whether it returned a valid result.
   DenseMap<Value *, unsigned int> ValueInfoNums;
   // OrderedBasicBlocks used during sorting uses
   DenseMap<const BasicBlock *, std::unique_ptr<OrderedBasicBlock>> OBBMap;
   // The set of edges along which we can only handle phi uses, due to critical
   // edges.
-  DenseSet<BasicBlockEdge> PhiUsesOnly;
+  DenseSet<std::pair<BasicBlock *, BasicBlock *>> PhiUsesOnly;
 };
 
 // This pass does eager building and then printing of PredicateInfo. It is used
 // by
 // the tests to be able to build, dump, and verify PredicateInfo.
 class PredicateInfoPrinterLegacyPass : public FunctionPass {
 public:
   PredicateInfoPrinterLegacyPass();

lib/Transforms/Scalar/NewGVN.cpp

const Expression *E =
                            lookupMemoryAccessEquiv(DefiningAccess));
   return E;
 }
 
 const Expression *
 NewGVN::performSymbolicPredicateInfoEvaluation(Instruction *I) {
   if (auto *PI = PredInfo->getPredicateInfoFor(I)) {
     DEBUG(dbgs() << "Found predicate info from instruction !\n");
+    if (auto *PWC = dyn_cast<PredicateWithComparison>(PI)) {
       auto *CopyOf = I->getOperand(0);
-    auto *Cmp = PI->Comparison;
+      auto *Cmp = PWC->Comparison;
       // If this is an assume predicate and a copy of a comparison, it must be
       // true.
       if (isa<PredicateAssume>(PI) && CopyOf == Cmp)
         return createConstantExpression(ConstantInt::getTrue(Cmp->getType()));
 
       Value *FirstOp = lookupOperandLeader(Cmp->getOperand(0));
       Value *SecondOp = lookupOperandLeader(Cmp->getOperand(1));
       // Sort the ops
       CmpInst::Predicate Predicate = Cmp->getPredicate();
       // FIXME: We should really be ranking them here
       if (getRank(FirstOp) > getRank(SecondOp)) {
         std::swap(FirstOp, SecondOp);
         Predicate = CmpInst::getSwappedPredicate(Predicate);
       }
 
       if (isa<PredicateAssume>(PI)) {
-      // If the comparison is true when the operands are equal, then we know the
+        // If the comparison is true when the operands are equal, then we know
+        // the
         // operands are equal, because assumes must always be true.
         if (CmpInst::isTrueWhenEqual(Predicate))
           if (auto *C = dyn_cast<Constant>(FirstOp))
             return createConstantExpression(C);
         return createVariableExpression(FirstOp);
       } else if (const auto *PBranch = dyn_cast<PredicateBranch>(PI)) {
         // If this is a copy of the comparison, it's value is determined by the
         // edge.
         if (CopyOf == Cmp) {
           if (PBranch->TrueEdge)
-          return createConstantExpression(ConstantInt::getTrue(Cmp->getType()));
-        return createConstantExpression(ConstantInt::getFalse(Cmp->getType()));
+            return createConstantExpression(
+                ConstantInt::getTrue(Cmp->getType()));
+          return createConstantExpression(
+              ConstantInt::getFalse(Cmp->getType()));
         } else if ((PBranch->TrueEdge && CmpInst::isTrueWhenEqual(Predicate)) ||
-                 (!PBranch->TrueEdge && CmpInst::isFalseWhenEqual(Predicate))) {
+                   (!PBranch->TrueEdge &&
+                    CmpInst::isFalseWhenEqual(Predicate))) {
           // If we are *not* a copy of the comparison, we may equal to the other
           // operand when
           // the predicate implies something about equality of operations.
           if (auto *C = dyn_cast<Constant>(FirstOp))
             return createConstantExpression(C);
           return createVariableExpression(FirstOp);
         } else if (((PBranch->TrueEdge && Predicate == CmpInst::FCMP_OEQ) ||
                     (!PBranch->TrueEdge && Predicate == CmpInst::FCMP_UNE)) &&
                    isa<ConstantFP>(FirstOp) &&
                    !cast<ConstantFP>(FirstOp)->isZero()) {
           return createConstantExpression(cast<Constant>(FirstOp));
         }
       }
+    } else if (auto *PSwitch = dyn_cast<PredicateSwitch>(PI)) {
+      // For switch statements, we know the value of the condition, and that is
+      // what Op must be.
+      assert(I->getOperand(0) == PSwitch->Switch->getCondition() &&
+             "Found a predicateswitch info that seems to be about a non-switch "
+             "condition");
+      assert(isa<Constant>(PSwitch->CaseValue));
+      return createConstantExpression(cast<Constant>(PSwitch->CaseValue));
+    }
   }
   return nullptr;
 }
 
 // Evaluate read only and pure calls, and create an expression result.
 const Expression *NewGVN::performSymbolicCallEvaluation(Instruction *I) {
   auto *CI = cast<CallInst>(I);
   if (auto *II = dyn_cast<IntrinsicInst>(I)) {
if (II && EI->getNumIndices() == 1 && *EI->idx_begin() == 0) {
         return createBinaryExpression(
             Opcode, EI->getType(), II->getArgOperand(0), II->getArgOperand(1));
       }
     }
   }
 
   return createAggregateValueExpression(I);
 }
 const Expression *NewGVN::performSymbolicCmpEvaluation(Instruction *I) {
   auto *CI = dyn_cast<CmpInst>(I);
   // See if our operands are equal to those of a previous predicate, and if so,
   // if it implies true or false.
   auto Op0 = lookupOperandLeader(CI->getOperand(0));
   auto Op1 = lookupOperandLeader(CI->getOperand(1));
   // Avoid processing the same info twice
   const PredicateBase *LastPredInfo = nullptr;
 
   // See if we know something about the comparison itself, like it is the target
const Expression *NewGVN::performSymbolicCmpEvaluation(Instruction *I) {
   // See if our operands have predicate info, so that we may be able to derive
   // something from a previous comparison.
   for (const auto &Op : CI->operands()) {
     auto *PI = PredInfo->getPredicateInfoFor(Op);
     if (const auto *PBranch = dyn_cast_or_null<PredicateBranch>(PI)) {
       if (PI == LastPredInfo)
         continue;
       LastPredInfo = PI;
       // TODO: Along the false edge, we may know more things too, like icmp of
       // same operands is false.
       //
       auto *BranchCond = PBranch->Comparison;
       if (lookupOperandLeader(BranchCond->getOperand(0)) == Op0 &&
           lookupOperandLeader(BranchCond->getOperand(1)) == Op1) {
         if (PBranch->TrueEdge) {
           // If we know the previous predicate is true and we are in the true
           // edge then we may be implied true or false.

lib/Transforms/Utils/PredicateInfo.cpp

INITIALIZE_PASS_BEGIN(PredicateInfoPrinterLegacyPass, "print-predicateinfo",
                       "PredicateInfo Printer", false, false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_END(PredicateInfoPrinterLegacyPass, "print-predicateinfo",
                     "PredicateInfo Printer", false, false)
 static cl::opt<bool> VerifyPredicateInfo(
     "verify-predicateinfo", cl::init(false), cl::Hidden,
     cl::desc("Verify PredicateInfo in legacy printer pass."));
+namespace {
+
+// Given a predicate info that is a type of branching terminator, get the
+// branching block.
+const BasicBlock *getBranchBlock(PredicateBase *PB) {
+  if (auto *PBranch = dyn_cast<PredicateBranch>(PB))
+    return PBranch->BranchBB;
+  if (auto *PSwitch = dyn_cast<PredicateSwitch>(PB))
+    return PSwitch->SwitchBB;
+  llvm_unreachable("Only branches and switches should have PHIOnly defs that "
+                   "require branch blocks");
+}
+
+// Given a predicate info that is a type of branching terminator, get the
+// branching terminator.
+static Instruction *getBranchTerminator(PredicateBase *PB) {
+  if (auto *PBranch = dyn_cast<PredicateBranch>(PB))
+    return PBranch->BranchBB->getTerminator();
+  if (auto *PSwitch = dyn_cast<PredicateSwitch>(PB))
+    return PSwitch->SwitchBB->getTerminator();
+  llvm_unreachable("Not a predicate info type we know how to handle");
+}
+
+// Given a predicate info that is a type of branching terminator, get the
+// edge this predicate info represents
+const std::pair<BasicBlock *, BasicBlock *>
+getBlockEdge(const PredicateBase *PB) {
+  if (auto *PBranch = dyn_cast<PredicateBranch>(PB))
+    return std::make_pair(PBranch->BranchBB, PBranch->SplitBB);
+  if (auto *PSwitch = dyn_cast<PredicateSwitch>(PB))
+    return std::make_pair(PSwitch->SwitchBB, PSwitch->TargetBB);
+  llvm_unreachable("Unhandle predicate info type");
+}
+}
+
 namespace llvm {
 namespace PredicateInfoClasses {
 enum LocalNum {
   // Operations that must appear first in the block.
   LN_First,
   // Operations that are somewhere in the middle of the block, and are sorted on
   // demand.
   LN_Middle,
bool operator()(const ValueDFS &A, const ValueDFS &B) const {
 
     if (!SameBlock || A.LocalNum != LN_Middle || B.LocalNum != LN_Middle)
       return std::tie(A.DFSIn, A.DFSOut, A.LocalNum, A.Def, A.Use) <
              std::tie(B.DFSIn, B.DFSOut, B.LocalNum, B.Def, B.Use);
     return localComesBefore(A, B);
   }
 
   // For a phi use, or a non-materialized def, return the edge it represents.
-  const std::pair<const BasicBlock *, const BasicBlock *>
+  const std::pair<BasicBlock *, BasicBlock *>
   getBlockEdge(const ValueDFS &VD) const {
     if (!VD.Def && VD.Use) {
       auto *PHI = cast<PHINode>(VD.Use->getUser());
       return std::make_pair(PHI->getIncomingBlock(*VD.Use), PHI->getParent());
     }
     // This is really a non-materialized def.
-    auto *PBranch = cast<PredicateBranch>(VD.PInfo);
-    return std::make_pair(PBranch->BranchBB, PBranch->SplitBB);
+    return ::getBlockEdge(VD.PInfo);
   }
 
   // For two phi related values, return the ordering.
   bool comparePHIRelated(const ValueDFS &A, const ValueDFS &B) const {
     auto &ABlockEdge = getBlockEdge(A);
     auto &BBlockEdge = getBlockEdge(B);
     if (ABlockEdge < BBlockEdge)
       return true;
bool PredicateInfo::stackIsInScope(const ValueDFSStack &Stack,
   // the defs they must go with so that we can know it's time to pop the stack
   // when we hit the end of the phi uses for a given def.
   if (Stack.back().PhiOnly) {
     if (!VDUse.Use)
       return false;
     auto *PHI = dyn_cast<PHINode>(VDUse.Use->getUser());
     if (!PHI)
       return false;
-    // The only phionly defs should be branch info.
-    auto *PBranch = dyn_cast<PredicateBranch>(Stack.back().PInfo);
-    assert(PBranch && "Only branches should have PHIOnly defs");
-    // Check edge
-    BasicBlock *EdgePred = PHI->getIncomingBlock(*VDUse.Use);
-    if (EdgePred != PBranch->BranchBB)
+    // The only phionly defs should be branch or switch info.
+    // Check that this is our edge.
+    const BasicBlock *EdgePred = PHI->getIncomingBlock(*VDUse.Use);
+    if (EdgePred != getBranchBlock(Stack.back().PInfo))
       return false;
   }
 
   return (VDUse.DFSIn >= Stack.back().DFSIn &&
           VDUse.DFSOut <= Stack.back().DFSOut);
 }
 
 void PredicateInfo::popStackUntilDFSScope(ValueDFSStack &Stack,
if (auto *Cmp = dyn_cast<CmpInst>(Comparison)) {
           if (!Succ->getSinglePredecessor())
             PhiUsesOnly.insert({BranchBB, Succ});
         }
       }
       CmpOperands.clear();
     }
   }
 }
+// Process a block terminating switch, and place relevant operations to be
+// renamed into OpsToRename.
+void PredicateInfo::processSwitch(SwitchInst *SI, BasicBlock *BranchBB,
+                                  SmallPtrSetImpl<Value *> &OpsToRename) {
+  Value *Op = SI->getCondition();
+  if ((!isa<Instruction>(Op) && !isa<Argument>(Op)) || Op->hasOneUse())
+    return;
+
+  // Remember how many outgoing edges there are to every successor.
+  SmallDenseMap<BasicBlock *, unsigned, 16> SwitchEdges;
+  for (unsigned i = 0, e = SI->getNumSuccessors(); i != e; ++i) {
+    BasicBlock *TargetBlock = SI->getSuccessor(i);
+    ++SwitchEdges[TargetBlock];
+  }
+
+  // Now propagate info for each case value
+  for (auto i = SI->case_begin(), e = SI->case_end(); i != e; ++i) {
+    BasicBlock *TargetBlock = i.getCaseSuccessor();
+    if (SwitchEdges.lookup(TargetBlock) == 1) {
+      OpsToRename.insert(Op);
+      auto &OperandInfo = getOrCreateValueInfo(Op);
+      PredicateSwitch *PS = new PredicateSwitch(
+          Op, SI->getParent(), TargetBlock, i.getCaseValue(), SI);
+      AllInfos.push_back(PS);
+      OperandInfo.Infos.push_back(PS);
+      if (!TargetBlock->getSinglePredecessor())
+        PhiUsesOnly.insert({BranchBB, TargetBlock});
+    }
+  }
+}
 
 // Build predicate info for our function
 void PredicateInfo::buildPredicateInfo() {
   DT.updateDFSNumbers();
   // Collect operands to rename from all conditional branch terminators, as well
   // as assume statements.
   SmallPtrSet<Value *, 8> OpsToRename;
   for (auto DTN : depth_first(DT.getRootNode())) {
     BasicBlock *BranchBB = DTN->getBlock();
     if (auto *BI = dyn_cast<BranchInst>(BranchBB->getTerminator())) {
       if (!BI->isConditional())
         continue;
       processBranch(BI, BranchBB, OpsToRename);
+    } else if (auto *SI = dyn_cast<SwitchInst>(BranchBB->getTerminator())) {
+      processSwitch(SI, BranchBB, OpsToRename);
     }
   }
   for (auto &Assume : AC.assumptions()) {
     if (auto *II = dyn_cast_or_null<IntrinsicInst>(Assume))
       processAssume(II, II->getParent(), OpsToRename);
   }
   // Now rename all our operations.
   renameUses(OpsToRename);
 }
+
+// Given the renaming stack, make all the operands currently on the stack real
+// by inserting them into the IR.  Return the last operation's value.
 Value *PredicateInfo::materializeStack(unsigned int &Counter,
                                        ValueDFSStack &RenameStack,
                                        Value *OrigOp) {
   // Find the first thing we have to materialize
   auto RevIter = RenameStack.rbegin();
   for (; RevIter != RenameStack.rend(); ++RevIter)
     if (RevIter->Def)
       break;
 
   size_t Start = RevIter - RenameStack.rbegin();
   // The maximum number of things we should be trying to materialize at once
   // right now is 4, depending on if we had an assume, a branch, and both used
   // and of conditions.
   for (auto RenameIter = RenameStack.end() - Start;
        RenameIter != RenameStack.end(); ++RenameIter) {
     auto *Op =
         RenameIter == RenameStack.begin() ? OrigOp : (RenameIter - 1)->Def;
     ValueDFS &Result = *RenameIter;
     auto *ValInfo = Result.PInfo;
     // For branches, we can just place the operand in the branch block before
     // the terminator.  For assume, we have to place it right before the assume
     // to ensure we dominate all of our uses.  Always insert right before the
     // relevant instruction (terminator, assume), so that we insert in proper
     // order in the case of multiple predicateinfo in the same block.
-    if (isa<PredicateBranch>(ValInfo)) {
-      auto *PBranch = cast<PredicateBranch>(ValInfo);
-      IRBuilder<> B(PBranch->BranchBB->getTerminator());
+    if (isa<PredicateBranch>(ValInfo) || isa<PredicateSwitch>(ValInfo)) {
+      IRBuilder<> B(getBranchTerminator(ValInfo));
       Function *IF = Intrinsic::getDeclaration(
           F.getParent(), Intrinsic::ssa_copy, Op->getType());
       Value *PIC = B.CreateCall(IF, Op, Op->getName() + "." + Twine(Counter++));
       PredicateMap.insert({PIC, ValInfo});
       Result.Def = PIC;
     } else {
       auto *PAssume = dyn_cast<PredicateAssume>(ValInfo);
       assert(PAssume &&
for (auto &PossibleCopy : ValueInfo.Infos) {
         VD.LocalNum = LN_Middle;
         DomTreeNode *DomNode = DT.getNode(PAssume->AssumeInst->getParent());
         if (!DomNode)
           continue;
         VD.DFSIn = DomNode->getDFSNumIn();
         VD.DFSOut = DomNode->getDFSNumOut();
         VD.PInfo = PossibleCopy;
         OrderedUses.push_back(VD);
-      } else if (const auto *PBranch =
-                     dyn_cast<PredicateBranch>(PossibleCopy)) {
+      } else if (isa<PredicateBranch>(PossibleCopy) ||
+                 isa<PredicateSwitch>(PossibleCopy)) {
         // If we can only do phi uses, we treat it like it's in the branch
         // block, and handle it specially. We know that it goes last, and only
         // dominate phi uses.
-        if (PhiUsesOnly.count({PBranch->BranchBB, PBranch->SplitBB})) {
+        auto BlockEdge = getBlockEdge(PossibleCopy);
+        if (PhiUsesOnly.count(BlockEdge)) {
           VD.LocalNum = LN_Last;
-          auto *DomNode = DT.getNode(PBranch->BranchBB);
+          auto *DomNode = DT.getNode(BlockEdge.first);
           if (DomNode) {
             VD.DFSIn = DomNode->getDFSNumIn();
             VD.DFSOut = DomNode->getDFSNumOut();
             VD.PInfo = PossibleCopy;
             VD.PhiOnly = true;
             OrderedUses.push_back(VD);
           }
         } else {
           // Otherwise, we are in the split block (even though we perform
           // insertion in the branch block).
           // Insert a possible copy at the split block and before the branch.
           VD.LocalNum = LN_First;
-          auto *DomNode = DT.getNode(PBranch->SplitBB);
+          auto *DomNode = DT.getNode(BlockEdge.second);
           if (DomNode) {
             VD.DFSIn = DomNode->getDFSNumIn();
             VD.DFSOut = DomNode->getDFSNumOut();
             VD.PInfo = PossibleCopy;
             OrderedUses.push_back(VD);
           }
         }
       }