[InstCombine] Don't fold freeze poison when it's used in shufflevector

[ManuelJBrito]

Summary

With this patch freeze undef/poison will no longer be folded into a constant if it's used as a vector operand in a shufflevector.
Currently it's folded into zeroinitializer (https://godbolt.org/z/a6168aecv)

This is necessary to fix avx cast intrinsics, which are currently being fixed in https://reviews.llvm.org/D143287 to use freeze(poison).

We need to maintain the freeze(poison). to have the correct lowering (https://godbolt.org/z/1ecM8roYh), because these are supposed to be no-op intrinsics.

[ManuelJBrito]

I do realize this is a bit hacky, any feedback is appreciated.

[spatel]

I haven't been following the front-end and back-end patches related to this, so I'm not sure where we stand.
Assuming the other parts are approved, we still need a regression test for this, and a code comment to explain why we're carving out this exception to the general rule.

[ManuelJBrito]

Add tests and code comment

[ManuelJBrito]

Alternatively I was thinking of moving this match into getUndefReplacement and let it be a case where the undef replacement does no replacing.
After reading the discussion here [ https://reviews.llvm.org/D123962 ], I think this approach makes the change cleaner and conceptually more sound.
Would this make sense?

[spatel]

In D143593#4118517, @ManuelJBrito wrote:

Alternatively I was thinking of moving this match into getUndefReplacement and let it be a case where the undef replacement does no replacing.
After reading the discussion here [ https://reviews.llvm.org/D123962 ], I think this approach makes the change cleaner and conceptually more sound.
Would this make sense?

Not sure if that's worth the effort - you'd have to adjust the caller too to allow no replacement?
It looks like the updated patch revision is on top of the previous revision rather than main, so that needs to be fixed.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

4022–4024

The indentation is off, and we don't want to mention any x86-specific hackery here.
Just say something like:

// This may improve codegen for shuffles that allow unspecified inputs.

llvm/test/Transforms/InstCombine/shufflevector_freezepoison.ll

3

This isn't what I was expecting. If we're bailing out for operand 1, then shouldn't we be symmetric and do the same for operand 0 (just in case that logically equivalent pattern was somehow created in the front-end)?

9

Why do all of the tests have an extra use? If we don't expect the behavior to change with extra uses, then it would be better if some tests have the extra use and some do not.

18

Are the CHECK lines manually written? Use utils/update_test_checks.py to auto-generate those lines instead.
You can pre-commit this test file to main with the baseline CHECKs, so we'll just see diffs in this patch.

[ManuelJBrito]

Thanks for the feedback !!

llvm/test/Transforms/InstCombine/shufflevector_freezepoison.ll

3

I matched only on op 1, because that's the pattern that is recognized by the backend.
As of right now the symmetric case isn't lower in the same manner : https://godbolt.org/z/K5EWfdbq8, but yes i do think bailing out on both operands makes the most sense.

[ManuelJBrito]

Address comments

[RKSimon]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

4029

I'm worried this will start to fail as soon as we have bitcasts inserted in the use chain somewhere - something that occurs a lot with x86 shuffles for instance.

[ManuelJBrito]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

4029

I'm afraid you are correct.
I'm thinking we can either disable the folding for bitcasts or check the usages of the bitcast and block the folding if it is used in a shufflevector.
Not sure if the first option is too aggressive. While the second option would require to visit the usages recursively.

[ManuelJBrito]

Handle bitcasts - I could not find any API for this so I made a little helper function.

[spatel]

I'm getting a bad feeling because the hack is already growing. :)

What if we squash any cast of freeze(poison) into the freeze itself as an addition to getUndefReplacement():
https://alive2.llvm.org/ce/z/o7J2JS

That would give us a far-fetched improvement on something like this:
https://alive2.llvm.org/ce/z/hC_Gws
...so that could be an independent patch before the shuffle hack.

[ManuelJBrito]

In D143593#4130108, @spatel wrote:

What if we squash any cast of freeze(poison) into the freeze itself as an addition to getUndefReplacement():
https://alive2.llvm.org/ce/z/o7J2JS

That would give us a far-fetched improvement on something like this:
https://alive2.llvm.org/ce/z/hC_Gws
...so that could be an independent patch before the shuffle hack.

I think this would work great! The main hurdle is that getUndefReplacement's users expect a Constant.
In particular Constant::replaceUndefsWith, so I would have to change the return type of getUndefReplacement to Value and also change Constant::replaceUndefsWith to a take a Value as a Replacement.
I'm not sure which option is more desirable, the current patch or your suggestion.

[spatel]

In D143593#4135928, @ManuelJBrito wrote:

In D143593#4130108, @spatel wrote:

What if we squash any cast of freeze(poison) into the freeze itself as an addition to getUndefReplacement():
https://alive2.llvm.org/ce/z/o7J2JS

That would give us a far-fetched improvement on something like this:
https://alive2.llvm.org/ce/z/hC_Gws
...so that could be an independent patch before the shuffle hack.

I think this would work great! The main hurdle is that getUndefReplacement's users expect a Constant.
In particular Constant::replaceUndefsWith, so I would have to change the return type of getUndefReplacement to Value and also change Constant::replaceUndefsWith to a take a Value as a Replacement.
I'm not sure which option is more desirable, the current patch or your suggestion.

Yes, this is a smaller patch. It's contained enough that it probably doesn't matter either way, so LGTM.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3962

Add a block comment for this function with something like:

/// Check if any direct or bitcast user of this value is a shuffle instruction.

[xbolva00]

This should be probably reviewed also by @nikic or @nlopes

[nikic]

I'm okay with this.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3966

Just check isa<ShuffleVectorInst>? Odd to check for specific operands if we don't care which one it is.

[ManuelJBrito]

Thanks !! I'll address these final comments in the commited patch.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

Op->replaceUsesWithIf(&FI, [&](Use &U) -> bool {
     bool Dominates = DT.dominates(&FI, U);
     Changed |= Dominates;
     return Dominates;
   });
 
   return Changed;
 }
 
+// Check if any direct or bitcast user of this value is a shuffle instruction.
+static bool isUsedWithinShuffleVector(Value *V) {
+  for (auto *U : V->users()) {
+    if (isa<ShuffleVectorInst>(U))
+      return true;
+    else if (match(U, m_BitCast(m_Specific(V))) && isUsedWithinShuffleVector(U))
+      return true;
+  }
+  return false;
+}
+
 Instruction *InstCombinerImpl::visitFreeze(FreezeInst &I) {
   Value *Op0 = I.getOperand(0);
 
   if (Value *V = simplifyFreezeInst(Op0, SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 
   // freeze (phi const, x) --> phi const, (freeze x)
   if (auto *PN = dyn_cast<PHINode>(Op0)) {
for (const auto *U : I.users()) {
         BestValue = C;
       else if (BestValue != C)
         BestValue = NullValue;
     }
     assert(BestValue && "Must have at least one use");
     return BestValue;
   };
 
-  if (match(Op0, m_Undef()))
+  if (match(Op0, m_Undef())) {
+    // Don't fold freeze(undef/poison) if it's used as a vector operand in
+    // a shuffle. This may improve codegen for shuffles that allow
+    // unspecified inputs.
+    if (isUsedWithinShuffleVector(&I))
+      return nullptr;
     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 uses of Op with freeze(Op).

llvm/test/Transforms/InstCombine/shufflevector_freezepoison.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -passes='instcombine' -S %s | FileCheck %s
 
 
 define <4 x double> @shuffle_op0_freeze_poison(<2 x double> %a) {
 ; CHECK-LABEL: @shuffle_op0_freeze_poison(
-; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x double> zeroinitializer, <2 x double> [[A:%.*]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    [[B:%.*]] = freeze <2 x double> poison
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x double> [[B]], <2 x double> [[A:%.*]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
 ; CHECK-NEXT:    ret <4 x double> [[SHUFFLE]]
 ;
   %b = freeze <2 x double> poison
   %shuffle = shufflevector <2 x double> %b, <2 x double> %a, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   ret <4 x double> %shuffle
 }
 
 define <4 x double> @shuffle_op1_freeze_poison(<2 x double> %a) {
 ; CHECK-LABEL: @shuffle_op1_freeze_poison(
-; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x double> [[A:%.*]], <2 x double> zeroinitializer, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    [[B:%.*]] = freeze <2 x double> poison
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x double> [[A:%.*]], <2 x double> [[B]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
 ; CHECK-NEXT:    ret <4 x double> [[SHUFFLE]]
 ;
   %b = freeze <2 x double> poison
   %shuffle = shufflevector <2 x double> %a, <2 x double> %b, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   ret <4 x double> %shuffle
 }
 
 define <4 x double> @shuffle_op0_freeze_poison_use(<2 x double> %a) {
 ; CHECK-LABEL: @shuffle_op0_freeze_poison_use(
-; CHECK-NEXT:    call void @use(<2 x double> zeroinitializer)
-; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x double> zeroinitializer, <2 x double> [[A:%.*]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    [[B:%.*]] = freeze <2 x double> poison
+; CHECK-NEXT:    call void @use(<2 x double> [[B]])
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x double> [[B]], <2 x double> [[A:%.*]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
 ; CHECK-NEXT:    ret <4 x double> [[SHUFFLE]]
 ;
   %b = freeze <2 x double> poison
   call void @use(<2 x double> %b)
   %shuffle = shufflevector <2 x double> %b, <2 x double> %a, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   ret <4 x double> %shuffle
 }
 
 define <4 x double> @shuffle_op1_freeze_poison_use(<2 x double> %a) {
 ; CHECK-LABEL: @shuffle_op1_freeze_poison_use(
-; CHECK-NEXT:    call void @use(<2 x double> zeroinitializer)
-; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x double> [[A:%.*]], <2 x double> zeroinitializer, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    [[B:%.*]] = freeze <2 x double> poison
+; CHECK-NEXT:    call void @use(<2 x double> [[B]])
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x double> [[A:%.*]], <2 x double> [[B]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
 ; CHECK-NEXT:    ret <4 x double> [[SHUFFLE]]
 ;
   %b = freeze <2 x double> poison
   call void @use(<2 x double> %b)
   %shuffle = shufflevector <2 x double> %a, <2 x double> %b, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   ret <4 x double> %shuffle
 }
 
 define <4 x double> @shuffle_op0_freeze_undef(<2 x double> %a) {
 ; CHECK-LABEL: @shuffle_op0_freeze_undef(
-; CHECK-NEXT:    call void @use(<2 x double> zeroinitializer)
-; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x double> zeroinitializer, <2 x double> [[A:%.*]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    [[B:%.*]] = freeze <2 x double> undef
+; CHECK-NEXT:    call void @use(<2 x double> [[B]])
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x double> [[B]], <2 x double> [[A:%.*]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
 ; CHECK-NEXT:    ret <4 x double> [[SHUFFLE]]
 ;
   %b = freeze <2 x double> undef
   call void @use(<2 x double> %b)
   %shuffle = shufflevector <2 x double> %b, <2 x double> %a, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   ret <4 x double> %shuffle
 }
 
 define <4 x double> @shuffle_op1_freeze_undef(<2 x double> %a) {
 ; CHECK-LABEL: @shuffle_op1_freeze_undef(
-; CHECK-NEXT:    call void @use(<2 x double> zeroinitializer)
-; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x double> [[A:%.*]], <2 x double> zeroinitializer, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    [[B:%.*]] = freeze <2 x double> undef
+; CHECK-NEXT:    call void @use(<2 x double> [[B]])
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x double> [[A:%.*]], <2 x double> [[B]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
 ; CHECK-NEXT:    ret <4 x double> [[SHUFFLE]]
 ;
   %b = freeze <2 x double> undef
   call void @use(<2 x double> %b)
   %shuffle = shufflevector <2 x double> %a, <2 x double> %b, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   ret <4 x double> %shuffle
 }
 
+define <4 x double> @shuffle_bc1(<2 x double>  %a)  {
+; CHECK-LABEL: @shuffle_bc1(
+; CHECK-NEXT:    [[B:%.*]] = freeze <4 x float> poison
+; CHECK-NEXT:    [[B1:%.*]] = bitcast <4 x float> [[B]] to <2 x double>
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x double> [[A:%.*]], <2 x double> [[B1]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    ret <4 x double> [[SHUFFLE]]
+;
+  %b = freeze <4 x float> poison
+  %b1 = bitcast <4 x float> %b to <2 x double>
+  %shuffle = shufflevector <2 x double> %a, <2 x double> %b1, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+  ret <4 x double> %shuffle
+}
+
+define <8 x float> @shuffle_bc2(<4 x float>  %a)  {
+; CHECK-LABEL: @shuffle_bc2(
+; CHECK-NEXT:    [[B:%.*]] = freeze <4 x float> poison
+; CHECK-NEXT:    [[B1:%.*]] = bitcast <4 x float> [[B]] to <2 x double>
+; CHECK-NEXT:    call void @use(<2 x double> [[B1]])
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x float> [[A:%.*]], <4 x float> [[B]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+; CHECK-NEXT:    ret <8 x float> [[SHUFFLE]]
+;
+  %b = freeze <4 x float> poison
+  %b1 = bitcast <4 x float> %b to <2 x double>
+  call void @use(<2 x double> %b1)
+  %b2 = bitcast <2 x double> %b1 to <4 x float>
+  %shuffle = shufflevector <4 x float> %a, <4 x float> %b2, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+  ret <8 x float> %shuffle
+}
+
 declare void @use(<2 x double>)