[IR] GEP: Fix byte-offsets in vectors of overaligned types

[jsilvanus]

Summary

Vectors contain their elements tightly packed together without
any padding bytes. If the elements have stricter-than-natural
alignment requirements, then the elements in the vector are smaller
than in for example an array or struct.

This fact was not accounted for when computing byte-offsets of GEPs.
For example, with i16:32 alignment, the byte-offset of the following
GEP was incorrectly computed as 8 bytes, while 4 bytes is correct:

getelementptr <3 x i16>, <3 x i16> *%ptr, i32 0, i32 2

To fix, instead of unconditionally using the elements' AllocSize,
use a new helper getElementSize(..) that returns the inner type's
size if the outer type is a vector, otherwise its AllocSize.

Add a unit test.

Also add an assertion to SROA:
SROA tries to generate "natural" GEPs for a given byte offset if possible,
instead of bitcasting to i8* and applying the offset there.
Add an assertion that such a generated natural GEP results in the correct
byte offset. This assertion would have failed prior to the GEP fix.

Update the LangRef to be more explicit on the layout of vectors of
overaligned elements.

The lit test is new. I'll precommit it if this patch is accepted.

[jsilvanus]

Adding reviewers, based on past changes to GetElementPtrTypeIterator.h.
Please feel free to reject or suggest other reviewers that might have a better background.

[nikic]

I think the general idea behind this change is correct. Vectors should indeed be tightly packed and ignore ABI alignment of the element type.

However, I'm afraid that the patch in the current form will make the situation worse by introducing inconsistent treatment in different places. We have a lot more code that is repeating basically the same pattern (getAllocSize on the GTI indexed type). BasicAA would be an obvious example, but I see quite a few more uses spread over the codebase from a grep on getTypeAllocSize(GTI.

I think that ideally we would add a gep_offset_iterator which exposes the GEP as an addition of scaled multiplies and constants only and hides the type-related details. This can be used in most places and would give a central place to update the offset calculation logic.

llvm/lib/Transforms/Scalar/SROA.cpp

1542

This entire function is unused if opaque pointers are used, so I don't think it makes much sense to add this assertion.

[jsilvanus]

Thanks for the review.
Your idea of adding a gep_offset_iterator seems plausible. I'll look into that.

As a heads-up notice, it may take a while until I can continue working on this.

[jsilvanus]

Abandoning this change in favor of avoiding GEPs into vectors instead of fixing their offsets.

I'll start with avoiding to generate vector GEPs in SROA.
In the future, we may forbid vector GEPs entirely, see https://discourse.llvm.org/t/status-of-geps-into-vectors-of-overaligned-elements/67497.

llvm/docs/LangRef.rst

 and a scalable property to represent vectors where the exact hardware
 vector length is unknown at compile time. Vector types are considered
 :ref:`first class <t_firstclass>`.
 
 :Memory Layout:
 
 In general vector elements are laid out in memory in the same way as
 :ref:`array types <t_array>`. Such an analogy works fine as long as the vector
-elements are byte sized. However, when the elements of the vector aren't byte
-sized it gets a bit more complicated. One way to describe the layout is by
-describing what happens when a vector such as <N x iM> is bitcasted to an
-integer type with N*M bits, and then following the rules for storing such an
-integer to memory.
+elements are byte sized and naturally aligned. Otherwise, it gets a bit more
+complicated. One way to describe the layout is by describing what happens
+when a vector such as <N x iM> is bitcasted to an integer type with N*M bits,
+and then following the rules for storing such an integer to memory.
 
 A bitcast from a vector type to a scalar integer type will see the elements
 being packed together (without padding). The order in which elements are
 inserted in the integer depends on endianess. For little endian element zero
 is put in the least significant bits of the integer, and for big endian
 element zero is put in the most significant bits.
 
 Using a vector such as ``<i4 1, i4 2, i4 3, i4 5>`` as an example, together

