[IR] Enable load/store/alloca for arrays of scalable vectors.

[paulwalker-arm]

I've raised https://discourse.llvm.org/t/rfc-enable-arrays-of-scalable-vector-types/72935 for discussion.

I'm still looking to tighten up the Verifier a little more because there's at least one case where loosing the "invalid array element type" error for ArrayType is allowing bogus IR to be accepted.

[nikic]

Unlike homogeneous scalable aggregates, this seems to allow arrays of scalable vectors in GEPs. I'd expect you need a lot more changes than this to correctly support that.

[paulwalker-arm]

In D158517#4609135, @nikic wrote:

Unlike homogeneous scalable aggregates, this seems to allow arrays of scalable vectors in GEPs. I'd expect you need a lot more changes than this to correctly support that.

I'll flesh out the testing to see what crops up.

[paulwalker-arm]

Extend fixes for various scalable vector code paths relating to getelementptr to also cover arrays of scalable vectors.

[paulwalker-arm]

review ping

[sdesmalen]

Does this mean we could also enable this for structs with scalable members?

llvm/lib/IR/Type.cpp

60–62

Could you update the description of isScalableTy in the header file. It already seems out of date because the description doesn't cover struct types with scalable vectors.

llvm/test/CodeGen/AArch64/alloca-load-store-scalable-array.ll

9

Can you also add a test for arrays with (nested) scalable arrays?

32

Is it worth having a test with insertvalue/extractvalue as well, just to make sure that it's handled correctly by ISel?

llvm/test/CodeGen/RISCV/rvv/alloca-load-store-scalable-array.ll

7

nit: please remove the space (also in the AArch64 test)

llvm/test/Transforms/GlobalOpt/2022-08-23-ScalableVectorArrayCrash.ll

10 ↗

