[InstCombine] Treat identical operands as one in pushFreezeToPreventPoisonFromPropagating (#145348)

To push a freeze through an instruction, only one operand may produce
poison. However, this currently fails for identical operands which are
treated as separate. This patch fixes this by treating them as a single
operand.

Author: c-rhodes

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -4899,13 +4899,14 @@ InstCombinerImpl::pushFreezeToPreventPoisonFromPropagating(FreezeInst &OrigFI) {
   // If operand is guaranteed not to be poison, there is no need to add freeze
   // to the operand. So we first find the operand that is not guaranteed to be
   // poison.
-  Use *MaybePoisonOperand = nullptr;
-  for (Use &U : OrigOpInst->operands()) {
-    if (isa<MetadataAsValue>(U.get()) ||
-        isGuaranteedNotToBeUndefOrPoison(U.get()))
+  Value *MaybePoisonOperand = nullptr;
+  for (Value *V : OrigOpInst->operands()) {
+    if (isa<MetadataAsValue>(V) || isGuaranteedNotToBeUndefOrPoison(V) ||
+        // Treat identical operands as a single operand.
+        (MaybePoisonOperand && MaybePoisonOperand == V))
       continue;
     if (!MaybePoisonOperand)
-      MaybePoisonOperand = &U;
+      MaybePoisonOperand = V;
     else
       return nullptr;
   }
@@ -4917,10 +4918,10 @@ InstCombinerImpl::pushFreezeToPreventPoisonFromPropagating(FreezeInst &OrigFI) {
     return OrigOp;
 
   Builder.SetInsertPoint(OrigOpInst);
-  auto *FrozenMaybePoisonOperand = Builder.CreateFreeze(
-      MaybePoisonOperand->get(), MaybePoisonOperand->get()->getName() + ".fr");
+  Value *FrozenMaybePoisonOperand = Builder.CreateFreeze(
+      MaybePoisonOperand, MaybePoisonOperand->getName() + ".fr");
 
-  replaceUse(*MaybePoisonOperand, FrozenMaybePoisonOperand);
+  OrigOpInst->replaceUsesOfWith(MaybePoisonOperand, FrozenMaybePoisonOperand);
   return OrigOp;
 }
 

llvm/test/Transforms/InstCombine/freeze.ll

@@ -142,6 +142,17 @@ define i32 @early_freeze_test3(i32 %v1) {
   ret i32 %v4.fr
 }
 
+define i32 @early_freeze_test4(i32 %v1) {
+; CHECK-LABEL: @early_freeze_test4(
+; CHECK-NEXT:    [[V2_FR:%.*]] = freeze i32 [[V2:%.*]]
+; CHECK-NEXT:    [[V3:%.*]] = mul i32 [[V2_FR]], [[V2_FR]]
+; CHECK-NEXT:    ret i32 [[V3]]
+;
+  %v2 = mul i32 %v1, %v1
+  %v2.fr = freeze i32 %v2
+  ret i32 %v2.fr
+}
+
 ; If replace all dominated uses of v to freeze(v).
 
 define void @freeze_dominated_uses_test1(i32 %v) {