[InstCombine] fold freeze of partial undef/poison vector constants

[spatel]

Summary

We can always replace the undef elements in a vector constant with regular constants to get rid of the freeze:
https://alive2.llvm.org/ce/z/nfRb4F

The select diffs show that we might do better by adjusting the logic for a frozen select condition, but I think those are still improvements.

[spatel]

It should not change the outcome, but there was an unintended change of the 'getTrue' call in the last revision.

[lebedev.ri]

Given that we are replacing freeze constant with constant, and selecting the constant based on all the users,
can we instead go over all the users, and perform replacement for each one separately, always picking the best choice?

[spatel]

In D123962#3457622, @lebedev.ri wrote:

Given that we are replacing freeze constant with constant, and selecting the constant based on all the users,
can we instead go over all the users, and perform replacement for each one separately, always picking the best choice?

Yes, we could use getBinOpIdentity or getSafeVectorConstantForBinop or something like that. But I think it would be bigger/independent of this change, and I haven't seen a need for it yet. Add a TODO?

[nikic]

For vector constants with some undef elements, wouldn't we be better off picking an existing (non-undef) vector element? That would be especially profitable if we can produce a splat vector constant.

In D123962#3457622, @lebedev.ri wrote:

Given that we are replacing freeze constant with constant, and selecting the constant based on all the users,
can we instead go over all the users, and perform replacement for each one separately, always picking the best choice?

Unless I'm misunderstanding what you're suggesting here, this is not legal: The value has to be chosen to be the same for all users.

[spatel]

In D123962#3457645, @nikic wrote:

For vector constants with some undef elements, wouldn't we be better off picking an existing (non-undef) vector element? That would be especially profitable if we can produce a splat vector constant.

In D123962#3457622, @lebedev.ri wrote:

Given that we are replacing freeze constant with constant, and selecting the constant based on all the users,
can we instead go over all the users, and perform replacement for each one separately, always picking the best choice?

Unless I'm misunderstanding what you're suggesting here, this is not legal: The value has to be chosen to be the same for all users.

Agree - looking back to D84948, this was discussed. The initial revision (this link is intended to show it):
https://reviews.llvm.org/D84948?vs=on&id=281923#toc
...added match of one-use freeze at various user sites; that would not work for multi-use.

So we could try harder for the one-use case by using binop identity or similar, but we need to end up with a common value for multi-use. Zero is the easy choice, but it looks like the current implementation could do better...for example, if all of the multiple users are "or" ops, we would generate a "-1".

[spatel]

Patch updated:

