[mlir] Don't inline calls from dead SCCs

[ftynse]

Summary

During iterative inlining of the functions in a multi-step call chain, the
inliner could add the same call operation several times to the worklist, which
led to use-after-free when this op was considered more than once.

Closes #52887.

[ftynse]

git add

[wsmoses]

mlir/lib/Transforms/Inliner.cpp

461

Could this be changed to a set vector to avoid an O(N) check on this line?

[wsmoses]

LGTM with a change to setvector.

[ftynse]

Use SmallSetVector.

mlir/lib/Transforms/Inliner.cpp

 /// 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,
                                       CallGraphSCC &currentSCC) {
   CallGraph &cg = inliner.cg;
   auto &calls = inliner.calls;
 
   // A set of dead nodes to remove after inlining.
-  SmallVector<CallGraphNode *, 1> deadNodes;
+  llvm::SmallSetVector<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;
 
     // Don't collect calls if the node is already dead.
     if (useList.isDead(node)) {
-      deadNodes.push_back(node);
+      deadNodes.insert(node);
     } else {
       collectCallOps(*node->getCallableRegion(), node, cg, inliner.symbolTable,
                      calls, /*traverseNestedCGNodes=*/false);
     }
   }
 
   // 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) {
+    if (deadNodes.contains(calls[i].sourceNode))
+      continue;
     ResolvedCall it = calls[i];
     bool doInline = shouldInline(it);
     CallOpInterface call = it.call;
     LLVM_DEBUG({
       if (doInline)
         llvm::dbgs() << "* Inlining call: " << call << "\n";
       else
         llvm::dbgs() << "* Not inlining call: " << call << "\n";
for (unsigned i = 0; i != calls.size(); ++i) {
     useList.mergeUsesAfterInlining(it.targetNode, it.sourceNode);
 
     // then erase the call.
     call.erase();
 
     // If we inlined in place, mark the node for deletion.
     if (inlineInPlace) {
       useList.eraseNode(it.targetNode);
-      deadNodes.push_back(it.targetNode);
+      deadNodes.insert(it.targetNode);
     }
   }
 
   for (CallGraphNode *node : deadNodes) {
     currentSCC.remove(node);
     inliner.markForDeletion(node);
   }
   calls.clear();

mlir/test/Transforms/inlining-repeated-use.mlir

+// RUN: mlir-opt -inline %s | FileCheck %s
+
+// This could crash the inliner, make sure it does not.
+
+func @A() {
+  call @B() { inA } : () -> ()
+  return
+}
+
+func @B() {
+  call @E() : () -> ()
+  return
+}
+
+func @C() {
+  call @D() : () -> ()
+  return
+}
+
+func private @D() {
+  call @B() { inD } : () -> ()
+  return
+}
+
+func @E() {
+  call @fabsf() : () -> ()
+  return
+}
+
+func private @fabsf()
+
+// CHECK: func @A() {
+// CHECK:   call @fabsf() : () -> ()
+// CHECK:   return
+// CHECK: }
+// CHECK: func @B() {
+// CHECK:   call @fabsf() : () -> ()
+// CHECK:   return
+// CHECK: }
+// CHECK: func @C() {
+// CHECK:   call @fabsf() : () -> ()
+// CHECK:   return
+// CHECK: }
+// CHECK: func @E() {
+// CHECK:   call @fabsf() : () -> ()
+// CHECK:   return
+// CHECK: }
+// CHECK: func private @fabsf()