[MLIR][mlir-link] Allow togling linkOnlyNeeded per module (#56)

Author: Jezurko

mlir/include/mlir/Linker/LLVMLinkerMixin.h

@@ -183,8 +183,8 @@ class LLVMLinkerMixin {
     bool alreadyDeclared =
         alreadyDefinedOrDeclared && derived.isDeclaration(pair.dst);
 
-    // Don't import globals that are already declared
-    if (derived.shouldLinkOnlyNeeded() && !alreadyDeclared)
+    // Don't import globals that are already defined
+    if (derived.shouldLinkOnlyNeeded() && !alreadyDeclared && !forDependency)
       return false;
 
     // Private dependencies are gonna be renamed and linked

mlir/include/mlir/Linker/Linker.h

@@ -59,9 +59,10 @@ class Linker {
   Linker(MLIRContext *context, const LinkerOptions &options = {})
       : context(context), options(options) {}
 
-  /// Add a module to be linked
   LogicalResult addModule(OwningOpRef<ModuleOp> src);
 
+  LogicalResult addModule(OwningOpRef<ModuleOp> src, bool onlyNeeded);
+
   /// Perform linking and materialization in the destination module.
   /// Returns the linked module.
   OwningOpRef<ModuleOp> link(bool sortSymbols = false);
@@ -71,16 +72,19 @@ class Linker {
   LogicalResult emitFileError(const Twine &fileName, const Twine &message);
   LogicalResult emitError(const Twine &message);
 
+  /// Return the flags controlling the linker behavior for the current module
+  unsigned getFlags() const;
+
 private:
   /// Setup the linker based on the first module
   LogicalResult initializeLinker(ModuleOp src);
 
+  /// Add a module to be linked
+  LogicalResult addModule(OwningOpRef<ModuleOp> src, unsigned flags);
+
   /// Obtain the linker interface for the given module
   ModuleLinkerInterface *getModuleLinkerInterface(ModuleOp op);
 
-  /// Return the flags controlling the linker behavior for the current module
-  unsigned getFlags() const;
-
   /// Preprocess the given module before linking with the given flags
   LogicalResult summarize(ModuleOp src, unsigned flags);
 

mlir/lib/Linker/Linker.cpp

@@ -40,8 +40,19 @@ ModuleLinkerInterface *Linker::getModuleLinkerInterface(ModuleOp op) {
 }
 
 LogicalResult Linker::addModule(OwningOpRef<ModuleOp> src) {
+  return addModule(std::move(src), getFlags());
+}
+
+LogicalResult Linker::addModule(OwningOpRef<ModuleOp> src, bool onlyNeeded) {
   unsigned flags = getFlags();
+  if (onlyNeeded)
+    flags |= LinkerFlags::LinkOnlyNeeded;
+  else
+    flags &= ~LinkerFlags::LinkOnlyNeeded;
+  return addModule(std::move(src), flags);
+}
 
+LogicalResult Linker::addModule(OwningOpRef<ModuleOp> src, unsigned flags) {
   ModuleOp mod = [&] {
     if (options.shouldKeepModulesAlive()) {
       modules.push_back(std::move(src));