[mlir] Add partial support for translating data layout

Add support for translating data layout specifications for integer and float
types between MLIR and LLVM IR. This is a first step towards removing the
string-based LLVM dialect data layout attribute on modules. The latter is still
available and will remain so until the first-class MLIR modeling can fully
replace it.

Depends On D120739

Reviewed By: wsmoses

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

mlir/include/mlir/Target/LLVMIR/Import.h

@@ -20,12 +20,13 @@
 
 // Forward-declare LLVM classes.
 namespace llvm {
+class DataLayout;
 class Module;
 } // namespace llvm
 
 namespace mlir {
 
-class DialectRegistry;
+class DataLayoutSpecInterface;
 class MLIRContext;
 class ModuleOp;
 
@@ -37,6 +38,11 @@ OwningOpRef<ModuleOp>
 translateLLVMIRToModule(std::unique_ptr<llvm::Module> llvmModule,
                         MLIRContext *context);
 
+/// Translate the given LLVM data layout into an MLIR equivalent using the DLTI
+/// dialect.
+DataLayoutSpecInterface translateDataLayout(const llvm::DataLayout &dataLayout,
+                                            MLIRContext *context);
+
 } // namespace mlir
 
 #endif // MLIR_TARGET_LLVMIR_IMPORT_H

mlir/lib/Target/LLVMIR/CMakeLists.txt

