[VFABI] Add support for vector functions that return struct types

llvm/include/llvm/Analysis/VectorUtils.h

@@ -18,6 +18,7 @@
 #include "llvm/Analysis/LoopAccessAnalysis.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/VFABIDemangler.h"
+#include "llvm/IR/VectorUtils.h"
 #include "llvm/Support/CheckedArithmetic.h"
 
 namespace llvm {
@@ -127,19 +128,6 @@ namespace Intrinsic {
 typedef unsigned ID;
 }
 
-/// A helper function for converting Scalar types to vector types. If
-/// the incoming type is void, we return void. If the EC represents a
-/// scalar, we return the scalar type.
-inline Type *ToVectorTy(Type *Scalar, ElementCount EC) {
-  if (Scalar->isVoidTy() || Scalar->isMetadataTy() || EC.isScalar())
-    return Scalar;
-  return VectorType::get(Scalar, EC);
-}
-
-inline Type *ToVectorTy(Type *Scalar, unsigned VF) {
-  return ToVectorTy(Scalar, ElementCount::getFixed(VF));
-}
-
 /// Identify if the intrinsic is trivially vectorizable.
 /// This method returns true if the intrinsic's argument types are all scalars
 /// for the scalar form of the intrinsic and all vectors (or scalars handled by

llvm/include/llvm/IR/StructWideningUtils.h

@@ -0,0 +1,30 @@
+//===--- StructWideningUtils.h - Utils for widening/narrowing struct types ===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_IR_STRUCTWIDENINGUTILS_H
+#define LLVM_IR_STRUCTWIDENINGUTILS_H
+
+#include "llvm/IR/DerivedTypes.h"
+
+namespace llvm {
+
+/// A helper for converting structs of scalar types to structs of vector types.
+/// Note: Only unpacked literal struct types are supported.
+Type *ToWideStructTy(StructType *StructTy, ElementCount EC);
+
+/// A helper for converting structs of vector types to structs of scalar types.
+/// Note: Only unpacked literal struct types are supported.
+Type *ToNarrowStructTy(StructType *StructTy);
+
+/// Returns true if `StructTy` is an unpacked literal struct where all elements
+/// are vectors of matching element count. This does not include empty structs.
+bool isWideStructTy(StructType *StructTy);
+
+} // namespace llvm
+
+#endif

llvm/include/llvm/IR/VectorUtils.h

@@ -0,0 +1,74 @@
+//===----------- VectorUtils.h - Vector type utility functions -*- C++ -*-====//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_IR_VECTORUTILS_H
+#define LLVM_IR_VECTORUTILS_H
+
+#include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/StructWideningUtils.h"
+
+namespace llvm {
+
+/// A helper function for converting Scalar types to vector types. If
+/// the incoming type is void, we return void. If the EC represents a
+/// scalar, we return the scalar type.
+inline Type *ToVectorTy(Type *Scalar, ElementCount EC) {
+  if (Scalar->isVoidTy() || Scalar->isMetadataTy() || EC.isScalar())
+    return Scalar;
+  return VectorType::get(Scalar, EC);
+}
+
+inline Type *ToVectorTy(Type *Scalar, unsigned VF) {
+  return ToVectorTy(Scalar, ElementCount::getFixed(VF));
+}
+
+/// A helper for converting to wider (vector) types. For scalar types, this is
+/// equivalent to calling `ToVectorTy`. For struct types, this returns a new
+/// struct where each element type has been widened to a vector type. Note: Only
+/// unpacked literal struct types are supported.
+inline Type *ToWideTy(Type *Ty, ElementCount EC) {
+  if (StructType *StructTy = dyn_cast<StructType>(Ty))
+    return ToWideStructTy(StructTy, EC);
+  return ToVectorTy(Ty, EC);
+}
+
+/// A helper for converting wide types to narrow (non-vector) types. For vector
+/// types, this is equivalent to calling .getScalarType(). For struct types,
+/// this returns a new struct where each element type has been converted to a
+/// scalar type. Note: Only unpacked literal struct types are supported.
+inline Type *ToNarrowTy(Type *Ty) {
+  if (StructType *StructTy = dyn_cast<StructType>(Ty))
+    return ToNarrowStructTy(StructTy);
+  return Ty->getScalarType();
+}
+
+/// Returns true if `Ty` is a vector type or a struct of vector types where all
+/// vector types share the same VF.
+inline bool isWideTy(Type *Ty) {
+  if (StructType *StructTy = dyn_cast<StructType>(Ty))
+    return isWideStructTy(StructTy);
+  return Ty->isVectorTy();
+}
+
+/// Returns the types contained in `Ty`. For struct types, it returns the
+/// elements, all other types are returned directly.
+inline ArrayRef<Type *> getContainedTypes(Type *const &Ty) {
+  if (auto *StructTy = dyn_cast<StructType>(Ty))
+    return StructTy->elements();
+  return ArrayRef<Type *>{&Ty, 1};
+}
+
+/// Returns the vectorization factor for a widened type.
+inline ElementCount getWideTypeVF(Type *Ty) {
+  assert(isWideTy(Ty) && "expected widened type");
+  return cast<VectorType>(getContainedTypes(Ty).front())->getElementCount();
+}
+
+} // namespace llvm
+
+#endif

llvm/lib/IR/CMakeLists.txt

@@ -64,6 +64,7 @@ add_llvm_component_library(LLVMCore
   PseudoProbe.cpp
   ReplaceConstant.cpp
   Statepoint.cpp
+  StructWideningUtils.cpp
   StructuralHash.cpp
   Type.cpp
   TypedPointerType.cpp

llvm/lib/IR/StructWideningUtils.cpp

@@ -0,0 +1,57 @@
+//===- StructWideningUtils.cpp - Utils for widening/narrowing struct types ===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/IR/StructWideningUtils.h"
+#include "llvm/ADT/SmallVectorExtras.h"
+#include "llvm/IR/VectorUtils.h"
+
+using namespace llvm;
+
+static bool isUnpackedStructLiteral(StructType *StructTy) {
+  return StructTy->isLiteral() && !StructTy->isPacked();
+}
+
+/// A helper for converting structs of scalar types to structs of vector types.
+/// Note: Only unpacked literal struct types are supported.
+Type *llvm::ToWideStructTy(StructType *StructTy, ElementCount EC) {
+  if (EC.isScalar())
+    return StructTy;
+  assert(isUnpackedStructLiteral(StructTy) &&
+         "expected unpacked struct literal");
+  return StructType::get(
+      StructTy->getContext(),
+      map_to_vector(StructTy->elements(), [&](Type *ElTy) -> Type * {
+        return VectorType::get(ElTy, EC);
+      }));
+}
+
+/// A helper for converting structs of vector types to structs of scalar types.
+/// Note: Only unpacked literal struct types are supported.
+Type *llvm::ToNarrowStructTy(StructType *StructTy) {
+  assert(isUnpackedStructLiteral(StructTy) &&
+         "expected unpacked struct literal");
+  return StructType::get(
+      StructTy->getContext(),
+      map_to_vector(StructTy->elements(), [](Type *ElTy) -> Type * {
+        return ElTy->getScalarType();
+      }));
+}
+
+/// Returns true if `StructTy` is an unpacked literal struct where all elements
+/// are vectors of matching element count. This does not include empty structs.
+bool llvm::isWideStructTy(StructType *StructTy) {
+  if (!isUnpackedStructLiteral(StructTy))
+    return false;
+  auto ElemTys = StructTy->elements();
+  if (ElemTys.empty() || !ElemTys.front()->isVectorTy())
+    return false;
+  ElementCount VF = cast<VectorType>(ElemTys.front())->getElementCount();
+  return all_of(ElemTys, [&](Type *Ty) {
+    return Ty->isVectorTy() && cast<VectorType>(Ty)->getElementCount() == VF;
+  });
+}

llvm/lib/IR/VFABIDemangler.cpp

@@ -11,6 +11,7 @@
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/VectorUtils.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include <limits>
@@ -346,12 +347,15 @@ getScalableECFromSignature(const FunctionType *Signature, const VFISAKind ISA,
   // Also check the return type if not void.
   Type *RetTy = Signature->getReturnType();
   if (!RetTy->isVoidTy()) {
-    std::optional<ElementCount> ReturnEC = getElementCountForTy(ISA, RetTy);
-    // If we have an unknown scalar element type we can't find a reasonable VF.
-    if (!ReturnEC)
-      return std::nullopt;
-    if (ElementCount::isKnownLT(*ReturnEC, MinEC))
-      MinEC = *ReturnEC;
+    for (Type *RetTy : getContainedTypes(RetTy)) {
+      std::optional<ElementCount> ReturnEC = getElementCountForTy(ISA, RetTy);
+      // If we have an unknown scalar element type we can't find a reasonable
+      // VF.
+      if (!ReturnEC)
+        return std::nullopt;
+      if (ElementCount::isKnownLT(*ReturnEC, MinEC))
+        MinEC = *ReturnEC;
+    }
   }
 
   // The SVE Vector function call ABI bases the VF on the widest element types
@@ -566,7 +570,7 @@ FunctionType *VFABI::createFunctionType(const VFInfo &Info,
 
   auto *RetTy = ScalarFTy->getReturnType();
   if (!RetTy->isVoidTy())
-    RetTy = VectorType::get(RetTy, VF);
+    RetTy = ToWideTy(RetTy, VF);
   return FunctionType::get(RetTy, VecTypes, false);
 }
 

llvm/unittests/IR/CMakeLists.txt

@@ -52,6 +52,7 @@ add_llvm_unittest(IRTests
   ValueTest.cpp
   VectorBuilderTest.cpp
   VectorTypesTest.cpp
+  VectorUtilsTest.cpp
   VerifierTest.cpp
   VFABIDemanglerTest.cpp
   VPIntrinsicTest.cpp

llvm/unittests/IR/VFABIDemanglerTest.cpp

@@ -753,6 +753,73 @@ TEST_F(VFABIParserTest, ParseVoidReturnTypeSVE) {
   EXPECT_EQ(VectorName, "vector_foo");
 }
 
+TEST_F(VFABIParserTest, ParseWideStructReturnTypeSVE) {
+  EXPECT_TRUE(
+      invokeParser("_ZGVsMxv_foo(vector_foo)", "{double, double}(float)"));
+  EXPECT_EQ(ISA, VFISAKind::SVE);
+  EXPECT_TRUE(isMasked());
+  FunctionType *FTy = FunctionType::get(
+      StructType::get(
+          VectorType::get(Type::getDoubleTy(Ctx), ElementCount::getScalable(2)),
+          VectorType::get(Type::getDoubleTy(Ctx),
+                          ElementCount::getScalable(2))),
+      {
+          VectorType::get(Type::getFloatTy(Ctx), ElementCount::getScalable(2)),
+          VectorType::get(Type::getInt1Ty(Ctx), ElementCount::getScalable(2)),
+      },
+      false);
+  EXPECT_EQ(getFunctionType(), FTy);
+  EXPECT_EQ(Parameters.size(), 2U);
+  EXPECT_EQ(Parameters[0], VFParameter({0, VFParamKind::Vector}));
+  EXPECT_EQ(Parameters[1], VFParameter({1, VFParamKind::GlobalPredicate}));
+  EXPECT_EQ(VF, ElementCount::getScalable(2));
+  EXPECT_EQ(ScalarName, "foo");
+  EXPECT_EQ(VectorName, "vector_foo");
+}
+
+TEST_F(VFABIParserTest, ParseWideStructMixedReturnTypeSVE) {
+  EXPECT_TRUE(invokeParser("_ZGVsMxv_foo(vector_foo)", "{float, i64}(float)"));
+  EXPECT_EQ(ISA, VFISAKind::SVE);
+  EXPECT_TRUE(isMasked());
+  FunctionType *FTy = FunctionType::get(
+      StructType::get(
+          VectorType::get(Type::getFloatTy(Ctx), ElementCount::getScalable(2)),
+          VectorType::get(Type::getInt64Ty(Ctx), ElementCount::getScalable(2))),
+      {
+          VectorType::get(Type::getFloatTy(Ctx), ElementCount::getScalable(2)),
+          VectorType::get(Type::getInt1Ty(Ctx), ElementCount::getScalable(2)),
+      },
+      false);
+  EXPECT_EQ(getFunctionType(), FTy);
+  EXPECT_EQ(Parameters.size(), 2U);
+  EXPECT_EQ(Parameters[0], VFParameter({0, VFParamKind::Vector}));
+  EXPECT_EQ(Parameters[1], VFParameter({1, VFParamKind::GlobalPredicate}));
+  EXPECT_EQ(VF, ElementCount::getScalable(2));
+  EXPECT_EQ(ScalarName, "foo");
+  EXPECT_EQ(VectorName, "vector_foo");
+}
+
+TEST_F(VFABIParserTest, ParseWideStructReturnTypeNEON) {
+  EXPECT_TRUE(
+      invokeParser("_ZGVnN2v_foo(vector_foo)", "{float, float}(float)"));
+  EXPECT_EQ(ISA, VFISAKind::AdvancedSIMD);
+  EXPECT_FALSE(isMasked());
+  FunctionType *FTy = FunctionType::get(
+      StructType::get(
+          VectorType::get(Type::getFloatTy(Ctx), ElementCount::getFixed(2)),
+          VectorType::get(Type::getFloatTy(Ctx), ElementCount::getFixed(2))),
+      {
+          VectorType::get(Type::getFloatTy(Ctx), ElementCount::getFixed(2)),
+      },
+      false);
+  EXPECT_EQ(getFunctionType(), FTy);
+  EXPECT_EQ(Parameters.size(), 1U);
+  EXPECT_EQ(Parameters[0], VFParameter({0, VFParamKind::Vector}));
+  EXPECT_EQ(VF, ElementCount::getFixed(2));
+  EXPECT_EQ(ScalarName, "foo");
+  EXPECT_EQ(VectorName, "vector_foo");
+}
+
 // Make sure we reject unsupported parameter types.
 TEST_F(VFABIParserTest, ParseUnsupportedElementTypeSVE) {
   EXPECT_FALSE(invokeParser("_ZGVsMxv_foo(vector_foo)", "void(i128)"));