[mlir][Inliner] Properly handle callgraph node deletion

[rriddle]

Summary

We previously weren't properly updating the SCC iterator when nodes were removed, leading to asan failures in certain situations. This commit adds a CallGraphSCC class and defers operation deletion until inlining has finished.

[jpienaar]

Nice, thanks for finding and fixing this!

mlir/lib/Transforms/Inliner.cpp

378

Is there a requirement on when this may be called? Or not really due to the helper class?

594

Would these just be "newly dead" ones or also ones that were dead before any inlinining? (E.g., does this overlap with symboldce?)

[rriddle]

mlir/lib/Transforms/Inliner.cpp

594

Yes, It includes those that are already dead, because it is "free" to detect when inlining.

mlir/lib/Transforms/Inliner.cpp

void CGUseList::decrementDiscardableUses(CGUser &uses) {
   for (auto &it : uses.innerUses)
     discardableSymNodeUses[it.first] -= it.second;
 }
 
 //===----------------------------------------------------------------------===//
 // CallGraph traversal
 //===----------------------------------------------------------------------===//
 
+namespace {
+/// This class represents a specific callgraph SCC.
+class CallGraphSCC {
+public:
+  CallGraphSCC(llvm::scc_iterator<const CallGraph *> &parentIterator)
+      : parentIterator(parentIterator) {}
+  /// Return a range over the nodes within this SCC.
+  std::vector<CallGraphNode *>::iterator begin() { return nodes.begin(); }
+  std::vector<CallGraphNode *>::iterator end() { return nodes.end(); }
+
+  /// Reset the nodes of this SCC with those provided.
+  void reset(const std::vector<CallGraphNode *> &newNodes) { nodes = newNodes; }
+
+  /// Remove the given node from this SCC.
+  void remove(CallGraphNode *node) {
+    auto it = llvm::find(nodes, node);
+    if (it != nodes.end()) {
+      nodes.erase(it);
+      parentIterator.ReplaceNode(node, nullptr);
+    }
+  }
+
+private:
+  std::vector<CallGraphNode *> nodes;
+  llvm::scc_iterator<const CallGraph *> &parentIterator;
+};
+} // end anonymous namespace
+
 /// Run a given transformation over the SCCs of the callgraph in a bottom up
 /// traversal.
-static void runTransformOnCGSCCs(
-    const CallGraph &cg,
-    function_ref<void(MutableArrayRef<CallGraphNode *>)> sccTransformer) {
-  std::vector<CallGraphNode *> currentSCCVec;
-  auto cgi = llvm::scc_begin(&cg);
+static void
+runTransformOnCGSCCs(const CallGraph &cg,
+                     function_ref<void(CallGraphSCC &)> sccTransformer) {
+  llvm::scc_iterator<const CallGraph *> cgi = llvm::scc_begin(&cg);
+  CallGraphSCC currentSCC(cgi);
   while (!cgi.isAtEnd()) {
     // Copy the current SCC and increment so that the transformer can modify the
     // SCC without invalidating our iterator.
-    currentSCCVec = *cgi;
+    currentSCC.reset(*cgi);
     ++cgi;
-    sccTransformer(currentSCCVec);
+    sccTransformer(currentSCC);
   }
 }
 
 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.
while (!(node = cg.lookupNode(region))) {
       region = region->getParentRegion();
       assert(region && "expected valid parent node");
     }
 
     collectCallOps(inlinedBlocks, node, cg, calls,
                    /*traverseNestedCGNodes=*/true);
   }
 
