[InstCombine] Allow freezing multiple out-of-loop values

[c-rhodes]

Extend foldFreezeIntoRecurrence to allow freezing multiple out-of-loop values. This is following on from #154336, which recently made the same change for a wider set of ops.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Cullen Rhodes (c-rhodes)

Changes

Extend foldFreezeIntoRecurrence to allow freezing multiple out-of-loop values. This is following on from #154336, which recently made the same change for a wider set of ops.

---

Full diff: https://github.com/llvm/llvm-project/pull/155638.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstructionCombining.cpp (+14-16)
• (modified) llvm/test/Transforms/InstCombine/freeze.ll (+4-3)

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 1c512ec1e21bb..e524839e2876e 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -5032,32 +5032,28 @@ Instruction *InstCombinerImpl::foldFreezeIntoRecurrence(FreezeInst &FI,
   // 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<Use *> StartUses;
   SmallVector<Value *> Worklist;
   for (Use &U : PN->incoming_values()) {
-    if (DT.dominates(PN->getParent(), PN->getIncomingBlock(U))) {
+    BasicBlock *StartBB = PN->getIncomingBlock(U);
+    Value *StartV = U.get();
+    if (DT.dominates(PN->getParent(), StartBB)) {
       // Add backedge value to worklist.
-      Worklist.push_back(U.get());
+      Worklist.push_back(StartV);
       continue;
     }
 
-    // Don't bother handling multiple start values.
-    if (StartU)
+    bool StartNeedsFreeze = !isGuaranteedNotToBeUndefOrPoison(StartV);
+    // We can't insert freeze if a start value is the result of the
+    // terminator (e.g. an invoke).
+    if (StartNeedsFreeze && StartBB->getTerminator() == StartV)
       return nullptr;
-    StartU = &U;
+    StartUses.push_back(&U);
   }
 
-  if (!StartU || Worklist.empty())
+  if (StartUses.empty() || Worklist.empty())
     return nullptr; // Not a recurrence.
 
-  Value *StartV = StartU->get();
-  BasicBlock *StartBB = PN->getIncomingBlock(*StartU);
-  bool StartNeedsFreeze = !isGuaranteedNotToBeUndefOrPoison(StartV);
-  // We can't insert freeze if the start value is the result of the
-  // terminator (e.g. an invoke).
-  if (StartNeedsFreeze && StartBB->getTerminator() == StartV)
-    return nullptr;
-
   SmallPtrSet<Value *, 32> Visited;
   SmallVector<Instruction *> DropFlags;
   while (!Worklist.empty()) {
@@ -5084,7 +5080,9 @@ Instruction *InstCombinerImpl::foldFreezeIntoRecurrence(FreezeInst &FI,
   for (Instruction *I : DropFlags)
     I->dropPoisonGeneratingAnnotations();
 
-  if (StartNeedsFreeze) {
+  for (Use *StartU : StartUses) {
+    Value *StartV = StartU->get();
+    BasicBlock *StartBB = PN->getIncomingBlock(*StartU);
     Builder.SetInsertPoint(StartBB->getTerminator());
     Value *FrozenStartV = Builder.CreateFreeze(StartV,
                                                StartV->getName() + ".fr");
diff --git a/llvm/test/Transforms/InstCombine/freeze.ll b/llvm/test/Transforms/InstCombine/freeze.ll
index b29421a655fa8..09fb2dc19912c 100644
--- a/llvm/test/Transforms/InstCombine/freeze.ll
+++ b/llvm/test/Transforms/InstCombine/freeze.ll
@@ -1106,13 +1106,14 @@ define void @fold_phi_multiple_start_values(i1 %c, i32 %init, i32 %init2, i32 %n
 ; CHECK-LABEL: define void @fold_phi_multiple_start_values(
 ; CHECK-SAME: i1 [[C:%.*]], i32 [[INIT:%.*]], i32 [[INIT2:%.*]], i32 [[N:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[INIT_FR:%.*]] = freeze i32 [[INIT]]
 ; CHECK-NEXT:    br i1 [[C]], label %[[IF:.*]], label %[[LOOP:.*]]
 ; CHECK:       [[IF]]:
+; CHECK-NEXT:    [[INIT2_FR:%.*]] = freeze i32 [[INIT2]]
 ; CHECK-NEXT:    br label %[[LOOP]]
 ; CHECK:       [[LOOP]]:
-; CHECK-NEXT:    [[I:%.*]] = phi i32 [ [[INIT]], %[[ENTRY]] ], [ [[INIT2]], %[[IF]] ], [ [[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]] ], [ [[INIT2_FR]], %[[IF]] ], [ [[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]]:

[c-rhodes]

@nikic I had a look at implementing this for foldFreezeIntoRecurrence as you mentioned on #154336 (comment). There's an existing test for this, with this patch we'd be able to remove freeze altogether here: https://godbolt.org/z/8qsf3zcWG

[nikic]

Can you please add a test where the two start values come from the same predecessor (e.g. dummy br with both successors equal)? You need to be careful about that because the phi needs to have identical values for the same predecessor.

By the way, this change isn't actually what I had in mind... I'd expect multiple starting values to get canonicalized away by loop simplify form. I think the more interesting extension here is the case where you need to freeze both the start value and the step value. So something like {start,+,step} where both values are not known non-poison.

[c-rhodes]

Can you please add a test where the two start values come from the same predecessor (e.g. dummy br with both successors equal)? You need to be careful about that because the phi needs to have identical values for the same predecessor.

Perhaps I'm misunderstanding, but that's not legal IR? https://godbolt.org/z/8MYM1Yf5b

By the way, this change isn't actually what I had in mind... I'd expect multiple starting values to get canonicalized away by loop simplify form. I think the more interesting extension here is the case where you need to freeze both the start value and the step value. So something like {start,+,step} where both values are not known non-poison.

ah ok, do you have an example test for this? Sounds like

llvm-project/llvm/test/Transforms/InstCombine/freeze.ll

Line 855 in 33b2c26

define void @fold_phi_drop_flags(i32 %init, i32 %n) {

but the flags get dropped there