llvm/include/llvm/IR/GetElementPtrTypeIterator.h

 #include <iterator>
 
 namespace llvm {
 
 template <typename ItTy = User::const_op_iterator>
 class generic_gep_type_iterator {
 
   ItTy OpIt;
-  PointerUnion<StructType *, Type *> CurTy;
+  PointerUnion<StructType *, VectorType *, Type *> CurTy;
 
   generic_gep_type_iterator() = default;
 
 public:
   using iterator_category = std::forward_iterator_tag;
   using value_type = Type *;
   using difference_type = std::ptrdiff_t;
   using pointer = value_type *;
public:
 
   // FIXME: Make this the iterator's operator*() after the 4.0 release.
   // operator*() had a different meaning in earlier releases, so we're
   // temporarily not giving this iterator an operator*() to avoid a subtle
   // semantics break.
   Type *getIndexedType() const {
     if (auto *T = CurTy.dyn_cast<Type *>())
       return T;
-    return CurTy.get<StructType *>()->getTypeAtIndex(getOperand());
+    if (auto *STy = CurTy.dyn_cast<StructType *>())
+      return STy->getTypeAtIndex(getOperand());
+    return CurTy.get<VectorType *>()->getElementType();
   }
 
   Value *getOperand() const { return const_cast<Value *>(&**OpIt); }
 
   generic_gep_type_iterator &operator++() { // Preincrement
     Type *Ty = getIndexedType();
     if (auto *ATy = dyn_cast<ArrayType>(Ty))
       CurTy = ATy->getElementType();
     else if (auto *VTy = dyn_cast<VectorType>(Ty))
-      CurTy = VTy->getElementType();
+      CurTy = VTy;
     else
       CurTy = dyn_cast<StructType>(Ty);
     ++OpIt;
     return *this;
   }
 
   generic_gep_type_iterator operator++(int) { // Postincrement
     generic_gep_type_iterator tmp = *this;
public:
   // FIXME: Most current users of this class are just interested in byte
   // offsets (a few need to know whether the outer type is a struct because
   // they are trying to replace a constant with a variable, which is only
   // legal for arrays, e.g. canReplaceOperandWithVariable in SimplifyCFG.cpp);
   // we should provide a more minimal API here that exposes not much more than
   // that.
 
   bool isStruct() const { return CurTy.is<StructType *>(); }
-  bool isSequential() const { return CurTy.is<Type *>(); }
+  bool isVector() const { return CurTy.is<VectorType *>(); }
+  bool isSequential() const { return !isStruct(); }
 
   StructType *getStructType() const { return CurTy.get<StructType *>(); }
 
   StructType *getStructTypeOrNull() const {
     return CurTy.dyn_cast<StructType *>();
   }
 };
 

llvm/include/llvm/IR/Operator.h

enum {
     // InRangeIndex: bits 1-6
   };
 
   void setIsInBounds(bool B) {
     SubclassOptionalData =
       (SubclassOptionalData & ~IsInBounds) | (B * IsInBounds);
   }
 
+  // Get the size of the indexed element in its containing outer type.
+  //
+  // If OuterType is a VectorType, the unpadded element size is returned,
+  // which must be byte-aligned.
+  // Otherwise (if OuterType is a StructType or ArrayType), the indexed
+  // element's AllocSize is returned.
+  //
+  // Useful to compute byte-based offsets of elements within the outer type.
+  static TypeSize getElementSize(const DataLayout &DL, Type *ElementTy,
+                                 bool OuterIsVector);
+
 public:
   /// Test whether this is an inbounds GEP, as defined by LangRef.html.
   bool isInBounds() const {
     return SubclassOptionalData & IsInBounds;
   }
 
   /// Returns the offset of the index with an inrange attachment, or None if
   /// none.

llvm/lib/IR/Operator.cpp

 }
 
 Type *GEPOperator::getResultElementType() const {
   if (auto *I = dyn_cast<GetElementPtrInst>(this))
     return I->getResultElementType();
   return cast<GetElementPtrConstantExpr>(this)->getResultElementType();
 }
 
+TypeSize GEPOperator::getElementSize(const DataLayout &DL, Type *ElementTy,
+                                     bool OuterIsVector) {
+  if (!OuterIsVector)
+    return DL.getTypeAllocSize(ElementTy);
+
+  auto BitSize = DL.getTypeSizeInBits(ElementTy);
+  assert(BitSize % 8 == 0 && "GEP element size must be byte-aligned!");
+  return {BitSize / 8, BitSize.isScalable()};
+}
+
 Align GEPOperator::getMaxPreservedAlignment(const DataLayout &DL) const {
   /// compute the worse possible offset for every level of the GEP et accumulate
   /// the minimum alignment into Result.
 
   Align Result = Align(llvm::Value::MaximumAlignment);
   for (gep_type_iterator GTI = gep_type_begin(this), GTE = gep_type_end(this);
        GTI != GTE; ++GTI) {
     int64_t Offset = 1;
     ConstantInt *OpC = dyn_cast<ConstantInt>(GTI.getOperand());
 
     if (StructType *STy = GTI.getStructTypeOrNull()) {
       const StructLayout *SL = DL.getStructLayout(STy);
       Offset = SL->getElementOffset(OpC->getZExtValue());
     } else {
       assert(GTI.isSequential() && "should be sequencial");
       /// If the index isn't know we take 1 because it is the index that will
       /// give the worse alignment of the offset.
       int64_t ElemCount = 1;
       if (OpC)
         ElemCount = OpC->getZExtValue();
-      Offset = DL.getTypeAllocSize(GTI.getIndexedType()) * ElemCount;
+      Offset =
+          getElementSize(DL, GTI.getIndexedType(), GTI.isVector()) * ElemCount;
     }
     Result = Align(MinAlign(Offset, Result.value()));
   }
   return Result;
 }
 
 bool GEPOperator::accumulateConstantOffset(
     const DataLayout &DL, APInt &Offset,
if (auto ConstOffset = dyn_cast<ConstantInt>(V)) {
         const StructLayout *SL = DL.getStructLayout(STy);
         // Element offset is in bytes.
         if (!AccumulateOffset(
                 APInt(Offset.getBitWidth(), SL->getElementOffset(ElementIdx)),
                 1))
           return false;
         continue;
       }
-      if (!AccumulateOffset(ConstOffset->getValue(),
-                            DL.getTypeAllocSize(GTI.getIndexedType())))
+      if (!AccumulateOffset(
+              ConstOffset->getValue(),
+              getElementSize(DL, GTI.getIndexedType(), GTI.isVector())))
         return false;
       continue;
     }
 
     // The operand is not constant, check if an external analysis was provided.
     // External analsis is not applicable to a struct type.
     if (!ExternalAnalysis || STy || ScalableType)
       return false;
     APInt AnalysisIndex;
     if (!ExternalAnalysis(*V, AnalysisIndex))
       return false;
     UsedExternalAnalysis = true;
-    if (!AccumulateOffset(AnalysisIndex,
-                          DL.getTypeAllocSize(GTI.getIndexedType())))
+    if (!AccumulateOffset(
+            AnalysisIndex,
+            getElementSize(DL, GTI.getIndexedType(), GTI.isVector())))
       return false;
   }
   return true;
 }
 
 bool GEPOperator::collectOffset(
     const DataLayout &DL, unsigned BitWidth,
     MapVector<Value *, APInt> &VariableOffsets,
if (auto ConstOffset = dyn_cast<ConstantInt>(V)) {
       if (STy) {
         unsigned ElementIdx = ConstOffset->getZExtValue();
         const StructLayout *SL = DL.getStructLayout(STy);
         // Element offset is in bytes.
         CollectConstantOffset(APInt(BitWidth, SL->getElementOffset(ElementIdx)),
                               1);
         continue;
       }
-      CollectConstantOffset(ConstOffset->getValue(),
-                            DL.getTypeAllocSize(GTI.getIndexedType()));
+      CollectConstantOffset(
+          ConstOffset->getValue(),
+          getElementSize(DL, GTI.getIndexedType(), GTI.isVector()));
       continue;
     }
 
     if (STy || ScalableType)
       return false;
-    APInt IndexedSize =
-        APInt(BitWidth, DL.getTypeAllocSize(GTI.getIndexedType()));
+    APInt IndexedSize = APInt(
+        BitWidth, getElementSize(DL, GTI.getIndexedType(), GTI.isVector()));
     // Insert an initial offset of 0 for V iff none exists already, then
     // increment the offset by IndexedSize.
     if (!IndexedSize.isZero()) {
       VariableOffsets.insert({V, APInt(BitWidth, 0)});
       VariableOffsets[V] += IndexedSize;
     }
   }
   return true;

