[LoopPredication] Account for critical edges when inserting assumes. PR26496

[mkazantsev]

Summary

Loop predication can insert assumes to preserve knowledge about some facts that
may otherwise be lost, because loop predication is a lossy transform. When a guard
is represented as branch by widenable condition, it should insert it in the guarded
block. However, if the guarded block has other predecessors than the guard block,
then the condition might not dominate it. Currently we generate invalid code here.

One possible fix here is to split critical edge and insert the assume there, but in this
case we had to modify CFG, which Loop Predication is not currently doing, and we
want to keep it that way.

The fix is to handle this case by inserting a Phi which takes Cond as input from the
guard block and true from any other blocks. This is valid in terms of IR and does
not introduce any new knowledge if we came from another block.

[nikic]

LGTM based on your description, but I'm not really familiar with this pass.

llvm/lib/Transforms/Scalar/LoopPredication.cpp

bool LoopPredication::widenWidenableBranchGuardConditions(
 
   // Emit the new guard condition
   IRBuilder<> Builder(findInsertPt(BI, Checks));
   Value *AllChecks = Builder.CreateAnd(Checks);
   auto *OldCond = BI->getCondition();
   BI->setCondition(AllChecks);
   if (InsertAssumesOfPredicatedGuardsConditions) {
     Builder.SetInsertPoint(IfTrueBB, IfTrueBB->getFirstInsertionPt());
-    Builder.CreateAssumption(Cond);
+    // If this block has other predecessors, we might not be able to use Cond.
+    // In this case, create a Phi where every other input is `true` and input
+    // from guard block is Cond.
+    Value *AssumeCond = Cond;
+    if (!IfTrueBB->getUniquePredecessor()) {
+      auto *GuardBB = BI->getParent();
+      auto *PN = Builder.CreatePHI(Cond->getType(), pred_size(IfTrueBB),
+                                   "assume.cond");
+      for (auto *Pred : predecessors(IfTrueBB))
+        PN->addIncoming(Pred == GuardBB ? Cond : Builder.getTrue(), Pred);
+      AssumeCond = PN;
+    }
+    Builder.CreateAssumption(AssumeCond);
   }
   RecursivelyDeleteTriviallyDeadInstructions(OldCond, nullptr /* TLI */, MSSAU);
   assert(isGuardAsWidenableBranch(BI) &&
          "Stopped being a guard after transform?");
 
   LLVM_DEBUG(dbgs() << "Widened checks = " << NumWidened << "\n");
   return true;
 }

llvm/test/Transforms/LoopPredication/pr61022.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -passes=loop-predication < %s 2>&1 | FileCheck %s
-; REQUIRES: asserts
-; XFAIL: *
 
 declare void @llvm.experimental.deoptimize.isVoid(...)
 
-; FIXME: Loop predication here inserts assume across the critical edge, and
-; it leads to malformed IR (assume's condition does not dominate it).
 define void @test_01(i1 %cond) {
-; CHECK-LABEL: test
+; CHECK-LABEL: @test_01(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[INST:%.*]] = call i1 @llvm.experimental.widenable.condition()
+; CHECK-NEXT:    [[TMP0:%.*]] = and i1 true, [[INST]]
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       unreached:
+; CHECK-NEXT:    unreachable
+; CHECK:       loop:
+; CHECK-NEXT:    [[INST3:%.*]] = phi i32 [ 0, [[BB:%.*]] ], [ [[INST4:%.*]], [[BACKEDGE:%.*]] ]
+; CHECK-NEXT:    [[INST4]] = add nsw i32 [[INST3]], 1
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[BACKEDGE]], label [[GUARD_BLOCK:%.*]]
+; CHECK:       normal_ret:
+; CHECK-NEXT:    ret void
+; CHECK:       backedge:
+; CHECK-NEXT:    [[ASSUME_COND:%.*]] = phi i1 [ [[INST9:%.*]], [[GUARD_BLOCK]] ], [ true, [[LOOP]] ]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[ASSUME_COND]])
+; CHECK-NEXT:    [[INST7:%.*]] = icmp sgt i32 [[INST3]], 137
+; CHECK-NEXT:    br i1 [[INST7]], label [[UNREACHED:%.*]], label [[LOOP]]
+; CHECK:       guard_block:
+; CHECK-NEXT:    [[INST9]] = icmp ult i32 [[INST4]], 10000
+; CHECK-NEXT:    br i1 [[TMP0]], label [[BACKEDGE]], label [[DEOPT:%.*]]
+; CHECK:       deopt:
+; CHECK-NEXT:    call void (...) @llvm.experimental.deoptimize.isVoid(i32 13) [ "deopt"() ]
+; CHECK-NEXT:    ret void
+; CHECK:       done:
+; CHECK-NEXT:    ret void
+;
 bb:
   %inst = call i1 @llvm.experimental.widenable.condition()
   br label %loop
 
 unreached:                                              ; preds = %backedge
   unreachable
 
 loop:                                              ; preds = %backedge, %bb