[DAGCombiner] Remove extra bitcasts surrounding vector shuffles

[RKSimon]

Summary

Patch to remove extra bitcasts from shuffles, this is often a legacy of XformToShuffleWithZero being used to combine bitmaskings (of float vectors bitcast to integer vectors) into shuffles:

bitcast(shuffle(bitcast(s0),bitcast(s1))) -> shuffle(s0,s1)

[qcolombet]

Hi Simon,

This looks mostly good, just one question.
Are the tests you added enough to cover all the code you add?
In particular, it is not obvious to me that we cover the case where the bitcast operand of the shuffle are undef, integer constant, or fp constant.
If this is indeed not covered, please add a test case.

Thanks,
-Quentin

test/CodeGen/X86/combine-or.ll

290

Please use instnamer to have easier to update tests.

[RKSimon]

Thanks Quentin,

I've added an extra test for peeking through the bitcasted masked constant vectors. The UNDEF implementation was never getting hit so I dropped that for simplicity (visitVECTOR_SHUFFLE has support for cases where the second operand is UNDEF already).

[qcolombet]

Thanks Simon.

LGTM.

Cheers,
Q.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitBITCAST(SDNode *N) {
 
   // bitconvert(build_pair(ld, ld)) -> ld iff load locations are consecutive.
   if (N0.getOpcode() == ISD::BUILD_PAIR) {
     SDValue CombineLD = CombineConsecutiveLoads(N0.getNode(), VT);
     if (CombineLD.getNode())
       return CombineLD;
   }
 
+  // Remove double bitcasts from shuffles - this is often a legacy of
+  // XformToShuffleWithZero being used to combine bitmaskings (of
+  // float vectors bitcast to integer vectors) into shuffles.
+  // bitcast(shuffle(bitcast(s0),bitcast(s1))) -> shuffle(s0,s1)
+  if (Level < AfterLegalizeDAG && TLI.isTypeLegal(VT) && VT.isVector() &&
+      N0->getOpcode() == ISD::VECTOR_SHUFFLE &&
+      VT.getVectorNumElements() >= N0.getValueType().getVectorNumElements() &&
+      !(VT.getVectorNumElements() % N0.getValueType().getVectorNumElements())) {
+    ShuffleVectorSDNode *SVN = cast<ShuffleVectorSDNode>(N0);
+
+    // If operands are a bitcast, peek through if it casts the original VT.
+    // If operands are a UNDEF or constant, just bitcast back to original VT.
+    auto PeekThroughBitcast = [&](SDValue Op) {
+      if (Op.getOpcode() == ISD::BITCAST &&
+          Op.getOperand(0)->getValueType(0) == VT)
+        return SDValue(Op.getOperand(0));
+      if (Op.getOpcode() == ISD::UNDEF ||
+          ISD::isBuildVectorOfConstantSDNodes(Op.getNode()) ||
+          ISD::isBuildVectorOfConstantFPSDNodes(Op.getNode()))
+        return DAG.getNode(ISD::BITCAST, SDLoc(N), VT, Op);
+      return SDValue();
+    };
+
+    SDValue SV0 = PeekThroughBitcast(N0->getOperand(0));
+    SDValue SV1 = PeekThroughBitcast(N0->getOperand(1));
+    if (!(SV0 && SV1))
+      return SDValue();
+
+    int MaskScale =
+        VT.getVectorNumElements() / N0.getValueType().getVectorNumElements();
+    SmallVector<int, 8> NewMask;
+    for (int M : SVN->getMask())
+      for (int i = 0; i != MaskScale; ++i)
+        NewMask.push_back(M < 0 ? -1 : M * MaskScale + i);
+
+    bool LegalMask = TLI.isShuffleMaskLegal(NewMask, VT);
+    if (!LegalMask) {
+      std::swap(SV0, SV1);
+      ShuffleVectorSDNode::commuteMask(NewMask);
+      LegalMask = TLI.isShuffleMaskLegal(NewMask, VT);
+    }
+
+    if (LegalMask)
+      return DAG.getVectorShuffle(VT, SDLoc(N), SV0, SV1, NewMask);
+  }
+
   return SDValue();
 }
 
 SDValue DAGCombiner::visitBUILD_PAIR(SDNode *N) {
   EVT VT = N->getValueType(0);
   return CombineConsecutiveLoads(N, VT);
 }
 
