[InstCombine] Merge foldFreezeIntoRecurrence into pushFreezeToPreventPoisonFromPropagating

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -5039,23 +5039,69 @@ InstCombinerImpl::pushFreezeToPreventPoisonFromPropagating(FreezeInst &OrigFI) {
                                    /*ConsiderFlagsAndMetadata*/ false);
   };
 
+  Use *StartUNeedsFreeze = nullptr;
+  auto CanPushFreezeThruPHI = [&](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;
+    bool HasBackedge = false;
+    for (Use &U : PN->incoming_values()) {
+      if (DT.dominates(PN->getParent(), PN->getIncomingBlock(U))) {
+        HasBackedge = true;
+        continue;
+      }
+
+      // Don't bother handling multiple start values.
+      if (StartU)
+        return false;
+      StartU = &U;
+    }
+
+    if (!StartU || !HasBackedge)
+      return false; // 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 false;
+
+    if (StartNeedsFreeze)
+      StartUNeedsFreeze = StartU;
+
+    return true;
+  };
+
   // Pushing freezes up long instruction chains can be expensive. Instead,
   // we directly push the freeze all the way to the leaves. However, we leave
   // deduplication of freezes on the same value for freezeOtherUses().
   Use *OrigUse = &OrigFI.getOperandUse(0);
+  auto *PN = dyn_cast<PHINode>(OrigUse->get());
+  if (PN && !CanPushFreezeThruPHI(PN))
+    return nullptr;
+
   SmallPtrSet<Instruction *, 8> Visited;
   SmallVector<Use *, 8> Worklist;
   Worklist.push_back(OrigUse);
   while (!Worklist.empty()) {
     auto *U = Worklist.pop_back_val();
     Value *V = U->get();
-    if (!CanPushFreeze(V)) {
+
+    if (!CanPushFreeze(V) && !isa<PHINode>(V)) {
       // If we can't push through the original instruction, abort the transform.
       if (U == OrigUse)
         return nullptr;
 
       auto *UserI = cast<Instruction>(U->getUser());
-      Builder.SetInsertPoint(UserI);
+      if (auto *UserPN = dyn_cast<PHINode>(UserI))
+        Builder.SetInsertPoint(UserPN->getIncomingBlock(*U)->getTerminator());
+      else
+        Builder.SetInsertPoint(UserI);
       Value *Frozen = Builder.CreateFreeze(V, V->getName() + ".fr");
       U->set(Frozen);
       continue;
@@ -5077,75 +5123,16 @@ InstCombinerImpl::pushFreezeToPreventPoisonFromPropagating(FreezeInst &OrigFI) {
     this->Worklist.add(I);
   }
 
-  return OrigUse->get();
-}
-
-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.
-
-  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()) {
-    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),
-                                     /*ConsiderFlagsAndMetadata*/ false))
-      return nullptr;
-
-    DropFlags.push_back(I);
-    append_range(Worklist, I->operands());
+  if (StartUNeedsFreeze) {
+    Value *StartV = StartUNeedsFreeze->get();
+    Builder.SetInsertPoint(
+        PN->getIncomingBlock(*StartUNeedsFreeze)->getTerminator());
+    Value *FrozenStartV =
+        Builder.CreateFreeze(StartV, StartV->getName() + ".fr");
+    replaceUse(*StartUNeedsFreeze, FrozenStartV);
   }
 
-  for (Instruction *I : DropFlags)
-    I->dropPoisonGeneratingAnnotations();
-
-  if (StartNeedsFreeze) {
-    Builder.SetInsertPoint(StartBB->getTerminator());
-    Value *FrozenStartV = Builder.CreateFreeze(StartV,
-                                               StartV->getName() + ".fr");
-    replaceUse(*StartU, FrozenStartV);
-  }
-  return replaceInstUsesWith(FI, PN);
+  return OrigUse->get();
 }
 
 bool InstCombinerImpl::freezeOtherUses(FreezeInst &FI) {
@@ -5209,8 +5196,6 @@ Instruction *InstCombinerImpl::visitFreeze(FreezeInst &I) {
   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))

llvm/test/Transforms/InstCombine/freeze.ll

@@ -273,13 +273,13 @@ define void @freeze_dominated_uses_catchswitch(i1 %c, i32 %x) personality ptr @_
 ; CHECK-NEXT:            to label %[[CLEANUP]] unwind label %[[CATCH_DISPATCH]]
 ; CHECK:       [[CATCH_DISPATCH]]:
 ; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ 0, %[[IF_THEN]] ], [ [[X]], %[[IF_ELSE]] ]
