[SimplifyCFG] Speculatively flatten CFG based on profiling metadata

[reames]

Summary

If we have a series of branches which are all unlikely to fail, we can possibly combine them into a single check on the fastpath combined with a bit of dispatch logic on the slowpath. We don't want to do this unconditionally since it requires speculating instructions past a branch, but if the profiling metadata on the branch indicates profitability, this can reduce the number of checks needed along the fast path.

The canonical example this is trying to handle is removing the second bounds check implied by the Java code: a[i] + a[i+1]. Note that it can currently only do so for really simple conditions and the values of a[i] can't be used anywhere except in the addition. (i.e. the load has to have been sunk already and not prevent speculation.) I plan on extending this transform over the next few days to handle alternate sequences.

Note that this builds on the "implies" function introduced in http://reviews.llvm.org/D13040. Please ignore its definition here and direct comments to that review. Based on Sanjoy's feedback, I'm going to introduce this logic into InstSimplify and refactor both patches based on that.

[reames]

Update based on http://reviews.llvm.org/D13074 (i.e. pulling the implication logic into InstSimplify for easy testing)

[hfinkel]

lib/Transforms/Utils/SimplifyCFG.cpp

2364

either check -> any of the checks

2393

I'd rather that this 100 ratio be a command-line parameter for easier experimentation.

2414

10 should be a command-line parameter too.

[reames]

lib/Transforms/Utils/SimplifyCFG.cpp

2364

Happy to make the requested change, but given this applies to a pair of branches, not quite clear why you think the new wording is more clear.

2393

Will do.

2414

Will do.

[hfinkel]

lib/Transforms/Utils/SimplifyCFG.cpp

2364

The comment starts with, "If we have a series of tests leading to a frequently executed fallthrough... we can test all the conditions at once...". You're right, but maybe we should just add something like, "Here, we deal with this one pair of tests at a time." So that it becomes clear that you're now talking about only two checks.

[sanjoy]

lib/Transforms/Utils/SimplifyCFG.cpp

2382

I'd drop the TryTo. LLVM usually names such functions directly (InlineFunction, UnrollLoop etc.).

2405

Why is this needed? Since you're changing PBI's condition to be BI's condition, and BI's condition implies PBI's condition, I don't see what gets speculatively executed.

2507

auto * is clearer, since the type is obvious from the dyn_cast.

[reames]

lib/Transforms/Utils/SimplifyCFG.cpp

2405

The problem is that instructions from BI's block are speculated above PBI. If you had some code like:
if (i > length) return;
%p = load a[i];
if (i + 1 > length) return;

After the transform, we'd end up with:
%p = load a[i];
if (i + 1 > length) {if (i > length) return; else return;}

We can't necessarily speculate the load above the first condition.

2507

Agreed, but for the record, I just moved this line.

[reames]

Address review comments by Hal and Sanjoy.

[sanjoy]

Minor nits inline. With those addressed, this LGTM (I'd still wait for @hfinkel to take a look).

lib/Transforms/Utils/SimplifyCFG.cpp

2354

Use dyn_cast_or_null to directly cast to Constant.

2385

Nit: spacing is a little off here.

2405

SGTM. :)

2412

Should CtxI be nullptr here? We're speculating it above PBI, so I think we cannot assume control dependence on PBI.

2417

Minor nit: InvokeInst should have been caught with isa<TerminatorInst>. In fact, we can pretty much assert(!isa<InvokeInst>(I)) since the only terminator allowed in BB is BI.

2424

[This is optional to fix, since the actual rewriting is fairly simple]

I'd put a small "before" and "after" IR snippet here, using the C++ identifiers in IR, like:

// Before:
//
//  br %whichCond, label %bb, label %fail1  ;; == PBB
//
// bb:
//  br %cond, label %success, label %fail2  ;; == BB
//

etc.

I think that will make the graph rewriting clearer.

2427

I'd call these FailBI and FailPBI.

[reames]

Respond to Sanjoy's review comments. Given he gave a contingent LGTM, I'm going to let this sit for a day in case anyone has final comments, then submit.

lib/Transforms/Utils/SimplifyCFG.cpp

 static cl::opt<bool>
 SinkCommon("simplifycfg-sink-common", cl::Hidden, cl::init(true),
        cl::desc("Sink common instructions down to the end block"));
 
 static cl::opt<bool> HoistCondStores(
     "simplifycfg-hoist-cond-stores", cl::Hidden, cl::init(true),
     cl::desc("Hoist conditional stores if an unconditional store precedes"));
 
