[MLIR][mlir-link] Share symbol table information between symbol linkers

Author: Jezurko

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

@@ -30,7 +30,8 @@ class LLVMSymbolLinkerInterface
   static uint32_t getAddressSpace(Operation *op);
   StringRef getSymbol(Operation *op) const override;
   Operation *materialize(Operation *src, link::LinkState &state) const override;
-  SmallVector<Operation *> dependencies(Operation *op) const override;
+  SmallVector<Operation *>
+  dependencies(Operation *op, SymbolTableCollection &collection) const override;
   LogicalResult initialize(ModuleOp src) override;
   LogicalResult finalize(ModuleOp dst) const override;
   Operation *appendGlobals(llvm::StringRef glob, link::LinkState &state);

mlir/include/mlir/Linker/LinkerInterface.h

@@ -144,7 +144,8 @@ class SymbolLinkerInterface : public LinkerInterface<SymbolLinkerInterface> {
   }
 
   /// Dependencies of the given operation required to be linked.
-  virtual SmallVector<Operation *> dependencies(Operation *op) const = 0;
+  virtual SmallVector<Operation *>
+  dependencies(Operation *op, SymbolTableCollection &collection) const = 0;
 
   void setFlags(unsigned flags) { this->flags = flags; }
 
@@ -199,7 +200,8 @@ class SymbolAttrLinkerInterface : public SymbolLinkerInterface {
   virtual LogicalResult verifyLinkageCompatibility(Conflict pair) const = 0;
 
   /// Dependencies of the given operation required to be linked.
-  SmallVector<Operation *> dependencies(Operation *op) const override;
+  SmallVector<Operation *>
+  dependencies(Operation *op, SymbolTableCollection &collection) const override;
 
 protected:
   // Operations that are to be linked with the original name.

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

@@ -256,10 +256,10 @@ LLVM::LLVMSymbolLinkerInterface::materialize(Operation *src,
   return SymbolAttrLinkerInterface::materialize(src, state);
 }
 
-SmallVector<Operation *>
-LLVM::LLVMSymbolLinkerInterface::dependencies(Operation *op) const {
+SmallVector<Operation *> LLVM::LLVMSymbolLinkerInterface::dependencies(
+    Operation *op, SymbolTableCollection &collection) const {
   Operation *module = op->getParentOfType<ModuleOp>();
-  SymbolTable st(module);
+  SymbolTable &st = collection.getSymbolTable(module);
   SmallVector<Operation *> result;
 
   auto insertDepIfExists = [&](auto symbolRef) -> void {

mlir/lib/IR/BuiltinLinkerInterface.cpp

@@ -68,7 +68,7 @@ class BuiltinLinkerInterface : public ModuleLinkerInterface {
       linker->registerForLink(op);
     }
 
-    for (Operation *dep : linker->dependencies(op)) {
+    for (Operation *dep : linker->dependencies(op, symbolTableCollection)) {
       if (summarize(dep, flags, /*forDependency=*/true).failed())
         return failure();
     }
@@ -87,6 +87,7 @@ class BuiltinLinkerInterface : public ModuleLinkerInterface {
 
 private:
   SymbolLinkerInterfaces symbolLinkers;
+  SymbolTableCollection symbolTableCollection;
 };
 
 //===----------------------------------------------------------------------===//

mlir/lib/Linker/LinkerInterface.cpp

@@ -161,11 +161,11 @@ LogicalResult SymbolAttrLinkerInterface::link(LinkState &state) const {
   return success();
 }
 
-SmallVector<Operation *>
-SymbolAttrLinkerInterface::dependencies(Operation *op) const {
+SmallVector<Operation *> SymbolAttrLinkerInterface::dependencies(
+    Operation *op, SymbolTableCollection &collection) const {
   // TODO: use something like SymbolTableAnalysis
   Operation *module = op->getParentOfType<ModuleOp>();
-  SymbolTable st(module);
+  SymbolTable &st = collection.getSymbolTable(module);
   SmallVector<Operation *> result;
   op->walk([&](SymbolUserOpInterface user) {
     if (user.getOperation() == op)