From 2f672e2ffa22d8c10279569de123a1e60d3aa00e Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko@google.com>
Date: Mon, 10 Jan 2022 11:39:06 +0100
Subject: [PATCH] [mlir] Don't inline calls from dead SCCs

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.

Reviewed By: wsmoses

Differential Revision: https://reviews.llvm.org/D116820
---
 mlir/lib/Transforms/Inliner.cpp               |  8 ++--
 .../Transforms/inlining-repeated-use.mlir     | 48 +++++++++++++++++++
 2 files changed, 53 insertions(+), 3 deletions(-)
 create mode 100644 mlir/test/Transforms/inlining-repeated-use.mlir

diff --git a/mlir/lib/Transforms/Inliner.cpp b/mlir/lib/Transforms/Inliner.cpp
index c0befb76eed0a..d1587d650f14a 100644
--- a/mlir/lib/Transforms/Inliner.cpp
+++ b/mlir/lib/Transforms/Inliner.cpp
@@ -435,7 +435,7 @@ static LogicalResult inlineCallsInSCC(Inliner &inliner, CGUseList &useList,
   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
@@ -446,7 +446,7 @@ static LogicalResult inlineCallsInSCC(Inliner &inliner, CGUseList &useList,
 
     // 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);
@@ -457,6 +457,8 @@ static LogicalResult inlineCallsInSCC(Inliner &inliner, CGUseList &useList,
   // 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;
@@ -493,7 +495,7 @@ static LogicalResult inlineCallsInSCC(Inliner &inliner, CGUseList &useList,
     // 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);
     }
   }
 
diff --git a/mlir/test/Transforms/inlining-repeated-use.mlir b/mlir/test/Transforms/inlining-repeated-use.mlir
new file mode 100644
index 0000000000000..72ea2f7e5236a
--- /dev/null
+++ b/mlir/test/Transforms/inlining-repeated-use.mlir
@@ -0,0 +1,48 @@
+// 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()