+static cl::opt<unsigned> SpeculativeFlattenBias(
+    "speculative-flatten-bias", cl::Hidden, cl::init(100),
+    cl::desc("Control how biased a branch needs to be to be considered rarely"
+             " failing for speculative flattening (default = 100)"));
+
+static cl::opt<unsigned> SpeculativeFlattenThreshold(
+    "speculative-flatten-threshold", cl::Hidden, cl::init(10),
+    cl::desc("Control how much speculation happens due to speculative"
+             " flattening (default = 10)"));
+
+
 STATISTIC(NumBitMaps, "Number of switch instructions turned into bitmaps");
 STATISTIC(NumLinearMaps, "Number of switch instructions turned into linear mapping");
 STATISTIC(NumLookupTables, "Number of switch instructions turned into lookup tables");
 STATISTIC(NumLookupTablesHoles, "Number of switch instructions turned into lookup tables (holes checked)");
 STATISTIC(NumTableCmpReuses, "Number of reused switch table lookup compares");
 STATISTIC(NumSinkCommons, "Number of common instructions sunk down to the end block");
 STATISTIC(NumSpeculations, "Number of speculative executed instructions");
 
for (BasicBlock::iterator I = BB->begin(), E = BB->end(); I != E; ++I)
       if (isa<DbgInfoIntrinsic>(*I))
         I->clone()->insertBefore(PBI);
 
     return true;
   }
   return false;
 }
 
+
+/// Return true if B is known to be implied by A.  A & B must be i1 (boolean)
+static bool implies(Value *A, Value *B, const DataLayout &DL) {
+  assert(A->getType()->isIntegerTy(1) && B->getType()->isIntegerTy(1));
+  // Note that the truth table for implication is the same as <=u on i1
+  // values. 
+  Value *Simplified = SimplifyICmpInst(ICmpInst::ICMP_ULE, A, B, DL);
+  if (!Simplified) return false;
+  Constant *Con = dyn_cast<Constant>(Simplified);
+  return Con && Con->isOneValue();
+}
+
+/// If we have a series of tests leading to a frequently executed fallthrough
+/// path and we can test all the conditions at once, we can execute a single
+/// test on the fast path and figure out which condition failed on the slow
+/// path.  This transformation considers a pair of branches at a time since
+/// recursively considering branches pairwise will cause an entire chain to
+/// collapse.  This transformation is code size neutral, but makes it
+/// dynamically more expensive to fail either check.  As such, we only want to
+/// do this if both checks are expected to essentially never fail.
+/// The key motivating examples are cases like:
+///   br (i < Length), in_bounds, fail1
+/// in_bounds:
+///   br (i+1 < Length), in_bounds2, fail2
+/// in_bounds2:
+///   ...
+///
+/// We can rewrite this as:
+///   br (i+1 < Length), in_bounds2, dispatch
+/// in_bounds2:
+///   ...
+/// dispatch:
+///   br (i < Length), fail2, fail1
+///
+/// TODO: we could consider duplicating some (non-speculatable) instructions
+/// from BI->getParent() down both paths
+static bool SpeculativelyFlattenCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
+                                                            const DataLayout &DL) {
+  auto *PredBB = PBI->getParent();
+  auto *BB = BI->getParent();
+
+  /// Is the failing path of this branch taken rarely if at all?
+  auto isRarelyUntaken = [](BranchInst *BI) {
+    uint64_t ProbTrue;
+    uint64_t ProbFalse;
+    if (!ExtractBranchMetadata(BI, ProbTrue, ProbFalse))
+      return false;
+    return ProbTrue > ProbFalse * SpeculativeFlattenBias;
+  };
+
+  if (PBI->getSuccessor(0) != BB ||
+      !isRarelyUntaken(PBI) || !isRarelyUntaken(BI) ||
+      !implies(BI->getCondition(), PBI->getCondition(), DL))
+    return false;
+
+  // TODO: The following code performs a similiar, but slightly distinct
+  // transformation to that done by SpeculativelyExecuteBB.  We should consider
+  // combining them at some point.
+
+  // Can we speculate everything in the given block (except for the terminator
+  // instruction) at the instruction boundary before 'At'?
+  unsigned SpeculationCost = 0;
+  for (Instruction &I : *BB) {
+    if (isa<TerminatorInst>(I)) break;
+    if (!isSafeToSpeculativelyExecute(&I, PBI))
+      return false;
+    SpeculationCost++;
+    // Only flatten relatively small BBs to avoid making the bad case terrible
+    if (SpeculationCost > SpeculativeFlattenThreshold ||
+        isa<CallInst>(I) || isa<InvokeInst>(I))
+      return false;
+  }
+
+  DEBUG(dbgs() << "Outlining slow path comparison: "
+        << *PBI->getCondition() << " implied by "
+        << *BI->getCondition() << "\n");
+  Value *WhichCond = PBI->getCondition();
+  auto *Success = BI->getSuccessor(0);
+  auto *Fail1 = PBI->getSuccessor(1);
+  auto *Fail2 =  BI->getSuccessor(1);
+  // Have PBI branch directly to the fast path using BI's condition, branch
+  // to this BI's block for the slow path dispatch
+  PBI->setSuccessor(0, Success);
+  PBI->setSuccessor(1, BB);
+  PBI->setCondition(BI->getCondition());
+  // Rewrite BI to distinguish between the two failing cases
+  BI->setSuccessor(0, Fail2);
+  BI->setSuccessor(1, Fail1);
+  BI->setCondition(WhichCond);
+  // Move all of the instructions from BI->getParent other than
+  // the terminator into the fastpath.  This requires speculating them past
+  // the original PBI branch, but we checked for that legality above.
+  // TODO: This doesn't handle dependent loads.  We could duplicate those
+  // down both paths, but that involves further code growth.  We need to
+  // figure out a good cost model here.
+  PredBB->getInstList().splice(PBI, BB->getInstList(),
+                               BB->begin(), std::prev(BB->end()));
+
+  // To be conservatively correct, drop all metadata on the rewritten
+  // branches.  TODO: update metadata
+  PBI->dropUnknownNonDebugMetadata();
+  BI->dropUnknownNonDebugMetadata();
+  return true;
+}
 /// If we have a conditional branch as a predecessor of another block,
 /// this function tries to simplify it.  We know
 /// that PBI and BI are both conditional branches, and BI is in one of the
 /// successor blocks of PBI - PBI branches to BI.
