[X86] Twist shuffle mask when fold HOP(SHUFFLE(X,Y),SHUFFLE(X,Y)) -> SHUFFLE(HOP(X,Y))

[pengfei]

Summary

This patch tries to fix PR50823.

The shuffle mask should be twisted twice before gotten the correct one due to the difference between inner HOP and outer.

[RKSimon]

Thanks for looking at this - I've been busy with other things this week and haven't really been keeping up with bug traffic!

llvm/test/CodeGen/X86/pr50823.ll

45

can some of this metadata be stripped out to simplify the test?

[lebedev.ri]

This has scaringly many magic numbers.

llvm/lib/Target/X86/X86ISelLowering.cpp

43639

Since the final shuffle has element type of i64/f64,
should this enforce that the source element type is less than that?

43650

Comment didn't save somehow.
Don't you want something like sizeof(i64)/sizeof(current elt size) here?

43664–43666

So what this checking is that ScaledMask0 == ScaledMask1 (i.e. scaled Mask0 and Mask1 are identical)?
Again, special-casing 8 seems weird.

[pengfei]

llvm/lib/Target/X86/X86ISelLowering.cpp

43639

The source type can only be i16, i32, f32 and f64.

43664–43666

Yes. Thanks for reminding.

[pengfei]

Address review comments.

Finally, I figured out the math here. The shuffle mask should be twisted twice before gotten the correct one. But the output happens to be identical when the input mask is <0, 2, 1, 3>. It confused me a long time and I think this is why the bug was hidden.

[pengfei]

llvm/test/CodeGen/X86/packss.ll

373

This is a bug. Verified here https://godbolt.org/z/x5Kns4Mha

[RKSimon]

llvm/test/CodeGen/X86/packss.ll

373

Did you manage to track this one down?

llvm/test/CodeGen/X86/pr50823.ll

7

unused declaration?

[pengfei]

llvm/test/CodeGen/X86/packss.ll

373

Oops, the words are misleading. I meant the previous generated code was buggy and this patch fixes it. I should have commented on the left side. :)

llvm/test/CodeGen/X86/pr50823.ll

7

Good catch!

[pengfei]

Address review comment.

[pengfei]

Fix Lint warning and rebase.

[RKSimon]

LGTM - confirmed the fix - thanks!

[pengfei]

In D104903#2857955, @RKSimon wrote:

LGTM - confirmed the fix - thanks!

Thanks for confirming it!

llvm/lib/Target/X86/X86ISelLowering.cpp

         Res = DAG.getVectorShuffle(ShufVT, DL, Res, Res, PostShuffle);
         return DAG.getBitcast(VT, Res);
       }
     }
   }
 
   // Attempt to fold HOP(SHUFFLE(X,Y),SHUFFLE(X,Y)) -> SHUFFLE(HOP(X,Y)).
   if (VT.is256BitVector() && Subtarget.hasInt256()) {
     SmallVector<int> Mask0, Mask1;
     SmallVector<SDValue> Ops0, Ops1;
     SmallVector<int, 2> ScaledMask0, ScaledMask1;
     if (getTargetShuffleInputs(BC0, Ops0, Mask0, DAG) && !isAnyZero(Mask0) &&
         getTargetShuffleInputs(BC1, Ops1, Mask1, DAG) && !isAnyZero(Mask1) &&
         !Ops0.empty() && !Ops1.empty() &&
         all_of(Ops0,
                [](SDValue Op) { return Op.getValueType().is256BitVector(); }) &&
         all_of(Ops1,
                [](SDValue Op) { return Op.getValueType().is256BitVector(); }) &&
         scaleShuffleElements(Mask0, 2, ScaledMask0) &&
         scaleShuffleElements(Mask1, 2, ScaledMask1)) {
       SDValue Op00 = peekThroughBitcasts(Ops0.front());
       SDValue Op10 = peekThroughBitcasts(Ops1.front());
       SDValue Op01 = peekThroughBitcasts(Ops0.back());
       SDValue Op11 = peekThroughBitcasts(Ops1.back());
       if ((Op00 == Op11) && (Op01 == Op10)) {
         std::swap(Op10, Op11);
         ShuffleVectorSDNode::commuteMask(ScaledMask1);
       }
       if ((Op00 == Op10) && (Op01 == Op11)) {
-        SmallVector<int, 4> ShuffleMask;
-        ShuffleMask.append(ScaledMask0.begin(), ScaledMask0.end());
-        ShuffleMask.append(ScaledMask1.begin(), ScaledMask1.end());
+        const int Map[4] = { 0, 2, 1, 3 };
+        SmallVector<int, 4> ShuffleMask(
+            {Map[ScaledMask0[0]], Map[ScaledMask1[0]], Map[ScaledMask0[1]],
+             Map[ScaledMask1[1]]});
         MVT ShufVT = VT.isFloatingPoint() ? MVT::v4f64 : MVT::v4i64;
         SDValue Res = DAG.getNode(Opcode, DL, VT, DAG.getBitcast(SrcVT, Op00),
                                   DAG.getBitcast(SrcVT, Op01));
         Res = DAG.getBitcast(ShufVT, Res);
         Res = DAG.getVectorShuffle(ShufVT, DL, Res, Res, ShuffleMask);
         return DAG.getBitcast(VT, Res);
       }
     }
   }
 
   return SDValue();

