[mlir] Move data layout from LLVMDialect to module Op attributes

[ftynse]

Summary

Legacy implementation of the LLVM dialect in MLIR contained an instance of
llvm::Module as it was required to parse LLVM IR types. The access to the data
layout of this module was exposed to the users for convenience, but in practice
this layout has always been the default one obtained by parsing an empty layout
description string. Current implementation of the dialect no longer relies on
wrapping LLVM IR types, but it kept an instance of DataLayout for
compatibility. This effectively forces a single data layout to be used across
all modules in a given MLIR context, which is not desirable. Remove DataLayout
from the LLVM dialect and attach it as a module attribute instead. Since MLIR
does not yet have support for data layouts, use the LLVM DataLayout in string
form with verification inside MLIR. Introduce the layout when converting a
module to the LLVM dialect and keep the default "" description for
compatibility.

This approach should be replaced with a proper MLIR-based data layout when it
becomes available, but provides an immediate solution to compiling modules with
different layouts, e.g. for GPUs.

This removes the need for LLVMDialectImpl, which is also removed.

Depends On D85650

[ftynse]

Forgot to git-add...

[aartbik]

Does this have impact on our default JIT pipeline?

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

129

can we remove the TODO here
(it is still true, but it will probably be transparent at this point, right?)

[ftynse]

In D85652#2211870, @aartbik wrote:

Does this have impact on our default JIT pipeline?

Immediately, no, since the data layout has always being the default-constructed object and the execution engine was blindly resetting that to the current arch. If we keep it as llvm.data_layout = "", we will obtain exactly the same behavior. But now we can also set the proper data layout earlier in the process and use it in transformations.

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

129

It still depends on the translation, but arguably it should not. This TODO can be removed when we can ask MLIR about the preferred alignment for a vector instead of asking LLVM.

[ftynse]

clang-tidy

[ftynse]

Rebase

mlir/include/mlir/Conversion/Passes.td

Option<"useBarePtrCallConv", "use-bare-ptr-memref-call-conv", "bool",
            "Replace FuncOp's MemRef arguments with bare pointers to the MemRef "
            "element types">,
     Option<"emitCWrappers", "emit-c-wrappers", "bool", /*default=*/"false",
            "Emit wrappers for C-compatible pointer-to-struct memref "
            "descriptors">,
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
            /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
            "Bitwidth of the index type, 0 to use size of machine word">,