-; CHECK-NEXT:    [[CS:%.*]] = catchswitch within none [label %[[CATCH:.*]], label %catch2] unwind to caller
+; CHECK-NEXT:    [[CS:%.*]] = catchswitch within none [label %[[CATCH:.*]], label %[[CATCH2:.*]]] unwind to caller
 ; CHECK:       [[CATCH]]:
 ; CHECK-NEXT:    [[CP:%.*]] = catchpad within [[CS]] [ptr null, i32 64, ptr null]
 ; CHECK-NEXT:    [[PHI_FREEZE:%.*]] = freeze i32 [[PHI]]
 ; CHECK-NEXT:    call void @use_i32(i32 [[PHI_FREEZE]]) [ "funclet"(token [[CP]]) ]
 ; CHECK-NEXT:    unreachable
-; CHECK:       [[CATCH2:.*:]]
+; CHECK:       [[CATCH2]]:
 ; CHECK-NEXT:    [[CP2:%.*]] = catchpad within [[CS]] [ptr null, i32 64, ptr null]
 ; CHECK-NEXT:    call void @use_i32(i32 [[PHI]]) [ "funclet"(token [[CP2]]) ]
 ; CHECK-NEXT:    unreachable
@@ -944,10 +944,10 @@ define void @fold_phi_gep_phi_offset(ptr %init, ptr %end, i64 noundef %n) {
 ; CHECK-LABEL: define void @fold_phi_gep_phi_offset(
 ; CHECK-SAME: ptr [[INIT:%.*]], ptr [[END:%.*]], i64 noundef [[N:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*]]:
-; CHECK-NEXT:    [[TMP0:%.*]] = freeze ptr [[INIT]]
+; CHECK-NEXT:    [[INIT_FR:%.*]] = freeze ptr [[INIT]]
 ; CHECK-NEXT:    br label %[[LOOP:.*]]
 ; CHECK:       [[LOOP]]:
-; CHECK-NEXT:    [[I:%.*]] = phi ptr [ [[TMP0]], %[[ENTRY]] ], [ [[I_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[I:%.*]] = phi ptr [ [[INIT_FR]], %[[ENTRY]] ], [ [[I_NEXT:%.*]], %[[LOOP]] ]
 ; CHECK-NEXT:    [[OFF:%.*]] = phi i64 [ [[N]], %[[ENTRY]] ], [ [[OFF_NEXT:%.*]], %[[LOOP]] ]
 ; CHECK-NEXT:    [[OFF_NEXT]] = shl i64 [[OFF]], 3
 ; CHECK-NEXT:    [[I_NEXT]] = getelementptr i8, ptr [[I]], i64 [[OFF_NEXT]]
@@ -978,10 +978,10 @@ define void @fold_phi_gep_inbounds_phi_offset(ptr %init, ptr %end, i64 noundef %
 ; CHECK-LABEL: define void @fold_phi_gep_inbounds_phi_offset(
 ; CHECK-SAME: ptr [[INIT:%.*]], ptr [[END:%.*]], i64 noundef [[N:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*]]:
-; CHECK-NEXT:    [[TMP0:%.*]] = freeze ptr [[INIT]]
+; CHECK-NEXT:    [[INIT_FR:%.*]] = freeze ptr [[INIT]]
 ; CHECK-NEXT:    br label %[[LOOP:.*]]
 ; CHECK:       [[LOOP]]:
-; CHECK-NEXT:    [[I:%.*]] = phi ptr [ [[TMP0]], %[[ENTRY]] ], [ [[I_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[I:%.*]] = phi ptr [ [[INIT_FR]], %[[ENTRY]] ], [ [[I_NEXT:%.*]], %[[LOOP]] ]
 ; CHECK-NEXT:    [[OFF:%.*]] = phi i64 [ [[N]], %[[ENTRY]] ], [ [[OFF_NEXT:%.*]], %[[LOOP]] ]
 ; CHECK-NEXT:    [[OFF_NEXT]] = shl i64 [[OFF]], 3
 ; CHECK-NEXT:    [[I_NEXT]] = getelementptr i8, ptr [[I]], i64 [[OFF_NEXT]]
@@ -1011,12 +1011,12 @@ define void @fold_phi_gep_phi_offset_multiple(ptr %init, ptr %end, i64 %n) {
 ; CHECK-LABEL: define void @fold_phi_gep_phi_offset_multiple(
 ; CHECK-SAME: ptr [[INIT:%.*]], ptr [[END:%.*]], i64 [[N:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*]]:
-; CHECK-NEXT:    [[TMP0:%.*]] = freeze ptr [[INIT]]
-; CHECK-NEXT:    [[TMP1:%.*]] = freeze i64 [[N]]
+; CHECK-NEXT:    [[INIT_FR:%.*]] = freeze ptr [[INIT]]
+; CHECK-NEXT:    [[N_FR:%.*]] = freeze i64 [[N]]
 ; CHECK-NEXT:    br label %[[LOOP:.*]]
 ; CHECK:       [[LOOP]]:
-; CHECK-NEXT:    [[I:%.*]] = phi ptr [ [[TMP0]], %[[ENTRY]] ], [ [[I_NEXT:%.*]], %[[LOOP]] ]
-; CHECK-NEXT:    [[OFF:%.*]] = phi i64 [ [[TMP1]], %[[ENTRY]] ], [ [[OFF_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[I:%.*]] = phi ptr [ [[INIT_FR]], %[[ENTRY]] ], [ [[I_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[OFF:%.*]] = phi i64 [ [[N_FR]], %[[ENTRY]] ], [ [[OFF_NEXT:%.*]], %[[LOOP]] ]
 ; CHECK-NEXT:    [[OFF_NEXT]] = shl i64 [[OFF]], 3
 ; CHECK-NEXT:    [[I_NEXT]] = getelementptr i8, ptr [[I]], i64 [[OFF_NEXT]]
 ; CHECK-NEXT:    [[COND:%.*]] = icmp eq ptr [[I_NEXT]], [[END]]
@@ -1189,6 +1189,64 @@ exit:
   ret void
 }
 
+declare ptr @get_ptr()
+
+; When the phi input comes from an invoke, we need to be careful the freeze
+; isn't pushed after the invoke.
+define void @fold_phi_noundef_start_value_with_invoke(ptr noundef %init, i1 %cond.0, i1 %cond.1) personality ptr undef {
+; CHECK-LABEL: define void @fold_phi_noundef_start_value_with_invoke(
+; CHECK-SAME: ptr noundef [[INIT:%.*]], i1 [[COND_0:%.*]], i1 [[COND_1:%.*]]) personality ptr undef {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    [[IV_0:%.*]] = phi ptr [ [[INIT]], %[[ENTRY]] ], [ [[IV_0_NEXT:%.*]], %[[LOOP_LATCH:.*]] ]
+; CHECK-NEXT:    br i1 [[COND_0]], label %[[LOOP_LATCH]], label %[[IF_ELSE:.*]]
+; CHECK:       [[IF_ELSE]]:
+; CHECK-NEXT:    [[IV_1:%.*]] = invoke ptr @get_ptr()
+; CHECK-NEXT:            to label %[[LOOP_LATCH]] unwind label %[[UNWIND:.*]]
+; CHECK:       [[LOOP_LATCH]]:
+; CHECK-NEXT:    [[IV_2:%.*]] = phi ptr [ [[IV_0]], %[[LOOP]] ], [ [[IV_1]], %[[IF_ELSE]] ]
+; CHECK-NEXT:    [[IV_2_FR:%.*]] = freeze ptr [[IV_2]]
+; CHECK-NEXT:    [[IV_2_FR_INT:%.*]] = ptrtoint ptr [[IV_2_FR]] to i64
+; CHECK-NEXT:    [[IV_0_INT:%.*]] = ptrtoint ptr [[IV_0]] to i64
+; CHECK-NEXT:    [[IDX:%.*]] = sub i64 [[IV_0_INT]], [[IV_2_FR_INT]]
+; CHECK-NEXT:    [[IV_0_NEXT]] = getelementptr i8, ptr [[IV_0]], i64 [[IDX]]
+; CHECK-NEXT:    br i1 [[COND_1]], label %[[EXIT:.*]], label %[[LOOP]]
+; CHECK:       [[UNWIND]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = landingpad i8
+; CHECK-NEXT:            cleanup
+; CHECK-NEXT:    unreachable
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %iv.0 = phi ptr [ %init, %entry ], [ %iv.0.next, %loop.latch ]
+  br i1 %cond.0, label %loop.latch, label %if.else
+
+if.else:
+  %iv.1 = invoke ptr @get_ptr()
+  to label %loop.latch unwind label %unwind
+
+loop.latch:
+  %iv.2 = phi ptr [ %iv.0, %loop ], [ %iv.1, %if.else ]
+  %iv.2.fr = freeze ptr %iv.2
+  %iv.2.fr.int = ptrtoint ptr %iv.2.fr to i64
+  %iv.0.int = ptrtoint ptr %iv.0 to i64
+  %idx = sub i64 %iv.0.int, %iv.2.fr.int
+  %iv.0.next = getelementptr i8, ptr %iv.0, i64 %idx
+  br i1 %cond.1, label %exit, label %loop
+
+unwind:
+  landingpad i8 cleanup
+  unreachable
+
+exit:
+  ret void
+}
+
 define void @fold_phi_invoke_start_value(i32 %n) personality ptr undef {
 ; CHECK-LABEL: define void @fold_phi_invoke_start_value(
 ; CHECK-SAME: i32 [[N:%.*]]) personality ptr undef {