| 1 | //===- LegalizeForExport.cpp - Prepare for translation to LLVM IR ---------===// |
| 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/LegalizeForExport.h" |
| 10 | |
| 11 | #include "mlir/Dialect/LLVMIR/LLVMDialect.h" |
| 12 | #include "mlir/Dialect/LLVMIR/Transforms/DIExpressionLegalization.h" |
| 13 | #include "mlir/IR/Block.h" |
| 14 | #include "mlir/IR/Builders.h" |
| 15 | #include "mlir/IR/BuiltinOps.h" |
| 16 | #include "mlir/Pass/Pass.h" |
| 17 | |
| 18 | namespace mlir { |
| 19 | namespace LLVM { |
| 20 | #define GEN_PASS_DEF_LLVMLEGALIZEFOREXPORTPASS |
| 21 | #include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc" |
| 22 | } // namespace LLVM |
| 23 | } // namespace mlir |
| 24 | |
| 25 | using namespace mlir; |
| 26 | |
| 27 | /// If the given block has the same successor with different arguments, |
| 28 | /// introduce dummy successor blocks so that all successors of the given block |
| 29 | /// are different. |
| 30 | static void ensureDistinctSuccessors(Block &bb) { |
| 31 | // Early exit if the block cannot have successors. |
| 32 | if (bb.empty() || !bb.back().mightHaveTrait<OpTrait::IsTerminator>()) |
| 33 | return; |
| 34 | |
| 35 | auto *terminator = bb.getTerminator(); |
| 36 | |
| 37 | // Find repeated successors with arguments. |
| 38 | llvm::SmallDenseMap<Block *, SmallVector<int, 4>> successorPositions; |
| 39 | for (int i = 0, e = terminator->getNumSuccessors(); i < e; ++i) { |
| 40 | Block *successor = terminator->getSuccessor(index: i); |
| 41 | // Blocks with no arguments are safe even if they appear multiple times |
| 42 | // because they don't need PHI nodes. |
| 43 | if (successor->getNumArguments() == 0) |
| 44 | continue; |
| 45 | successorPositions[successor].push_back(i); |
| 46 | } |
| 47 | |
| 48 | // If a successor appears for the second or more time in the terminator, |
| 49 | // create a new dummy block that unconditionally branches to the original |
| 50 | // destination, and retarget the terminator to branch to this new block. |
| 51 | // There is no need to pass arguments to the dummy block because it will be |
| 52 | // dominated by the original block and can therefore use any values defined in |
| 53 | // the original block. |
| 54 | OpBuilder builder(terminator->getContext()); |
| 55 | for (const auto &successor : successorPositions) { |
| 56 | // Start from the second occurrence of a block in the successor list. |
| 57 | for (int position : llvm::drop_begin(successor.second, 1)) { |
| 58 | Block *dummyBlock = builder.createBlock(bb.getParent()); |
| 59 | terminator->setSuccessor(dummyBlock, position); |
| 60 | for (BlockArgument arg : successor.first->getArguments()) |
| 61 | dummyBlock->addArgument(arg.getType(), arg.getLoc()); |
| 62 | builder.create<LLVM::BrOp>(terminator->getLoc(), |
| 63 | dummyBlock->getArguments(), successor.first); |
| 64 | } |
| 65 | } |
| 66 | } |
| 67 | |
| 68 | void mlir::LLVM::ensureDistinctSuccessors(Operation *op) { |
| 69 | op->walk(callback: [](Operation *nested) { |
| 70 | for (Region ®ion : llvm::make_early_inc_range(Range: nested->getRegions())) { |
| 71 | for (Block &block : llvm::make_early_inc_range(Range&: region)) { |
| 72 | ::ensureDistinctSuccessors(bb&: block); |
| 73 | } |
| 74 | } |
| 75 | }); |
| 76 | } |
| 77 | |
| 78 | namespace { |
| 79 | struct LegalizeForExportPass |
| 80 | : public LLVM::impl::LLVMLegalizeForExportPassBase<LegalizeForExportPass> { |
| 81 | void runOnOperation() override { |
| 82 | LLVM::ensureDistinctSuccessors(op: getOperation()); |
| 83 | LLVM::legalizeDIExpressionsRecursively(op: getOperation()); |
| 84 | } |
| 85 | }; |
| 86 | } // namespace |
| 87 | |