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

[rkayaith]

Summary

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

[rkayaith]

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

7

This ignores the idx field of the pointer spec. I'm not sure how to handle that, since we don't support per-addrspace index widths in MLIR. Would it be correct to just set all address spaces to use the type size of index? And how should that work on the import side, when translating datalayout -> DLTI?

[rkayaith]

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

7

@ftynse any thoughts here?

[ftynse]

Could we also add a test with non-default address space to see how it looks in DLTI?

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

7

I think we can support the field, at least for roundtripping. At the LLVM dialect level, there is no index type anymore. Any pointer address manipulation that treats them as integers is better off using whatever LLVM expects it to use rather than the converted index type.

[rkayaith]

add test using non-default address space

[rkayaith]

Support translating the (optional) 'idx' layout value as well

[rkayaith]

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

7

I've added support for exporting the 4th element of dense attribute as the idx field. Looks like the verifier already allowed an optional fourth element so it didn't need any updates, though I didn't see it being used anywhere until now.

I had to make extractPointerSpecValue public to be able to use it during translation.

[ftynse]

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

538

Nit: carry over the documentation comment from DLEntryPos here.

539

[rkayaith]

update comments

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

 Type getScalableVectorType(Type elementType, unsigned numElements);
 
 /// Returns the size of the given primitive LLVM dialect-compatible type
 /// (including vectors) in bits, for example, the size of i16 is 16 and
 /// the size of vector<4xi16> is 64. Returns 0 for non-primitive
 /// (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
 
 #endif // MLIR_DIALECT_LLVMIR_LLVMTYPES_H_

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

 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;
 
-/// 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
 /// 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) {
+                          PtrDLEntryPos 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);
+    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
+    return pos == PtrDLEntryPos::Size ? kDefaultPointerSizeBits
                                       : kDefaultPointerAlignment;
   }
 
   return llvm::None;
 }
 
 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
   // space.
   if (isOpaque())
     return dataLayout.getTypeSizeInBits(get(getContext()));
   return dataLayout.getTypeSizeInBits(get(getElementType()));
 }
 
 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())
     return dataLayout.getTypeABIAlignment(get(getContext()));
   return dataLayout.getTypeABIAlignment(get(getElementType()));
 }
 
 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())
     return dataLayout.getTypePreferredAlignment(get(getContext()));
   return dataLayout.getTypePreferredAlignment(get(getElementType()));
 }
 
 bool LLVMPointerType::areCompatible(DataLayoutEntryListRef oldLayout,
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);
+      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;
   }
   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() && !key.getElementType().isInteger(8)) {
       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";
     }
   }
   return success();
 }
 
 void LLVMPointerType::walkImmediateSubElements(

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 
 using namespace mlir;
 using namespace mlir::LLVM;
 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) {
   if (!loc)
     loc = UnknownLoc::get(attribute.getContext());
 
   // Translate the endianness attribute.
   std::string llvmDataLayout;
if (key.getValue() == DLTIDialect::kDataLayoutEndiannessKey) {
       layoutStream.flush();
       continue;
     }
     emitError(*loc) << "unsupported data layout key " << key;
     return failure();
   }
 
   // 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)
+    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 "
-                      << 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();
-            });
-    if (failed(prefix))
-      return failure();
-
+                         << intType;
+                layoutStream << "i";
+              } else {
+                layoutStream << "f";
+              }
               unsigned size = dataLayout.getTypeSizeInBits(type);
               unsigned abi = dataLayout.getTypeABIAlignment(type) * 8u;
-    unsigned preferred = dataLayout.getTypePreferredAlignment(type) * 8u;
-    layoutStream << "-" << *prefix << size << ":" << abi;
+              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(result))
+      return failure();
   }
   layoutStream.flush();
   StringRef layoutSpec(llvmDataLayout);
   if (layoutSpec.startswith("-"))
     layoutSpec = layoutSpec.drop_front();
 
   return llvm::DataLayout(layoutSpec);
 }

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

 // RUN: mlir-translate -mlir-to-llvmir %s -split-input-file -verify-diagnostics | FileCheck %s
 
 // CHECK: target datalayout
 // 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
   }
 }
 
 // -----