[SelectionDAG][WebAssembly] Recognize splat value ABS operations

[tlively]

Summary

This patch gives SelectionDAG::isSplatValue the ability to look
through ABS nodes, which allows the combine introduced in D83602 to
scalarize ABS operations when only a single lane is needed.

WebAssembly does not support ABS natively as a scalar operation, but
we still want it to be scalarized and expanded, so this patch also
adds a custom combine in the WebAssembly backend that scalarizes ABS
specifically when it is a splat and only one lane is needed. This is
useful when scalarizing WebAssembly shift values, which often contain
ABS operations in Halide output.

Depends on D83602.

[RKSimon]

one minor in SelectionDAG - but a wasm dev needs to looks at the target code

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

2317

Unnecessary braces around the if()

And if you move the getOperand(0) inside the isSplatValue() call then the other braces can go as well.

[tlively]

• Remove superfluous parentheses

[aheejin]

Wouldn't we want to add ISD::ABS to scalarizeSplatValue method in DAGCombiner.cpp too? WebAssembly wouldn't benefit from it, but it can benefit other targets that have a scalar abs instruction, and it is also consistent.

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

1749

Why is this?

1758

Does this function run before DAGCombiner::visitEXTRACT_VECTOR_ELT or after it? Or does this replace it? If this runs once before or after it, we wouldn't be able to optimize patterns that mixes abs with add or sub, such as abs(add(abs(v1) + abs(v2))), right?

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

2316

I don't think you need the fall through

case ISD::ABS:
  return isSplatValue(V.getOperand(0), DemandedElts, UndefElts);

[RKSimon]

@tlively Not sure what plan you have for this patch, but I've created D98778 that handles the generic ISD::ABS splat handling in SelectionDAG::isSplatValue

[RKSimon]

In D83605#2631635, @RKSimon wrote:

@tlively Not sure what plan you have for this patch, but I've created D98778 that handles the generic ISD::ABS splat handling in SelectionDAG::isSplatValue

The generic ISD::ABS splat handling has landed - it could be that we can generically fold binop(splat(),splat()) -> splat(binop()) as well

[tlively]

Thanks for working on this @RKSimon! FWIW, I have no problem tossing out or overhauling these patches if related work lands upstream, so please don't hesitate to make them redundant.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

