[MLIR][mlir-link] Make linker interfaces multithreaded

Author: frabert

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

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

mlir/include/mlir/Linker/LLVMLinkerMixin.h

@@ -155,6 +155,10 @@ class LLVMLinkerMixin {
     return static_cast<const DerivedLinkerInterface &>(*this);
   }
 
+  DerivedLinkerInterface &getDerived() {
+    return static_cast<DerivedLinkerInterface &>(*this);
+  }
+
 public:
   bool isDeclarationForLinker(Operation *op) const {
     const DerivedLinkerInterface &derived = getDerived();

mlir/include/mlir/Linker/Linker.h

@@ -12,6 +12,7 @@
 #include "mlir/IR/BuiltinOps.h"
 
 #include "mlir/Linker/LinkerInterface.h"
+#include <mutex>
 namespace mlir::link {
 
 /// These are gathered alphabetically sorted linker options
@@ -101,6 +102,9 @@ class Linker {
 
   /// Modules registry used if `keepModulesAlive` is true
   std::vector<OwningOpRef<ModuleOp>> modules;
+
+  /// Mutex to protect modules vector and initialization during parallel addModule calls
+  std::mutex linkerMutex;
 };
 
 } // namespace mlir::link

mlir/include/mlir/Linker/LinkerInterface.h

@@ -21,6 +21,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/Error.h"
 #include <memory>
+#include <mutex>
 
 namespace mlir::link {
 
@@ -37,8 +38,9 @@ enum LinkerFlags {
 class LinkState {
 public:
   LinkState(ModuleOp dst, mlir::SymbolTableCollection &symbolTableCollection)
-      : mapping(std::make_shared<IRMapping>()), builder(dst.getBodyRegion()),
-        symbolTableCollection(symbolTableCollection), moduleMaps() {}
+      : mapping(std::make_shared<IRMapping>()), mutex(std::make_shared<std::mutex>()),
+        builder(dst.getBodyRegion()), symbolTableCollection(symbolTableCollection),
+        moduleMaps() {}
 
   Operation *clone(Operation *src);
   Operation *cloneWithoutRegions(Operation *src);
@@ -49,7 +51,7 @@ class LinkState {
 
   LinkState nest(ModuleOp submod) const;
 
-  IRMapping &getMapping();
+  std::pair<IRMapping &, std::mutex &> getMapping();
   SymbolTableCollection &getSymbolTableCollection() {
     return symbolTableCollection;
   }
@@ -63,11 +65,14 @@ class LinkState {
 private:
   // Private constructor used by nest()
   LinkState(ModuleOp dst, std::shared_ptr<IRMapping> mapping,
+            std::shared_ptr<std::mutex> mutex,
             SymbolTableCollection &symbolTableCollection)
-      : mapping(std::move(mapping)), builder(dst.getBodyRegion()),
+      : mapping(std::move(mapping)), mutex(std::move(mutex)),
+        builder(dst.getBodyRegion()),
         symbolTableCollection(symbolTableCollection), moduleMaps() {}
 
   std::shared_ptr<IRMapping> mapping;
+  std::shared_ptr<std::mutex> mutex;
   OpBuilder builder;
   SymbolTableCollection &symbolTableCollection;
   DenseMap<ModuleOp, SymbolUserMap> moduleMaps;
@@ -95,7 +100,7 @@ class LinkerInterface : public DialectInterface::Base<ConcreteType> {
   virtual LogicalResult finalize(ModuleOp dst) const { return success(); }
 
   /// Link operations from current summary using state builder
-  virtual LogicalResult link(LinkState &state) const = 0;
+  virtual LogicalResult link(LinkState &state) = 0;
 };
 
 //===----------------------------------------------------------------------===//
@@ -140,7 +145,7 @@ class SymbolLinkerInterface : public LinkerInterface<SymbolLinkerInterface> {
   virtual void registerForLink(Operation *op, SymbolTableCollection &collection) = 0;
 
   /// Materialize new operation for the given conflict src operation.
-  virtual Operation *materialize(Operation *src, LinkState &state) const {
+  virtual Operation *materialize(Operation *src, LinkState &state) {
     return state.clone(src);
   }
 
@@ -178,7 +183,7 @@ class SymbolAttrLinkerInterface : public SymbolLinkerInterface {
   using SymbolLinkerInterface::SymbolLinkerInterface;
 
   /// Link operations from current summary using state builder
-  LogicalResult link(LinkState &state) const override;
+  LogicalResult link(LinkState &state) override;
 
   /// Returns the symbol for the given operation.
   StringRef getSymbol(Operation *op) const override;
@@ -211,6 +216,9 @@ class SymbolAttrLinkerInterface : public SymbolLinkerInterface {
 
   // Operations that are to be linked with unique names.
   SetVector<Operation *> uniqued;
+
+  // Mutex to protect summary and uniqued during parallel summarization.
+  mutable std::mutex summaryMutex;
 };
 
 //===----------------------------------------------------------------------===//

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

@@ -246,8 +246,8 @@ StringRef LLVM::LLVMSymbolLinkerInterface::getSymbol(Operation *op) const {
 
 Operation *
 LLVM::LLVMSymbolLinkerInterface::materialize(Operation *src,
-                                             LinkState &state) const {
-  auto derived = LinkerMixin::getDerived();
+                                             LinkState &state) {
+  auto &derived = LinkerMixin::getDerived();
   // empty append means that we either have single module or that something went
   // wrong
   if (isAppendingLinkage(derived.getLinkage(src)) && !append.empty()) {
@@ -427,7 +427,7 @@ getAppendedOpWithInitRegion(llvm::ArrayRef<mlir::Operation *> globs,
   auto elemType = originalType.getElementType();
   size_t elemCount = 0;
 
-  IRMapping &mapping = state.getMapping();
+  auto [mapping, mutex] = state.getMapping();
   auto builder = OpBuilder(targetRegion);
   std::vector<Value> values;
   std::vector<std::vector<int64_t>> positions;
@@ -442,8 +442,12 @@ getAppendedOpWithInitRegion(llvm::ArrayRef<mlir::Operation *> globs,
       if (isa<LLVM::UndefOp, LLVM::ReturnOp, LLVM::InsertValueOp>(op))
         continue;
 
-      Operation *cloned = builder.clone(op, mapping);
-      mapping.map(&op, cloned);
+      Operation *cloned;
+      {
+        std::lock_guard<std::mutex> lock(mutex);
+        cloned = builder.clone(op, mapping);
+        mapping.map(&op, cloned);
+      }
       // LLVM dialect does not have multiple result operations
       // zero result operation should not appear in this context
       // unless its a return which we skip

mlir/lib/IR/BuiltinLinkerInterface.cpp

@@ -12,6 +12,7 @@
 
 #include "mlir/IR/BuiltinLinkerInterface.h"
 #include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/Threading.h"
 #include "mlir/Linker/LinkerInterface.h"
 
 using namespace mlir;
@@ -36,16 +37,18 @@ class BuiltinLinkerInterface : public ModuleLinkerInterface {
 
   LogicalResult summarize(ModuleOp src, unsigned flags,
                           SymbolTableCollection &collection) override {
-    WalkResult result = src.walk([&](Operation *op) {
-      if (op == src)
-        return WalkResult::advance();
-
-      if (summarize(op, flags, /*forDependency=*/false, collection).failed())
-        return WalkResult::interrupt();
-      return WalkResult::advance();
+    // Collect all operations to process in parallel
+    SmallVector<Operation *> ops;
+    src.walk([&](Operation *op) {
+      if (op != src)
+        ops.push_back(op);
     });
 
-    return failure(result.wasInterrupted());
+    // Process operations in parallel
+    return failableParallelForEach(
+        src.getContext(), ops, [&](Operation *op) {
+          return summarize(op, flags, /*forDependency=*/false, collection);
+        });
   }
 
   LogicalResult summarize(Operation *op, unsigned flags, bool forDependency,
@@ -70,15 +73,15 @@ class BuiltinLinkerInterface : public ModuleLinkerInterface {
       linker->registerForLink(op, collection);
     }
 
-    for (Operation *dep : linker->dependencies(op, symbolTableCollection)) {
-      if (summarize(dep, flags, /*forDependency=*/true, collection).failed())
-        return failure();
-    }
+    SmallVector<Operation *> deps = linker->dependencies(op, symbolTableCollection);
+    auto res = failableParallelForEach(getContext(), deps, [&](Operation *dep) {
+      return summarize(dep, flags, /*forDependency=*/true, collection);
+    });
 
-    return success();
+    return res;
   }
 
-  LogicalResult link(LinkState &state) const override {
+  LogicalResult link(LinkState &state) override {
     return symbolLinkers.link(state);
   }
 

mlir/lib/Linker/Linker.cpp

@@ -53,21 +53,26 @@ LogicalResult Linker::addModule(OwningOpRef<ModuleOp> src, bool onlyNeeded) {
 }
 
 LogicalResult Linker::addModule(OwningOpRef<ModuleOp> src, unsigned flags) {
-  ModuleOp mod = [&] {
+  ModuleOp mod;
+
+  {
+    std::lock_guard<std::mutex> lock(linkerMutex);
+
     if (options.shouldKeepModulesAlive()) {
       modules.push_back(std::move(src));
-      return modules.back().get();
+      mod = modules.back().get();
+    } else {
+      mod = src.get();
     }
-    return src.get();
-  }();
 
-  // If this is the first module, setup the linker based on it
-  if (!composite) {
-    if (failed(initializeLinker(mod)))
-      return failure();
+    // If this is the first module, setup the linker based on it
+    if (!composite) {
+      if (failed(initializeLinker(mod)))
+        return failure();
 
-    // We always override from source for the first module.
-    flags &= LinkerFlags::OverrideFromSrc;
+      // We always override from source for the first module.
+      flags &= LinkerFlags::OverrideFromSrc;
+    }
   }
 
   return summarize(mod, flags);