[InstCombine] Support load-store forwarding with scalable vector types

[david-arm]

Summary

This patch fixes an invalid TypeSize->uint64_t implicit conversion in
FoldReinterpretLoadFromConst. We now check to see if the offset is
greater than or equal to the known minimum size of the constant. In
addition, I have added support for scalable vector splats in
ReadDataFromGlobal, since we can support load-store forwarding
when we know the offset is safe.

Tests added here:

Transforms/InstCombine/load-store-forward.ll

[nikic]

llvm/lib/Analysis/ConstantFolding.cpp

601–602

This is saying that accessing a scalable vector past its *minimum* size is undefined behavior. This doesn't look right to me. You probably want to return nullptr in that case, not UndefValue.

[david-arm]

• Return nullptr if the offset exceeds the known minimum value for a scalable vector.

[david-arm]

llvm/lib/Analysis/ConstantFolding.cpp

601–602

Well spotted, thanks @nikic !

[nikic]

llvm/lib/Analysis/ConstantFolding.cpp

601–602

I think the code in ReadDataFromGlobal suffers from a similar issue, where a read that starts before the min size but crossed over it will leave the remaining bytes as zero.

[david-arm]

llvm/lib/Analysis/ConstantFolding.cpp

601–602

I couldn't seem to create a test for that because whenever the loaded type is larger than the known minimum size of the constant type we don't perform the optimisation at all. Also, I wasn't able to construct a test case where the start offset is non-zero because then we don't even get into this function (I tried using getelementptr to load from an offset of the store ptr). If you have any suggestions on how to this cross-over case I'm happy to add them? Alternatively, I could just add an assert for that case?

[nikic]

llvm/lib/Analysis/ConstantFolding.cpp

601–602

Are scalable vectors supported as globals? If so, the more straightforward way to test this code is to perform a load from a scalable vector constant, in which case you can also test GEP with offset. The InstCombine code you're targeting currently only allows zero offsets.

[sdesmalen]

llvm/lib/Analysis/ConstantFolding.cpp

502

Can you split out all changes to make ReadDataFromGlobal work on scalable splats into a separate patch? I think the changes aren't trivial, and we want a fix for the assertion-failures in LLVM 14. The improvements to make this optimisation work for scalable vectors is less urgent.

512–513

I'd really like to avoid things like this, where a variable is named NumElts and we use the "Minimum known number of elements" from a scalable vector.

To avoid relying on the VectorType's ElementCount entirely, you can define the number of elements based on BytesLeft, i.e.:

Type *EltTy;
if (auto *AT = dyn_cast<ArrayType>(C->getType()))
  EltTy = AT->getElementType();
else
  EltTy = cast<VectorType>(C->getType())->getElementType();

uint64_t NumElts =
    std::max<uint64_t>(1, (8 * BytesLeft) / EltTy->getScalarSizeInBits());

521

Instead of checking for SplatVal and the special case for scalable vectors, I'd rather you just add this support to Constant::getAggregateElement directly.

[david-arm]

llvm/lib/Analysis/ConstantFolding.cpp

521

Hi @sdesmalen, I'm not sure this is right to be honest. It seems like we already have separate Constant::getAggregateElement and Constant::getSplatValue functions for a reason. If I understand you correctly, what you're asking for is basically something that does:

Constant *Constant::getAggregateElement(unsigned Elt) {
  if (Constant *SplatVal = getSplatValue(false))
    return SplatVal;
  .. existing code ..
}

If we do this then doesn't it also make sense to simply get rid of getSplatValue and replace all uses with getAggregateElement, rather than have two differently named interfaces that do the same thing?

[david-arm]

llvm/lib/Analysis/ConstantFolding.cpp

512–513

So the problem with this approach is that currently BytesLeft can be larger than the type you're reading from. For example, in test/Transforms/InstCombine/strncpy-3.ll we try to read 4 bytes from an array of 3 bytes. In your example code above that means NumElts would be 4 and we end up crashing. That means we still have to use the number of elements in the type to limit the reads/writes to only those that are valid. Perhaps in the first instance where we call ReadDataFromGlobal I can clamp the value of BytesLeft to the known minimum size of the type?

