[MLIR][mlir-link] Fix conflict resolution

Author: xlauko

mlir/include/mlir/Linker/IRMover.h

@@ -25,6 +25,8 @@ struct ConflictPair {
   static ConflictPair noConflict(Operation *src) { return {nullptr, src}; }
 };
 
+using Summary = llvm::StringMap<ConflictPair>;
+
 struct IRMover {
 
   ModuleOp composite;
@@ -36,7 +38,7 @@ struct IRMover {
 
   explicit IRMover(ModuleOp composite) : composite(composite) {}
 
-  LogicalResult move(ArrayRef<ConflictPair> valuesToLink);
+  LogicalResult move(const Summary &summary);
 
 private:
   Operation * remap(ConflictPair pair);
@@ -46,26 +48,4 @@ struct IRMover {
 
 } // namespace mlir::link
 
-namespace llvm {
-
-template <>
-struct DenseMapInfo<mlir::link::ConflictPair> {
-  static mlir::link::ConflictPair getEmptyKey() {
-    auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
-    return {{}, static_cast<mlir::Operation*>(pointer)};
-  }
-  static mlir::link::ConflictPair getTombstoneKey() {
-    auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
-    return {{}, static_cast<mlir::Operation *>(pointer)};
-  }
-  static unsigned getHashValue(mlir::link::ConflictPair val) {
-    return DenseMapInfo<const mlir::Operation *>::getHashValue(val.src);
-  }
-  static bool isEqual(mlir::link::ConflictPair lhs, mlir::link::ConflictPair rhs) {
-    return lhs.src == rhs.src;
-  }
-};
-
-} // namespace llvm
-
 #endif // MLIR_LINKER_IRMOVER_H

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

@@ -167,25 +167,30 @@ class LLVMSymbolLinkerInterface : public SymbolLinkerInterface {
 
   StringRef getSymbol(Operation *op) const override { return symbol(op); }
 
+  ConflictPair join(ConflictPair existing, Operation *src) const {
+    assert(!existing.hasConflict() && "expected non-conflicting pair");
+    if (isLocalLinkage(getLinkage(existing.src)))
+      return ConflictPair::noConflict(existing.src);
+    return {existing.src, src};
+  }
+
   ConflictPair findConflict(Operation *src) const override {
     assert(canBeLinked(src) && "expected linkable operation");
 
     if (isLocalLinkage(getLinkage(src)))
       return ConflictPair::noConflict(src);
 
-    // TODO: make lookup through module state
     if (auto it = summary.find(getSymbol(src)); it != summary.end()) {
-      Operation *dst = it->second;
-      if (dst != src && !isLocalLinkage(getLinkage(dst))) {
-        return {dst, src};
-      }
+      return join(it->second, src);
     }
 
     return ConflictPair::noConflict(src);
   }
 
   bool isLinkNeeded(ConflictPair pair, bool forDependency) const override {
     assert(canBeLinked(pair.src) && "expected linkable operation");
+    if (pair.src == pair.dst)
+      return false;
 
     LLVM::Linkage srcLinkage = getLinkage(pair.src);
 
@@ -320,21 +325,13 @@ class LLVMSymbolLinkerInterface : public SymbolLinkerInterface {
       valuesToClone.insert(*linkFromSrc ? pair.dst : pair.src);
 
     if (*linkFromSrc)
-      registerForLink(pair);
+      summary[getSymbol(pair.src)] = pair;
     return success();
   }
 
   void registerForLink(Operation *op) override {
     assert(canBeLinked(op) && "expected linkable operation");
-    registerForLink(ConflictPair::noConflict(op));
-  }
-
-  void registerForLink(ConflictPair pair) {
-    StringRef sym = getSymbol(pair.src);
-    assert(!summary.contains(sym) && "expected unique symbol");
-    summary[sym] = pair.src;
-
-    valuesToLink.insert(pair);
+    summary[getSymbol(op)] = ConflictPair::noConflict(op);
   }
 
   LogicalResult initialize(ModuleOp src) override {
@@ -354,7 +351,7 @@ class LLVMSymbolLinkerInterface : public SymbolLinkerInterface {
     // }
 
     IRMover mover(dst);
-    return mover.move(valuesToLink.getArrayRef());
+    return mover.move(summary);
   }
 
   Operation *prototype(ConflictPair pair) const {
@@ -374,10 +371,10 @@ class LLVMSymbolLinkerInterface : public SymbolLinkerInterface {
   Operation *materialize(ConflictPair pair, ModuleOp dst) const override {
     // Make definition if destination does not have one or has only declaration
     bool forDefinition = !pair.dst || isDeclaration(pair.dst);
-    if (!forDefinition)
+    if (!forDefinition || isDeclaration(pair.src))
       return prototype(pair);
 
-    // Definition laready exists
+    // Definition already exists
     if (pair.dst && !isDeclaration(pair.dst))
       return pair.dst;
 
@@ -411,13 +408,11 @@ class LLVMSymbolLinkerInterface : public SymbolLinkerInterface {
   }
 
 private:
-  llvm::StringMap<Operation *> summary;
-
   std::unique_ptr<SymbolTable> symbolTable;
 
   SetVector<Operation *> valuesToClone;
 
-  SetVector<ConflictPair> valuesToLink;
+  llvm::StringMap<ConflictPair> summary;
 };
 
 //===----------------------------------------------------------------------===//

mlir/lib/Linker/IRMover.cpp

@@ -14,14 +14,18 @@
 using namespace mlir;
 using namespace mlir::link;
 
-LogicalResult IRMover::move(ArrayRef<ConflictPair> valuesToLink) {
-  worklist.insert(worklist.end(), valuesToLink.rbegin(), valuesToLink.rend());
+LogicalResult IRMover::move(const Summary &summary) {
+  worklist.reserve(summary.size());
+  for (const auto &[_, pair] : summary) {
+    worklist.push_back(pair);
+  }
 
   while (!worklist.empty()) {
     ConflictPair pair = worklist.back();
     worklist.pop_back();
 
-    assert(!mapping.contains(pair.src) && "expected no mapping for source");
+    if (mapping.contains(pair.src))
+      continue;
 
     if (!remap(pair))
       return failure();

mlir/test/mlir-link/linkage-b.mlir

@@ -1,4 +1,5 @@
 // RUN: mlir-link -sort-symbols -split-input-file %s | FileCheck %s
+// REQUIRES: mlir-link-common, mlir-link-weak
 
 // CHECK: llvm.mlir.global common @test1_a(0 : i8) {addr_space = 0 : i32} : i8
 // CHECK: llvm.mlir.global external @test2_a(0 : i8) {addr_space = 0 : i32} : i8

mlir/test/mlir-link/linkage-c.mlir

@@ -1,4 +1,5 @@
 // RUN: mlir-link -sort-symbols -split-input-file %s | FileCheck %s
+// REQUIRES: mlir-link-common, mlir-link-weak
 
 // CHECK: llvm.mlir.global common @test1_a(0 : i8) {addr_space = 0 : i32} : i8
 // CHECK: llvm.mlir.global external @test2_a(0 : i8) {addr_space = 0 : i32} : i8