[InstCombine] Push freeze through recurrence phi

[nikic]

Summary

We really want to push freezes through recurrence phis, so that we freeze only the start value, rather than the IV value on every iteration. foldOpIntoPhi() already handles this for the case where the transfer function doesn't produce poison, e.g. %iv.next = add %iv, 1. However, this does not work if nowrap flags are present, e.g. the very common %iv.next = add nuw %iv, 1 case.

This patch adds a fold that pushes freeze instructions to the start value by checking whether all backedge values will be non-poison after poison generating flags have been dropped. This allows pushing freezes out of loops in most cases. I suspect that this also obsoletes the CanonicalizeFreezeInLoops pass, and we can probably drop it.

Fixes https://github.com/llvm/llvm-project/issues/56048.

[nlopes]

llvm/test/Transforms/InstCombine/freeze.ll

746

This case is annoying because if undef didn't exist*, we could get rid of the freeze here altogether.
The reason is that if %init=poison, then %init.fr becomes any value we want.
Since the add below is nuw, we can pick a convenient value that makes the addition overflow and thus poison.
Thus, if %init is poison, %cond becomes poison, and the loop triggers UB.

• The argument still applies to full undef values. But it does not for values that are partially undef (e.g., (undef & 1) | (x & ~1)). We've never enabled partial undefs in Alive2 because they are quite expensive and we only had some artificial examples on when they are needed. But this example shows it well.

We can decide to ignore partial undefs here, or at least leave a FIXME for that day when we remove undef from LLVM 🤣

[nikic]

llvm/test/Transforms/InstCombine/freeze.ll

746

I think there are two aspects here: The first is the dropped nowrap flags, and the other the freeze on the start value.

If we drop support for undef, then we'll automatically start preserving nowrap flags here (I checked that this happens if we replace isGuaranteedNotToBeUndefOrPoison with isGuaranteedNotToBePoison, which can exploit poison propagation + branch on poison UB).

The argument you make that would allow dropping the freeze entirely is a lot more subtle, and depends on the exact instructions being used. I think it would not work if the increment is a variable, for example (because it could be zero, in which case there might not be a choice of frozen value that results in a branch on poison). I tried to demonstrate this with https://alive2.llvm.org/ce/z/w9eiFn, but alive happily accepts this. I would have thought that init = poison, step = 0 is an obvious counter-example, where tgt causes branch on poison, while src is well-defined.

[nlopes]

LGTM

llvm/test/Transforms/InstCombine/freeze.ll

746

That stretches a bit the capabilities of Alive2 regarding loops. FWIW, an increment of zero gives an infinite loop, so we may be fine.
Anyway, let's leave that discussion for when undef is gone :)

[aqjune]

I suspect that this also obsoletes the CanonicalizeFreezeInLoops pass, and we can probably drop it.

ATM CanonicalizeFreezeInLoops targets patterns that are weakly related but different from this patch:
(1) It deals with freeze(add nsw (phi, 1)). (foldFreezeIntoRecurrence is called only for freeze(phi) in this patch)
(2) Freeze instructions that are not finally fed into the branch condition are hoisted as well - this is to allow successful SCEV analysis for all variables inside a loop
So immediately dropping it might not be possible.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3803

Not sure whether this can happen; when does this condition meet?

[nlopes]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3803

I was also wondering about that, but I guess it may happen with malformed BBs (e.g., with just PHIs but no terminator). Unreachable BBs are allowed to be malformed in LLVM IR (which is a bit unfortunate).

[aqjune]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3803

Oh okay, it makes sense to me.

[nikic]

Handle invoke terminator, remove unnecessary catchswitch handling.

[nikic]

In D127960#3590870, @aqjune wrote:

I suspect that this also obsoletes the CanonicalizeFreezeInLoops pass, and we can probably drop it.

ATM CanonicalizeFreezeInLoops targets patterns that are weakly related but different from this patch:
(1) It deals with freeze(add nsw (phi, 1)). (foldFreezeIntoRecurrence is called only for freeze(phi) in this patch)

InstCombine handles this case as well: pushFreezeToPreventPoisonFromPropagating() will convert freeze(add nsw (phi, 1)) into add(freeze(phi), 1), and from there we will push the freeze out of the loop.

(2) Freeze instructions that are not finally fed into the branch condition are hoisted as well - this is to allow successful SCEV analysis for all variables inside a loop

This code doesn't depend on the branch condition either, so I think it should be fine. I'll run the CanonicalizeFreezeInLoops tests through InstCombine to check whether we're missing any cases.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3803

Unreachable code still requires basic blocks to have terminators. What this was handling is a special case where no instructions can be inserted into a block for the case where an EH pad is also a terminator. This happens with catchswitch block pads, which can only contain phis and catchswitch. However, after trying to construct a test case for this, I don't think this is relevant for this code, because the successor of catchswitch can't be a loop (it would be another EH pad). So I've dropped this.

Instead I've added a check for invoke terminators, which we do need to handle (@fold_phi_invoke_start_value).

[nikic]

