MLIR  21.0.0git
Inliner.cpp
Go to the documentation of this file.
1 //===- Inliner.cpp ---- SCC-based inliner ---------------------------------===//
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 // This file implements Inliner that uses a basic inlining
10 // algorithm that operates bottom up over the Strongly Connect Components(SCCs)
11 // of the CallGraph. This enables a more incremental propagation of inlining
12 // decisions from the leafs to the roots of the callgraph.
13 //
14 //===----------------------------------------------------------------------===//
15 
17 #include "mlir/IR/Threading.h"
20 #include "mlir/Pass/Pass.h"
23 #include "llvm/ADT/SCCIterator.h"
24 #include "llvm/ADT/STLExtras.h"
25 #include "llvm/ADT/SmallPtrSet.h"
26 #include "llvm/Support/Debug.h"
27 
28 #define DEBUG_TYPE "inlining"
29 
30 using namespace mlir;
31 
33 
34 //===----------------------------------------------------------------------===//
35 // Symbol Use Tracking
36 //===----------------------------------------------------------------------===//
37 
38 /// Walk all of the used symbol callgraph nodes referenced with the given op.
40  Operation *op, CallGraph &cg, SymbolTableCollection &symbolTable,
42  function_ref<void(CallGraphNode *, Operation *)> callback) {
43  auto symbolUses = SymbolTable::getSymbolUses(op);
44  assert(symbolUses && "expected uses to be valid");
45 
46  Operation *symbolTableOp = op->getParentOp();
47  for (const SymbolTable::SymbolUse &use : *symbolUses) {
48  auto refIt = resolvedRefs.insert({use.getSymbolRef(), nullptr});
49  CallGraphNode *&node = refIt.first->second;
50 
51  // If this is the first instance of this reference, try to resolve a
52  // callgraph node for it.
53  if (refIt.second) {
54  auto *symbolOp = symbolTable.lookupNearestSymbolFrom(symbolTableOp,
55  use.getSymbolRef());
56  auto callableOp = dyn_cast_or_null<CallableOpInterface>(symbolOp);
57  if (!callableOp)
58  continue;
59  node = cg.lookupNode(callableOp.getCallableRegion());
60  }
61  if (node)
62  callback(node, use.getUser());
63  }
64 }
65 
66 //===----------------------------------------------------------------------===//
67 // CGUseList
68 //===----------------------------------------------------------------------===//
69 
70 namespace {
71 /// This struct tracks the uses of callgraph nodes that can be dropped when
72 /// use_empty. It directly tracks and manages a use-list for all of the
73 /// call-graph nodes. This is necessary because many callgraph nodes are
74 /// referenced by SymbolRefAttr, which has no mechanism akin to the SSA `Use`
75 /// class.
76 struct CGUseList {
77  /// This struct tracks the uses of callgraph nodes within a specific
78  /// operation.
79  struct CGUser {
80  /// Any nodes referenced in the top-level attribute list of this user. We
81  /// use a set here because the number of references does not matter.
82  DenseSet<CallGraphNode *> topLevelUses;
83 
84  /// Uses of nodes referenced by nested operations.
86  };
87 
88  CGUseList(Operation *op, CallGraph &cg, SymbolTableCollection &symbolTable);
89 
90  /// Drop uses of nodes referred to by the given call operation that resides
91  /// within 'userNode'.
92  void dropCallUses(CallGraphNode *userNode, Operation *callOp, CallGraph &cg);
93 
94  /// Remove the given node from the use list.
95  void eraseNode(CallGraphNode *node);
96 
97  /// Returns true if the given callgraph node has no uses and can be pruned.
98  bool isDead(CallGraphNode *node) const;
99 
100  /// Returns true if the given callgraph node has a single use and can be
101  /// discarded.
102  bool hasOneUseAndDiscardable(CallGraphNode *node) const;
103 
104  /// Recompute the uses held by the given callgraph node.
105  void recomputeUses(CallGraphNode *node, CallGraph &cg);
106 
107  /// Merge the uses of 'lhs' with the uses of the 'rhs' after inlining a copy
108  /// of 'lhs' into 'rhs'.
109  void mergeUsesAfterInlining(CallGraphNode *lhs, CallGraphNode *rhs);
110 
111 private:
112  /// Decrement the uses of discardable nodes referenced by the given user.
113  void decrementDiscardableUses(CGUser &uses);
114 
115  /// A mapping between a discardable callgraph node (that is a symbol) and the
116  /// number of uses for this node.
117  DenseMap<CallGraphNode *, int> discardableSymNodeUses;
118 
119  /// A mapping between a callgraph node and the symbol callgraph nodes that it
120  /// uses.
122 
123  /// A symbol table to use when resolving call lookups.
124  SymbolTableCollection &symbolTable;
125 };
126 } // namespace
127 
128 CGUseList::CGUseList(Operation *op, CallGraph &cg,
129  SymbolTableCollection &symbolTable)
130  : symbolTable(symbolTable) {
131  /// A set of callgraph nodes that are always known to be live during inlining.
132  DenseMap<Attribute, CallGraphNode *> alwaysLiveNodes;
133 
134  // Walk each of the symbol tables looking for discardable callgraph nodes.
135  auto walkFn = [&](Operation *symbolTableOp, bool allUsesVisible) {
136  for (Operation &op : symbolTableOp->getRegion(0).getOps()) {
137  // If this is a callgraph operation, check to see if it is discardable.
138  if (auto callable = dyn_cast<CallableOpInterface>(&op)) {
139  if (auto *node = cg.lookupNode(callable.getCallableRegion())) {
140  SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(&op);
141  if (symbol && (allUsesVisible || symbol.isPrivate()) &&
142  symbol.canDiscardOnUseEmpty()) {
143  discardableSymNodeUses.try_emplace(node, 0);
144  }
145  continue;
146  }
147  }
148  // Otherwise, check for any referenced nodes. These will be always-live.
149  walkReferencedSymbolNodes(&op, cg, symbolTable, alwaysLiveNodes,
150  [](CallGraphNode *, Operation *) {});
151  }
152  };
153  SymbolTable::walkSymbolTables(op, /*allSymUsesVisible=*/!op->getBlock(),
154  walkFn);
155 
156  // Drop the use information for any discardable nodes that are always live.
157  for (auto &it : alwaysLiveNodes)
158  discardableSymNodeUses.erase(it.second);
159 
160  // Compute the uses for each of the callable nodes in the graph.
161  for (CallGraphNode *node : cg)
162  recomputeUses(node, cg);
163 }
164 
165 void CGUseList::dropCallUses(CallGraphNode *userNode, Operation *callOp,
166  CallGraph &cg) {
167  auto &userRefs = nodeUses[userNode].innerUses;
168  auto walkFn = [&](CallGraphNode *node, Operation *user) {
169  auto parentIt = userRefs.find(node);
170  if (parentIt == userRefs.end())
171  return;
172  --parentIt->second;
173  --discardableSymNodeUses[node];
174  };
176  walkReferencedSymbolNodes(callOp, cg, symbolTable, resolvedRefs, walkFn);
177 }
178 
179 void CGUseList::eraseNode(CallGraphNode *node) {
180  // Drop all child nodes.
181  for (auto &edge : *node)
182  if (edge.isChild())
183  eraseNode(edge.getTarget());
184 
185  // Drop the uses held by this node and erase it.
186  auto useIt = nodeUses.find(node);
187  assert(useIt != nodeUses.end() && "expected node to be valid");
188  decrementDiscardableUses(useIt->getSecond());
189  nodeUses.erase(useIt);
190  discardableSymNodeUses.erase(node);
191 }
192 
193 bool CGUseList::isDead(CallGraphNode *node) const {
194  // If the parent operation isn't a symbol, simply check normal SSA deadness.
195  Operation *nodeOp = node->getCallableRegion()->getParentOp();
196  if (!isa<SymbolOpInterface>(nodeOp))
197  return isMemoryEffectFree(nodeOp) && nodeOp->use_empty();
198 
199  // Otherwise, check the number of symbol uses.
200  auto symbolIt = discardableSymNodeUses.find(node);
201  return symbolIt != discardableSymNodeUses.end() && symbolIt->second == 0;
202 }
203 
204 bool CGUseList::hasOneUseAndDiscardable(CallGraphNode *node) const {
205  // If this isn't a symbol node, check for side-effects and SSA use count.
206  Operation *nodeOp = node->getCallableRegion()->getParentOp();
207  if (!isa<SymbolOpInterface>(nodeOp))
208  return isMemoryEffectFree(nodeOp) && nodeOp->hasOneUse();
209 
210  // Otherwise, check the number of symbol uses.
211  auto symbolIt = discardableSymNodeUses.find(node);
212  return symbolIt != discardableSymNodeUses.end() && symbolIt->second == 1;
213 }
214 
215 void CGUseList::recomputeUses(CallGraphNode *node, CallGraph &cg) {
216  Operation *parentOp = node->getCallableRegion()->getParentOp();
217  CGUser &uses = nodeUses[node];
218  decrementDiscardableUses(uses);
219 
220  // Collect the new discardable uses within this node.
221  uses = CGUser();
223  auto walkFn = [&](CallGraphNode *refNode, Operation *user) {
224  auto discardSymIt = discardableSymNodeUses.find(refNode);
225  if (discardSymIt == discardableSymNodeUses.end())
226  return;
227 
228  if (user != parentOp)
229  ++uses.innerUses[refNode];
230  else if (!uses.topLevelUses.insert(refNode).second)
231  return;
232  ++discardSymIt->second;
233  };
234  walkReferencedSymbolNodes(parentOp, cg, symbolTable, resolvedRefs, walkFn);
235 }
236 
237 void CGUseList::mergeUsesAfterInlining(CallGraphNode *lhs, CallGraphNode *rhs) {
238  auto &lhsUses = nodeUses[lhs], &rhsUses = nodeUses[rhs];
239  for (auto &useIt : lhsUses.innerUses) {
240  rhsUses.innerUses[useIt.first] += useIt.second;
241  discardableSymNodeUses[useIt.first] += useIt.second;
242  }
243 }
244 
245 void CGUseList::decrementDiscardableUses(CGUser &uses) {
246  for (CallGraphNode *node : uses.topLevelUses)
247  --discardableSymNodeUses[node];
248  for (auto &it : uses.innerUses)
249  discardableSymNodeUses[it.first] -= it.second;
250 }
251 
252 //===----------------------------------------------------------------------===//
253 // CallGraph traversal
254 //===----------------------------------------------------------------------===//
255 
256 namespace {
257 /// This class represents a specific callgraph SCC.
258 class CallGraphSCC {
259 public:
260  CallGraphSCC(llvm::scc_iterator<const CallGraph *> &parentIterator)
261  : parentIterator(parentIterator) {}
262  /// Return a range over the nodes within this SCC.
263  std::vector<CallGraphNode *>::iterator begin() { return nodes.begin(); }
264  std::vector<CallGraphNode *>::iterator end() { return nodes.end(); }
265 
266  /// Reset the nodes of this SCC with those provided.
267  void reset(const std::vector<CallGraphNode *> &newNodes) { nodes = newNodes; }
268 
269  /// Remove the given node from this SCC.
270  void remove(CallGraphNode *node) {
271  auto it = llvm::find(nodes, node);
272  if (it != nodes.end()) {
273  nodes.erase(it);
274  parentIterator.ReplaceNode(node, nullptr);
275  }
276  }
277 
278 private:
279  std::vector<CallGraphNode *> nodes;
280  llvm::scc_iterator<const CallGraph *> &parentIterator;
281 };
282 } // namespace
283 
284 /// Run a given transformation over the SCCs of the callgraph in a bottom up
285 /// traversal.
286 static LogicalResult runTransformOnCGSCCs(
287  const CallGraph &cg,
288  function_ref<LogicalResult(CallGraphSCC &)> sccTransformer) {
289  llvm::scc_iterator<const CallGraph *> cgi = llvm::scc_begin(&cg);
290  CallGraphSCC currentSCC(cgi);
291  while (!cgi.isAtEnd()) {
292  // Copy the current SCC and increment so that the transformer can modify the
293  // SCC without invalidating our iterator.
294  currentSCC.reset(*cgi);
295  ++cgi;
296  if (failed(sccTransformer(currentSCC)))
297  return failure();
298  }
299  return success();
300 }
301 
302 /// Collect all of the callable operations within the given range of blocks. If
303 /// `traverseNestedCGNodes` is true, this will also collect call operations
304 /// inside of nested callgraph nodes.
306  CallGraphNode *sourceNode, CallGraph &cg,
307  SymbolTableCollection &symbolTable,
309  bool traverseNestedCGNodes) {
311  auto addToWorklist = [&](CallGraphNode *node,
313  for (Block &block : blocks)
314  worklist.emplace_back(&block, node);
315  };
316 
317  addToWorklist(sourceNode, blocks);
318  while (!worklist.empty()) {
319  Block *block;
320  std::tie(block, sourceNode) = worklist.pop_back_val();
321 
322  for (Operation &op : *block) {
323  if (auto call = dyn_cast<CallOpInterface>(op)) {
324  // TODO: Support inlining nested call references.
325  CallInterfaceCallable callable = call.getCallableForCallee();
326  if (SymbolRefAttr symRef = dyn_cast<SymbolRefAttr>(callable)) {
327  if (!isa<FlatSymbolRefAttr>(symRef))
328  continue;
329  }
330 
331  CallGraphNode *targetNode = cg.resolveCallable(call, symbolTable);
332  if (!targetNode->isExternal())
333  calls.emplace_back(call, sourceNode, targetNode);
334  continue;
335  }
336 
337  // If this is not a call, traverse the nested regions. If
338  // `traverseNestedCGNodes` is false, then don't traverse nested call graph
339  // regions.
340  for (auto &nestedRegion : op.getRegions()) {
341  CallGraphNode *nestedNode = cg.lookupNode(&nestedRegion);
342  if (traverseNestedCGNodes || !nestedNode)
343  addToWorklist(nestedNode ? nestedNode : sourceNode, nestedRegion);
344  }
345  }
346  }
347 }
348 
349 //===----------------------------------------------------------------------===//
350 // InlinerInterfaceImpl
351 //===----------------------------------------------------------------------===//
352 
353 #ifndef NDEBUG
354 static std::string getNodeName(CallOpInterface op) {
355  if (llvm::dyn_cast_if_present<SymbolRefAttr>(op.getCallableForCallee()))
356  return debugString(op);
357  return "_unnamed_callee_";
358 }
359 #endif
360 
361 /// Return true if the specified `inlineHistoryID` indicates an inline history
362 /// that already includes `node`.
364  CallGraphNode *node, std::optional<size_t> inlineHistoryID,
365  MutableArrayRef<std::pair<CallGraphNode *, std::optional<size_t>>>
366  inlineHistory) {
367  while (inlineHistoryID.has_value()) {
368  assert(*inlineHistoryID < inlineHistory.size() &&
369  "Invalid inline history ID");
370  if (inlineHistory[*inlineHistoryID].first == node)
371  return true;
372  inlineHistoryID = inlineHistory[*inlineHistoryID].second;
373  }
374  return false;
375 }
376 
377 namespace {
378 /// This class provides a specialization of the main inlining interface.
379 struct InlinerInterfaceImpl : public InlinerInterface {
380  InlinerInterfaceImpl(MLIRContext *context, CallGraph &cg,
381  SymbolTableCollection &symbolTable)
382  : InlinerInterface(context), cg(cg), symbolTable(symbolTable) {}
383 
384  /// Process a set of blocks that have been inlined. This callback is invoked
385  /// *before* inlined terminator operations have been processed.
386  void
388  // Find the closest callgraph node from the first block.
389  CallGraphNode *node;
390  Region *region = inlinedBlocks.begin()->getParent();
391  while (!(node = cg.lookupNode(region))) {
392  region = region->getParentRegion();
393  assert(region && "expected valid parent node");
394  }
395 
396  collectCallOps(inlinedBlocks, node, cg, symbolTable, calls,
397  /*traverseNestedCGNodes=*/true);
398  }
399 
400  /// Mark the given callgraph node for deletion.
401  void markForDeletion(CallGraphNode *node) { deadNodes.insert(node); }
402 
403  /// This method properly disposes of callables that became dead during
404  /// inlining. This should not be called while iterating over the SCCs.
405  void eraseDeadCallables() {
406  for (CallGraphNode *node : deadNodes)
407  node->getCallableRegion()->getParentOp()->erase();
408  }
409 
410  /// The set of callables known to be dead.
412 
413  /// The current set of call instructions to consider for inlining.
415 
416  /// The callgraph being operated on.
417  CallGraph &cg;
418 
419  /// A symbol table to use when resolving call lookups.
420  SymbolTableCollection &symbolTable;
421 };
422 } // namespace
423 
424 namespace mlir {
425 
427 public:
428  Impl(Inliner &inliner) : inliner(inliner) {}
429 
430  /// Attempt to inline calls within the given scc, and run simplifications,
431  /// until a fixed point is reached. This allows for the inlining of newly
432  /// devirtualized calls. Returns failure if there was a fatal error during
433  /// inlining.
434  LogicalResult inlineSCC(InlinerInterfaceImpl &inlinerIface,
435  CGUseList &useList, CallGraphSCC &currentSCC,
436  MLIRContext *context);
437 
438 private:
439  /// Optimize the nodes within the given SCC with one of the held optimization
440  /// pass pipelines. Returns failure if an error occurred during the
441  /// optimization of the SCC, success otherwise.
442  LogicalResult optimizeSCC(CallGraph &cg, CGUseList &useList,
443  CallGraphSCC &currentSCC, MLIRContext *context);
444 
445  /// Optimize the nodes within the given SCC in parallel. Returns failure if an
446  /// error occurred during the optimization of the SCC, success otherwise.
447  LogicalResult optimizeSCCAsync(MutableArrayRef<CallGraphNode *> nodesToVisit,
448  MLIRContext *context);
449 
450  /// Optimize the given callable node with one of the pass managers provided
451  /// with `pipelines`, or the generic pre-inline pipeline. Returns failure if
452  /// an error occurred during the optimization of the callable, success
453  /// otherwise.
454  LogicalResult optimizeCallable(CallGraphNode *node,
455  llvm::StringMap<OpPassManager> &pipelines);
456 
457  /// Attempt to inline calls within the given scc. This function returns
458  /// success if any calls were inlined, failure otherwise.
459  LogicalResult inlineCallsInSCC(InlinerInterfaceImpl &inlinerIface,
460  CGUseList &useList, CallGraphSCC &currentSCC);
461 
462  /// Returns true if the given call should be inlined.
463  bool shouldInline(ResolvedCall &resolvedCall);
464 
465 private:
466  Inliner &inliner;
468 };
469 
470 LogicalResult Inliner::Impl::inlineSCC(InlinerInterfaceImpl &inlinerIface,
471  CGUseList &useList,
472  CallGraphSCC &currentSCC,
473  MLIRContext *context) {
474  // Continuously simplify and inline until we either reach a fixed point, or
475  // hit the maximum iteration count. Simplifying early helps to refine the cost
476  // model, and in future iterations may devirtualize new calls.
477  unsigned iterationCount = 0;
478  do {
479  if (failed(optimizeSCC(inlinerIface.cg, useList, currentSCC, context)))
480  return failure();
481  if (failed(inlineCallsInSCC(inlinerIface, useList, currentSCC)))
482  break;
483  } while (++iterationCount < inliner.config.getMaxInliningIterations());
484  return success();
485 }
486 
487 LogicalResult Inliner::Impl::optimizeSCC(CallGraph &cg, CGUseList &useList,
488  CallGraphSCC &currentSCC,
489  MLIRContext *context) {
490  // Collect the sets of nodes to simplify.
491  SmallVector<CallGraphNode *, 4> nodesToVisit;
492  for (auto *node : currentSCC) {
493  if (node->isExternal())
494  continue;
495 
496  // Don't simplify nodes with children. Nodes with children require special
497  // handling as we may remove the node during simplification. In the future,
498  // we should be able to handle this case with proper node deletion tracking.
499  if (node->hasChildren())
500  continue;
501 
502  // We also won't apply simplifications to nodes that can't have passes
503  // scheduled on them.
504  auto *region = node->getCallableRegion();
506  continue;
507  nodesToVisit.push_back(node);
508  }
509  if (nodesToVisit.empty())
510  return success();
511 
512  // Optimize each of the nodes within the SCC in parallel.
513  if (failed(optimizeSCCAsync(nodesToVisit, context)))
514  return failure();
515 
516  // Recompute the uses held by each of the nodes.
517  for (CallGraphNode *node : nodesToVisit)
518  useList.recomputeUses(node, cg);
519  return success();
520 }
521 
522 LogicalResult
523 Inliner::Impl::optimizeSCCAsync(MutableArrayRef<CallGraphNode *> nodesToVisit,
524  MLIRContext *ctx) {
525  // We must maintain a fixed pool of pass managers which is at least as large
526  // as the maximum parallelism of the failableParallelForEach below.
527  // Note: The number of pass managers here needs to remain constant
528  // to prevent issues with pass instrumentations that rely on having the same
529  // pass manager for the main thread.
530  size_t numThreads = ctx->getNumThreads();
531  const auto &opPipelines = inliner.config.getOpPipelines();
532  if (pipelines.size() < numThreads) {
533  pipelines.reserve(numThreads);
534  pipelines.resize(numThreads, opPipelines);
535  }
536 
537  // Ensure an analysis manager has been constructed for each of the nodes.
538  // This prevents thread races when running the nested pipelines.
539  for (CallGraphNode *node : nodesToVisit)
540  inliner.am.nest(node->getCallableRegion()->getParentOp());
541 
542  // An atomic failure variable for the async executors.
543  std::vector<std::atomic<bool>> activePMs(pipelines.size());
544  std::fill(activePMs.begin(), activePMs.end(), false);
545  return failableParallelForEach(ctx, nodesToVisit, [&](CallGraphNode *node) {
546  // Find a pass manager for this operation.
547  auto it = llvm::find_if(activePMs, [](std::atomic<bool> &isActive) {
548  bool expectedInactive = false;
549  return isActive.compare_exchange_strong(expectedInactive, true);
550  });
551  assert(it != activePMs.end() &&
552  "could not find inactive pass manager for thread");
553  unsigned pmIndex = it - activePMs.begin();
554 
555  // Optimize this callable node.
556  LogicalResult result = optimizeCallable(node, pipelines[pmIndex]);
557 
558  // Reset the active bit for this pass manager.
559  activePMs[pmIndex].store(false);
560  return result;
561  });
562 }
563 
564 LogicalResult
565 Inliner::Impl::optimizeCallable(CallGraphNode *node,
566  llvm::StringMap<OpPassManager> &pipelines) {
567  Operation *callable = node->getCallableRegion()->getParentOp();
568  StringRef opName = callable->getName().getStringRef();
569  auto pipelineIt = pipelines.find(opName);
570  const auto &defaultPipeline = inliner.config.getDefaultPipeline();
571  if (pipelineIt == pipelines.end()) {
572  // If a pipeline didn't exist, use the generic pipeline if possible.
573  if (!defaultPipeline)
574  return success();
575 
576  OpPassManager defaultPM(opName);
577  defaultPipeline(defaultPM);
578  pipelineIt = pipelines.try_emplace(opName, std::move(defaultPM)).first;
579  }
580  return inliner.runPipelineHelper(inliner.pass, pipelineIt->second, callable);
581 }
582 
583 /// Attempt to inline calls within the given scc. This function returns
584 /// success if any calls were inlined, failure otherwise.
585 LogicalResult
586 Inliner::Impl::inlineCallsInSCC(InlinerInterfaceImpl &inlinerIface,
587  CGUseList &useList, CallGraphSCC &currentSCC) {
588  CallGraph &cg = inlinerIface.cg;
589  auto &calls = inlinerIface.calls;
590 
591  // A set of dead nodes to remove after inlining.
592  llvm::SmallSetVector<CallGraphNode *, 1> deadNodes;
593 
594  // Collect all of the direct calls within the nodes of the current SCC. We
595  // don't traverse nested callgraph nodes, because they are handled separately
596  // likely within a different SCC.
597  for (CallGraphNode *node : currentSCC) {
598  if (node->isExternal())
599  continue;
600 
601  // Don't collect calls if the node is already dead.
602  if (useList.isDead(node)) {
603  deadNodes.insert(node);
604  } else {
605  collectCallOps(*node->getCallableRegion(), node, cg,
606  inlinerIface.symbolTable, calls,
607  /*traverseNestedCGNodes=*/false);
608  }
609  }
610 
611  // When inlining a callee produces new call sites, we want to keep track of
612  // the fact that they were inlined from the callee. This allows us to avoid
613  // infinite inlining.
614  using InlineHistoryT = std::optional<size_t>;
616  std::vector<InlineHistoryT> callHistory(calls.size(), InlineHistoryT{});
617 
618  LLVM_DEBUG({
619  llvm::dbgs() << "* Inliner: Initial calls in SCC are: {\n";
620  for (unsigned i = 0, e = calls.size(); i < e; ++i)
621  llvm::dbgs() << " " << i << ". " << calls[i].call << ",\n";
622  llvm::dbgs() << "}\n";
623  });
624 
625  // Try to inline each of the call operations. Don't cache the end iterator
626  // here as more calls may be added during inlining.
627  bool inlinedAnyCalls = false;
628  for (unsigned i = 0; i < calls.size(); ++i) {
629  if (deadNodes.contains(calls[i].sourceNode))
630  continue;
631  ResolvedCall it = calls[i];
632 
633  InlineHistoryT inlineHistoryID = callHistory[i];
634  bool inHistory =
635  inlineHistoryIncludes(it.targetNode, inlineHistoryID, inlineHistory);
636  bool doInline = !inHistory && shouldInline(it);
637  CallOpInterface call = it.call;
638  LLVM_DEBUG({
639  if (doInline)
640  llvm::dbgs() << "* Inlining call: " << i << ". " << call << "\n";
641  else
642  llvm::dbgs() << "* Not inlining call: " << i << ". " << call << "\n";
643  });
644  if (!doInline)
645  continue;
646 
647  unsigned prevSize = calls.size();
648  Region *targetRegion = it.targetNode->getCallableRegion();
649 
650  // If this is the last call to the target node and the node is discardable,
651  // then inline it in-place and delete the node if successful.
652  bool inlineInPlace = useList.hasOneUseAndDiscardable(it.targetNode);
653 
654  LogicalResult inlineResult =
655  inlineCall(inlinerIface, inliner.config.getCloneCallback(), call,
656  cast<CallableOpInterface>(targetRegion->getParentOp()),
657  targetRegion, /*shouldCloneInlinedRegion=*/!inlineInPlace);
658  if (failed(inlineResult)) {
659  LLVM_DEBUG(llvm::dbgs() << "** Failed to inline\n");
660  continue;
661  }
662  inlinedAnyCalls = true;
663 
664  // Create a inline history entry for this inlined call, so that we remember
665  // that new callsites came about due to inlining Callee.
666  InlineHistoryT newInlineHistoryID{inlineHistory.size()};
667  inlineHistory.push_back(std::make_pair(it.targetNode, inlineHistoryID));
668 
669  auto historyToString = [](InlineHistoryT h) {
670  return h.has_value() ? std::to_string(*h) : "root";
671  };
672  (void)historyToString;
673  LLVM_DEBUG(llvm::dbgs()
674  << "* new inlineHistory entry: " << newInlineHistoryID << ". ["
675  << getNodeName(call) << ", " << historyToString(inlineHistoryID)
676  << "]\n");
677 
678  for (unsigned k = prevSize; k != calls.size(); ++k) {
679  callHistory.push_back(newInlineHistoryID);
680  LLVM_DEBUG(llvm::dbgs() << "* new call " << k << " {" << calls[i].call
681  << "}\n with historyID = " << newInlineHistoryID
682  << ", added due to inlining of\n call {" << call
683  << "}\n with historyID = "
684  << historyToString(inlineHistoryID) << "\n");
685  }
686 
687  // If the inlining was successful, Merge the new uses into the source node.
688  useList.dropCallUses(it.sourceNode, call.getOperation(), cg);
689  useList.mergeUsesAfterInlining(it.targetNode, it.sourceNode);
690 
691  // then erase the call.
692  call.erase();
693 
694  // If we inlined in place, mark the node for deletion.
695  if (inlineInPlace) {
696  useList.eraseNode(it.targetNode);
697  deadNodes.insert(it.targetNode);
698  }
699  }
700 
701  for (CallGraphNode *node : deadNodes) {
702  currentSCC.remove(node);
703  inlinerIface.markForDeletion(node);
704  }
705  calls.clear();
706  return success(inlinedAnyCalls);
707 }
708 
709 /// Returns true if the given call should be inlined.
710 bool Inliner::Impl::shouldInline(ResolvedCall &resolvedCall) {
711  // Don't allow inlining terminator calls. We currently don't support this
712  // case.
713  if (resolvedCall.call->hasTrait<OpTrait::IsTerminator>())
714  return false;
715 
716  // Don't allow inlining if the target is a self-recursive function.
717  // Don't allow inlining if the call graph is like A->B->A.
718  if (llvm::count_if(*resolvedCall.targetNode,
719  [&](CallGraphNode::Edge const &edge) -> bool {
720  return edge.getTarget() == resolvedCall.targetNode ||
721  edge.getTarget() == resolvedCall.sourceNode;
722  }) > 0)
723  return false;
724 
725  // Don't allow inlining if the target is an ancestor of the call. This
726  // prevents inlining recursively.
727  Region *callableRegion = resolvedCall.targetNode->getCallableRegion();
728  if (callableRegion->isAncestor(resolvedCall.call->getParentRegion()))
729  return false;
730 
731  // Don't allow inlining if the callee has multiple blocks (unstructured
732  // control flow) but we cannot be sure that the caller region supports that.
733  if (!inliner.config.getCanHandleMultipleBlocks()) {
734  bool calleeHasMultipleBlocks =
735  llvm::hasNItemsOrMore(*callableRegion, /*N=*/2);
736  // If both parent ops have the same type, it is safe to inline. Otherwise,
737  // decide based on whether the op has the SingleBlock trait or not.
738  // Note: This check does currently not account for
739  // SizedRegion/MaxSizedRegion.
740  auto callerRegionSupportsMultipleBlocks = [&]() {
741  return callableRegion->getParentOp()->getName() ==
742  resolvedCall.call->getParentOp()->getName() ||
743  !resolvedCall.call->getParentOp()
744  ->mightHaveTrait<OpTrait::SingleBlock>();
745  };
746  if (calleeHasMultipleBlocks && !callerRegionSupportsMultipleBlocks())
747  return false;
748  }
749 
750  if (!inliner.isProfitableToInline(resolvedCall))
751  return false;
752 
753  // Otherwise, inline.
754  return true;
755 }
756 
757 LogicalResult Inliner::doInlining() {
758  Impl impl(*this);
759  auto *context = op->getContext();
760  // Run the inline transform in post-order over the SCCs in the callgraph.
761  SymbolTableCollection symbolTable;
762  // FIXME: some clean-up can be done for the arguments
763  // of the Impl's methods, if the inlinerIface and useList
764  // become the states of the Impl.
765  InlinerInterfaceImpl inlinerIface(context, cg, symbolTable);
766  CGUseList useList(op, cg, symbolTable);
767  LogicalResult result = runTransformOnCGSCCs(cg, [&](CallGraphSCC &scc) {
768  return impl.inlineSCC(inlinerIface, useList, scc, context);
769  });
770  if (failed(result))
771  return result;
772 
773  // After inlining, make sure to erase any callables proven to be dead.
774  inlinerIface.eraseDeadCallables();
775  return success();
776 }
777 } // namespace mlir
static void collectCallOps(iterator_range< Region::iterator > blocks, CallGraphNode *sourceNode, CallGraph &cg, SymbolTableCollection &symbolTable, SmallVectorImpl< ResolvedCall > &calls, bool traverseNestedCGNodes)
Collect all of the callable operations within the given range of blocks.
Definition: Inliner.cpp:305
Inliner::ResolvedCall ResolvedCall
Definition: Inliner.cpp:32
static void walkReferencedSymbolNodes(Operation *op, CallGraph &cg, SymbolTableCollection &symbolTable, DenseMap< Attribute, CallGraphNode * > &resolvedRefs, function_ref< void(CallGraphNode *, Operation *)> callback)
Walk all of the used symbol callgraph nodes referenced with the given op.
Definition: Inliner.cpp:39
static bool inlineHistoryIncludes(CallGraphNode *node, std::optional< size_t > inlineHistoryID, MutableArrayRef< std::pair< CallGraphNode *, std::optional< size_t >>> inlineHistory)
Return true if the specified inlineHistoryID indicates an inline history that already includes node.
Definition: Inliner.cpp:363
static std::string getNodeName(CallOpInterface op)
Definition: Inliner.cpp:354
static LogicalResult runTransformOnCGSCCs(const CallGraph &cg, function_ref< LogicalResult(CallGraphSCC &)> sccTransformer)
Run a given transformation over the SCCs of the callgraph in a bottom up traversal.
Definition: Inliner.cpp:286
Block represents an ordered list of Operations.
Definition: Block.h:33
This class represents a directed edge between two nodes in the callgraph.
Definition: CallGraph.h:43
This class represents a single callable in the callgraph.
Definition: CallGraph.h:40
bool isExternal() const
Returns true if this node is an external node.
Definition: CallGraph.cpp:32
bool hasChildren() const
Returns true if this node has any child edges.
Definition: CallGraph.cpp:59
Region * getCallableRegion() const
Returns the callable region this node represents.
Definition: CallGraph.cpp:36
CallGraphNode * resolveCallable(CallOpInterface call, SymbolTableCollection &symbolTable) const
Resolve the callable for given callee to a node in the callgraph, or the external node if a valid nod...
Definition: CallGraph.cpp:147
CallGraphNode * lookupNode(Region *region) const
Lookup a call graph node for the given region, or nullptr if none is registered.
Definition: CallGraph.cpp:139
unsigned getMaxInliningIterations() const
Definition: Inliner.h:42
This interface provides the hooks into the inlining interface.
virtual void processInlinedBlocks(iterator_range< Region::iterator > inlinedBlocks)
Process a set of blocks that have been inlined.
LogicalResult inlineSCC(InlinerInterfaceImpl &inlinerIface, CGUseList &useList, CallGraphSCC &currentSCC, MLIRContext *context)
Attempt to inline calls within the given scc, and run simplifications, until a fixed point is reached...
Definition: Inliner.cpp:470
Impl(Inliner &inliner)
Definition: Inliner.cpp:428
This is an implementation of the inliner that operates bottom up over the Strongly Connected Componen...
Definition: Inliner.h:103
LogicalResult doInlining()
Perform inlining on a OpTrait::SymbolTable operation.
Definition: Inliner.cpp:757
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
unsigned getNumThreads()
Return the number of threads used by the thread pool in this context.
This class represents a pass manager that runs passes on either a specific operation type,...
Definition: PassManager.h:47
This class provides the API for ops that are known to be isolated from above.
This class provides the API for ops that are known to be terminators.
Definition: OpDefinition.h:772
StringRef getStringRef() const
Return the name of this operation. This always succeeds.
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
bool use_empty()
Returns true if this operation has no uses.
Definition: Operation.h:853
bool hasTrait()
Returns true if the operation was registered with a particular trait, e.g.
Definition: Operation.h:750
bool hasOneUse()
Returns true if this operation has exactly one use.
Definition: Operation.h:850
MLIRContext * getContext()
Return the context this operation is associated with.
Definition: Operation.h:216
Operation * getParentOp()
Returns the closest surrounding operation that contains this operation or nullptr if this is a top-le...
Definition: Operation.h:234
Block * getBlock()
Returns the operation block that contains this operation.
Definition: Operation.h:213
Region & getRegion(unsigned index)
Returns the region held by this operation at position 'index'.
Definition: Operation.h:687
MutableArrayRef< Region > getRegions()
Returns the regions held by this operation.
Definition: Operation.h:677
OperationName getName()
The name of an operation is the key identifier for it.
Definition: Operation.h:119
void erase()
Remove this operation from its parent block and delete it.
Definition: Operation.cpp:539
This class contains a list of basic blocks and a link to the parent operation it is attached to.
Definition: Region.h:26
iterator_range< OpIterator > getOps()
Definition: Region.h:172
Region * getParentRegion()
Return the region containing this region or nullptr if the region is attached to a top-level operatio...
Definition: Region.cpp:45
bool isAncestor(Region *other)
Return true if this region is ancestor of the other region.
Definition: Region.h:222
Operation * getParentOp()
Return the parent operation this region is attached to.
Definition: Region.h:200
iterator begin()
Definition: Region.h:55
This class represents a collection of SymbolTables.
Definition: SymbolTable.h:283
Operation * lookupNearestSymbolFrom(Operation *from, StringAttr symbol)
Returns the operation registered with the given symbol name within the closest parent operation of,...
This class represents a specific symbol use.
Definition: SymbolTable.h:183
static std::optional< UseRange > getSymbolUses(Operation *from)
Get an iterator range for all of the uses, for any symbol, that are nested within the given operation...
Include the generated interface declarations.
LogicalResult failableParallelForEach(MLIRContext *context, IteratorT begin, IteratorT end, FuncT &&func)
Invoke the given function on the elements between [begin, end) asynchronously.
Definition: Threading.h:36
static std::string debugString(T &&op)
bool isMemoryEffectFree(Operation *op)
Returns true if the given operation is free of memory effects.
LogicalResult inlineCall(InlinerInterface &interface, function_ref< InlinerInterface::CloneCallbackSigTy > cloneCallback, CallOpInterface call, CallableOpInterface callable, Region *src, bool shouldCloneInlinedRegion=true)
This function inlines a given region, 'src', of a callable operation, 'callable', into the location d...
A callable is either a symbol, or an SSA value, that is referenced by a call-like operation.
This struct represents a resolved call to a given callgraph node.
Definition: Inliner.h:109
CallGraphNode * sourceNode
Definition: Inliner.h:114
CallOpInterface call
Definition: Inliner.h:113
CallGraphNode * targetNode
Definition: Inliner.h:114
This class provides APIs and verifiers for ops with regions having a single block.
Definition: OpDefinition.h:880