1. Added an explanatory comment for the position of this fold.
2. Added TODO comment for enhancement.
3. Added an assert for "BestValue" (suggested in D84948, but didn't make it in).

These all apply to the existing code, so that could be an NFC pre-commit if preferred.

[nikic]

What about

For vector constants with some undef elements, wouldn't we be better off picking an existing (non-undef) vector element? That would be especially profitable if we can produce a splat vector constant.

?

I still think that would make more sense for the vector case.

[spatel]

In D123962#3458814, @nikic wrote:

What about

For vector constants with some undef elements, wouldn't we be better off picking an existing (non-undef) vector element? That would be especially profitable if we can produce a splat vector constant.

?

I still think that would make more sense for the vector case.

Sorry, I didn't see that comment initially. It's complicated to decide what qualifies as "best" constant for a vector. Yes, a splat is generally good, but a zero can be better if it allows a follow-on transform like narrowing. I don't have any evidence to support one way or the other yet.

The motivation for this patch is largely theoretical at this point - I was looking at making the transform from https://github.com/llvm/llvm-project/issues/47012 safe via freeze and saw that it would cause vector test regressions because there was no support for reducing vector constants at all.

[spatel]

Ping.

[nlopes]

LGTM. It's better than keeping freeze around. If we get concrete examples later of non-ideal "best" constants we can improve.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

Instruction *InstCombinerImpl::visitFreeze(FreezeInst &I) {
   if (auto *PN = dyn_cast<PHINode>(Op0)) {
     if (Instruction *NV = foldOpIntoPhi(I, PN))
       return NV;
   }
 
   if (Value *NI = pushFreezeToPreventPoisonFromPropagating(I))
     return replaceInstUsesWith(I, NI);
 
-  if (match(Op0, m_Undef())) {
-    // 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
-    //                       the condition true.
-    // - default: pick 0
+  auto getUndefReplacement = [&I](Type *Ty) {
     Constant *BestValue = nullptr;
-    Constant *NullValue = Constant::getNullValue(I.getType());
+    Constant *NullValue = Constant::getNullValue(Ty);
     for (const auto *U : I.users()) {
       Constant *C = NullValue;
-
       if (match(U, m_Or(m_Value(), m_Value())))
-        C = ConstantInt::getAllOnesValue(I.getType());
+        C = ConstantInt::getAllOnesValue(Ty);
       else if (match(U, m_Select(m_Specific(&I), m_Constant(), m_Value())))
-        C = ConstantInt::getTrue(I.getType());
+        C = ConstantInt::getTrue(Ty);
 
       if (!BestValue)
         BestValue = C;
       else if (BestValue != C)
         BestValue = NullValue;
     }
+    return BestValue;
+  };
 
-    return replaceInstUsesWith(I, BestValue);
+  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))
     return &I;
 
   return nullptr;
 }

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

 }
 
 define <2 x i32> @vec_undef(i1 %cond) {
 ; CHECK-LABEL: @vec_undef(
 ; CHECK-NEXT:    br i1 [[COND:%.*]], label [[A:%.*]], label [[B:%.*]]
 ; CHECK:       A:
 ; CHECK-NEXT:    br label [[C:%.*]]
 ; CHECK:       B:
-; CHECK-NEXT:    [[PHI_FR:%.*]] = freeze <2 x i32> <i32 2, i32 undef>
 ; CHECK-NEXT:    br label [[C]]
 ; CHECK:       C:
-; CHECK-NEXT:    [[Y:%.*]] = phi <2 x i32> [ <i32 0, i32 1>, [[A]] ], [ [[PHI_FR]], [[B]] ]
+; CHECK-NEXT:    [[Y:%.*]] = phi <2 x i32> [ <i32 0, i32 1>, [[A]] ], [ <i32 2, i32 0>, [[B]] ]
 ; CHECK-NEXT:    ret <2 x i32> [[Y]]
 ;
   br i1 %cond, label %A, label %B
 A:
   br label %C
 B:
   br label %C
 C:

llvm/test/Transforms/InstCombine/freeze.ll

;
   %a = select i1 %f, i32 %x, i32 32 ; prefers %f to be false
   %b = select i1 %f, i1 true, i1 %y ; prefers %f to be true
   call void @use_i32_i1(i32 %a, i1 %b)
   ret void
 }
 
 define <3 x i4> @partial_undef_vec() {
 ; CHECK-LABEL: @partial_undef_vec(
-; CHECK-NEXT:    [[F:%.*]] = freeze <3 x i4> <i4 poison, i4 1, i4 undef>
-; CHECK-NEXT:    ret <3 x i4> [[F]]
+; CHECK-NEXT:    ret <3 x i4> <i4 0, i4 1, i4 0>
 ;
   %f = freeze <3 x i4> <i4 poison, i4 1, i4 undef>
   ret <3 x i4> %f
 }
 
 ; Move the freeze forward to prevent poison from spreading.
 
 define i32 @early_freeze_test1(i32 %x, i32 %y) {

llvm/test/Transforms/InstCombine/select.ll

 ;
   %cond.fr = freeze <2 x i1> <i1 undef, i1 undef>
   %s = select <2 x i1> %cond.fr, <2 x i8> <i8 1, i8 2>, <2 x i8> %x
   ret <2 x i8> %s
 }
 
 define <2 x i8> @partial_cond_freeze_constant_true_val_vec(<2 x i8> %x) {
 ; CHECK-LABEL: @partial_cond_freeze_constant_true_val_vec(
-; CHECK-NEXT:    [[COND_FR:%.*]] = freeze <2 x i1> <i1 true, i1 undef>
-; CHECK-NEXT:    [[S:%.*]] = select <2 x i1> [[COND_FR]], <2 x i8> <i8 1, i8 2>, <2 x i8> [[X:%.*]]
-; CHECK-NEXT:    ret <2 x i8> [[S]]
+; CHECK-NEXT:    ret <2 x i8> <i8 1, i8 2>
 ;
   %cond.fr = freeze <2 x i1> <i1 true, i1 undef>
   %s = select <2 x i1> %cond.fr, <2 x i8> <i8 1, i8 2>, <2 x i8> %x
   ret <2 x i8> %s
 }
 
 define <2 x i8> @partial_cond_freeze_constant_false_val_vec(<2 x i8> %x) {
 ; CHECK-LABEL: @partial_cond_freeze_constant_false_val_vec(
-; CHECK-NEXT:    [[COND_FR:%.*]] = freeze <2 x i1> <i1 true, i1 undef>
-; CHECK-NEXT:    [[S:%.*]] = select <2 x i1> [[COND_FR]], <2 x i8> [[X:%.*]], <2 x i8> <i8 1, i8 2>
-; CHECK-NEXT:    ret <2 x i8> [[S]]
+; CHECK-NEXT:    [[S1:%.*]] = insertelement <2 x i8> [[X:%.*]], i8 2, i64 1
+; CHECK-NEXT:    ret <2 x i8> [[S1]]
 ;
   %cond.fr = freeze <2 x i1> <i1 true, i1 undef>
   %s = select <2 x i1> %cond.fr, <2 x i8> %x, <2 x i8> <i8 1, i8 2>
   ret <2 x i8> %s
 }
 
 define <2 x i8> @partial_cond_freeze_both_arms_constant_vec() {
 ; CHECK-LABEL: @partial_cond_freeze_both_arms_constant_vec(
-; CHECK-NEXT:    [[COND_FR:%.*]] = freeze <2 x i1> <i1 false, i1 undef>
-; CHECK-NEXT:    [[S:%.*]] = select <2 x i1> [[COND_FR]], <2 x i8> <i8 1, i8 2>, <2 x i8> <i8 42, i8 43>
-; CHECK-NEXT:    ret <2 x i8> [[S]]
+; CHECK-NEXT:    ret <2 x i8> <i8 42, i8 2>
 ;
   %cond.fr = freeze <2 x i1> <i1 false, i1 undef>
   %s = select <2 x i1> %cond.fr, <2 x i8> <i8 1, i8 2>, <2 x i8> <i8 42, i8 43>
   ret <2 x i8> %s
 }
 
 declare void @foo2(i8, i8)