[mlir] Implement `DataLayoutTypeInterface` for `LLVMArrayType`

[zero9178]

Summary

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.

[zero9178]

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

71

This is basically equal to llvm::DataLayouts getTypeAllocSize. I am not sure whether MLIR wants to support over aligning types in which case we'd need an equivalent of it in DataLayout and co. Maybe worth a future TODO

[ftynse]

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

66

nit: kBitsInByte?

71

Ultimately, the LLVM dialect should support anything expressible in LLVM IR. The rest of MLIR is free to make any decisions.

[zero9178]

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

66

Addressed in final commit

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

 
 //===----------------------------------------------------------------------===//
 // LLVMArrayType.
 //===----------------------------------------------------------------------===//
 
 /// 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;
   using Base::getChecked;
 
   /// Checks if the given type can be used inside an array type.
   static bool isValidElementType(Type type);
 
   /// Gets or creates an instance of LLVM dialect array type containing
   /// `numElements` of `elementType`, in the same context as `elementType`.
   static LLVMArrayType get(Type elementType, unsigned numElements);
   static LLVMArrayType getChecked(function_ref<InFlightDiagnostic()> emitError,
                                   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;
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMFunctionType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect function type. It consists of a single return type (unlike MLIR
 /// which can have multiple), a list of parameter types and can optionally be

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

 LLVMArrayType
 LLVMArrayType::getChecked(function_ref<InFlightDiagnostic()> emitError,
                           Type elementType, unsigned numElements) {
   assert(elementType && "expected non-null subtype");
   return Base::getChecked(emitError, elementType.getContext(), elementType,
                           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,
                       Type elementType, unsigned numElements) {
   if (!isValidElementType(elementType))
     return emitError() << "invalid array element type: " << elementType;
   return success();
 }
 
+unsigned LLVMArrayType::getTypeSizeInBits(const DataLayout &dataLayout,
+                                          DataLayoutEntryListRef params) const {
+  return 8 * 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.
 //===----------------------------------------------------------------------===//
 
 bool LLVMFunctionType::isValidArgumentType(Type type) {
   return !type.isa<LLVMVoidType, LLVMFunctionType>();
 }
 

mlir/test/Dialect/LLVMIR/layout.mlir

module attributes { dlti.dl_spec = #dlti.dl_spec<
   #dlti.dl_entry<!llvm.struct<()>, dense<[64, 32]> : vector<2xi32>>
 >} {
   func @struct() {
     return
   }
 }
 
 // -----
+
+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
+    }
+}