[IR] Allow scalable vectors in structs to support intrinsics returning multiple values.

RISC-V would like to use a struct of scalable vectors to return multiple
values from intrinsics. This woud also be needed for target independent
intrinsics like llvm.sadd.overflow.

This patch removes the existing restriction for this. I've modified
StructType::isSized to consider a struct containing scalable vectors
as unsized so the verifier won't allow loads/stores/allocas of these
structs.

Reviewed By: sdesmalen

Differential Revision: https://reviews.llvm.org/D94142

llvm/docs/LangRef.rst

@@ -704,7 +704,9 @@ Variables and aliases can have a
 :ref:`Thread Local Storage Model <tls_model>`.
 
 :ref:`Scalable vectors <t_vector>` cannot be global variables or members of
-structs or arrays because their size is unknown at compile time.
+arrays because their size is unknown at compile time. They are allowed in
+structs to facilitate intrinsics returning multiple values. Structs containing
+scalable vectors cannot be used in loads, stores, allocas, or GEPs.
 
 Syntax::
 

llvm/include/llvm/IR/DerivedTypes.h

@@ -284,6 +284,9 @@ public:
   /// isSized - Return true if this is a sized type.
   bool isSized(SmallPtrSetImpl<Type *> *Visited = nullptr) const;
 
+  /// Returns true if this struct contains a scalable vector.
+  bool containsScalableVectorType() const;
+
   /// Return true if this is a named struct that has a non-empty name.
   bool hasName() const { return SymbolTableEntry != nullptr; }
 

llvm/lib/CodeGen/Analysis.cpp

