[AArch64] Use DUP for BUILD_VECTOR with few different elements.

[fhahn]

Summary

If most elements of BUILD_VECTOR are the same, with a few different
elements, it is better to use DUP for the common elements and
INSERT_VECTOR_ELT for the different elements.

Currently this transform is guarded quite restrictively to only trigger
in clearly beneficial cases.

With D90176, the lowering for patterns originating from code like
float32x4_t y = {a,a,a,0}; (common in 3D apps) are lowered even
better (unnecessary fmov is removed).

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9156

I think NumElts >= 4 is redundant? If NumElts == 2, there is no value of NumDifferentLanes that makes this pass.

llvm/test/CodeGen/AArch64/arm64-vector-insertion.ll

217–218

The test coverage here seems a bit lacking.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9192

This can be removed.

9204

Could this just create a DUP? Or is this so that it does not have to figure out what kind of DUP/DUPLANE to use?

llvm/test/CodeGen/AArch64/arm64-vector-insertion.ll

33–215

Comment can be updated too

[fhahn]

Remove unnecessary check, spurious comment and fixed TODO, thanks!

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9204

We could, but I am not sure how much extra logic would be required to handle constant and non-constant splats and other cases optimized by LowerBUILD_VECTOR. This seems to also be in-line with other parts of this function that also delegate back to LowerBUILD_VECTOR.

llvm/test/CodeGen/AArch64/arm64-vector-insertion.ll

33–215

yes, removed!

217–218

I added more positive and negative tests in 5febc535c4f8.

[efriedma]

Maybe add a testcase for -0 in case someone else makes the same mistake? Otherwise LGTM

[efriedma]

Maybe add a testcase for -0 in case someone else makes the same mistake?

Err, I guess the -0 thing was the other patch; nevermind. ;)

[fhahn]

Thank you very much for the review!

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

SDValue AArch64TargetLowering::LowerBUILD_VECTOR(SDValue Op,
   SDLoc dl(Op);
   unsigned NumElts = VT.getVectorNumElements();
   bool isOnlyLowElement = true;
   bool usesOnlyOneValue = true;
   bool usesOnlyOneConstantValue = true;
   bool isConstant = true;
   bool AllLanesExtractElt = true;
   unsigned NumConstantLanes = 0;
+  unsigned NumDifferentLanes = 0;
+  unsigned NumUndefLanes = 0;
   SDValue Value;
   SDValue ConstantValue;
   for (unsigned i = 0; i < NumElts; ++i) {
     SDValue V = Op.getOperand(i);
     if (V.getOpcode() != ISD::EXTRACT_VECTOR_ELT)
       AllLanesExtractElt = false;
-    if (V.isUndef())
+    if (V.isUndef()) {
+      ++NumUndefLanes;
       continue;
+    }
     if (i > 0)
       isOnlyLowElement = false;
     if (!isa<ConstantFPSDNode>(V) && !isa<ConstantSDNode>(V))
       isConstant = false;
 
     if (isa<ConstantSDNode>(V) || isa<ConstantFPSDNode>(V)) {
       ++NumConstantLanes;
       if (!ConstantValue.getNode())
         ConstantValue = V;
       else if (ConstantValue != V)
         usesOnlyOneConstantValue = false;
     }
 
     if (!Value.getNode())
       Value = V;
-    else if (V != Value)
+    else if (V != Value) {
       usesOnlyOneValue = false;
+      ++NumDifferentLanes;
+    }
   }
 
   if (!Value.getNode()) {
     LLVM_DEBUG(
         dbgs() << "LowerBUILD_VECTOR: value undefined, creating undef node\n");
     return DAG.getUNDEF(VT);
   }
 
if (VT.getVectorElementType().isFloatingPoint()) {
       LLVM_DEBUG(dbgs() << "LowerBUILD_VECTOR: trying to lower new vector: ";
                  Val.dump(););
       Val = LowerBUILD_VECTOR(Val, DAG);
       if (Val.getNode())
         return DAG.getNode(ISD::BITCAST, dl, VT, Val);
     }
   }
 
+  // If we need to insert a small number of different non-constant elements and
+  // the vector width is sufficiently large, prefer using DUP with the common
+  // value and INSERT_VECTOR_ELT for the different lanes. If DUP is preferred,
+  // skip the constant lane handling below.
+  bool PreferDUPAndInsert =
+      !isConstant && NumDifferentLanes >= 1 &&
+      NumDifferentLanes < ((NumElts - NumUndefLanes) / 2) &&
+      NumDifferentLanes >= NumConstantLanes;
+
   // If there was only one constant value used and for more than one lane,
   // start by splatting that value, then replace the non-constant lanes. This
   // is better than the default, which will perform a separate initialization
   // for each lane.