(On Diff #554391)

nit: can you remove the attributes from the test that aren't strictly necessary?

12 ↗

(On Diff #554391)

No action expected here, but this looks like dodgy IR, which I think might be the purpose of the test? At compile-time we don't know whether sizeof([4 x <vscale x 2 x double>]) exceeds 768. I can see that from a store perspective, the pointer is just an opaque pointer that's being stored to. I would have expected 'proper' IR to insert explicit conversions using extractelement/llvm.vector.extract/insertelement/llvm.vector.insert to implement such a store, rather than using this.

[paulwalker-arm]

Rebased, updated function description of isScalableTy and added more tests.

[paulwalker-arm]

llvm/test/Transforms/GlobalOpt/2022-08-23-ScalableVectorArrayCrash.ll

10 ↗

(On Diff #554391)

Do you mind if I don't because it's a clone of an existing test (2022-08-23-ScalableVectorCrash.ll) and I'd rather they remain in sync. I can follow up and clean both together if that's ok?

12 ↗

(On Diff #554391)

I don't believe the IR is dodgy, it just assumes the IR writer knows what they're doing. This array example is not possible from C/C++ today but considering just scalable vectors you're free to access global data via the SVE builtins without any form of casting or address/range validation. In fact this is the only way to access global data given such types cannot be declared as scalable.

[awarzynski]

@benmxwl-arm and I tested this change in the context of MLIR and things just work ™, thanks Paul! In general, once this lands it will unblock a very important path for us for supporting scalable auto-vec in MLIR 🙏🏻

[paulwalker-arm]

review ping

[sdesmalen]

Thanks for adding the extra tests.

llvm/test/CodeGen/AArch64/alloca-load-store-scalable-array.ll

7

nit: feel free to ignore, but maybe change these names to something descriptive, e.g. my_subtype -> my_nxv2f64 and my_type -> 3x_my_nxv2f64?

llvm/test/Transforms/GlobalOpt/2022-08-23-ScalableVectorArrayCrash.ll

10 ↗

(On Diff #554391)

Sure, cleaning it as a follow-up patch is fine to me.

llvm/docs/LangRef.rst

 Globals can also have a :ref:`DLL storage class <dllstorageclass>`,
 an optional :ref:`runtime preemption specifier <runtime_preemption_model>`,
 an optional :ref:`global attributes <glattrs>` and
 an optional list of attached :ref:`metadata <metadata>`.
 
 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
-arrays because their size is unknown at compile time. They are allowed in
-structs to facilitate intrinsics returning multiple values. Generally, structs
-containing scalable vectors are not considered "sized" and cannot be used in
-loads, stores, allocas, or GEPs. The only exception to this rule is for structs
-that contain scalable vectors of the same type (e.g. ``{<vscale x 2 x i32>,
-<vscale x 2 x i32>}`` contains the same type while ``{<vscale x 2 x i32>,
-<vscale x 2 x i64>}`` doesn't). These kinds of structs (we may call them
-homogeneous scalable vector structs) are considered sized and can be used in
-loads, stores, allocas, but not GEPs.
+Globals cannot be or contain :ref:`Scalable vectors <t_vector>` because their
+size is unknown at compile time. They are allowed in structs to facilitate
+intrinsics returning multiple values. Generally, structs containing scalable
+vectors are not considered "sized" and cannot be used in loads, stores, allocas,
+or GEPs. The only exception to this rule is for structs that contain scalable
+vectors of the same type (e.g. ``{<vscale x 2 x i32>, <vscale x 2 x i32>}``
+contains the same type while ``{<vscale x 2 x i32>, <vscale x 2 x i64>}``
+doesn't). These kinds of structs (we may call them homogeneous scalable vector
+structs) are considered sized and can be used in loads, stores, allocas, but
+not GEPs.
 
 Syntax::
 
       @<GlobalVarName> = [Linkage] [PreemptionSpecifier] [Visibility]
                          [DLLStorageClass] [ThreadLocal]
                          [(unnamed_addr|local_unnamed_addr)] [AddrSpace]
                          [ExternallyInitialized]
                          <global | constant> <Type> [<InitializerConstant>]

llvm/lib/Analysis/InstructionSimplify.cpp

if (isa<PoisonValue>(Ptr) ||
       any_of(Indices, [](const auto *V) { return isa<PoisonValue>(V); }))
     return PoisonValue::get(GEPTy);
 
   // getelementptr undef, idx -> undef
   if (Q.isUndefValue(Ptr))
     return UndefValue::get(GEPTy);
 
   bool IsScalableVec =
-      isa<ScalableVectorType>(SrcTy) || any_of(Indices, [](const Value *V) {
+      SrcTy->isScalableTy() || any_of(Indices, [](const Value *V) {
         return isa<ScalableVectorType>(V->getType());
       });
 
   if (Indices.size() == 1) {
     // getelementptr P, 0 -> P.
     if (match(Indices[0], m_Zero()) && Ptr->getType() == GEPTy)
       return Ptr;
 

llvm/lib/IR/Type.cpp

default:
     return nullptr;
   }
 }
 
 bool Type::isIntegerTy(unsigned Bitwidth) const {
   return isIntegerTy() && cast<IntegerType>(this)->getBitWidth() == Bitwidth;
 }
 
 bool Type::isScalableTy() const {
+  if (const auto *ATy = dyn_cast<ArrayType>(this))
+    return ATy->getElementType()->isScalableTy();
   if (const auto *STy = dyn_cast<StructType>(this)) {
     SmallPtrSet<Type *, 4> Visited;
     return STy->containsScalableVectorType(&Visited);
   }
   return getTypeID() == ScalableVectorTyID || isScalableTargetExtTy();
 }
 
 const fltSemantics &Type::getFltSemantics() const {
ArrayType *ArrayType::get(Type *ElementType, uint64_t NumElements) {
   if (!Entry)
     Entry = new (pImpl->Alloc) ArrayType(ElementType, NumElements);
   return Entry;
 }
 
 bool ArrayType::isValidElementType(Type *ElemTy) {
   return !ElemTy->isVoidTy() && !ElemTy->isLabelTy() &&
          !ElemTy->isMetadataTy() && !ElemTy->isFunctionTy() &&
-         !ElemTy->isTokenTy() && !ElemTy->isX86_AMXTy() &&
-         !isa<ScalableVectorType>(ElemTy);
+         !ElemTy->isTokenTy() && !ElemTy->isX86_AMXTy();
 }
 
 //===----------------------------------------------------------------------===//
 //                          VectorType Implementation
 //===----------------------------------------------------------------------===//
 
 VectorType::VectorType(Type *ElType, unsigned EQ, Type::TypeID TID)
     : Type(ElType->getContext(), TID), ContainedType(ElType),

llvm/lib/IR/Verifier.cpp

for (auto *MD : MDs) {
     if (auto *GVE = dyn_cast<DIGlobalVariableExpression>(MD))
       visitDIGlobalVariableExpression(*GVE);
     else
       CheckDI(false, "!dbg attachment of global variable must be a "
                      "DIGlobalVariableExpression");
   }
 
   // Scalable vectors cannot be global variables, since we don't know
-  // 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.
-  Check(!isa<ScalableVectorType>(GV.getValueType()),
-        "Globals cannot contain scalable vectors", &GV);
-
-  if (auto *STy = dyn_cast<StructType>(GV.getValueType())) {
-    SmallPtrSet<Type *, 4> Visited;
-    Check(!STy->containsScalableVectorType(&Visited),
-          "Globals cannot contain scalable vectors", &GV);
-  }
+  // the runtime size.
+  Check(!GV.getValueType()->isScalableTy(),
+        "Globals cannot contain scalable types", &GV);
 
   // Check if it's a target extension type that disallows being used as a
   // global.
   if (auto *TTy = dyn_cast<TargetExtType>(GV.getValueType()))
     Check(TTy->hasProperty(TargetExtType::CanBeGlobal),
           "Global @" + GV.getName() + " has illegal target extension type",
           TTy);
 

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

if (auto *AT = dyn_cast<ArrayType>(T)) {
     // Bail out if the array is too large. Ideally we would like to optimize
     // arrays of arbitrary size but this has a terrible impact on compile time.
     // The threshold here is chosen arbitrarily, maybe needs a little bit of
     // tuning.
     if (NumElements > IC.MaxArraySizeForCombine)
       return nullptr;
 
     const DataLayout &DL = IC.getDataLayout();
-    auto EltSize = DL.getTypeAllocSize(ET);
+    TypeSize EltSize = DL.getTypeAllocSize(ET);
     const auto Align = LI.getAlign();
 
     auto *Addr = LI.getPointerOperand();
     auto *IdxType = Type::getInt64Ty(T->getContext());
     auto *Zero = ConstantInt::get(IdxType, 0);
 
     Value *V = PoisonValue::get(T);
-    uint64_t Offset = 0;
+    TypeSize Offset = TypeSize::get(0, ET->isScalableTy());
     for (uint64_t i = 0; i < NumElements; i++) {
       Value *Indices[2] = {
         Zero,
         ConstantInt::get(IdxType, i),
       };
       auto *Ptr = IC.Builder.CreateInBoundsGEP(AT, Addr, ArrayRef(Indices),
                                                Name + ".elt");
+      auto EltAlign = commonAlignment(Align, Offset.getKnownMinValue());
       auto *L = IC.Builder.CreateAlignedLoad(AT->getElementType(), Ptr,
-                                             commonAlignment(Align, Offset),
-                                             Name + ".unpack");
+                                             EltAlign, Name + ".unpack");
       L->setAAMetadata(LI.getAAMetadata());
       V = IC.Builder.CreateInsertValue(V, L, i);
       Offset += EltSize;
     }
 
     V->setName(Name);
     return IC.replaceInstUsesWith(LI, V);
   }
if (auto *AT = dyn_cast<ArrayType>(T)) {
     // Bail out if the array is too large. Ideally we would like to optimize
     // arrays of arbitrary size but this has a terrible impact on compile time.
     // The threshold here is chosen arbitrarily, maybe needs a little bit of
     // tuning.
     if (NumElements > IC.MaxArraySizeForCombine)
       return false;
 
     const DataLayout &DL = IC.getDataLayout();
-    auto EltSize = DL.getTypeAllocSize(AT->getElementType());
+    TypeSize EltSize = DL.getTypeAllocSize(AT->getElementType());
     const auto Align = SI.getAlign();
 
     SmallString<16> EltName = V->getName();
     EltName += ".elt";
     auto *Addr = SI.getPointerOperand();
     SmallString<16> AddrName = Addr->getName();
     AddrName += ".repack";
 
     auto *IdxType = Type::getInt64Ty(T->getContext());
     auto *Zero = ConstantInt::get(IdxType, 0);
 
-    uint64_t Offset = 0;
+    TypeSize Offset = TypeSize::get(0, AT->getElementType()->isScalableTy());
     for (uint64_t i = 0; i < NumElements; i++) {
       Value *Indices[2] = {
         Zero,
         ConstantInt::get(IdxType, i),
       };
       auto *Ptr =
           IC.Builder.CreateInBoundsGEP(AT, Addr, ArrayRef(Indices), AddrName);
       auto *Val = IC.Builder.CreateExtractValue(V, i, EltName);
-      auto EltAlign = commonAlignment(Align, Offset);
+      auto EltAlign = commonAlignment(Align, Offset.getKnownMinValue());
       Instruction *NS = IC.Builder.CreateAlignedStore(Val, Ptr, EltAlign);
       NS->setAAMetadata(SI.getAAMetadata());
       Offset += EltSize;
     }
 
     return true;
   }
 

llvm/test/CodeGen/AArch64/alloca-load-store-scalable-array.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; RUN: llc < %s | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+%my_type = type [ 3 x <vscale x 2 x double>]
+
+define void @test(ptr %addr) #0 {
+; CHECK-LABEL: test:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    addvl sp, sp, #-3
+; CHECK-NEXT:    .cfi_escape 0x0f, 0x0c, 0x8f, 0x00, 0x11, 0x10, 0x22, 0x11, 0x18, 0x92, 0x2e, 0x00, 0x1e, 0x22 // sp + 16 + 24 * VG
+; CHECK-NEXT:    .cfi_offset w29, -16
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    ld1d { z0.d }, p0/z, [x0, #1, mul vl]
+; CHECK-NEXT:    ld1d { z1.d }, p0/z, [x0, #2, mul vl]
+; CHECK-NEXT:    ld1d { z2.d }, p0/z, [x0]
+; CHECK-NEXT:    st1d { z2.d }, p0, [sp]
+; CHECK-NEXT:    st1d { z1.d }, p0, [sp, #2, mul vl]
+; CHECK-NEXT:    st1d { z0.d }, p0, [sp, #1, mul vl]
+; CHECK-NEXT:    addvl sp, sp, #3
+; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+entry:
+  %ret = alloca %my_type, align 8
+  %val = load %my_type, ptr %addr
+  store %my_type %val, ptr %ret, align 8
+  ret void
+}
+
+attributes #0 = { "target-features"="+sve" }

llvm/test/CodeGen/RISCV/rvv/alloca-load-store-scalable-array.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; RUN: llc -mtriple=riscv64 -mattr=+v -verify-machineinstrs --riscv-no-aliases < %s | FileCheck %s
+
+target triple = "riscv64-unknown-unknown-elf"
+
+%my_type = type [ 3 x <vscale x 1 x double>]
+
+define void @test(ptr %addr) {
+; CHECK-LABEL: test:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    addi sp, sp, -16
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    csrrs a1, vlenb, zero
+; CHECK-NEXT:    slli a1, a1, 2
+; CHECK-NEXT:    sub sp, sp, a1
+; CHECK-NEXT:    .cfi_escape 0x0f, 0x0d, 0x72, 0x00, 0x11, 0x10, 0x22, 0x11, 0x04, 0x92, 0xa2, 0x38, 0x00, 0x1e, 0x22 # sp + 16 + 4 * vlenb
+; CHECK-NEXT:    csrrs a1, vlenb, zero
+; CHECK-NEXT:    add a2, a0, a1
+; CHECK-NEXT:    vl1re64.v v8, (a2)
+; CHECK-NEXT:    slli a2, a1, 1
+; CHECK-NEXT:    vl1re64.v v9, (a0)
+; CHECK-NEXT:    add a0, a0, a2
+; CHECK-NEXT:    vl1re64.v v10, (a0)
+; CHECK-NEXT:    addi a0, sp, 16
+; CHECK-NEXT:    vs1r.v v9, (a0)
+; CHECK-NEXT:    add a2, a0, a2
+; CHECK-NEXT:    vs1r.v v10, (a2)
+; CHECK-NEXT:    add a0, a0, a1
+; CHECK-NEXT:    vs1r.v v8, (a0)
+; CHECK-NEXT:    csrrs a0, vlenb, zero
+; CHECK-NEXT:    slli a0, a0, 2
+; CHECK-NEXT:    add sp, sp, a0
+; CHECK-NEXT:    addi sp, sp, 16
+; CHECK-NEXT:    jalr zero, 0(ra)
+entry:
+  %ret = alloca %my_type, align 8
+  %val = load %my_type, ptr %addr
+  store %my_type %val, ptr %ret, align 8
+  ret void
+}

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

-; RUN: not opt -S -passes=verify < %s 2>&1 | FileCheck %s
-
-;; Arrays cannot contain scalable vectors; make sure we detect them even
-;; when nested inside other aggregates.
-
-%ty = type { i64, [4 x <vscale x 256 x i1>] }
-; CHECK: error: invalid array element type
-; CHECK: %ty = type { i64, [4 x <vscale x 256 x i1>] }

llvm/test/Transforms/InstCombine/scalable-vector-array.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
+; RUN: opt -passes=instcombine -S < %s | FileCheck %s
+
+define <vscale x 4 x i32> @load(ptr %x) {
+; CHECK-LABEL: define <vscale x 4 x i32> @load
+; CHECK-SAME: (ptr [[X:%.*]]) {
+; CHECK-NEXT:    [[A_ELT1:%.*]] = getelementptr inbounds [2 x <vscale x 4 x i32>], ptr [[X]], i64 0, i64 1
+; CHECK-NEXT:    [[A_UNPACK2:%.*]] = load <vscale x 4 x i32>, ptr [[A_ELT1]], align 16
+; CHECK-NEXT:    ret <vscale x 4 x i32> [[A_UNPACK2]]
+;
+  %a = load [2 x <vscale x 4 x i32>], ptr %x
+  %b = extractvalue [2 x <vscale x 4 x i32>] %a, 1
+  ret <vscale x 4 x i32> %b
+}
+
+define void @store(ptr %x, <vscale x 4 x i32> %y, <vscale x 4 x i32> %z) {
+; CHECK-LABEL: define void @store
+; CHECK-SAME: (ptr [[X:%.*]], <vscale x 4 x i32> [[Y:%.*]], <vscale x 4 x i32> [[Z:%.*]]) {
+; CHECK-NEXT:    store <vscale x 4 x i32> [[Y]], ptr [[X]], align 16
+; CHECK-NEXT:    [[X_REPACK1:%.*]] = getelementptr inbounds [2 x <vscale x 4 x i32>], ptr [[X]], i64 0, i64 1
+; CHECK-NEXT:    store <vscale x 4 x i32> [[Z]], ptr [[X_REPACK1]], align 16
+; CHECK-NEXT:    ret void
+;
+  %a = insertvalue [2 x <vscale x 4 x i32>] poison, <vscale x 4 x i32> %y, 0
+  %b = insertvalue [2 x <vscale x 4 x i32>] %a, <vscale x 4 x i32> %z, 1
+  store [2 x <vscale x 4 x i32>] %b, ptr %x
+  ret void
+}

llvm/test/Transforms/SROA/scalable-vector-array.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
+; RUN: opt < %s -passes='sroa<preserve-cfg>' -S | FileCheck %s
+; RUN: opt < %s -passes='sroa<modify-cfg>' -S | FileCheck %s
+
+; This test checks that SROA runs mem2reg on arrays of scalable vectors.
+
+define [ 2 x <vscale x 4 x i32> ] @alloca(<vscale x 4 x i32> %x, <vscale x 4 x i32> %y) {
+; CHECK-LABEL: define [2 x <vscale x 4 x i32>] @alloca
+; CHECK-SAME: (<vscale x 4 x i32> [[X:%.*]], <vscale x 4 x i32> [[Y:%.*]]) {
+; CHECK-NEXT:    [[AGG0:%.*]] = insertvalue [2 x <vscale x 4 x i32>] poison, <vscale x 4 x i32> [[X]], 0
+; CHECK-NEXT:    [[AGG1:%.*]] = insertvalue [2 x <vscale x 4 x i32>] [[AGG0]], <vscale x 4 x i32> [[Y]], 1
+; CHECK-NEXT:    ret [2 x <vscale x 4 x i32>] [[AGG1]]
+;
+  %addr = alloca [ 2 x <vscale x 4 x i32> ], align 4
+  %agg0 = insertvalue [ 2 x <vscale x 4 x i32> ] poison, <vscale x 4 x i32> %x, 0
+  %agg1 = insertvalue [ 2 x <vscale x 4 x i32> ] %agg0, <vscale x 4 x i32> %y, 1
+  store [ 2 x <vscale x 4 x i32> ] %agg1, ptr %addr, align 4
+  %val = load [ 2 x <vscale x 4 x i32> ], ptr %addr, align 4
+  ret [ 2 x <vscale x 4 x i32> ] %val
+}

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

 ; RUN: not opt -S -passes=verify < %s 2>&1 | FileCheck %s
 
 ;; Global variables cannot be scalable vectors, since we don't
 ;; know the size at compile time.
 
-; CHECK: Globals cannot contain scalable vectors
+; CHECK: Globals cannot contain scalable types
 ; CHECK-NEXT: ptr @ScalableVecGlobal
 @ScalableVecGlobal = global <vscale x 4 x i32> zeroinitializer
 
-; CHECK-NEXT: Globals cannot contain scalable vectors
+; CHECK-NEXT: Globals cannot contain scalable types
+; CHECK-NEXT: ptr @ScalableVecArrayGlobal
+@ScalableVecArrayGlobal = global [ 8 x  <vscale x 4 x i32> ] zeroinitializer
+
+; CHECK-NEXT: Globals cannot contain scalable types
 ; CHECK-NEXT: ptr @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 ptr zeroinitializer