[SimplifyCFG] Do not split blocks if breaking a loop

[junbuml]

Summary

Do not insert a new block, if it could transform a loop into an irreducible loop.

[junbuml]

Kindly ping. Please let me know any comment.

[efriedma]

Another option here would be to make SimplifyCFG insert a preheader for loops which don't have one. That seems better in terms of canonicalization.

[efriedma]

Err, nevermind... I somehow thought we were preserving preheaders already in SimplifyCFG, but I guess we aren't.

I'll think about this a bit more.

[efriedma]

Added some code review comments.

At a higher level, we should consider changing SimplifyCFG so it doesn't fight LoopSimplify over the canonical form of a loop... but just fixing the bug is fine for now.

lib/Transforms/Utils/SimplifyCFG.cpp

172

Do we need to do any additional work to keep this map up-to-date? The predecessors of the loop header could get folded.

1800

BBEnd instead of BI1->getSuccessor(0)?

1802

*irreducible

include/llvm/Transforms/Utils/Local.h

 class DIBuilder;
 class DominatorTree;
 class LazyValueInfo;
 
 template<typename T> class SmallVectorImpl;
 
 typedef SmallVector<DbgValueInst *, 1> DbgValueList;
 
+typedef DenseMap<const BasicBlock *, SmallPtrSet<const BasicBlock *, 4>>
+    LoopHeaderList;
+
 //===----------------------------------------------------------------------===//
 //  Local constant propagation.
 //
 
 /// If a terminator instruction is predicated on a constant value, convert it
 /// into an unconditional branch to the constant destination.
 /// This is a nontrivial operation because the successors of this basic block
 /// must have their PHI nodes updated.

 /// example, it adjusts branches to branches to eliminate the extra hop, it
 /// eliminates unreachable basic blocks, and does other "peephole" optimization
 /// of the CFG.  It returns true if a modification was made, possibly deleting
 /// the basic block that was pointed to. LoopHeaders is an optional input
 /// parameter, providing the set of loop header that SimplifyCFG should not
 /// eliminate.
 bool SimplifyCFG(BasicBlock *BB, const TargetTransformInfo &TTI,
                  unsigned BonusInstThreshold, AssumptionCache *AC = nullptr,
-                 SmallPtrSetImpl<BasicBlock *> *LoopHeaders = nullptr);
+                 LoopHeaderList *LoopHeaders = nullptr);
 
 /// This function is used to flatten a CFG. For example, it uses parallel-and
 /// and parallel-or mode to collapse if-conditions and merge if-regions with
 /// identical statements.
 bool FlattenCFG(BasicBlock *BB, AliasAnalysis *AA = nullptr);
 
 /// If this basic block is ONLY a setcc and a branch, and if a predecessor
 /// branches to us and one of our successors, fold the setcc into the