llvm/lib/Analysis/ConstantFolding.cpp

 namespace {
 
 /// Recursive helper to read bits out of global. C is the constant being copied
 /// out of. ByteOffset is an offset into C. CurPtr is the pointer to copy
 /// results into and BytesLeft is the number of bytes left in
 /// the CurPtr buffer. DL is the DataLayout.
 bool ReadDataFromGlobal(Constant *C, uint64_t ByteOffset, unsigned char *CurPtr,
                         unsigned BytesLeft, const DataLayout &DL) {
-  assert(ByteOffset <= DL.getTypeAllocSize(C->getType()) &&
+  assert(ByteOffset <= DL.getTypeAllocSize(C->getType()).getKnownMinValue() &&
          "Out of range access");
 
   // If this element is zero or undefined, we can just return since *CurPtr is
   // zero initialized.
   if (isa<ConstantAggregateZero>(C) || isa<UndefValue>(C))
     return true;
 
   if (auto *CI = dyn_cast<ConstantInt>(C)) {
while (true) {
       CurPtr += NextEltOffset - CurEltOffset - ByteOffset;
       BytesLeft -= NextEltOffset - CurEltOffset - ByteOffset;
       ByteOffset = 0;
       CurEltOffset = NextEltOffset;
     }
     // not reached.
   }
 
-  if (isa<ConstantArray>(C) || isa<ConstantVector>(C) ||
+  Constant *SplatVal =
+      isa<VectorType>(C->getType()) ? C->getSplatValue(false) : nullptr;
+  if (SplatVal || isa<ConstantArray>(C) || isa<ConstantVector>(C) ||
       isa<ConstantDataSequential>(C)) {
     uint64_t NumElts;
     Type *EltTy;
     if (auto *AT = dyn_cast<ArrayType>(C->getType())) {
       NumElts = AT->getNumElements();
       EltTy = AT->getElementType();
     } else {
-      NumElts = cast<FixedVectorType>(C->getType())->getNumElements();
-      EltTy = cast<FixedVectorType>(C->getType())->getElementType();
+      NumElts =
+          cast<VectorType>(C->getType())->getElementCount().getKnownMinValue();
+      EltTy = cast<VectorType>(C->getType())->getElementType();
     }
-    uint64_t EltSize = DL.getTypeAllocSize(EltTy);
+    uint64_t EltSize = DL.getTypeAllocSize(EltTy).getFixedSize();
     uint64_t Index = ByteOffset / EltSize;
     uint64_t Offset = ByteOffset - Index * EltSize;
 
     for (; Index != NumElts; ++Index) {
-      if (!ReadDataFromGlobal(C->getAggregateElement(Index), Offset, CurPtr,
-                              BytesLeft, DL))
+      Constant *El = SplatVal ? SplatVal : C->getAggregateElement(Index);
+      if (!ReadDataFromGlobal(El, Offset, CurPtr, BytesLeft, DL))
         return false;
 
       uint64_t BytesWritten = EltSize - Offset;
       assert(BytesWritten <= EltSize && "Not indexing into this element?");
       if (BytesWritten >= BytesLeft)
         return true;
 
       Offset = 0;
if (!IntType) {
     }
     return nullptr;
   }
 
   unsigned BytesLoaded = (IntType->getBitWidth() + 7) / 8;
   if (BytesLoaded > 32 || BytesLoaded == 0)
     return nullptr;
 
-  int64_t InitializerSize = DL.getTypeAllocSize(C->getType()).getFixedSize();
+  TypeSize InitializerSize = DL.getTypeAllocSize(C->getType());
 
   // If we're not accessing anything in this constant, the result is undefined.
   if (Offset <= -1 * static_cast<int64_t>(BytesLoaded))
     return UndefValue::get(IntType);
 
   // If we're not accessing anything in this constant, the result is undefined.
-  if (Offset >= InitializerSize)
-    return UndefValue::get(IntType);
+  if (Offset >= static_cast<int64_t>(InitializerSize.getKnownMinValue()))
+    return InitializerSize.isScalable() ? nullptr : UndefValue::get(IntType);
 
   unsigned char RawBytes[32] = {0};
   unsigned char *CurPtr = RawBytes;
   unsigned BytesLeft = BytesLoaded;
 
   // If we're loading off the beginning of the global, some bytes may be valid.
   if (Offset < 0) {
     CurPtr += -Offset;

llvm/test/Transforms/InstCombine/load-store-forward.ll

;
   %p2 = bitcast i32* %p to <2 x i32>*
   store <2 x i32> <i32 1, i32 2>, <2 x i32>* %p2
   %p3 = bitcast i32* %p to i8*
   %p4 = getelementptr i8, i8* %p3, i64 2
   %p5 = bitcast i8* %p4 to i32*
   %load = load i32, i32* %p5, align 2
   ret i32 %load
 }
+
+define i32 @load_i32_store_nxv4i32(i32* %a) vscale_range(1,16) {
+; CHECK-LABEL: @load_i32_store_nxv4i32(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i32* [[A:%.*]] to <vscale x 4 x i32>*
+; CHECK-NEXT:    store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* [[TMP0]], align 16
+; CHECK-NEXT:    ret i32 1
+;
+entry:
+  %0 = bitcast i32* %a to <vscale x 4 x i32>*
+  store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* %0, align 16
+  %1 = load i32, i32* %a, align 4
+  ret i32 %1
+}
+
+define i64 @load_i64_store_nxv8i8(i8* %a) vscale_range(1,16) {
+; CHECK-LABEL: @load_i64_store_nxv8i8(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i8* [[A:%.*]] to <vscale x 8 x i8>*
+; CHECK-NEXT:    store <vscale x 8 x i8> shufflevector (<vscale x 8 x i8> insertelement (<vscale x 8 x i8> poison, i8 1, i32 0), <vscale x 8 x i8> poison, <vscale x 8 x i32> zeroinitializer), <vscale x 8 x i8>* [[TMP0]], align 16
+; CHECK-NEXT:    ret i64 72340172838076673
+;
+entry:
+  %0 = bitcast i8* %a to <vscale x 8 x i8>*
+  store <vscale x 8 x i8> shufflevector (<vscale x 8 x i8> insertelement (<vscale x 8 x i8> poison, i8 1, i32 0), <vscale x 8 x i8> poison, <vscale x 8 x i32> zeroinitializer), <vscale x 8 x i8>* %0, align 16
+  %a2 = bitcast i8* %a to i64*
+  %load = load i64, i64* %a2, align 8
+  ret i64 %load
+}
+
+define i64 @load_i64_store_nxv4i32(i32* %a) vscale_range(1,16) {
+; CHECK-LABEL: @load_i64_store_nxv4i32(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i32* [[A:%.*]] to <vscale x 4 x i32>*
+; CHECK-NEXT:    store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* [[TMP0]], align 16
+; CHECK-NEXT:    ret i64 4294967297
+;
+entry:
+  %0 = bitcast i32* %a to <vscale x 4 x i32>*
+  store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* %0, align 16
+  %a2 = bitcast i32* %a to i64*
+  %load = load i64, i64* %a2, align 8
+  ret i64 %load
+}
+
+define i8 @load_i8_store_nxv4i32(i32* %a) vscale_range(1,16) {
+; LITTLE-LABEL: @load_i8_store_nxv4i32(
+; LITTLE-NEXT:  entry:
+; LITTLE-NEXT:    [[TMP0:%.*]] = bitcast i32* [[A:%.*]] to <vscale x 4 x i32>*
+; LITTLE-NEXT:    store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* [[TMP0]], align 16
+; LITTLE-NEXT:    ret i8 1
+;
+; BIG-LABEL: @load_i8_store_nxv4i32(
+; BIG-NEXT:  entry:
+; BIG-NEXT:    [[TMP0:%.*]] = bitcast i32* [[A:%.*]] to <vscale x 4 x i32>*
+; BIG-NEXT:    store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* [[TMP0]], align 16
+; BIG-NEXT:    ret i8 0
+;
+entry:
+  %0 = bitcast i32* %a to <vscale x 4 x i32>*
+  store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* %0, align 16
+  %a2 = bitcast i32* %a to i8*
+  %load = load i8, i8* %a2, align 1
+  ret i8 %load
+}
+
+define float @load_f32_store_nxv4f32(float* %a) vscale_range(1,16) {
+; CHECK-LABEL: @load_f32_store_nxv4f32(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast float* [[A:%.*]] to <vscale x 4 x float>*
+; CHECK-NEXT:    store <vscale x 4 x float> shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> poison, float 1.000000e+00, i64 0), <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x float>* [[TMP0]], align 16
+; CHECK-NEXT:    ret float 1.000000e+00
+;
+entry:
+  %0 = bitcast float* %a to <vscale x 4 x float>*
+  store <vscale x 4 x float> shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> poison, float 1.0, i64 0), <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x float>* %0, align 16
+  %1 = load float, float* %a, align 4
+  ret float %1
+}
+
+define i32 @load_i32_store_nxv4f32(float* %a) vscale_range(1,16) {
+; CHECK-LABEL: @load_i32_store_nxv4f32(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast float* [[A:%.*]] to <vscale x 4 x float>*
+; CHECK-NEXT:    store <vscale x 4 x float> shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> poison, float 1.000000e+00, i64 0), <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x float>* [[TMP0]], align 16
+; CHECK-NEXT:    ret i32 1065353216
+;
+entry:
+  %0 = bitcast float* %a to <vscale x 4 x float>*
+  store <vscale x 4 x float> shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> poison, float 1.0, i64 0), <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x float>* %0, align 16
+  %a2 = bitcast float* %a to i32*
+  %load = load i32, i32* %a2, align 4
+  ret i32 %load
+}
+
+define <4 x i32> @load_v4i32_store_nxv4i32(i32* %a) vscale_range(1,16) {
+; CHECK-LABEL: @load_v4i32_store_nxv4i32(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i32* [[A:%.*]] to <vscale x 4 x i32>*
+; CHECK-NEXT:    store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* [[TMP0]], align 16
+; CHECK-NEXT:    ret <4 x i32> <i32 1, i32 1, i32 1, i32 1>
+;
+entry:
+  %0 = bitcast i32* %a to <vscale x 4 x i32>*
+  store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* %0, align 16
+  %1 = bitcast i32* %a to <4 x i32>*
+  %2 = load <4 x i32>, <4 x i32>* %1, align 16
+  ret <4 x i32> %2
+}
+
+define <4 x i16> @load_v4i16_store_nxv4i32(i32* %a) vscale_range(1,16) {
+; LITTLE-LABEL: @load_v4i16_store_nxv4i32(
+; LITTLE-NEXT:  entry:
+; LITTLE-NEXT:    [[TMP0:%.*]] = bitcast i32* [[A:%.*]] to <vscale x 4 x i32>*
+; LITTLE-NEXT:    store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* [[TMP0]], align 16
+; LITTLE-NEXT:    ret <4 x i16> <i16 1, i16 0, i16 1, i16 0>
+;
+; BIG-LABEL: @load_v4i16_store_nxv4i32(
+; BIG-NEXT:  entry:
+; BIG-NEXT:    [[TMP0:%.*]] = bitcast i32* [[A:%.*]] to <vscale x 4 x i32>*
+; BIG-NEXT:    store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* [[TMP0]], align 16
+; BIG-NEXT:    ret <4 x i16> <i16 0, i16 1, i16 0, i16 1>
+;
+entry:
+  %0 = bitcast i32* %a to <vscale x 4 x i32>*
+  store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* %0, align 16
+  %1 = bitcast i32* %a to <4 x i16>*
+  %2 = load <4 x i16>, <4 x i16>* %1, align 16
+  ret <4 x i16> %2
+}
+
+; Loaded data type exceeds the known minimum size of the store.
+define i64 @load_i64_store_nxv4i8(i8* %a) vscale_range(1,16) {
+; CHECK-LABEL: @load_i64_store_nxv4i8(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i8* [[A:%.*]] to <vscale x 4 x i8>*
+; CHECK-NEXT:    store <vscale x 4 x i8> shufflevector (<vscale x 4 x i8> insertelement (<vscale x 4 x i8> poison, i8 1, i32 0), <vscale x 4 x i8> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i8>* [[TMP0]], align 16
+; CHECK-NEXT:    [[A2:%.*]] = bitcast i8* [[A]] to i64*
+; CHECK-NEXT:    [[LOAD:%.*]] = load i64, i64* [[A2]], align 8
+; CHECK-NEXT:    ret i64 [[LOAD]]
+;
+entry:
+  %0 = bitcast i8* %a to <vscale x 4 x i8>*
+  store <vscale x 4 x i8> shufflevector (<vscale x 4 x i8> insertelement (<vscale x 4 x i8> poison, i8 1, i32 0), <vscale x 4 x i8> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i8>* %0, align 16
+  %a2 = bitcast i8* %a to i64*
+  %load = load i64, i64* %a2, align 8
+  ret i64 %load
+}
+
+; Loaded data size is unknown - we cannot guarantee it won't
+; exceed the store size.
+define <vscale x 4 x i8> @load_nxv4i8_store_nxv4i32(i32* %a) vscale_range(1,16) {
+; CHECK-LABEL: @load_nxv4i8_store_nxv4i32(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i32* [[A:%.*]] to <vscale x 4 x i32>*
+; CHECK-NEXT:    store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* [[TMP0]], align 16
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i32* [[A]] to <vscale x 4 x i8>*
+; CHECK-NEXT:    [[TMP2:%.*]] = load <vscale x 4 x i8>, <vscale x 4 x i8>* [[TMP1]], align 16
+; CHECK-NEXT:    ret <vscale x 4 x i8> [[TMP2]]
+;
+entry:
+  %0 = bitcast i32* %a to <vscale x 4 x i32>*
+  store <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer), <vscale x 4 x i32>* %0, align 16
+  %1 = bitcast i32* %a to <vscale x 4 x i8>*
+  %2 = load <vscale x 4 x i8>, <vscale x 4 x i8>* %1, align 16
+  ret <vscale x 4 x i8> %2
+}