bool SelectionDAG::isSplatValue(SDValue V, const APInt &DemandedElts,
   if (!VT.isScalableVector() && !DemandedElts)
     return false; // No demanded elts, better to assume we don't know anything.
 
   // Deal with some common cases here that work for both fixed and scalable
   // vector types.
   switch (V.getOpcode()) {
   case ISD::SPLAT_VECTOR:
     return true;
+  case ISD::ABS:
+    if (isSplatValue(V.getOperand(0), DemandedElts, UndefElts))
+      return true;
+    break;
   case ISD::ADD:
   case ISD::SUB:
   case ISD::AND: {
     APInt UndefLHS, UndefRHS;
     SDValue LHS = V.getOperand(0);
     SDValue RHS = V.getOperand(1);
     if (isSplatValue(LHS, DemandedElts, UndefLHS) &&
         isSplatValue(RHS, DemandedElts, UndefRHS)) {
       UndefElts = UndefLHS | UndefRHS;

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

if (Subtarget->hasSIMD128())
         setOperationAction(Op, T, Expand);
   }
 
   // SIMD-specific configuration
   if (Subtarget->hasSIMD128()) {
     // Hoist bitcasts out of shuffles
     setTargetDAGCombine(ISD::VECTOR_SHUFFLE);
 
+    // Scalarize extracts of abs splats
+    setTargetDAGCombine(ISD::EXTRACT_VECTOR_ELT);
+
     // 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);

 }
 
 //===----------------------------------------------------------------------===//
 //   Custom DAG combine hooks
 //===----------------------------------------------------------------------===//
 static SDValue
 performVECTOR_SHUFFLECombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI) {
   auto &DAG = DCI.DAG;
-  auto Shuffle = cast<ShuffleVectorSDNode>(N);
+  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.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
+performEXTRACT_VECTOR_ELTCombine(SDNode *N,
+                                 TargetLowering::DAGCombinerInfo &DCI) {
+  auto &DAG = DCI.DAG;
+  SDLoc DL(N);
+  EVT ScalarVT = N->getValueType(0);
+  SDValue VecOp = N->getOperand(0);
+  SDValue IndexOp = N->getOperand(1);
+
+  // Scalarize extracts of splatted ABS nodes even though we do not support a
+  // scalar ABS instruction. Emulating it with scalars is still probably better
+  // than using vector ops that would otherwise be unnecessary.
+  if (VecOp.getOpcode() == ISD::ABS && VecOp.hasOneUse() &&
+      !DCI.isAfterLegalizeDAG() &&
+      DAG.isSplatValue(VecOp, /*AllowUndefs=*/true)) {
+    SDValue AbsOp = VecOp.getOperand(0);
+    SDValue NewExt =
+        DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, ScalarVT, AbsOp, IndexOp);
+    return DAG.getNode(ISD::ABS, DL, ScalarVT, NewExt);
+  }
+
+  return SDValue();
+}
+
 SDValue
 WebAssemblyTargetLowering::PerformDAGCombine(SDNode *N,
                                              DAGCombinerInfo &DCI) const {
   switch (N->getOpcode()) {
   default:
     return SDValue();
   case ISD::VECTOR_SHUFFLE:
     return performVECTOR_SHUFFLECombine(N, DCI);
+  case ISD::EXTRACT_VECTOR_ELT:
+    return performEXTRACT_VECTOR_ELTCombine(N, DCI);
   }
 }

llvm/test/CodeGen/WebAssembly/simd-shift-complex-splats.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 SIMD shifts can be lowered correctly even with shift
 ; values that are more complex than plain splats.
 
 target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
 target triple = "wasm32-unknown-unknown"
 
-;; TODO: Optimize this further by scalarizing the add
-
 ; CHECK-LABEL: shl_add:
 ; CHECK-NEXT: .functype shl_add (v128, i32, i32) -> (v128)
 ; CHECK-NEXT: i32.add $push0=, $1, $2
 ; CHECK-NEXT: i8x16.shl $push1=, $0, $pop0
 ; CHECK-NEXT: return $pop1
 define <16 x i8> @shl_add(<16 x i8> %v, i8 %a, i8 %b) {
   %t1 = insertelement <16 x i8> undef, i8 %a, i32 0
   %va = shufflevector <16 x i8> %t1, <16 x i8> undef, <16 x i32> zeroinitializer
   %t2 = insertelement <16 x i8> undef, i8 %b, i32 0
   %vb = shufflevector <16 x i8> %t2, <16 x i8> undef, <16 x i32> zeroinitializer
   %shift = add <16 x i8> %va, %vb
   %r = shl <16 x i8> %v, %shift
   ret <16 x i8> %r
 }
 
 ; CHECK-LABEL: shl_abs:
 ; CHECK-NEXT: .functype shl_abs (v128, i32) -> (v128)
-; CHECK-NEXT: i8x16.extract_lane_u $push8=, $0, 0
-; CHECK-NEXT: i8x16.splat $push0=, $1
-; CHECK-NEXT: i8x16.abs $push98=, $pop0
-; CHECK-NEXT: local.tee $push97=, $2=, $pop98
-; CHECK-NEXT: i8x16.extract_lane_u $push6=, $pop97, 0
-; CHECK-NEXT: i32.const $push2=, 7
-; CHECK-NEXT: i32.and $push7=, $pop6, $pop2
-; CHECK-NEXT: i32.shl $push9=, $pop8, $pop7
-; CHECK-NEXT: i8x16.splat $push10=, $pop9
-; CHECK-NEXT: i8x16.extract_lane_u $push4=, $0, 1
-; CHECK-NEXT: i8x16.extract_lane_u $push1=, $2, 1
-; CHECK-NEXT: i32.const $push96=, 7
-; CHECK-NEXT: i32.and $push3=, $pop1, $pop96
-; CHECK-NEXT: i32.shl $push5=, $pop4, $pop3
-; CHECK-NEXT: i8x16.replace_lane $push11=, $pop10, 1, $pop5
-; ...
-; CHECK:      i8x16.extract_lane_u $push79=, $0, 15
-; CHECK-NEXT: i8x16.extract_lane_u $push77=, $2, 15
-; CHECK-NEXT: i32.const $push82=, 7
-; CHECK-NEXT: i32.and $push78=, $pop77, $pop82
-; CHECK-NEXT: i32.shl $push80=, $pop79, $pop78
-; CHECK-NEXT: i8x16.replace_lane $push81=, $pop76, 15, $pop80
-; CHECK-NEXT: return $pop81
+; CHECK-NEXT: i32.const $push0=, 31
+; CHECK-NEXT: i32.shr_s $push5=, $1, $pop0
+; CHECK-NEXT: local.tee $push4=, $2=, $pop5
+; CHECK-NEXT: i32.add $push1=, $1, $pop4
+; CHECK-NEXT: i32.xor $push2=, $pop1, $2
+; CHECK-NEXT: i8x16.shl $push3=, $0, $pop2
+; CHECK-NEXT: return $pop3
 define <16 x i8> @shl_abs(<16 x i8> %v, i8 %a) {
   %t1 = insertelement <16 x i8> undef, i8 %a, i32 0
   %va = shufflevector <16 x i8> %t1, <16 x i8> undef, <16 x i32> zeroinitializer
   %nva = sub <16 x i8> zeroinitializer, %va
   %c = icmp sgt <16 x i8> %va, zeroinitializer
   %shift = select <16 x i1> %c, <16 x i8> %va, <16 x i8> %nva
   %r = shl <16 x i8> %v, %shift
   ret <16 x i8> %r
 }
 
+;; TODO: Complete scalarization by removing the unnecessary shuffle
+
 ; CHECK-LABEL: shl_abs_add:
 ; CHECK-NEXT: .functype shl_abs_add (v128, i32, i32) -> (v128)
-; CHECK-NEXT: i8x16.extract_lane_u $push11=, $0, 0
 ; CHECK-NEXT: i8x16.splat $push1=, $1
 ; CHECK-NEXT: i8x16.splat $push0=, $2
 ; CHECK-NEXT: i8x16.add $push2=, $pop1, $pop0
 ; CHECK-NEXT: v8x16.shuffle $push3=, $pop2, $0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
-; CHECK-NEXT: i8x16.abs $push101=, $pop3
-; CHECK-NEXT: local.tee $push100=, $3=, $pop101
-; CHECK-NEXT: i8x16.extract_lane_u $push9=, $pop100, 0
-; CHECK-NEXT: i32.const $push5=, 7
-; CHECK-NEXT: i32.and $push10=, $pop9, $pop5
-; CHECK-NEXT: i32.shl $push12=, $pop11, $pop10
-; CHECK-NEXT: i8x16.splat $push13=, $pop12
-; CHECK-NEXT: i8x16.extract_lane_u $push7=, $0, 1
-; CHECK-NEXT: i8x16.extract_lane_u $push4=, $3, 1
-; CHECK-NEXT: i32.const $push99=, 7
-; CHECK-NEXT: i32.and $push6=, $pop4, $pop99
-; CHECK-NEXT: i32.shl $push8=, $pop7, $pop6
-; CHECK-NEXT: i8x16.replace_lane $push14=, $pop13, 1, $pop8
-; ...
-; CHECK:      i8x16.extract_lane_u $push82=, $0, 15
-; CHECK-NEXT: i8x16.extract_lane_u $push80=, $3, 15
-; CHECK-NEXT: i32.const $push85=, 7
-; CHECK-NEXT: i32.and $push81=, $pop80, $pop85
-; CHECK-NEXT: i32.shl $push83=, $pop82, $pop81
-; CHECK-NEXT: i8x16.replace_lane $push84=, $pop79, 15, $pop83
-; CHECK-NEXT: return $pop84
+; CHECK-NEXT: i8x16.extract_lane_u $push11=, $pop3, 0
+; CHECK-NEXT: local.tee $push10=, $1=, $pop11
+; CHECK-NEXT: i32.const $push4=, 31
+; CHECK-NEXT: i32.shr_s $push9=, $1, $pop4
+; CHECK-NEXT: local.tee $push8=, $1=, $pop9
+; CHECK-NEXT: i32.add $push5=, $pop10, $pop8
+; CHECK-NEXT: i32.xor $push6=, $pop5, $1
+; CHECK-NEXT: i8x16.shl $push7=, $0, $pop6
+; CHECK-NEXT: return $pop7
 define <16 x i8> @shl_abs_add(<16 x i8> %v, i8 %a, i8 %b) {
   %t1 = insertelement <16 x i8> undef, i8 %a, i32 0
   %va = shufflevector <16 x i8> %t1, <16 x i8> undef, <16 x i32> zeroinitializer
   %t2 = insertelement <16 x i8> undef, i8 %b, i32 0
   %vb = shufflevector <16 x i8> %t2, <16 x i8> undef, <16 x i32> zeroinitializer
   %vadd = add <16 x i8> %va, %vb
   %nvadd = sub <16 x i8> zeroinitializer, %vadd
   %c = icmp sgt <16 x i8> %vadd, zeroinitializer
   %shift = select <16 x i1> %c, <16 x i8> %vadd, <16 x i8> %nvadd
   %r = shl <16 x i8> %v, %shift
   ret <16 x i8> %r
 }