llvm/lib/Transforms/Scalar/SROA.cpp

 /// possible. We recurse by decreasing Offset, adding the appropriate index to
 /// Indices, and setting Ty to the result subtype.
 ///
 /// If no natural GEP can be constructed, this function returns null.
 static Value *getNaturalGEPWithOffset(IRBuilderTy &IRB, const DataLayout &DL,
                                       Value *Ptr, APInt Offset, Type *TargetTy,
                                       SmallVectorImpl<Value *> &Indices,
                                       const Twine &NamePrefix) {
+#ifndef NDEBUG
+  APInt OrigOffset = Offset;
+#endif
   PointerType *Ty = cast<PointerType>(Ptr->getType());
 
   // 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->getNonOpaquePointerElementType();
   if (!ElementTy->isSized())
     return nullptr; // We can't GEP through an unsized element.
 
   SmallVector<APInt> IntIndices = DL.getGEPIndicesForOffset(ElementTy, Offset);
   if (Offset != 0)
     return nullptr;
 
   for (const APInt &Index : IntIndices)
     Indices.push_back(IRB.getInt(Index));
-  return getNaturalGEPWithType(IRB, DL, Ptr, ElementTy, TargetTy, Indices,
-                               NamePrefix);
+  Value *Result = getNaturalGEPWithType(IRB, DL, Ptr, ElementTy, TargetTy,
+                                        Indices, NamePrefix);
+#ifndef NDEBUG
+  auto *GEP = dyn_cast<GetElementPtrInst>(Result);
+  if (GEP && GEP->getPointerOperand() == Ptr) {
+    APInt GEPOffset(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
+    assert(GEP->accumulateConstantOffset(DL, GEPOffset) &&
+           "Expected GEP with constant offset!");
+    assert(APInt::isSameValue(GEPOffset, OrigOffset) &&
+           "GEP has incorrect offset!");
+  }
+#endif
+
+  return Result;
 }
 
 /// Compute an adjusted pointer from Ptr by Offset bytes where the
 /// resulting pointer has PointerTy.
 ///
 /// This tries very hard to compute a "natural" GEP which arrives at the offset
 /// and produces the pointer type desired. Where it cannot, it will try to use
 /// the natural GEP to arrive at the offset and bitcast to the type. Where that

llvm/test/Transforms/SROA/overaligned-datalayout.ll

 ; After storing to the alloca, bitcast to an i8,
 ; apply a GEP, load the ptr, and return as result.
 ; Because the i8 GEP ends up at the start of an i16
 ; within the vector, SROA transforms it to a "natural GEP"
 ; on the vector.
 %VecStruct = type { <4 x i16> }
 define i8 @test_vector_bitcast_i8() {
 ; OVERALIGNED-LABEL: @test_vector_bitcast_i8(
-; OVERALIGNED-NEXT:    ret i8 poison
+; OVERALIGNED-NEXT:    [[ALLOCA_SROA_0:%.*]] = alloca <4 x i16>, align 8
+; OVERALIGNED-NEXT:    store <4 x i16> <i16 0, i16 1, i16 2, i16 3>, <4 x i16>* [[ALLOCA_SROA_0]], align 8
+; OVERALIGNED-NEXT:    [[ALLOCA_SROA_0_6_I8PTROFFSET_SROA_IDX1:%.*]] = getelementptr inbounds <4 x i16>, <4 x i16>* [[ALLOCA_SROA_0]], i64 0, i64 3
+; OVERALIGNED-NEXT:    [[ALLOCA_SROA_0_6_I8PTROFFSET_SROA_CAST2:%.*]] = bitcast i16* [[ALLOCA_SROA_0_6_I8PTROFFSET_SROA_IDX1]] to i8*
+; OVERALIGNED-NEXT:    [[ALLOCA_SROA_0_6_ALLOCA_SROA_0_6_RES:%.*]] = load i8, i8* [[ALLOCA_SROA_0_6_I8PTROFFSET_SROA_CAST2]], align 2
+; OVERALIGNED-NEXT:    ret i8 [[ALLOCA_SROA_0_6_ALLOCA_SROA_0_6_RES]]
 ;
 ; NATURAL-LABEL: @test_vector_bitcast_i8(
 ; NATURAL-NEXT:    [[ALLOCA_SROA_0:%.*]] = alloca <4 x i16>, align 8
 ; NATURAL-NEXT:    store <4 x i16> <i16 0, i16 1, i16 2, i16 3>, <4 x i16>* [[ALLOCA_SROA_0]], align 8
 ; NATURAL-NEXT:    [[ALLOCA_SROA_0_6_I8PTROFFSET_SROA_IDX1:%.*]] = getelementptr inbounds <4 x i16>, <4 x i16>* [[ALLOCA_SROA_0]], i64 0, i64 3
 ; NATURAL-NEXT:    [[ALLOCA_SROA_0_6_I8PTROFFSET_SROA_CAST2:%.*]] = bitcast i16* [[ALLOCA_SROA_0_6_I8PTROFFSET_SROA_IDX1]] to i8*
 ; NATURAL-NEXT:    [[ALLOCA_SROA_0_6_ALLOCA_SROA_0_6_RES:%.*]] = load i8, i8* [[ALLOCA_SROA_0_6_I8PTROFFSET_SROA_CAST2]], align 2
 ; NATURAL-NEXT:    ret i8 [[ALLOCA_SROA_0_6_ALLOCA_SROA_0_6_RES]]

llvm/unittests/IR/InstructionsTest.cpp

TEST(InstructionsTest, VectorGep) {
   delete BB0;
 
   delete ICmp0;
   delete ICmp1;
   delete PtrVecA;
   delete PtrVecB;
 }
 
+TEST(InstructionsTest, GepOffsets) {
+  // Test byte-based offsets of GEPs into vectors and arrays,
+  // including the case of overaligned element types.
+  LLVMContext C;
+  DataLayout DefaultDL(
+      "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f3"
+      "2:32:32-f64:64:64-v64:64:64-v128:128:128-a:0:64-s:64:64-f80"
+      ":128:128-n8:16:32:64-S128");
+  DataLayout AlignMin32DL(
+      "e-p:64:64:64-i1:8:8-i8:8:8-i16:32:32-i32:32:32-i64:64:64-f3"
+      "2:32:32-f64:64:64-v64:64:64-v128:128:128-a:0:64-s:64:64-f80"
+      ":128:128-n8:16:32:64-S128");
+  DataLayout AlignMin64DL(
+      "e-p:64:64:64-i1:8:8-i8:8:8-i16:64:64-i32:64:64-i64:64:64-f3"
+      "2:64:64-f64:64:64-v64:64:64-v128:128:128-a:0:64-s:64:64-f80"
+      ":128:128-n8:16:32:64-S128");
+
+  for (uint64_t ElemBitWidth : {8, 16, 24, 32, 64}) {
+    IntegerType *ElemTy = IntegerType::get(C, ElemBitWidth);
+    EXPECT_EQ(DefaultDL.getTypeSizeInBits(ElemTy), ElemBitWidth);
+
+    {
+      // Check GEP into vector
+      VectorType *VecTy = FixedVectorType::get(ElemTy, 8);
+      Constant *VectorNullPtr = Constant::getNullValue(VecTy->getPointerTo());
+      std::unique_ptr<llvm::GetElementPtrInst> Gep(GetElementPtrInst::Create(
+          VecTy, VectorNullPtr,
+          {ConstantInt::get(Type::getInt32Ty(C), 0),
+           ConstantInt::get(Type::getInt32Ty(C), 1)}));
+
+      for (const DataLayout *DL : {&DefaultDL, &AlignMin32DL, &AlignMin64DL}) {
+        EXPECT_EQ(DL->getTypeSizeInBits(VecTy), ElemBitWidth * 8);
+        APInt GEPOffset(DL->getIndexTypeSizeInBits(Gep->getType()), 0);
+        EXPECT_TRUE(Gep->accumulateConstantOffset(*DL, GEPOffset));
+        EXPECT_EQ(DL->getTypeSizeInBits(ElemTy), ElemBitWidth);
+        EXPECT_EQ(GEPOffset.getZExtValue(), ElemBitWidth / 8);
+      }
+    }
+    {
+      // Check GEP into array
+      ArrayType *ArrTy = ArrayType::get(ElemTy, 8);
+      Constant *ArrayNullPtr = Constant::getNullValue(ArrTy->getPointerTo());
+      std::unique_ptr<llvm::GetElementPtrInst> Gep(GetElementPtrInst::Create(
+          ArrTy, ArrayNullPtr,
+          {ConstantInt::get(Type::getInt32Ty(C), 0),
+           ConstantInt::get(Type::getInt32Ty(C), 1)}));
+
+      for (const DataLayout *DL : {&DefaultDL, &AlignMin32DL, &AlignMin64DL}) {
+        EXPECT_EQ(DL->getTypeSizeInBits(ArrTy),
+                  DL->getTypeAllocSizeInBits(ElemTy) * 8);
+        APInt GEPOffset(DL->getIndexTypeSizeInBits(Gep->getType()), 0);
+        EXPECT_TRUE(Gep->accumulateConstantOffset(*DL, GEPOffset));
+        EXPECT_GE(DL->getTypeAllocSizeInBits(ElemTy), ElemBitWidth);
+        EXPECT_EQ(GEPOffset.getZExtValue(), DL->getTypeAllocSize(ElemTy));
+      }
+    }
+  }
+}
+
 TEST(InstructionsTest, FPMathOperator) {
   LLVMContext Context;
   IRBuilder<> Builder(Context);
   MDBuilder MDHelper(Context);
   Instruction *I = Builder.CreatePHI(Builder.getDoubleTy(), 0);
   MDNode *MD1 = MDHelper.createFPMath(1.0);
   Value *V1 = Builder.CreateFAdd(I, I, "", MD1);
   EXPECT_TRUE(isa<FPMathOperator>(V1));