[LoopPredication] Fix where we generate widened condition. PR61963

[anna]

Summary

Loop predication's predicateLoopExit pass does two incorrect things:

It sinks the widenable call into the loop, thereby converting an invariant condition to a variant one
It widens the widenable call at a branch thereby converting it into a loop-varying one.

The latter is problematic when the branch may have been a loop-invariant
branch and prior optimizations (such as indvars) may have relied on this
fact, and updated the deopt state accordingly.

Now, when we widen this with a loop-varying condition, the deopt state
is no longer correct. See
https://github.com/llvm/llvm-project/issues/61963.

[mkazantsev]

I generally agree with the change in code, but can we shrink the test? I'm pretty sure most of things here are not really required.

llvm/test/Transforms/LoopPredication/pr61963.ll

11

Is it possible to reduce the test via bugpoint? Just assert whenever you move non-loop-invariant into loop in this transform and let it do the job.

[anna]

llvm/test/Transforms/LoopPredication/pr61963.ll

11

Yea, I had reduced this with bugpoint using exactly similar assert. The problem is it removes all the deopt calls as well and makes it oversimplified (we can see the problem, but it's not that clear, since the deopt states are missing). I can hack bug point to avoid removing these deopt calls. Will try that or manually reducing this further.

[anna]

simplified testcase.

[anna]

renamed some blocks in test and simplified further.

[apilipenko]

Given that it's the second time we have to fix a similar issue, we should make a common utility for widening that makes it impossible to emit a widened condition in an incorrect context.

[anna]

rebased over landed testcase.

[anna]

In D147662#4249473, @apilipenko wrote:

Given that it's the second time we have to fix a similar issue, we should make a common utility for widening that makes it impossible to emit a widened condition in an incorrect context.

Yep, I think it works naturally with the guard widening fix as well (as a follow-up commit). One thing though - this does not prevent from directly identifying an incorrect context unless the utility is used (in short, a verifier would be best).

Also, we should add asserts at two utilities in GuardUtils: https://github.com/llvm/llvm-project/blob/main/llvm/lib/Transforms/Utils/GuardUtils.cpp#L82

1. setWidenableBranchCond(BranchInst *WidenableBR, Value *NewCond)
2. widenWidenableBranch(BranchInst *WidenableBR, Value *NewCond)

Assert that we didn't convert a loop-invariant WidenableBR into a loop-variant one. We add an optional LI to this API, but LI is available in both these passes which use the utilities. Will try this out.

[mkazantsev]

LG, thanks

llvm/test/Transforms/LoopPredication/pr61963.ll

89

I guess we can also just leave 1 value here (the one that was replaced with invariant) rather than keeping all of them.

llvm/lib/Transforms/Scalar/LoopPredication.cpp

if (parseWidenableBranch(BI, Cond, WC, IfTrueBB, IfFalseBB) &&
         L->contains(IfTrueBB)) {
       WC->set(ConstantInt::getTrue(IfTrueBB->getContext()));
       ChangedLoop = true;
     }
   }
   if (ChangedLoop)
     SE->forgetLoop(L);
 
+  // The insertion point for the widening should be at the widenably call, not
+  // at the WidenableBR. If we do this at the widenableBR, we can incorrectly
+  // change a loop-invariant condition to a loop-varying one.
+  auto *IP = cast<Instruction>(WidenableBR->getCondition());
+
   // The use of umin(all analyzeable exits) instead of latch is subtle, but
   // important for profitability.  We may have a loop which hasn't been fully
   // canonicalized just yet.  If the exit we chose to widen is provably never
   // taken, we want the widened form to *also* be provably never taken.  We
   // can't guarantee this as a current unanalyzeable exit may later become
   // analyzeable, but we can at least avoid the obvious cases.
   const SCEV *MinEC = getMinAnalyzeableBackedgeTakenCount(*SE, *DT, L);
   if (isa<SCEVCouldNotCompute>(MinEC) || MinEC->getType()->isPointerTy() ||
       !SE->isLoopInvariant(MinEC, L) ||
-      !Rewriter.isSafeToExpandAt(MinEC, WidenableBR))
+      !Rewriter.isSafeToExpandAt(MinEC, IP))
     return ChangedLoop;
 
