[mlir] Implement `DataLayoutTypeInterface` for `LLVMStructType`

[zero9178]

Summary

Using this implementation of the interface it is possible to query the size, ABI alignment as well as the preferred alignment of a struct. It should yield the same results as LLVMs llvm::DataLayout on an equivalent llvm::StructType, including for packed structs.

Additionally it is also possible to to increase the ABI and preferred alignment using a data layout entry with the type llvm.struct<()>, which serves the same functionality as the a: component in LLVMs data layout string.

Depends on https://reviews.llvm.org/D115596

[ftynse]

Nice, thanks!

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

 /// caller does not intend to ever set the body (e.g. forward-declarations of
 /// structs from another module) and wants to disallow further modification of
 /// the body. For intentionally-opaque structs or non-opaque structs with the
 /// body, one is not allowed to set another body (however, one can set exactly
 /// the same body).
 ///
 /// Note that the packedness of the struct takes place in uniquing of literal
 /// structs, but does not in uniquing of identified structs.
-class LLVMStructType : public Type::TypeBase<LLVMStructType, Type,
-                                             detail::LLVMStructTypeStorage> {
+class LLVMStructType
+    : public Type::TypeBase<LLVMStructType, Type, detail::LLVMStructTypeStorage,
+                            DataLayoutTypeInterface::Trait> {
 public:
   /// Inherit base constructors.
   using Base::Base;
 
   /// Checks if the given type can be contained in a structure type.
   static bool isValidElementType(Type type);
 
   /// Gets or creates an identified struct with the given name in the provided
public:
   /// not be set, e.g. if the struct already has a body or if it was marked as
   /// intentionally opaque. This might happen in a multi-threaded context when a
   /// different thread modified the struct after it was created. Most callers
   /// are likely to assert this always succeeds, but it is possible to implement
   /// a local renaming scheme based on the result of this call.
   LogicalResult setBody(ArrayRef<Type> types, bool isPacked);
 
   /// Checks if a struct is packed.
-  bool isPacked();
+  bool isPacked() const;
 
   /// Checks if a struct is identified.
-  bool isIdentified();
+  bool isIdentified() const;
 
   /// Checks if a struct is opaque.
   bool isOpaque();
 
   /// Checks if a struct is initialized.
   bool isInitialized();
 
   /// Returns the name of an identified struct.
   StringRef getName();
 
   /// Returns the list of element types contained in a non-opaque struct.
-  ArrayRef<Type> getBody();
+  ArrayRef<Type> getBody() const;
 
   /// Verifies that the type about to be constructed is well-formed.
   static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
                               StringRef, bool);
   static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
                               ArrayRef<Type> types, bool);
+
+  /// Hooks for DataLayoutTypeInterface. Should not be called directly. Obtain a
+  /// DataLayout instance and query it instead.
+  unsigned getTypeSizeInBits(const DataLayout &dataLayout,
+                             DataLayoutEntryListRef params) const;
+
+  unsigned getABIAlignment(const DataLayout &dataLayout,
+                           DataLayoutEntryListRef params) const;
+
+  unsigned getPreferredAlignment(const DataLayout &dataLayout,
+                                 DataLayoutEntryListRef params) const;
+
+  bool areCompatible(DataLayoutEntryListRef oldLayout,
+                     DataLayoutEntryListRef newLayout) const;
+
+  LogicalResult verifyEntries(DataLayoutEntryListRef entries,
+                              Location loc) const;
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMVectorType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect vector type, represents a sequence of elements that can be
 /// processed as one, typically in SIMD context. This is a base class for fixed

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

 
 LogicalResult LLVMStructType::setBody(ArrayRef<Type> types, bool isPacked) {
   assert(isIdentified() && "can only set bodies of identified structs");
   assert(llvm::all_of(types, LLVMStructType::isValidElementType) &&
          "expected valid body types");
   return Base::mutate(types, isPacked);
 }
 