-static bool SimplifyCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI) {
+static bool SimplifyCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
+                                           const DataLayout &DL) {
   assert(PBI->isConditional() && BI->isConditional());
   BasicBlock *BB = BI->getParent();
 
   // If this block ends with a branch instruction, and if there is a
   // predecessor that ends on a branch of the same condition, make
   // this conditional branch redundant.
   if (PBI->getCondition() == BI->getCondition() &&
       PBI->getSuccessor(0) != PBI->getSuccessor(1)) {
if (BlockIsSimpleEnoughToThreadThrough(BB)) {
         }
       }
 
       BI->setCondition(NewPN);
       return true;
     }
   }
 
+  if (auto *CE = dyn_cast<ConstantExpr>(BI->getCondition()))
+    if (CE->canTrap())
+      return false;
+
+  if (SpeculativelyFlattenCondBranchToCondBranch(PBI, BI, DL))
+    return true;
+
   // If this is a conditional branch in an empty block, and if any
   // predecessors are a conditional branch to one of our destinations,
   // fold the conditions into logical ops and one cond br.
   BasicBlock::iterator BBI = BB->begin();
   // Ignore dbg intrinsics.
   while (isa<DbgInfoIntrinsic>(BBI))
     ++BBI;
   if (&*BBI != BI)
     return false;
 
-
-  if (ConstantExpr *CE = dyn_cast<ConstantExpr>(BI->getCondition()))
-    if (CE->canTrap())
-      return false;
-
   int PBIOp, BIOp;
   if (PBI->getSuccessor(0) == BI->getSuccessor(0))
     PBIOp = BIOp = 0;
   else if (PBI->getSuccessor(0) == BI->getSuccessor(1))
     PBIOp = 0, BIOp = 1;
   else if (PBI->getSuccessor(1) == BI->getSuccessor(0))
     PBIOp = 1, BIOp = 0;
   else if (PBI->getSuccessor(1) == BI->getSuccessor(1))
