[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
This commit is contained in:
Alex Zinenko
2022-01-10 11:39:06 +01:00
parent e92d63b467
commit 2f672e2ffa
2 changed files with 53 additions and 3 deletions

View File

@@ -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);
}
}

View File

@@ -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()