[MLIR][mlir-link] Add COMDAT resolution to LLVM dialect linker.

Author: Jezurko

clang/include/clang/CIR/Interfaces/CIRLinkerInterface.h

@@ -45,8 +45,13 @@ class CIRSymbolLinkerInterface
 
   static bool isComdat(Operation *op);
 
-  static std::optional<mlir::link::ComdatSelector>
-  getComdatSelector(Operation *op);
+  static bool hasComdat(Operation *op);
+
+  static const link::Comdat *getComdatResolution(Operation *op);
+
+  static bool selectedByComdat(Operation *op);
+
+  static void updateNoDeduplicate(Operation *op);
 
   static Visibility getVisibility(Operation *op);
 

clang/lib/CIR/Dialect/IR/CIRLinkerInterface.cpp

@@ -99,12 +99,23 @@ bool CIRSymbolLinkerInterface::isComdat(Operation *op) {
   return false;
 }
 
-std::optional<link::ComdatSelector>
-CIRSymbolLinkerInterface::getComdatSelector(Operation *op) {
-  // TODO(frabert): Extracting comdat info from CIR is not implemented yet
-  return std::nullopt;
+bool CIRSymbolLinkerInterface::hasComdat(Operation *op) {
+  // TODO: Extracting comdat info from CIR is not implemented yet
+  return false;
+}
+
+const link::Comdat *
+CIRSymbolLinkerInterface::getComdatResolution(Operation *op) {
+  return nullptr;
 }
 
+bool CIRSymbolLinkerInterface::selectedByComdat(Operation *op) {
+  // TODO: Extracting comdat info from CIR is not implemented yet
+  llvm_unreachable("comdat resolution not implemented for CIR");
+}
+
+void CIRSymbolLinkerInterface::updateNoDeduplicate(Operation *op) {}
+
 Visibility CIRSymbolLinkerInterface::getVisibility(Operation *op) {
   if (auto gv = dyn_cast<GlobalOp>(op))
     return toLLVMVisibility(gv.getGlobalVisibility());

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

@@ -16,7 +16,9 @@ class LLVMSymbolLinkerInterface
   static Visibility getVisibility(Operation *op);
   static void setVisibility(Operation *op, Visibility visibility);
   static bool isComdat(Operation *op);
-  static std::optional<link::ComdatSelector> getComdatSelector(Operation *op);
+  static bool hasComdat(Operation *op);
+  static SymbolRefAttr getComdatSymbol(Operation *op);
+  static LLVM::comdat::Comdat getComdatSelectionKind(Operation *op);
   static bool isDeclaration(Operation *op);
   static unsigned getBitWidth(Operation *op);
   static UnnamedAddr getUnnamedAddr(Operation *op);
@@ -34,7 +36,19 @@ class LLVMSymbolLinkerInterface
   dependencies(Operation *op, SymbolTableCollection &collection) const override;
   LogicalResult initialize(ModuleOp src) override;
   LogicalResult finalize(ModuleOp dst) const override;
+  LogicalResult moduleOpSummary(ModuleOp src,
+                                SymbolTableCollection &collection) override;
   Operation *appendGlobals(llvm::StringRef glob, link::LinkState &state);
+  Operation *appendComdatOps(ArrayRef<Operation *> globs, LLVM::ComdatOp comdat,
+                             link::LinkState &state);
+  link::ComdatResolution
+  computeComdatResolution(Operation *, SymbolTableCollection &, link::Comdat *);
+  LogicalResult resolveComdats(ModuleOp srcMod,
+                               SymbolTableCollection &collection);
+  const link::Comdat *getComdatResolution(Operation *op) const;
+  bool selectedByComdat(Operation *op) const;
+  void dropReplacedComdat(Operation *op) const;
+  static void updateNoDeduplicate(Operation *op);
 
   template <typename structor_t>
   Operation *appendGlobalStructors(link::LinkState &state) {
@@ -79,22 +93,27 @@ class LLVMSymbolLinkerInterface
       ArrayRef<Attribute> priorities = structor.getPriorities().getValue();
       ArrayRef<Attribute> data = structor.getData().getValue();
 
-      for (auto [idx, dataAttr] : llvm::enumerate(data)) {
+      for (auto [idx, structor] : llvm::enumerate(structorList)) {
+        auto structorSymbol = cast<FlatSymbolRefAttr>(structor);
+        // Skip constructors not included based on COMDAT
+        if (!summary.contains(structorSymbol.getValue()))
+          continue;
+
+        auto dataAttr = data[idx];
         // data value is either #llvm.zero or symbol ref
         // if it is zero, we always have to include the value
         // if it is a symbol ref, we have to check if the symbol
         // from the same module is being used
-        //
         if (auto globalSymbol = dyn_cast<FlatSymbolRefAttr>(dataAttr)) {
-          auto globalOp = summary.lookup(globalSymbol.getValue());
+          Operation *globalOp = summary.lookup(globalSymbol.getValue());
           assert(globalOp && "structor referenced global not in summary?");
           // globals are definde at module level
           if (globalOp->getParentOp() != op->getParentOp())
             continue;
         }
 
         newData.push_back(dataAttr);
-        newStructorList.push_back(structorList[idx]);
+        newStructorList.push_back(structor);
         newPriorities.push_back(priorities[idx]);
       }
     }
@@ -125,6 +144,7 @@ class LLVMSymbolLinkerInterface
 private:
   DataLayoutSpecInterface dtla = {};
   TargetSystemSpecInterface targetSys = {};
+  llvm::StringMap<link::Comdat> comdatResolution;
 };
 
 } // namespace LLVM

mlir/include/mlir/Linker/LLVMLinkerMixin.h

@@ -139,9 +139,17 @@ static UnnamedAddr getMinUnnamedAddr(UnnamedAddr lhs, UnnamedAddr rhs) {
 
 using ComdatKind = LLVM::comdat::Comdat;
 
-struct ComdatSelector {
-  StringRef name;
+struct Comdat {
   ComdatKind kind;
+  Operation *selectorOp;
+  llvm::SmallPtrSet<Operation *, 2> users;
+};
+
+enum class ComdatResolution {
+  LinkFromSrc,
+  LinkFromDst,
+  LinkFromBoth,
+  Failure,
 };
 
 //===----------------------------------------------------------------------===//
@@ -172,6 +180,20 @@ class LLVMLinkerMixin {
     if (derived.isComdat(pair.src))
       return true;
 
+    // Thrown away symbol can affect the visibility
+    if (pair.dst) {
+      Visibility srcVisibility = derived.getVisibility(pair.src);
+      Visibility dstVisibility = derived.getVisibility(pair.dst);
+      Visibility visibility = getMinVisibility(srcVisibility, dstVisibility);
+
+      derived.setVisibility(pair.src, visibility);
+      derived.setVisibility(pair.dst, visibility);
+    }
+    if (derived.hasComdat(pair.src)) {
+      // operations with COMDAT are selected as a group
+      return derived.selectedByComdat(pair.src);
+    }
+
     Linkage srcLinkage = derived.getLinkage(pair.src);
 
     // Always import variables with appending linkage.
@@ -218,6 +240,10 @@ class LLVMLinkerMixin {
       return pair.src->emitError(error) << " dst: " << pair.dst->getLoc();
     };
 
+    if (derived.isComdat(pair.src) != derived.isComdat(pair.dst)) {
+      return linkError("Linking ComdatOp with non-comdat op");
+    }
+
     Linkage srcLinkage = derived.getLinkage(pair.src);
     Linkage dstLinkage = derived.getLinkage(pair.dst);
 
@@ -340,6 +366,11 @@ class LLVMLinkerMixin {
     if (isWeakForLinker(srcLinkage)) {
       assert(!isExternalWeakLinkage(dstLinkage));
       assert(!isAvailableExternallyLinkage(dstLinkage));
+      const Comdat *comdat = derived.getComdatResolution(pair.src);
+      if (comdat && comdat->kind == ComdatKind::NoDeduplicate) {
+        derived.updateNoDeduplicate(pair.src);
+        return ConflictResolution::LinkFromBothAndRenameSrc;
+      }
       if (isLinkOnceLinkage(dstLinkage) && isWeakLinkage(srcLinkage)) {
         return ConflictResolution::LinkFromSrc;
       }
@@ -349,38 +380,12 @@ class LLVMLinkerMixin {
 
     if (isWeakForLinker(dstLinkage)) {
       assert(isExternalLinkage(srcLinkage));
-      return ConflictResolution::LinkFromSrc;
-    }
-
-    std::optional<ComdatSelector> srcComdatSel =
-        derived.getComdatSelector(pair.src);
-    std::optional<ComdatSelector> dstComdatSel =
-        derived.getComdatSelector(pair.dst);
-    if (srcComdatSel.has_value() && dstComdatSel.has_value()) {
-      auto srcComdatName = srcComdatSel->name;
-      auto dstComdatName = dstComdatSel->name;
-      auto srcComdat = srcComdatSel->kind;
-      auto dstComdat = dstComdatSel->kind;
-      if (srcComdatName != dstComdatName) {
-        llvm_unreachable("Comdat selector names don't match");
-      }
-      if (srcComdat != dstComdat) {
-        llvm_unreachable("Comdat selector kinds don't match");
-      }
-
-      if (srcComdat == mlir::LLVM::comdat::Comdat::Any) {
-        return ConflictResolution::LinkFromDst;
-      }
-      if (srcComdat == mlir::LLVM::comdat::Comdat::NoDeduplicate) {
-        return ConflictResolution::Failure;
+      const Comdat *comdat = derived.getComdatResolution(pair.dst);
+      if (comdat && comdat->kind == ComdatKind::NoDeduplicate) {
+        derived.updateNoDeduplicate(pair.dst);
+        return ConflictResolution::LinkFromBothAndRenameDst;
       }
-      if (srcComdat == mlir::LLVM::comdat::Comdat::ExactMatch) {
-        return ConflictResolution::LinkFromDst;
-      }
-      if (srcComdat == mlir::LLVM::comdat::Comdat::SameSize) {
-        return ConflictResolution::LinkFromDst;
-      }
-      llvm_unreachable("unimplemented comdat kind");
+      return ConflictResolution::LinkFromSrc;
     }
 
     // If we reach here, we have two external definitions that can't be resolved

mlir/include/mlir/Linker/LinkerInterface.h

@@ -60,6 +60,8 @@ class LinkState {
     return builder.create<Op>(location, std::forward<Args>(args)...);
   }
 
+  OpBuilder &getBuilder() { return builder; };
+
 private:
   // Private constructor used by nest()
   LinkState(ModuleOp dst, std::shared_ptr<IRMapping> mapping)
@@ -145,7 +147,8 @@ class SymbolLinkerInterface : public LinkerInterface<SymbolLinkerInterface> {
 
   /// 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) {
+  virtual LogicalResult moduleOpSummary(ModuleOp module,
+                                        SymbolTableCollection &collection) {
     return success();
   }
 
@@ -282,9 +285,10 @@ class SymbolLinkerInterfaces {
     return Conflict::noConflict(src);
   }
 
-  LogicalResult moduleOpSummary(ModuleOp src) {
+  LogicalResult moduleOpSummary(ModuleOp src,
+                                SymbolTableCollection &collection) {
     for (SymbolLinkerInterface *linker : interfaces) {
-      if (failed(linker->moduleOpSummary(src)))
+      if (failed(linker->moduleOpSummary(src, collection)))
         return failure();
     }
     return success();