[DAGCombiner] Teach scalarizeBinOpOfSplats handle scalable splat.

[jacquesguan]

Summary

This revision supports to scalarize a binary operation of two scalable
splat vectors.

[RKSimon]

precommit the new tests and rebase so the patch shows the diff?

[jacquesguan]

rebase the precommit test.

[jacquesguan]

ping.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23501

Not sure we should be creating SPLAT_VECTOR on targets that don't support SPLAT_VECTOR. At the very least it's a waste of compile time to convert it back to a BUILD_VECTOR.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23471

Use DAG.isSplatValue instead?

23501

+1 at the very least we need a hasOperation() check for the getSplatVector path and getBuildVector fallback

[jacquesguan]

Address comment.

[jacquesguan]

Fix test case.

[jacquesguan]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23471

I tried to use DAG.isSplatValue, but this way would affect many no splat cases of X86 target, some of them are negative.

23501

Done, I only build splat_vector when is scalable now.

[jacquesguan]

ping.

[RKSimon]

LGTM with one minor

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23471

OK - please can you add a TODO comment? It sounds like something that needs revisiting in the future.

[jacquesguan]

Address comment.

[abrachet]

We're seeing failures after this revision https://ci.chromium.org/ui/p/fuchsia/builders/toolchain.ci/clang-linux-x64/b8808242343748276609/overview, as well as upstream bots https://lab.llvm.org/buildbot/#/builders/188/builds/17108

[jacquesguan]

Fix tests.

[jacquesguan]

In D122791#3661574, @abrachet wrote:

We're seeing failures after this revision https://ci.chromium.org/ui/p/fuchsia/builders/toolchain.ci/clang-linux-x64/b8808242343748276609/overview, as well as upstream bots https://lab.llvm.org/buildbot/#/builders/188/builds/17108

Thanks a lot, and I have fixed the tests.

[jacquesguan]

rebase main

[jacquesguan]

recommit in https://reviews.llvm.org/rGe60eb7053dc91a23cc3b96652aeda7513dd4b082

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

static SDValue scalarizeBinOpOfSplats(SDNode *N, SelectionDAG &DAG,
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
 
   // TODO: Remove/replace the extract cost check? If the elements are available
   //       as scalars, then there may be no extract cost. Should we ask if
   //       inserting a scalar back into a vector is cheap instead?
   int Index0, Index1;
   SDValue Src0 = DAG.getSplatSourceVector(N0, Index0);
   SDValue Src1 = DAG.getSplatSourceVector(N1, Index1);
+  // Extract element from splat_vector should be free.
+  // TODO: use DAG.isSplatValue instead?
+  bool IsBothSplatVector = N0.getOpcode() == ISD::SPLAT_VECTOR &&
+                           N1.getOpcode() == ISD::SPLAT_VECTOR;
   if (!Src0 || !Src1 || Index0 != Index1 ||
       Src0.getValueType().getVectorElementType() != EltVT ||
       Src1.getValueType().getVectorElementType() != EltVT ||
-      !TLI.isExtractVecEltCheap(VT, Index0) ||
+      !(IsBothSplatVector || TLI.isExtractVecEltCheap(VT, Index0)) ||
       !TLI.isOperationLegalOrCustom(Opcode, EltVT))
     return SDValue();
 
   SDValue IndexC = DAG.getVectorIdxConstant(Index0, DL);
   SDValue X = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, EltVT, Src0, IndexC);
   SDValue Y = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, EltVT, Src1, IndexC);
   SDValue ScalarBO = DAG.getNode(Opcode, DL, EltVT, X, Y, N->getFlags());
 
   // If all lanes but 1 are undefined, no need to splat the scalar result.
   // TODO: Keep track of undefs and use that info in the general case.
   if (N0.getOpcode() == ISD::BUILD_VECTOR && N0.getOpcode() == N1.getOpcode() &&
       count_if(N0->ops(), [](SDValue V) { return !V.isUndef(); }) == 1 &&
       count_if(N1->ops(), [](SDValue V) { return !V.isUndef(); }) == 1) {
     // bo (build_vec ..undef, X, undef...), (build_vec ..undef, Y, undef...) -->
     // build_vec ..undef, (bo X, Y), undef...
     SmallVector<SDValue, 8> Ops(VT.getVectorNumElements(), DAG.getUNDEF(EltVT));
     Ops[Index0] = ScalarBO;
     return DAG.getBuildVector(VT, DL, Ops);
   }
 
   // bo (splat X, Index), (splat Y, Index) --> splat (bo X, Y), Index
