[mlir] make LLVMPointerType implement the data layout type interface

[ftynse]

Summary

This brings us closer to replacing the LLVM data layout string with a
first-class layout modeling in MLIR.

Depends On D103945

[ftynse]

Add forgotten verify-diagnostics to the test.

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

 // correspond to the LLVM IR type system.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LLVMIR_LLVMTYPES_H_
 #define MLIR_DIALECT_LLVMIR_LLVMTYPES_H_
 
 #include "mlir/IR/Types.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
 
 namespace llvm {
 class ElementCount;
 class TypeSize;
 } // namespace llvm
 
 namespace mlir {
 

 //===----------------------------------------------------------------------===//
 // LLVMPointerType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect pointer type. This type typically represents a reference to an
 /// object in memory. It is parameterized by the element type and the address
 /// space.
 class LLVMPointerType : public Type::TypeBase<LLVMPointerType, Type,
-                                              detail::LLVMPointerTypeStorage> {
+                                              detail::LLVMPointerTypeStorage,
+                                              DataLayoutTypeInterface::Trait> {
 public:
   /// Inherit base constructors.
   using Base::Base;
   using Base::getChecked;
 
   /// Checks if the given type can have a pointer type pointing to it.
   static bool isValidElementType(Type type);
 
   /// Gets or creates an instance of LLVM dialect pointer type pointing to an
   /// object of `pointee` type in the given address space. The pointer type is
   /// created in the same context as `pointee`.
   static LLVMPointerType get(Type pointee, unsigned addressSpace = 0);
   static LLVMPointerType
   getChecked(function_ref<InFlightDiagnostic()> emitError, Type pointee,
              unsigned addressSpace = 0);
 
   /// Returns the pointed-to type.
-  Type getElementType();
+  Type getElementType() const;
 
   /// Returns the address space of the pointer.
-  unsigned getAddressSpace();
+  unsigned getAddressSpace() const;
 
   /// Verifies that the type about to be constructed is well-formed.
   static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
                               Type pointee, unsigned);
+
+  /// 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;
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMStructType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect structure type representing a collection of different-typed
 /// elements manipulated together. Structured can optionally be packed, meaning

mlir/lib/Dialect/LLVMIR/CMakeLists.txt

add_mlir_dialect_library(MLIRLLVMIR
   AsmParser
   BitReader
   BitWriter
   Core
 
   LINK_LIBS PUBLIC
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
+  MLIRDataLayoutInterfaces
   MLIRIR
   MLIRSideEffectInterfaces
   MLIRSupport
   )
 
 add_mlir_dialect_library(MLIRNVVMIR
   IR/NVVMDialect.cpp
 

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

 
 LLVMPointerType
 LLVMPointerType::getChecked(function_ref<InFlightDiagnostic()> emitError,
                             Type pointee, unsigned addressSpace) {
   return Base::getChecked(emitError, pointee.getContext(), pointee,
                           addressSpace);
 }
 
-Type LLVMPointerType::getElementType() { return getImpl()->pointeeType; }
+Type LLVMPointerType::getElementType() const { return getImpl()->pointeeType; }
 
-unsigned LLVMPointerType::getAddressSpace() { return getImpl()->addressSpace; }
+unsigned LLVMPointerType::getAddressSpace() const {
+  return getImpl()->addressSpace;
+}
 
 LogicalResult
 LLVMPointerType::verify(function_ref<InFlightDiagnostic()> emitError,
                         Type pointee, unsigned) {
   if (!isValidElementType(pointee))
     return emitError() << "invalid pointer element type: " << pointee;
   return success();
 }
 
+namespace {
+/// The positions of different values in the data layout entry.
+enum class DLEntryPos { Size = 0, Abi = 1, Preferred = 2, Address = 3 };
+} // namespace
+
+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.
+static unsigned extractPointerSpecValue(Attribute attr, DLEntryPos pos) {
+  return attr.cast<DenseIntElementsAttr>().getValue<unsigned>(
+      static_cast<unsigned>(pos));
+}
+
+/// Returns the part of the data layout entry that corresponds to `pos` for the
+/// given `type` by interpreting the list of entries `params`. For the pointer
+/// type in the default address space, returns the default value if the entries
+/// do not provide a custom one, for other address spaces returns None.
+static Optional<unsigned>
+getPointerDataLayoutEntry(DataLayoutEntryListRef params, LLVMPointerType type,
+                          DLEntryPos pos) {
+  // First, look for the entry for the pointer in the current address space.
+  Attribute currentEntry;
+  for (DataLayoutEntryInterface entry : params) {
+    if (!entry.isTypeEntry())
+      continue;
+    if (entry.getKey().get<Type>().cast<LLVMPointerType>().getAddressSpace() ==
+        type.getAddressSpace()) {
+      currentEntry = entry.getValue();
+      break;
+    }
+  }
+  if (currentEntry) {
+    return extractPointerSpecValue(currentEntry, pos) /
+           (pos == DLEntryPos::Size ? 1 : kBitsInByte);
+  }
+
+  // If not found, and this is the pointer to the default memory space, assume
+  // 64-bit pointers.
+  if (type.getAddressSpace() == 0) {
+    return pos == DLEntryPos::Size ? kDefaultPointerSizeBits
+                                   : kDefaultPointerAlignment;
+  }
+
+  return llvm::None;
+}
+
+unsigned
+LLVMPointerType::getTypeSizeInBits(const DataLayout &dataLayout,
+                                   DataLayoutEntryListRef params) const {
+  if (Optional<unsigned> size =
+          getPointerDataLayoutEntry(params, *this, DLEntryPos::Size))
+    return *size;
+
+  // For other memory spaces, use the size of the pointer to the default memory
+  // space.
+  return dataLayout.getTypeSizeInBits(get(getElementType()));
+}
+
+unsigned LLVMPointerType::getABIAlignment(const DataLayout &dataLayout,
+                                          DataLayoutEntryListRef params) const {
+  if (Optional<unsigned> alignment =
+          getPointerDataLayoutEntry(params, *this, DLEntryPos::Abi))
+    return *alignment;
+
+  return dataLayout.getTypeABIAlignment(get(getElementType()));
+}
+
+unsigned
+LLVMPointerType::getPreferredAlignment(const DataLayout &dataLayout,
+                                       DataLayoutEntryListRef params) const {
+  if (Optional<unsigned> alignment =
+          getPointerDataLayoutEntry(params, *this, DLEntryPos::Preferred))
+    return *alignment;
+
+  return dataLayout.getTypePreferredAlignment(get(getElementType()));
+}
+
+bool LLVMPointerType::areCompatible(DataLayoutEntryListRef oldLayout,
+                                    DataLayoutEntryListRef newLayout) const {
+  for (DataLayoutEntryInterface newEntry : newLayout) {
+    if (!newEntry.isTypeEntry())
+      continue;
+    unsigned size = kDefaultPointerSizeBits;
+    unsigned abi = kDefaultPointerAlignment;
+    auto newType = newEntry.getKey().get<Type>().cast<LLVMPointerType>();
+    auto it = llvm::find_if(oldLayout, [&](DataLayoutEntryInterface entry) {
+      if (auto type = entry.getKey().dyn_cast<Type>()) {
+        return type.cast<LLVMPointerType>().getAddressSpace() ==
+               newType.getAddressSpace();
+      }
+      return false;
+    });
+    if (it == oldLayout.end()) {
+      llvm::find_if(oldLayout, [&](DataLayoutEntryInterface entry) {
+        if (auto type = entry.getKey().dyn_cast<Type>()) {
+          return type.cast<LLVMPointerType>().getAddressSpace() == 0;
+        }
+        return false;
+      });
+    }
+    if (it != oldLayout.end()) {
+      size = extractPointerSpecValue(*it, DLEntryPos::Size);
+      abi = extractPointerSpecValue(*it, DLEntryPos::Abi);
+    }
+
+    Attribute newSpec = newEntry.getValue().cast<DenseIntElementsAttr>();
+    unsigned newSize = extractPointerSpecValue(newSpec, DLEntryPos::Size);
+    unsigned newAbi = extractPointerSpecValue(newSpec, DLEntryPos::Abi);
+    if (size != newSize || abi < newAbi || abi % newAbi != 0)
+      return false;
+  }
+  return true;
+}
+
+LogicalResult LLVMPointerType::verifyEntries(DataLayoutEntryListRef entries,
+                                             Location loc) const {
+  for (DataLayoutEntryInterface entry : entries) {
+    if (!entry.isTypeEntry())
+      continue;
+    auto key = entry.getKey().get<Type>().cast<LLVMPointerType>();
+    auto values = entry.getValue().dyn_cast<DenseIntElementsAttr>();
+    if (!values || (values.size() != 3 && values.size() != 4)) {
+      return emitError(loc)
+             << "expected layout attribute for " << entry.getKey().get<Type>()
+             << " to be a dense integer elements attribute with 3 or 4 "
+                "elements";
+    }
+    if (!key.getElementType().isInteger(8)) {
+      return emitError(loc) << "unexpected layout attribute for pointer to "
+                            << key.getElementType();
+    }
+    if (extractPointerSpecValue(values, DLEntryPos::Abi) >
+        extractPointerSpecValue(values, DLEntryPos::Preferred)) {
+      return emitError(loc) << "preferred alignment is expected to be at least "
+                               "as large as ABI alignment";
+    }
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Struct type.
 //===----------------------------------------------------------------------===//
 
 bool LLVMStructType::isValidElementType(Type type) {
   return !type.isa<LLVMVoidType, LLVMLabelType, LLVMMetadataType,
                    LLVMFunctionType, LLVMTokenType, LLVMScalableVectorType>();
 }

mlir/test/Dialect/LLVMIR/layout.mlir

+// RUN: mlir-opt --test-data-layout-query --split-input-file --verify-diagnostics %s | FileCheck %s
+
+module {
+  // CHECK: @no_spec
+  func @no_spec() {
+    // CHECK: alignment = 8
+    // CHECK: bitsize = 64
+    // CHECK: preferred = 8
+    // CHECK: size = 8
+    "test.data_layout_query"() : () -> !llvm.ptr<i8>
+    // CHECK: alignment = 8
+    // CHECK: bitsize = 64
+    // CHECK: preferred = 8
+    // CHECK: size = 8
+    "test.data_layout_query"() : () -> !llvm.ptr<i32>
+    // CHECK: alignment = 8
+    // CHECK: bitsize = 64
+    // CHECK: preferred = 8
+    // CHECK: size = 8
+    "test.data_layout_query"() : () -> !llvm.ptr<bf16>
+    // CHECK: alignment = 8
+    // CHECK: bitsize = 64
+    // CHECK: preferred = 8
+    // CHECK: size = 8
+    "test.data_layout_query"() : () -> !llvm.ptr<!llvm.ptr<i8>>
+    // CHECK: alignment = 8
+    // CHECK: bitsize = 64
+    // CHECK: preferred = 8
+    // CHECK: size = 8
+    "test.data_layout_query"() : () -> !llvm.ptr<i8, 3>
+    // CHECK: alignment = 8
+    // CHECK: bitsize = 64
+    // CHECK: preferred = 8
+    // CHECK: size = 8
+    "test.data_layout_query"() : () -> !llvm.ptr<i8, 5>
+    return
+  }
+}
+
+// -----
+
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!llvm.ptr<i8>, dense<[32, 32, 64]> : vector<3xi32>>,
+  #dlti.dl_entry<!llvm.ptr<i8, 5>, dense<[64, 64, 64]> : vector<3xi32>>
+>} {
+  // CHECK: @spec
+  func @spec() {
+    // CHECK: alignment = 4
+    // CHECK: bitsize = 32
+    // CHECK: preferred = 8
+    // CHECK: size = 4
+    "test.data_layout_query"() : () -> !llvm.ptr<i8>
+    // CHECK: alignment = 4
+    // CHECK: bitsize = 32
+    // CHECK: preferred = 8
+    // CHECK: size = 4
+    "test.data_layout_query"() : () -> !llvm.ptr<i32>
+    // CHECK: alignment = 4
+    // CHECK: bitsize = 32
+    // CHECK: preferred = 8
+    // CHECK: size = 4
+    "test.data_layout_query"() : () -> !llvm.ptr<bf16>
+    // CHECK: alignment = 4
+    // CHECK: bitsize = 32
+    // CHECK: preferred = 8
+    // CHECK: size = 4
+    "test.data_layout_query"() : () -> !llvm.ptr<!llvm.ptr<i8>>
+    // CHECK: alignment = 4
+    // CHECK: bitsize = 32
+    // CHECK: preferred = 8
+    // CHECK: size = 4
+    "test.data_layout_query"() : () -> !llvm.ptr<i8, 3>
+    // CHECK: alignment = 8
+    // CHECK: bitsize = 64
+    // CHECK: preferred = 8
+    // CHECK: size = 8
+    "test.data_layout_query"() : () -> !llvm.ptr<i8, 5>
+    return
+  }
+}
+
+// -----
+
+// expected-error@below {{unexpected layout attribute for pointer to 'i32'}}
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!llvm.ptr<i32>, dense<[64, 64, 64]> : vector<3xi32>>
+>} {
+  func @pointer() {
+    return
+  }
+}
+
+// -----
+
+// expected-error@below {{expected layout attribute for '!llvm.ptr<i8>' to be a dense integer elements attribute with 3 or 4 elements}}
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!llvm.ptr<i8>, dense<[64.0, 64.0, 64.0]> : vector<3xf32>>
+>} {
+  func @pointer() {
+    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.ptr<i8>, dense<[64, 64, 32]> : vector<3xi32>>
+>} {
+  func @pointer() {
+    return
+  }
+}