[InstCombine] Freeze other uses of frozen value

[nikic]

Summary

If there is a freeze %x, we currently replace all other uses of %x with freeze %x -- as long as they are dominated by the freeze instruction. This patch extends this behavior to cases where we did not originally dominate the use by moving the freeze instruction to the nearest common dominator.

The motivation can be seen in test @combine_and_after_freezing_uses: Canonicalizing everything to freeze %x allows folds that are based on value identity (i.e. same operand occurring in two places) to trigger. This also covers the case from D125248.

[mtrofin]

could you also add the test from D125248 - can't hurt to have more.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3838–3843

wouldn't going to the def of Op be sufficient and compile-time cheaper?

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3838–3843

Yes, that would work as well. The reason I went for the nearest common dominator is to avoid hoisting the freeze more than necessary, e.g. out of a loop. Though now that I write that, LICM would do that anyway, so probably that concern doesn't make sense.

[nikic]

Insert the freeze directly after the definition of the op, rather than at the nearest common dominator of the uses.

Finding the correct insert position "directly after the def" is annoyingly tricky...

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

public:
   Instruction *visitInsertValueInst(InsertValueInst &IV);
   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 freezeDominatedUses(FreezeInst &FI);
+  bool freezeOtherUses(FreezeInst &FI);
   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) {
   auto *FrozenMaybePoisonOperand = new FreezeInst(
       MaybePoisonOperand->get(), MaybePoisonOperand->get()->getName() + ".fr");
 
   replaceUse(*MaybePoisonOperand, FrozenMaybePoisonOperand);
   FrozenMaybePoisonOperand->insertBefore(OrigOpInst);
   return OrigOp;
 }
 
-bool InstCombinerImpl::freezeDominatedUses(FreezeInst &FI) {
+bool InstCombinerImpl::freezeOtherUses(FreezeInst &FI) {
   Value *Op = FI.getOperand(0);
 
-  if (isa<Constant>(Op))
+  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
+  // *all* uses if the operand is an invoke/callbr and the use is in a phi on
+  // the normal/default destination. This is why the domination check in the
+  // replacement below is still necessary.
+  Instruction *MoveBefore = nullptr;
+  if (isa<Argument>(Op)) {
+    MoveBefore = &FI.getFunction()->getEntryBlock().front();
+    while (isa<AllocaInst>(MoveBefore))
+      MoveBefore = MoveBefore->getNextNode();
+  } else if (auto *PN = dyn_cast<PHINode>(Op)) {
+    MoveBefore = PN->getParent()->getFirstNonPHI();
+  } else if (auto *II = dyn_cast<InvokeInst>(Op)) {
+    MoveBefore = II->getNormalDest()->getFirstNonPHI();
+  } else if (auto *CB = dyn_cast<CallBrInst>(Op)) {
+    MoveBefore = CB->getDefaultDest()->getFirstNonPHI();
+  } else {
+    auto *I = cast<Instruction>(Op);
+    assert(!I->isTerminator() && "Cannot be a terminator");
+    MoveBefore = I->getNextNode();
+  }
+
   bool Changed = false;
+  if (&FI != MoveBefore) {
+    FI.moveBefore(MoveBefore);
+    Changed = true;
+  }
+
   Op->replaceUsesWithIf(&FI, [&](Use &U) -> bool {
     bool Dominates = DT.dominates(&FI, U);
     Changed |= Dominates;
     return Dominates;
   });
 
   return Changed;
 }
if (match(Op0, m_Undef()))
     return replaceInstUsesWith(I, getUndefReplacement(I.getType()));
 
   Constant *C;
   if (match(Op0, m_Constant(C)) && C->containsUndefOrPoisonElement()) {
     Constant *ReplaceC = getUndefReplacement(I.getType()->getScalarType());
     return replaceInstUsesWith(I, Constant::replaceUndefsWith(C, ReplaceC));
   }
 
-  // Replace all dominated uses of Op to freeze(Op).
-  if (freezeDominatedUses(I))
+  // Replace uses of Op with freeze(Op).
+  if (freezeOtherUses(I))
     return &I;
 
   return nullptr;
 }
 
 /// Check for case where the call writes to an otherwise dead alloca.  This
 /// shows up for unused out-params in idiomatic C/C++ code.   Note that this
 /// helper *only* analyzes the write; doesn't check any other legality aspect.

llvm/test/Transforms/InstCombine/freeze.ll

;
   call void @use_i32(i32 %v.fr)
   ret void
 }
 
 define void @freeze_dominated_uses_test2(i32 %v) {
 ; CHECK-LABEL: @freeze_dominated_uses_test2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[V_FR:%.*]] = freeze i32 [[V:%.*]]
 ; CHECK-NEXT:    call void @use_p32(i32* nonnull [[A]])
