[mlir] Add partial support for translating data layout

[ftynse]

Summary

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

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

 
 #include "mlir/IR/OwningOpRef.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/StringRef.h"
 #include <memory>
 
 // Forward-declare LLVM classes.
 namespace llvm {
+class DataLayout;
 class Module;
 } // namespace llvm
 
 namespace mlir {
 
-class DialectRegistry;
+class DataLayoutSpecInterface;
 class MLIRContext;
 class ModuleOp;
 
 /// Convert the given LLVM module into MLIR's LLVM dialect.  The LLVM context is
 /// extracted from the registered LLVM IR dialect. In case of error, report it
 /// to the error handler registered with the MLIR context, if any (obtained from
 /// the MLIR module), and return `{}`.
 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

add_mlir_translation_library(MLIRTargetLLVMIRExport
   intrinsics_gen
 
   LINK_COMPONENTS
   Core
   FrontendOpenMP
   TransformUtils
 
   LINK_LIBS PUBLIC
+  MLIRDLTI
   MLIRLLVMIR
   MLIRLLVMIRTransforms
   MLIRTranslation
   )
 
 add_mlir_translation_library(MLIRToLLVMIRTranslationRegistration
   ConvertToLLVMIR.cpp
 
add_mlir_translation_library(MLIRTargetLLVMIRImport
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Target/LLVMIR
 
   LINK_COMPONENTS
   Core
   IRReader
 
   LINK_LIBS PUBLIC
+  MLIRDLTI
   MLIRLLVMIR
   MLIRTranslation
   )

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

 //===- ConvertFromLLVMIR.cpp - MLIR to LLVM IR conversion -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a translation between LLVM IR and the MLIR LLVM dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"

 // operator << (raw_ostream&) defined.
 static std::string diag(llvm::Value &v) {
   std::string s;
   llvm::raw_string_ostream os(s);
   os << 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 {
 public:
   Importer(MLIRContext *context, ModuleOp module)
       : b(context), context(context), module(module),
         unknownLoc(FileLineColLoc::get(context, "imported-bitcode", 0, 0)),
         typeTranslator(*context) {
LogicalResult Importer::processBasicBlock(llvm::BasicBlock *bb, Block *block) {
   }
   return success();
 }
 
 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))
       return {};
   }
   for (llvm::Function &f : llvmModule->functions()) {
     if (failed(deserializer.processFunction(&f)))
       return {};

mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp

 //===- ConvertToLLVMIR.cpp - MLIR to LLVM IR conversion -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a translation between the MLIR LLVM dialect and LLVM IR.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Target/LLVMIR/Dialect/All.h"
 #include "mlir/Target/LLVMIR/Export.h"
 #include "mlir/Translation.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 
 using namespace mlir;
 
 namespace mlir {
 void registerToLLVMIRTranslation() {
   TranslateFromMLIRRegistration registration(
       "mlir-to-llvmir",
       [](ModuleOp module, raw_ostream &output) {
         llvm::LLVMContext llvmContext;
         auto llvmModule = translateModuleToLLVMIR(module, llvmContext);
         if (!llvmModule)
           return failure();
 
         llvmModule->print(output, nullptr);
         return success();
       },
       [](DialectRegistry &registry) {
+        registry.insert<DLTIDialect>();
         registerAllToLLVMIRTranslations(registry);
       });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 //===- ModuleTranslation.cpp - MLIR to LLVM conversion --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the translation between an MLIR LLVM dialect module and
 // the corresponding LLVMIR module. It only handles core LLVM IR operations.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/RegionGraphTraits.h"
 #include "mlir/Support/LLVM.h"

 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
 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>
+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
 /// sequential types. Reports errors at `loc` and returns nullptr on error.
 static llvm::Constant *
 buildSequentialConstant(ArrayRef<llvm::Constant *> &constants,
                         ArrayRef<int64_t> shape, llvm::Type *type,
                         Location loc) {

 void ModuleTranslation::StackFrame::anchor() {}
 
 static std::unique_ptr<llvm::Module>
 prepareLLVMModule(Operation *m, llvm::LLVMContext &llvmContext,
                   StringRef name) {
   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());
 
   // Inject declarations for `malloc` and `free` functions that can be used in
   // memref allocation/deallocation coming from standard ops lowering.
   llvm::IRBuilder<> builder(llvmContext);
   llvmModule->getOrInsertFunction("malloc", builder.getInt8PtrTy(),
                                   builder.getInt64Ty());
   llvmModule->getOrInsertFunction("free", builder.getVoidTy(),
                                   builder.getInt8PtrTy());
 
   return llvmModule;
 }
 
 std::unique_ptr<llvm::Module>
 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);
 
   ModuleTranslation translator(module, std::move(llvmModule));
   if (failed(translator.convertFunctionSignatures()))
     return nullptr;
   if (failed(translator.convertGlobals()))
     return nullptr;

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

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

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

cc_library(
     hdrs = [
         "include/mlir/Target/LLVMIR/Export.h",
         "include/mlir/Target/LLVMIR/LLVMTranslationInterface.h",
         "include/mlir/Target/LLVMIR/ModuleTranslation.h",
         "include/mlir/Target/LLVMIR/TypeToLLVM.h",
     ],
     includes = ["include"],
     deps = [
+        ":DLTIDialect",
         ":IR",
         ":LLVMConversionIncGen",
         ":LLVMDialect",
         ":LLVMIRTransforms",
         ":OpenMPDialect",
         ":Support",
         "//llvm:Core",
         "//llvm:FrontendOpenMP",

 )
 
 cc_library(
     name = "ToLLVMIRTranslationRegistration",
     srcs = ["lib/Target/LLVMIR/ConvertToLLVMIR.cpp"],
     includes = ["include"],
     deps = [
         ":AllToLLVMIRTranslations",
+        ":DLTIDialect",
         ":IR",
         ":ToLLVMIRTranslation",
         ":Translation",
         "//llvm:Core",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "FromLLVMIRTranslation",
     srcs = [
         "lib/Target/LLVMIR/ConvertFromLLVMIR.cpp",
         "lib/Target/LLVMIR/TypeFromLLVM.cpp",
     ],
     hdrs = [
         "include/mlir/Target/LLVMIR/Import.h",
         "include/mlir/Target/LLVMIR/TypeFromLLVM.h",
     ],
     includes = ["include"],
     deps = [
+        ":DLTIDialect",
         ":IR",
         ":LLVMConversionIncGen",
         ":LLVMDialect",
         ":Support",
         ":Translation",
         "//llvm:Core",
         "//llvm:IRReader",
         "//llvm:Support",