[WebAssembly] Scalarize extract_vector_elt of binops

[tlively]

Summary

Override the shouldScalarizeBinop target lowering hook using the same
implementation used in the x86 backend. This causes extract_vector_elts of
vector binary ops to be scalarized if the scalarized version would be supported.

[aheejin]

Nice!

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.h

WebAssemblyTargetLowering(const TargetMachine &TM,
                             const WebAssemblySubtarget &STI);
 
 private:
   /// Keep a pointer to the WebAssemblySubtarget around so that we can make the
   /// right decision when generating code for different targets.
   const WebAssemblySubtarget *Subtarget;
 
   AtomicExpansionKind shouldExpandAtomicRMWInIR(AtomicRMWInst *) const override;
+  bool shouldScalarizeBinop(SDValue VecOp) const override;
   FastISel *createFastISel(FunctionLoweringInfo &FuncInfo,
                            const TargetLibraryInfo *LibInfo) const override;
   MVT getScalarShiftAmountTy(const DataLayout &DL, EVT) const override;
   MachineBasicBlock *
   EmitInstrWithCustomInserter(MachineInstr &MI,
                               MachineBasicBlock *MBB) const override;
   const char *getTargetNodeName(unsigned Opcode) const override;
   std::pair<unsigned, const TargetRegisterClass *>

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

WebAssemblyTargetLowering::shouldExpandAtomicRMWInIR(AtomicRMWInst *AI) const {
   case AtomicRMWInst::Xchg:
     return AtomicExpansionKind::None;
   default:
     break;
   }
   return AtomicExpansionKind::CmpXChg;
 }
 
