[Transforms][SROA] Skip uses of allocas where the type is scalable

[c-rhodes]

Summary

When splitting an alloca of a fixed-length vector (VLS) and making a
slice for each load and store, skip uses where the VLS is bitcasted to a
scalable vector (VLA).

This is relevant in the implementation of the 'arm_sve_vector_bits'
attribute that is used to define VLS types. See D85128 for more
information.

[david-arm]

llvm/test/Transforms/SROA/scalable-vectors.ll

56

I think this is undefined behaviour for anything other than 512 fixed length vectors. Is it possible to amend the opt line to specify the fixed length vector size?

[c-rhodes]

llvm/test/Transforms/SROA/scalable-vectors.ll

56

I think this is undefined behaviour for anything other than 512 fixed length vectors. Is it possible to amend the opt line to specify the fixed length vector size?

Good point, I can add the -aarch64-sve-vector-bits-min= flag.

[c-rhodes]

Specify -aarch64-sve-vector-bits-min=512 in test.

[c-rhodes]

llvm/test/Transforms/SROA/scalable-vectors.ll

56

I think this is undefined behaviour for anything other than 512 fixed length vectors. Is it possible to amend the opt line to specify the fixed length vector size?

Good point, I can add the -aarch64-sve-vector-bits-min= flag.

I've added new check lines for these tests that specify -aarch64-sve-vector-bits-min=512. I also tried specifying it in an attribute group but I wasn't sure it was having any effect since it didn't complain when I put garbage in there.

[efriedma]

LGTM

[efriedma]

I also tried specifying it in an attribute group but I wasn't sure it was having any effect since it didn't complain when I put garbage in there.

Err, right... aarch64-sve-vector-bits-min is a backend flag; it's not going to have any effect on SROA. (I'm not sure what effect it could have, anyway.) So probably the extra run lines aren't useful.

[c-rhodes]

In D85725#2211150, @efriedma wrote:

I also tried specifying it in an attribute group but I wasn't sure it was having any effect since it didn't complain when I put garbage in there.

Err, right... aarch64-sve-vector-bits-min is a backend flag; it's not going to have any effect on SROA. (I'm not sure what effect it could have, anyway.) So probably the extra run lines aren't useful.

Ah ok, I'll remove them then.

llvm/lib/Transforms/Scalar/SROA.cpp