if (PHINode *PN = dyn_cast<PHINode>(BI->getCondition()))
     if (PN->getParent() == BI->getParent())
       if (FoldCondBranchOnPHI(BI, DL))
         return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
 
   // Scan predecessor blocks for conditional branches.
   for (pred_iterator PI = pred_begin(BB), E = pred_end(BB); PI != E; ++PI)
     if (BranchInst *PBI = dyn_cast<BranchInst>((*PI)->getTerminator()))
       if (PBI != BI && PBI->isConditional())
-        if (SimplifyCondBranchToCondBranch(PBI, BI))
+        if (SimplifyCondBranchToCondBranch(PBI, BI, DL))
           return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
 
   return false;
 }
 
 /// Check if passing a value to an instruction will cause undefined behavior.
 static bool passingValueIsAlwaysUndefined(Value *V, Instruction *I) {
   Constant *C = dyn_cast<Constant>(V);

test/Transforms/SimplifyCFG/fast-fallthrough.ll

+; RUN: opt -S %s -simplifycfg | FileCheck %s
+
+define void @test(i32 %length.i, i32 %i) {
+; CHECK-LABEL: @test
+  %iplus1 = add nsw i32 %i, 1
+  %var29 = icmp slt i32 %i, %length.i
+  %var30 = icmp slt i32 %iplus1, %length.i
+; CHECK: br i1 %var30, label %in_bounds, label %next
+  br i1 %var29, label %next, label %out_of_bounds, !prof !{!"branch_weights", i32 1000, i32 0}
+
+next:
+; CHECK-LABEL: next:
+; CHECK: br i1 %var29, label %out_of_bounds2, label %out_of_bounds
+  br i1 %var30, label %in_bounds, label %out_of_bounds2, !prof !{!"branch_weights", i32 1000, i32 0}
+
+in_bounds:
+  ret void
+
+out_of_bounds:
+  call void @foo(i64 0)
+  unreachable
+
+out_of_bounds2:
+  call void @foo(i64 1)
+  unreachable
+}
+
+define void @test2(i32 %length.i, i32 %i) {
+; CHECK-LABEL: @test2
+  %var29 = icmp slt i32 %i, %length.i
+; CHECK: br i1 %var30, label %in_bounds, label %next
+  br i1 %var29, label %next, label %out_of_bounds, !prof !{!"branch_weights", i32 1000, i32 0}
+
+next:
+; CHECK-LABEL: next:
+; CHECK: br i1 %var29, label %out_of_bounds2, label %out_of_bounds
+  %iplus1 = add nsw i32 %i, 1
+  %var30 = icmp slt i32 %iplus1, %length.i
+  br i1 %var30, label %in_bounds, label %out_of_bounds2, !prof !{!"branch_weights", i32 1000, i32 0}
+
+in_bounds:
+  ret void
+
+out_of_bounds:
+  call void @foo(i64 0)
+  unreachable
+
+out_of_bounds2:
+  call void @foo(i64 1)
+  unreachable
+}
+
+; As written, this one can't trigger today.  It would require us to duplicate
+; the %val1 load down two paths and that's not implemented yet.
+define i64 @test3(i32 %length.i, i32 %i, i64* %base) {
+; CHECK-LABEL: @test3
+  %var29 = icmp slt i32 %i, %length.i
+; CHECK: br i1 %var29, label %next, label %out_of_bounds
+  br i1 %var29, label %next, label %out_of_bounds, !prof !{!"branch_weights", i32 1000, i32 0}
+
+next:
+; CHECK-LABEL: next:
+  %addr1 = getelementptr i64, i64* %base, i32 %i
+  %val1 = load i64, i64* %addr1
+  %iplus1 = add nsw i32 %i, 1
+  %var30 = icmp slt i32 %iplus1, %length.i
+; CHECK: br i1 %var30, label %in_bounds, label %out_of_bounds2
+  br i1 %var30, label %in_bounds, label %out_of_bounds2, !prof !{!"branch_weights", i32 1000, i32 0}
+
+in_bounds:
+  %addr2 = getelementptr i64, i64* %base, i32 %iplus1
+  %val2 = load i64, i64* %addr2
+  %res = sub i64 %val1, %val2
+  ret i64 %res
+
+out_of_bounds:
+  call void @foo(i64 0)
+  unreachable
+
+out_of_bounds2:
+  call void @foo(i64 %val1)
+  unreachable
+}
+
+declare void @foo(i64)
+