[mlir] Implement DataLayoutTypeInterface for LLVMStructType

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 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.

Differential Revision: https://reviews.llvm.org/D115600

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

@@ -226,8 +226,9 @@ class LLVMPointerType : public Type::TypeBase<LLVMPointerType, Type,
 ///
 /// 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;
@@ -282,10 +283,10 @@ class LLVMStructType : public Type::TypeBase<LLVMStructType, Type,
   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();
@@ -297,13 +298,30 @@ class LLVMStructType : public Type::TypeBase<LLVMStructType, Type,
   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;
 };
 
 //===----------------------------------------------------------------------===//

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

@@ -361,15 +361,15 @@ LogicalResult LLVMStructType::setBody(ArrayRef<Type> types, bool isPacked) {
   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();
 }
@@ -389,6 +389,147 @@ LLVMStructType::verify(function_ref<InFlightDiagnostic()> emitError,
   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.
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/LLVMIR/layout.mlir

@@ -111,3 +111,136 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     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
+  }
+}
+
+// -----