[Mlir-commits] [mlir] [mlir] Add pass to add comdat to all linkonce functions (PR #65270)

[Mehdi Amini]

================
@@ -0,0 +1,78 @@
+//===- AddComdats.cpp - Prepare for translation to LLVM IR ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/LLVMIR/Transforms/AddComdats.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir {
+namespace LLVM {
+#define GEN_PASS_DEF_LLVMADDCOMDATS
+#include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc"
+} // namespace LLVM
+} // namespace mlir
+
+using namespace mlir;
+
+static void addComdat(LLVM::LLVMFuncOp &op, PatternRewriter &rewriter,
+                      ModuleOp &module) {
+  const char *comdatName = "__llvm_comdat";
+  mlir::LLVM::ComdatOp comdatOp =
+      module.lookupSymbol<mlir::LLVM::ComdatOp>(comdatName);
+  if (!comdatOp) {
+    PatternRewriter::InsertionGuard guard(rewriter);
+    rewriter.setInsertionPointToStart(module.getBody());
+    comdatOp =
+        rewriter.create<mlir::LLVM::ComdatOp>(module.getLoc(), comdatName);
+  }
+
+  PatternRewriter::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPointToStart(&comdatOp.getBody().back());
+  auto selectorOp = rewriter.create<mlir::LLVM::ComdatSelectorOp>(
+      comdatOp.getLoc(), op.getSymName(), mlir::LLVM::comdat::Comdat::Any);
+  rewriter.updateRootInPlace(op, [&]() {
+    op.setComdatAttr(mlir::SymbolRefAttr::get(
+        rewriter.getContext(), comdatName,
+        mlir::FlatSymbolRefAttr::get(selectorOp.getSymNameAttr())));
+  });
+}
+
+struct AddComdat : public OpRewritePattern<LLVM::LLVMFuncOp> {
+  using OpRewritePattern<LLVM::LLVMFuncOp>::OpRewritePattern;
+
+private:
+  LogicalResult matchAndRewrite(LLVM::LLVMFuncOp op,
+                                PatternRewriter &rewriter) const override {
+    if (op.getComdat() || (op.getLinkage() != LLVM::Linkage::Linkonce &&
+                           op.getLinkage() != LLVM::Linkage::LinkonceODR))
+      return failure();
+
+    auto mod = op->getParentOfType<mlir::ModuleOp>();
+    addComdat(op, rewriter, mod);
+    return success();
+  }
+};
+
+namespace {
+struct AddComdatsPass : public LLVM::impl::LLVMAddComdatsBase<AddComdatsPass> {
+  void runOnOperation() override {
+    MLIRContext *ctx = &getContext();
+    RewritePatternSet patterns(ctx);
+    patterns.add<AddComdat>(ctx);
+    if (failed(
+            applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
+      return signalPassFailure();
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> LLVM::createAddComdatsPass() {
+  return std::make_unique<AddComdatsPass>();
+}
----------------
joker-eph wrote:

(same as in the header, this will be part of the include line 18)

https://github.com/llvm/llvm-project/pull/65270