-; CHECK-NEXT:    call void @use_i32(i32 [[V:%.*]])
-; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[V]], 0
+; CHECK-NEXT:    call void @use_i32(i32 [[V_FR]])
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[V_FR]], 0
 ; CHECK-NEXT:    br i1 [[COND]], label [[BB0:%.*]], label [[BB1:%.*]]
 ; CHECK:       bb0:
-; CHECK-NEXT:    [[V_FR:%.*]] = freeze i32 [[V]]
 ; CHECK-NEXT:    call void @use_i32(i32 [[V_FR]])
 ; CHECK-NEXT:    call void @use_i32(i32 [[V_FR]])
 ; CHECK-NEXT:    br label [[END:%.*]]
 ; CHECK:       bb1:
-; CHECK-NEXT:    call void @use_i32(i32 [[V]])
+; CHECK-NEXT:    call void @use_i32(i32 [[V_FR]])
 ; CHECK-NEXT:    br label [[END]]
 ; CHECK:       end:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %a = alloca i32
   call void @use_p32(i32* %a)
   call void @use_i32(i32 %v)

 }
 
 declare i32 @get_i32()
 
 define i32 @freeze_use_in_different_branches(i1 %c) {
 ; CHECK-LABEL: @freeze_use_in_different_branches(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[X:%.*]] = call i32 @get_i32()
+; CHECK-NEXT:    [[FR:%.*]] = freeze i32 [[X]]
 ; CHECK-NEXT:    call void @use_i32(i32 0)
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF:%.*]], label [[ELSE:%.*]]
 ; CHECK:       if:
-; CHECK-NEXT:    call void @use_i32(i32 [[X]])
+; CHECK-NEXT:    call void @use_i32(i32 [[FR]])
 ; CHECK-NEXT:    ret i32 0
 ; CHECK:       else:
-; CHECK-NEXT:    call void @use_i32(i32 [[X]])
-; CHECK-NEXT:    [[FR:%.*]] = freeze i32 [[X]]
+; CHECK-NEXT:    call void @use_i32(i32 [[FR]])
 ; CHECK-NEXT:    call void @use_i32(i32 [[FR]])
 ; CHECK-NEXT:    ret i32 1
 ;
 entry:
   %x = call i32 @get_i32()
   call void @use_i32(i32 0)
   br i1 %c, label %if, label %else
 
 if:
   call void @use_i32(i32 %x)
   ret i32 0
 
 else:
   call void @use_i32(i32 %x)
   %fr = freeze i32 %x
   call void @use_i32(i32 %fr)
   ret i32 1
 }
 
 define i32 @freeze_phi_use(i1 %c) {
 ; CHECK-LABEL: @freeze_phi_use(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[X:%.*]] = call i32 @get_i32()
+; CHECK-NEXT:    [[FR:%.*]] = freeze i32 [[X]]
 ; CHECK-NEXT:    call void @use_i32(i32 0)
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF:%.*]], label [[JOIN:%.*]]
 ; CHECK:       if:
 ; CHECK-NEXT:    br label [[JOIN]]
 ; CHECK:       join:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ [[X]], [[IF]] ], [ 0, [[ENTRY:%.*]] ]
