xref: /llvm-project/mlir/lib/Dialect/LLVMIR/Transforms/AddComdats.cpp (revision a6857156df9a73720fbd9633067d1a61c32dd74e)
1 //===- AddComdats.cpp - Add comdats to linkonce functions -----------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/Dialect/LLVMIR/Transforms/AddComdats.h"
10 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
11 #include "mlir/Pass/Pass.h"
12 
13 namespace mlir {
14 namespace LLVM {
15 #define GEN_PASS_DEF_LLVMADDCOMDATS
16 #include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc"
17 } // namespace LLVM
18 } // namespace mlir
19 
20 using namespace mlir;
21 
addComdat(LLVM::LLVMFuncOp & op,OpBuilder & builder,SymbolTable & symbolTable,ModuleOp & module)22 static void addComdat(LLVM::LLVMFuncOp &op, OpBuilder &builder,
23                       SymbolTable &symbolTable, ModuleOp &module) {
24   const char *comdatName = "__llvm_comdat";
25   mlir::LLVM::ComdatOp comdatOp =
26       symbolTable.lookup<mlir::LLVM::ComdatOp>(comdatName);
27   if (!comdatOp) {
28     PatternRewriter::InsertionGuard guard(builder);
29     builder.setInsertionPointToStart(module.getBody());
30     comdatOp =
31         builder.create<mlir::LLVM::ComdatOp>(module.getLoc(), comdatName);
32     symbolTable.insert(comdatOp);
33   }
34 
35   PatternRewriter::InsertionGuard guard(builder);
36   builder.setInsertionPointToStart(&comdatOp.getBody().back());
37   auto selectorOp = builder.create<mlir::LLVM::ComdatSelectorOp>(
38       comdatOp.getLoc(), op.getSymName(), mlir::LLVM::comdat::Comdat::Any);
39   op.setComdatAttr(mlir::SymbolRefAttr::get(
40       builder.getContext(), comdatName,
41       mlir::FlatSymbolRefAttr::get(selectorOp.getSymNameAttr())));
42 }
43 
44 namespace {
45 struct AddComdatsPass : public LLVM::impl::LLVMAddComdatsBase<AddComdatsPass> {
runOnOperation__anon7371f9f80111::AddComdatsPass46   void runOnOperation() override {
47     OpBuilder builder{&getContext()};
48     ModuleOp mod = getOperation();
49 
50     std::unique_ptr<SymbolTable> symbolTable;
51     auto getSymTab = [&]() -> SymbolTable & {
52       if (!symbolTable)
53         symbolTable = std::make_unique<SymbolTable>(mod);
54       return *symbolTable;
55     };
56     for (auto op : mod.getBody()->getOps<LLVM::LLVMFuncOp>()) {
57       if (op.getLinkage() == LLVM::Linkage::Linkonce ||
58           op.getLinkage() == LLVM::Linkage::LinkonceODR) {
59         addComdat(op, builder, getSymTab(), mod);
60       }
61     }
62   }
63 };
64 } // namespace
65