[mlir][Func] Extract datalayout string attribute setting as a separate module pass

FuncToLLVM uses the data layout string attribute in 3 different ways:
1. LowerToLLVMOptions options(&getContext(), getAnalysis<DataLayoutAnalysis>().getAtOrAbove(m));
2. options.dataLayout = llvm::DataLayout(this->dataLayout);
3. m->setAttr(..., this->dataLayout));

The 3rd way is unrelated to the other 2 and occurs after conversion, making it confusing.
This revision separates this post-hoc module annotation functionality into its own pass.
The convert-func-to-llvm pass loses its `data-layout` option and instead recovers it from
the `llvm.data_layout` attribute attached to the module, when present.

In the future, `LowerToLLVMOptions options(&getContext(), getAnalysis<DataLayoutAnalysis>().getAtOrAbove(m))` and
`options.dataLayout = llvm::DataLayout(dataLayout);` should be unified.

Reviewed By: ftynse, mehdi_amini

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

Author: nicolasvasilache

mlir/include/mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h

@@ -16,6 +16,7 @@ namespace mlir {
 class Pass;
 
 #define GEN_PASS_DECL_CONVERTFUNCTOLLVMPASS
+#define GEN_PASS_DECL_SETLLVMMODULEDATALAYOUTPASS
 #include "mlir/Conversion/Passes.h.inc"
 
 } // namespace mlir

mlir/include/mlir/Conversion/Passes.td

@@ -341,6 +341,21 @@ def ConvertControlFlowToSPIRV : Pass<"convert-cf-to-spirv"> {
 // FuncToLLVM
 //===----------------------------------------------------------------------===//
 
+def SetLLVMModuleDataLayoutPass : Pass<"set-llvm-module-datalayout", "ModuleOp"> {
+  let summary = "Attach a datalayout string as a module attribute";
+  let description = [{
+    Verify that the dataLayout string is a valid LLVM datalayout string and
+    attach it as an attribute `LLVMDialect::getDataLayoutAttrName()` to the
+    module, overriding the existing one.
+  }];
+  let options = [
+    Option<"dataLayout", "data-layout", "std::string",
+           /*default=*/"\"\"",
+           "String description (LLVM format) of the data layout that is "
+           "expected on the produced module">,
+  ];
+}
+
 def ConvertFuncToLLVMPass : Pass<"convert-func-to-llvm", "ModuleOp"> {
   let summary = "Convert from the Func dialect to the LLVM dialect";
   let description = [{
@@ -357,6 +372,11 @@ def ConvertFuncToLLVMPass : Pass<"convert-func-to-llvm", "ModuleOp"> {
     IR dialect operations, the pass will fail.  Any LLVM IR operations or types
     already present in the IR will be kept as is.
 
+    An LLVM datalayout string can be attached as an attribute to the module on
+    which the pass anchors. Such an attribute is attached by calling the 
+    set-module-datalayout pass. If present, an llvm::DataLayout object is 
+    created from this attribute and used in the conversion to LLVM.
+
     #### Output IR
 
     Functions converted to LLVM IR. Function arguments types are converted
@@ -374,10 +394,6 @@ def ConvertFuncToLLVMPass : Pass<"convert-func-to-llvm", "ModuleOp"> {
     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">,
     Option<"useOpaquePointers", "use-opaque-pointers", "bool",
                        /*default=*/"true", "Generate LLVM IR using opaque pointers "
                        "instead of typed pointers">,

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

@@ -43,13 +43,15 @@
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Type.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FormatVariadic.h"
 #include <algorithm>
 #include <functional>
 
 namespace mlir {
 #define GEN_PASS_DEF_CONVERTFUNCTOLLVMPASS
+#define GEN_PASS_DEF_SETLLVMMODULEDATALAYOUTPASS
 #include "mlir/Conversion/Passes.h.inc"
 } // namespace mlir
 
@@ -746,23 +748,29 @@ struct ConvertFuncToLLVMPass
 
   /// Run the dialect converter on the module.
   void runOnOperation() override {
+    ModuleOp m = getOperation();
+    StringRef dataLayout;
+    auto dataLayoutAttr = dyn_cast_or_null<StringAttr>(
+        m->getAttr(LLVM::LLVMDialect::getDataLayoutAttrName()));
+    if (dataLayoutAttr)
+      dataLayout = dataLayoutAttr.getValue();
+
     if (failed(LLVM::LLVMDialect::verifyDataLayoutString(
-            this->dataLayout, [this](const Twine &message) {
+            dataLayout, [this](const Twine &message) {
               getOperation().emitError() << message.str();
             }))) {
       signalPassFailure();
       return;
     }
 
-    ModuleOp m = getOperation();
     const auto &dataLayoutAnalysis = getAnalysis<DataLayoutAnalysis>();
 
     LowerToLLVMOptions options(&getContext(),
                                dataLayoutAnalysis.getAtOrAbove(m));
     options.useBarePtrCallConv = useBarePtrCallConv;
     if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
       options.overrideIndexBitwidth(indexBitwidth);
-    options.dataLayout = llvm::DataLayout(this->dataLayout);
+    options.dataLayout = llvm::DataLayout(dataLayout);
     options.useOpaquePointers = useOpaquePointers;
 
     LLVMTypeConverter typeConverter(&getContext(), options,
@@ -778,7 +786,24 @@ struct ConvertFuncToLLVMPass
     LLVMConversionTarget target(getContext());
     if (failed(applyPartialConversion(m, target, std::move(patterns))))
       signalPassFailure();
+  }
+};
+
+struct SetLLVMModuleDataLayoutPass
+    : public impl::SetLLVMModuleDataLayoutPassBase<
+          SetLLVMModuleDataLayoutPass> {
+  using Base::Base;
 
+  /// Run the dialect converter on the module.
+  void runOnOperation() override {
+    if (failed(LLVM::LLVMDialect::verifyDataLayoutString(
+            this->dataLayout, [this](const Twine &message) {
+              getOperation().emitError() << message.str();
+            }))) {
+      signalPassFailure();
+      return;
+    }
+    ModuleOp m = getOperation();
     m->setAttr(LLVM::LLVMDialect::getDataLayoutAttrName(),
                StringAttr::get(m.getContext(), this->dataLayout));
   }

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

@@ -1,5 +1,7 @@
-// RUN: mlir-opt -convert-func-to-llvm='use-opaque-pointers=1' %s | FileCheck %s
-// RUN-32: mlir-opt -convert-func-to-llvm='data-layout=p:32:32:32 use-opaque-pointers=1' %s | FileCheck %s
+// RUN: mlir-opt -set-llvm-module-datalayout -convert-func-to-llvm='use-opaque-pointers=1' %s | FileCheck %s
+
+// RUN-32: mlir-opt -set-llvm-module-datalayout='data-layout=p:32:32:32' -convert-func-to-llvm='use-opaque-pointers=1' %s \
+// RUN-32: | FileCheck %s
 
 // CHECK: module attributes {llvm.data_layout = ""}
 // CHECK-32: module attributes {llvm.data_layout ="p:32:32:32"}