[DAG] Fold neg(bvsplat(neg(x)) -> bvsplat(x)

[dmgreen]

Summary

This add as a fold of sub(0, buildvector_splat(sub(0, x))) -> buildvector_splat(x). This is something that can come up in the lowering of right shifts under AArch64, where we generate a shift left of a negated number.

[RKSimon]

Funnily enough I was wondering about this pattern the other day as a followup to D98778...

Should we always be folding unaryop(splat(x)) -> splat(unaryop(x)) if the unaryop is legal/custom on the scalar type? And then maybe extend that to binop(splat(x),splat(y)) -> splat(binop(x,y)) as well?

[dmgreen]

Rebase over D103756

[dmgreen]

In D103755#2800920, @RKSimon wrote:

Funnily enough I was wondering about this pattern the other day as a followup to D98778...

Should we always be folding unaryop(splat(x)) -> splat(unaryop(x)) if the unaryop is legal/custom on the scalar type? And then maybe extend that to binop(splat(x),splat(y)) -> splat(binop(x,y)) as well?

Maybe. Not sure. I always find it difficult to see when optimizations like that would be universally beneficial over a wide range of very different architectures, considering how different they can be. I can see that it would make this more general, but it seems easy to think of cases where it would make things worse.

For this case I think it would need to work with the truncated type, not the scalar element type. For a v8i16 the scalar element type would not be legal under AArch64. If it was using the element type it would only handle half of the tests changed here.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

3313

Can SelectionDAG::getSplatValue() be used here? If so, then this doesn't need to be limited to BUILD_VECTOR, as it seems this fold would work equally well for scalable vectors (which use SPLAT_VECTOR).

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

3313

Thanks for taking a look. Yeah, seems to work OK. I'll make it so.

[dmgreen]

Use getSplatValue.

[RKSimon]

LGTM - cheers

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

3312

nit: s/bvsplat/splat/g

llvm/test/CodeGen/AArch64/neon-shift-neg.ll

254

It seems odd to have a scalable-vector test in neon-shift-neg.ll ? (I actually can't find this test in latest HEAD, am I looking at the right diff?)

[dmgreen]

llvm/test/CodeGen/AArch64/neon-shift-neg.ll

254

Yeah.. The same pattern of negated shifts does not apply for SVE (which makes it a little less useful). I've added a more direct test in 77ae9b364a9d9b99501163761313cefbb345cea7.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (isNullOrNullSplat(N0)) {
     }
 
     // Convert 0 - abs(x).
     SDValue Result;
     if (N1->getOpcode() == ISD::ABS &&
         !TLI.isOperationLegalOrCustom(ISD::ABS, VT) &&
         TLI.expandABS(N1.getNode(), Result, DAG, true))
       return Result;
+
+    // Fold neg(bvsplat(neg(x)) -> bvsplat(x)
+    if (N1.getOpcode() == ISD::BUILD_VECTOR &&
+        llvm::all_of(N1->ops(),
+                     [&](SDValue Op) { return Op == N1.getOperand(0); }) &&
+        N1.getOperand(0)->getOpcode() == ISD::SUB &&
+        isNullConstant(N1.getOperand(0)->getOperand(0)))
+      return DAG.getSplatBuildVector(VT, DL, N1.getOperand(0)->getOperand(1));
   }
 
   // Canonicalize (sub -1, x) -> ~x, i.e. (xor x, -1)
   if (isAllOnesOrAllOnesSplat(N0))
     return DAG.getNode(ISD::XOR, DL, VT, N1, N0);
 
   // fold (A - (0-B)) -> A+B
   if (N1.getOpcode() == ISD::SUB && isNullOrNullSplat(N1.getOperand(0)))

