[mlir] Add support for detecting single use callables in the Inliner.

[rriddle]

Summary

This is somewhat complex(annoying) as it involves directly tracking the uses within each of the callgraph nodes, and updating them as needed during inlining. The benefit of this is that we can have a more exact cost model, enable inlining some otherwise non-inlinable cases, and also ensure that newly dead callables are properly disposed of.

[rriddle]

gentle ping.

[mehdi_amini]

mlir/lib/Transforms/Inliner.cpp

192

I don't quite get the !op->getBlock()?'
(top-level module isn't hidden but isn't in a block)

245

Nit: if you move this above the if you can use it in the if condition (and the if block becomes a one liner)

[rriddle]

Resolve comments

[rriddle]

mlir/lib/Transforms/Inliner.cpp

192

Renamed to allSymUsesVisibile to make it clear. It means that the uses are hidden from other operations, i.e. can't be referenced. It essentially means that we can assume that all symbol uses are visible to the use list.

mlir/examples/toy/Ch4/toyc.cpp

int dumpMLIR() {
 
   if (enableOpt) {
     mlir::PassManager pm(&context);
     // Apply any generic pass manager command line options and run the pipeline.
     applyPassManagerCLOptions(pm);
 
     // Inline all functions into main and then delete them.
     pm.addPass(mlir::createInlinerPass());
-    pm.addPass(mlir::createSymbolDCEPass());
 
     // Now that there is only one function, we can infer the shapes of each of
     // the operations.
     mlir::OpPassManager &optPM = pm.nest<mlir::FuncOp>();
     optPM.addPass(mlir::toy::createShapeInferencePass());
     optPM.addPass(mlir::createCanonicalizerPass());
     optPM.addPass(mlir::createCSEPass());
 

mlir/examples/toy/Ch5/toyc.cpp

int dumpMLIR() {
   applyPassManagerCLOptions(pm);
 
   // Check to see what granularity of MLIR we are compiling to.
   bool isLoweringToAffine = emitAction >= Action::DumpMLIRAffine;
 
   if (enableOpt || isLoweringToAffine) {
     // Inline all functions into main and then delete them.
     pm.addPass(mlir::createInlinerPass());
-    pm.addPass(mlir::createSymbolDCEPass());
 
     // Now that there is only one function, we can infer the shapes of each of
     // the operations.
     mlir::OpPassManager &optPM = pm.nest<mlir::FuncOp>();
     optPM.addPass(mlir::toy::createShapeInferencePass());
     optPM.addPass(mlir::createCanonicalizerPass());
     optPM.addPass(mlir::createCSEPass());
   }

mlir/examples/toy/Ch6/toyc.cpp

int loadAndProcessMLIR(mlir::MLIRContext &context,
 
   // Check to see what granularity of MLIR we are compiling to.
   bool isLoweringToAffine = emitAction >= Action::DumpMLIRAffine;
   bool isLoweringToLLVM = emitAction >= Action::DumpMLIRLLVM;
 
   if (enableOpt || isLoweringToAffine) {
     // Inline all functions into main and then delete them.
     pm.addPass(mlir::createInlinerPass());
-    pm.addPass(mlir::createSymbolDCEPass());
 
     // Now that there is only one function, we can infer the shapes of each of
     // the operations.
     mlir::OpPassManager &optPM = pm.nest<mlir::FuncOp>();
     optPM.addPass(mlir::toy::createShapeInferencePass());
     optPM.addPass(mlir::createCanonicalizerPass());
     optPM.addPass(mlir::createCSEPass());
   }

mlir/examples/toy/Ch7/toyc.cpp

int loadAndProcessMLIR(mlir::MLIRContext &context,
 
   // Check to see what granularity of MLIR we are compiling to.
   bool isLoweringToAffine = emitAction >= Action::DumpMLIRAffine;
   bool isLoweringToLLVM = emitAction >= Action::DumpMLIRLLVM;
 
   if (enableOpt || isLoweringToAffine) {
     // Inline all functions into main and then delete them.
     pm.addPass(mlir::createInlinerPass());
-    pm.addPass(mlir::createSymbolDCEPass());
 
     // Now that there is only one function, we can infer the shapes of each of
     // the operations.
     mlir::OpPassManager &optPM = pm.nest<mlir::FuncOp>();
     optPM.addPass(mlir::createCanonicalizerPass());
     optPM.addPass(mlir::toy::createShapeInferencePass());
     optPM.addPass(mlir::createCanonicalizerPass());
     optPM.addPass(mlir::createCSEPass());

mlir/include/mlir/Analysis/CallGraph.h

public:
   CallGraphNode *getExternalNode() const {
     return const_cast<CallGraphNode *>(&externalNode);
   }
 
   /// Resolve the callable for given callee to a node in the callgraph, or the
   /// external node if a valid node was not resolved.
   CallGraphNode *resolveCallable(CallOpInterface call) const;
 
+  /// Erase the given node from the callgraph.
+  void eraseNode(CallGraphNode *node);
+
   /// An iterator over the nodes of the graph.
   using iterator = NodeIterator;
   iterator begin() const { return nodes.begin(); }
   iterator end() const { return nodes.end(); }
 
   /// Dump the graph in a human readable format.
   void dump() const;
   void print(raw_ostream &os) const;

mlir/lib/Analysis/CallGraph.cpp

CallGraphNode *CallGraph::resolveCallable(CallOpInterface call) const {
   if (auto callableOp = dyn_cast_or_null<CallableOpInterface>(callable))
     if (auto *node = lookupNode(callableOp.getCallableRegion()))
       return node;
 
   // If we don't have a valid direct region, this is an external call.
   return getExternalNode();
 }
 
+/// Erase the given node from the callgraph.
+void CallGraph::eraseNode(CallGraphNode *node) {
+  // Erase any children of this node first.
+  if (node->hasChildren()) {
+    for (const CallGraphNode::Edge &edge : llvm::make_early_inc_range(*node))
+      if (edge.isChild())
+        eraseNode(edge.getTarget());
+  }
+  // Erase any edges to this node from any other nodes.
+  for (auto &it : nodes) {
+    it.second->edges.remove_if([node](const CallGraphNode::Edge &edge) {
+      return edge.getTarget() == node;
+    });
+  }
+  nodes.erase(node->getCallableRegion());
+}
+
 //===----------------------------------------------------------------------===//
 // Printing
 
 /// Dump the graph in a human readable format.
 void CallGraph::dump() const { print(llvm::errs()); }
 void CallGraph::print(raw_ostream &os) const {
   os << "// ---- CallGraph ----\n";
 

mlir/lib/Transforms/Inliner.cpp

 //===- Inliner.cpp - Pass to inline function calls ------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a basic inlining algorithm that operates bottom up over
 // the Strongly Connect Components(SCCs) of the CallGraph. This enables a more
 // incremental propagation of inlining decisions from the leafs to the roots of
 // the callgraph.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/CallGraph.h"
-#include "mlir/IR/Builders.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/SideEffects.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/SCCIterator.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Parallel.h"
 
 #define DEBUG_TYPE "inlining"
 
 using namespace mlir;
 
 static llvm::cl::opt<bool> disableCanonicalization(
     "mlir-disable-inline-simplify",
     llvm::cl::desc("Disable running simplifications during inlining"),
     llvm::cl::ReallyHidden, llvm::cl::init(false));
 
 static llvm::cl::opt<unsigned> maxInliningIterations(
     "mlir-max-inline-iterations",
     llvm::cl::desc("Maximum number of iterations when inlining within an SCC"),
     llvm::cl::ReallyHidden, llvm::cl::init(4));
 
 //===----------------------------------------------------------------------===//
+// Symbol Use Tracking
+//===----------------------------------------------------------------------===//
+
+/// Returns true if this operation can be discarded if it is a symbol and has no
+/// uses. 'allUsesVisible' corresponds to if the parent symbol table is hidden
+/// from above.
+static bool canDiscardSymbolOnUseEmpty(Operation *op, bool allUsesVisible) {
+  if (!SymbolTable::isSymbol(op))
+    return false;
+
+  // TODO: This is essentially the same logic from SymbolDCE. Remove this when
+  // we have a 'Symbol' interface.
+  // Private symbols are always initially considered dead.
+  SymbolTable::Visibility visibility = SymbolTable::getSymbolVisibility(op);
+  if (visibility == mlir::SymbolTable::Visibility::Private)
+    return true;
+  // We only include nested visibility here if all uses are visible.
+  if (allUsesVisible && visibility == SymbolTable::Visibility::Nested)
+    return true;
+  // Otherwise, public symbols are never removable.
+  return false;
+}
+
+/// Walk all of the symbol table operations nested with 'op' along with a
+/// boolean signifying if the symbols within can be treated as if all uses are
+/// visible. The provided callback is invoked with the symbol table operation,
+/// and a boolean signaling if all of the uses within the symbol table are
+/// visible.
+static void walkSymbolTables(Operation *op, bool allSymUsesVisible,
+                             function_ref<void(Operation *, bool)> callback) {
+  if (op->hasTrait<OpTrait::SymbolTable>()) {
+    allSymUsesVisible = allSymUsesVisible || !SymbolTable::isSymbol(op) ||
+                        SymbolTable::getSymbolVisibility(op) ==
+                            SymbolTable::Visibility::Private;
+    callback(op, allSymUsesVisible);
+  } else {
+    // Otherwise if 'op' is not a symbol table, any nested symbols are
+    // guaranteed to be hidden.
+    allSymUsesVisible = true;
+  }
+
+  for (Region &region : op->getRegions())
+    for (Block &block : region)
+      for (Operation &nested : block)
+        walkSymbolTables(&nested, allSymUsesVisible, callback);
+}
+
+/// Walk all of the used symbol callgraph nodes referenced with the given op.
+static void walkReferencedSymbolNodes(
+    Operation *op, CallGraph &cg,
+    DenseMap<Attribute, CallGraphNode *> &resolvedRefs,
+    function_ref<void(CallGraphNode *, Operation *)> callback) {
+  auto symbolUses = SymbolTable::getSymbolUses(op);
+  assert(symbolUses && "expected uses to be valid");
+
+  Operation *symbolTableOp = op->getParentOp();
+  for (const SymbolTable::SymbolUse &use : *symbolUses) {
+    auto refIt = resolvedRefs.insert({use.getSymbolRef(), nullptr});
+    CallGraphNode *&node = refIt.first->second;
+
+    // If this is the first instance of this reference, try to resolve a
+    // callgraph node for it.
+    if (refIt.second) {
+      auto *symbolOp = SymbolTable::lookupNearestSymbolFrom(symbolTableOp,
+                                                            use.getSymbolRef());
+      auto callableOp = dyn_cast_or_null<CallableOpInterface>(symbolOp);
+      if (!callableOp)
+        continue;
+      node = cg.lookupNode(callableOp.getCallableRegion());
+    }
+    if (node)
+      callback(node, use.getUser());
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// CGUseList
+
+namespace {
+/// This struct tracks the uses of callgraph nodes that can be dropped when
+/// use_empty. It directly tracks and manages a use-list for all of the
+/// call-graph nodes. This is necessary because many callgraph nodes are
+/// referenced by SymbolRefAttr, which has no mechanism akin to the SSA `Use`
+/// class.
+struct CGUseList {
+  /// This struct tracks the uses of callgraph nodes within a specific
+  /// operation.
+  struct CGUser {
+    /// Any nodes referenced in the top-level attribute list of this user. We
+    /// use a set here because the number of references does not matter.
+    DenseSet<CallGraphNode *> topLevelUses;
+
+    /// Uses of nodes referenced by nested operations.
+    DenseMap<CallGraphNode *, int> innerUses;
+  };
+
+  CGUseList(Operation *op, CallGraph &cg);
+
+  /// Drop uses of nodes referred to by the given call operation that resides
+  /// within 'userNode'.
+  void dropCallUses(CallGraphNode *userNode, Operation *callOp, CallGraph &cg);
+
+  /// Remove the given node from the use list.
+  void eraseNode(CallGraphNode *node);
+
+  /// Returns true if the given callgraph node has no uses and can be pruned.
+  bool isDead(CallGraphNode *node) const;
+
+  /// Returns true if the given callgraph node has a single use and can be
+  /// discarded.
+  bool hasOneUseAndDiscardable(CallGraphNode *node) const;
+
+  /// Recompute the uses held by the given callgraph node.
+  void recomputeUses(CallGraphNode *node, CallGraph &cg);
+
+  /// Merge the uses of 'lhs' with the uses of the 'rhs' after inlining a copy
+  /// of 'lhs' into 'rhs'.
+  void mergeUsesAfterInlining(CallGraphNode *lhs, CallGraphNode *rhs);
+
+private:
+  /// Decrement the uses of discardable nodes referenced by the given user.
+  void decrementDiscardableUses(CGUser &uses);
+
+  /// A mapping between a discardable callgraph node (that is a symbol) and the
+  /// number of uses for this node.
+  DenseMap<CallGraphNode *, int> discardableSymNodeUses;
+  /// A mapping between a callgraph node and the symbol callgraph nodes that it
+  /// uses.
+  DenseMap<CallGraphNode *, CGUser> nodeUses;
+};
+} // end anonymous namespace
+
+CGUseList::CGUseList(Operation *op, CallGraph &cg) {
+  /// A set of callgraph nodes that are always known to be live during inlining.
+  DenseMap<Attribute, CallGraphNode *> alwaysLiveNodes;
+
+  // Walk each of the symbol tables looking for discardable callgraph nodes.
+  auto walkFn = [&](Operation *symbolTableOp, bool allUsesVisible) {
+    for (Block &block : symbolTableOp->getRegion(0)) {
+      for (Operation &op : block) {
+        // If this is a callgraph operation, check to see if it is discardable.
+        if (auto callable = dyn_cast<CallableOpInterface>(&op)) {
+          if (auto *node = cg.lookupNode(callable.getCallableRegion())) {
+            if (canDiscardSymbolOnUseEmpty(&op, allUsesVisible))
+              discardableSymNodeUses.try_emplace(node, 0);
+            continue;
+          }
+        }
+        // Otherwise, check for any referenced nodes. These will be always-live.
+        walkReferencedSymbolNodes(&op, cg, alwaysLiveNodes,
+                                  [](CallGraphNode *, Operation *) {});
+      }
+    }
+  };
+  walkSymbolTables(op, /*allSymUsesVisible=*/!op->getBlock(), walkFn);
+
+  // Drop the use information for any discardable nodes that are always live.
+  for (auto &it : alwaysLiveNodes)
+    discardableSymNodeUses.erase(it.second);
+
+  // Compute the uses for each of the callable nodes in the graph.
+  for (CallGraphNode *node : cg)
+    recomputeUses(node, cg);
+}
+
+void CGUseList::dropCallUses(CallGraphNode *userNode, Operation *callOp,
+                             CallGraph &cg) {
+  auto &userRefs = nodeUses[userNode].innerUses;
+  auto walkFn = [&](CallGraphNode *node, Operation *user) {
+    auto parentIt = userRefs.find(node);
+    if (parentIt == userRefs.end())
+      return;
+    --parentIt->second;
+    --discardableSymNodeUses[node];
+  };
+  DenseMap<Attribute, CallGraphNode *> resolvedRefs;
+  walkReferencedSymbolNodes(callOp, cg, resolvedRefs, walkFn);
+}
+
+void CGUseList::eraseNode(CallGraphNode *node) {
+  // Drop all child nodes.
+  for (auto &edge : *node)
+    if (edge.isChild())
+      eraseNode(edge.getTarget());
+
+  // Drop the uses held by this node and erase it.
+  auto useIt = nodeUses.find(node);
+  assert(useIt != nodeUses.end() && "expected node to be valid");
+  decrementDiscardableUses(useIt->getSecond());
+  nodeUses.erase(useIt);
+  discardableSymNodeUses.erase(node);
+}
+
+bool CGUseList::isDead(CallGraphNode *node) const {
+  // If the parent operation isn't a symbol, simply check normal SSA deadness.
+  Operation *nodeOp = node->getCallableRegion()->getParentOp();
+  if (!SymbolTable::isSymbol(nodeOp))
+    return MemoryEffectOpInterface::hasNoEffect(nodeOp) && nodeOp->use_empty();
+
+  // Otherwise, check the number of symbol uses.
+  auto symbolIt = discardableSymNodeUses.find(node);
+  return symbolIt != discardableSymNodeUses.end() && symbolIt->second == 0;
+}
+
+bool CGUseList::hasOneUseAndDiscardable(CallGraphNode *node) const {
+  // If this isn't a symbol node, check for side-effects and SSA use count.
+  Operation *nodeOp = node->getCallableRegion()->getParentOp();
+  if (!SymbolTable::isSymbol(nodeOp))
+    return MemoryEffectOpInterface::hasNoEffect(nodeOp) && nodeOp->hasOneUse();
+
+  // Otherwise, check the number of symbol uses.
+  auto symbolIt = discardableSymNodeUses.find(node);
+  return symbolIt != discardableSymNodeUses.end() && symbolIt->second == 1;
+}
+
+void CGUseList::recomputeUses(CallGraphNode *node, CallGraph &cg) {
+  Operation *parentOp = node->getCallableRegion()->getParentOp();
+  CGUser &uses = nodeUses[node];
+  decrementDiscardableUses(uses);
+
+  // Collect the new discardable uses within this node.
+  uses = CGUser();
+  DenseMap<Attribute, CallGraphNode *> resolvedRefs;
+  auto walkFn = [&](CallGraphNode *refNode, Operation *user) {
+    auto discardSymIt = discardableSymNodeUses.find(refNode);
+    if (discardSymIt == discardableSymNodeUses.end())
+      return;
+
+    if (user != parentOp)
+      ++uses.innerUses[refNode];
+    else if (!uses.topLevelUses.insert(refNode).second)
+      return;
+    ++discardSymIt->second;
+  };
+  walkReferencedSymbolNodes(parentOp, cg, resolvedRefs, walkFn);
+}
+
+void CGUseList::mergeUsesAfterInlining(CallGraphNode *lhs, CallGraphNode *rhs) {
+  auto &lhsUses = nodeUses[lhs], &rhsUses = nodeUses[rhs];
+  for (auto &useIt : lhsUses.innerUses) {
+    rhsUses.innerUses[useIt.first] += useIt.second;
+    discardableSymNodeUses[useIt.first] += useIt.second;
+  }
+}
+
+void CGUseList::decrementDiscardableUses(CGUser &uses) {
+  for (CallGraphNode *node : uses.topLevelUses)
+    --discardableSymNodeUses[node];
+  for (auto &it : uses.innerUses)
+    discardableSymNodeUses[it.first] -= it.second;
+}
+
+//===----------------------------------------------------------------------===//
 // CallGraph traversal
 //===----------------------------------------------------------------------===//
 
 /// Run a given transformation over the SCCs of the callgraph in a bottom up
 /// traversal.
 static void runTransformOnCGSCCs(
     const CallGraph &cg,
-    function_ref<void(ArrayRef<CallGraphNode *>)> sccTransformer) {
+    function_ref<void(MutableArrayRef<CallGraphNode *>)> sccTransformer) {
   std::vector<CallGraphNode *> currentSCCVec;
   auto cgi = llvm::scc_begin(&cg);
   while (!cgi.isAtEnd()) {
     // Copy the current SCC and increment so that the transformer can modify the
     // SCC without invalidating our iterator.
     currentSCCVec = *cgi;
     ++cgi;
     sccTransformer(currentSCCVec);
   }
 }
 
 namespace {
 /// This struct represents a resolved call to a given callgraph node. Given that
 /// the call does not actually contain a direct reference to the
 /// Region(CallGraphNode) that it is dispatching to, we need to resolve them
 /// explicitly.
 struct ResolvedCall {
-  ResolvedCall(CallOpInterface call, CallGraphNode *targetNode)
-      : call(call), targetNode(targetNode) {}
+  ResolvedCall(CallOpInterface call, CallGraphNode *sourceNode,
+               CallGraphNode *targetNode)
+      : call(call), sourceNode(sourceNode), targetNode(targetNode) {}
   CallOpInterface call;
-  CallGraphNode *targetNode;
+  CallGraphNode *sourceNode, *targetNode;
 };
 } // end anonymous namespace
 
 /// Collect all of the callable operations within the given range of blocks. If
 /// `traverseNestedCGNodes` is true, this will also collect call operations
 /// inside of nested callgraph nodes.
 static void collectCallOps(iterator_range<Region::iterator> blocks,
-                           CallGraph &cg, SmallVectorImpl<ResolvedCall> &calls,
+                           CallGraphNode *sourceNode, CallGraph &cg,
+                           SmallVectorImpl<ResolvedCall> &calls,
                            bool traverseNestedCGNodes) {
-  SmallVector<Block *, 8> worklist;
-  auto addToWorklist = [&](iterator_range<Region::iterator> blocks) {
+  SmallVector<std::pair<Block *, CallGraphNode *>, 8> worklist;
+  auto addToWorklist = [&](CallGraphNode *node,
+                           iterator_range<Region::iterator> blocks) {
     for (Block &block : blocks)
-      worklist.push_back(&block);
+      worklist.emplace_back(&block, node);
   };
 
-  addToWorklist(blocks);
+  addToWorklist(sourceNode, blocks);
   while (!worklist.empty()) {
-    for (Operation &op : *worklist.pop_back_val()) {
+    Block *block;
+    std::tie(block, sourceNode) = worklist.pop_back_val();
+
+    for (Operation &op : *block) {
       if (auto call = dyn_cast<CallOpInterface>(op)) {
         // TODO(riverriddle) Support inlining nested call references.
         CallInterfaceCallable callable = call.getCallableForCallee();
         if (SymbolRefAttr symRef = callable.dyn_cast<SymbolRefAttr>()) {
           if (!symRef.isa<FlatSymbolRefAttr>())
             continue;
         }
 
-        CallGraphNode *node = cg.resolveCallable(call);
-        if (!node->isExternal())
-          calls.emplace_back(call, node);
+        CallGraphNode *targetNode = cg.resolveCallable(call);
+        if (!targetNode->isExternal())
+          calls.emplace_back(call, sourceNode, targetNode);
         continue;
       }
 
       // If this is not a call, traverse the nested regions. If
       // `traverseNestedCGNodes` is false, then don't traverse nested call graph
       // regions.
-      for (auto &nestedRegion : op.getRegions())
-        if (traverseNestedCGNodes || !cg.lookupNode(&nestedRegion))
-          addToWorklist(nestedRegion);
+      for (auto &nestedRegion : op.getRegions()) {
+        CallGraphNode *nestedNode = cg.lookupNode(&nestedRegion);
+        if (traverseNestedCGNodes || !nestedNode)
+          addToWorklist(nestedNode ? nestedNode : sourceNode, nestedRegion);
+      }
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Inliner
 //===----------------------------------------------------------------------===//
 namespace {
 /// This class provides a specialization of the main inlining interface.
 struct Inliner : public InlinerInterface {
   Inliner(MLIRContext *context, CallGraph &cg)
       : InlinerInterface(context), cg(cg) {}
 
   /// Process a set of blocks that have been inlined. This callback is invoked
   /// *before* inlined terminator operations have been processed.
   void
   processInlinedBlocks(iterator_range<Region::iterator> inlinedBlocks) final {
-    collectCallOps(inlinedBlocks, cg, calls, /*traverseNestedCGNodes=*/true);
+    // Find the closest callgraph node from the first block.
+    CallGraphNode *node;
+    Region *region = inlinedBlocks.begin()->getParent();
+    while (!(node = cg.lookupNode(region))) {
+      region = region->getParentRegion();
+      assert(region && "expected valid parent node");
+    }
+
+    collectCallOps(inlinedBlocks, node, cg, calls,
+                   /*traverseNestedCGNodes=*/true);
   }
 
   /// The current set of call instructions to consider for inlining.
   SmallVector<ResolvedCall, 8> calls;
 
   /// The callgraph being operated on.
   CallGraph &cg;
 };
static bool shouldInline(ResolvedCall &resolvedCall) {
   if (resolvedCall.targetNode->getCallableRegion()->isAncestor(
           resolvedCall.call.getParentRegion()))
     return false;
 
   // Otherwise, inline.
   return true;
 }
 
+/// Delete the given node and remove it from the current scc and the callgraph.
+static void deleteNode(CallGraphNode *node, CGUseList &useList, CallGraph &cg,
+                       MutableArrayRef<CallGraphNode *> currentSCC) {
+  // Erase the parent operation and remove it from the various lists.
+  node->getCallableRegion()->getParentOp()->erase();
+  cg.eraseNode(node);
+
+  // Replace this node in the currentSCC with the external node.
+  auto it = llvm::find(currentSCC, node);
+  if (it != currentSCC.end())
+    *it = cg.getExternalNode();
+}
+
 /// Attempt to inline calls within the given scc. This function returns
 /// success if any calls were inlined, failure otherwise.
-static LogicalResult inlineCallsInSCC(Inliner &inliner,
-                                      ArrayRef<CallGraphNode *> currentSCC) {
+static LogicalResult
+inlineCallsInSCC(Inliner &inliner, CGUseList &useList,
+                 MutableArrayRef<CallGraphNode *> currentSCC) {
   CallGraph &cg = inliner.cg;
   auto &calls = inliner.calls;
 
   // Collect all of the direct calls within the nodes of the current SCC. We
   // don't traverse nested callgraph nodes, because they are handled separately
   // likely within a different SCC.
-  for (auto *node : currentSCC) {
-    if (!node->isExternal())
-      collectCallOps(*node->getCallableRegion(), cg, calls,
+  for (CallGraphNode *node : currentSCC) {
+    if (node->isExternal())
+      continue;
+
+    // If this node is dead, just delete it now.
+    if (useList.isDead(node))
+      deleteNode(node, useList, cg, currentSCC);
+    else
+      collectCallOps(*node->getCallableRegion(), node, cg, calls,
                      /*traverseNestedCGNodes=*/false);
   }
   if (calls.empty())
     return failure();
 
+  // A set of dead nodes to remove after inlining.
+  SmallVector<CallGraphNode *, 1> deadNodes;
+
   // Try to inline each of the call operations. Don't cache the end iterator
   // here as more calls may be added during inlining.
   bool inlinedAnyCalls = false;
   for (unsigned i = 0; i != calls.size(); ++i) {
     ResolvedCall &it = calls[i];
     LLVM_DEBUG({
       llvm::dbgs() << "* Considering inlining call: ";
       it.call.dump();
     });
     if (!shouldInline(it))
       continue;
-
     CallOpInterface call = it.call;
     Region *targetRegion = it.targetNode->getCallableRegion();
+
+    // If this is the last call to the target node and the node is discardable,
+    // then inline it in-place and delete the node if successful.
+    bool inlineInPlace = useList.hasOneUseAndDiscardable(it.targetNode);
+
     LogicalResult inlineResult = inlineCall(
         inliner, call, cast<CallableOpInterface>(targetRegion->getParentOp()),
-        targetRegion);
+        targetRegion, /*shouldCloneInlinedRegion=*/!inlineInPlace);
     if (failed(inlineResult))
       continue;
+    inlinedAnyCalls = true;
 
-    // If the inlining was successful, then erase the call.
+    // If the inlining was successful, Merge the new uses into the source node.
+    useList.dropCallUses(it.sourceNode, call.getOperation(), cg);
+    useList.mergeUsesAfterInlining(it.targetNode, it.sourceNode);
+
+    // then erase the call.
     call.erase();
-    inlinedAnyCalls = true;
+
+    // If we inlined in place, mark the node for deletion.
+    if (inlineInPlace) {
+      useList.eraseNode(it.targetNode);
+      deadNodes.push_back(it.targetNode);
+    }
   }
+
+  for (CallGraphNode *node : deadNodes)
+    deleteNode(node, useList, cg, currentSCC);
   calls.clear();
   return success(inlinedAnyCalls);
 }
 
 /// Canonicalize the nodes within the given SCC with the given set of
 /// canonicalization patterns.
-static void canonicalizeSCC(CallGraph &cg, ArrayRef<CallGraphNode *> currentSCC,
+static void canonicalizeSCC(CallGraph &cg, CGUseList &useList,
+                            MutableArrayRef<CallGraphNode *> currentSCC,
                             MLIRContext *context,
                             const OwningRewritePatternList &canonPatterns) {
   // Collect the sets of nodes to canonicalize.
   SmallVector<CallGraphNode *, 4> nodesToCanonicalize;
   for (auto *node : currentSCC) {
     // Don't canonicalize the external node, it has no valid callable region.
     if (node->isExternal())
       continue;
llvm::parallel::for_each_n(
         // Apply the canonicalization patterns to this region.
         auto *node = nodesToCanonicalize[index];
         applyPatternsGreedily(*node->getCallableRegion(), canonPatterns);
 
         // Make sure to reset the order ID for the diagnostic handler, as this
         // thread may be used in a different context.
         canonicalizationHandler.eraseOrderIDForThread();
       });
+
+  // Recompute the uses held by each of the nodes.
+  for (CallGraphNode *node : nodesToCanonicalize)
+    useList.recomputeUses(node, cg);
 }
 
 /// Attempt to inline calls within the given scc, and run canonicalizations with
 /// the given patterns, until a fixed point is reached. This allows for the
 /// inlining of newly devirtualized calls.
-static void inlineSCC(Inliner &inliner, ArrayRef<CallGraphNode *> currentSCC,
+static void inlineSCC(Inliner &inliner, CGUseList &useList,
+                      MutableArrayRef<CallGraphNode *> currentSCC,
                       MLIRContext *context,
                       const OwningRewritePatternList &canonPatterns) {
   // If we successfully inlined any calls, run some simplifications on the
   // nodes of the scc. Continue attempting to inline until we reach a fixed
   // point, or a maximum iteration count. We canonicalize here as it may
   // devirtualize new calls, as well as give us a better cost model.
   unsigned iterationCount = 0;
-  while (succeeded(inlineCallsInSCC(inliner, currentSCC))) {
+  while (succeeded(inlineCallsInSCC(inliner, useList, currentSCC))) {
     // If we aren't allowing simplifications or the max iteration count was
     // reached, then bail out early.
     if (disableCanonicalization || ++iterationCount >= maxInliningIterations)
       break;
-    canonicalizeSCC(inliner.cg, currentSCC, context, canonPatterns);
+    canonicalizeSCC(inliner.cg, useList, currentSCC, context, canonPatterns);
   }
 }
 
 //===----------------------------------------------------------------------===//
 // InlinerPass
 //===----------------------------------------------------------------------===//
 
-// TODO(riverriddle) This pass should currently only be used for basic testing
-// of inlining functionality.
 namespace {
 struct InlinerPass : public OperationPass<InlinerPass> {
   void runOnOperation() override {
     CallGraph &cg = getAnalysis<CallGraph>();
     auto *context = &getContext();
 
     // The inliner should only be run on operations that define a symbol table,
     // as the callgraph will need to resolve references.
     Operation *op = getOperation();
     if (!op->hasTrait<OpTrait::SymbolTable>()) {
       op->emitOpError() << " was scheduled to run under the inliner, but does "
                            "not define a symbol table";
       return signalPassFailure();
     }
 
     // Collect a set of canonicalization patterns to use when simplifying
     // callable regions within an SCC.
     OwningRewritePatternList canonPatterns;
     for (auto *op : context->getRegisteredOperations())
       op->getCanonicalizationPatterns(canonPatterns, context);
 
     // Run the inline transform in post-order over the SCCs in the callgraph.
     Inliner inliner(context, cg);
-    runTransformOnCGSCCs(cg, [&](ArrayRef<CallGraphNode *> scc) {
-      inlineSCC(inliner, scc, context, canonPatterns);
+    CGUseList useList(getOperation(), cg);
+    runTransformOnCGSCCs(cg, [&](MutableArrayRef<CallGraphNode *> scc) {
+      inlineSCC(inliner, useList, scc, context, canonPatterns);
     });
   }
 };
 } // end anonymous namespace
 
 std::unique_ptr<Pass> mlir::createInlinerPass() {
   return std::make_unique<InlinerPass>();
 }
 
 static PassRegistration<InlinerPass> pass("inline", "Inline function calls");

mlir/test/Transforms/inlining-dce.mlir

+// RUN: mlir-opt %s -inline | FileCheck %s
+
+// This file tests the callgraph dead code elimination performed by the inliner.
+
+// Function is already dead.
+// CHECK-NOT: func @dead_function
+func @dead_function() attributes {sym_visibility = "private"} {
+  return
+}
+
+// Function becomes dead after inlining.
+// CHECK-NOT: func @dead_function_b
+func @dead_function_b() attributes {sym_visibility = "private"} {
+  return
+}
+
+// CHECK: func @live_function()
+func @live_function() {
+  call @dead_function_b() : () -> ()
+  return
+}
+
+// Same as above, but a transitive example.
+
+// CHECK: func @live_function_b
+func @live_function_b() {
+  return
+}
+// CHECK-NOT: func @dead_function_c
+func @dead_function_c() attributes {sym_visibility = "private"} {
+  call @live_function_b() : () -> ()
+  return
+}
+// CHECK-NOT: func @dead_function_d
+func @dead_function_d() attributes {sym_visibility = "private"} {
+  call @dead_function_c() : () -> ()
+  call @dead_function_c() : () -> ()
+  return
+}
+// CHECK: func @live_function_c
+func @live_function_c() {
+  call @dead_function_c() : () -> ()
+  call @dead_function_d() : () -> ()
+  return
+}
+
+// Function is referenced by non-callable top-level user.
+// CHECK: func @live_function_d
+func @live_function_d() attributes {sym_visibility = "private"} {
+  return
+}
+
+"live.user"() {use = @live_function_d} : () -> ()