-  // Subtlety: We need to avoid inserting additional uses of the WC.  We know
-  // that it can only have one transitive use at the moment, and thus moving
-  // that use to just before the branch and inserting code before it and then
-  // modifying the operand is legal.
-  auto *IP = cast<Instruction>(WidenableBR->getCondition());
-  // Here we unconditionally modify the IR, so after this point we should return
-  // only `true`!
-  IP->moveBefore(WidenableBR);
-  if (MSSAU)
-    if (auto *MUD = MSSAU->getMemorySSA()->getMemoryAccess(IP))
-       MSSAU->moveToPlace(MUD, WidenableBR->getParent(),
-                          MemorySSA::BeforeTerminator);
   Rewriter.setInsertPoint(IP);
   IRBuilder<> B(IP);
 
   bool InvalidateLoop = false;
   Value *MinECV = nullptr; // lazily generated if needed
   for (BasicBlock *ExitingBB : ExitingBlocks) {
     // If our exiting block exits multiple loops, we can only rewrite the
     // innermost one.  Otherwise, we're changing how many times the innermost

llvm/test/Transforms/LoopPredication/pr61963.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
 ; RUN: opt -S -passes=loop-predication < %s 2>&1 | FileCheck %s
 
 ; Do not convert branch in loop_outer block on the widenable_cond11 to a
 ; loop-varying one.
 ; It will result in a miscompile.
 ; deopt9 will have incorrect deopt state (it currently uses init_val because
 ; indvars identified that if that exit is taken, it will be taken on first
 ; iteration, since widenable_cond11 is a loop-invariant condition).
 define i32 @foo(ptr addrspace(1) %arg) {
 ; CHECK-LABEL: define i32 @foo
 ; CHECK-SAME: (ptr addrspace(1) [[ARG:%.*]]) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[INIT_VAL:%.*]] = load i32, ptr addrspace(1) [[ARG]], align 4
+; CHECK-NEXT:    [[WIDENABLE_COND11:%.*]] = call i1 @llvm.experimental.widenable.condition()
 ; CHECK-NEXT:    br label [[LOOP_OUTER:%.*]]
 ; CHECK:       loop_outer:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ [[PHI36:%.*]], [[OUTER_LOOP_LATCH:%.*]] ], [ 42, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[PHI21:%.*]] = phi i32 [ [[ADD39:%.*]], [[OUTER_LOOP_LATCH]] ], [ [[INIT_VAL]], [[ENTRY]] ]
 ; CHECK-NEXT:    [[ADD27:%.*]] = add i32 [[IV]], 1
 ; CHECK-NEXT:    [[ICMP28:%.*]] = icmp eq i32 [[ADD27]], 60
-; CHECK-NEXT:    [[TMP0:%.*]] = sub i32 -2, [[IV]]
-; CHECK-NEXT:    [[UMIN:%.*]] = call i32 @llvm.umin.i32(i32 [[TMP0]], i32 8)
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ugt i32 [[TMP0]], [[UMIN]]
-; CHECK-NEXT:    [[TMP2:%.*]] = freeze i1 [[TMP1]]
-; CHECK-NEXT:    [[WIDENABLE_COND11:%.*]] = call i1 @llvm.experimental.widenable.condition()
-; CHECK-NEXT:    [[TMP3:%.*]] = and i1 [[TMP2]], [[WIDENABLE_COND11]]
-; CHECK-NEXT:    br i1 [[TMP3]], label [[INNER_LOOP_PH:%.*]], label [[DEOPT9:%.*]]
+; CHECK-NEXT:    br i1 [[WIDENABLE_COND11]], label [[INNER_LOOP_PH:%.*]], label [[DEOPT9:%.*]]
 ; CHECK:       inner_loop_ph:
 ; CHECK-NEXT:    store atomic i32 606, ptr addrspace(1) [[ARG]] unordered, align 4
 ; CHECK-NEXT:    br label [[INNER_LOOP:%.*]]
 ; CHECK:       inner_loop:
 ; CHECK-NEXT:    [[PHI43:%.*]] = phi i32 [ 1, [[INNER_LOOP_PH]] ], [ [[ADD55:%.*]], [[INNER_LOOP_LATCH:%.*]] ]
 ; CHECK-NEXT:    [[PHI44:%.*]] = phi i32 [ [[ADD27]], [[INNER_LOOP_PH]] ], [ [[ADD48:%.*]], [[INNER_LOOP_LATCH]] ]
 ; CHECK-NEXT:    [[ADD48]] = add i32 [[PHI44]], 1
 ; CHECK-NEXT:    [[ICMP49:%.*]] = icmp eq i32 [[ADD48]], 0