Okay, I just checked the CanonicalizeFreezeInLoops tests, and there is one case that it handles but InstCombine doesn't: https://github.com/llvm/llvm-project/blob/c6b88cb9184fc6c17bb3d9c2c9568646dc46638d/llvm/test/Transforms/CanonicalizeFreezeInLoops/onephi.ll#L361 This is the case where the step value needs to be frozen.

[aqjune]

Thank you for confirming it!

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

public:
   Instruction *visitInsertElementInst(InsertElementInst &IE);
   Instruction *visitExtractElementInst(ExtractElementInst &EI);
   Instruction *visitShuffleVectorInst(ShuffleVectorInst &SVI);
   Instruction *visitExtractValueInst(ExtractValueInst &EV);
   Instruction *visitLandingPadInst(LandingPadInst &LI);
   Instruction *visitVAEndInst(VAEndInst &I);
   Value *pushFreezeToPreventPoisonFromPropagating(FreezeInst &FI);
   bool freezeOtherUses(FreezeInst &FI);
+  Instruction *foldFreezeIntoRecurrence(FreezeInst &I, PHINode *PN);
   Instruction *visitFreeze(FreezeInst &I);
 
   /// Specify what to return for unhandled instructions.
   Instruction *visitInstruction(Instruction &I) { return nullptr; }
 
   /// True when DB dominates all uses of DI except UI.
   /// UI must be in the same block as DI.
   /// The routine checks that the DI parent and DB are different.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

InstCombinerImpl::pushFreezeToPreventPoisonFromPropagating(FreezeInst &OrigFI) {
   Builder.SetInsertPoint(OrigOpInst);
   auto *FrozenMaybePoisonOperand = Builder.CreateFreeze(
       MaybePoisonOperand->get(), MaybePoisonOperand->get()->getName() + ".fr");
 
   replaceUse(*MaybePoisonOperand, FrozenMaybePoisonOperand);
   return OrigOp;
 }
 