+    Option<"dataLayout", "data-layout", "std::string",
+           /*default=*/"\"\"",
+           "String description (LLVM format) of the data layout that is "
+           "expected on the produced module">
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // StandardToSPIRV
 //===----------------------------------------------------------------------===//
 
 def LegalizeStandardForSPIRV : Pass<"legalize-std-for-spirv"> {

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h

public:
   /// C/C++ ABI lowering related to struct argument passing.
   Value promoteOneMemRefDescriptor(Location loc, Value operand,
                                    OpBuilder &builder);
 
   /// Converts the function type to a C-compatible format, in particular using
   /// pointers to memref descriptors for arguments.
   LLVM::LLVMType convertFunctionTypeCWrapper(FunctionType type);
 
+  /// Returns the data layout to use during and after conversion.
+  const llvm::DataLayout &getDataLayout() { return options.dataLayout; }
+
   /// Gets the LLVM representation of the index type. The returned type is an
   /// integer type with the size configured for this type converter.
   LLVM::LLVMType getIndexType();
 
   /// Gets the bitwidth of the index type when converted to LLVM.
   unsigned getIndexTypeBitwidth() { return options.indexBitwidth; }
 
   /// Gets the pointer bitwidth.

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h

 //===- ConvertStandardToLLVMPass.h - Pass entrypoint ------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_CONVERSION_STANDARDTOLLVM_CONVERTSTANDARDTOLLVMPASS_H_
 #define MLIR_CONVERSION_STANDARDTOLLVM_CONVERTSTANDARDTOLLVMPASS_H_
 
+#include "llvm/IR/DataLayout.h"
+
 #include <memory>
 
 namespace mlir {
 class LLVMTypeConverter;
 class ModuleOp;
 template <typename T>
 class OperationPass;
 class OwningRewritePatternList;
 
 /// Value to pass as bitwidth for the index type when the converter is expected
 /// to derive the bitwidth from the LLVM data layout.
 static constexpr unsigned kDeriveIndexBitwidthFromDataLayout = 0;
 
 /// Options to control the Standard dialect to LLVM lowering. The struct is used
 /// to share lowering options between passes, patterns, and type converter.
 struct LowerToLLVMOptions {
   bool useBarePtrCallConv = false;
   bool emitCWrappers = false;
   unsigned indexBitwidth = kDeriveIndexBitwidthFromDataLayout;
   /// Use aligned_alloc for heap allocations.
   bool useAlignedAlloc = false;
 
+  /// The data layout of the module to produce. This must be consistent with the
+  /// data layout used in the upper levels of the lowering pipeline.
+  // TODO: this should be replaced by MLIR data layout when one exists.
+  llvm::DataLayout dataLayout = llvm::DataLayout("");
+
   /// Get a statically allocated copy of the default LowerToLLVMOptions.
   static const LowerToLLVMOptions &getDefaultOptions() {
     static LowerToLLVMOptions options;
     return options;
   }
 };
 
 /// Collect a set of patterns to convert memory-related operations from the

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

 #define LLVMIR_OP_BASE
 
 include "mlir/IR/OpBase.td"
 
 def LLVM_Dialect : Dialect {
   let name = "llvm";
   let cppNamespace = "LLVM";
   let hasRegionArgAttrVerify = 1;
+  let hasOperationAttrVerify = 1;
   let extraClassDeclaration = [{
-    ~LLVMDialect();
-    const llvm::DataLayout &getDataLayout();
+    /// Name of the data layout attributes.
+    static StringRef getDataLayoutAttrName() { return "llvm.data_layout"; }
 
-  private:
-    friend LLVMType;
-
-    // This can't be a unique_ptr because the ctor is generated inline
-    // in the class definition at the moment.
-    detail::LLVMDialectImpl *impl;
+    /// Verifies if the given string is a well-formed data layout descriptor.
+    /// Uses `reportError` to report errors.
+    static LogicalResult verifyDataLayoutString(
+        StringRef descr, llvm::function_ref<void (const Twine &)> reportError);
   }];
 }
 
 // LLVM IR type wrapped in MLIR.
 def LLVM_Type : DialectType<LLVM_Dialect,
                             CPred<"$_self.isa<::mlir::LLVM::LLVMType>()">,
                             "LLVM dialect type">;
 

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 
 /// Create an LLVMTypeConverter using custom LowerToLLVMOptions.
 LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
                                      const LowerToLLVMOptions &options)
     : llvmDialect(ctx->getRegisteredDialect<LLVM::LLVMDialect>()),
       options(options) {
   assert(llvmDialect && "LLVM IR dialect is not registered");
   if (options.indexBitwidth == kDeriveIndexBitwidthFromDataLayout)
-    this->options.indexBitwidth =
-        llvmDialect->getDataLayout().getPointerSizeInBits();
+    this->options.indexBitwidth = options.dataLayout.getPointerSizeInBits();
 
   // Register conversions for the standard types.
   addConversion([&](ComplexType type) { return convertComplexType(type); });
   addConversion([&](FloatType type) { return convertFloatType(type); });
   addConversion([&](FunctionType type) { return convertFunctionType(type); });
   addConversion([&](IndexType type) { return convertIndexType(type); });
   addConversion([&](IntegerType type) { return convertIntegerType(type); });
   addConversion([&](MemRefType type) { return convertMemRefType(type); });
MLIRContext &LLVMTypeConverter::getContext() {
   return *getDialect()->getContext();
 }
 
 LLVM::LLVMType LLVMTypeConverter::getIndexType() {
   return LLVM::LLVMType::getIntNTy(&getContext(), getIndexTypeBitwidth());
 }
 
 unsigned LLVMTypeConverter::getPointerBitwidth(unsigned addressSpace) {
-  return llvmDialect->getDataLayout().getPointerSizeInBits(addressSpace);
+  return options.dataLayout.getPointerSizeInBits(addressSpace);
 }
 
 Type LLVMTypeConverter::convertIndexType(IndexType type) {
   return getIndexType();
 }
 
 Type LLVMTypeConverter::convertIntegerType(IntegerType type) {
   return LLVM::LLVMType::getIntNTy(&getContext(), type.getWidth());
SmallVector<Value, 4> LLVMTypeConverter::promoteOperands(Location loc,
   return promotedOperands;
 }
 
 namespace {
 /// A pass converting MLIR operations into the LLVM IR dialect.
 struct LLVMLoweringPass : public ConvertStandardToLLVMBase<LLVMLoweringPass> {
   LLVMLoweringPass() = default;
   LLVMLoweringPass(bool useBarePtrCallConv, bool emitCWrappers,
-                   unsigned indexBitwidth, bool useAlignedAlloc) {
+                   unsigned indexBitwidth, bool useAlignedAlloc,
+                   const llvm::DataLayout &dataLayout) {
     this->useBarePtrCallConv = useBarePtrCallConv;
     this->emitCWrappers = emitCWrappers;
     this->indexBitwidth = indexBitwidth;
     this->useAlignedAlloc = useAlignedAlloc;
+    this->dataLayout = dataLayout.getStringRepresentation();
   }
 
   /// Run the dialect converter on the module.
   void runOnOperation() override {
     if (useBarePtrCallConv && emitCWrappers) {
       getOperation().emitError()
           << "incompatible conversion options: bare-pointer calling convention "
              "and C wrapper emission";
       signalPassFailure();
       return;
     }
+    if (failed(LLVM::LLVMDialect::verifyDataLayoutString(
+            this->dataLayout, [this](const Twine &message) {
+              getOperation().emitError() << message.str();
+            }))) {
+      signalPassFailure();
+      return;
+    }
 
     ModuleOp m = getOperation();
 
     LowerToLLVMOptions options = {useBarePtrCallConv, emitCWrappers,
-                                  indexBitwidth, useAlignedAlloc};
+                                  indexBitwidth, useAlignedAlloc,
+                                  llvm::DataLayout(this->dataLayout)};
     LLVMTypeConverter typeConverter(&getContext(), options);
 
     OwningRewritePatternList patterns;
     populateStdToLLVMConversionPatterns(typeConverter, patterns);
 
     LLVMConversionTarget target(getContext());
     if (failed(applyPartialConversion(m, target, patterns)))
       signalPassFailure();
+    m.setAttr(LLVM::LLVMDialect::getDataLayoutAttrName(),
+              StringAttr::get(this->dataLayout, m.getContext()));
   }
 };
 } // end namespace
 
 mlir::LLVMConversionTarget::LLVMConversionTarget(MLIRContext &ctx)
     : ConversionTarget(ctx) {
   this->addLegalDialect<LLVM::LLVMDialect>();
   this->addIllegalOp<LLVM::DialectCastOp>();
   this->addIllegalOp<TanhOp>();
 }
 
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createLowerToLLVMPass(const LowerToLLVMOptions &options) {
   return std::make_unique<LLVMLoweringPass>(
       options.useBarePtrCallConv, options.emitCWrappers, options.indexBitwidth,
-      options.useAlignedAlloc);
+      options.useAlignedAlloc, options.dataLayout);
 }

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

 template <typename T>
 LogicalResult getMemRefAlignment(LLVMTypeConverter &typeConverter, T op,
                                  unsigned &align) {
   Type elementTy =
       typeConverter.convertType(op.getMemRefType().getElementType());
   if (!elementTy)
     return failure();
 
   // TODO: this should use the MLIR data layout when it becomes available and
   // stop depending on translation.
-  LLVM::LLVMDialect *dialect = typeConverter.getDialect();
   llvm::LLVMContext llvmContext;
   align = LLVM::TypeToLLVMIRTranslator(llvmContext)
               .getPreferredAlignment(elementTy.cast<LLVM::LLVMType>(),
-                                     dialect->getDataLayout());
+                                     typeConverter.getDataLayout());
   return success();
 }
 
 // Helper that returns the base address of a memref.
 LogicalResult getBase(ConversionPatternRewriter &rewriter, Location loc,
                       Value memref, MemRefType memRefType, Value &base) {
   // Inspect stride and offset structure.
   //

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

return op.emitOpError("can be given only acquire, release, acq_rel, "
                           "and seq_cst orderings");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // LLVMDialect initialization, type parsing, and registration.
 //===----------------------------------------------------------------------===//
 
-namespace mlir {
-namespace LLVM {
-namespace detail {
-struct LLVMDialectImpl {
-  LLVMDialectImpl() : layout("") {}
-
-  /// Default data layout to use.
-  // TODO: this should be moved to some Op equivalent to LLVM module and
-  // eventually replaced with a proper MLIR data layout.
-  llvm::DataLayout layout;
-};
-} // end namespace detail
-} // end namespace LLVM
-} // end namespace mlir
-
 void LLVMDialect::initialize() {
-  impl = new detail::LLVMDialectImpl();
   // clang-format off
   addTypes<LLVMVoidType,
            LLVMHalfType,
            LLVMBFloatType,
            LLVMFloatType,
            LLVMDoubleType,
            LLVMFP128Type,
            LLVMX86FP80Type,

 #define GET_OP_LIST
 #include "mlir/Dialect/LLVMIR/LLVMOps.cpp.inc"
       >();
 
   // Support unknown operations because not all LLVM operations are registered.
   allowUnknownOperations();
 }
 
-LLVMDialect::~LLVMDialect() { delete impl; }
-
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/LLVMOps.cpp.inc"
 
-const llvm::DataLayout &LLVMDialect::getDataLayout() { return impl->layout; }
-
 /// Parse a type registered to this dialect.
 Type LLVMDialect::parseType(DialectAsmParser &parser) const {
   return detail::parseType(parser);
 }
 
 /// Print a type registered to this dialect.
 void LLVMDialect::printType(Type type, DialectAsmPrinter &os) const {
   return detail::printType(type.cast<LLVMType>(), os);
 }
 
+LogicalResult LLVMDialect::verifyDataLayoutString(
+    StringRef descr, llvm::function_ref<void(const Twine &)> reportError) {
+  llvm::Expected<llvm::DataLayout> maybeDataLayout =
+      llvm::DataLayout::parse(descr);
+  if (maybeDataLayout)
+    return success();
+
+  std::string message;
+  llvm::raw_string_ostream messageStream(message);
+  llvm::logAllUnhandledErrors(maybeDataLayout.takeError(), messageStream);
+  reportError("invalid data layout descriptor: " + messageStream.str());
+  return failure();
+}
+
+/// Verify LLVM dialect attributes.
+LogicalResult LLVMDialect::verifyOperationAttribute(Operation *op,
+                                                    NamedAttribute attr) {
+  // If the data layout attribute is present, it must use the LLVM data layout
+  // syntax. Try parsing it and report errors in case of failure. Users of this
+  // attribute may assume it is well-formed and can pass it to the (asserting)
+  // llvm::DataLayout constructor.
+  if (attr.first.strref() != LLVM::LLVMDialect::getDataLayoutAttrName())
+    return success();
+  if (auto stringAttr = attr.second.dyn_cast<StringAttr>())
+    return verifyDataLayoutString(
+        stringAttr.getValue(),
+        [op](const Twine &message) { op->emitOpError() << message.str(); });
+
+  return op->emitOpError() << "expected '"
+                           << LLVM::LLVMDialect::getDataLayoutAttrName()
+                           << "' to be a string attribute";
+}
+
 /// Verify LLVMIR function argument attributes.
 LogicalResult LLVMDialect::verifyRegionArgAttribute(Operation *op,
                                                     unsigned regionIdx,
                                                     unsigned argIdx,
                                                     NamedAttribute argAttr) {
   // Check that llvm.noalias is a boolean attribute.
   if (argAttr.first == "llvm.noalias" && !argAttr.second.isa<BoolAttr>())
     return op->emitError()

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

for (Value v : values) {
     assert(valueMapping.count(v) && "referencing undefined value");
     remapped.push_back(valueMapping.lookup(v));
   }
   return remapped;
 }
 
 std::unique_ptr<llvm::Module> ModuleTranslation::prepareLLVMModule(
     Operation *m, llvm::LLVMContext &llvmContext, StringRef name) {
-  auto *dialect = m->getContext()->getRegisteredDialect<LLVM::LLVMDialect>();
-  assert(dialect && "LLVM dialect must be registered");
-
   auto llvmModule = std::make_unique<llvm::Module>(name, llvmContext);
-  llvmModule->setDataLayout(dialect->getDataLayout());
+
+  if (auto dataLayoutAttr =
+          m->getAttr(LLVM::LLVMDialect::getDataLayoutAttrName()))
+    llvmModule->setDataLayout(dataLayoutAttr.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;
 }

mlir/test/Conversion/StandardToLLVM/convert-data-layout.mlir

+// RUN: mlir-opt -convert-std-to-llvm %s | FileCheck %s
+// RUN-32: mlir-opt -convert-std-to-llvm='data-layout=p:32:32:32' %s | FileCheck %s
+
+// CHECK: module attributes {llvm.data_layout = ""}
+// CHECK-32: module attributes {llvm.data_layout ="p:32:32:32"}
+module {}

mlir/test/Dialect/LLVMIR/invalid.mlir

 }
 
 // -----
 
 func @invalid_ordering_in_fence() {
   // expected-error @+1 {{can be given only acquire, release, acq_rel, and seq_cst orderings}}
   llvm.fence syncscope("agent") monotonic
 }
+
+// -----
+
+// expected-error @+1 {{invalid data layout descriptor}}
+module attributes {llvm.data_layout = "#vjkr32"} {
+  func @invalid_data_layout()
+}

mlir/test/Target/llvmir.mlir

llvm.func @nontemoral_store_and_load() {
   // CHECK: !nontemporal ![[NODE:[0-9]+]]
   llvm.store %val, %0 {nontemporal} : !llvm.ptr<i32>
   // CHECK: !nontemporal ![[NODE]]
   %1 = llvm.load %0 {nontemporal} : !llvm.ptr<i32>
   llvm.return
 }
 
 // CHECK: ![[NODE]] = !{i32 1}
+
+// -----
+
+// Check that the translation does not crash in absence of a data layout.
+module {
+  // CHECK: declare void @module_default_layout
+  llvm.func @module_default_layout()
+}
+
+// -----
+
+// CHECK: target datalayout = "E"
+module attributes {llvm.data_layout = "E"} {
+  llvm.func @module_big_endian()
+}
+