[MLIR][LLVM][DLTI] Handle data layout token 'n32:64' (#141299)

Author: bcardosolopes

mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td

@@ -95,6 +95,9 @@ def DLTI_DataLayoutSpecAttr :
     /// Returns the function pointer alignment identifier.
     StringAttr getFunctionPointerAlignmentIdentifier(MLIRContext *context) const;
 
+    /// Returns the legal int widths identifier.
+    StringAttr getLegalIntWidthsIdentifier(MLIRContext *context) const;
+
     /// Returns the attribute associated with the key.
     FailureOr<Attribute> query(DataLayoutEntryKey key) {
       return ::llvm::cast<mlir::DataLayoutSpecInterface>(*this).queryHelper(key);

mlir/include/mlir/Dialect/DLTI/DLTIBase.td

@@ -70,6 +70,9 @@ def DLTI_Dialect : Dialect {
 
     constexpr const static ::llvm::StringLiteral
     kDataLayoutFunctionPointerAlignmentKey = "dlti.function_pointer_alignment";
+
+    constexpr const static ::llvm::StringLiteral
+    kDataLayoutLegalIntWidthsKey = "dlti.legal_int_widths";
   }];
 
   let useDefaultAttributePrinterParser = 1;

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

@@ -105,6 +105,10 @@ uint64_t getDefaultStackAlignment(DataLayoutEntryInterface entry);
 /// the DataLayoutInterface if specified, otherwise returns the default.
 Attribute getDefaultFunctionPointerAlignment(DataLayoutEntryInterface entry);
 
+/// Default handler for the legal int widths request. Dispatches to the
+/// DataLayoutInterface if specified, otherwise returns the default.
+Attribute getDefaultLegalIntWidths(DataLayoutEntryInterface entry);
+
 /// Returns the value of the property from the specified DataLayoutEntry. If the
 /// property is missing from the entry, returns std::nullopt.
 std::optional<Attribute> getDevicePropertyValue(DataLayoutEntryInterface entry);
@@ -266,6 +270,9 @@ class DataLayout {
   /// Returns function pointer alignment.
   Attribute getFunctionPointerAlignment() const;
 
+  /// Returns the legal int widths.
+  Attribute getLegalIntWidths() const;
+
   /// Returns the value of the specified property if the property is defined for
   /// the given device ID, otherwise returns std::nullopt.
   std::optional<Attribute>
@@ -312,6 +319,8 @@ class DataLayout {
   mutable std::optional<uint64_t> stackAlignment;
   /// Cache for function pointer alignment.
   mutable std::optional<Attribute> functionPointerAlignment;
+  /// Cache for legal int widths.
+  mutable std::optional<Attribute> legalIntWidths;
 };
 
 } // namespace mlir

mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

@@ -177,6 +177,12 @@ def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface", [DLTIQuer
       /*methodName=*/"getFunctionPointerAlignmentIdentifier",
       /*args=*/(ins "::mlir::MLIRContext *":$context)
     >,
+    InterfaceMethod<
+      /*description=*/"Returns the legal int widths identifier.",
+      /*retTy=*/"::mlir::StringAttr",
+      /*methodName=*/"getLegalIntWidthsIdentifier",
+      /*args=*/(ins "::mlir::MLIRContext *":$context)
+    >,
     // Implementations may override this if they have an efficient lookup
     // mechanism.
     InterfaceMethod<
@@ -571,6 +577,18 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
         return ::mlir::detail::getDefaultFunctionPointerAlignment(entry);
       }]
     >,
+    StaticInterfaceMethod<
+      /*description=*/"Returns the legal int widths, each width in bits computed "
+                      "using the relevant entries. The data layout object "
+                      "can be used for recursive queries.",
+      /*retTy=*/"Attribute",
+      /*methodName=*/"getLegalIntWidths",
+      /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::detail::getDefaultLegalIntWidths(entry);
+      }]
+    >,
     StaticInterfaceMethod<
       /*description=*/"Returns the value of the property, if the property is "
                       "defined. Otherwise, it returns std::nullopt.",

mlir/lib/Dialect/DLTI/DLTI.cpp

@@ -426,6 +426,12 @@ StringAttr DataLayoutSpecAttr::getFunctionPointerAlignmentIdentifier(
       DLTIDialect::kDataLayoutFunctionPointerAlignmentKey);
 }
 
+StringAttr
+DataLayoutSpecAttr::getLegalIntWidthsIdentifier(MLIRContext *context) const {
+  return Builder(context).getStringAttr(
+      DLTIDialect::kDataLayoutLegalIntWidthsKey);
+}
+
 /// Parses an attribute with syntax:
 ///   dl-spec-attr ::= `#dlti.` `dl_spec` `<` entry-list `>`
 ///   entry-list ::= | entry | entry `,` entry-list