+bool WebAssemblyTargetLowering::shouldScalarizeBinop(SDValue VecOp) const {
+  // Implementation copied from X86TargetLowering.
+  unsigned Opc = VecOp.getOpcode();
+
+  // Assume target opcodes can't be scalarized.
+  // TODO - do we have any exceptions?
+  if (Opc >= ISD::BUILTIN_OP_END)
+    return false;
+
+  // If the vector op is not supported, try to convert to scalar.
+  EVT VecVT = VecOp.getValueType();
+  if (!isOperationLegalOrCustomOrPromote(Opc, VecVT))
+    return true;
+
+  // If the vector op is supported, but the scalar op is not, the transform may
+  // not be worthwhile.
+  EVT ScalarVT = VecVT.getScalarType();
+  return isOperationLegalOrCustomOrPromote(Opc, ScalarVT);
+}
+
 FastISel *WebAssemblyTargetLowering::createFastISel(
     FunctionLoweringInfo &FuncInfo, const TargetLibraryInfo *LibInfo) const {
   return WebAssembly::createFastISel(FuncInfo, LibInfo);
 }
 
 MVT WebAssemblyTargetLowering::getScalarShiftAmountTy(const DataLayout & /*DL*/,
                                                       EVT VT) const {
   unsigned BitWidth = NextPowerOf2(VT.getSizeInBits() - 1);

llvm/test/CodeGen/WebAssembly/masked-shifts.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -verify-machineinstrs -mattr=+simd128 | FileCheck %s
 
 ;; Check that masked shift counts are optimized out.
 
-;; TODO: optimize the *_late functions.
-
 target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
 target triple = "wasm32-unknown-unknown"
 
 define i32 @shl_i32(i32 %v, i32 %x) {
 ; CHECK-LABEL: shl_i32:
 ; CHECK:         .functype shl_i32 (i32, i32) -> (i32)
 ; CHECK-NEXT:  # %bb.0:
 ; CHECK-NEXT:    local.get 0

 }
 
 define <4 x i32> @shl_v4i32_late(<4 x i32> %v, i32 %x) {
 ; CHECK-LABEL: shl_v4i32_late:
 ; CHECK:         .functype shl_v4i32_late (v128, i32) -> (v128)
 ; CHECK-NEXT:  # %bb.0:
 ; CHECK-NEXT:    local.get 0
 ; CHECK-NEXT:    local.get 1
-; CHECK-NEXT:    i32x4.splat
-; CHECK-NEXT:    v128.const 31, 31, 31, 31
-; CHECK-NEXT:    v128.and
-; CHECK-NEXT:    i32x4.extract_lane 0
 ; CHECK-NEXT:    i32x4.shl
 ; CHECK-NEXT:    # fallthrough-return
   %t = insertelement <4 x i32> undef, i32 %x, i32 0
   %s = shufflevector <4 x i32> %t, <4 x i32> undef,
     <4 x i32> <i32 0, i32 0, i32 0, i32 0>
   %m = and <4 x i32> %s, <i32 31, i32 31, i32 31, i32 31>
   %a = shl <4 x i32> %v, %m
   ret <4 x i32> %a

 }
 
 define <4 x i32> @ashr_v4i32_late(<4 x i32> %v, i32 %x) {
 ; CHECK-LABEL: ashr_v4i32_late:
 ; CHECK:         .functype ashr_v4i32_late (v128, i32) -> (v128)
 ; CHECK-NEXT:  # %bb.0:
 ; CHECK-NEXT:    local.get 0
 ; CHECK-NEXT:    local.get 1
-; CHECK-NEXT:    i32x4.splat
-; CHECK-NEXT:    v128.const 31, 31, 31, 31
-; CHECK-NEXT:    v128.and
-; CHECK-NEXT:    i32x4.extract_lane 0
 ; CHECK-NEXT:    i32x4.shr_s
 ; CHECK-NEXT:    # fallthrough-return
   %t = insertelement <4 x i32> undef, i32 %x, i32 0
   %s = shufflevector <4 x i32> %t, <4 x i32> undef,
     <4 x i32> <i32 0, i32 0, i32 0, i32 0>
   %m = and <4 x i32> %s, <i32 31, i32 31, i32 31, i32 31>
   %a = ashr <4 x i32> %v, %m
   ret <4 x i32> %a

 }
 
 define <4 x i32> @lshr_v4i32_late(<4 x i32> %v, i32 %x) {
 ; CHECK-LABEL: lshr_v4i32_late:
 ; CHECK:         .functype lshr_v4i32_late (v128, i32) -> (v128)
 ; CHECK-NEXT:  # %bb.0:
 ; CHECK-NEXT:    local.get 0
 ; CHECK-NEXT:    local.get 1
-; CHECK-NEXT:    i32x4.splat
-; CHECK-NEXT:    v128.const 31, 31, 31, 31
-; CHECK-NEXT:    v128.and
-; CHECK-NEXT:    i32x4.extract_lane 0
 ; CHECK-NEXT:    i32x4.shr_u
 ; CHECK-NEXT:    # fallthrough-return
   %t = insertelement <4 x i32> undef, i32 %x, i32 0
   %s = shufflevector <4 x i32> %t, <4 x i32> undef,
     <4 x i32> <i32 0, i32 0, i32 0, i32 0>
   %m = and <4 x i32> %s, <i32 31, i32 31, i32 31, i32 31>
   %a = lshr <4 x i32> %v, %m
   ret <4 x i32> %a

 }
 
 define <2 x i64> @shl_v2i64_late(<2 x i64> %v, i64 %x) {
 ; CHECK-LABEL: shl_v2i64_late:
 ; CHECK:         .functype shl_v2i64_late (v128, i64) -> (v128)
 ; CHECK-NEXT:  # %bb.0:
 ; CHECK-NEXT:    local.get 0
 ; CHECK-NEXT:    local.get 1
-; CHECK-NEXT:    i64x2.splat
-; CHECK-NEXT:    v128.const 63, 63
-; CHECK-NEXT:    v128.and
-; CHECK-NEXT:    i64x2.extract_lane 0
 ; CHECK-NEXT:    i32.wrap_i64
 ; CHECK-NEXT:    i64x2.shl
 ; CHECK-NEXT:    # fallthrough-return
   %t = insertelement <2 x i64> undef, i64 %x, i32 0
   %s = shufflevector <2 x i64> %t, <2 x i64> undef, <2 x i32> <i32 0, i32 0>
   %m = and <2 x i64> %s, <i64 63, i64 63>
   %a = shl <2 x i64> %v, %m
   ret <2 x i64> %a

 }
 
 define <2 x i64> @ashr_v2i64_late(<2 x i64> %v, i64 %x) {
 ; CHECK-LABEL: ashr_v2i64_late:
 ; CHECK:         .functype ashr_v2i64_late (v128, i64) -> (v128)
 ; CHECK-NEXT:  # %bb.0:
 ; CHECK-NEXT:    local.get 0
 ; CHECK-NEXT:    local.get 1
-; CHECK-NEXT:    i64x2.splat
-; CHECK-NEXT:    v128.const 63, 63
-; CHECK-NEXT:    v128.and
-; CHECK-NEXT:    i64x2.extract_lane 0
 ; CHECK-NEXT:    i32.wrap_i64
 ; CHECK-NEXT:    i64x2.shr_s
 ; CHECK-NEXT:    # fallthrough-return
   %t = insertelement <2 x i64> undef, i64 %x, i32 0
   %s = shufflevector <2 x i64> %t, <2 x i64> undef, <2 x i32> <i32 0, i32 0>
   %m = and <2 x i64> %s, <i64 63, i64 63>
   %a = ashr <2 x i64> %v, %m
   ret <2 x i64> %a

 }
 
 define <2 x i64> @lshr_v2i64_late(<2 x i64> %v, i64 %x) {
 ; CHECK-LABEL: lshr_v2i64_late:
 ; CHECK:         .functype lshr_v2i64_late (v128, i64) -> (v128)
 ; CHECK-NEXT:  # %bb.0:
 ; CHECK-NEXT:    local.get 0
 ; CHECK-NEXT:    local.get 1
-; CHECK-NEXT:    i64x2.splat
-; CHECK-NEXT:    v128.const 63, 63
-; CHECK-NEXT:    v128.and
-; CHECK-NEXT:    i64x2.extract_lane 0
 ; CHECK-NEXT:    i32.wrap_i64
 ; CHECK-NEXT:    i64x2.shr_u
 ; CHECK-NEXT:    # fallthrough-return
   %t = insertelement <2 x i64> undef, i64 %x, i32 0
   %s = shufflevector <2 x i64> %t, <2 x i64> undef, <2 x i32> <i32 0, i32 0>
   %m = and <2 x i64> %s, <i64 63, i64 63>
   %a = lshr <2 x i64> %v, %m
   ret <2 x i64> %a
 }