[MLIR][mlir-link] Add hook to perform per Module computation before summary

This allows the linker interface to pre-compute some information that might be necessary for the linking.

Author: Jezurko

mlir/include/mlir/Linker/LinkerInterface.h

@@ -18,6 +18,7 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/DialectInterface.h"
 #include "mlir/IR/IRMapping.h"
+#include "mlir/IR/Threading.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/Error.h"
 #include <memory>
@@ -149,6 +150,12 @@ class SymbolLinkerInterface : public LinkerInterface<SymbolLinkerInterface> {
     return state.clone(src);
   }
 
+  /// Perform tasks that need to be computed on whole-module basis before actual summary.
+  /// E.g. Pre-compute COMDAT resolution before actually linking the modules.
+  virtual LogicalResult moduleOpSummary(ModuleOp module) {
+    return success();
+  }
+
   /// Dependencies of the given operation required to be linked.
   virtual SmallVector<Operation *>
   dependencies(Operation *op, SymbolTableCollection &collection) const = 0;
@@ -286,6 +293,13 @@ class SymbolLinkerInterfaces {
     return Conflict::noConflict(src);
   }
 
+  LogicalResult moduleOpSummary(ModuleOp src) {
+    return failableParallelForEach(src.getContext(), interfaces,
+                                   [src](SymbolLinkerInterface *linker) {
+      return linker->moduleOpSummary(src);
+    });
+  }
+
 private:
   SetVector<SymbolLinkerInterface *> interfaces;
 };

mlir/lib/IR/BuiltinLinkerInterface.cpp

@@ -39,11 +39,20 @@ class BuiltinLinkerInterface : public ModuleLinkerInterface {
                           SymbolTableCollection &collection) override {
     // Collect all operations to process in parallel
     SmallVector<Operation *> ops;
-    src.walk([&](Operation *op) {
-      if (op != src)
-        ops.push_back(op);
+    WalkResult result = src.walk([&](Operation *op) {
+      if (op == src) {
+        if (symbolLinkers.moduleOpSummary(src).failed())
+          return WalkResult::interrupt();
+        return WalkResult::advance();
+      }
+      ops.push_back(op);
+      return WalkResult::advance();
     });
 
+    if (result.wasInterrupted()) {
+      return failure();
+    }
+
     // Process operations in parallel
     return failableParallelForEach(
         src.getContext(), ops, [&](Operation *op) {