[WebAssembly] Optimize splats of bitcasted vectors

[tlively]

Summary

This new custom DAG combine fixes a codegen issue with the
wasm_simd128.h intrinsics. Clang lowers the

return (v128_t)(__f32x4){__a, __a, __a, __a};

body of f32x4_splat to a splat shuffle of a bitcasted vector, as seen
in the new simd-shuffle-bitcast.ll test. The bitcast interfered with
the target-independent DAG combine that combines splat shuffles into
BUILD_VECTOR nodes, so this patch introduces a new custom DAG combine
to hoist the bitcast out of the shuffle, allowing the
target-independent combine to work as intended.

[dschuff]

You know, it's pretty amazing that WebAssemblyISelLowering.cpp is as small as it is.

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.h

private:
   SDValue LowerCopyToReg(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerIntrinsic(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerSIGN_EXTEND_INREG(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerBUILD_VECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVECTOR_SHUFFLE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerSETCC(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerAccessVectorElement(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerShift(SDValue Op, SelectionDAG &DAG) const;
+
+  // Custom DAG combine hooks
+  SDValue
+  PerformDAGCombine(SDNode *N,
+                    TargetLowering::DAGCombinerInfo &DCI) const override;
 };
 
 namespace WebAssembly {
 FastISel *createFastISel(FunctionLoweringInfo &funcInfo,
                          const TargetLibraryInfo *libInfo);
 } // end namespace WebAssembly
 
 } // end namespace llvm
 
 #endif

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

for (auto T : {MVT::i32, MVT::i64})
       setOperationAction(Op, T, Expand);
     if (Subtarget->hasSIMD128())
       for (auto T : {MVT::v16i8, MVT::v8i16, MVT::v4i32, MVT::v2i64})
         setOperationAction(Op, T, Expand);
   }
 
   // SIMD-specific configuration
   if (Subtarget->hasSIMD128()) {
+    // Hoist bitcasts out of shuffles
+    setTargetDAGCombine(ISD::VECTOR_SHUFFLE);
+
     // Support saturating add for i8x16 and i16x8
     for (auto Op : {ISD::SADDSAT, ISD::UADDSAT})
       for (auto T : {MVT::v16i8, MVT::v8i16})
         setOperationAction(Op, T, Legal);
 
     // Support integer abs
     for (auto T : {MVT::v16i8, MVT::v8i16, MVT::v4i32})
       setOperationAction(ISD::ABS, T, Legal);
default:
     llvm_unreachable("unexpected opcode");
   }
   APInt Shift = SplatConst->getAPIntValue().zextOrTrunc(32);
   return DAG.getNode(Opcode, DL, Op.getValueType(), Op.getOperand(0),
                      DAG.getConstant(Shift, DL, MVT::i32));
 }
 
 //===----------------------------------------------------------------------===//
-//                          WebAssembly Optimization Hooks
+//   Custom DAG combine hooks
 //===----------------------------------------------------------------------===//
+static SDValue
+performVECTOR_SHUFFLECombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI) {
+  auto &DAG = DCI.DAG;
+  auto Shuffle = cast<ShuffleVectorSDNode>(N);
+
+  // Hoist vector bitcasts that don't change the number of lanes out of unary
+  // shuffles, where they are less likely to get in the way of other combines.
+  // (shuffle (vNxT1 (bitcast (vNxT0 x))),  undef, mask) ->
+  //  (vNxT1 (bitcast (vNxt0 (shuffle x, undef, mask))))
+  SDValue Bitcast = N->getOperand(0);
+  if (Bitcast.getOpcode() != ISD::BITCAST)
+    return SDValue();
+  if (!N->getOperand(1).isUndef())
+    return SDValue();
+  SDValue CastOp = Bitcast.getOperand(0);
+  MVT SrcType = CastOp.getSimpleValueType();
+  MVT DstType = Bitcast.getSimpleValueType();
+  if (SrcType.getVectorNumElements() != DstType.getVectorNumElements())
+    return SDValue();
+  SDValue NewShuffle = DAG.getVectorShuffle(
+      SrcType, SDLoc(N), CastOp, DAG.getUNDEF(SrcType), Shuffle->getMask());
+  return DAG.getBitcast(DstType, NewShuffle);
+}
+
+SDValue
+WebAssemblyTargetLowering::PerformDAGCombine(SDNode *N,
+                                             DAGCombinerInfo &DCI) const {
+  switch (N->getOpcode()) {
+  default:
+    return SDValue();
+  case ISD::VECTOR_SHUFFLE:
+    return performVECTOR_SHUFFLECombine(N, DCI);
+  }
+}

llvm/test/CodeGen/WebAssembly/simd-shuffle-bitcast.ll

+; RUN: llc < %s -asm-verbose=false -verify-machineinstrs -disable-wasm-fallthrough-return-opt -wasm-disable-explicit-locals -wasm-keep-registers -mattr=+simd128 | FileCheck %s
+
+; Test that a splat shuffle of an fp-to-int bitcasted vector correctly
+; optimizes and lowers to a single splat instruction. Without a custom
+; DAG combine, this ends up doing both a splat and a shuffle.
+
+target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+target triple = "wasm32-unknown-unkown"
+
+; CHECK-LABEL: f32x4_splat:
+; CHECK-NEXT: .functype f32x4_splat (f32) -> (v128){{$}}
+; CHECK-NEXT: f32x4.splat $push[[R:[0-9]+]]=, $0{{$}}
+; CHECK-NEXT: return $pop[[R]]{{$}}
+define <4 x i32> @f32x4_splat(float %x) {
+  %vecinit = insertelement <4 x float> undef, float %x, i32 0
+  %a = bitcast <4 x float> %vecinit to <4 x i32>
+  %b = shufflevector <4 x i32> %a, <4 x i32> undef, <4 x i32> zeroinitializer
+  ret <4 x i32> %b
+}