-bool LLVMStructType::isPacked() { return getImpl()->isPacked(); }
-bool LLVMStructType::isIdentified() { return getImpl()->isIdentified(); }
+bool LLVMStructType::isPacked() const { return getImpl()->isPacked(); }
+bool LLVMStructType::isIdentified() const { return getImpl()->isIdentified(); }
 bool LLVMStructType::isOpaque() {
   return getImpl()->isIdentified() &&
          (getImpl()->isOpaque() || !getImpl()->isInitialized());
 }
 bool LLVMStructType::isInitialized() { return getImpl()->isInitialized(); }
 StringRef LLVMStructType::getName() { return getImpl()->getIdentifier(); }
-ArrayRef<Type> LLVMStructType::getBody() {
+ArrayRef<Type> LLVMStructType::getBody() const {
   return isIdentified() ? getImpl()->getIdentifiedStructBody()
                         : getImpl()->getTypeList();
 }
 
 LogicalResult LLVMStructType::verify(function_ref<InFlightDiagnostic()>,
                                      StringRef, bool) {
   return success();
 }
 
 LogicalResult
 LLVMStructType::verify(function_ref<InFlightDiagnostic()> emitError,
                        ArrayRef<Type> types, bool) {
   for (Type t : types)
     if (!isValidElementType(t))
       return emitError() << "invalid LLVM structure element type: " << t;
 
   return success();
 }
 
+unsigned
+LLVMStructType::getTypeSizeInBits(const DataLayout &dataLayout,
+                                  DataLayoutEntryListRef params) const {
+  unsigned structSize = 0;
+  unsigned structAlignment = 1;
+  for (Type element : getBody()) {
+    unsigned elementAlignment =
+        isPacked() ? 1 : dataLayout.getTypeABIAlignment(element);
+    // Add padding to the struct size to align it to the abi alignment of the
+    // element type before than adding the size of the element
+    structSize = llvm::alignTo(structSize, elementAlignment);
+    structSize += dataLayout.getTypeSize(element);
+
+    // The alignment requirement of a struct is equal to the strictest alignment
+    // requirement of its elements.
+    structAlignment = std::max(elementAlignment, structAlignment);
+  }
+  // At the end, add padding to the struct to satisfy its own alignment
+  // requirement. Otherwise structs inside of arrays would be misaligned.
+  structSize = llvm::alignTo(structSize, structAlignment);
+  return structSize * kBitsInByte;
+}
+
+namespace {
+enum class StructDLEntryPos { Abi = 0, Preferred = 1 };
+}
+
+static Optional<unsigned>
+getStructDataLayoutEntry(DataLayoutEntryListRef params, LLVMStructType type,
+                         StructDLEntryPos pos) {
+  auto currentEntry = llvm::find_if(params, [](DataLayoutEntryInterface entry) {
+    return entry.isTypeEntry();
+  });
+  if (currentEntry == params.end())
+    return llvm::None;
+
+  auto attr = currentEntry->getValue().cast<DenseIntElementsAttr>();
+  if (pos == StructDLEntryPos::Preferred &&
+      attr.size() <= static_cast<unsigned>(StructDLEntryPos::Preferred))
+    // If no preferred was specified, fall back to abi alignment
+    pos = StructDLEntryPos::Abi;
+
+  return attr.getValues<unsigned>()[static_cast<unsigned>(pos)];
+}
+
+static unsigned calculateStructAlignment(const DataLayout &dataLayout,
+                                         DataLayoutEntryListRef params,
+                                         LLVMStructType type,
+                                         StructDLEntryPos pos) {
+  // Packed structs always have an abi alignment of 1
+  if (pos == StructDLEntryPos::Abi && type.isPacked()) {
+    return 1;
+  }
+
+  // The alignment requirement of a struct is equal to the strictest alignment
+  // requirement of its elements.
+  unsigned structAlignment = 1;
+  for (Type iter : type.getBody()) {
+    structAlignment =
+        std::max(dataLayout.getTypeABIAlignment(iter), structAlignment);
+  }
+
+  // Entries are only allowed to be stricter than the required alignment
+  if (Optional<unsigned> entryResult =
+          getStructDataLayoutEntry(params, type, pos))
+    return std::max(*entryResult / kBitsInByte, structAlignment);
+
+  return structAlignment;
+}
+
+unsigned LLVMStructType::getABIAlignment(const DataLayout &dataLayout,
+                                         DataLayoutEntryListRef params) const {
+  return calculateStructAlignment(dataLayout, params, *this,
+                                  StructDLEntryPos::Abi);
+}
+
+unsigned
+LLVMStructType::getPreferredAlignment(const DataLayout &dataLayout,
+                                      DataLayoutEntryListRef params) const {
+  return calculateStructAlignment(dataLayout, params, *this,
+                                  StructDLEntryPos::Preferred);
+}
+
+static unsigned extractStructSpecValue(Attribute attr, StructDLEntryPos pos) {
+  return attr.cast<DenseIntElementsAttr>()
+      .getValues<unsigned>()[static_cast<unsigned>(pos)];
+}
+
+bool LLVMStructType::areCompatible(DataLayoutEntryListRef oldLayout,
+                                   DataLayoutEntryListRef newLayout) const {
+  for (DataLayoutEntryInterface newEntry : newLayout) {
+    if (!newEntry.isTypeEntry())
+      continue;
+
+    auto previousEntry =
+        llvm::find_if(oldLayout, [](DataLayoutEntryInterface entry) {
+          return entry.isTypeEntry();
+        });
+    if (previousEntry == oldLayout.end())
+      continue;
+
+    unsigned abi = extractStructSpecValue(previousEntry->getValue(),
+                                          StructDLEntryPos::Abi);
+    unsigned newAbi =
+        extractStructSpecValue(newEntry.getValue(), StructDLEntryPos::Abi);
+    if (abi < newAbi || abi % newAbi != 0)
+      return false;
+  }
+  return true;
+}
+
+LogicalResult LLVMStructType::verifyEntries(DataLayoutEntryListRef entries,
+                                            Location loc) const {
+  for (DataLayoutEntryInterface entry : entries) {
+    if (!entry.isTypeEntry())
+      continue;
+
+    auto key = entry.getKey().get<Type>().cast<LLVMStructType>();
+    auto values = entry.getValue().dyn_cast<DenseIntElementsAttr>();
+    if (!values || (values.size() != 2 && values.size() != 1)) {
+      return emitError(loc)
+             << "expected layout attribute for " << entry.getKey().get<Type>()
+             << " to be a dense integer elements attribute of 1 or 2 elements";
+    }
+
+    if (key.isIdentified() || !key.getBody().empty()) {
+      return emitError(loc) << "unexpected layout attribute for struct " << key;
+    }
+
+    if (values.size() == 1)
+      continue;
+
+    if (extractStructSpecValue(values, StructDLEntryPos::Abi) >
+        extractStructSpecValue(values, StructDLEntryPos::Preferred)) {
+      return emitError(loc) << "preferred alignment is expected to be at least "
+                               "as large as ABI alignment";
+    }
+  }
+  return mlir::success();
+}
+
 //===----------------------------------------------------------------------===//
 // Vector types.
 //===----------------------------------------------------------------------===//
 
 /// Verifies that the type about to be constructed is well-formed.
 template <typename VecTy>
 static LogicalResult
 verifyVectorConstructionInvariants(function_ref<InFlightDiagnostic()> emitError,

mlir/test/Dialect/LLVMIR/layout.mlir

 // expected-error@below {{preferred alignment is expected to be at least as large as ABI alignment}}
 module attributes { dlti.dl_spec = #dlti.dl_spec<
   #dlti.dl_entry<!llvm.ptr<i8>, dense<[64, 64, 32]> : vector<3xi32>>
 >} {
   func @pointer() {
     return
   }
 }
+
+// -----
+
+module {
+    // CHECK: @no_spec
+    func @no_spec() {
+        // simple case
+        // CHECK: alignment = 4
+        // CHECK: bitsize = 32
+        // CHECK: preferred = 4
+        // CHECK: size = 4
+        "test.data_layout_query"() : () -> !llvm.struct<(i32)>
+
+        // padding inbetween
+        // CHECK: alignment = 8
+        // CHECK: bitsize = 128
+        // CHECK: preferred = 8
+        // CHECK: size = 16
+        "test.data_layout_query"() : () -> !llvm.struct<(i32, f64)>
+
+        // padding at end of struct
+        // CHECK: alignment = 8
+        // CHECK: bitsize = 128
+        // CHECK: preferred = 8
+        // CHECK: size = 16
+        "test.data_layout_query"() : () -> !llvm.struct<(f64, i32)>
+
+         // packed
+         // CHECK: alignment = 1
+         // CHECK: bitsize = 96
+         // CHECK: preferred = 8
+         // CHECK: size = 12
+         "test.data_layout_query"() : () -> !llvm.struct<packed (f64, i32)>
+
+         // empty
+         // CHECK: alignment = 1
+         // CHECK: bitsize = 0
+         // CHECK: preferred = 1
+         // CHECK: size = 0
+         "test.data_layout_query"() : () -> !llvm.struct<()>
+         return
+    }
+}
+
+// -----
+
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!llvm.struct<()>, dense<[32, 32]> : vector<2xi32>>
+>} {
+    // CHECK: @spec
+    func @spec() {
+        // Strict alignment is applied
+        // CHECK: alignment = 4
+        // CHECK: bitsize = 16
+        // CHECK: preferred = 4
+        // CHECK: size = 2
+        "test.data_layout_query"() : () -> !llvm.struct<(i16)>
+
+        // No impact on structs that have stricter requirements
+        // CHECK: alignment = 8
+        // CHECK: bitsize = 128
+        // CHECK: preferred = 8
+        // CHECK: size = 16
+        "test.data_layout_query"() : () -> !llvm.struct<(i32, f64)>
+
+         // Only the preferred alignment of structs is affected
+         // CHECK: alignment = 1
+         // CHECK: bitsize = 32
+         // CHECK: preferred = 4
+         // CHECK: size = 4
+         "test.data_layout_query"() : () -> !llvm.struct<packed (i16, i16)>
+
+         // empty
+         // CHECK: alignment = 4
+         // CHECK: bitsize = 0
+         // CHECK: preferred = 4
+         // CHECK: size = 0
+         "test.data_layout_query"() : () -> !llvm.struct<()>
+         return
+    }
+}
+
+// -----
+
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!llvm.struct<()>, dense<[32]> : vector<1xi32>>
+>} {
+    // CHECK: @spec_without_preferred
+    func @spec_without_preferred() {
+        // abi alignment is applied to both preferred and abi
+        // CHECK: alignment = 4
+        // CHECK: bitsize = 16
+        // CHECK: preferred = 4
+        // CHECK: size = 2
+        "test.data_layout_query"() : () -> !llvm.struct<(i16)>
+        return
+    }
+}
+
+// -----
+
+// expected-error@below {{unexpected layout attribute for struct '!llvm.struct<(i8)>'}}
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!llvm.struct<(i8)>, dense<[64, 64]> : vector<2xi32>>
+>} {
+  func @struct() {
+    return
+  }
+}
+
+// -----
+
+// expected-error@below {{expected layout attribute for '!llvm.struct<()>' to be a dense integer elements attribute of 1 or 2 elements}}
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!llvm.struct<()>, dense<[64, 64, 64]> : vector<3xi32>>
+>} {
+  func @struct() {
+    return
+  }
+}
+
+// -----
+
+// expected-error@below {{preferred alignment is expected to be at least as large as ABI alignment}}
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!llvm.struct<()>, dense<[64, 32]> : vector<2xi32>>
+>} {
+  func @struct() {
+    return
+  }
+}
+
+// -----