[SCCP] Resolve indirect branch target when possible.

Summary:
Resolve indirect branch target when possible.
This potentially eliminates more basicblocks and result in better evaluation for phi and other things.

Reviewers: davide, efriedma, sanjoy

Subscribers: llvm-commits

Differential Revision: https://reviews.llvm.org/D30322

llvm-svn: 299830

2 files changed, 147 insertions, 8 deletions

diff --git a/llvm/lib/Transforms/Scalar/SCCP.cpp b/llvm/lib/Transforms/Scalar/SCCP.cpp
index 2448e79679b..f2d9f162892 100644
--- a/llvm/lib/Transforms/Scalar/SCCP.cpp
+++ b/llvm/lib/Transforms/Scalar/SCCP.cpp
@@ -140,6 +140,14 @@ public:
     return nullptr;
   }
 
+  /// getBlockAddress - If this is a constant with a BlockAddress value, return
+  /// it, otherwise return null.
+  BlockAddress *getBlockAddress() const {
+    if (isConstant())
+      return dyn_cast<BlockAddress>(getConstant());
+    return nullptr;
+  }
+
   void markForcedConstant(Constant *V) {
     assert(isUnknown() && "Can't force a defined value!");
     Val.setInt(forcedconstant);
@@ -600,10 +608,32 @@ void SCCPSolver::getFeasibleSuccessors(TerminatorInst &TI,
     return;
   }
 
-  // TODO: This could be improved if the operand is a [cast of a] BlockAddress.
-  if (isa<IndirectBrInst>(&TI)) {
-    // Just mark all destinations executable!
-    Succs.assign(TI.getNumSuccessors(), true);
+  // In case of indirect branch and its address is a blockaddress, we mark
+  // the target as executable.
+  if (auto *IBR = dyn_cast<IndirectBrInst>(&TI)) {
+    // Casts are folded by visitCastInst.
+    LatticeVal IBRValue = getValueState(IBR->getAddress());
+    BlockAddress *Addr = IBRValue.getBlockAddress();
+    if (!Addr) {   // Overdefined or unknown condition?
+      // All destinations are executable!
+      if (!IBRValue.isUnknown())
+        Succs.assign(TI.getNumSuccessors(), true);
+      return;
+    }
+
+    BasicBlock* T = Addr->getBasicBlock();
+    assert(Addr->getFunction() == T->getParent() &&
+           "Block address of a different function ?");
+    for (unsigned i = 0; i < IBR->getNumSuccessors(); ++i) {
+      // This is the target.
+      if (IBR->getDestination(i) == T) {
+        Succs[i] = true;
+        return;
+      }
+    }
+
+    // If we didn't find our destination in the IBR successor list, then we
+    // have undefined behavior. Its ok to assume no successor is executable.
     return;
   }
 
@@ -656,10 +686,18 @@ bool SCCPSolver::isEdgeFeasible(BasicBlock *From, BasicBlock *To) {
     return SI->findCaseValue(CI).getCaseSuccessor() == To;
   }
 
-  // Just mark all destinations executable!
-  // TODO: This could be improved if the operand is a [cast of a] BlockAddress.
-  if (isa<IndirectBrInst>(TI))
-    return true;
+  // In case of indirect branch and its address is a blockaddress, we mark
+  // the target as executable.
+  if (auto *IBR = dyn_cast<IndirectBrInst>(TI)) {
+    LatticeVal IBRValue = getValueState(IBR->getAddress());
+    BlockAddress *Addr = IBRValue.getBlockAddress();
+
+    if (!Addr)
+      return !IBRValue.isUnknown();
+
+    // At this point, the indirectbr is branching on a blockaddress.
+    return Addr->getBasicBlock() == To;
+  }
 
   DEBUG(dbgs() << "Unknown terminator instruction: " << *TI << '\n');
   llvm_unreachable("SCCP: Don't know how to handle this terminator!");
@@ -1484,6 +1522,31 @@ bool SCCPSolver::ResolvedUndefsIn(Function &F) {
       return true;
     }
 
+   if (auto *IBR = dyn_cast<IndirectBrInst>(TI)) {
+      // Indirect branch with no successor ?. Its ok to assume it branches
+      // to no target.
+      if (IBR->getNumSuccessors() < 1)
+        continue;
+
+      if (!getValueState(IBR->getAddress()).isUnknown())
+        continue;
+
+      // If the input to SCCP is actually branch on undef, fix the undef to
+      // the first successor of the indirect branch.
+      if (isa<UndefValue>(IBR->getAddress())) {
+        IBR->setAddress(BlockAddress::get(IBR->getSuccessor(0)));
+        markEdgeExecutable(&BB, IBR->getSuccessor(0));
+        return true;
+      }
+
+      // Otherwise, it is a branch on a symbolic value which is currently
+      // considered to be undef.  Handle this by forcing the input value to the
+      // branch to the first successor.
+      markForcedConstant(IBR->getAddress(),
+                         BlockAddress::get(IBR->getSuccessor(0)));
+      return true;
+    }
+
     if (auto *SI = dyn_cast<SwitchInst>(TI)) {
       if (!SI->getNumCases() || !getValueState(SI->getCondition()).isUnknown())
         continue;
diff --git a/llvm/test/Transforms/SCCP/indirectbr.ll b/llvm/test/Transforms/SCCP/indirectbr.ll
new file mode 100644
index 00000000000..b977961ca49
--- /dev/null
+++ b/llvm/test/Transforms/SCCP/indirectbr.ll
@@ -0,0 +1,76 @@
+; RUN: opt -S -sccp < %s | FileCheck %s
+
+declare void @BB0_f()
+declare void @BB1_f()
+
+; Make sure we can eliminate what is in BB0 as we know that the indirectbr is going to BB1.
+;
+; CHECK-LABEL: define void @indbrtest1(
+; CHECK-NOT: call void @BB0_f()
+; CHECK: ret void
+define void @indbrtest1() {
+entry:
+  indirectbr i8* blockaddress(@indbrtest1, %BB1), [label %BB0, label %BB1]
+BB0:
+  call void @BB0_f()
+  br label %BB1
+BB1:
+  call void @BB1_f()
+  ret void
+}
+
+; Make sure we can eliminate what is in BB0 as we know that the indirectbr is going to BB1
+; by looking through the casts. The casts should be folded away when they are visited
+; before the indirectbr instruction.
+;
+; CHECK-LABEL: define void @indbrtest2(
+; CHECK-NOT: call void @BB0_f()
+; CHECK: ret void
+define void @indbrtest2() {
+entry:
+  %a = ptrtoint i8* blockaddress(@indbrtest2, %BB1) to i64
+  %b = inttoptr i64 %a to i8*
+  %c = bitcast i8* %b to i8*
+  indirectbr i8* %b, [label %BB0, label %BB1]
+BB0:
+  call void @BB0_f()
+  br label %BB1
+BB1:
+  call void @BB1_f()
+  ret void
+}
+
+; Make sure we can not eliminate BB0 as we do not know the target of the indirectbr.
+;
+; CHECK-LABEL: define void @indbrtest3(
+; CHECK: call void @BB0_f()
+; CHECK: ret void
+define void @indbrtest3(i8** %Q) {
+entry:
+  %t = load i8*, i8** %Q
+  indirectbr i8* %t, [label %BB0, label %BB1]
+BB0:
+  call void @BB0_f()
+  br label %BB1
+BB1:
+  call void @BB1_f()
+  ret void
+}
+
+; Make sure we eliminate BB1 as we pick the first successor on undef.
+;
+; CHECK-LABEL: define void @indbrtest4(
+; CHECK: call void @BB0_f()
+; CHECK: ret void
+define void @indbrtest4(i8** %Q) {
+entry:
+  indirectbr i8* undef, [label %BB0, label %BB1]
+BB0:
+  call void @BB0_f()
+  br label %BB1
+BB1:
+  call void @BB1_f()
+  ret void
+}
+
+