@@ -88,19 +88,25 @@ void llvm::ComputeValueVTs(const TargetLowering &TLI, const DataLayout &DL,
                            uint64_t StartingOffset) {
   // Given a struct type, recursively traverse the elements.
   if (StructType *STy = dyn_cast<StructType>(Ty)) {
-    const StructLayout *SL = DL.getStructLayout(STy);
+    // If the Offsets aren't needed, don't query the struct layout. This allows
+    // us to support structs with scalable vectors for operations that don't
+    // need offsets.
+    const StructLayout *SL = Offsets ? DL.getStructLayout(STy) : nullptr;
     for (StructType::element_iterator EB = STy->element_begin(),
                                       EI = EB,
                                       EE = STy->element_end();
-         EI != EE; ++EI)
+         EI != EE; ++EI) {
+      // Don't compute the element offset if we didn't get a StructLayout above.
+      uint64_t EltOffset = SL ? SL->getElementOffset(EI - EB) : 0;
       ComputeValueVTs(TLI, DL, *EI, ValueVTs, MemVTs, Offsets,
-                      StartingOffset + SL->getElementOffset(EI - EB));
+                      StartingOffset + EltOffset);
+    }
     return;
   }
   // Given an array type, recursively traverse the elements.
   if (ArrayType *ATy = dyn_cast<ArrayType>(Ty)) {
     Type *EltTy = ATy->getElementType();
-    uint64_t EltSize = DL.getTypeAllocSize(EltTy);
+    uint64_t EltSize = DL.getTypeAllocSize(EltTy).getFixedValue();
     for (unsigned i = 0, e = ATy->getNumElements(); i != e; ++i)
       ComputeValueVTs(TLI, DL, EltTy, ValueVTs, MemVTs, Offsets,
                       StartingOffset + i * EltSize);
@@ -131,16 +137,21 @@ void llvm::computeValueLLTs(const DataLayout &DL, Type &Ty,
                             uint64_t StartingOffset) {
   // Given a struct type, recursively traverse the elements.
   if (StructType *STy = dyn_cast<StructType>(&Ty)) {
-    const StructLayout *SL = DL.getStructLayout(STy);
-    for (unsigned I = 0, E = STy->getNumElements(); I != E; ++I)
+    // If the Offsets aren't needed, don't query the struct layout. This allows
+    // us to support structs with scalable vectors for operations that don't
+    // need offsets.
+    const StructLayout *SL = Offsets ? DL.getStructLayout(STy) : nullptr;
+    for (unsigned I = 0, E = STy->getNumElements(); I != E; ++I) {
+      uint64_t EltOffset = SL ? SL->getElementOffset(I) : 0;
       computeValueLLTs(DL, *STy->getElementType(I), ValueTys, Offsets,
-                       StartingOffset + SL->getElementOffset(I));
+                       StartingOffset + EltOffset);
+    }
     return;
   }
   // Given an array type, recursively traverse the elements.
   if (ArrayType *ATy = dyn_cast<ArrayType>(&Ty)) {
     Type *EltTy = ATy->getElementType();
-    uint64_t EltSize = DL.getTypeAllocSize(EltTy);
+    uint64_t EltSize = DL.getTypeAllocSize(EltTy).getFixedValue();
     for (unsigned i = 0, e = ATy->getNumElements(); i != e; ++i)
       computeValueLLTs(DL, *EltTy, ValueTys, Offsets,
                        StartingOffset + i * EltSize);

llvm/lib/IR/DataLayout.cpp

@@ -65,7 +65,8 @@ StructLayout::StructLayout(StructType *ST, const DataLayout &DL) {
     StructAlignment = std::max(TyAlign, StructAlignment);
 
     MemberOffsets[i] = StructSize;
-    StructSize += DL.getTypeAllocSize(Ty); // Consume space for this data item
+    // Consume space for this data item
+    StructSize += DL.getTypeAllocSize(Ty).getFixedValue();
   }
 
   // Add padding to the end of the struct so that it could be put in an array

llvm/lib/IR/Type.cpp

@@ -390,6 +390,18 @@ StructType *StructType::get(LLVMContext &Context, ArrayRef<Type*> ETypes,
   return ST;
 }
 
+bool StructType::containsScalableVectorType() const {
+  for (Type *Ty : elements()) {
+    if (isa<ScalableVectorType>(Ty))
+      return true;
+    if (auto *STy = dyn_cast<StructType>(Ty))
+      if (STy->containsScalableVectorType())
+        return true;
+  }
+
+  return false;
+}
+
 void StructType::setBody(ArrayRef<Type*> Elements, bool isPacked) {
   assert(isOpaque() && "Struct body already set!");
 
@@ -509,9 +521,14 @@ bool StructType::isSized(SmallPtrSetImpl<Type*> *Visited) const {
   // Okay, our struct is sized if all of the elements are, but if one of the
   // elements is opaque, the struct isn't sized *yet*, but may become sized in
   // the future, so just bail out without caching.
-  for (element_iterator I = element_begin(), E = element_end(); I != E; ++I)
-    if (!(*I)->isSized(Visited))
+  for (Type *Ty : elements()) {
+    // If the struct contains a scalable vector type, don't consider it sized.
+    // This prevents it from being used in loads/stores/allocas/GEPs.
+    if (isa<ScalableVectorType>(Ty))
       return false;
+    if (!Ty->isSized(Visited))
+      return false;
+  }
 
   // Here we cheat a bit and cast away const-ness. The goal is to memoize when
   // we find a sized type, as types can only move from opaque to sized, not the
@@ -531,7 +548,7 @@ StringRef StructType::getName() const {
 bool StructType::isValidElementType(Type *ElemTy) {
   return !ElemTy->isVoidTy() && !ElemTy->isLabelTy() &&
          !ElemTy->isMetadataTy() && !ElemTy->isFunctionTy() &&
-         !ElemTy->isTokenTy() && !isa<ScalableVectorType>(ElemTy);
+         !ElemTy->isTokenTy();
 }
 
 bool StructType::isLayoutIdentical(StructType *Other) const {

llvm/lib/IR/Verifier.cpp

@@ -714,12 +714,16 @@ void Verifier::visitGlobalVariable(const GlobalVariable &GV) {
   }
 
   // Scalable vectors cannot be global variables, since we don't know
-  // the runtime size. If the global is a struct or an array containing
-  // scalable vectors, that will be caught by the isValidElementType methods
-  // in StructType or ArrayType instead.
+  // the runtime size. If the global is an array containing scalable vectors,
+  // that will be caught by the isValidElementType methods in StructType or
+  // ArrayType instead.
   Assert(!isa<ScalableVectorType>(GV.getValueType()),
          "Globals cannot contain scalable vectors", &GV);
 
+  if (auto *STy = dyn_cast<StructType>(GV.getValueType()))
+    Assert(!STy->containsScalableVectorType(),
+           "Globals cannot contain scalable vectors", &GV);
+
   if (!GV.hasInitializer()) {
     visitGlobalValue(GV);
     return;

llvm/test/CodeGen/RISCV/scalable-vector-struct.ll

@@ -0,0 +1,25 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=riscv32 -mattr=+experimental-v -verify-machineinstrs < %s | FileCheck %s
+
+; This demonstrates that we can pass a struct containing scalable vectors across
+; a basic block.
+
+define i32 @foo({ {<vscale x 2 x i32>, <vscale x 2 x i32>}, i32 } %x, <vscale x 2 x i32>* %y, <vscale x 2 x i32>* %z) {
+; CHECK-LABEL: foo:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a3, zero, e32,m1,ta,mu
+; CHECK-NEXT:    vse32.v v16, (a1)
+; CHECK-NEXT:    vse32.v v17, (a2)
+; CHECK-NEXT:    ret
+entry:
+  br label %return
+
+return:
+  %a = extractvalue { {<vscale x 2 x i32>, <vscale x 2 x i32>}, i32 } %x, 1
+  %b = extractvalue { {<vscale x 2 x i32>, <vscale x 2 x i32>}, i32 } %x, 0, 0
+  %c = extractvalue { {<vscale x 2 x i32>, <vscale x 2 x i32>}, i32 } %x, 0, 1
+  store <vscale x 2 x i32> %b, <vscale x 2 x i32>* %y
+  store <vscale x 2 x i32> %c, <vscale x 2 x i32>* %z
+
+  ret i32 %a
+}

llvm/test/Other/scalable-vector-struct-intrinsic.ll

@@ -0,0 +1,18 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -verify < %s 2>&1 | FileCheck %s
+
+; Make sure we allow scalable vectors in structs for returning multiple
+; values from intrinsics.
+
+declare { <vscale x 2 x i32>, <vscale x 2 x i1> } @llvm.sadd.with.overflow.nxv2i32(<vscale x 2 x i32>, <vscale x 2 x i32>)
+
+define <vscale x 2 x i32> @foo(<vscale x 2 x i32> %x, <vscale x 2 x i32> %y) {
+; CHECK-LABEL: @foo(
+; CHECK-NEXT:    [[A:%.*]] = call { <vscale x 2 x i32>, <vscale x 2 x i1> } @llvm.sadd.with.overflow.nxv2i32(<vscale x 2 x i32> [[X:%.*]], <vscale x 2 x i32> [[Y:%.*]])
+; CHECK-NEXT:    [[B:%.*]] = extractvalue { <vscale x 2 x i32>, <vscale x 2 x i1> } [[A]], 0
+; CHECK-NEXT:    ret <vscale x 2 x i32> [[B]]
+;
+  %a = call { <vscale x 2 x i32>, <vscale x 2 x i1> } @llvm.sadd.with.overflow.nxv2i32(<vscale x 2 x i32> %x, <vscale x 2 x i32> %y)
+  %b = extractvalue { <vscale x 2 x i32>, <vscale x 2 x i1> } %a, 0
+  ret <vscale x 2 x i32> %b
+}

llvm/test/Other/scalable-vector-struct.ll

@@ -1,8 +0,0 @@
-; RUN: not opt -S -verify < %s 2>&1 | FileCheck %s
-
-;; Structs cannot contain scalable vectors; make sure we detect them even
-;; when nested inside other aggregates.
-
-%ty = type [2 x { i32, <vscale x 1 x i32> }]
-; CHECK: error: invalid element type for struct
-; CHECK: %ty = type [2 x { i32, <vscale x 1 x i32> }]

llvm/test/Verifier/scalable-global-vars.ll

@@ -7,6 +7,10 @@
 ; CHECK-NEXT: <vscale x 4 x i32>* @ScalableVecGlobal
 @ScalableVecGlobal = global <vscale x 4 x i32> zeroinitializer
 
+; CHECK-NEXT: Globals cannot contain scalable vectors
+; CHECK-NEXT: { i32, <vscale x 4 x i32> }* @ScalableVecStructGlobal
+@ScalableVecStructGlobal = global { i32,  <vscale x 4 x i32> } zeroinitializer
+
 ;; Global _pointers_ to scalable vectors are fine
 ; CHECK-NOT: Globals cannot contain scalable vectors
-@ScalableVecPtr = global <vscale x 8 x i16>* zeroinitializer
+@ScalableVecPtr = global <vscale x 8 x i16>* zeroinitializer

llvm/test/Verifier/scalable-vector-struct-alloca.ll

@@ -0,0 +1,7 @@
+; RUN: not opt -S -verify < %s 2>&1 | FileCheck %s
+
+define void @alloca() {
+; CHECK: error: Cannot allocate unsized type
+  %a = alloca { i32, <vscale x 1 x i32> }
+  ret void
+}

llvm/test/Verifier/scalable-vector-struct-load.ll

@@ -0,0 +1,8 @@
+; RUN: not opt -S -verify < %s 2>&1 | FileCheck %s
+
+define <vscale x 1 x i32> @load({ i32, <vscale x 1 x i32> }* %x) {
+; CHECK: error: loading unsized types is not allowed
+  %a = load { i32, <vscale x 1 x i32> }, { i32, <vscale x 1 x i32> }* %x
+  %b = extractvalue { i32, <vscale x 1 x i32> } %a, 1
+  ret <vscale x 1 x i32> %b
+}

llvm/test/Verifier/scalable-vector-struct-store.ll

@@ -0,0 +1,9 @@
+; RUN: not opt -S -verify < %s 2>&1 | FileCheck %s
+
+define void @store({ i32, <vscale x 1 x i32> }* %x, i32 %y, <vscale x 1 x i32> %z) {
+; CHECK: error: storing unsized types is not allowed
+  %a = insertvalue { i32, <vscale x 1 x i32> } undef, i32 %y, 0
+  %b = insertvalue { i32, <vscale x 1 x i32> } %a,  <vscale x 1 x i32> %z, 1
+  store { i32, <vscale x 1 x i32> } %b, { i32, <vscale x 1 x i32> }* %x
+  ret void
+}