+Instruction *InstCombinerImpl::foldFreezeIntoRecurrence(FreezeInst &FI,
+                                                        PHINode *PN) {
+  // Detect whether this is a recurrence with a start value and some number of
+  // backedge values. We'll check whether we can push the freeze through the
+  // backedge values (possibly dropping poison flags along the way) until we
+  // reach the phi again. In that case, we can move the freeze to the start
+  // value.
+  Use *StartU = nullptr;
+  SmallVector<Value *> Worklist;
+  for (Use &U : PN->incoming_values()) {
+    if (DT.dominates(PN->getParent(), PN->getIncomingBlock(U))) {
+      // Add backedge value to worklist.
+      Worklist.push_back(U.get());
+      continue;
+    }
+
+    // Don't bother handling multiple start values.
+    if (StartU)
+      return nullptr;
+    StartU = &U;
+  }
+
+  if (!StartU || Worklist.empty())
+    return nullptr; // Not a recurrence.
+
+  BasicBlock *StartBB = PN->getIncomingBlock(*StartU);
+  if (StartBB->getFirstInsertionPt() == StartBB->end())
+    return nullptr; // Cannot insert into this block.
+
+  SmallPtrSet<Value *, 32> Visited;
+  SmallVector<Instruction *> DropFlags;
+  while (!Worklist.empty()) {
+    Value *V = Worklist.pop_back_val();
+    if (!Visited.insert(V).second)
+      continue;
+
+    if (Visited.size() > 32)
+      return nullptr; // Limit the total number of values we inspect.
+
+    // Assume that PN is non-poison, because it will be after the transform.
+    if (V == PN || isGuaranteedNotToBeUndefOrPoison(V))
+      continue;
+
+    Instruction *I = dyn_cast<Instruction>(V);
+    if (!I || canCreateUndefOrPoison(cast<Operator>(I),
+                                     /*ConsiderFlags*/ false))
+      return nullptr;
+
+    DropFlags.push_back(I);
+    append_range(Worklist, I->operands());
+  }
+
+  for (Instruction *I : DropFlags)
+    I->dropPoisonGeneratingFlags();
+
+  Builder.SetInsertPoint(StartBB->getTerminator());
+  Value *FrozenStartV = Builder.CreateFreeze(StartU->get(),
+                                             StartU->get()->getName() + ".fr");
+  replaceUse(*StartU, FrozenStartV);
+  return replaceInstUsesWith(FI, PN);
+}
+
 bool InstCombinerImpl::freezeOtherUses(FreezeInst &FI) {
   Value *Op = FI.getOperand(0);
 
   if (isa<Constant>(Op) || Op->hasOneUse())
     return false;
 
   // Move the freeze directly after the definition of its operand, so that
   // it dominates the maximum number of uses. Note that it may not dominate
Instruction *InstCombinerImpl::visitFreeze(FreezeInst &I) {
 
   if (Value *V = simplifyFreezeInst(Op0, SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 
   // freeze (phi const, x) --> phi const, (freeze x)
   if (auto *PN = dyn_cast<PHINode>(Op0)) {
     if (Instruction *NV = foldOpIntoPhi(I, PN))
       return NV;
+    if (Instruction *NV = foldFreezeIntoRecurrence(I, PN))
+      return NV;
   }
 
   if (Value *NI = pushFreezeToPreventPoisonFromPropagating(I))
     return replaceInstUsesWith(I, NI);
 
   // If I is freeze(undef), check its uses and fold it to a fixed constant.
   // - or: pick -1
   // - select's condition: if the true value is constant, choose it by making

llvm/test/Transforms/InstCombine/freeze.ll

loop:                                             ; preds = %loop, %entry
   %i.next = add i32 %i.fr, 1
   %cond = icmp eq i32 %i.next, %n
   br i1 %cond, label %loop, label %exit
 
 exit:                                             ; preds = %loop
   ret void
 }
 
-define void @fold_phi_neg_flags(i32 %init, i32 %n) {
-; CHECK-LABEL: @fold_phi_neg_flags(
+define void @fold_phi_drop_flags(i32 %init, i32 %n) {
+; CHECK-LABEL: @fold_phi_drop_flags(
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[INIT_FR:%.*]] = freeze i32 [[INIT:%.*]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[I:%.*]] = phi i32 [ [[INIT:%.*]], [[ENTRY:%.*]] ], [ [[I_NEXT:%.*]], [[LOOP]] ]
-; CHECK-NEXT:    [[I_FR:%.*]] = freeze i32 [[I]]
-; CHECK-NEXT:    [[I_NEXT]] = add nuw nsw i32 [[I_FR]], 1
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ [[INIT_FR]], [[ENTRY:%.*]] ], [ [[I_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[I_NEXT]] = add i32 [[I]], 1
 ; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[I_NEXT]], [[N:%.*]]
 ; CHECK-NEXT:    br i1 [[COND]], label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %loop
 

 
 exit:                                             ; preds = %loop
   ret void
 }
 
 define void @fold_phi_multiple_insts(i32 %init, i32 %n) {
 ; CHECK-LABEL: @fold_phi_multiple_insts(
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[INIT_FR:%.*]] = freeze i32 [[INIT:%.*]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[I:%.*]] = phi i32 [ [[INIT:%.*]], [[ENTRY:%.*]] ], [ [[I_NEXT:%.*]], [[LOOP]] ]
-; CHECK-NEXT:    [[I_FR:%.*]] = freeze i32 [[I]]
-; CHECK-NEXT:    [[I_SQ:%.*]] = mul nuw nsw i32 [[I_FR]], [[I_FR]]
-; CHECK-NEXT:    [[I_NEXT]] = add nuw nsw i32 [[I_SQ]], 1
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ [[INIT_FR]], [[ENTRY:%.*]] ], [ [[I_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[I_SQ:%.*]] = mul i32 [[I]], [[I]]
+; CHECK-NEXT:    [[I_NEXT]] = add i32 [[I_SQ]], 1
 ; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[I_NEXT]], [[N:%.*]]
 ; CHECK-NEXT:    br i1 [[COND]], label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %loop
 
 loop:                                             ; preds = %loop, %entry
   %i = phi i32 [ %init, %entry ], [ %i.next, %loop ]
   %i.fr = freeze i32 %i
   %i.sq = mul nsw nuw i32 %i.fr, %i.fr
   %i.next = add nsw nuw i32 %i.sq, 1
   %cond = icmp eq i32 %i.next, %n
   br i1 %cond, label %loop, label %exit
 
 exit:                                             ; preds = %loop
   ret void
 }
 
 define void @fold_phi_multiple_back_edges(i32 %init, i32 %n) {
 ; CHECK-LABEL: @fold_phi_multiple_back_edges(
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[INIT_FR:%.*]] = freeze i32 [[INIT:%.*]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[I:%.*]] = phi i32 [ [[INIT:%.*]], [[ENTRY:%.*]] ], [ [[I_NEXT:%.*]], [[LOOP]] ], [ [[I_NEXT2:%.*]], [[LOOP_LATCH2:%.*]] ]
-; CHECK-NEXT:    [[I_FR:%.*]] = freeze i32 [[I]]
-; CHECK-NEXT:    [[I_NEXT]] = add nuw nsw i32 [[I_FR]], 1
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ [[INIT_FR]], [[ENTRY:%.*]] ], [ [[I_NEXT:%.*]], [[LOOP]] ], [ [[I_NEXT2:%.*]], [[LOOP_LATCH2:%.*]] ]
+; CHECK-NEXT:    [[I_NEXT]] = add i32 [[I]], 1
 ; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[I_NEXT]], [[N:%.*]]
 ; CHECK-NEXT:    br i1 [[COND]], label [[LOOP]], label [[LOOP_LATCH2]]
 ; CHECK:       loop.latch2:
-; CHECK-NEXT:    [[I_NEXT2]] = add nuw nsw i32 [[I_FR]], 2
+; CHECK-NEXT:    [[I_NEXT2]] = add i32 [[I]], 2
 ; CHECK-NEXT:    br i1 false, label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %loop
 
 loop:                                             ; preds = %loop, %entry