+  if (VT.isScalableVector())
+    return DAG.getSplatVector(VT, DL, ScalarBO);
   SmallVector<SDValue, 8> Ops(VT.getVectorNumElements(), ScalarBO);
   return DAG.getBuildVector(VT, DL, Ops);
 }
 
 /// Visit a binary vector operation, like ADD.
 SDValue DAGCombiner::SimplifyVBinOp(SDNode *N, const SDLoc &DL) {
   EVT VT = N->getValueType(0);
   assert(VT.isVector() && "SimplifyVBinOp only works on vectors!");
 
   SDValue LHS = N->getOperand(0);

llvm/test/CodeGen/RISCV/rvv/vadd-sdnode.ll

 ; RV32-NEXT:    sw a2, 8(sp)
 ; RV32-NEXT:    vlse64.v v16, (a0), zero
 ; RV32-NEXT:    vadd.vv v8, v8, v16
 ; RV32-NEXT:    addi sp, sp, 16
 ; RV32-NEXT:    ret
 ;
 ; RV64-LABEL: vadd_xx_nxv8i64:
 ; RV64:       # %bb.0:
-; RV64-NEXT:    vsetvli a2, zero, e64, m8, ta, mu
+; RV64-NEXT:    add a0, a0, a1
+; RV64-NEXT:    vsetvli a1, zero, e64, m8, ta, mu
 ; RV64-NEXT:    vmv.v.x v8, a0
-; RV64-NEXT:    vadd.vx v8, v8, a1
 ; RV64-NEXT:    ret
   %head1 = insertelement <vscale x 8 x i64> poison, i64 %a, i32 0
   %splat1 = shufflevector <vscale x 8 x i64> %head1, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
   %head2 = insertelement <vscale x 8 x i64> poison, i64 %b, i32 0
   %splat2 = shufflevector <vscale x 8 x i64> %head2, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
   %v = add <vscale x 8 x i64> %splat1, %splat2
   ret <vscale x 8 x i64> %v
 }

llvm/test/CodeGen/RISCV/rvv/vand-sdnode.ll

 ; RV32-NEXT:    sw a2, 8(sp)
 ; RV32-NEXT:    vlse64.v v16, (a0), zero
 ; RV32-NEXT:    vand.vv v8, v8, v16
 ; RV32-NEXT:    addi sp, sp, 16
 ; RV32-NEXT:    ret
 ;
 ; RV64-LABEL: vand_xx_nxv8i64:
 ; RV64:       # %bb.0:
-; RV64-NEXT:    vsetvli a2, zero, e64, m8, ta, mu
+; RV64-NEXT:    and a0, a0, a1
+; RV64-NEXT:    vsetvli a1, zero, e64, m8, ta, mu
 ; RV64-NEXT:    vmv.v.x v8, a0
-; RV64-NEXT:    vand.vx v8, v8, a1
 ; RV64-NEXT:    ret
   %head1 = insertelement <vscale x 8 x i64> poison, i64 %a, i32 0
   %splat1 = shufflevector <vscale x 8 x i64> %head1, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
   %head2 = insertelement <vscale x 8 x i64> poison, i64 %b, i32 0
   %splat2 = shufflevector <vscale x 8 x i64> %head2, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
   %v = and <vscale x 8 x i64> %splat1, %splat2
   ret <vscale x 8 x i64> %v
 }

llvm/test/CodeGen/RISCV/rvv/vmul-sdnode.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=riscv32 -mattr=+v -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,RV32
-; RUN: llc -mtriple=riscv64 -mattr=+v -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,RV64
+; RUN: llc -mtriple=riscv64 -mattr=+v -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,RV64,RV64NOM
 
 ; RUN: llc -mtriple=riscv32 -mattr=+v,+m -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,RV32
-; RUN: llc -mtriple=riscv64 -mattr=+v,+m -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,RV64
+; RUN: llc -mtriple=riscv64 -mattr=+v,+m -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,RV64,RV64M
 
 define <vscale x 1 x i8> @vmul_vv_nxv1i8(<vscale x 1 x i8> %va, <vscale x 1 x i8> %vb) {
 ; CHECK-LABEL: vmul_vv_nxv1i8:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e8, mf8, ta, mu
 ; CHECK-NEXT:    vmul.vv v8, v8, v9
 ; CHECK-NEXT:    ret
   %vc = mul <vscale x 1 x i8> %va, %vb

 ; RV32-NEXT:    vlse64.v v8, (a0), zero
 ; RV32-NEXT:    sw a3, 12(sp)
 ; RV32-NEXT:    sw a2, 8(sp)
 ; RV32-NEXT:    vlse64.v v16, (a0), zero
 ; RV32-NEXT:    vmul.vv v8, v8, v16
 ; RV32-NEXT:    addi sp, sp, 16
 ; RV32-NEXT:    ret
 ;
-; RV64-LABEL: vmul_xx_nxv8i64:
-; RV64:       # %bb.0:
-; RV64-NEXT:    vsetvli a2, zero, e64, m8, ta, mu
-; RV64-NEXT:    vmv.v.x v8, a0
-; RV64-NEXT:    vmul.vx v8, v8, a1
-; RV64-NEXT:    ret
+; RV64NOM-LABEL: vmul_xx_nxv8i64:
+; RV64NOM:       # %bb.0:
+; RV64NOM-NEXT:    vsetvli a2, zero, e64, m8, ta, mu
+; RV64NOM-NEXT:    vmv.v.x v8, a0
+; RV64NOM-NEXT:    vmul.vx v8, v8, a1
+; RV64NOM-NEXT:    ret
+;
+; RV64M-LABEL: vmul_xx_nxv8i64:
+; RV64M:       # %bb.0:
+; RV64M-NEXT:    mul a0, a0, a1
+; RV64M-NEXT:    vsetvli a1, zero, e64, m8, ta, mu
+; RV64M-NEXT:    vmv.v.x v8, a0
+; RV64M-NEXT:    ret
   %head1 = insertelement <vscale x 8 x i64> poison, i64 %a, i32 0
   %splat1 = shufflevector <vscale x 8 x i64> %head1, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
   %head2 = insertelement <vscale x 8 x i64> poison, i64 %b, i32 0
   %splat2 = shufflevector <vscale x 8 x i64> %head2, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
   %v = mul <vscale x 8 x i64> %splat1, %splat2
   ret <vscale x 8 x i64> %v
 }