+  /// Mark the given callgraph node for deletion.
+  void markForDeletion(CallGraphNode *node) { deadNodes.insert(node); }
+
+  /// This method properly disposes of callables that became dead during
+  /// inlining.
+  void eraseDeadCallables() {
+    for (CallGraphNode *node : deadNodes)
+      node->getCallableRegion()->getParentOp()->erase();
+  }
+
+  /// The set of callables known to be dead.
+  SmallPtrSet<CallGraphNode *, 8> deadNodes;
+
   /// The current set of call instructions to consider for inlining.
   SmallVector<ResolvedCall, 8> calls;
 
   /// The callgraph being operated on.
   CallGraph &cg;
 };
 } // namespace
 
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, CGUseList &useList,
-                 MutableArrayRef<CallGraphNode *> currentSCC) {
+static LogicalResult inlineCallsInSCC(Inliner &inliner, CGUseList &useList,
+                                      CallGraphSCC &currentSCC) {
   CallGraph &cg = inliner.cg;
   auto &calls = inliner.calls;
 
+  // A set of dead nodes to remove after inlining.
+  SmallVector<CallGraphNode *, 1> deadNodes;
+
   // 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 (CallGraphNode *node : currentSCC) {
     if (node->isExternal())
       continue;
 
-    // If this node is dead, just delete it now.
+    // Don't collect calls if the node is already dead.
     if (useList.isDead(node))
-      deleteNode(node, useList, cg, currentSCC);
+      deadNodes.push_back(node);
     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];
     bool doInline = shouldInline(it);
     CallOpInterface call = it.call;
for (unsigned i = 0; i != calls.size(); ++i) {
 
     // 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);
+  for (CallGraphNode *node : deadNodes) {
+    currentSCC.remove(node);
+    inliner.markForDeletion(node);
+  }
   calls.clear();
   return success(inlinedAnyCalls);
 }
 
 /// Canonicalize the nodes within the given SCC with the given set of
 /// canonicalization patterns.
 static void canonicalizeSCC(CallGraph &cg, CGUseList &useList,
-                            MutableArrayRef<CallGraphNode *> currentSCC,
-                            MLIRContext *context,
+                            CallGraphSCC &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;
 

 
 namespace {
 struct InlinerPass : public InlinerBase<InlinerPass> {
   void runOnOperation() override;
 
   /// 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.
-  void inlineSCC(Inliner &inliner, CGUseList &useList,
-                 MutableArrayRef<CallGraphNode *> currentSCC,
+  void inlineSCC(Inliner &inliner, CGUseList &useList, CallGraphSCC &currentSCC,
                  MLIRContext *context,
                  const OwningRewritePatternList &canonPatterns);
 };
 } // end anonymous namespace
 
 void InlinerPass::runOnOperation() {
   CallGraph &cg = getAnalysis<CallGraph>();
   auto *context = &getContext();
void InlinerPass::runOnOperation() {
   // 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);
   CGUseList useList(getOperation(), cg);
-  runTransformOnCGSCCs(cg, [&](MutableArrayRef<CallGraphNode *> scc) {
+  runTransformOnCGSCCs(cg, [&](CallGraphSCC &scc) {
     inlineSCC(inliner, useList, scc, context, canonPatterns);
   });
+
+  // After inlining, make sure to erase any callables proven to be dead.
+  inliner.eraseDeadCallables();
 }
 
 void InlinerPass::inlineSCC(Inliner &inliner, CGUseList &useList,
-                            MutableArrayRef<CallGraphNode *> currentSCC,
-                            MLIRContext *context,
+                            CallGraphSCC &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, useList, currentSCC))) {
     // If we aren't allowing simplifications or the max iteration count was

mlir/test/Transforms/inlining-dce.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s -inline | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -inline -split-input-file | 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 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} : () -> ()
+
+// -----
+
+// This test checks that the inliner can properly handle the deletion of
+// functions in different SCCs that are referenced by calls materialized during
+// canonicalization.
+// CHECK: func @live_function_e
+func @live_function_e() {
+  call @dead_function_e() : () -> ()
+  return
+}
+// CHECK-NOT: func @dead_function_e
+func @dead_function_e() -> () attributes {sym_visibility = "private"} {
+  "test.fold_to_call_op"() {callee=@dead_function_f} : () -> ()
+  return
+}
+// CHECK-NOT: func @dead_function_f
+func @dead_function_f() attributes {sym_visibility = "private"} {
+  return
+}

mlir/test/lib/Dialect/Test/TestDialect.cpp

 
 Optional<MutableOperandRange>
 TestBranchOp::getMutableSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
   return targetOperandsMutable();
 }
 
 //===----------------------------------------------------------------------===//
+// TestFoldToCallOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct FoldToCallOpPattern : public OpRewritePattern<FoldToCallOp> {
+  using OpRewritePattern<FoldToCallOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(FoldToCallOp op,
+                                PatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<CallOp>(op, ArrayRef<Type>(), op.calleeAttr(),
+                                        ValueRange());
+    return success();
+  }
+};
+} // end anonymous namespace
+
+void FoldToCallOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert<FoldToCallOpPattern>(context);
+}
+
+//===----------------------------------------------------------------------===//
 // Test IsolatedRegionOp - parse passthrough region arguments.
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseIsolatedRegionOp(OpAsmParser &parser,
                                          OperationState &result) {
   OpAsmParser::OperandType argInfo;
   Type argType = parser.getBuilder().getIndexType();
 

mlir/test/lib/Dialect/Test/TestOps.td

def FunctionalRegionOp : TEST_Op<"functional_region_op",
   let extraClassDeclaration = [{
     Region *getCallableRegion() { return &body(); }
     ArrayRef<Type> getCallableResults() {
       return getType().cast<FunctionType>().getResults();
     }
   }];
 }
 
+
+def FoldToCallOp : TEST_Op<"fold_to_call_op"> {
+  let arguments = (ins FlatSymbolRefAttr:$callee);
+  let hasCanonicalizer = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Test Traits
 //===----------------------------------------------------------------------===//
 
 def SameOperandElementTypeOp : TEST_Op<"same_operand_element_type",
     [SameOperandsElementType]> {
   let arguments = (ins AnyType, AnyType);
   let results = (outs AnyType);