Search for live blocks backward from uses. Add limit for searching

Signed-off-by: John Lu <[email protected]>

Author: LU-JOHN

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -198,6 +198,11 @@ static cl::opt<unsigned> MaxSwitchCasesPerResult(
     "max-switch-cases-per-result", cl::Hidden, cl::init(16),
     cl::desc("Limit cases to analyze when converting a switch to select"));
 
+static cl::opt<unsigned> MaxJumpThreadingLiveBlocks(
+    "max-jump-threading-live-blocks", cl::Hidden, cl::init(24),
+    cl::desc("Limit number of blocks a define in a threaded block is allowed "
+             "to be live in"));
+
 STATISTIC(NumBitMaps, "Number of switch instructions turned into bitmaps");
 STATISTIC(NumLinearMaps,
           "Number of switch instructions turned into linear mapping");
@@ -3457,24 +3462,25 @@ bool SimplifyCFGOpt::speculativelyExecuteBB(BranchInst *BI,
 
 using BlocksSet = SmallPtrSet<BasicBlock *, 8>;
 
-static bool reachesUsed(BasicBlock *BB, BlocksSet &ReachesNonLocalUses,
-                        BlocksSet &VisitedBlocksSet) {
-  if (ReachesNonLocalUses.contains(BB))
+// Return false if number of blocks searched is too much.
+static bool findReaching(BasicBlock *BB, BasicBlock *DefBB,
+                         BlocksSet &ReachesNonLocalUses) {
+  if (BB == DefBB)
+    return true;
+  if (!ReachesNonLocalUses.insert(BB).second)
     return true;
-  if (!VisitedBlocksSet.insert(BB).second)
+
+  if (ReachesNonLocalUses.size() > MaxJumpThreadingLiveBlocks)
     return false;
-  for (BasicBlock *Succ : successors(BB))
-    if (reachesUsed(Succ, ReachesNonLocalUses, VisitedBlocksSet)) {
-      ReachesNonLocalUses.insert(BB);
-      return true;
-    }
-  return false;
+  for (BasicBlock *Pred : predecessors(BB))
+    if (!findReaching(Pred, DefBB, ReachesNonLocalUses))
+      return false;
+  return true;
 }
 
 /// Return true if we can thread a branch across this block.
-static bool
-blockIsSimpleEnoughToThreadThrough(BasicBlock *BB,
-                                   BlocksSet &UsedInNonLocalBlocksSet) {
+static bool blockIsSimpleEnoughToThreadThrough(BasicBlock *BB,
+                                               BlocksSet &NonLocalUseBlocks) {
   int Size = 0;
   EphemeralValueTracker EphTracker;
 
@@ -3503,7 +3509,7 @@ blockIsSimpleEnoughToThreadThrough(BasicBlock *BB,
         if (isa<PHINode>(UI))
           return false;
       } else
-        UsedInNonLocalBlocksSet.insert(UsedInBB);
+        NonLocalUseBlocks.insert(UsedInBB);
     }
 
     // Looks ok, continue checking.
@@ -3565,10 +3571,19 @@ foldCondBranchOnValueKnownInPredecessorImpl(BranchInst *BI, DomTreeUpdater *DTU,
   // Check that the block is small enough and record which non-local blocks use
   // values defined in the block.
 
-  BlocksSet ReachesNonLocalUses;
-  if (!blockIsSimpleEnoughToThreadThrough(BB, ReachesNonLocalUses))
+  BlocksSet NonLocalUseBlocks;
+  BlocksSet ReachesNonLocalUseBlocks;
+  if (!blockIsSimpleEnoughToThreadThrough(BB, NonLocalUseBlocks))
     return false;
 
+  if (NonLocalUseBlocks.contains(BI->getSuccessor(0)) &&
+      NonLocalUseBlocks.contains(BI->getSuccessor(1)))
+    return false;
+
+  for (BasicBlock *UseBB : NonLocalUseBlocks)
+    if (!findReaching(UseBB, BB, ReachesNonLocalUseBlocks))
+      return false;
+
   for (const auto &Pair : KnownValues) {
     ConstantInt *CB = Pair.first;
     ArrayRef<BasicBlock *> PredBBs = Pair.second.getArrayRef();
@@ -3586,8 +3601,7 @@ foldCondBranchOnValueKnownInPredecessorImpl(BranchInst *BI, DomTreeUpdater *DTU,
       continue;
 
     // Only revector to RealDest if no values defined in BB are live.
-    BlocksSet VisitedBlocksSet;
-    if (reachesUsed(RealDest, ReachesNonLocalUses, VisitedBlocksSet))
+    if (ReachesNonLocalUseBlocks.contains(RealDest))
       continue;
 
     LLVM_DEBUG({

llvm/test/Transforms/SimplifyCFG/jump-threading-max-jump-threading-live-blocks.ll

@@ -0,0 +1,95 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes=simplifycfg -S -max-jump-threading-live-blocks=3 < %s | FileCheck %s --check-prefixes=CHECK_LIMIT_3
+; RUN: opt -passes=simplifycfg -S -max-jump-threading-live-blocks=4 < %s | FileCheck %s --check-prefixes=CHECK_LIMIT_4
+
+; Test option -max-jump-threading-live-blocks=<num>
+
+define void @testB(ptr %ptrA, ptr %ptrB, i64 %a, i64 %b, i64 %c) {
+; CHECK_LIMIT_3-LABEL: define void @testB(
+; CHECK_LIMIT_3-SAME: ptr [[PTRA:%.*]], ptr [[PTRB:%.*]], i64 [[A:%.*]], i64 [[B:%.*]], i64 [[C:%.*]]) {
+; CHECK_LIMIT_3-NEXT:  [[MAINA:.*]]:
+; CHECK_LIMIT_3-NEXT:    [[COND:%.*]] = icmp slt i64 [[A]], [[B]]
+; CHECK_LIMIT_3-NEXT:    br i1 [[COND]], label %[[IFA:.*]], label %[[MAINB:.*]]
+; CHECK_LIMIT_3:       [[IFA]]:
+; CHECK_LIMIT_3-NEXT:    [[TMP0:%.*]] = load i64, ptr [[PTRA]], align 4
+; CHECK_LIMIT_3-NEXT:    br label %[[MAINB]]
+; CHECK_LIMIT_3:       [[MAINB]]:
+; CHECK_LIMIT_3-NEXT:    [[VALUE:%.*]] = phi i64 [ [[TMP0]], %[[IFA]] ], [ 0, %[[MAINA]] ]
+; CHECK_LIMIT_3-NEXT:    br i1 [[COND]], label %[[IFB:.*]], label %[[MAINC:.*]]
+; CHECK_LIMIT_3:       [[IFB]]:
+; CHECK_LIMIT_3-NEXT:    [[COND2:%.*]] = icmp slt i64 [[A]], [[C]]
+; CHECK_LIMIT_3-NEXT:    br i1 [[COND2]], label %[[IFB_ARM1:.*]], label %[[IFB_ARM2:.*]]
+; CHECK_LIMIT_3:       [[IFB_ARM1]]:
+; CHECK_LIMIT_3-NEXT:    [[PTR_ARM1:%.*]] = getelementptr i64, ptr [[PTRB]], i64 8
+; CHECK_LIMIT_3-NEXT:    store i128 0, ptr [[PTR_ARM1]], align 4
+; CHECK_LIMIT_3-NEXT:    br label %[[IFB_JOIN:.*]]
+; CHECK_LIMIT_3:       [[IFB_ARM2]]:
+; CHECK_LIMIT_3-NEXT:    [[PTR_ARM2:%.*]] = getelementptr i64, ptr [[PTRB]], i64 16
+; CHECK_LIMIT_3-NEXT:    store i128 0, ptr [[PTR_ARM2]], align 4
+; CHECK_LIMIT_3-NEXT:    br label %[[IFB_JOIN]]
+; CHECK_LIMIT_3:       [[IFB_JOIN]]:
+; CHECK_LIMIT_3-NEXT:    [[PTRC:%.*]] = phi ptr [ [[PTR_ARM1]], %[[IFB_ARM1]] ], [ [[PTR_ARM2]], %[[IFB_ARM2]] ]
+; CHECK_LIMIT_3-NEXT:    store i64 [[VALUE]], ptr [[PTRC]], align 4
+; CHECK_LIMIT_3-NEXT:    br label %[[MAINC]]
+; CHECK_LIMIT_3:       [[MAINC]]:
+; CHECK_LIMIT_3-NEXT:    ret void
+;
+; CHECK_LIMIT_4-LABEL: define void @testB(
+; CHECK_LIMIT_4-SAME: ptr [[PTRA:%.*]], ptr [[PTRB:%.*]], i64 [[A:%.*]], i64 [[B:%.*]], i64 [[C:%.*]]) {
+; CHECK_LIMIT_4-NEXT:  [[MAINA:.*:]]
+; CHECK_LIMIT_4-NEXT:    [[COND:%.*]] = icmp slt i64 [[A]], [[B]]
+; CHECK_LIMIT_4-NEXT:    br i1 [[COND]], label %[[IFA:.*]], label %[[MAINC:.*]]
+; CHECK_LIMIT_4:       [[IFA]]:
+; CHECK_LIMIT_4-NEXT:    [[TMP0:%.*]] = load i64, ptr [[PTRA]], align 4
+; CHECK_LIMIT_4-NEXT:    [[COND2:%.*]] = icmp slt i64 [[A]], [[C]]
+; CHECK_LIMIT_4-NEXT:    br i1 [[COND2]], label %[[IFB_ARM1:.*]], label %[[IFB_ARM2:.*]]
+; CHECK_LIMIT_4:       [[IFB_ARM1]]:
+; CHECK_LIMIT_4-NEXT:    [[PTR_ARM1:%.*]] = getelementptr i64, ptr [[PTRB]], i64 8
+; CHECK_LIMIT_4-NEXT:    store i128 0, ptr [[PTR_ARM1]], align 4
+; CHECK_LIMIT_4-NEXT:    br label %[[IFB_JOIN:.*]]
+; CHECK_LIMIT_4:       [[IFB_ARM2]]:
+; CHECK_LIMIT_4-NEXT:    [[PTR_ARM2:%.*]] = getelementptr i64, ptr [[PTRB]], i64 16
+; CHECK_LIMIT_4-NEXT:    store i128 0, ptr [[PTR_ARM2]], align 4
+; CHECK_LIMIT_4-NEXT:    br label %[[IFB_JOIN]]
+; CHECK_LIMIT_4:       [[IFB_JOIN]]:
+; CHECK_LIMIT_4-NEXT:    [[PTRC:%.*]] = phi ptr [ [[PTR_ARM1]], %[[IFB_ARM1]] ], [ [[PTR_ARM2]], %[[IFB_ARM2]] ]
+; CHECK_LIMIT_4-NEXT:    store i64 [[TMP0]], ptr [[PTRC]], align 4
+; CHECK_LIMIT_4-NEXT:    br label %[[MAINC]]
+; CHECK_LIMIT_4:       [[MAINC]]:
+; CHECK_LIMIT_4-NEXT:    ret void
+;
+mainA:
+  %cond = icmp slt i64 %a, %b
+  br i1 %cond, label %ifA, label %mainB
+
+ifA:
+  %518 = load i64, ptr %ptrA
+  br label %mainB
+
+; Use of %value is not in either immediate destination of mainB.
+mainB:
+  %value = phi i64 [ %518, %ifA ], [ zeroinitializer, %mainA ]
+  br i1 %cond, label %ifB, label %mainC
+
+ifB:
+  %cond2 = icmp slt i64 %a, %c
+  br i1 %cond2, label %ifB_arm1, label %ifB_arm2
+
+ifB_arm1:
+  %ptr_arm1 = getelementptr i64, ptr %ptrB, i64 8
+  store i128 0, ptr %ptr_arm1
+  br label %ifB_join
+
+ifB_arm2:
+  %ptr_arm2 = getelementptr i64, ptr %ptrB, i64 16
+  store i128 0, ptr %ptr_arm2
+  br label %ifB_join
+
+ifB_join:
+  %ptrC = phi ptr [ %ptr_arm1, %ifB_arm1 ], [ %ptr_arm2, %ifB_arm2 ]
+  store i64 %value, ptr %ptrC
+  br label %mainC
+
+mainC:
+  ret void
+}