[mlir][mlir-link] Copy in datalayout in llvm dialect linker (#77)

* [mlir][mlir-link] Copy in datalayout in llvm dialect linker

Author: Jezurko

mlir/include/mlir/Dialect/LLVMIR/LLVMLinkerInterface.h

@@ -31,6 +31,8 @@ class LLVMSymbolLinkerInterface
   StringRef getSymbol(Operation *op) const override;
   Operation *materialize(Operation *src, link::LinkState &state) const override;
   SmallVector<Operation *> dependencies(Operation *op) const override;
+  LogicalResult initialize(ModuleOp src) override;
+  LogicalResult finalize(ModuleOp dst) const override;
   Operation *appendGlobals(llvm::StringRef glob, link::LinkState &state);
 
   template <typename structor_t>
@@ -118,6 +120,10 @@ class LLVMSymbolLinkerInterface
     structor.setDataAttr(newDataAttr);
     return cloned;
   }
+
+private:
+  DataLayoutSpecInterface dtla = {};
+  TargetSystemSpecInterface targetSys = {};
 };
 
 } // namespace LLVM

mlir/include/mlir/Linker/LinkerInterface.h

@@ -87,9 +87,12 @@ class LinkerInterface : public DialectInterface::Base<ConcreteType> {
   LinkerInterface(Dialect *dialect)
       : DialectInterface::Base<ConcreteType>(dialect) {}
 
-  /// Runs initialization of alinker before summarization for the given module
+  /// Runs initialization of a linker before summarization for the given module
   virtual LogicalResult initialize(ModuleOp src) { return success(); }
 
+  /// Runs finalization of a linker after linking for the given module
+  virtual LogicalResult finalize(ModuleOp dst) const { return success(); }
+
   /// Link operations from current summary using state builder
   virtual LogicalResult link(LinkState &state) const = 0;
 };
@@ -255,6 +258,14 @@ class SymbolLinkerInterfaces {
     return success();
   }
 
+  LogicalResult finalize(ModuleOp dst) const {
+    for (SymbolLinkerInterface *linker : interfaces) {
+      if (failed(linker->finalize(dst)))
+        return failure();
+    }
+    return success();
+  }
+
   Conflict findConflict(Operation *src) const {
     for (SymbolLinkerInterface *linker : interfaces) {
       if (auto pair = linker->findConflict(src); pair.hasConflict())

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

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/LLVMIR/LLVMLinkerInterface.h"
+#include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Linker/LLVMLinkerMixin.h"
 #include "mlir/Linker/LinkerInterface.h"
@@ -309,6 +310,23 @@ LLVM::LLVMSymbolLinkerInterface::dependencies(Operation *op) const {
   return result;
 }
 
+LogicalResult LLVM::LLVMSymbolLinkerInterface::initialize(ModuleOp src) {
+  dtla = src.getDataLayoutSpec();
+  targetSys = src.getTargetSystemSpec();
+  return success();
+}
+
+LogicalResult LLVM::LLVMSymbolLinkerInterface::finalize(ModuleOp dst) const {
+  SmallVector<NamedAttribute, 2> newAttrs;
+  // The names are currently hardcoded for dlti dialect
+  // Nice solution would be preferable
+  if (dtla)
+    dst->setAttr(DataLayoutSpecAttr::name, dyn_cast<Attribute>(dtla));
+  if (targetSys)
+    dst->setAttr(TargetSystemSpecAttr::name, dyn_cast<Attribute>(targetSys));
+  return success();
+}
+
 static std::pair<Attribute, Type>
 getAppendedArrayAttr(llvm::ArrayRef<mlir::Operation *> globs,
                      LinkState &state) {

mlir/lib/IR/BuiltinLinkerInterface.cpp

@@ -30,6 +30,10 @@ class BuiltinLinkerInterface : public ModuleLinkerInterface {
     return symbolLinkers.initialize(src);
   }
 
+  LogicalResult finalize(ModuleOp dst) const override {
+    return symbolLinkers.finalize(dst);
+  }
+
   LogicalResult summarize(ModuleOp src, unsigned flags) override {
     WalkResult result = src.walk([&](Operation *op) {
       if (op == src)

mlir/lib/Linker/Linker.cpp

@@ -109,6 +109,10 @@ OwningOpRef<ModuleOp> Linker::link(bool sortSymbols) {
       symbol->moveBefore(&mod.front());
     }
   }
+  ModuleLinkerInterface *iface = getModuleLinkerInterface(composite.get());
+
+  if (failed(iface->finalize(composite.get())))
+    return nullptr;
 
   return std::move(composite);
 }