[mlir] Implement `DataLayoutTypeInterface` for `LLVMArrayType`

Implementation of the interface allows querying the size and alignments of an LLVMArrayType as well as query the size and alignment of a struct containing an LLVMArrayType.
The implementation should yield the same results as llvm::DataLayout, including support for over aligned element types.
There is no customization point for adjusting an arrays alignment; it is simply taken from the element type.

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

3 files changed, 92 insertions, 7 deletions

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h b/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h
index 399542317953..3d6a6fdd7078 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h
@@ -71,8 +71,9 @@ DEFINE_TRIVIAL_LLVM_TYPE(LLVMMetadataType);
 /// LLVM dialect array type. It is an aggregate type representing consecutive
 /// elements in memory, parameterized by the number of elements and the element
 /// type.
-class LLVMArrayType : public Type::TypeBase<LLVMArrayType, Type,
-                                            detail::LLVMTypeAndSizeStorage> {
+class LLVMArrayType
+    : public Type::TypeBase<LLVMArrayType, Type, detail::LLVMTypeAndSizeStorage,
+                            DataLayoutTypeInterface::Trait> {
 public:
   /// Inherit base constructors.
   using Base::Base;
@@ -88,14 +89,28 @@ public:
                                   Type elementType, unsigned numElements);
 
   /// Returns the element type of the array.
-  Type getElementType();
+  Type getElementType() const;
 
   /// Returns the number of elements in the array type.
-  unsigned getNumElements();
+  unsigned getNumElements() const;
 
   /// Verifies that the type about to be constructed is well-formed.
   static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
                               Type elementType, unsigned numElements);
+
+  /// 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 getTypeSize(const DataLayout &dataLayout,
+                       DataLayoutEntryListRef params) const;
+
+  unsigned getABIAlignment(const DataLayout &dataLayout,
+                           DataLayoutEntryListRef params) const;
+
+  unsigned getPreferredAlignment(const DataLayout &dataLayout,
+                                 DataLayoutEntryListRef params) const;
 };
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
index 6751825fdfe6..87dc3b8d4f10 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
@@ -25,6 +25,8 @@
 using namespace mlir;
 using namespace mlir::LLVM;
 
+constexpr const static unsigned kBitsInByte = 8;
+
 //===----------------------------------------------------------------------===//
 // Array type.
 //===----------------------------------------------------------------------===//
@@ -47,9 +49,11 @@ LLVMArrayType::getChecked(function_ref<InFlightDiagnostic()> emitError,
                           numElements);
 }
 
-Type LLVMArrayType::getElementType() { return getImpl()->elementType; }
+Type LLVMArrayType::getElementType() const { return getImpl()->elementType; }
 
-unsigned LLVMArrayType::getNumElements() { return getImpl()->numElements; }
+unsigned LLVMArrayType::getNumElements() const {
+  return getImpl()->numElements;
+}
 
 LogicalResult
 LLVMArrayType::verify(function_ref<InFlightDiagnostic()> emitError,
@@ -59,6 +63,29 @@ LLVMArrayType::verify(function_ref<InFlightDiagnostic()> emitError,
   return success();
 }
 
+unsigned LLVMArrayType::getTypeSizeInBits(const DataLayout &dataLayout,
+                                          DataLayoutEntryListRef params) const {
+  return kBitsInByte * getTypeSize(dataLayout, params);
+}
+
+unsigned LLVMArrayType::getTypeSize(const DataLayout &dataLayout,
+                                    DataLayoutEntryListRef params) const {
+  return llvm::alignTo(dataLayout.getTypeSize(getElementType()),
+                       dataLayout.getTypeABIAlignment(getElementType())) *
+         getNumElements();
+}
+
+unsigned LLVMArrayType::getABIAlignment(const DataLayout &dataLayout,
+                                        DataLayoutEntryListRef params) const {
+  return dataLayout.getTypeABIAlignment(getElementType());
+}
+
+unsigned
+LLVMArrayType::getPreferredAlignment(const DataLayout &dataLayout,
+                                     DataLayoutEntryListRef params) const {
+  return dataLayout.getTypePreferredAlignment(getElementType());
+}
+
 //===----------------------------------------------------------------------===//
 // Function type.
 //===----------------------------------------------------------------------===//
@@ -159,7 +186,6 @@ enum class DLEntryPos { Size = 0, Abi = 1, Preferred = 2, Address = 3 };
 
 constexpr const static unsigned kDefaultPointerSizeBits = 64;
 constexpr const static unsigned kDefaultPointerAlignment = 8;
-constexpr const static unsigned kBitsInByte = 8;
 
 /// Returns the value that corresponds to named position `pos` from the
 /// attribute `attr` assuming it's a dense integer elements attribute.
diff --git a/mlir/test/Dialect/LLVMIR/layout.mlir b/mlir/test/Dialect/LLVMIR/layout.mlir
index a64724dfaa71..dc6d5a91376a 100644
--- a/mlir/test/Dialect/LLVMIR/layout.mlir
+++ b/mlir/test/Dialect/LLVMIR/layout.mlir
@@ -244,3 +244,47 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
 }
 
 // -----
+
+module {
+    // CHECK: @arrays
+    func @arrays() {
+        // simple case
+        // CHECK: alignment = 4
+        // CHECK: bitsize = 64
+        // CHECK: preferred = 4
+        // CHECK: size = 8
+        "test.data_layout_query"() : () -> !llvm.array<2 x i32>
+
+        // size 0
+        // CHECK: alignment = 8
+        // CHECK: bitsize = 0
+        // CHECK: preferred = 8
+        // CHECK: size = 0
+        "test.data_layout_query"() : () -> !llvm.array<0 x f64>
+
+        // alignment info matches element type
+        // CHECK: alignment = 4
+        // CHECK: bitsize = 64
+        // CHECK: preferred = 8
+        // CHECK: size = 8
+        "test.data_layout_query"() : () -> !llvm.array<1 x i64>
+        return
+    }
+}
+
+// -----
+
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!llvm.struct<()>, dense<[64]> : vector<1xi32>>
+>} {
+    // CHECK: @overaligned
+    func @overaligned() {
+        // Over aligned element types are respected
+        // CHECK: alignment = 8
+        // CHECK: bitsize = 128
+        // CHECK: preferred = 8
+        // CHECK: size = 16
+        "test.data_layout_query"() : () -> !llvm.array<2 x struct<(i8)>>
+         return
+    }
+}