[WebAssembly] Optimize splats of bitcasted vectors

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.

Reviewers: aheejin, dschuff

Subscribers: sbc100, jgravelle-google, hiraditya, sunfish, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D80021

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

@@ -120,6 +120,9 @@ WebAssemblyTargetLowering::WebAssemblyTargetLowering(
 
   // 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})
@@ -1703,5 +1706,39 @@ SDValue WebAssemblyTargetLowering::LowerShift(SDValue Op,
 }
 
 //===----------------------------------------------------------------------===//
-//                          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/lib/Target/WebAssembly/WebAssemblyISelLowering.h

@@ -118,6 +118,11 @@ private:
   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 {

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

@@ -0,0 +1,19 @@
+; 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
+}