[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 ↗

(On Diff #23992)

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.

llvm/trunk/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 (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);
 }
 

llvm/trunk/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
+  %bc1 = bitcast <2 x double> %a0 to <2 x i64>
+  %bc2 = bitcast <2 x double> %a1 to <2 x i64>
+  %and1 = and <2 x i64> %bc1, <i64 0, i64 -1>
+  %and2 = and <2 x i64> %bc2, <i64 -1, i64 0>
+  %or = or <2 x i64> %and1, %and2
+  %bc3 = bitcast <2 x i64> %or to <2 x double>
+  ret <2 x double> %bc3
+}
+
+
+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
+  %bc1 = bitcast <4 x float> %a0 to <4 x i32>
+  %bc2 = bitcast <4 x float> %a1 to <4 x i32>
+  %and1 = and <4 x i32> %bc1, <i32 0, i32 -1, i32 -1, i32 0>
+  %and2 = and <4 x i32> %bc2, <i32 -1, i32 0, i32 0, i32 -1>
+  %or = or <4 x i32> %and1, %and2
+  %bc3 = bitcast <4 x i32> %or to <4 x float>
+  ret <4 x float> %bc3
+}
+
+
+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
+  %bc1 = bitcast <4 x float> %a0 to <2 x i64>
+  %bc2 = bitcast <4 x float> %a1 to <2 x i64>
+  %and1 = and <2 x i64> %bc1, <i64 0, i64 -1>
+  %and2 = and <2 x i64> %bc2, <i64 -1, i64 0>
+  %or = or <2 x i64> %and1, %and2
+  %bc3 = bitcast <2 x i64> %or to <4 x float>
+  ret <4 x float> %bc3
+}
+
+
+define <4 x float> @test25(<4 x float> %a0) {
+; CHECK-LABEL: test25:
+; CHECK:       # BB#0:
+; CHECK-NEXT:    blendps {{.*#+}} xmm0 = mem[0],xmm0[1,2],mem[3]
+; CHECK-NEXT:    retq
+  %bc1 = bitcast <4 x float> %a0 to <4 x i32>
+  %bc2 = bitcast <4 x float> <float 1.0, float 1.0, float 1.0, float 1.0> to <4 x i32>
+  %and1 = and <4 x i32> %bc1, <i32 0, i32 -1, i32 -1, i32 0>
+  %and2 = and <4 x i32> %bc2, <i32 -1, i32 0, i32 0, i32 -1>
+  %or = or <4 x i32> %and1, %and2
+  %bc3 = bitcast <4 x i32> %or to <4 x float>
+  ret <4 x float> %bc3
+}
+
+
 ; 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
 }