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

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 should 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.

Differential Revision: https://reviews.llvm.org/D144859
Reviewed By: nikic, skatkov

diff --git a/llvm/lib/Transforms/Scalar/LoopPredication.cpp b/llvm/lib/Transforms/Scalar/LoopPredication.cpp
index 49c0fff..1724d80 100644 (file)
--- a/llvm/lib/Transforms/Scalar/LoopPredication.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopPredication.cpp
@@ -863,7 +863,19 @@ bool LoopPredication::widenWidenableBranchGuardConditions(
   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) &&

diff --git a/llvm/test/Transforms/LoopPredication/pr61022.ll b/llvm/test/Transforms/LoopPredication/pr61022.ll
index 0381368..33c49dd 100644 (file)
--- a/llvm/test/Transforms/LoopPredication/pr61022.ll
+++ b/llvm/test/Transforms/LoopPredication/pr61022.ll
@@ -1,13 +1,36 @@
+; 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