@@ -27,6 +27,7 @@ add_mlir_translation_library(MLIRTargetLLVMIRExport
   TransformUtils
 
   LINK_LIBS PUBLIC
+  MLIRDLTI
   MLIRLLVMIR
   MLIRLLVMIRTransforms
   MLIRTranslation
@@ -59,6 +60,7 @@ add_mlir_translation_library(MLIRTargetLLVMIRImport
   IRReader
 
   LINK_LIBS PUBLIC
+  MLIRDLTI
   MLIRLLVMIR
   MLIRTranslation
   )

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

@@ -12,11 +12,13 @@
 
 #include "mlir/Target/LLVMIR/Import.h"
 
+#include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/MLIRContext.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Target/LLVMIR/TypeFromLLVM.h"
 #include "mlir/Translation.h"
 
@@ -47,6 +49,91 @@ static std::string diag(llvm::Value &v) {
   return os.str();
 }
 
+/// Creates an attribute containing ABI and preferred alignment numbers parsed
+/// a string. The string may be either "abi:preferred" or just "abi". In the
+/// latter case, the prefrred alignment is considered equal to ABI alignment.
+static DenseIntElementsAttr parseDataLayoutAlignment(MLIRContext &ctx,
+                                                     StringRef spec) {
+  auto i32 = IntegerType::get(&ctx, 32);
+
+  StringRef abiString, preferredString;
+  std::tie(abiString, preferredString) = spec.split(':');
+  int abi, preferred;
+  if (abiString.getAsInteger(/*Radix=*/10, abi))
+    return nullptr;
+
+  if (preferredString.empty())
+    preferred = abi;
+  else if (preferredString.getAsInteger(/*Radix=*/10, preferred))
+    return nullptr;
+
+  return DenseIntElementsAttr::get(VectorType::get({2}, i32), {abi, preferred});
+}
+
+/// Returns a supported MLIR floating point type of the given bit width or null
+/// if the bit width is not supported.
+static FloatType getDLFloatType(MLIRContext &ctx, int32_t bitwidth) {
+  switch (bitwidth) {
+  case 16:
+    return FloatType::getF16(&ctx);
+  case 32:
+    return FloatType::getF32(&ctx);
+  case 64:
+    return FloatType::getF64(&ctx);
+  case 80:
+    return FloatType::getF80(&ctx);
+  case 128:
+    return FloatType::getF128(&ctx);
+  default:
+    return nullptr;
+  }
+}
+
+DataLayoutSpecInterface
+mlir::translateDataLayout(const llvm::DataLayout &dataLayout,
+                          MLIRContext *context) {
+  assert(context && "expected MLIR context");
+  StringRef layout = dataLayout.getStringRepresentation();
+  SmallVector<DataLayoutEntryInterface> entries;
+  while (!layout.empty()) {
+    // Split at '-'.
+    std::pair<StringRef, StringRef> split = layout.split('-');
+    StringRef current;
+    std::tie(current, layout) = split;
+
+    // Split at ':'.
+    StringRef kind, spec;
+    std::tie(kind, spec) = current.split(':');
+
+    char symbol = kind.front();
+    StringRef parameter = kind.substr(1);
+
+    if (symbol == 'i' || symbol == 'f') {
+      unsigned bitwidth;
+      if (parameter.getAsInteger(/*Radix=*/10, bitwidth))
+        return nullptr;
+      DenseIntElementsAttr params = parseDataLayoutAlignment(*context, spec);
+      if (!params)
+        return nullptr;
+      auto entry = DataLayoutEntryAttr::get(
+          symbol == 'i' ? static_cast<Type>(IntegerType::get(context, bitwidth))
+                        : getDLFloatType(*context, bitwidth),
+          params);
+      entries.emplace_back(entry);
+    } else if (symbol == 'e' || symbol == 'E') {
+      auto value = StringAttr::get(
+          context, symbol == 'e' ? DLTIDialect::kDataLayoutEndiannessLittle
+                                 : DLTIDialect::kDataLayoutEndiannessBig);
+      auto entry = DataLayoutEntryAttr::get(
+          StringAttr::get(context, DLTIDialect::kDataLayoutEndiannessKey),
+          value);
+      entries.emplace_back(entry);
+    }
+  }
+
+  return DataLayoutSpecAttr::get(context, entries);
+}
+
 // Handles importing globals and functions from an LLVM module.
 namespace {
 class Importer {
@@ -862,9 +949,19 @@ OwningOpRef<ModuleOp>
 mlir::translateLLVMIRToModule(std::unique_ptr<llvm::Module> llvmModule,
                               MLIRContext *context) {
   context->loadDialect<LLVMDialect>();
+  context->loadDialect<DLTIDialect>();
   OwningOpRef<ModuleOp> module(ModuleOp::create(
       FileLineColLoc::get(context, "", /*line=*/0, /*column=*/0)));
 
+  DataLayoutSpecInterface dlSpec =
+      translateDataLayout(llvmModule->getDataLayout(), context);
+  if (!dlSpec) {
+    emitError(UnknownLoc::get(context), "can't translate data layout");
+    return {};
+  }
+
+  module.get()->setAttr(DLTIDialect::kDataLayoutAttrName, dlSpec);
+
   Importer deserializer(context, module.get());
   for (llvm::GlobalVariable &gv : llvmModule->globals()) {
     if (!deserializer.processGlobal(&gv))

mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Target/LLVMIR/Dialect/All.h"
 #include "mlir/Target/LLVMIR/Export.h"
@@ -33,6 +34,7 @@ void registerToLLVMIRTranslation() {
         return success();
       },
       [](DialectRegistry &registry) {
+        registry.insert<DLTIDialect>();
         registerAllToLLVMIRTranslations(registry);
       });
 }

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -14,6 +14,7 @@
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 
 #include "DebugTranslation.h"
+#include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/Transforms/LegalizeForExport.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
@@ -50,6 +51,76 @@ 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>
+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;
+  llvm::raw_string_ostream layoutStream(llvmDataLayout);
+  for (DataLayoutEntryInterface entry : attribute.getEntries()) {
+    auto key = entry.getKey().dyn_cast<StringAttr>();
+    if (!key)
+      continue;
+    if (key.getValue() == DLTIDialect::kDataLayoutEndiannessKey) {
+      auto value = entry.getValue().cast<StringAttr>();
+      bool isLittleEndian =
+          value.getValue() == DLTIDialect::kDataLayoutEndiannessLittle;
+      layoutStream << (isLittleEndian ? "e" : "E");
+      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.
+  for (DataLayoutEntryInterface entry : attribute.getEntries()) {
+    auto type = entry.getKey().dyn_cast<Type>();
+    if (!type)
+      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();
+            });
+    if (failed(prefix))
+      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);
+  if (layoutSpec.startswith("-"))
+    layoutSpec = layoutSpec.drop_front();
+
+  return llvm::DataLayout(layoutSpec);
+}
+
 /// Builds a constant of a sequential LLVM type `type`, potentially containing
 /// other sequential types recursively, from the individual constant values
 /// provided in `constants`. `shape` contains the number of elements in nested
@@ -1010,8 +1081,25 @@ prepareLLVMModule(Operation *m, llvm::LLVMContext &llvmContext,
   m->getContext()->getOrLoadDialect<LLVM::LLVMDialect>();
   auto llvmModule = std::make_unique<llvm::Module>(name, llvmContext);
   if (auto dataLayoutAttr =
-          m->getAttr(LLVM::LLVMDialect::getDataLayoutAttrName()))
+          m->getAttr(LLVM::LLVMDialect::getDataLayoutAttrName())) {
     llvmModule->setDataLayout(dataLayoutAttr.cast<StringAttr>().getValue());
+  } else {
+    FailureOr<llvm::DataLayout> llvmDataLayout(llvm::DataLayout(""));
+    if (auto iface = dyn_cast<DataLayoutOpInterface>(m)) {
+      if (DataLayoutSpecInterface spec = iface.getDataLayoutSpec()) {
+        llvmDataLayout =
+            translateDataLayout(spec, DataLayout(iface), m->getLoc());
+      }
+    } else if (auto mod = dyn_cast<ModuleOp>(m)) {
+      if (DataLayoutSpecInterface spec = mod.getDataLayoutSpec()) {
+        llvmDataLayout =
+            translateDataLayout(spec, DataLayout(mod), m->getLoc());
+      }
+    }
+    if (failed(llvmDataLayout))
+      return nullptr;
+    llvmModule->setDataLayout(*llvmDataLayout);
+  }
   if (auto targetTripleAttr =
           m->getAttr(LLVM::LLVMDialect::getTargetTripleAttrName()))
     llvmModule->setTargetTriple(targetTripleAttr.cast<StringAttr>().getValue());
@@ -1032,8 +1120,11 @@ mlir::translateModuleToLLVMIR(Operation *module, llvm::LLVMContext &llvmContext,
                               StringRef name) {
   if (!satisfiesLLVMModule(module))
     return nullptr;
+
   std::unique_ptr<llvm::Module> llvmModule =
       prepareLLVMModule(module, llvmContext, name);
+  if (!llvmModule)
+    return nullptr;
 
   LLVM::ensureDistinctSuccessors(module);
 

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

@@ -0,0 +1,10 @@
+; RUN: mlir-translate -import-llvm %s | FileCheck %s
+
+; CHECK: dlti.dl_spec = 
+; CHECK: #dlti.dl_spec<
+; CHECK:   #dlti.dl_entry<"dlti.endianness", "little">
+; CHECK:   #dlti.dl_entry<i64, dense<64> : vector<2xi32>>
+; CHECK:   #dlti.dl_entry<f80, dense<128> : vector<2xi32>>
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+declare void @foo()

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

@@ -0,0 +1,36 @@
+// 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
+module attributes {dlti.dl_spec = #dlti.dl_spec<
+#dlti.dl_entry<"dlti.endianness", "big">,
+#dlti.dl_entry<i64, dense<[64,128]> : vector<2xi32>>,
+#dlti.dl_entry<f80, dense<[128,256]> : vector<2xi32>>
+>} {
+  llvm.func @foo() {
+    llvm.return
+  }
+}
+
+// -----
+
+// expected-error@below {{unsupported data layout for non-signless integer 'ui64'}}
+module attributes {dlti.dl_spec = #dlti.dl_spec<
+#dlti.dl_entry<ui64, dense<[64,128]> : vector<2xi32>>>
+} {}
+
+// -----
+
+// expected-error@below {{unsupported type in data layout: 'bf16'}}
+module attributes {dlti.dl_spec = #dlti.dl_spec<
+#dlti.dl_entry<bf16, dense<[64,128]> : vector<2xi32>>>
+} {}
+
+// -----
+
+// expected-error@below {{unsupported data layout key "foo"}}
+module attributes {dlti.dl_spec = #dlti.dl_spec<
+#dlti.dl_entry<"foo", dense<[64,128]> : vector<2xi32>>>
+} {}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

@@ -5542,6 +5542,7 @@ cc_library(
     ],
     includes = ["include"],
     deps = [
+        ":DLTIDialect",
         ":IR",
         ":LLVMConversionIncGen",
         ":LLVMDialect",
@@ -5725,6 +5726,7 @@ cc_library(
     includes = ["include"],
     deps = [
         ":AllToLLVMIRTranslations",
+        ":DLTIDialect",
         ":IR",
         ":ToLLVMIRTranslation",
         ":Translation",
@@ -5745,6 +5747,7 @@ cc_library(
     ],
     includes = ["include"],
     deps = [
+        ":DLTIDialect",
         ":IR",
         ":LLVMConversionIncGen",
         ":LLVMDialect",