[llvm] [AArch64][SVE] Support lowering fixed-length BUILD_VECTORS to ZIPs (PR #111698)

[Sander de Smalen via llvm-commits]

================
@@ -14384,24 +14385,67 @@ static SDValue ConstantBuildVector(SDValue Op, SelectionDAG &DAG,
   return SDValue();
 }
 
-SDValue AArch64TargetLowering::LowerBUILD_VECTOR(SDValue Op,
-                                                 SelectionDAG &DAG) const {
+SDValue AArch64TargetLowering::LowerFixedLengthBuildVectorToSVE(
+    SDValue Op, SelectionDAG &DAG) const {
   EVT VT = Op.getValueType();
+  SDLoc DL(Op);
+  EVT ContainerVT = getContainerForFixedLengthVector(DAG, VT);
+  auto *BVN = cast<BuildVectorSDNode>(Op);
 
-  if (useSVEForFixedLengthVectorVT(VT, !Subtarget->isNeonAvailable())) {
-    if (auto SeqInfo = cast<BuildVectorSDNode>(Op)->isConstantSequence()) {
-      SDLoc DL(Op);
-      EVT ContainerVT = getContainerForFixedLengthVector(DAG, VT);
-      SDValue Start = DAG.getConstant(SeqInfo->first, DL, ContainerVT);
-      SDValue Steps = DAG.getStepVector(DL, ContainerVT, SeqInfo->second);
-      SDValue Seq = DAG.getNode(ISD::ADD, DL, ContainerVT, Start, Steps);
-      return convertFromScalableVector(DAG, Op.getValueType(), Seq);
-    }
+  if (auto SeqInfo = BVN->isConstantSequence()) {
+    SDValue Start = DAG.getConstant(SeqInfo->first, DL, ContainerVT);
+    SDValue Steps = DAG.getStepVector(DL, ContainerVT, SeqInfo->second);
+    SDValue Seq = DAG.getNode(ISD::ADD, DL, ContainerVT, Start, Steps);
+    return convertFromScalableVector(DAG, VT, Seq);
+  }
 
-    // Revert to common legalisation for all other variants.
+  if (!VT.isPow2VectorType() || VT.getFixedSizeInBits() > 128 ||
+      VT.getVectorNumElements() <= 1 || BVN->isConstant())
     return SDValue();
+
+  // Lower (pow2) BUILD_VECTORS that are <= 128-bit to a sequence of ZIP1s.
+  EVT ZipVT = ContainerVT;
+  SDValue ZeroI64 = DAG.getConstant(0, DL, MVT::i64);
+  SmallVector<SDValue, 16> Intermediates =
+      llvm::map_to_vector<16>(Op->op_values(), [&](SDValue Op) {
+        SDValue Undef = DAG.getUNDEF(ZipVT);
+        return Op.isUndef() ? Undef
+                            : DAG.getNode(ISD::INSERT_VECTOR_ELT, DL, ZipVT,
+                                          Undef, Op, ZeroI64);
+      });
+
+  while (Intermediates.size() > 1) {
+    auto ToZipVT = [&](SDValue Op) { return DAG.getBitcast(ZipVT, Op); };
+    for (unsigned I = 0; I < Intermediates.size(); I += 2) {
+      SDValue Op0 = Intermediates[I + 0];
+      SDValue Op1 = Intermediates[I + 1];
+      Intermediates[I / 2] = Op1.isUndef()
+                                 ? Op0
+                                 : DAG.getNode(AArch64ISD::ZIP1, DL, ZipVT,
+                                               ToZipVT(Op0), ToZipVT(Op1));
+    }
----------------
sdesmalen-arm wrote:

nit:
```suggestion
    for (unsigned I = 0; I < Intermediates.size(); I += 2) {
      SDValue Op0 = DAG.getBitcast(ZipVT, Intermediates[I + 0]);
      SDValue Op1 = DAG.getBitcast(ZipVT, Intermediates[I + 1]);
      Intermediates[I / 2] =
          Op1.isUndef() ? Op0
                        : DAG.getNode(AArch64ISD::ZIP1, DL, ZipVT, Op0, Op1);
    }
```

https://github.com/llvm/llvm-project/pull/111698