void visitLoadInst(LoadInst &LI) {
 
     if (!IsOffsetKnown)
       return PI.setAborted(&LI);
 
     if (LI.isVolatile() &&
         LI.getPointerAddressSpace() != DL.getAllocaAddrSpace())
       return PI.setAborted(&LI);
 
+    if (isa<ScalableVectorType>(LI.getType()))
+      return PI.setAborted(&LI);
+
     uint64_t Size = DL.getTypeStoreSize(LI.getType()).getFixedSize();
     return handleLoadOrStore(LI.getType(), LI, Offset, Size, LI.isVolatile());
   }
 
   void visitStoreInst(StoreInst &SI) {
     Value *ValOp = SI.getValueOperand();
     if (ValOp == *U)
       return PI.setEscapedAndAborted(&SI);
     if (!IsOffsetKnown)
       return PI.setAborted(&SI);
 
     if (SI.isVolatile() &&
         SI.getPointerAddressSpace() != DL.getAllocaAddrSpace())
       return PI.setAborted(&SI);
 
+    if (isa<ScalableVectorType>(ValOp->getType()))
+      return PI.setAborted(&SI);
+
     uint64_t Size = DL.getTypeStoreSize(ValOp->getType()).getFixedSize();
 
     // If this memory access can be shown to *statically* extend outside the
     // bounds of the allocation, it's behavior is undefined, so simply
     // ignore it. Note that this is more strict than the generic clamping
     // behavior of insertUse. We also try to handle cases which might run the
     // risk of overflow.
     // FIXME: We should instead consider the pointer to have escaped if this
static Value *getNaturalGEPWithOffset(IRBuilderTy &IRB, const DataLayout &DL,
   // Don't consider any GEPs through an i8* as natural unless the TargetTy is
   // an i8.
   if (Ty == IRB.getInt8PtrTy(Ty->getAddressSpace()) && TargetTy->isIntegerTy(8))
     return nullptr;
 
   Type *ElementTy = Ty->getElementType();
   if (!ElementTy->isSized())
     return nullptr; // We can't GEP through an unsized element.
+  if (isa<ScalableVectorType>(ElementTy))
+    return nullptr;
   APInt ElementSize(Offset.getBitWidth(),
                     DL.getTypeAllocSize(ElementTy).getFixedSize());
   if (ElementSize == 0)
     return nullptr; // Zero-length arrays can't help us build a natural GEP.
   APInt NumSkippedElements = Offset.sdiv(ElementSize);
 
   Offset -= NumSkippedElements * ElementSize;
   Indices.push_back(IRB.getInt(NumSkippedElements));

llvm/test/Transforms/SROA/scalable-vectors.ll

 ; RUN: opt < %s -sroa -S | FileCheck %s
 ; RUN: opt < %s -passes=sroa -S | FileCheck %s
+; RUN: opt -aarch64-sve-vector-bits-min=512 < %s -sroa -S | FileCheck %s --check-prefix=CHECK-VLS
+; RUN: opt -aarch64-sve-vector-bits-min=512 < %s -passes=sroa -S | FileCheck %s --check-prefix=CHECK-VLS
 
 ; This test checks that SROA runs mem2reg on scalable vectors.
 
 define <vscale x 16 x i1> @alloca_nxv16i1(<vscale x 16 x i1> %pg) {
 ; CHECK-LABEL: alloca_nxv16i1
 ; CHECK-NEXT: ret <vscale x 16 x i1> %pg
   %pg.addr = alloca <vscale x 16 x i1>
   store <vscale x 16 x i1> %pg, <vscale x 16 x i1>* %pg.addr

 ; CHECK-NEXT: store volatile <vscale x 16 x i8> %vec, <vscale x 16 x i8>* %vec.addr
 ; CHECK-NEXT: %1 = load volatile <vscale x 16 x i8>, <vscale x 16 x i8>* %vec.addr
 ; CHECK-NEXT: ret <vscale x 16 x i8> %1
   %vec.addr = alloca <vscale x 16 x i8>
   store volatile <vscale x 16 x i8> %vec, <vscale x 16 x i8>* %vec.addr
   %1 = load volatile <vscale x 16 x i8>, <vscale x 16 x i8>* %vec.addr
   ret <vscale x 16 x i8> %1
 }
+
+; Test we bail out when using an alloca of a fixed-length vector (VLS) that was
+; bitcasted to a scalable vector.
+define <vscale x 4 x i32> @cast_alloca_to_svint32_t(<vscale x 4 x i32> %type.coerce) {
+; CHECK-VLS-LABEL: cast_alloca_to_svint32_t
+; CHECK-VLS-NEXT: %type = alloca <16 x i32>, align 64
+; CHECK-VLS-NEXT: %type.addr = alloca <16 x i32>, align 64
+; CHECK-VLS-NEXT: %1 = bitcast <16 x i32>* %type to <vscale x 4 x i32>*
+; CHECK-VLS-NEXT: store <vscale x 4 x i32> %type.coerce, <vscale x 4 x i32>* %1, align 16
+; CHECK-VLS-NEXT: %type1 = load <16 x i32>, <16 x i32>* %type, align 64
+; CHECK-VLS-NEXT: store <16 x i32> %type1, <16 x i32>* %type.addr, align 64
+; CHECK-VLS-NEXT: %2 = load <16 x i32>, <16 x i32>* %type.addr, align 64
+; CHECK-VLS-NEXT: %3 = bitcast <16 x i32>* %type.addr to <vscale x 4 x i32>*
+; CHECK-VLS-NEXT: %4 = load <vscale x 4 x i32>, <vscale x 4 x i32>* %3, align 16
+; CHECK-VLS-NEXT: ret <vscale x 4 x i32> %4
+  %type = alloca <16 x i32>
+  %type.addr = alloca <16 x i32>
+  %1 = bitcast <16 x i32>* %type to <vscale x 4 x i32>*
+  store <vscale x 4 x i32> %type.coerce, <vscale x 4 x i32>* %1
+  %type1 = load <16 x i32>, <16 x i32>* %type
+  store <16 x i32> %type1, <16 x i32>* %type.addr
+  %2 = load <16 x i32>, <16 x i32>* %type.addr
+  %3 = bitcast <16 x i32>* %type.addr to <vscale x 4 x i32>*
+  %4 = load <vscale x 4 x i32>, <vscale x 4 x i32>* %3
+  ret <vscale x 4 x i32> %4
+}
+
+; When casting from VLA to VLS via memory check we bail out when producing a
+; GEP where the element type is a scalable vector.
+define <vscale x 4 x i32> @cast_alloca_from_svint32_t() {
+; CHECK-VLS-LABEL: cast_alloca_from_svint32_t
+; CHECK-VLS-NEXT: %retval.coerce = alloca <vscale x 4 x i32>, align 16
+; CHECK-VLS-NEXT: %1 = bitcast <vscale x 4 x i32>* %retval.coerce to i8*
+; CHECK-VLS-NEXT: %retval.0..sroa_cast = bitcast i8* %1 to <16 x i32>*
+; CHECK-VLS-NEXT: store <16 x i32> undef, <16 x i32>* %retval.0..sroa_cast, align 16
+; CHECK-VLS-NEXT: %2 = load <vscale x 4 x i32>, <vscale x 4 x i32>* %retval.coerce, align 16
+; CHECK-VLS-NEXT: ret <vscale x 4 x i32> %2
+  %retval = alloca <16 x i32>
+  %retval.coerce = alloca <vscale x 4 x i32>
+  %1 = bitcast <vscale x 4 x i32>* %retval.coerce to i8*
+  %2 = bitcast <16 x i32>* %retval to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 16 %1, i8* align 16 %2, i64 64, i1 false)
+  %3 = load <vscale x 4 x i32>, <vscale x 4 x i32>* %retval.coerce
+  ret <vscale x 4 x i32> %3
+}
+
+declare void @llvm.memcpy.p0i8.p0i8.i64(i8* nocapture, i8* nocapture, i64, i1) nounwind