[clang] [llvm] [LLVMABI][AARCH64] Implement Pure Scalable Type handling (PR #227504)

[Madhur Amilkanthwar via llvm-commits]

================
@@ -339,6 +375,175 @@ bool AArch64TargetInfo::isIllegalVectorType(const Type *Ty) const {
   return false;
 }
 
+// Expand a memory type into a sequence with an element for each non-record,
+// non-array member of the type, with the exception of the padding types, which
+// are retained.
+void AArch64TargetInfo::flattenType(
+    const Type *Ty, SmallVectorImpl<const Type *> &Flattened) const {
+  if (ArgInfo::isPaddingForCoerceAndExpand(Ty)) {
+    Flattened.push_back(Ty);
+    return;
+  }
+
+  if (const auto *AT = dyn_cast<ArrayType>(Ty)) {
+    uint64_t NElt = AT->getNumElements();
+    if (NElt == 0)
+      return;
+
+    SmallVector<const Type *, 4> EltFlattened;
+    flattenType(AT->getElementType(), EltFlattened);
+
+    for (uint64_t I = 0; I < NElt; ++I)
+      llvm::append_range(Flattened, EltFlattened);
+    return;
+  }
+
+  if (const auto *RT = dyn_cast<RecordType>(Ty)) {
+    for (const FieldInfo &Field : RT->getFields())
+      flattenType(Field.FieldType, Flattened);
+    return;
+  }
+
+  Flattened.push_back(Ty);
+}
+
+ArgInfo AArch64TargetInfo::coerceAndExpandPureScalableAggregate(
+    const Type *Ty, bool IsNamedArg, unsigned NVec, unsigned NPred,
+    const SmallVectorImpl<const Type *> &UnpaddedCoerceToSeq, unsigned &NSRN,
+    unsigned &NPRN) const {
+  // An unnamed argument, or one that does not fit in the remaining Z or P
+  // registers, is passed indirectly and does not consume those registers.
+  if (!IsNamedArg || NSRN + NVec > 8 || NPRN + NPred > 4)
+    return getNaturalAlignIndirect(Ty, getAllocaAddrSpace(), /*ByVal=*/false);
+
+  NSRN += NVec;
+  NPRN += NPred;
+
+  // A sizeless SVE tuple is already one register per member.
+  if (Ty->isSVESizelessType())
+    return ArgInfo::getDirect();
+
+  assert(!UnpaddedCoerceToSeq.empty() && "pure scalable type has no members");
+  const Type *UnpaddedCoerceToType =
+      UnpaddedCoerceToSeq.size() == 1
+          ? UnpaddedCoerceToSeq[0]
+          : getStructOfTypes(UnpaddedCoerceToSeq, /*Packed=*/true);
+
+  SmallVector<const Type *, 8> CoerceToSeq;
+  flattenType(convertTypeForMem(Ty), CoerceToSeq);
+  return ArgInfo::getCoerceAndExpand(
+      getStructOfTypes(CoerceToSeq, /*Packed=*/false), UnpaddedCoerceToType);
+}
+
+// A Pure Scalable Type (AAPCS64) is passed in Z and P registers. On success
+// NVec and NPred are how many of each it needs, and CoerceToSeq has one
+// scalable vector per register. A sequence longer than 12 is rejected so the
+// caller treats the type as a large composite.
+bool AArch64TargetInfo::passAsPureScalableType(
+    const Type *Ty, unsigned &NVec, unsigned &NPred,
+    SmallVectorImpl<const Type *> &CoerceToSeq) const {
+  if (const auto *AT = dyn_cast<ArrayType>(Ty)) {
+    if (AT->isMatrixType())
+      return false;
+
+    uint64_t NElt = AT->getNumElements();
+    if (NElt == 0)
+      return false;
+
+    unsigned NV = 0, NP = 0;
+    SmallVector<const Type *, 4> EltCoerceToSeq;
+    if (!passAsPureScalableType(AT->getElementType(), NV, NP, EltCoerceToSeq))
+      return false;
+
+    if (CoerceToSeq.size() + NElt * EltCoerceToSeq.size() > 12)
+      return false;
+
+    for (uint64_t I = 0; I < NElt; ++I)
+      llvm::append_range(CoerceToSeq, EltCoerceToSeq);
+
+    NVec += NElt * NV;
+    NPred += NElt * NP;
+    return true;
+  }
+
+  if (const auto *RT = dyn_cast<RecordType>(Ty)) {
+    if (getRecordArgABI(RT) != RAA_Default)
+      return false;
+    // Pure scalable types are never unions and never contain unions.
+    if (RT->isUnion())
+      return false;
+
+    // Direct virtual bases are not in getBaseClasses(). A record that has
+    // one cannot be passed in registers, and getRecordArgABI rejected it
+    // above.
+    for (const FieldInfo &Base : RT->getBaseClasses()) {
+      if (Base.FieldType->isEmptyRecord())
+        continue;
+      if (!passAsPureScalableType(Base.FieldType, NVec, NPred, CoerceToSeq))
+        return false;
+    }
+    for (const FieldInfo &Field : RT->getFields()) {
+      if (Field.isEmpty())
+        continue;
----------------
madhur13490 wrote:

I think flexible array members slip through here. QualTypeMapper maps T y[] to a zero-element array (clang/lib/CodeGen/QualTypeMapper.cpp:349-350), and FieldInfo::isEmpty() treats zero-element arrays as empty (llvm/lib/ABI/Types.cpp:91). So the field is skipped and the struct is classified as a PST. Clang doesn't skip it, because an incomplete array isn't a constant array in isEmptyField, so Clang rejects the struct as a PST. struct { fixed_int32_t x; fixed_int32_t y[]; } gives Kind mismatch (expected: Direct, given: CoerceAndExpand) at 128 bits and expected: Indirect at 256 and 512. Maybe bail out early on RT->hasFlexibleArrayMember()?

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