CIRCT  19.0.0git
InsertMergeBlocks.cpp
Go to the documentation of this file.
1 //===- InsertMergeBlocks.cpp - Insert Merge Blocks --------------*- C++ -*-===//
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 
10 #include "mlir/Analysis/CFGLoopInfo.h"
11 #include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
12 #include "mlir/Conversion/LLVMCommon/Pattern.h"
13 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
14 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
15 #include "mlir/Dialect/Func/IR/FuncOps.h"
16 #include "mlir/IR/Dominance.h"
17 #include "mlir/Pass/Pass.h"
18 #include "mlir/Transforms/DialectConversion.h"
19 #include "llvm/ADT/TypeSwitch.h"
20 
21 namespace circt {
22 #define GEN_PASS_DEF_INSERTMERGEBLOCKS
23 #include "circt/Transforms/Passes.h.inc"
24 } // namespace circt
25 
26 using namespace mlir;
27 using namespace circt;
28 
29 /// Replaces the branching to oldDest of with an equivalent operation that
30 /// instead branches to newDest.
31 static LogicalResult changeBranchTarget(Block *block, Block *oldDest,
32  Block *newDest,
33  ConversionPatternRewriter &rewriter) {
34  rewriter.setInsertionPointToEnd(block);
35  auto term = block->getTerminator();
36  return llvm::TypeSwitch<Operation *, LogicalResult>(term)
37  .Case<cf::BranchOp>([&](auto branchOp) {
38  rewriter.replaceOpWithNewOp<cf::BranchOp>(branchOp, newDest,
39  branchOp->getOperands());
40  return success();
41  })
42  .Case<cf::CondBranchOp>([&](auto condBr) {
43  auto cond = condBr.getCondition();
44 
45  Block *trueDest = condBr.getTrueDest();
46  Block *falseDest = condBr.getFalseDest();
47 
48  // Change to the correct destination.
49  if (trueDest == oldDest)
50  trueDest = newDest;
51 
52  if (falseDest == oldDest)
53  falseDest = newDest;
54 
55  rewriter.replaceOpWithNewOp<cf::CondBranchOp>(
56  condBr, cond, trueDest, condBr.getTrueOperands(), falseDest,
57  condBr.getFalseOperands());
58  return success();
59  })
60  .Default([&](Operation *op) {
61  return op->emitError("Unexpected terminator that cannot be handled.");
62  });
63 }
64 
65 /// Creates a new intermediate block that b1 and b2 branch to. The new block
66 /// branches to their common successor oldSucc.
67 static FailureOr<Block *> buildMergeBlock(Block *b1, Block *b2, Block *oldSucc,
68  ConversionPatternRewriter &rewriter) {
69  auto blockArgTypes = oldSucc->getArgumentTypes();
70  SmallVector<Location> argLocs(blockArgTypes.size(), rewriter.getUnknownLoc());
71 
72  Block *res = rewriter.createBlock(oldSucc, blockArgTypes, argLocs);
73  rewriter.create<cf::BranchOp>(rewriter.getUnknownLoc(), oldSucc,
74  res->getArguments());
75 
76  if (failed(changeBranchTarget(b1, oldSucc, res, rewriter)))
77  return failure();
78  if (failed(changeBranchTarget(b2, oldSucc, res, rewriter)))
79  return failure();
80 
81  return res;
82 }
83 
84 namespace {
85 /// A dual CFG that contracts cycles into single logical blocks.
86 struct DualGraph {
87  DualGraph(Region &r, CFGLoopInfo &loopInfo);
88 
89  size_t getNumPredecessors(Block *b) { return predCnts.lookup(b); }
90  void getPredecessors(Block *b, SmallVectorImpl<Block *> &res);
91 
92  size_t getNumSuccessors(Block *b) { return succMap.lookup(b).size(); }
93  ArrayRef<Block *> getSuccessors(Block *b) {
94  return succMap.find(b)->getSecond();
95  }
96 
97  // If the block is part of a contracted block, the header of the contracted
98  // block is returned. Otherwise, the block itself is returned.
99  Block *lookupDualBlock(Block *b);
100  DenseMap<Block *, size_t> getPredCountMapCopy() { return predCnts; }
101 
102 private:
103  CFGLoopInfo &loopInfo;
104 
105  DenseMap<Block *, SmallVector<Block *>> succMap;
106  DenseMap<Block *, size_t> predCnts;
107 };
108 } // namespace
109 
110 DualGraph::DualGraph(Region &r, CFGLoopInfo &loopInfo)
111  : loopInfo(loopInfo), succMap(), predCnts() {
112  for (Block &b : r) {
113  CFGLoop *loop = loopInfo.getLoopFor(&b);
114 
115  if (loop && loop->getHeader() != &b)
116  continue;
117 
118  // Create and get a new succ map entry for the current block.
119  SmallVector<Block *> &succs =
120  succMap.try_emplace(&b, SmallVector<Block *>()).first->getSecond();
121 
122  // NOTE: This assumes that there is only one exiting node, i.e., not
123  // two blocks from the same loop can be predecessors of one block.
124  unsigned predCnt = 0;
125  for (auto *pred : b.getPredecessors())
126  if (!loop || !loop->contains(pred))
127  predCnt++;
128 
129  if (loop && loop->getHeader() == &b)
130  loop->getExitBlocks(succs);
131  else
132  llvm::copy(b.getSuccessors(), std::back_inserter(succs));
133 
134  predCnts.try_emplace(&b, predCnt);
135  }
136 }
137 
138 Block *DualGraph::lookupDualBlock(Block *b) {
139  CFGLoop *loop = loopInfo.getLoopFor(b);
140  if (!loop)
141  return b;
142 
143  return loop->getHeader();
144 }
145 
146 void DualGraph::getPredecessors(Block *b, SmallVectorImpl<Block *> &res) {
147  CFGLoop *loop = loopInfo.getLoopFor(b);
148  assert((!loop || loop->getHeader() == b) &&
149  "can only get predecessors of blocks in the graph");
150 
151  for (auto *pred : b->getPredecessors()) {
152  if (loop && loop->contains(pred))
153  continue;
154 
155  if (CFGLoop *predLoop = loopInfo.getLoopFor(pred)) {
156  assert(predLoop->getExitBlock() &&
157  "multiple exit blocks are not yet supported");
158  res.push_back(predLoop->getHeader());
159  continue;
160  }
161  res.push_back(pred);
162  }
163 }
164 
165 namespace {
166 using BlockToBlockMap = DenseMap<Block *, Block *>;
167 /// A helper class to store the split block information gathered during analysis
168 /// of the CFG.
169 struct SplitInfo {
170  /// Points to the last split block that dominates the block.
171  BlockToBlockMap in;
172  /// Either points to the last split block or to itself, if the block itself is
173  /// a split block.
174  BlockToBlockMap out;
175 };
176 } // namespace
177 
178 /// Builds a binary merge block tree for the predecessors of currBlock.
179 static LogicalResult buildMergeBlocks(Block *currBlock, SplitInfo &splitInfo,
180  Block *predDom,
181  ConversionPatternRewriter &rewriter,
182  DualGraph &graph) {
183  SmallVector<Block *> preds;
184  llvm::copy(currBlock->getPredecessors(), std::back_inserter(preds));
185 
186  // Map from split blocks to blocks that descend from it.
187  DenseMap<Block *, Block *> predsToConsider;
188 
189  while (!preds.empty()) {
190  Block *pred = preds.pop_back_val();
191  Block *splitBlock = splitInfo.out.lookup(graph.lookupDualBlock(pred));
192  if (splitBlock == predDom)
193  // Needs no additional merge block, as this directly descends from the
194  // correct split block.
195  continue;
196 
197  if (predsToConsider.count(splitBlock) == 0) {
198  // No other block with the same split block was found yet, so just store
199  // it and wait for a match.
200  predsToConsider.try_emplace(splitBlock, pred);
201  continue;
202  }
203 
204  // Found a pair, so insert a new merge block for them.
205  Block *other = predsToConsider.lookup(splitBlock);
206  predsToConsider.erase(splitBlock);
207 
209  buildMergeBlock(pred, other, currBlock, rewriter);
210  if (failed(mergeBlock))
211  return failure();
212 
213  // Update info for the newly created block.
214  Block *splitIn = splitInfo.in.lookup(splitBlock);
215  splitInfo.in.try_emplace(*mergeBlock, splitIn);
216  // By construction, this block has only one successor, therefore, out == in.
217  splitInfo.out.try_emplace(*mergeBlock, splitIn);
218 
219  preds.push_back(*mergeBlock);
220  }
221  if (!predsToConsider.empty())
222  return currBlock->getParentOp()->emitError(
223  "irregular control flow is not yet supported");
224  return success();
225 }
226 
227 /// Checks preconditions of this transformation.
228 static LogicalResult preconditionCheck(Region &r, CFGLoopInfo &loopInfo) {
229  for (auto &info : loopInfo.getTopLevelLoops())
230  // Does only return a block if it is the only exit block.
231  if (!info->getExitBlock())
232  return r.getParentOp()->emitError(
233  "multiple exit blocks are not yet supported");
234 
235  return success();
236 }
237 
238 /// Insert additional blocks that serve as counterparts to the blocks that
239 /// diverged the control flow.
240 /// The resulting merge block tree is guaranteed to be a binary tree.
241 ///
242 /// This transformation does not affect any blocks that are part of a loop as it
243 /// treats a loop as one logical block.
244 /// Irregular control flow is not supported and results in a failed
245 /// transformation.
246 LogicalResult circt::insertMergeBlocks(Region &r,
247  ConversionPatternRewriter &rewriter) {
248  Block *entry = &r.front();
249  DominanceInfo domInfo(r.getParentOp());
250 
251  CFGLoopInfo loopInfo(domInfo.getDomTree(&r));
252  if (failed(preconditionCheck(r, loopInfo)))
253  return failure();
254 
255  // Traversing the graph in topological order can be simply done with a stack.
256  SmallVector<Block *> stack;
257  stack.push_back(entry);
258 
259  // Holds the graph that contains the relevant blocks. It for example contracts
260  // loops into one block to preserve a DAG structure.
261  DualGraph graph(r, loopInfo);
262 
263  // Counts the amount of predecessors remaining.
264  auto predsToVisit = graph.getPredCountMapCopy();
265 
266  SplitInfo splitInfo;
267 
268  while (!stack.empty()) {
269  Block *currBlock = stack.pop_back_val();
270 
271  Block *in = nullptr;
272  Block *out = nullptr;
273 
274  bool isMergeBlock = graph.getNumPredecessors(currBlock) > 1;
275  bool isSplitBlock = graph.getNumSuccessors(currBlock) > 1;
276 
277  SmallVector<Block *> preds;
278  graph.getPredecessors(currBlock, preds);
279 
280  if (isMergeBlock) {
281  Block *predDom = currBlock;
282  for (auto *pred : preds) {
283  predDom = domInfo.findNearestCommonDominator(predDom, pred);
284  }
285 
286  if (failed(
287  buildMergeBlocks(currBlock, splitInfo, predDom, rewriter, graph)))
288  return failure();
289 
290  // The sub-CFG created by the predDom (split block) and the current merge
291  // block can logically be treated like a single block, thus their "in"s
292  // are the same.
293  in = splitInfo.in.lookup(predDom);
294  } else if (!preds.empty()) {
295  Block *pred = preds.front();
296 
297  in = splitInfo.out.lookup(pred);
298  }
299 
300  if (isSplitBlock)
301  out = currBlock;
302  else
303  out = in;
304 
305  splitInfo.in.try_emplace(currBlock, in);
306  splitInfo.out.try_emplace(currBlock, out);
307 
308  for (auto *succ : graph.getSuccessors(currBlock)) {
309  auto it = predsToVisit.find(succ);
310  unsigned predsRemaining = --(it->getSecond());
311  // Pushing the block on the stack once all it's successors were visited
312  // ensures a topological traversal.
313  if (predsRemaining == 0)
314  stack.push_back(succ);
315  }
316  }
317 
318  return success();
319 }
320 
321 namespace {
322 
323 using PtrSet = SmallPtrSet<Operation *, 4>;
324 
325 struct FuncOpPattern : public OpConversionPattern<func::FuncOp> {
326 
327  FuncOpPattern(PtrSet &rewrittenFuncs, MLIRContext *ctx)
328  : OpConversionPattern(ctx), rewrittenFuncs(rewrittenFuncs) {}
329 
330  LogicalResult
331  matchAndRewrite(func::FuncOp op, OpAdaptor adaptor,
332  ConversionPatternRewriter &rewriter) const override {
333  rewriter.startOpModification(op);
334 
335  if (!op.isExternal())
336  if (failed(insertMergeBlocks(op.getRegion(), rewriter))) {
337  rewriter.cancelOpModification(op);
338  return failure();
339  }
340 
341  rewriter.finalizeOpModification(op);
342  rewrittenFuncs.insert(op);
343 
344  return success();
345  }
346 
347 private:
348  PtrSet &rewrittenFuncs;
349 };
350 
351 struct InsertMergeBlocksPass
352  : public circt::impl::InsertMergeBlocksBase<InsertMergeBlocksPass> {
353 public:
354  void runOnOperation() override {
355  auto *ctx = &getContext();
356  RewritePatternSet patterns(ctx);
357  // Remembers traversed functions to only apply the conversion once.
358  PtrSet rewrittenFuncs;
359  patterns.add<FuncOpPattern>(rewrittenFuncs, ctx);
360 
361  ConversionTarget target(*ctx);
362  target.addDynamicallyLegalOp<func::FuncOp>(
363  [&](func::FuncOp func) { return rewrittenFuncs.contains(func); });
364  target.addLegalDialect<cf::ControlFlowDialect>();
365 
366  if (applyPartialConversion(getOperation(), target, std::move(patterns))
367  .failed())
368  signalPassFailure();
369  }
370 };
371 
372 } // namespace
373 
374 namespace circt {
375 std::unique_ptr<mlir::Pass> createInsertMergeBlocksPass() {
376  return std::make_unique<InsertMergeBlocksPass>();
377 }
378 } // namespace circt
assert(baseType &&"element must be base type")
static void mergeBlock(Block &destination, Block::iterator insertPoint, Block &source)
Move all operations from a source block in to a destination block.
Definition: ExpandWhens.cpp:28
static LogicalResult changeBranchTarget(Block *block, Block *oldDest, Block *newDest, ConversionPatternRewriter &rewriter)
Replaces the branching to oldDest of with an equivalent operation that instead branches to newDest.
static LogicalResult preconditionCheck(Region &r, CFGLoopInfo &loopInfo)
Checks preconditions of this transformation.
static LogicalResult buildMergeBlocks(Block *currBlock, SplitInfo &splitInfo, Block *predDom, ConversionPatternRewriter &rewriter, DualGraph &graph)
Builds a binary merge block tree for the predecessors of currBlock.
static FailureOr< Block * > buildMergeBlock(Block *b1, Block *b2, Block *oldSucc, ConversionPatternRewriter &rewriter)
Creates a new intermediate block that b1 and b2 branch to.
The InstanceGraph op interface, see InstanceGraphInterface.td for more details.
Definition: DebugAnalysis.h:21
std::unique_ptr< mlir::Pass > createInsertMergeBlocksPass()
mlir::LogicalResult insertMergeBlocks(mlir::Region &r, mlir::ConversionPatternRewriter &rewriter)
Insert additional blocks that serve as counterparts to the blocks that diverged the control flow.