llvm/test/CodeGen/RISCV/rvv/vor-sdnode.ll

 ; RV32-NEXT:    sw a2, 8(sp)
 ; RV32-NEXT:    vlse64.v v16, (a0), zero
 ; RV32-NEXT:    vor.vv v8, v8, v16
 ; RV32-NEXT:    addi sp, sp, 16
 ; RV32-NEXT:    ret
 ;
 ; RV64-LABEL: vor_xx_nxv8i64:
 ; RV64:       # %bb.0:
-; RV64-NEXT:    vsetvli a2, zero, e64, m8, ta, mu
+; RV64-NEXT:    or a0, a0, a1
+; RV64-NEXT:    vsetvli a1, zero, e64, m8, ta, mu
 ; RV64-NEXT:    vmv.v.x v8, a0
-; RV64-NEXT:    vor.vx v8, v8, a1
 ; RV64-NEXT:    ret
   %head1 = insertelement <vscale x 8 x i64> poison, i64 %a, i32 0
   %splat1 = shufflevector <vscale x 8 x i64> %head1, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
   %head2 = insertelement <vscale x 8 x i64> poison, i64 %b, i32 0
   %splat2 = shufflevector <vscale x 8 x i64> %head2, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
   %v = or <vscale x 8 x i64> %splat1, %splat2
   ret <vscale x 8 x i64> %v
 }

llvm/test/CodeGen/RISCV/rvv/vsub-sdnode.ll

 ; RV32-NEXT:    sw a2, 8(sp)
 ; RV32-NEXT:    vlse64.v v16, (a0), zero
 ; RV32-NEXT:    vsub.vv v8, v8, v16
 ; RV32-NEXT:    addi sp, sp, 16
 ; RV32-NEXT:    ret
 ;
 ; RV64-LABEL: vsub_xx_nxv8i64:
 ; RV64:       # %bb.0:
-; RV64-NEXT:    vsetvli a2, zero, e64, m8, ta, mu
+; RV64-NEXT:    sub a0, a0, a1
+; RV64-NEXT:    vsetvli a1, zero, e64, m8, ta, mu
 ; RV64-NEXT:    vmv.v.x v8, a0
-; RV64-NEXT:    vsub.vx v8, v8, a1
 ; RV64-NEXT:    ret
   %head1 = insertelement <vscale x 8 x i64> poison, i64 %a, i32 0
   %splat1 = shufflevector <vscale x 8 x i64> %head1, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
   %head2 = insertelement <vscale x 8 x i64> poison, i64 %b, i32 0
   %splat2 = shufflevector <vscale x 8 x i64> %head2, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
   %v = sub <vscale x 8 x i64> %splat1, %splat2
   ret <vscale x 8 x i64> %v
 }

llvm/test/CodeGen/RISCV/rvv/vxor-sdnode.ll

 ; RV32-NEXT:    sw a2, 8(sp)
 ; RV32-NEXT:    vlse64.v v16, (a0), zero
 ; RV32-NEXT:    vxor.vv v8, v8, v16
 ; RV32-NEXT:    addi sp, sp, 16
 ; RV32-NEXT:    ret
 ;
 ; RV64-LABEL: vxor_xx_nxv8i64:
 ; RV64:       # %bb.0:
-; RV64-NEXT:    vsetvli a2, zero, e64, m8, ta, mu
+; RV64-NEXT:    xor a0, a0, a1
+; RV64-NEXT:    vsetvli a1, zero, e64, m8, ta, mu
 ; RV64-NEXT:    vmv.v.x v8, a0
-; RV64-NEXT:    vxor.vx v8, v8, a1
 ; RV64-NEXT:    ret
   %head1 = insertelement <vscale x 8 x i64> poison, i64 %a, i32 0
   %splat1 = shufflevector <vscale x 8 x i64> %head1, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
   %head2 = insertelement <vscale x 8 x i64> poison, i64 %b, i32 0
   %splat2 = shufflevector <vscale x 8 x i64> %head2, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
   %v = xor <vscale x 8 x i64> %splat1, %splat2
   ret <vscale x 8 x i64> %v
 }