-; CHECK-NEXT:    br i1 false, label [[DEOPT57:%.*]], label [[INNER_LOOP_LATCH]]
+; CHECK-NEXT:    br i1 [[ICMP49]], label [[DEOPT57:%.*]], label [[INNER_LOOP_LATCH]]
 ; CHECK:       inner_loop_latch:
 ; CHECK-NEXT:    store atomic i32 606, ptr addrspace(1) [[ARG]] unordered, align 4
 ; CHECK-NEXT:    [[ADD55]] = add nuw nsw i32 [[PHI43]], 1
 ; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp eq i32 [[ADD55]], 10
 ; CHECK-NEXT:    br i1 [[EXITCOND]], label [[OUTER_LOOP_LATCH]], label [[INNER_LOOP]]
 ; CHECK:       outer_loop_latch:
 ; CHECK-NEXT:    [[PHI36]] = phi i32 [ [[ADD48]], [[INNER_LOOP_LATCH]] ]
 ; CHECK-NEXT:    [[ADD39]] = add i32 [[PHI21]], 1
outer_loop_latch:                                             ; preds = %inner_loop_latch
   br label %loop_outer
 
 deopt9:                                           ; preds = %loop_outer
   %lcssa = phi i32 [ %init_val, %loop_outer ]
   %call53 = call i32 (...) @llvm.experimental.deoptimize.i32(i32 13) [ "deopt"(i32 606, i32 %lcssa) ]
   ret i32 %call53
 
 deopt57:                                             ; preds = %inner_loop
   %call62 = call i32 (...) @llvm.experimental.deoptimize.i32(i32 12) [ "deopt"(i32 7) ]
   ret i32 %call62
 }
 
 declare i32 @llvm.experimental.deoptimize.i32(...)
 
 ; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(inaccessiblemem: readwrite)
 declare noundef i1 @llvm.experimental.widenable.condition()

llvm/test/Transforms/LoopPredication/predicate-exits.ll

 exit:
   %result = phi i32 [ %loop.acc.next, %guarded2 ]
   ret i32 %result
 }
 
 define void @test_memssa() {
 ; CHECK-LABEL: @test_memssa(
 ; CHECK-NEXT:  bb:
-; CHECK-NEXT:    [[TMP1:%.*]] = call i1 @llvm.experimental.widenable.condition()
 ; CHECK-NEXT:    [[TMP:%.*]] = call i1 @llvm.experimental.widenable.condition()
+; CHECK-NEXT:    [[TMP1:%.*]] = call i1 @llvm.experimental.widenable.condition()
 ; CHECK-NEXT:    br i1 [[TMP]], label [[BB3:%.*]], label [[BB2:%.*]]
 ; CHECK:       bb2:
 ; CHECK-NEXT:    unreachable
 ; CHECK:       bb3:
 ; CHECK-NEXT:    br label [[BB4:%.*]]
 ; CHECK:       bb4:
 ; CHECK-NEXT:    [[TMP5:%.*]] = phi i32 [ [[TMP7:%.*]], [[BB6:%.*]] ], [ 0, [[BB3]] ]
 ; CHECK-NEXT:    br i1 true, label [[BB10:%.*]], label [[BB6]]