-  if (NumConstantLanes > 0 && usesOnlyOneConstantValue) {
+  if (!PreferDUPAndInsert && NumConstantLanes > 0 && usesOnlyOneConstantValue) {
     // Firstly, try to materialize the splat constant.
     SDValue Vec = DAG.getSplatBuildVector(VT, dl, ConstantValue),
             Val = ConstantBuildVector(Vec, DAG);
     if (!Val) {
       // Otherwise, materialize the constant and splat it.
       Val = DAG.getNode(AArch64ISD::DUP, dl, VT, ConstantValue);
       DAG.ReplaceAllUsesWith(Vec.getNode(), &Val);
     }
SDValue AArch64TargetLowering::LowerBUILD_VECTOR(SDValue Op,
 
   // This will generate a load from the constant pool.
   if (isConstant) {
     LLVM_DEBUG(
         dbgs() << "LowerBUILD_VECTOR: all elements are constant, use default "
                   "expansion\n");
     return SDValue();
   }
 
   // Empirical tests suggest this is rarely worth it for vectors of length <= 2.
   if (NumElts >= 4) {
     if (SDValue shuffle = ReconstructShuffle(Op, DAG))
       return shuffle;
   }
 
+  if (PreferDUPAndInsert) {
+    // First, build a constant vector with the common element.
+    SmallVector<SDValue, 8> Ops;
+    for (unsigned I = 0; I < NumElts; ++I)
+      Ops.push_back(Value);
+    SDValue NewVector = LowerBUILD_VECTOR(DAG.getBuildVector(VT, dl, Ops), DAG);
+    // Next, insert the elements that do not match the common value.
+    for (unsigned I = 0; I < NumElts; ++I)
+      if (Op.getOperand(I) != Value)
+        NewVector =
+            DAG.getNode(ISD::INSERT_VECTOR_ELT, dl, VT, NewVector,
+                        Op.getOperand(I), DAG.getConstant(I, dl, MVT::i64));
+
+    return NewVector;
+  }
+
   // If all else fails, just use a sequence of INSERT_VECTOR_ELT when we
   // know the default expansion would otherwise fall back on something even
   // worse. For a vector with one or two non-undef values, that's
   // scalar_to_vector for the elements followed by a shuffle (provided the
   // shuffle is valid for the target) and materialization element by element
   // on the stack followed by a load for everything else.
   if (!isConstant && !usesOnlyOneValue) {
     LLVM_DEBUG(

llvm/test/CodeGen/AArch64/arm64-vector-insertion.ll

 ; CHECK-NEXT:    str q1, [x0]
 ; CHECK-NEXT:    ret
 entry:
   %0 = insertelement <4 x float> <float undef, float 1.000000e+00, float 1.000000e+00, float 1.000000e+00>, float %a, i32 0
   %1 = bitcast float* %x to <4 x float>*
   store <4 x float> %0, <4 x float>* %1, align 16
   ret void
 }
 
-; TODO: This should jsut be a dup + clearing lane 4.
-
 define <16 x i8> @test_insert_v16i8_insert_1(i8 %a) {
 ; CHECK-LABEL: test_insert_v16i8_insert_1:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi.2d v0, #0000000000000000
 ; CHECK-NEXT:    mov.b v0[14], w0
 ; CHECK-NEXT:    ret
   %v.0 = insertelement <16 x i8> zeroinitializer, i8 %a, i32 14
   ret <16 x i8> %v.0
 }
 
 define <16 x i8> @test_insert_v16i8_insert_2(i8 %a) {
 ; CHECK-LABEL: test_insert_v16i8_insert_2:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi.2d v0, #0000000000000000
 ; CHECK-NEXT:    mov.b v0[1], w0
 ; CHECK-NEXT:    mov.b v0[2], w0
 ; CHECK-NEXT:    ret
   %v.0 = insertelement <16 x i8> zeroinitializer, i8 %a, i32 2
   %v.1 = insertelement <16 x i8> %v.0, i8 %a, i32 1
   ret <16 x i8> %v.1
 }
 
 define <16 x i8> @test_insert_v16i8_insert_2_undef_base(i8 %a) {
 ; CHECK-LABEL: test_insert_v16i8_insert_2_undef_base:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    movi.2d v0, #0000000000000000
-; CHECK-NEXT:    mov.b v0[0], w0
-; CHECK-NEXT:    mov.b v0[1], w0
-; CHECK-NEXT:    mov.b v0[2], w0
-; CHECK-NEXT:    mov.b v0[3], w0
-; CHECK-NEXT:    mov.b v0[4], w0
-; CHECK-NEXT:    mov.b v0[6], w0
-; CHECK-NEXT:    mov.b v0[7], w0
-; CHECK-NEXT:    mov.b v0[10], w0
-; CHECK-NEXT:    mov.b v0[11], w0
-; CHECK-NEXT:    mov.b v0[12], w0
-; CHECK-NEXT:    mov.b v0[13], w0
-; CHECK-NEXT:    mov.b v0[14], w0
-; CHECK-NEXT:    mov.b v0[15], w0
+; CHECK-NEXT:    dup.16b v0, w0
+; CHECK-NEXT:    mov.b v0[5], wzr
+; CHECK-NEXT:    mov.b v0[9], wzr
 ; CHECK-NEXT:    ret
   %v.0 = insertelement <16 x i8> <i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 0, i8 undef, i8 undef, i8 undef, i8 0, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef>  , i8 %a, i32 0
   %v.1 = insertelement <16 x i8> %v.0, i8 %a, i32 1
   %v.2 = insertelement <16 x i8> %v.1, i8 %a, i32 2
   %v.3 = insertelement <16 x i8> %v.2, i8 %a, i32 3
   %v.4 = insertelement <16 x i8> %v.3, i8 %a, i32 4
   %v.6 = insertelement <16 x i8> %v.4, i8 %a, i32 6
   %v.7 = insertelement <16 x i8> %v.6, i8 %a, i32 7
   %v.8 = insertelement <16 x i8> %v.7, i8 %a, i32 8
   %v.10 = insertelement <16 x i8> %v.7, i8 %a, i32 10
   %v.11 = insertelement <16 x i8> %v.10, i8 %a, i32 11
   %v.12 = insertelement <16 x i8> %v.11, i8 %a, i32 12
   %v.13 = insertelement <16 x i8> %v.12, i8 %a, i32 13
   %v.14 = insertelement <16 x i8> %v.13, i8 %a, i32 14
   %v.15 = insertelement <16 x i8> %v.14, i8 %a, i32 15
   ret <16 x i8> %v.15
 }
 
 define <16 x i8> @test_insert_v16i8_insert_2_undef_base_different_valeus(i8 %a, i8 %b) {
 ; CHECK-LABEL: test_insert_v16i8_insert_2_undef_base_different_valeus:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    movi.2d v0, #0000000000000000
-; CHECK-NEXT:    mov.b v0[0], w0
-; CHECK-NEXT:    mov.b v0[1], w0
+; CHECK-NEXT:    dup.16b v0, w0
 ; CHECK-NEXT:    mov.b v0[2], w1
-; CHECK-NEXT:    mov.b v0[3], w0
-; CHECK-NEXT:    mov.b v0[4], w0
-; CHECK-NEXT:    mov.b v0[6], w0
+; CHECK-NEXT:    mov.b v0[5], wzr
 ; CHECK-NEXT:    mov.b v0[7], w1
-; CHECK-NEXT:    mov.b v0[10], w0
-; CHECK-NEXT:    mov.b v0[11], w0
+; CHECK-NEXT:    mov.b v0[9], wzr
 ; CHECK-NEXT:    mov.b v0[12], w1
-; CHECK-NEXT:    mov.b v0[13], w0
-; CHECK-NEXT:    mov.b v0[14], w0
 ; CHECK-NEXT:    mov.b v0[15], w1
 ; CHECK-NEXT:    ret
   %v.0 = insertelement <16 x i8> <i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 0, i8 undef, i8 undef, i8 undef, i8 0, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef>  , i8 %a, i32 0
   %v.1 = insertelement <16 x i8> %v.0, i8 %a, i32 1
   %v.2 = insertelement <16 x i8> %v.1, i8 %b, i32 2
   %v.3 = insertelement <16 x i8> %v.2, i8 %a, i32 3
   %v.4 = insertelement <16 x i8> %v.3, i8 %a, i32 4
   %v.6 = insertelement <16 x i8> %v.4, i8 %a, i32 6
   %v.7 = insertelement <16 x i8> %v.6, i8 %b, i32 7
   %v.8 = insertelement <16 x i8> %v.7, i8 %a, i32 8
   %v.10 = insertelement <16 x i8> %v.7, i8 %a, i32 10
   %v.11 = insertelement <16 x i8> %v.10, i8 %a, i32 11
   %v.12 = insertelement <16 x i8> %v.11, i8 %b, i32 12
   %v.13 = insertelement <16 x i8> %v.12, i8 %a, i32 13
   %v.14 = insertelement <16 x i8> %v.13, i8 %a, i32 14
   %v.15 = insertelement <16 x i8> %v.14, i8 %b, i32 15
   ret <16 x i8> %v.15
 }
 
 define <8 x half> @test_insert_v8f16_insert_1(half %a) {
 ; CHECK-LABEL: test_insert_v8f16_insert_1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    movi.2d v1, #0000000000000000
 ; CHECK-NEXT:    // kill: def $h0 killed $h0 def $q0
-; CHECK-NEXT:    mov.h v1[0], v0[0]
-; CHECK-NEXT:    mov.h v1[1], v0[0]
-; CHECK-NEXT:    mov.h v1[2], v0[0]
-; CHECK-NEXT:    mov.h v1[3], v0[0]
-; CHECK-NEXT:    mov.h v1[4], v0[0]
-; CHECK-NEXT:    mov.h v1[5], v0[0]
-; CHECK-NEXT:    mov.h v1[6], v0[0]
-; CHECK-NEXT:    mov.16b v0, v1
+; CHECK-NEXT:    adrp x8, .LCPI6_0
+; CHECK-NEXT:    dup.8h v0, v0[0]
+; CHECK-NEXT:    add x8, x8, :lo12:.LCPI6_0
+; CHECK-NEXT:    ld1.h { v0 }[7], [x8]
 ; CHECK-NEXT:    ret
   %v.0 = insertelement <8 x half> <half undef, half undef, half undef, half undef, half undef, half undef, half undef, half 0.0>, half %a, i32 0
   %v.1 = insertelement <8 x half> %v.0, half %a, i32 1
   %v.2 = insertelement <8 x half> %v.1, half %a, i32 2
   %v.3 = insertelement <8 x half> %v.2, half %a, i32 3
   %v.4 = insertelement <8 x half> %v.3, half %a, i32 4
   %v.5 = insertelement <8 x half> %v.4, half %a, i32 5
   %v.6 = insertelement <8 x half> %v.5, half %a, i32 6
   ret <8 x half> %v.6
 }
 
 
 define <8 x half> @test_insert_v8f16_insert_2(half %a) {
 ; CHECK-LABEL: test_insert_v8f16_insert_2:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi.2d v1, #0000000000000000
 ; CHECK-NEXT:    // kill: def $h0 killed $h0 def $q0
 ; CHECK-NEXT:    mov.h v1[1], v0[0]
 ; CHECK-NEXT:    mov.h v1[2], v0[0]
 ; CHECK-NEXT:    mov.16b v0, v1
 ; CHECK-NEXT:    ret
   %v.0 = insertelement <8 x half> zeroinitializer, half %a, i32 2
   %v.1 = insertelement <8 x half> %v.0, half %a, i32 1
   ret <8 x half> %v.1
 }
 
 define <8 x i16> @test_insert_v8i16_insert_2(i16 %a) {
 ; CHECK-LABEL: test_insert_v8i16_insert_2:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    movi.2d v0, #0000000000000000
-; CHECK-NEXT:    mov.h v0[0], w0
-; CHECK-NEXT:    mov.h v0[1], w0
-; CHECK-NEXT:    mov.h v0[2], w0
-; CHECK-NEXT:    mov.h v0[4], w0
-; CHECK-NEXT:    mov.h v0[5], w0
-; CHECK-NEXT:    mov.h v0[6], w0
+; CHECK-NEXT:    dup.8h v0, w0
+; CHECK-NEXT:    mov.h v0[3], wzr
+; CHECK-NEXT:    mov.h v0[7], wzr
 ; CHECK-NEXT:    ret
   %v.0 = insertelement <8 x i16> <i16 undef, i16 undef, i16 undef, i16 0, i16 undef, i16 undef, i16 undef, i16 0>, i16 %a, i32 0
   %v.1 = insertelement <8 x i16> %v.0, i16 %a, i32 1
   %v.2 = insertelement <8 x i16> %v.1, i16 %a, i32 2
   %v.3 = insertelement <8 x i16> %v.2, i16 %a, i32 4
   %v.4 = insertelement <8 x i16> %v.3, i16 %a, i32 5
   %v.5 = insertelement <8 x i16> %v.4, i16 %a, i32 6
   ret <8 x i16> %v.5
 }
 
 define <8 x i16> @test_insert_v8i16_insert_3(i16 %a) {
 ; CHECK-LABEL: test_insert_v8i16_insert_3:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    movi.2d v0, #0000000000000000
-; CHECK-NEXT:    mov.h v0[0], w0
-; CHECK-NEXT:    mov.h v0[2], w0
-; CHECK-NEXT:    mov.h v0[4], w0
-; CHECK-NEXT:    mov.h v0[5], w0
-; CHECK-NEXT:    mov.h v0[6], w0
+; CHECK-NEXT:    dup.8h v0, w0
+; CHECK-NEXT:    mov.h v0[1], wzr
+; CHECK-NEXT:    mov.h v0[3], wzr
+; CHECK-NEXT:    mov.h v0[7], wzr
 ; CHECK-NEXT:    ret
   %v.0 = insertelement <8 x i16> <i16 undef, i16 0, i16 undef, i16 0, i16 undef, i16 undef, i16 undef, i16 0>, i16 %a, i32 0
   %v.2 = insertelement <8 x i16> %v.0, i16 %a, i32 2
   %v.3 = insertelement <8 x i16> %v.2, i16 %a, i32 4
   %v.4 = insertelement <8 x i16> %v.3, i16 %a, i32 5
   %v.5 = insertelement <8 x i16> %v.4, i16 %a, i32 6
   ret <8 x i16> %v.5
 }
 
 define <8 x i16> @test_insert_v8i16_insert_4(i16 %a) {
 ; CHECK-LABEL: test_insert_v8i16_insert_4:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi.2d v0, #0000000000000000
 ; CHECK-NEXT:    mov.h v0[0], w0
 ; CHECK-NEXT:    mov.h v0[2], w0
 ; CHECK-NEXT:    mov.h v0[4], w0
 ; CHECK-NEXT:    mov.h v0[5], w0
 ; CHECK-NEXT:    ret
   %v.0 = insertelement <8 x i16> <i16 undef, i16 0, i16 undef, i16 0, i16 undef, i16 undef, i16 0, i16 0>, i16 %a, i32 0
   %v.2 = insertelement <8 x i16> %v.0, i16 %a, i32 2
   %v.3 = insertelement <8 x i16> %v.2, i16 %a, i32 4
   %v.4 = insertelement <8 x i16> %v.3, i16 %a, i32 5
   ret <8 x i16> %v.4
 }
 
 define <8 x i16> @test_insert_v8i16_insert_5(i16 %a) {
 ; CHECK-LABEL: test_insert_v8i16_insert_5:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi.2d v0, #0000000000000000
 ; CHECK-NEXT:    mov.h v0[0], w0
 ; CHECK-NEXT:    mov.h v0[4], w0
 ; CHECK-NEXT:    mov.h v0[5], w0
 ; CHECK-NEXT:    ret
   %v.0 = insertelement <8 x i16> <i16 undef, i16 0, i16 0, i16 0, i16 undef, i16 undef, i16 0, i16 0>, i16 %a, i32 0
   %v.3 = insertelement <8 x i16> %v.0, i16 %a, i32 4
   %v.4 = insertelement <8 x i16> %v.3, i16 %a, i32 5
   ret <8 x i16> %v.4
 }
 
 define <2 x float> @test_insert_v2f32_undef_zero_vector(float %a) {
 ; CHECK-LABEL: test_insert_v2f32_undef_zero_vector:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi d1, #0000000000000000
 ; CHECK-NEXT:    // kill: def $s0 killed $s0 def $q0
 ; CHECK-NEXT:    mov.s v1[1], v0[0]
 ; CHECK-NEXT:    mov.16b v0, v1
 ; CHECK-NEXT:    ret
   %v.0 = insertelement <2 x float> <float 0.000000e+00, float undef>, float %a, i32 1
   ret <2 x float> %v.0
 }
 
 define <4 x float> @test_insert_3_f32_undef_zero_vector(float %a) {
 ; CHECK-LABEL: test_insert_3_f32_undef_zero_vector:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    movi.2d v1, #0000000000000000
 ; CHECK-NEXT:    // kill: def $s0 killed $s0 def $q0
-; CHECK-NEXT:    mov.s v1[0], v0[0]
-; CHECK-NEXT:    mov.s v1[1], v0[0]
-; CHECK-NEXT:    mov.s v1[2], v0[0]
-; CHECK-NEXT:    mov.16b v0, v1
+; CHECK-NEXT:    fmov s1, wzr
+; CHECK-NEXT:    dup.4s v0, v0[0]
+; CHECK-NEXT:    mov.s v0[3], v1[0]
 ; CHECK-NEXT:    ret
   %v.0 = insertelement <4 x float> <float undef, float undef, float undef, float 0.000000e+00>, float %a, i32 0
   %v.1 = insertelement <4 x float> %v.0, float %a, i32 1
   %v.2 = insertelement <4 x float> %v.1, float %a, i32 2
   ret <4 x float> %v.2
 }
 
 define <4 x float> @test_insert_3_f32_undef(float %a) {