[mlir][llvm] Support pointer entries in data layout translation

This adds support for pointer DLTI entries in LLVMIR export, e.g.
```
// translated to: p0:32:64:128
#dlti.dl_entry<!llvm.ptr, dense<[32,64,128]> : vector<3xi32>>
// translated to: p1:32:32:32:64
#dlti.dl_entry<!llvm.ptr<1>, dense<[32,32,32,64]> : vector<4xi32>>
```

Reviewed By: ftynse

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

Author: rkayaith

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

@@ -535,6 +535,15 @@ Type getScalableVectorType(Type elementType, unsigned numElements);
 /// (aggregates such as struct) or types that don't have a size (such as void).
 llvm::TypeSize getPrimitiveTypeSizeInBits(Type type);
 
+/// The positions of different values in the data layout entry for pointers.
+enum class PtrDLEntryPos { Size = 0, Abi = 1, Preferred = 2, Index = 3 };
+
+/// Returns the value that corresponds to named position `pos` from the
+/// data layout entry `attr` assuming it's a dense integer elements attribute.
+/// Returns `None` if `pos` is not present in the entry.
+/// Currently only `PtrDLEntryPos::Index` is optional, and all other positions
+/// may be assumed to be present.
+Optional<unsigned> extractPointerSpecValue(Attribute attr, PtrDLEntryPos pos);
 } // namespace LLVM
 } // namespace mlir
 

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

