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

[Max Kazantsev via llvm-commits]

Author: Max Kazantsev
Date: 2023-02-27T18:26:17+07:00
New Revision: a18ce47a3e0065d0a26b370a73440411b76d8493

URL: https://github.com/llvm/llvm-project/commit/a18ce47a3e0065d0a26b370a73440411b76d8493
DIFF: https://github.com/llvm/llvm-project/commit/a18ce47a3e0065d0a26b370a73440411b76d8493.diff

LOG: [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

Added: 
    

Modified: 
    llvm/lib/Transforms/Scalar/LoopPredication.cpp
    llvm/test/Transforms/LoopPredication/pr61022.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/LoopPredication.cpp b/llvm/lib/Transforms/Scalar/LoopPredication.cpp
index 49c0fff84d817..1724d809458df 100644
--- 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 038136800c1cb..33c49ddd7fbe0 100644
--- 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