[IR]: Add generic_gep_type_iterator::getSequentialElementStride

This prepares a fix to GEP offset computations on vectors
of overaligned elements.

We have many places that analyze GEP offsets using GEP iterators
with the following pattern:

  GTI = gep_type_begin(ElemTy, Indices),
  GTE = gep_type_end(ElemTy, Indices);
  for (; GTI != GTE; ++GTI) {
    if (StructType *STy = GTI.getStructTypeOrNull()) {
       // handle struct
       [..]
    } else {
      // handle sequential (outmost index, array, vector):
      auto Stride = DL.getTypeAllocSize(GTI.getIndexedType());
      Offset += Index * Size;
    }
  }

This is incorrect for vectors of types whose bit size does not
equal its alloc size (e.g. overaligned types), as vectors
always bit-pack their elements.

This patch introduces new functions generic_gep_type_iterator::isVector()
and generic_gep_type_iterator::getSequentialElementStride(const DataLayout &)
to fix these patterns without having to teach all these places about
the specifics of vector bit layouts. With these helpers, the pattern
above can be fixed by replacing the stride computation:

      auto Stride = GTI.getSequentialElementStride(DL);

Author: jasilvanus

llvm/include/llvm/IR/GetElementPtrTypeIterator.h

@@ -16,6 +16,7 @@
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/PointerUnion.h"
+#include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/User.h"
@@ -30,7 +31,39 @@ template <typename ItTy = User::const_op_iterator>
 class generic_gep_type_iterator {
 
   ItTy OpIt;
-  PointerUnion<StructType *, Type *> CurTy;
+  // We use two different mechanisms to store the type a GEP index applies to.
+  // In some cases, we need to know the outer aggregate type the index is
+  // applied within, e.g. a struct. In such cases, we store the aggregate type
+  // in the iterator, and derive the element type on the fly.
+  //
+  // However, this is not always possible, because for the outermost index there
+  // is no containing type. In such cases, or if the containing type is not
+  // relevant, e.g. for arrays, the element type is stored as Type* in CurTy.
+  //
+  // If CurTy contains a Type* value, this does not imply anything about the
+  // type itself, because it is the element type and not the outer type.
+  // In particular, Type* can be a struct type.
+  //
+  // Consider this example:
+  //
+  //    %my.struct = type { i32, [ 4 x float ] }
+  //    [...]
+  //    %gep = getelementptr %my.struct, ptr %ptr, i32 10, i32 1, 32 3
+  //
+  // Iterating over the indices of this GEP, CurTy will contain the following
+  // values:
+  //    * i32 10: The outer index always operates on the GEP value type.
+  //              CurTy contains a Type*       pointing at `%my.struct`.
+  //    * i32 1:  This index is within a struct.
+  //              CurTy contains a StructType* pointing at `%my.struct`.
+  //    * i32 3:  This index is within an array. We reuse the "flat" indexing
+  //              for arrays which is also used in the top level GEP index.
+  //              CurTy contains a Type*       pointing at `float`.
+  //
+  // Vectors are handled separately because the layout of vectors is different
+  // for overaligned elements: Vectors are always bit-packed, whereas arrays
+  // respect ABI alignment of the elements.
+  PointerUnion<StructType *, VectorType *, Type *> CurTy;
 
   generic_gep_type_iterator() = default;
 
@@ -69,6 +102,8 @@ class generic_gep_type_iterator {
   Type *getIndexedType() const {
     if (auto *T = dyn_cast_if_present<Type *>(CurTy))
       return T;
+    if (auto *VT = dyn_cast_if_present<VectorType *>(CurTy))
+      return VT->getElementType();
     return cast<StructType *>(CurTy)->getTypeAtIndex(getOperand());
   }
 
@@ -79,7 +114,7 @@ class generic_gep_type_iterator {
     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;
@@ -108,7 +143,23 @@ class generic_gep_type_iterator {
   // that.
 
   bool isStruct() const { return isa<StructType *>(CurTy); }
-  bool isSequential() const { return isa<Type *>(CurTy); }
+  bool isVector() const { return isa<VectorType *>(CurTy); }
+  bool isSequential() const { return !isStruct(); }
+
+  // For sequential GEP indices (all except those into structs), the index value
+  // can be translated into a byte offset by multiplying with an element stride.
+  // This function returns this stride, which both depends on the element type,
+  // and the containing aggregate type, as vectors always tightly bit-pack their
+  // elements.
+  TypeSize getSequentialElementStride(const DataLayout &DL) const {
+    assert(isSequential());
+    Type *ElemTy = getIndexedType();
+    TypeSize ElemSizeInBits = isVector() ? DL.getTypeSizeInBits(ElemTy)
+                                         : DL.getTypeAllocSizeInBits(ElemTy);
+    // Check for invalid GEPs that are not byte-addressable.
+    assert(ElemSizeInBits.isKnownMultipleOf(8));
+    return ElemSizeInBits.divideCoefficientBy(8);
+  }
 
   StructType *getStructType() const { return cast<StructType *>(CurTy); }