[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

9157

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

41

The test coverage here seems a bit lacking.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9193

This can be removed.

9205

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

Comment can be updated too

[fhahn]

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

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9205

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

yes, removed!

41

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) && NumElts >= 4 &&
+      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

entry:
 
   ; CHECK-LABEL: test1f
   ; CHECK: fmov.4s v[[TEMP:[0-9]+]], #1.0
   ; CHECK: mov.s v[[TEMP]][0], v0[0]
   ; CHECK: str q[[TEMP]], [x0]
   ; CHECK: ret
 }
 
 ; TODO: This should jsut be a dup + clearing lane 4.
 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
-; 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
 ;
 entry:
   %0 = insertelement <4 x float> <float undef, float undef, float undef, float 0.000000e+00>, float %a, i32 0
   %1 = insertelement <4 x float> %0, float %a, i32 1
   %vecinit3 = insertelement <4 x float> %1, float %a, i32 2
   ret <4 x float> %vecinit3
 }