[WebAssembly] Custom combine splat build_vectors into swizzles

[tlively]

Summary

Some splat build_vectors can be lowered to swizzles, but only if all
of their lanes but one are undefined. This patch adds a custom combine
to turn these build_vectors into swizzles before their undefined lanes
are combined away by the combine introduced in D83606.

Depends on D83606.

[aheejin]

LGTM, but I posted a question on whether we need performBUILD_VECTORCombine in D83606.

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

performVECTOR_SHUFFLECombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI) {
   if (!SrcType.is128BitVector() ||
       SrcType.getVectorNumElements() != DstType.getVectorNumElements())
     return SDValue();
   SDValue NewShuffle = DAG.getVectorShuffle(
       SrcType, SDLoc(N), CastOp, DAG.getUNDEF(SrcType), Shuffle->getMask());
   return DAG.getBitcast(DstType, NewShuffle);
 }
 
+static SDValue combineSwizzleSplat(BuildVectorSDNode *N, SelectionDAG &DAG,
+                                   const BitVector &UndefElts) {
+  // This splat can only be lowered to a swizzle if all lanes but one are undef.
+  //
+  // We are looking for the following pattern:
+  //   (insert undef,
+  //     (extract Vec1,
+  //       (sext (extract Vec2, Index))
+  //     ),
+  //     Index
+  //   )
+  //
+  // To combine to: (swizzle vec1, vec2)
+  if (N->getValueType(0) != MVT::v16i8)
+    return SDValue();
+  if (UndefElts.count() != 15)
+    return SDValue();
+  unsigned Index1 = UndefElts.find_first_unset();
+  auto Extract1 = N->getOperand(Index1);
+  if (Extract1.getOpcode() != ISD::EXTRACT_VECTOR_ELT)
+    return SDValue();
+  auto Vec1 = Extract1.getOperand(0);
+  if (Vec1.getValueType() != MVT::v16i8)
+    return SDValue();
+  auto SExt = Extract1.getOperand(1);
+  if (SExt.getOpcode() != ISD::SIGN_EXTEND)
+    return SDValue();
+  auto Extract2 = SExt.getOperand(0);
+  if (Extract2.getOpcode() != ISD::EXTRACT_VECTOR_ELT)
+    return SDValue();
+  auto Vec2 = Extract2.getOperand(0);
+  if (Vec2.getValueType() != MVT::v16i8)
+    return SDValue();
+  if (Extract2.getOperand(1).getOpcode() != ISD::Constant)
+    return SDValue();
+  unsigned Index2 = Extract2.getConstantOperandVal(1);
+  if (Index1 != Index2)
+    return SDValue();
+
+  return DAG.getNode(WebAssemblyISD::SWIZZLE, SDLoc(N), MVT::v16i8, Vec1, Vec2);
+}
+
 static SDValue
 performBUILD_VECTORCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI) {
   auto &DAG = DCI.DAG;
   auto *Build = cast<BuildVectorSDNode>(N);
 
   // Remove undef lanes from splats. They don't allow us to make any extra
   // optimizations and they can inhibit splat scalarization combines.
   BitVector UndefElts;
-  if (SDValue SplatVal = Build->getSplatValue(&UndefElts))
-    if (UndefElts.any())
+  if (SDValue SplatVal = Build->getSplatValue(&UndefElts)) {
+    if (UndefElts.any()) {
+      // If this BUILD_VECTOR can be implemented as a swizzle, perform that
+      // transformation now because once we remove the undef lanes we will not
+      // be able to recover the swizzle.
+      if (auto Swizzle = combineSwizzleSplat(Build, DAG, UndefElts))
+        return Swizzle;
       return DAG.getSplatBuildVector(N->getValueType(0), SDLoc(N), SplatVal);
+    }
+  }
 
   return SDValue();
 }
 
 static SDValue
 performEXTRACT_VECTOR_ELTCombine(SDNode *N,
                                  TargetLowering::DAGCombinerInfo &DCI) {
   auto &DAG = DCI.DAG;

llvm/test/CodeGen/WebAssembly/simd-build-vector.ll

define <8 x i16> @splat_common_arg_i16x8(i16 %a, i16 %b, i16 %c) {
   %v3 = insertelement <8 x i16> %v2, i16 %c, i32 3
   %v4 = insertelement <8 x i16> %v3, i16 %b, i32 4
   %v5 = insertelement <8 x i16> %v4, i16 %c, i32 5
   %v6 = insertelement <8 x i16> %v5, i16 %c, i32 6
   %v7 = insertelement <8 x i16> %v6, i16 %b, i32 7
   ret <8 x i16> %v7
 }
 
-;; TODO: This should be a swizzle, but will need a custom combine to
-;; preempt the combine that removes undef lanes from splat
-;; build_vectors, since swizzle lowering depends on those lanes being
-;; undef.
-
 ; CHECK-LABEL: swizzle_one_i8x16:
 ; CHECK-NEXT:  .functype       swizzle_one_i8x16 (v128, v128) -> (v128)
-; CHECK-NOT:  v8x16.swizzle
+; CHECK-NEXT:  v8x16.swizzle   $push[[L0:[0-9]+]]=, $0, $1
+; CHECK-NEXT:  return          $pop[[L0]]
 define <16 x i8> @swizzle_one_i8x16(<16 x i8> %src, <16 x i8> %mask) {
   %m0 = extractelement <16 x i8> %mask, i32 0
   %s0 = extractelement <16 x i8> %src, i8 %m0
   %v0 = insertelement <16 x i8> undef, i8 %s0, i32 0
   ret <16 x i8> %v0
 }
 
 ; CHECK-LABEL: swizzle_two_i8x16: