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

[nicolasvasilache]

Summary

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.

[nicolasvasilache]

Update test

[nicolasvasilache]

Rebase

[mehdi_amini]

mlir/include/mlir/Conversion/Passes.td

348

Mention override existing one?

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

789

That's changing the behavior of the pass ConvertFuncToLLVM pass, might be desirable, but how is it OK to change the datalayout after the fact?

[nicolasvasilache]

Address comment

[nicolasvasilache]

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

789

The "change the datalayout after the fact" behavior is already present before this revision.
I would certainly want that cleaned up.

For now I am just trying to separate pieces, there is another cleanup of "untested therefore unused" datalayout string usage coming up.

[mehdi_amini]

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

789

One different I see is that before, the ConvertFuncToLLVM would use a data layout for performing the conversion, and then set it on the Module, while now the ConvertFuncToLLVM still takes a data layout option but it is lost in the IR.

[mehdi_amini]

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

789

(I guess I'm still trying to figure out the direction here)

[nicolasvasilache]

Update

[nicolasvasilache]

Update

[nicolasvasilache]

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

789

Good point, I reshuffled a bit so that the setting of the llvm.datalayout string attribute on the module is explicit.
convert-func-to-llvm can then lose its data-layout string option and recover it from the module, when present, to build LowerToLLVMOptions::dataLayout and at least have some consistency there.

This way, we at least won't need to make all xxx-to-llvm passes grow a data-layout string option to be consistent with convert-func-to-llvm.

Note however that consistency is a big word here:

1. no other LowerToLLVMOptions::dataLayout object instance is set up today outside of llvm.datalayout: the "data layout from string" is fundamentally inconsistent across xxx-to-llvm passes today.
2. the "MLIR DataLayout mechanism" is currently only used to set up LowerToLLVMOptions::indexBitwidth (but not LowerToLLVMOptions::dataLayout) (see here ). Index bitwidth at the module level is not very relevant as we already experience with GPU global / shared address spaces.
3. TypeConverter is the common place from which conversion patterns extract the llvm.datalayout (via an LowerToLLVMOptions::dataLayout); it also has a dataLayoutAnalysis so it combines both ways of specifying data layouts (LLVM's and MLIR's). Unfortunately, we know the TypeConverter across passes is very inconsistent already as per this discussion.

I believe this is a step towards cleaning some of this up.

[ftynse]

mlir/include/mlir/Conversion/Passes.td

343

Nit: I'd prefer explicitly mentioning llvm in the name here, e.g., set-module-llvm-datalayout.

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

753

Nit: we shouldn't need llvm:: here.

789

IIRC, the direction was to eventually stop using the LLVM string entirely in the conversions and have it constructed from the MLIR data layout. It can be either at conversion time or even at translation time to avoid mismatch. Don't remember why this has never happened, probably the lack of support for some of the non-size parts of the data layout that are available on GPU.

[mehdi_amini]

Works for me, if @ftynse is OK :)

[nicolasvasilache]

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

789

Yes and https://reviews.llvm.org/D157726 actually shows that the data layout itself is currently not used (or at least not tested) beyond being forwarded to LLVM.

[ftynse]

Okay for me.

[nicolasvasilache]

Address comments

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

 
 #include <memory>
 #include <string>
 
 namespace mlir {
 class Pass;
 
 #define GEN_PASS_DECL_CONVERTFUNCTOLLVMPASS
+#define GEN_PASS_DECL_SETMODULEDATALAYOUTPASS
 #include "mlir/Conversion/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_FUNCTOLLVM_CONVERTFUNCTOLLVMPASS_H_

mlir/include/mlir/Conversion/Passes.td

Option<"emulateLT32BitScalarTypes", "emulate-lt-32-bit-scalar-types",
            " the target">
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // FuncToLLVM
 //===----------------------------------------------------------------------===//
 
+def SetModuleDataLayoutPass : Pass<"set-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 = [{
     Convert Func dialect operations into the LLVM IR dialect operations.
 
     #### Input invariant
 
     -   no `tensor` types;
     -   all `vector` are one-dimensional;
     -   all blocks are reachable by following the successors of the first basic
         block;
 
     If other operations are present and their results are required by the LLVM
     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
     one-to-one. Function results are converted one-to-one and, in case more than
     1 value is returned, packed into an LLVM IR struct type. Function calls and
     returns are updated accordingly. Block argument types are updated to use
     LLVM IR types.
   }];
   let dependentDialects = ["LLVM::LLVMDialect"];
   let options = [
     Option<"useBarePtrCallConv", "use-bare-ptr-memref-call-conv", "bool",
            /*default=*/"false",
            "Replace FuncOp's MemRef arguments with bare pointers to the MemRef "
            "element types">,
     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">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // FuncToSPIRV

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

 #include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #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_SETMODULEDATALAYOUTPASS
 #include "mlir/Conversion/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
 
 #define PASS_NAME "convert-func-to-llvm"
 
 static constexpr StringRef varargsAttrName = "func.varargs";

 namespace {
 /// A pass converting Func operations into the LLVM IR dialect.
 struct ConvertFuncToLLVMPass
     : public impl::ConvertFuncToLLVMPassBase<ConvertFuncToLLVMPass> {
   using Base::Base;
 
   /// Run the dialect converter on the module.
   void runOnOperation() override {
+    ModuleOp m = getOperation();
+    StringRef dataLayout;
+    auto dataLayoutAttr = llvm::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,
                                     &dataLayoutAnalysis);
 
     RewritePatternSet patterns(&getContext());
     populateFuncToLLVMConversionPatterns(typeConverter, patterns);
 
     // TODO: Remove these in favor of their dedicated conversion passes.
     arith::populateArithToLLVMConversionPatterns(typeConverter, patterns);
     cf::populateControlFlowToLLVMConversionPatterns(typeConverter, patterns);
 
     LLVMConversionTarget target(getContext());
     if (failed(applyPartialConversion(m, target, std::move(patterns))))
       signalPassFailure();
+  }
+};
+
+struct SetModuleDataLayoutPass
+    : public impl::SetModuleDataLayoutPassBase<SetModuleDataLayoutPass> {
+  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));
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // ConvertToLLVMPatternInterface implementation

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

-// 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 -convert-func-to-llvm='use-opaque-pointers=1' -set-module-datalayout %s | FileCheck %s
+
+// RUN-32: mlir-opt -set-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"}
 module {}