if (N0.getOpcode() == ISD::FMUL) {
       // Fold scalars or any vector constants (not just splats).
       // This fold is done in general by InstCombine, but extra fmul insts
       // may have been generated during lowering.
       SDValue N00 = N0.getOperand(0);
       SDValue N01 = N0.getOperand(1);
       auto *BV1 = dyn_cast<BuildVectorSDNode>(N1);
       auto *BV00 = dyn_cast<BuildVectorSDNode>(N00);
       auto *BV01 = dyn_cast<BuildVectorSDNode>(N01);
-      
       // Check 1: Make sure that the first operand of the inner multiply is NOT
       // a constant. Otherwise, we may induce infinite looping.
       if (!(isConstOrConstSplatFP(N00) || (BV00 && BV00->isConstant()))) {
         // Check 2: Make sure that the second operand of the inner multiply and
         // the second operand of the outer multiply are constants.
         if ((N1CFP && isConstOrConstSplatFP(N01)) ||
             (BV1 && BV01 && BV1->isConstant() && BV01->isConstant())) {
           SDLoc SL(N);
if (LastOpVT.isFloatingPoint())
       AnyFP = true;
     else if (LastOpVT.isInteger())
       AnyInteger = true;
     else
       return SDValue();
   }
 
   // If any of the operands is a floating point scalar bitcast to a vector,
-  // use floating point types throughout, and bitcast everything.  
+  // use floating point types throughout, and bitcast everything.
   // Replace UNDEFs by another scalar UNDEF node, of the final desired type.
   if (AnyFP) {
     SVT = EVT::getFloatingPointVT(OpVT.getSizeInBits());
     ScalarUndef = DAG.getNode(ISD::UNDEF, DL, SVT);
     if (AnyInteger) {
       for (SDValue &Op : Ops) {
         if (Op.getValueType() == SVT)
           continue;

 /// e.g. AND V, <0xffffffff, 0, 0xffffffff, 0>. ==>
 ///      vector_shuffle V, Zero, <0, 4, 2, 4>
 SDValue DAGCombiner::XformToShuffleWithZero(SDNode *N) {
   EVT VT = N->getValueType(0);
   SDValue LHS = N->getOperand(0);
   SDValue RHS = N->getOperand(1);
   SDLoc dl(N);
 
-  // Make sure we're not running after operation legalization where it 
+  // Make sure we're not running after operation legalization where it
   // may have custom lowered the vector shuffles.
   if (LegalOperations)
     return SDValue();
 
   if (N->getOpcode() != ISD::AND)
     return SDValue();
 
   if (RHS.getOpcode() == ISD::BITCAST)

test/CodeGen/X86/combine-or.ll

 ; CHECK-NEXT:    pslldq {{.*#+}} xmm0 = zero,zero,zero,zero,zero,zero,zero,zero,xmm0[0,1,2,3,4,5,6,7]
 ; CHECK-NEXT:    retq
   %shuf1 = shufflevector <2 x i64> %a, <2 x i64> zeroinitializer, <2 x i32><i32 2, i32 0>
   %shuf2 = shufflevector <2 x i64> %b, <2 x i64> zeroinitializer, <2 x i32><i32 2, i32 0>
   %or = or <2 x i64> %shuf1, %shuf2
   ret <2 x i64> %or
 }
 
+
+; Verify that the dag-combiner keeps the correct domain for float/double vectors
+; bitcast to use the mask-or blend combine.
+
+define <2 x double> @test22(<2 x double> %a0, <2 x double> %a1) {
+; CHECK-LABEL: test22:
+; CHECK:       # BB#0:
+; CHECK-NEXT:    blendpd {{.*#+}} xmm0 = xmm1[0],xmm0[1]
+; CHECK-NEXT:    retq
+  %1 = bitcast <2 x double> %a0 to <2 x i64>
+  %2 = bitcast <2 x double> %a1 to <2 x i64>
+  %3 = and <2 x i64> %1, <i64 0, i64 -1>
+  %4 = and <2 x i64> %2, <i64 -1, i64 0>
+  %5 = or <2 x i64> %3, %4
+  %6 = bitcast <2 x i64> %5 to <2 x double>
+  ret <2 x double> %6
+}
+
+
+define <4 x float> @test23(<4 x float> %a0, <4 x float> %a1) {
+; CHECK-LABEL: test23:
+; CHECK:       # BB#0:
+; CHECK-NEXT:    blendps {{.*#+}} xmm0 = xmm1[0],xmm0[1,2],xmm1[3]
+; CHECK-NEXT:    retq
+  %1 = bitcast <4 x float> %a0 to <4 x i32>
+  %2 = bitcast <4 x float> %a1 to <4 x i32>
+  %3 = and <4 x i32> %1, <i32 0, i32 -1, i32 -1, i32 0>
+  %4 = and <4 x i32> %2, <i32 -1, i32 0, i32 0, i32 -1>
+  %5 = or <4 x i32> %3, %4
+  %6 = bitcast <4 x i32> %5 to <4 x float>
+  ret <4 x float> %6
+}
+
+
+define <4 x float> @test24(<4 x float> %a0, <4 x float> %a1) {
+; CHECK-LABEL: test24:
+; CHECK:       # BB#0:
+; CHECK-NEXT:    blendpd {{.*#+}} xmm0 = xmm1[0],xmm0[1]
+; CHECK-NEXT:    retq
+  %1 = bitcast <4 x float> %a0 to <2 x i64>
+  %2 = bitcast <4 x float> %a1 to <2 x i64>
+  %3 = and <2 x i64> %1, <i64 0, i64 -1>
+  %4 = and <2 x i64> %2, <i64 -1, i64 0>
+  %5 = or <2 x i64> %3, %4
+  %6 = bitcast <2 x i64> %5 to <4 x float>
+  ret <4 x float> %6
+}
+
+
 ; Verify that the DAGCombiner doesn't crash in the attempt to check if a shuffle
 ; with illegal type has a legal mask. Method 'isShuffleMaskLegal' only knows how to
 ; handle legal vector value types.
 define <4 x i8> @test_crash(<4 x i8> %a, <4 x i8> %b) {
 ; CHECK-LABEL: test_crash:
 ; CHECK:       # BB#0:
 ; CHECK-NEXT:    pblendw {{.*#+}} xmm0 = xmm1[0,1,2,3],xmm0[4,5,6,7]
 ; CHECK-NEXT:    retq
   %shuf1 = shufflevector <4 x i8> %a, <4 x i8> zeroinitializer, <4 x i32><i32 4, i32 4, i32 2, i32 3>
   %shuf2 = shufflevector <4 x i8> %b, <4 x i8> zeroinitializer, <4 x i32><i32 0, i32 1, i32 4, i32 4>
   %or = or <4 x i8> %shuf1, %shuf2
   ret <4 x i8> %or
 }