@@ -632,6 +638,7 @@ class TargetDataLayoutInterface : public DataLayoutDialectInterface {
         entryName == DLTIDialect::kDataLayoutGlobalMemorySpaceKey ||
         entryName == DLTIDialect::kDataLayoutStackAlignmentKey ||
         entryName == DLTIDialect::kDataLayoutFunctionPointerAlignmentKey ||
+        entryName == DLTIDialect::kDataLayoutLegalIntWidthsKey ||
         entryName == DLTIDialect::kDataLayoutManglingModeKey)
       return success();
     return emitError(loc) << "unknown data layout entry name: " << entryName;

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

@@ -321,6 +321,16 @@ Attribute mlir::detail::getDefaultFunctionPointerAlignment(
   return entry.getValue();
 }
 
+// Returns the legal int widths if specified in the given entry. If the entry is
+// empty the default legal int widths represented by an empty attribute is
+// returned.
+Attribute
+mlir::detail::getDefaultLegalIntWidths(DataLayoutEntryInterface entry) {
+  if (entry == DataLayoutEntryInterface())
+    return Attribute();
+  return entry.getValue();
+}
+
 std::optional<Attribute>
 mlir::detail::getDevicePropertyValue(DataLayoutEntryInterface entry) {
   if (entry == DataLayoutEntryInterface())
@@ -736,6 +746,22 @@ Attribute mlir::DataLayout::getFunctionPointerAlignment() const {
   return *functionPointerAlignment;
 }
 
+Attribute mlir::DataLayout::getLegalIntWidths() const {
+  checkValid();
+  if (legalIntWidths)
+    return *legalIntWidths;
+  DataLayoutEntryInterface entry;
+  if (originalLayout)
+    entry = originalLayout.getSpecForIdentifier(
+        originalLayout.getLegalIntWidthsIdentifier(
+            originalLayout.getContext()));
+  if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
+    legalIntWidths = iface.getLegalIntWidths(entry);
+  else
+    legalIntWidths = detail::getDefaultLegalIntWidths(entry);
+  return *legalIntWidths;
+}
+
 std::optional<Attribute> mlir::DataLayout::getDevicePropertyValue(
     TargetSystemSpecInterface::DeviceID deviceID,
     StringAttr propertyName) const {

mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp

@@ -63,20 +63,25 @@ FailureOr<uint64_t> DataLayoutImporter::tryToParseInt(StringRef &token) const {
   return parameter;
 }
 
-FailureOr<SmallVector<uint64_t>>
-DataLayoutImporter::tryToParseIntList(StringRef token) const {
+template <class T>
+static FailureOr<SmallVector<T>> tryToParseIntListImpl(StringRef token) {
   SmallVector<StringRef> tokens;
   token.consume_front(":");
   token.split(tokens, ':');
 
   // Parse an integer list.
-  SmallVector<uint64_t> results(tokens.size());
+  SmallVector<T> results(tokens.size());
   for (auto [result, token] : llvm::zip(results, tokens))
     if (token.getAsInteger(/*Radix=*/10, result))
       return failure();
   return results;
 }
 
+FailureOr<SmallVector<uint64_t>>
+DataLayoutImporter::tryToParseIntList(StringRef token) const {
+  return tryToParseIntListImpl<uint64_t>(token);
+}
+
 FailureOr<DenseIntElementsAttr>
 DataLayoutImporter::tryToParseAlignment(StringRef token) const {
   FailureOr<SmallVector<uint64_t>> alignment = tryToParseIntList(token);
@@ -251,6 +256,25 @@ LogicalResult DataLayoutImporter::tryToEmplaceFunctionPointerAlignmentEntry(
   return success();
 }
 
+LogicalResult
+DataLayoutImporter::tryToEmplaceLegalIntWidthsEntry(StringRef token) {
+  auto key =
+      StringAttr::get(context, DLTIDialect::kDataLayoutLegalIntWidthsKey);
+  if (keyEntries.count(key))
+    return success();
+
+  FailureOr<SmallVector<int32_t>> intWidths =
+      tryToParseIntListImpl<int32_t>(token);
+  if (failed(intWidths) || intWidths->empty())
+    return failure();
+
+  OpBuilder builder(context);
+  keyEntries.try_emplace(
+      key,
+      DataLayoutEntryAttr::get(key, builder.getDenseI32ArrayAttr(*intWidths)));
+  return success();
+}
+
 void DataLayoutImporter::translateDataLayout(
     const llvm::DataLayout &llvmDataLayout) {
   dataLayout = {};
@@ -360,6 +384,12 @@ void DataLayoutImporter::translateDataLayout(
         return;
       continue;
     }
+    // Parse native integer widths specifications.
+    if (*prefix == "n") {
+      if (failed(tryToEmplaceLegalIntWidthsEntry(token)))
+        return;
+      continue;
+    }
     // Parse function pointer alignment specifications.
     // Note that prefix here is "Fn" or "Fi", not a single character.
     if (prefix->starts_with("F")) {

mlir/lib/Target/LLVMIR/DataLayoutImporter.h

@@ -113,6 +113,9 @@ class DataLayoutImporter {
   tryToEmplaceFunctionPointerAlignmentEntry(StringRef fnPtrAlignEntry,
                                             StringRef token);
 
+  /// Adds legal int widths entry if there is none yet.
+  LogicalResult tryToEmplaceLegalIntWidthsEntry(StringRef token);
+
   std::string layoutStr = {};
   StringRef lastToken = {};
   SmallVector<StringRef> unhandledTokens;

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -35,6 +35,7 @@
 #include "mlir/Target/LLVMIR/TypeToLLVM.h"
 
 #include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
@@ -52,6 +53,7 @@
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
@@ -250,6 +252,13 @@ translateDataLayout(DataLayoutSpecInterface attribute,
                    << alignment;
       continue;
     }
+    if (key.getValue() == DLTIDialect::kDataLayoutLegalIntWidthsKey) {
+      layoutStream << "-n";
+      llvm::interleave(
+          cast<DenseI32ArrayAttr>(entry.getValue()).asArrayRef(), layoutStream,
+          [&](int32_t val) { layoutStream << val; }, ":");
+      continue;
+    }
     emitError(*loc) << "unsupported data layout key " << key;
     return failure();
   }

mlir/test/Dialect/LLVMIR/layout.mlir

@@ -12,6 +12,7 @@ module {
     // CHECK-SAME: #dlti.function_pointer_alignment<0, function_dependent = false>,
     // CHECK: global_memory_space = 0
     // CHECK: index = 64
+    // CHECK: legal_int_widths = array<i32>
     // CHECK: mangling_mode = ""
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 0
@@ -27,6 +28,7 @@ module {
     // CHECK-SAME: #dlti.function_pointer_alignment<0, function_dependent = false>,
     // CHECK: global_memory_space = 0
     // CHECK: index = 64
+    // CHECK: legal_int_widths = array<i32>
     // CHECK: mangling_mode = ""
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 0
@@ -42,6 +44,7 @@ module {
     // CHECK-SAME: #dlti.function_pointer_alignment<0, function_dependent = false>,
     // CHECK: global_memory_space = 0
     // CHECK: index = 64
+    // CHECK: legal_int_widths = array<i32>
     // CHECK: mangling_mode = ""
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 0
@@ -65,6 +68,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
   #dlti.dl_entry<"dlti.program_memory_space", 3 : ui64>,
   #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>,
   #dlti.dl_entry<"dlti.mangling_mode", "e">,
+  #dlti.dl_entry<"dlti.legal_int_widths", array<i32: 32, 64>>,
   #dlti.dl_entry<"dlti.function_pointer_alignment",
                  "#dlti.function_pointer_alignment<32, function_dependent = true>">
 >} {
@@ -79,6 +83,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK-SAME: "#dlti.function_pointer_alignment<32, function_dependent = true>",
     // CHECK: global_memory_space = 2
     // CHECK: index = 32
+    // CHECK: legal_int_widths = array<i32: 32, 64>
     // CHECK: mangling_mode = "e"
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 3
@@ -94,6 +99,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK-SAME: "#dlti.function_pointer_alignment<32, function_dependent = true>",
     // CHECK: global_memory_space = 2
     // CHECK: index = 32
+    // CHECK: legal_int_widths = array<i32: 32, 64>
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 3
     // CHECK: size = 4
@@ -108,6 +114,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK-SAME: "#dlti.function_pointer_alignment<32, function_dependent = true>",
     // CHECK: global_memory_space = 2
     // CHECK: index = 64
+    // CHECK: legal_int_widths = array<i32: 32, 64>
     // CHECK: mangling_mode = "e"
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 3
@@ -123,6 +130,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK-SAME: "#dlti.function_pointer_alignment<32, function_dependent = true>",
     // CHECK: global_memory_space = 2
     // CHECK: index = 24
+    // CHECK: legal_int_widths = array<i32: 32, 64>
     // CHECK: mangling_mode = "e"
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 3