-; CHECK-NEXT:    [[FR:%.*]] = freeze i32 [[X]]
+; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ [[FR]], [[IF]] ], [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    call void @use_i32(i32 [[FR]])
 ; CHECK-NEXT:    ret i32 [[PHI]]
 ;
 entry:
   %x = call i32 @get_i32()
   call void @use_i32(i32 0)
   br i1 %c, label %if, label %join
 

 ; CHECK-LABEL: @freeze_phi_followed_by_phi(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF:%.*]], label [[JOIN:%.*]]
 ; CHECK:       if:
 ; CHECK-NEXT:    br label [[JOIN]]
 ; CHECK:       join:
 ; CHECK-NEXT:    [[X:%.*]] = phi i32 [ [[Y:%.*]], [[IF]] ], [ [[Z:%.*]], [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ [[Z]], [[IF]] ], [ [[Y]], [[ENTRY]] ]
-; CHECK-NEXT:    call void @use_i32(i32 [[X]])
 ; CHECK-NEXT:    [[FR:%.*]] = freeze i32 [[X]]
 ; CHECK-NEXT:    call void @use_i32(i32 [[FR]])
+; CHECK-NEXT:    call void @use_i32(i32 [[FR]])
 ; CHECK-NEXT:    ret i32 [[PHI]]
 ;
 entry:
   br i1 %c, label %if, label %join
 
 if:
   br label %join
 

 
 define i32 @freeze_invoke_use_after_phi(i1 %c) personality i8* undef {
 ; CHECK-LABEL: @freeze_invoke_use_after_phi(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[X:%.*]] = invoke i32 @get_i32()
 ; CHECK-NEXT:    to label [[INVOKE_CONT:%.*]] unwind label [[INVOKE_UNWIND:%.*]]
 ; CHECK:       invoke.cont:
 ; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ [[X]], [[ENTRY:%.*]] ], [ 0, [[INVOKE_CONT]] ]
-; CHECK-NEXT:    call void @use_i32(i32 [[X]])
 ; CHECK-NEXT:    [[FR:%.*]] = freeze i32 [[X]]
 ; CHECK-NEXT:    call void @use_i32(i32 [[FR]])
+; CHECK-NEXT:    call void @use_i32(i32 [[FR]])
 ; CHECK-NEXT:    call void @use_i32(i32 [[PHI]])
 ; CHECK-NEXT:    br label [[INVOKE_CONT]]
 ; CHECK:       invoke.unwind:
 ; CHECK-NEXT:    [[TMP0:%.*]] = landingpad i8
 ; CHECK-NEXT:    cleanup
 ; CHECK-NEXT:    unreachable
 ;
 entry:

 
 define i32 @freeze_callbr_use_after_phi(i1 %c) {
 ; CHECK-LABEL: @freeze_callbr_use_after_phi(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[X:%.*]] = callbr i32 asm sideeffect "", "=r"() #[[ATTR0:[0-9]+]]
 ; CHECK-NEXT:    to label [[CALLBR_CONT:%.*]] []
 ; CHECK:       callbr.cont:
 ; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ [[X]], [[ENTRY:%.*]] ], [ 0, [[CALLBR_CONT]] ]
-; CHECK-NEXT:    call void @use_i32(i32 [[X]])
 ; CHECK-NEXT:    [[FR:%.*]] = freeze i32 [[X]]
 ; CHECK-NEXT:    call void @use_i32(i32 [[FR]])
+; CHECK-NEXT:    call void @use_i32(i32 [[FR]])
 ; CHECK-NEXT:    call void @use_i32(i32 [[PHI]])
 ; CHECK-NEXT:    br label [[CALLBR_CONT]]
 ;
 entry:
   %x = callbr i32 asm sideeffect "", "=r"()
   to label %callbr.cont []
 
 callbr.cont:
   %phi = phi i32 [ %x, %entry ], [ 0, %callbr.cont ]
   call void @use_i32(i32 %x)
   %fr = freeze i32 %x
   call void @use_i32(i32 %fr)
   call void @use_i32(i32 %phi)
   br label %callbr.cont
 }
 
 define i1 @combine_and_after_freezing_uses(i32 %x) {
 ; CHECK-LABEL: @combine_and_after_freezing_uses(
-; CHECK-NEXT:    [[AND1:%.*]] = and i32 [[X:%.*]], 4
-; CHECK-NEXT:    [[CMP1:%.*]] = icmp ne i32 [[AND1]], 0
-; CHECK-NEXT:    [[X_FR:%.*]] = freeze i32 [[X]]
-; CHECK-NEXT:    [[AND2:%.*]] = and i32 [[X_FR]], 11
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp eq i32 [[AND2]], 11
-; CHECK-NEXT:    [[AND:%.*]] = and i1 [[CMP1]], [[CMP2]]
-; CHECK-NEXT:    ret i1 [[AND]]
+; CHECK-NEXT:    [[X_FR:%.*]] = freeze i32 [[X:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X_FR]], 15
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i32 [[TMP1]], 15
+; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %and1 = and i32 %x, 4
   %cmp1 = icmp ne i32 %and1, 0
   %x.fr = freeze i32 %x
   %and2 = and i32 %x.fr, 11
   %cmp2 = icmp eq i32 %and2, 11
   %and = and i1 %cmp1, %cmp2
   ret i1 %and
 }
 
 declare i1 @mock_use(i64, i64)
 define i1 @fully_propagate_freeze(i32 %0, i32 noundef %1) {
 ; CHECK-LABEL: @fully_propagate_freeze(
-; CHECK-NEXT:    [[DR:%.*]] = lshr i32 [[TMP0:%.*]], 2
+; CHECK-NEXT:    [[DOTFR:%.*]] = freeze i32 [[TMP0:%.*]]
+; CHECK-NEXT:    [[DR:%.*]] = lshr i32 [[DOTFR]], 2
 ; CHECK-NEXT:    [[IDX1:%.*]] = zext i32 [[DR]] to i64
-; CHECK-NEXT:    [[DR_FR:%.*]] = freeze i32 [[DR]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[DR_FR]], 1
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw i32 [[DR]], 1
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[ADD]], [[TMP1:%.*]]
-; CHECK-NEXT:    [[IDX2:%.*]] = zext i32 [[DR_FR]] to i64
+; CHECK-NEXT:    [[IDX2:%.*]] = zext i32 [[DR]] to i64
 ; CHECK-NEXT:    [[V:%.*]] = call i1 @mock_use(i64 [[IDX1]], i64 [[IDX2]])
 ; CHECK-NEXT:    [[RET:%.*]] = and i1 [[V]], [[CMP]]
 ; CHECK-NEXT:    ret i1 [[RET]]
 ;
   %dr = lshr i32 %0, 2
   %idx1 = zext i32 %dr to i64
   %add = add i32 %dr, 1
   %cmp = icmp slt i32 %add, %1

llvm/test/Transforms/InstCombine/onehot_merge.ll

;
   %or = or i1 %t3, %t5
   ret i1 %or
 }
 
 define i1 @foo1_and_extra_use_shl2_logical(i32 %k, i32 %c1, i32 %c2, i32* %p) {
 ; CHECK-LABEL: @foo1_and_extra_use_shl2_logical(
 ; CHECK-NEXT:    [[T0:%.*]] = shl i32 1, [[C1:%.*]]
 ; CHECK-NEXT:    [[T1:%.*]] = shl i32 1, [[C2:%.*]]
-; CHECK-NEXT:    store i32 [[T1]], i32* [[P:%.*]], align 4
 ; CHECK-NEXT:    [[TMP1:%.*]] = freeze i32 [[T1]]
+; CHECK-NEXT:    store i32 [[TMP1]], i32* [[P:%.*]], align 4
 ; CHECK-NEXT:    [[TMP2:%.*]] = or i32 [[T0]], [[TMP1]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = and i32 [[TMP2]], [[K:%.*]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = icmp ne i32 [[TMP3]], [[TMP2]]
 ; CHECK-NEXT:    ret i1 [[TMP4]]
 ;
   %t0 = shl i32 1, %c1
   %t1 = shl i32 1, %c2
   store i32 %t1, i32* %p  ; extra use of shl
;
   %or = or i1 %t3, %t5
   ret i1 %or
 }
 
 define i1 @foo1_and_extra_use_and2_logical(i32 %k, i32 %c1, i32 %c2, i32* %p) {
 ; CHECK-LABEL: @foo1_and_extra_use_and2_logical(
 ; CHECK-NEXT:    [[T0:%.*]] = shl i32 1, [[C1:%.*]]
 ; CHECK-NEXT:    [[T1:%.*]] = shl i32 1, [[C2:%.*]]
-; CHECK-NEXT:    [[T4:%.*]] = and i32 [[T1]], [[K:%.*]]
-; CHECK-NEXT:    store i32 [[T4]], i32* [[P:%.*]], align 4
 ; CHECK-NEXT:    [[TMP1:%.*]] = freeze i32 [[T1]]
+; CHECK-NEXT:    [[T4:%.*]] = and i32 [[TMP1]], [[K:%.*]]
+; CHECK-NEXT:    store i32 [[T4]], i32* [[P:%.*]], align 4
 ; CHECK-NEXT:    [[TMP2:%.*]] = or i32 [[T0]], [[TMP1]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = and i32 [[TMP2]], [[K]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = icmp ne i32 [[TMP3]], [[TMP2]]
 ; CHECK-NEXT:    ret i1 [[TMP4]]
 ;
   %t0 = shl i32 1, %c1
   %t1 = shl i32 1, %c2
   %t2 = and i32 %t0, %k
;
   %or = or i1 %t3, %t5
   ret i1 %or
 }
 
 define i1 @foo1_and_extra_use_cmp2_logical(i32 %k, i32 %c1, i32 %c2, i1* %p) {
 ; CHECK-LABEL: @foo1_and_extra_use_cmp2_logical(
 ; CHECK-NEXT:    [[T0:%.*]] = shl i32 1, [[C1:%.*]]
 ; CHECK-NEXT:    [[T1:%.*]] = shl i32 1, [[C2:%.*]]
-; CHECK-NEXT:    [[T4:%.*]] = and i32 [[T1]], [[K:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i32 [[T1]]
+; CHECK-NEXT:    [[T4:%.*]] = and i32 [[TMP1]], [[K:%.*]]
 ; CHECK-NEXT:    [[T5:%.*]] = icmp eq i32 [[T4]], 0
 ; CHECK-NEXT:    store i1 [[T5]], i1* [[P:%.*]], align 1
-; CHECK-NEXT:    [[TMP1:%.*]] = freeze i32 [[T1]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = or i32 [[T0]], [[TMP1]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = and i32 [[TMP2]], [[K]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = icmp ne i32 [[TMP3]], [[TMP2]]
 ; CHECK-NEXT:    ret i1 [[TMP4]]
 ;
   %t0 = shl i32 1, %c1
   %t1 = shl i32 1, %c2
   %t2 = and i32 %t0, %k