[MLIR][mlir-link] Link required symbols even if they are only declarations

Author: xlauko

mlir/include/mlir/Linker/LinkerInterface.h

@@ -75,7 +75,7 @@ class SymbolLinkerInterface : public LinkerInterface<SymbolLinkerInterface> {
   virtual StringRef getSymbol(Operation *op) const = 0;
 
   /// Determines if an operation should be linked into the destination module.
-  virtual bool isLinkNeeded(ConflictPair pair) const = 0;
+  virtual bool isLinkNeeded(ConflictPair pair, bool forDependency) const = 0;
 
   /// Checks if an operation conflicts with existing linked operations.
   virtual ConflictPair findConflict(Operation *src) const = 0;
@@ -89,6 +89,9 @@ class SymbolLinkerInterface : public LinkerInterface<SymbolLinkerInterface> {
   /// Link the operations in the source module into the destination module.
   virtual Operation *materialize(ConflictPair pair, ModuleOp dst) const = 0;
 
+  /// Dependencies of the given operation required to be linked.
+  virtual SmallVector<Operation *> dependencies(Operation *op) const = 0;
+
   void setFlags(unsigned flags) { this->flags = flags; }
 
   bool shouldLinkOnlyNeeded() const {

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

@@ -184,7 +184,7 @@ class LLVMSymbolLinkerInterface : public SymbolLinkerInterface {
     return ConflictPair::noConflict(src);
   }
 
-  bool isLinkNeeded(ConflictPair pair) const override {
+  bool isLinkNeeded(ConflictPair pair, bool forDependency) const override {
     assert(canBeLinked(pair.src) && "expected linkable operation");
 
     LLVM::Linkage srcLinkage = getLinkage(pair.src);
@@ -194,19 +194,28 @@ class LLVMSymbolLinkerInterface : public SymbolLinkerInterface {
       return true;
     }
 
-    if (shouldLinkOnlyNeeded()) {
-      // Don't import globals that are already defined
-      if (!pair.dst || !isDeclaration(pair.dst))
-        return false;
-    }
+    bool alreadyDeclared = pair.dst && isDeclaration(pair.dst);
+
+    // Don't import globals that are already defined
+    if (shouldLinkOnlyNeeded() && !alreadyDeclared)
+      return false;
+
+    // Always import dependencies that are not yet defined or declared
+    if (forDependency && !pair.dst)
+      return true;
+
     if (isDeclaration(pair.src))
       return false;
 
-    bool keepOnlyInSource = isLocalLinkage(srcLinkage) ||
-                            isLinkOnceLinkage(srcLinkage) ||
-                            isAvailableExternallyLinkage(srcLinkage);
+    if (pair.hasConflict())
+      return true;
 
-    return pair.dst || shouldOverrideFromSrc() || !keepOnlyInSource;
+    if (shouldOverrideFromSrc())
+      return true;
+
+    // linkage specifies to keep operation only in source
+    return !(isLocalLinkage(srcLinkage) || isLinkOnceLinkage(srcLinkage) ||
+             isAvailableExternallyLinkage(srcLinkage));
   }
 
   FailureOr<bool> shouldLinkFromSource(ConflictPair pair) const {
@@ -385,6 +394,22 @@ class LLVMSymbolLinkerInterface : public SymbolLinkerInterface {
     return pair.dst;
   }
 
+  SmallVector<Operation *> dependencies(Operation *op) const override {
+    SmallVector<Operation *> result;
+
+    Operation *symbolTableOp = symbolTable->getOp();
+    op->walk([&](SymbolUserOpInterface user) {
+      if (user.getOperation() == op)
+        return;
+
+      if (SymbolRefAttr symbol = user.getUserSymbol())
+        if (Operation *dep = symbolTable->lookupSymbolIn(symbolTableOp, symbol))
+          result.push_back(dep);
+    });
+
+    return result;
+  }
+
 private:
   llvm::StringMap<Operation *> summary;
 

mlir/lib/IR/BuiltinLinkerInterface.cpp

@@ -35,33 +35,44 @@ class BuiltinLinkerInterface : public ModuleLinkerInterface {
       if (op == src)
         return WalkResult::advance();
 
-      auto linker = dyn_cast<SymbolLinkerInterface>(op->getDialect());
-      if (!linker)
-        return WalkResult::advance();
+      if (summarize(op, flags, /*forDependency=*/false).failed())
+        return WalkResult::interrupt();
+      return WalkResult::advance();
+    });
 
-      // TODO do this in init
-      linker->setFlags(flags);
+    return failure(result.wasInterrupted());
+  }
 
-      if (!linker->canBeLinked(op))
-        return WalkResult::advance();
+  LogicalResult summarize(Operation *op, unsigned flags, bool forDependency) {
+    auto linker = dyn_cast<SymbolLinkerInterface>(op->getDialect());
+    if (!linker)
+      return success();
 
-      ConflictPair conflict = linker->findConflict(op);
-      if (!linker->isLinkNeeded(conflict))
-        return WalkResult::advance();
+    linker->setFlags(flags);
+
+    if (!linker->canBeLinked(op))
+      return success();
 
-      if (conflict.hasConflict())
-        return failed(linker->resolveConflict(conflict))
-                   ? WalkResult::interrupt()
-                   : WalkResult::advance();
+    ConflictPair conflict = linker->findConflict(op);
+    if (!linker->isLinkNeeded(conflict, forDependency))
+      return success();
 
-      // TODO rename: registerForLink
+    if (conflict.hasConflict()) {
+      if (linker->resolveConflict(conflict).failed())
+        return failure();
+    } else {
       linker->registerForLink(op);
-      return WalkResult::advance();
-    });
+    }
 
-    // TODO deal with references
+    if (forDependency)
+      return success();
 
-    return failure(result.wasInterrupted());
+    for (Operation *dep : linker->dependencies(op)) {
+      if (summarize(dep, flags, /*forDependency=*/true).failed())
+        return failure();
+    }
+
+    return success();
   }
 
   LogicalResult link(ModuleOp dst) const override {