[AArch64][SelectionDAG] Support passing/returning scalable vectors with unusual types.

[efriedma]

Summary

This adds handling for two cases:

1. A scalable vector where the element type is promoted.
2. A scalable vector where the element count is odd (or more generally, not divisble by the element count of the part type).

(Some element types still don't work; for example, <vscale x 2 x i128>, or <vscale x 2 x fp128>.)

The change to SplitVecRes_INSERT_SUBVECTOR needs to be considered in the context of the discussion of "widening" a scalable vector.

[sdesmalen]

The patch is a bit tricky to follow because it has many changing parts, so it's not always clear that all the parts are covered by the tests, but I trust you've made all these changes for a reason.
The changes themselves seem sensible to me, so I'm happy to accept it. As the title suggests, the types are all quite unusual, so if there's parts that aren't bullet proof we can fix those when we run into them.

Does D106028 depend on this patch? (I'm thinking mostly about the widening)

[efriedma]

In D105591#2920820, @sdesmalen wrote:

The patch is a bit tricky to follow because it has many changing parts, so it's not always clear that all the parts are covered by the tests, but I trust you've made all these changes for a reason.

I'll experiment a bit to see if there's some way to separate out the legalization changes; the llvm.experimental.vector.extract tricks I figured out for D106028 might be usable for that. The rest is sort of tied together.

Does D106028 depend on this patch? (I'm thinking mostly about the widening)

There isn't an explicit dependency. The SplitVecRes_INSERT_SUBVECTOR here is a less general version of the one in D106028.

[efriedma]

In D105591#2920910, @efriedma wrote:

In D105591#2920820, @sdesmalen wrote:

The patch is a bit tricky to follow because it has many changing parts, so it's not always clear that all the parts are covered by the tests, but I trust you've made all these changes for a reason.

I'll experiment a bit to see if there's some way to separate out the legalization changes; the llvm.experimental.vector.extract tricks I figured out for D106028 might be usable for that. The rest is sort of tied together.

Oh, right... if I try that, I get into issues with the extract/insert themselves, which end up needing more code to make them work.

I think it's reasonable to commit this as a single unit, even if it might theoretically be possible to split it more.

[efriedma]

Oh, oops, just realized the LegalizeVectorTypes aren't necessary here, at least for the most basic cases. I'll commit without them.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

static SDValue getCopyFromPartsVector(SelectionDAG &DAG, const SDLoc &DL,
 
   // There is now one part, held in Val.  Correct it to match ValueVT.
   EVT PartEVT = Val.getValueType();
 
   if (PartEVT == ValueVT)
     return Val;
 
   if (PartEVT.isVector()) {
+    // Vector/Vector bitcast.
+    if (ValueVT.getSizeInBits() == PartEVT.getSizeInBits())
+      return DAG.getNode(ISD::BITCAST, DL, ValueVT, Val);
+
     // If the element type of the source/dest vectors are the same, but the
     // parts vector has more elements than the value vector, then we have a
     // vector widening case (e.g. <2 x float> -> <4 x float>).  Extract the
     // elements we want.
-    if (PartEVT.getVectorElementType() == ValueVT.getVectorElementType()) {
+    if (PartEVT.getVectorElementCount() != ValueVT.getVectorElementCount()) {
       assert((PartEVT.getVectorElementCount().getKnownMinValue() >
               ValueVT.getVectorElementCount().getKnownMinValue()) &&
              (PartEVT.getVectorElementCount().isScalable() ==
               ValueVT.getVectorElementCount().isScalable()) &&
              "Cannot narrow, it would be a lossy transformation");
-      return DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, ValueVT, Val,
-                         DAG.getVectorIdxConstant(0, DL));
+      PartEVT =
+          EVT::getVectorVT(*DAG.getContext(), PartEVT.getVectorElementType(),
+                           ValueVT.getVectorElementCount());
+      Val = DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, PartEVT, Val,
+      if (PartEVT == ValueVT)
+        return Val;
     }
 
-    // Vector/Vector bitcast.
-    if (ValueVT.getSizeInBits() == PartEVT.getSizeInBits())
-      return DAG.getNode(ISD::BITCAST, DL, ValueVT, Val);
-
-    assert(PartEVT.getVectorElementCount() == ValueVT.getVectorElementCount() &&
-      "Cannot handle this kind of promotion");
     // Promoted vector extract
     return DAG.getAnyExtOrTrunc(Val, DL, ValueVT);
-
   }
 
   // Trivial bitcast if the types are the same size and the destination
   // vector type is legal.
   if (PartEVT.getSizeInBits() == ValueVT.getSizeInBits() &&
       TLI.isTypeLegal(ValueVT))
     return DAG.getNode(ISD::BITCAST, DL, ValueVT, Val);
 
static void getCopyToPartsVector(SelectionDAG &DAG, const SDLoc &DL,
   EVT BuiltVectorTy = EVT::getVectorVT(
       *DAG.getContext(), IntermediateVT.getScalarType(), DestEltCnt.getValue());
 
   if (ValueVT == BuiltVectorTy) {
     // Nothing to do.
   } else if (ValueVT.getSizeInBits() == BuiltVectorTy.getSizeInBits()) {
     // Bitconvert vector->vector case.
     Val = DAG.getNode(ISD::BITCAST, DL, BuiltVectorTy, Val);
-  } else if (SDValue Widened =
-                 widenVectorToPartType(DAG, Val, DL, BuiltVectorTy)) {
+  } else {
+    if (BuiltVectorTy.getVectorElementType().bitsGT(
+            ValueVT.getVectorElementType())) {
+      // Integer promotion.
+      ValueVT = EVT::getVectorVT(*DAG.getContext(),
+                                 BuiltVectorTy.getVectorElementType(),
+                                 ValueVT.getVectorElementCount());
+      Val = DAG.getNode(ISD::ANY_EXTEND, DL, ValueVT, Val);
+    }
+
+    if (SDValue Widened = widenVectorToPartType(DAG, Val, DL, BuiltVectorTy)) {
       Val = Widened;
-  } else if (BuiltVectorTy.getVectorElementType().bitsGE(
-                 ValueVT.getVectorElementType()) &&
-             BuiltVectorTy.getVectorElementCount() ==
-                 ValueVT.getVectorElementCount()) {
-    // Promoted vector extract
-    Val = DAG.getAnyExtOrTrunc(Val, DL, BuiltVectorTy);
+    }
   }
 
   assert(Val.getValueType() == BuiltVectorTy && "Unexpected vector value type");
 
   // Split the vector into intermediate operands.
   SmallVector<SDValue, 8> Ops(NumIntermediates);
   for (unsigned i = 0; i != NumIntermediates; ++i) {
     if (IntermediateVT.isVector()) {

llvm/lib/CodeGen/TargetLoweringBase.cpp

unsigned TargetLoweringBase::getVectorTypeBreakdown(LLVMContext &Context,
 
   // Figure out the right, legal destination reg to copy into.
   EVT EltTy = VT.getVectorElementType();
 
   unsigned NumVectorRegs = 1;
 
   // Scalable vectors cannot be scalarized, so handle the legalisation of the
   // types like done elsewhere in SelectionDAG.
-  if (VT.isScalableVector() && !isPowerOf2_32(EltCnt.getKnownMinValue())) {
+  if (EltCnt.isScalable()) {
     LegalizeKind LK;
     EVT PartVT = VT;
     do {
       // Iterate until we've found a legal (part) type to hold VT.
       LK = getTypeConversion(Context, PartVT);
       PartVT = LK.second;
     } while (LK.first != TypeLegal);
 
-    NumIntermediates = VT.getVectorElementCount().getKnownMinValue() /
-                       PartVT.getVectorElementCount().getKnownMinValue();
+    if (!PartVT.isVector()) {
+      report_fatal_error(
+          "Don't know how to legalize this scalable vector type");
+    }
 
-    // FIXME: This code needs to be extended to handle more complex vector
-    // breakdowns, like nxv7i64 -> nxv8i64 -> 4 x nxv2i64. Currently the only
-    // supported cases are vectors that are broken down into equal parts
-    // such as nxv6i64 -> 3 x nxv2i64.
-    assert((PartVT.getVectorElementCount() * NumIntermediates) ==
-               VT.getVectorElementCount() &&
-           "Expected an integer multiple of PartVT");
+    NumIntermediates =
+        divideCeil(VT.getVectorElementCount().getKnownMinValue(),
+                   PartVT.getVectorElementCount().getKnownMinValue());
     IntermediateVT = PartVT;
     RegisterVT = getRegisterType(Context, IntermediateVT);
     return NumIntermediates;
   }
 
   // FIXME: We don't support non-power-of-2-sized vectors for now.  Ideally
   // we could break down into LHS/RHS like LegalizeDAG does.
   if (!isPowerOf2_32(EltCnt.getKnownMinValue())) {

llvm/test/CodeGen/AArch64/sve-breakdown-scalable-vectortype.ll

 ; CHECK-NEXT:    ret
   br i1 %b, label %L1, label %L2
 L1:
   call aarch64_sve_vector_pcs void @bar()
   ret <vscale x 8 x double> undef
 L2:
   ret <vscale x 8 x double> %illegal
 }
+
+define <vscale x 8 x i63> @wide_8i63(i1 %b, <vscale x 16 x i8> %legal, <vscale x 8 x i63> %illegal) nounwind {
+; CHECK-LABEL: wide_8i63:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    tbnz w0, #0, .LBB21_2
+; CHECK-NEXT:  // %bb.1: // %L2
+; CHECK-NEXT:    mov z0.d, z1.d
+; CHECK-NEXT:    mov z1.d, z2.d
+; CHECK-NEXT:    mov z2.d, z3.d
+; CHECK-NEXT:    mov z3.d, z4.d
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB21_2: // %L1
+; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    bl bar
+  br i1 %b, label %L1, label %L2
+L1:
+  call aarch64_sve_vector_pcs void @bar()
+  unreachable
+L2:
+  ret <vscale x 8 x i63> %illegal
+}
+
+define <vscale x 7 x i63> @wide_7i63(i1 %b, <vscale x 16 x i8> %legal, <vscale x 7 x i63> %illegal) nounwind {
+; CHECK-LABEL: wide_7i63:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    tbnz w0, #0, .LBB22_2
+; CHECK-NEXT:  // %bb.1: // %L2
+; CHECK-NEXT:    mov z0.d, z1.d
+; CHECK-NEXT:    mov z1.d, z2.d
+; CHECK-NEXT:    mov z2.d, z3.d
+; CHECK-NEXT:    mov z3.d, z4.d
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB22_2: // %L1
+; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    bl bar
+  br i1 %b, label %L1, label %L2
+L1:
+  call aarch64_sve_vector_pcs void @bar()
+  unreachable
+L2:
+  ret <vscale x 7 x i63> %illegal
+}
+
+define <vscale x 7 x i31> @wide_7i31(i1 %b, <vscale x 16 x i8> %legal, <vscale x 7 x i31> %illegal) nounwind {
+; CHECK-LABEL: wide_7i31:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    tbnz w0, #0, .LBB23_2
+; CHECK-NEXT:  // %bb.1: // %L2
+; CHECK-NEXT:    mov z0.d, z1.d
+; CHECK-NEXT:    mov z1.d, z2.d
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB23_2: // %L1
+; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    bl bar
+  br i1 %b, label %L1, label %L2
+L1:
+  call aarch64_sve_vector_pcs void @bar()
+  unreachable
+L2:
+  ret <vscale x 7 x i31> %illegal
+}