[JumpThread] Do not break loop structures when splitting blocks

[junbuml]

Summary

This change prevent SimplifyPartiallyRedundantLoad() from breaking a loop structure if splitting predecessors turn a loop into an irreducible loop.

[haicheng]

include/llvm/Transforms/Scalar/JumpThreading.h

68

Please add a comment to explain how the map works.

[junbuml]

Kindly ping. Please let me know any comment.

[junbuml]

Kindly ping. Please let me know any comment.

include/llvm/Transforms/Scalar/JumpThreading.h

 /// revectored to the false side of the second if.
 ///
 class JumpThreadingPass : public PassInfoMixin<JumpThreadingPass> {
   TargetLibraryInfo *TLI;
   LazyValueInfo *LVI;
   std::unique_ptr<BlockFrequencyInfo> BFI;
   std::unique_ptr<BranchProbabilityInfo> BPI;
   bool HasProfileData = false;
+
+  // Map a loop header to its predecessors from the backedge.
 #ifdef NDEBUG
-  SmallPtrSet<const BasicBlock *, 16> LoopHeaders;
+  DenseMap<const BasicBlock *, SmallPtrSet<const BasicBlock *, 16>> LoopHeaders;
 #else
-  SmallSet<AssertingVH<const BasicBlock>, 16> LoopHeaders;
+  DenseMap<AssertingVH<const BasicBlock>, SmallPtrSet<const BasicBlock *, 16>>
+      LoopHeaders;
 #endif
+
   DenseSet<std::pair<Value *, BasicBlock *>> RecursionSet;
 
   unsigned BBDupThreshold;
 
   // RAII helper for updating the recursion stack.
   struct RecursionSetRemover {
     DenseSet<std::pair<Value *, BasicBlock *>> &TheSet;
     std::pair<Value *, BasicBlock *> ThePair;

lib/Transforms/Scalar/JumpThreading.cpp

 /// enough to track all of these properties and keep it up-to-date as the CFG
 /// mutates, so we don't allow any of these transformations.
 ///
 void JumpThreadingPass::FindLoopHeaders(Function &F) {
   SmallVector<std::pair<const BasicBlock*,const BasicBlock*>, 32> Edges;
   FindFunctionBackedges(F, Edges);
 
   for (const auto &Edge : Edges)
-    LoopHeaders.insert(Edge.second);
+    LoopHeaders[Edge.second].insert(Edge.first);
 }
 
 /// getKnownConstant - Helper method to determine if we can thread over a
 /// terminator with the given value as its condition, and if so what value to
 /// use for that. What kind of value this is depends on whether we want an
 /// integer or a block address, but an undef is always accepted.
 /// Returns null if Val is null or not an appropriate constant.
 static Constant *getKnownConstant(Value *Val, ConstantPreference Preference) {
bool JumpThreadingPass::ProcessBlock(BasicBlock *BB) {
   // successor, merge the blocks.  This encourages recursive jump threading
   // because now the condition in this block can be threaded through
   // predecessors of our predecessor block.
   if (BasicBlock *SinglePred = BB->getSinglePredecessor()) {
     const TerminatorInst *TI = SinglePred->getTerminator();
     if (!TI->isExceptional() && TI->getNumSuccessors() == 1 &&
         SinglePred != BB && !hasAddressTakenAndUsed(BB)) {
       // If SinglePred was a loop header, BB becomes one.
-      if (LoopHeaders.erase(SinglePred))
-        LoopHeaders.insert(BB);
+      if (LoopHeaders.count(SinglePred)) {
+        LoopHeaders[BB] = LoopHeaders[SinglePred];
+        LoopHeaders.erase(SinglePred);
+      }
 
       LVI->eraseBlock(SinglePred);
       MergeBasicBlockIntoOnlyPred(BB);
 
       return true;
     }
   }
 
for (BasicBlock *P : predecessors(LoadBB)) {
       // If the predecessor is an indirect goto, we can't split the edge.
       if (isa<IndirectBrInst>(P->getTerminator()))
         return false;
 
       if (!AvailablePredSet.count(P))
         PredsToSplit.push_back(P);
     }
 
-    // Split them out to their own block.
-    UnavailablePred = SplitBlockPreds(LoadBB, PredsToSplit, "thread-pre-split");
+    // Split them out to their own block unless it mess up loop strutures.
+    if (!(UnavailablePred =
+              SplitBlockPreds(LoadBB, PredsToSplit, "thread-pre-split")))
+      return false;
   }
 
   // If the value isn't available in all predecessors, then there will be
   // exactly one where it isn't available.  Insert a load on that edge and add
   // it to the AvailablePreds list.
   if (UnavailablePred) {
     assert(UnavailablePred->getTerminator()->getNumSuccessors() == 1 &&
            "Can't handle critical edge here!");

 }
 
 /// Create a new basic block that will be the predecessor of BB and successor of
 /// all blocks in Preds. When profile data is available, update the frequency of
 /// this new block.
 BasicBlock *JumpThreadingPass::SplitBlockPreds(BasicBlock *BB,
                                                ArrayRef<BasicBlock *> Preds,
                                                const char *Suffix) {
+  bool HasBackedge = false;
+  bool HasNonBackedge = false;
+
+  // If we are splitting a backedge and an edge from outside of the loop is also
+  // hooked to the new block, the backedge will no longer be dominated by the
+  // header. In such case, we should not split it to keep it as a valid loop.
+  if (LoopHeaders.count(BB)) {
+    for (auto PredBB : Preds) {
+      if (LoopHeaders[BB].count(PredBB))
+        HasBackedge = true;
+      else
+        HasNonBackedge = true;
+      if (HasNonBackedge && HasBackedge)
+        return nullptr;
+    }
+  }
+
   // Collect the frequencies of all predecessors of BB, which will be used to
   // update the edge weight on BB->SuccBB.
   BlockFrequency PredBBFreq(0);
   if (HasProfileData)
     for (auto Pred : Preds)
       PredBBFreq += BFI->getBlockFreq(Pred) * BPI->getEdgeProbability(Pred, BB);
 
   BasicBlock *PredBB = SplitBlockPredecessors(BB, Preds, Suffix);

test/Transforms/JumpThreading/keep-loop.ll

+; RUN: opt < %s -jump-threading -loop-unswitch -S | FileCheck %s
+
+; This tests jump threading do not turn a proper loop structures into an
+; irreducible loop by SimplifyPartiallyRedundantLoad(). This test contain
+; obvious unswithch opportunity which should be handled by loop unswitch pass.
+;
+; CHECK-LABEL: @fn_keep_loop
+; CHECK-LABEL: while.cond2.us:
+; CHECK: br i1 true, label %if.then.us, label %if.end.us
+; CHECK-LABELif.then.us:
+; CHECK: call void @fn3(i64 %l3.us)
+; CHECK-LABEL: while.cond2:
+; CHECK: br i1 false, label %if.then, label %if.end
+; CHECK-LABEL: if.then:
+; CHECK: call void @fn3(i64 %l3)
+
+define i32 @fn_keep_loop(i1 %c0, i1 %c1, i64* %ptr0, i64* %ptr1) {
+entry:
+  br i1 %c0, label %while.cond, label %sw.epilog
+
+while.cond:                                       ; preds = %land.rhs, %entry
+  %magicptr = load i64, i64* %ptr1, align 8
+  %c = icmp eq i64 %magicptr, 0
+  br i1 %c, label %while.cond2, label %land.rhs
+
+land.rhs:                                         ; preds = %while.cond
+  %arrayidx = getelementptr inbounds i64, i64* %ptr0, i64 1
+  %l2 = load i64, i64* %arrayidx, align 4
+  %tobool1 = icmp eq i64 %l2, 0
+  br i1 %tobool1, label %while.cond2, label %while.cond
+
+while.cond2:
+  %l3 = load i64, i64* %ptr1, align 8
+  br i1 %c1, label %if.then, label %if.end
+
+if.then:
+  call void @fn3(i64 %l3)
+  br label %if.end
+
+if.end:
+  %c2 = icmp eq i64 %l3,  0
+  br i1 %c2, label %sw.epilog, label %while.body4
+
+while.body4:                                      ; preds = %while.cond2
+  call void @fn2(i64 %l3)
+  br label %while.cond2
+
+sw.epilog:                                        ; preds = %while.cond2, %entry
+  ret i32 undef
+}
+
+declare void @fn2(i64)
+declare void @fn3(i64)