@@ -229,19 +229,16 @@ LLVMPointerType::verify(function_ref<InFlightDiagnostic()> emitError,
   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;
 
-/// 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>()
-      .getValues<unsigned>()[static_cast<unsigned>(pos)];
+Optional<unsigned> mlir::LLVM::extractPointerSpecValue(Attribute attr,
+                                                       PtrDLEntryPos pos) {
+  auto spec = attr.cast<DenseIntElementsAttr>();
+  auto idx = static_cast<unsigned>(pos);
+  if (idx >= spec.size())
+    return None;
+  return spec.getValues<unsigned>()[idx];
 }
 
 /// Returns the part of the data layout entry that corresponds to `pos` for the
@@ -250,7 +247,7 @@ static unsigned extractPointerSpecValue(Attribute attr, DLEntryPos pos) {
 /// do not provide a custom one, for other address spaces returns None.
 static Optional<unsigned>
 getPointerDataLayoutEntry(DataLayoutEntryListRef params, LLVMPointerType type,
-                          DLEntryPos pos) {
+                          PtrDLEntryPos pos) {
   // First, look for the entry for the pointer in the current address space.
   Attribute currentEntry;
   for (DataLayoutEntryInterface entry : params) {
@@ -263,15 +260,15 @@ getPointerDataLayoutEntry(DataLayoutEntryListRef params, LLVMPointerType type,
     }
   }
   if (currentEntry) {
-    return extractPointerSpecValue(currentEntry, pos) /
-           (pos == DLEntryPos::Size ? 1 : kBitsInByte);
+    return *extractPointerSpecValue(currentEntry, pos) /
+           (pos == PtrDLEntryPos::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 pos == PtrDLEntryPos::Size ? kDefaultPointerSizeBits
+                                      : kDefaultPointerAlignment;
   }
 
   return llvm::None;
@@ -281,7 +278,7 @@ unsigned
 LLVMPointerType::getTypeSizeInBits(const DataLayout &dataLayout,
                                    DataLayoutEntryListRef params) const {
   if (Optional<unsigned> size =
-          getPointerDataLayoutEntry(params, *this, DLEntryPos::Size))
+          getPointerDataLayoutEntry(params, *this, PtrDLEntryPos::Size))
     return *size;
 
   // For other memory spaces, use the size of the pointer to the default memory
@@ -294,7 +291,7 @@ LLVMPointerType::getTypeSizeInBits(const DataLayout &dataLayout,
 unsigned LLVMPointerType::getABIAlignment(const DataLayout &dataLayout,
                                           DataLayoutEntryListRef params) const {
   if (Optional<unsigned> alignment =
-          getPointerDataLayoutEntry(params, *this, DLEntryPos::Abi))
+          getPointerDataLayoutEntry(params, *this, PtrDLEntryPos::Abi))
     return *alignment;
 
   if (isOpaque())
@@ -306,7 +303,7 @@ unsigned
 LLVMPointerType::getPreferredAlignment(const DataLayout &dataLayout,
                                        DataLayoutEntryListRef params) const {
   if (Optional<unsigned> alignment =
-          getPointerDataLayoutEntry(params, *this, DLEntryPos::Preferred))
+          getPointerDataLayoutEntry(params, *this, PtrDLEntryPos::Preferred))
     return *alignment;
 
   if (isOpaque())
@@ -339,13 +336,13 @@ bool LLVMPointerType::areCompatible(DataLayoutEntryListRef oldLayout,
       });
     }
     if (it != oldLayout.end()) {
-      size = extractPointerSpecValue(*it, DLEntryPos::Size);
-      abi = extractPointerSpecValue(*it, DLEntryPos::Abi);
+      size = *extractPointerSpecValue(*it, PtrDLEntryPos::Size);
+      abi = *extractPointerSpecValue(*it, PtrDLEntryPos::Abi);
     }
 
     Attribute newSpec = newEntry.getValue().cast<DenseIntElementsAttr>();
-    unsigned newSize = extractPointerSpecValue(newSpec, DLEntryPos::Size);
-    unsigned newAbi = extractPointerSpecValue(newSpec, DLEntryPos::Abi);
+    unsigned newSize = *extractPointerSpecValue(newSpec, PtrDLEntryPos::Size);
+    unsigned newAbi = *extractPointerSpecValue(newSpec, PtrDLEntryPos::Abi);
     if (size != newSize || abi < newAbi || abi % newAbi != 0)
       return false;
   }
@@ -369,8 +366,8 @@ LogicalResult LLVMPointerType::verifyEntries(DataLayoutEntryListRef entries,
       return emitError(loc) << "unexpected layout attribute for pointer to "
                             << key.getElementType();
     }
-    if (extractPointerSpecValue(values, DLEntryPos::Abi) >
-        extractPointerSpecValue(values, DLEntryPos::Preferred)) {
+    if (extractPointerSpecValue(values, PtrDLEntryPos::Abi) >
+        extractPointerSpecValue(values, PtrDLEntryPos::Preferred)) {
       return emitError(loc) << "preferred alignment is expected to be at least "
                                "as large as ABI alignment";
     }

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -52,8 +52,8 @@ using namespace mlir::LLVM::detail;
 #include "mlir/Dialect/LLVMIR/LLVMConversionEnumsToLLVM.inc"
 
 /// Translates the given data layout spec attribute to the LLVM IR data layout.
-/// Only integer, float and endianness entries are currently supported.
-FailureOr<llvm::DataLayout>
+/// Only integer, float, pointer and endianness entries are currently supported.
+static FailureOr<llvm::DataLayout>
 translateDataLayout(DataLayoutSpecInterface attribute,
                     const DataLayout &dataLayout,
                     Optional<Location> loc = llvm::None) {
@@ -80,41 +80,56 @@ translateDataLayout(DataLayoutSpecInterface attribute,
   }
 
   // Go through the list of entries to check which types are explicitly
-  // specified in entries. Don't use the entries directly though but query the
-  // data from the layout.
+  // specified in entries. Where possible, data layout queries are used instead
+  // of directly inspecting the entries.
   for (DataLayoutEntryInterface entry : attribute.getEntries()) {
     auto type = entry.getKey().dyn_cast<Type>();
     if (!type)
       continue;
     // Data layout for the index type is irrelevant at this point.
     if (type.isa<IndexType>())
       continue;
-    FailureOr<std::string> prefix =
-        llvm::TypeSwitch<Type, FailureOr<std::string>>(type)
-            .Case<IntegerType>(
-                [loc](IntegerType integerType) -> FailureOr<std::string> {
-                  if (integerType.getSignedness() == IntegerType::Signless)
-                    return std::string("i");
-                  emitError(*loc)
-                      << "unsupported data layout for non-signless integer "
-                      << integerType;
-                  return failure();
-                })
-            .Case<Float16Type, Float32Type, Float64Type, Float80Type,
-                  Float128Type>([](Type) { return std::string("f"); })
-            .Default([loc](Type type) -> FailureOr<std::string> {
-              emitError(*loc) << "unsupported type in data layout: " << type;
-              return failure();
+    layoutStream << "-";
+    LogicalResult result =
+        llvm::TypeSwitch<Type, LogicalResult>(type)
+            .Case<IntegerType, Float16Type, Float32Type, Float64Type,
+                  Float80Type, Float128Type>([&](Type type) -> LogicalResult {
+              if (auto intType = type.dyn_cast<IntegerType>()) {
+                if (intType.getSignedness() != IntegerType::Signless)
+                  return emitError(*loc)
+                         << "unsupported data layout for non-signless integer "
+                         << intType;
+                layoutStream << "i";
+              } else {
+                layoutStream << "f";
+              }
+              unsigned size = dataLayout.getTypeSizeInBits(type);
+              unsigned abi = dataLayout.getTypeABIAlignment(type) * 8u;
+              unsigned preferred =
+                  dataLayout.getTypePreferredAlignment(type) * 8u;
+              layoutStream << size << ":" << abi;
+              if (abi != preferred)
+                layoutStream << ":" << preferred;
+              return success();
+            })
+            .Case([&](LLVMPointerType ptrType) {
+              layoutStream << "p" << ptrType.getAddressSpace() << ":";
+              unsigned size = dataLayout.getTypeSizeInBits(type);
+              unsigned abi = dataLayout.getTypeABIAlignment(type) * 8u;
+              unsigned preferred =
+                  dataLayout.getTypePreferredAlignment(type) * 8u;
+              layoutStream << size << ":" << abi << ":" << preferred;
+              if (Optional<unsigned> index = extractPointerSpecValue(
+                      entry.getValue(), PtrDLEntryPos::Index))
+                layoutStream << ":" << *index;
+              return success();
+            })
+            .Default([loc](Type type) {
+              return emitError(*loc)
+                     << "unsupported type in data layout: " << type;
             });
-    if (failed(prefix))
+    if (failed(result))
       return failure();
-
-    unsigned size = dataLayout.getTypeSizeInBits(type);
-    unsigned abi = dataLayout.getTypeABIAlignment(type) * 8u;
-    unsigned preferred = dataLayout.getTypePreferredAlignment(type) * 8u;
-    layoutStream << "-" << *prefix << size << ":" << abi;
-    if (abi != preferred)
-      layoutStream << ":" << preferred;
   }
   layoutStream.flush();
   StringRef layoutSpec(llvmDataLayout);

mlir/test/Target/LLVMIR/data-layout.mlir

@@ -4,11 +4,15 @@
 // CHECK: E-
 // CHECK: i64:64:128
 // CHECK: f80:128:256
+// CHECK: p0:32:64:128
+// CHECK: p1:32:32:32:64
 module attributes {dlti.dl_spec = #dlti.dl_spec<
 #dlti.dl_entry<"dlti.endianness", "big">,
 #dlti.dl_entry<index, 64>,
 #dlti.dl_entry<i64, dense<[64,128]> : vector<2xi32>>,
-#dlti.dl_entry<f80, dense<[128,256]> : vector<2xi32>>
+#dlti.dl_entry<f80, dense<[128,256]> : vector<2xi32>>,
+#dlti.dl_entry<!llvm.ptr, dense<[32,64,128]> : vector<3xi32>>,
+#dlti.dl_entry<!llvm.ptr<1>, dense<[32,32,32,64]> : vector<4xi32>>
 >} {
   llvm.func @foo() {
     llvm.return