llvm/test/CodeGen/AArch64/neon-shift-neg.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=aarch64-none-linux-gnu -mattr=+neon | FileCheck %s
 
 define <2 x i64> @shr64x2(<2 x i64> %a, i64 %b) {
 ; CHECK-LABEL: shr64x2:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    neg x8, x0
-; CHECK-NEXT:    dup v1.2d, x8
-; CHECK-NEXT:    neg v1.2d, v1.2d
+; CHECK-NEXT:    dup v1.2d, x0
 ; CHECK-NEXT:    sshl v0.2d, v0.2d, v1.2d
 ; CHECK-NEXT:    ret
 entry:
   %sub = sub nsw i64 0, %b
   %splat.splatinsert = insertelement <2 x i64> poison, i64 %sub, i32 0
   %splat.splat = shufflevector <2 x i64> %splat.splatinsert, <2 x i64> poison, <2 x i32> zeroinitializer
   %shr = ashr <2 x i64> %a, %splat.splat
   ret <2 x i64> %shr
 }
 
 define <4 x i32> @shr32x4(<4 x i32> %a, i32 %b) {
 ; CHECK-LABEL: shr32x4:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    neg w8, w0
-; CHECK-NEXT:    dup v1.4s, w8
-; CHECK-NEXT:    neg v1.4s, v1.4s
+; CHECK-NEXT:    dup v1.4s, w0
 ; CHECK-NEXT:    sshl v0.4s, v0.4s, v1.4s
 ; CHECK-NEXT:    ret
 entry:
   %sub = sub nsw i32 0, %b
   %splat.splatinsert = insertelement <4 x i32> poison, i32 %sub, i32 0
   %splat.splat = shufflevector <4 x i32> %splat.splatinsert, <4 x i32> poison, <4 x i32> zeroinitializer
   %shr = ashr <4 x i32> %a, %splat.splat
   ret <4 x i32> %shr
 }
 
 define <4 x i32> @shr32x4undef(<4 x i32> %a, i32 %b) {
 ; CHECK-LABEL: shr32x4undef:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    neg w8, w0
-; CHECK-NEXT:    dup v1.4s, w8
-; CHECK-NEXT:    neg v1.4s, v1.4s
+; CHECK-NEXT:    dup v1.4s, w0
 ; CHECK-NEXT:    sshl v0.4s, v0.4s, v1.4s
 ; CHECK-NEXT:    ret
 entry:
   %sub = sub nsw i32 0, %b
   %splat.splatinsert = insertelement <4 x i32> poison, i32 %sub, i32 0
   %splat.splat = shufflevector <4 x i32> %splat.splatinsert, <4 x i32> poison, <4 x i32> <i32 undef, i32 0, i32 0, i32 0>
   %shr = ashr <4 x i32> %a, %splat.splat
   ret <4 x i32> %shr
entry:
   %sh_prom = shufflevector <16 x i8> %0, <16 x i8> undef, <16 x i32> zeroinitializer
   %shr = ashr <16 x i8> %a, %sh_prom
   ret <16 x i8> %shr
 }
 
 define <1 x i64> @shr64x1(<1 x i64> %a, i64 %b) {
 ; CHECK-LABEL: shr64x1:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    neg x8, x0
-; CHECK-NEXT:    fmov d1, x8
-; CHECK-NEXT:    neg d1, d1
+; CHECK-NEXT:    fmov d1, x0
 ; CHECK-NEXT:    sshl d0, d0, d1
 ; CHECK-NEXT:    ret
 entry:
   %sub = sub nsw i64 0, %b
   %splat.splatinsert = insertelement <1 x i64> poison, i64 %sub, i32 0
   %shr = ashr <1 x i64> %a, %splat.splatinsert
   ret <1 x i64> %shr
 }
 
 define <2 x i32> @shr32x2(<2 x i32> %a, i32 %b) {
 ; CHECK-LABEL: shr32x2:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    neg w8, w0
-; CHECK-NEXT:    dup v1.2s, w8
-; CHECK-NEXT:    neg v1.2s, v1.2s
+; CHECK-NEXT:    dup v1.2s, w0
 ; CHECK-NEXT:    sshl v0.2s, v0.2s, v1.2s
 ; CHECK-NEXT:    ret
 entry:
   %sub = sub nsw i32 0, %b
   %splat.splatinsert = insertelement <2 x i32> poison, i32 %sub, i32 0
   %splat.splat = shufflevector <2 x i32> %splat.splatinsert, <2 x i32> poison, <2 x i32> zeroinitializer
   %shr = ashr <2 x i32> %a, %splat.splat
   ret <2 x i32> %shr

 ; CHECK-NEXT:    ushl v0.8b, v0.8b, v1.8b
 ; CHECK-NEXT:    ret
 entry:
   %sub = sub i8 0, %b
   %0 = insertelement <8 x i8> undef, i8 %sub, i32 0
   %sh_prom = shufflevector <8 x i8> %0, <8 x i8> undef, <8 x i32> zeroinitializer
   %shl = shl <8 x i8> %a, %sh_prom
   ret <8 x i8> %shl
 }