llvm/test/CodeGen/X86/haddsub-undef.ll

 ; AVX1-FAST-NEXT:    vinsertf128 $1, %xmm1, %ymm0, %ymm0
 ; AVX1-FAST-NEXT:    vhaddpd %ymm2, %ymm0, %ymm0
 ; AVX1-FAST-NEXT:    retq
 ;
 ; AVX512-FAST-LABEL: PR34724_add_v4f64_u123:
 ; AVX512-FAST:       # %bb.0:
 ; AVX512-FAST-NEXT:    vextractf128 $1, %ymm0, %xmm0
 ; AVX512-FAST-NEXT:    vhaddpd %ymm1, %ymm0, %ymm0
-; AVX512-FAST-NEXT:    vpermpd {{.*#+}} ymm0 = ymm0[0,2,0,3]
+; AVX512-FAST-NEXT:    vpermpd {{.*#+}} ymm0 = ymm0[0,0,1,3]
 ; AVX512-FAST-NEXT:    retq
   %3 = shufflevector <4 x double> %0, <4 x double> %1, <2 x i32> <i32 2, i32 4>
   %4 = shufflevector <4 x double> %0, <4 x double> %1, <2 x i32> <i32 3, i32 5>
   %5 = fadd <2 x double> %3, %4
   %6 = extractelement <2 x double> %5, i32 0
   %7 = insertelement <4 x double> undef, double %6, i32 1
   %8 = extractelement <2 x double> %5, i32 1
   %9 = insertelement <4 x double> %7, double %8, i32 2

 ; AVX1-FAST-NEXT:    vperm2f128 {{.*#+}} ymm2 = ymm0[2,3],ymm1[2,3]
 ; AVX1-FAST-NEXT:    vblendpd {{.*#+}} ymm0 = ymm0[0,1],ymm1[2,3]
 ; AVX1-FAST-NEXT:    vhaddpd %ymm2, %ymm0, %ymm0
 ; AVX1-FAST-NEXT:    retq
 ;
 ; AVX512-FAST-LABEL: PR34724_add_v4f64_01u3:
 ; AVX512-FAST:       # %bb.0:
 ; AVX512-FAST-NEXT:    vhaddpd %ymm1, %ymm0, %ymm0
-; AVX512-FAST-NEXT:    vpermpd {{.*#+}} ymm0 = ymm0[0,3,1,3]
+; AVX512-FAST-NEXT:    vpermpd {{.*#+}} ymm0 = ymm0[0,2,3,3]
 ; AVX512-FAST-NEXT:    retq
   %3 = shufflevector <4 x double> %0, <4 x double> undef, <2 x i32> <i32 0, i32 2>
   %4 = shufflevector <4 x double> %0, <4 x double> undef, <2 x i32> <i32 1, i32 3>
   %5 = fadd <2 x double> %3, %4
   %6 = extractelement <2 x double> %5, i32 0
   %7 = insertelement <4 x double> undef, double %6, i32 0
   %8 = extractelement <2 x double> %5, i32 1
   %9 = insertelement <4 x double> %7, double %8, i32 1

llvm/test/CodeGen/X86/packss.ll

 ; AVX1-NEXT:    vinsertf128 $1, %xmm0, %ymm1, %ymm0
 ; AVX1-NEXT:    ret{{[l|q]}}
 ;
 ; AVX2-LABEL: packsswb_icmp_zero_trunc_256:
 ; AVX2:       # %bb.0:
 ; AVX2-NEXT:    vpxor %xmm1, %xmm1, %xmm1
 ; AVX2-NEXT:    vpcmpeqw %ymm1, %ymm0, %ymm0
 ; AVX2-NEXT:    vpacksswb %ymm0, %ymm1, %ymm0
-; AVX2-NEXT:    vpermq {{.*#+}} ymm0 = ymm0[0,0,2,3]
+; AVX2-NEXT:    vpermq {{.*#+}} ymm0 = ymm0[0,1,0,3]
 ; AVX2-NEXT:    ret{{[l|q]}}
   %1 = icmp eq <16 x i16> %a0, zeroinitializer
   %2 = sext <16 x i1> %1 to <16 x i16>
   %3 = shufflevector <16 x i16> zeroinitializer, <16 x i16> %2, <32 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 16, i32 17, i32 18, i32 19, i32 20, i32 21, i32 22, i32 23, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15, i32 24, i32 25, i32 26, i32 27, i32 28, i32 29, i32 30, i32 31>
   %4 = trunc <32 x i16> %3 to <32 x i8>
   ret <32 x i8> %4
 }

llvm/test/CodeGen/X86/pr50823.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown -mcpu=core-avx2 | FileCheck %s
+
+%v8_uniform_FVector3 = type { float, float, float }
+
+declare <8 x float> @llvm.x86.avx.hadd.ps.256(<8 x float>, <8 x float>)
+
+define void @foo(%v8_uniform_FVector3* %Out, float* %In, <8 x i32> %__mask) {
+; CHECK-LABEL: foo:
+; CHECK:       # %bb.0: # %allocas
+; CHECK-NEXT:    vmovups (%rsi), %xmm0
+; CHECK-NEXT:    vhaddps 32(%rsi), %xmm0, %xmm0
+; CHECK-NEXT:    vpermpd {{.*#+}} ymm0 = ymm0[0,0,1,1]
+; CHECK-NEXT:    vhaddps %ymm0, %ymm0, %ymm0
+; CHECK-NEXT:    vextractf128 $1, %ymm0, %xmm1
+; CHECK-NEXT:    vaddss %xmm1, %xmm0, %xmm0
+; CHECK-NEXT:    vmovss %xmm0, (%rdi)
+; CHECK-NEXT:    vzeroupper
+; CHECK-NEXT:    retq
+allocas:
+  %ptr_cast_for_load = bitcast float* %In to <8 x float>*
+  %ptr_masked_load74 = load <8 x float>, <8 x float>* %ptr_cast_for_load, align 4
+  %ptr8096 = getelementptr float, float* %In, i64 8
+  %ptr_cast_for_load81 = bitcast float* %ptr8096 to <8 x float>*
+  %ptr80_masked_load82 = load <8 x float>, <8 x float>* %ptr_cast_for_load81, align 4
+  %ret_7.i.i = shufflevector <8 x float> %ptr_masked_load74, <8 x float> %ptr80_masked_load82, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 8, i32 9, i32 10, i32 11>
+  %Out_load19 = getelementptr %v8_uniform_FVector3, %v8_uniform_FVector3* %Out, i64 0, i32 0
+  %v1.i.i100 = tail call <8 x float> @llvm.x86.avx.hadd.ps.256(<8 x float> %ret_7.i.i, <8 x float> %ret_7.i.i)
+  %v2.i.i101 = tail call <8 x float> @llvm.x86.avx.hadd.ps.256(<8 x float> %v1.i.i100, <8 x float> %v1.i.i100)
+  %scalar1.i.i102 = extractelement <8 x float> %v2.i.i101, i32 0
+  %scalar2.i.i103 = extractelement <8 x float> %v2.i.i101, i32 4
+  %sum.i.i104 = fadd float %scalar1.i.i102, %scalar2.i.i103
+  store float %sum.i.i104, float* %Out_load19, align 4
+  ret void
+}