lib/Transforms/Scalar/SimplifyCFGPass.cpp

 static bool iterativelySimplifyCFG(Function &F, const TargetTransformInfo &TTI,
                                    AssumptionCache *AC,
                                    unsigned BonusInstThreshold) {
   bool Changed = false;
   bool LocalChange = true;
 
   SmallVector<std::pair<const BasicBlock *, const BasicBlock *>, 32> Edges;
   FindFunctionBackedges(F, Edges);
-  SmallPtrSet<BasicBlock *, 16> LoopHeaders;
+  LoopHeaderList LoopHeaders;
+
   for (unsigned i = 0, e = Edges.size(); i != e; ++i)
-    LoopHeaders.insert(const_cast<BasicBlock *>(Edges[i].second));
+    LoopHeaders[Edges[i].second].insert(Edges[i].first);
 
   while (LocalChange) {
     LocalChange = false;
 
     // Loop over all of the basic blocks and remove them if they are unneeded.
     for (Function::iterator BBIt = F.begin(); BBIt != F.end(); ) {
       if (SimplifyCFG(&*BBIt++, TTI, BonusInstThreshold, AC, &LoopHeaders)) {
         LocalChange = true;

lib/Transforms/Utils/SimplifyCFG.cpp

struct ValueEqualityComparisonCase {
   bool operator==(BasicBlock *RHSDest) const { return Dest == RHSDest; }
 };
 
 class SimplifyCFGOpt {
   const TargetTransformInfo &TTI;
   const DataLayout &DL;
   unsigned BonusInstThreshold;
   AssumptionCache *AC;
-  SmallPtrSetImpl<BasicBlock *> *LoopHeaders;
+  LoopHeaderList *LoopHeaders;
   Value *isValueEqualityComparison(TerminatorInst *TI);
   BasicBlock *GetValueEqualityComparisonCases(
       TerminatorInst *TI, std::vector<ValueEqualityComparisonCase> &Cases);
   bool SimplifyEqualityComparisonWithOnlyPredecessor(TerminatorInst *TI,
                                                      BasicBlock *Pred,
                                                      IRBuilder<> &Builder);
   bool FoldValueComparisonIntoPredecessors(TerminatorInst *TI,
                                            IRBuilder<> &Builder);
 
   bool SimplifyReturn(ReturnInst *RI, IRBuilder<> &Builder);
   bool SimplifyResume(ResumeInst *RI, IRBuilder<> &Builder);
   bool SimplifySingleResume(ResumeInst *RI);
   bool SimplifyCommonResume(ResumeInst *RI);
   bool SimplifyCleanupReturn(CleanupReturnInst *RI);
   bool SimplifyUnreachable(UnreachableInst *UI);
   bool SimplifySwitch(SwitchInst *SI, IRBuilder<> &Builder);
   bool SimplifyIndirectBr(IndirectBrInst *IBI);
   bool SimplifyUncondBranch(BranchInst *BI, IRBuilder<> &Builder);
   bool SimplifyCondBranch(BranchInst *BI, IRBuilder<> &Builder);
 
 public:
   SimplifyCFGOpt(const TargetTransformInfo &TTI, const DataLayout &DL,
                  unsigned BonusInstThreshold, AssumptionCache *AC,
-                 SmallPtrSetImpl<BasicBlock *> *LoopHeaders)
+                 LoopHeaderList *LoopHeaders)
       : TTI(TTI), DL(DL), BonusInstThreshold(BonusInstThreshold), AC(AC),
         LoopHeaders(LoopHeaders) {}
 
   bool run(BasicBlock *BB);
 };
 
 } // end anonymous namespace
 
namespace {
   };
 
 } // end anonymous namespace
 
 /// Given an unconditional branch that goes to BBEnd,
 /// check whether BBEnd has only two predecessors and the other predecessor
 /// ends with an unconditional branch. If it is true, sink any common code
 /// in the two predecessors to BBEnd.
-static bool SinkThenElseCodeToEnd(BranchInst *BI1) {
+static bool SinkThenElseCodeToEnd(BranchInst *BI1,
+                                  LoopHeaderList *LoopHeaders) {
   assert(BI1->isUnconditional());
   BasicBlock *BBEnd = BI1->getSuccessor(0);
 
   // We support two situations:
   //   (1) all incoming arcs are unconditional
   //   (2) one incoming arc is conditional
   //
   // (2) is very common in switch defaults and
while (ProfitableToSinkInstruction(LRI) && Idx < ScanIdx) {
       ++Idx;
     }
     if (!Profitable)
       return false;
 
     DEBUG(dbgs() << "SINK: Splitting edge\n");
     // We have a conditional edge and we're going to sink some instructions.
     // Insert a new block postdominating all blocks we're going to sink from.
+
+    if (LoopHeaders && LoopHeaders->count(BI1->getSuccessor(0))) {
+      // If we are handling a loop header and both backedges and non-backedges
+      // are hooked to the new block, this will make the loop irreduciable
+      // because the new block is not dominated by the loop header. In such
+      // case, we should not split it to keep it as a valid loop.
+      bool HasBackedge = false;
+      bool HasNonBackedge = false;
+      SmallPtrSet<const BasicBlock *, 4> &BackEdges =
+          (*LoopHeaders)[BI1->getSuccessor(0)];
+      for (auto PredBB : UnconditionalPreds) {
+        if (BackEdges.count(PredBB))
+          HasBackedge = true;
+        else
+          HasNonBackedge = true;
+        if (HasNonBackedge && HasBackedge)
+          return false;
+      }
+    }
+
     if (!SplitBlockPredecessors(BI1->getSuccessor(0), UnconditionalPreds,
                                 ".sink.split"))
       // Edges couldn't be split.
       return false;
     Changed = true;
   }
 
   // Now that we've analyzed all potential sinking candidates, perform the
static bool TryToMergeLandingPad(LandingPadInst *LPad, BranchInst *BI,
   }
   return false;
 }
 
 bool SimplifyCFGOpt::SimplifyUncondBranch(BranchInst *BI,
                                           IRBuilder<> &Builder) {
   BasicBlock *BB = BI->getParent();
 
-  if (SinkCommon && SinkThenElseCodeToEnd(BI))
+  if (SinkCommon && SinkThenElseCodeToEnd(BI, LoopHeaders))
     return true;
 
   // If the Terminator is the only non-phi instruction, simplify the block.
   // if LoopHeader is provided, check if the block is a loop header
   // (This is for early invocations before loop simplify and vectorization
   // to keep canonical loop forms for nested loops.
   // These blocks can be eliminated when the pass is invoked later
   // in the back-end.)

 
 /// This function is used to do simplification of a CFG.
 /// For example, it adjusts branches to branches to eliminate the extra hop,
 /// eliminates unreachable basic blocks, and does other "peephole" optimization
 /// of the CFG.  It returns true if a modification was made.
 ///
 bool llvm::SimplifyCFG(BasicBlock *BB, const TargetTransformInfo &TTI,
                        unsigned BonusInstThreshold, AssumptionCache *AC,
-                       SmallPtrSetImpl<BasicBlock *> *LoopHeaders) {
+                       LoopHeaderList *LoopHeaders) {
   return SimplifyCFGOpt(TTI, BB->getModule()->getDataLayout(),
                         BonusInstThreshold, AC, LoopHeaders)
       .run(BB);
 }

test/Transforms/SimplifyCFG/sink-common-code.ll

if.end:
   %r6.addr.0 = phi i32 [ %0, %if.then ], [ %1, %if.else ]
   ret i32 %r6.addr.0
 }
 
 ; CHECK-LABEL: @test_inline_asm1(
 ; CHECK: call i32 asm "rorl $2, $0", "=&r,0,n,~{dirflag},~{fpsr},~{flags}"(i32 %r6, i32 8)
 ; CHECK: call i32 asm "rorl $2, $0", "=&r,0,n,~{dirflag},~{fpsr},~{flags}"(i32 %r6, i32 6)
 
+; Check that simplifycfg doesn't transform the loop (%loop.header) into an irreduciable loop.
+
+define zeroext i1 @test_keeploop(i1 zeroext %flag, i32 %blksA, i32 %blksB, i32 %nblks, i8* %p) {
+entry:
+  br i1 %flag, label %if.then, label %loop.header
+
+if.then:
+  %cmp = icmp uge i32 %blksA, %nblks
+  %frombool1 = zext i1 %cmp to i8
+  store i8 %frombool1, i8* %p
+  br label %loop.header
+
+loop.header:
+  %c = call i1 @f2()
+  br i1 %c, label %loop.inc,  label %loop.exit
+
+loop.inc:
+  %add = add i32 %nblks, %blksB
+  %cmp2 = icmp ule i32 %add, %blksA
+  %frombool3 = zext i1 %cmp2 to i8
+  store i8 %frombool3, i8* %p
+  br label %loop.header
+
+loop.exit:
+  ret i1 true
+}
+declare i1 @f2()
+
+; CHECK-LABEL: @test_keeploop
+; CHECK-LABE: if.then:
+; CHECK: br label %loop.header
+; CHECK-LABE: loop.inc:
+; CHECK: br label %loop.header
+
 declare i32 @call_target()
 
 define void @test_operand_bundles(i1 %cond, i32* %ptr) {
 entry:
   br i1 %cond, label %left, label %right
 
 left:
   %val0 = call i32 @call_target() [ "deopt"(i32 10) ]