[mlir][SliceAnalysis] Fix stack overflow in graph regions #139694
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir Author: Martin Erhart (maerhart) ChangesThis analysis currently just crashes when applied to a graph region that has a use-def cycle. This PR fixes that by keeping track of the operations the DFS has already visited when following use-def edges and stopping once we visit an operation again. Full diff: https://github.com/llvm/llvm-project/pull/139694.diff 3 Files Affected:
diff --git a/mlir/include/mlir/Analysis/SliceAnalysis.h b/mlir/include/mlir/Analysis/SliceAnalysis.h
index 3b731e8bb1c22..89f570a7c7115 100644
--- a/mlir/include/mlir/Analysis/SliceAnalysis.h
+++ b/mlir/include/mlir/Analysis/SliceAnalysis.h
@@ -65,8 +65,9 @@ using ForwardSliceOptions = SliceOptions;
///
/// The implementation traverses the use chains in postorder traversal for
/// efficiency reasons: if an operation is already in `forwardSlice`, no
-/// need to traverse its uses again. Since use-def chains form a DAG, this
-/// terminates.
+/// need to traverse its uses again. In the presence of use-def cycles in a
+/// graph region, the traversal stops at the first operation that was already
+/// visited (which is not added to the slice anymore).
///
/// Upon return to the root call, `forwardSlice` is filled with a
/// postorder list of uses (i.e. a reverse topological order). To get a proper
@@ -114,8 +115,9 @@ void getForwardSlice(Value root, SetVector<Operation *> *forwardSlice,
///
/// The implementation traverses the def chains in postorder traversal for
/// efficiency reasons: if an operation is already in `backwardSlice`, no
-/// need to traverse its definitions again. Since useuse-def chains form a DAG,
-/// this terminates.
+/// need to traverse its definitions again. In the presence of use-def cycles
+/// in a graph region, the traversal stops at the first operation that was
+/// already visited (which is not added to the slice anymore).
///
/// Upon return to the root call, `backwardSlice` is filled with a
/// postorder list of defs. This happens to be a topological order, from the
diff --git a/mlir/lib/Analysis/SliceAnalysis.cpp b/mlir/lib/Analysis/SliceAnalysis.cpp
index 5aebb19e3a86e..bbdab5a6c863c 100644
--- a/mlir/lib/Analysis/SliceAnalysis.cpp
+++ b/mlir/lib/Analysis/SliceAnalysis.cpp
@@ -27,7 +27,8 @@
using namespace mlir;
static void
-getForwardSliceImpl(Operation *op, SetVector<Operation *> *forwardSlice,
+getForwardSliceImpl(Operation *op, DenseSet<Operation *> &visited,
+ SetVector<Operation *> *forwardSlice,
const SliceOptions::TransitiveFilter &filter = nullptr) {
if (!op)
return;
@@ -41,20 +42,41 @@ getForwardSliceImpl(Operation *op, SetVector<Operation *> *forwardSlice,
for (Region ®ion : op->getRegions())
for (Block &block : region)
for (Operation &blockOp : block)
- if (forwardSlice->count(&blockOp) == 0)
- getForwardSliceImpl(&blockOp, forwardSlice, filter);
- for (Value result : op->getResults()) {
- for (Operation *userOp : result.getUsers())
- if (forwardSlice->count(userOp) == 0)
- getForwardSliceImpl(userOp, forwardSlice, filter);
- }
+ if (forwardSlice->count(&blockOp) == 0) {
+ // We don't have to check if the 'blockOp' is already visited because
+ // there cannot be a traversal path from this nested op to the parent
+ // and thus a cycle cannot be closed here. We still have to mark it
+ // as visited to stop before visiting this operation again if it is
+ // part of a cycle.
+ visited.insert(&blockOp);
+ getForwardSliceImpl(&blockOp, visited, forwardSlice, filter);
+ visited.erase(&blockOp);
+ }
+
+ for (Value result : op->getResults())
+ for (Operation *userOp : result.getUsers()) {
+ // A cycle can only occur within a basic block (not across regions or
+ // basic blocks) because the parent region must be a graph region, graph
+ // regions are restricted to always have 0 or 1 blocks, and there cannot
+ // be a def-use edge from a nested operation to an operation in an
+ // ancestor region. Therefore, we don't have to but may use the same
+ // 'visited' set across regions/blocks as long as we remove operations
+ // from the set again when the DFS traverses back from the leaf to the
+ // root.
+ if (forwardSlice->count(userOp) == 0 && visited.insert(userOp).second)
+ getForwardSliceImpl(userOp, visited, forwardSlice, filter);
+
+ visited.erase(userOp);
+ }
forwardSlice->insert(op);
}
void mlir::getForwardSlice(Operation *op, SetVector<Operation *> *forwardSlice,
const ForwardSliceOptions &options) {
- getForwardSliceImpl(op, forwardSlice, options.filter);
+ DenseSet<Operation *> visited;
+ visited.insert(op);
+ getForwardSliceImpl(op, visited, forwardSlice, options.filter);
if (!options.inclusive) {
// Don't insert the top level operation, we just queried on it and don't
// want it in the results.
@@ -70,8 +92,12 @@ void mlir::getForwardSlice(Operation *op, SetVector<Operation *> *forwardSlice,
void mlir::getForwardSlice(Value root, SetVector<Operation *> *forwardSlice,
const SliceOptions &options) {
- for (Operation *user : root.getUsers())
- getForwardSliceImpl(user, forwardSlice, options.filter);
+ DenseSet<Operation *> visited;
+ for (Operation *user : root.getUsers()) {
+ visited.insert(user);
+ getForwardSliceImpl(user, visited, forwardSlice, options.filter);
+ visited.erase(user);
+ }
// Reverse to get back the actual topological order.
// std::reverse does not work out of the box on SetVector and I want an
@@ -80,7 +106,7 @@ void mlir::getForwardSlice(Value root, SetVector<Operation *> *forwardSlice,
forwardSlice->insert(v.rbegin(), v.rend());
}
-static void getBackwardSliceImpl(Operation *op,
+static void getBackwardSliceImpl(Operation *op, DenseSet<Operation *> &visited,
SetVector<Operation *> *backwardSlice,
const BackwardSliceOptions &options) {
if (!op || op->hasTrait<OpTrait::IsIsolatedFromAbove>())
@@ -94,8 +120,11 @@ static void getBackwardSliceImpl(Operation *op,
auto processValue = [&](Value value) {
if (auto *definingOp = value.getDefiningOp()) {
- if (backwardSlice->count(definingOp) == 0)
- getBackwardSliceImpl(definingOp, backwardSlice, options);
+ if (backwardSlice->count(definingOp) == 0 &&
+ visited.insert(definingOp).second)
+ getBackwardSliceImpl(definingOp, visited, backwardSlice, options);
+
+ visited.erase(definingOp);
} else if (auto blockArg = dyn_cast<BlockArgument>(value)) {
if (options.omitBlockArguments)
return;
@@ -108,7 +137,7 @@ static void getBackwardSliceImpl(Operation *op,
if (parentOp && backwardSlice->count(parentOp) == 0) {
assert(parentOp->getNumRegions() == 1 &&
llvm::hasSingleElement(parentOp->getRegion(0).getBlocks()));
- getBackwardSliceImpl(parentOp, backwardSlice, options);
+ getBackwardSliceImpl(parentOp, visited, backwardSlice, options);
}
} else {
llvm_unreachable("No definingOp and not a block argument.");
@@ -138,7 +167,9 @@ static void getBackwardSliceImpl(Operation *op,
void mlir::getBackwardSlice(Operation *op,
SetVector<Operation *> *backwardSlice,
const BackwardSliceOptions &options) {
- getBackwardSliceImpl(op, backwardSlice, options);
+ DenseSet<Operation *> visited;
+ visited.insert(op);
+ getBackwardSliceImpl(op, visited, backwardSlice, options);
if (!options.inclusive) {
// Don't insert the top level operation, we just queried on it and don't
diff --git a/mlir/test/Dialect/Affine/slicing-utils.mlir b/mlir/test/Dialect/Affine/slicing-utils.mlir
index 0848a924b9d96..c53667a98cfbe 100644
--- a/mlir/test/Dialect/Affine/slicing-utils.mlir
+++ b/mlir/test/Dialect/Affine/slicing-utils.mlir
@@ -292,3 +292,26 @@ func.func @slicing_test_multiple_return(%arg0: index) -> (index, index) {
%0:2 = "slicing-test-op"(%arg0, %arg0): (index, index) -> (index, index)
return %0#0, %0#1 : index, index
}
+
+// -----
+
+// FWD-LABEL: graph_region_with_cycle
+// BWD-LABEL: graph_region_with_cycle
+// FWDBWD-LABEL: graph_region_with_cycle
+func.func @graph_region_with_cycle() {
+ test.isolated_graph_region {
+ // FWD: matched: [[V0:%.+]] = "slicing-test-op"([[V1:%.+]]) : (i1) -> i1 forward static slice:
+ // FWD: [[V1]] = "slicing-test-op"([[V0]]) : (i1) -> i1
+ // FWD: matched: [[V1]] = "slicing-test-op"([[V0]]) : (i1) -> i1 forward static slice:
+ // FWD: [[V0]] = "slicing-test-op"([[V1]]) : (i1) -> i1
+
+ // BWD: matched: [[V0:%.+]] = "slicing-test-op"([[V1:%.+]]) : (i1) -> i1 backward static slice:
+ // BWD: [[V1]] = "slicing-test-op"([[V0]]) : (i1) -> i1
+ // BWD: matched: [[V1]] = "slicing-test-op"([[V0]]) : (i1) -> i1 backward static slice:
+ // BWD: [[V0]] = "slicing-test-op"([[V1]]) : (i1) -> i1
+ %0 = "slicing-test-op"(%1) : (i1) -> i1
+ %1 = "slicing-test-op"(%0) : (i1) -> i1
+ }
+
+ return
+}
|
|
@llvm/pr-subscribers-mlir-affine Author: Martin Erhart (maerhart) ChangesThis analysis currently just crashes when applied to a graph region that has a use-def cycle. This PR fixes that by keeping track of the operations the DFS has already visited when following use-def edges and stopping once we visit an operation again. Full diff: https://github.com/llvm/llvm-project/pull/139694.diff 3 Files Affected:
diff --git a/mlir/include/mlir/Analysis/SliceAnalysis.h b/mlir/include/mlir/Analysis/SliceAnalysis.h
index 3b731e8bb1c22..89f570a7c7115 100644
--- a/mlir/include/mlir/Analysis/SliceAnalysis.h
+++ b/mlir/include/mlir/Analysis/SliceAnalysis.h
@@ -65,8 +65,9 @@ using ForwardSliceOptions = SliceOptions;
///
/// The implementation traverses the use chains in postorder traversal for
/// efficiency reasons: if an operation is already in `forwardSlice`, no
-/// need to traverse its uses again. Since use-def chains form a DAG, this
-/// terminates.
+/// need to traverse its uses again. In the presence of use-def cycles in a
+/// graph region, the traversal stops at the first operation that was already
+/// visited (which is not added to the slice anymore).
///
/// Upon return to the root call, `forwardSlice` is filled with a
/// postorder list of uses (i.e. a reverse topological order). To get a proper
@@ -114,8 +115,9 @@ void getForwardSlice(Value root, SetVector<Operation *> *forwardSlice,
///
/// The implementation traverses the def chains in postorder traversal for
/// efficiency reasons: if an operation is already in `backwardSlice`, no
-/// need to traverse its definitions again. Since useuse-def chains form a DAG,
-/// this terminates.
+/// need to traverse its definitions again. In the presence of use-def cycles
+/// in a graph region, the traversal stops at the first operation that was
+/// already visited (which is not added to the slice anymore).
///
/// Upon return to the root call, `backwardSlice` is filled with a
/// postorder list of defs. This happens to be a topological order, from the
diff --git a/mlir/lib/Analysis/SliceAnalysis.cpp b/mlir/lib/Analysis/SliceAnalysis.cpp
index 5aebb19e3a86e..bbdab5a6c863c 100644
--- a/mlir/lib/Analysis/SliceAnalysis.cpp
+++ b/mlir/lib/Analysis/SliceAnalysis.cpp
@@ -27,7 +27,8 @@
using namespace mlir;
static void
-getForwardSliceImpl(Operation *op, SetVector<Operation *> *forwardSlice,
+getForwardSliceImpl(Operation *op, DenseSet<Operation *> &visited,
+ SetVector<Operation *> *forwardSlice,
const SliceOptions::TransitiveFilter &filter = nullptr) {
if (!op)
return;
@@ -41,20 +42,41 @@ getForwardSliceImpl(Operation *op, SetVector<Operation *> *forwardSlice,
for (Region ®ion : op->getRegions())
for (Block &block : region)
for (Operation &blockOp : block)
- if (forwardSlice->count(&blockOp) == 0)
- getForwardSliceImpl(&blockOp, forwardSlice, filter);
- for (Value result : op->getResults()) {
- for (Operation *userOp : result.getUsers())
- if (forwardSlice->count(userOp) == 0)
- getForwardSliceImpl(userOp, forwardSlice, filter);
- }
+ if (forwardSlice->count(&blockOp) == 0) {
+ // We don't have to check if the 'blockOp' is already visited because
+ // there cannot be a traversal path from this nested op to the parent
+ // and thus a cycle cannot be closed here. We still have to mark it
+ // as visited to stop before visiting this operation again if it is
+ // part of a cycle.
+ visited.insert(&blockOp);
+ getForwardSliceImpl(&blockOp, visited, forwardSlice, filter);
+ visited.erase(&blockOp);
+ }
+
+ for (Value result : op->getResults())
+ for (Operation *userOp : result.getUsers()) {
+ // A cycle can only occur within a basic block (not across regions or
+ // basic blocks) because the parent region must be a graph region, graph
+ // regions are restricted to always have 0 or 1 blocks, and there cannot
+ // be a def-use edge from a nested operation to an operation in an
+ // ancestor region. Therefore, we don't have to but may use the same
+ // 'visited' set across regions/blocks as long as we remove operations
+ // from the set again when the DFS traverses back from the leaf to the
+ // root.
+ if (forwardSlice->count(userOp) == 0 && visited.insert(userOp).second)
+ getForwardSliceImpl(userOp, visited, forwardSlice, filter);
+
+ visited.erase(userOp);
+ }
forwardSlice->insert(op);
}
void mlir::getForwardSlice(Operation *op, SetVector<Operation *> *forwardSlice,
const ForwardSliceOptions &options) {
- getForwardSliceImpl(op, forwardSlice, options.filter);
+ DenseSet<Operation *> visited;
+ visited.insert(op);
+ getForwardSliceImpl(op, visited, forwardSlice, options.filter);
if (!options.inclusive) {
// Don't insert the top level operation, we just queried on it and don't
// want it in the results.
@@ -70,8 +92,12 @@ void mlir::getForwardSlice(Operation *op, SetVector<Operation *> *forwardSlice,
void mlir::getForwardSlice(Value root, SetVector<Operation *> *forwardSlice,
const SliceOptions &options) {
- for (Operation *user : root.getUsers())
- getForwardSliceImpl(user, forwardSlice, options.filter);
+ DenseSet<Operation *> visited;
+ for (Operation *user : root.getUsers()) {
+ visited.insert(user);
+ getForwardSliceImpl(user, visited, forwardSlice, options.filter);
+ visited.erase(user);
+ }
// Reverse to get back the actual topological order.
// std::reverse does not work out of the box on SetVector and I want an
@@ -80,7 +106,7 @@ void mlir::getForwardSlice(Value root, SetVector<Operation *> *forwardSlice,
forwardSlice->insert(v.rbegin(), v.rend());
}
-static void getBackwardSliceImpl(Operation *op,
+static void getBackwardSliceImpl(Operation *op, DenseSet<Operation *> &visited,
SetVector<Operation *> *backwardSlice,
const BackwardSliceOptions &options) {
if (!op || op->hasTrait<OpTrait::IsIsolatedFromAbove>())
@@ -94,8 +120,11 @@ static void getBackwardSliceImpl(Operation *op,
auto processValue = [&](Value value) {
if (auto *definingOp = value.getDefiningOp()) {
- if (backwardSlice->count(definingOp) == 0)
- getBackwardSliceImpl(definingOp, backwardSlice, options);
+ if (backwardSlice->count(definingOp) == 0 &&
+ visited.insert(definingOp).second)
+ getBackwardSliceImpl(definingOp, visited, backwardSlice, options);
+
+ visited.erase(definingOp);
} else if (auto blockArg = dyn_cast<BlockArgument>(value)) {
if (options.omitBlockArguments)
return;
@@ -108,7 +137,7 @@ static void getBackwardSliceImpl(Operation *op,
if (parentOp && backwardSlice->count(parentOp) == 0) {
assert(parentOp->getNumRegions() == 1 &&
llvm::hasSingleElement(parentOp->getRegion(0).getBlocks()));
- getBackwardSliceImpl(parentOp, backwardSlice, options);
+ getBackwardSliceImpl(parentOp, visited, backwardSlice, options);
}
} else {
llvm_unreachable("No definingOp and not a block argument.");
@@ -138,7 +167,9 @@ static void getBackwardSliceImpl(Operation *op,
void mlir::getBackwardSlice(Operation *op,
SetVector<Operation *> *backwardSlice,
const BackwardSliceOptions &options) {
- getBackwardSliceImpl(op, backwardSlice, options);
+ DenseSet<Operation *> visited;
+ visited.insert(op);
+ getBackwardSliceImpl(op, visited, backwardSlice, options);
if (!options.inclusive) {
// Don't insert the top level operation, we just queried on it and don't
diff --git a/mlir/test/Dialect/Affine/slicing-utils.mlir b/mlir/test/Dialect/Affine/slicing-utils.mlir
index 0848a924b9d96..c53667a98cfbe 100644
--- a/mlir/test/Dialect/Affine/slicing-utils.mlir
+++ b/mlir/test/Dialect/Affine/slicing-utils.mlir
@@ -292,3 +292,26 @@ func.func @slicing_test_multiple_return(%arg0: index) -> (index, index) {
%0:2 = "slicing-test-op"(%arg0, %arg0): (index, index) -> (index, index)
return %0#0, %0#1 : index, index
}
+
+// -----
+
+// FWD-LABEL: graph_region_with_cycle
+// BWD-LABEL: graph_region_with_cycle
+// FWDBWD-LABEL: graph_region_with_cycle
+func.func @graph_region_with_cycle() {
+ test.isolated_graph_region {
+ // FWD: matched: [[V0:%.+]] = "slicing-test-op"([[V1:%.+]]) : (i1) -> i1 forward static slice:
+ // FWD: [[V1]] = "slicing-test-op"([[V0]]) : (i1) -> i1
+ // FWD: matched: [[V1]] = "slicing-test-op"([[V0]]) : (i1) -> i1 forward static slice:
+ // FWD: [[V0]] = "slicing-test-op"([[V1]]) : (i1) -> i1
+
+ // BWD: matched: [[V0:%.+]] = "slicing-test-op"([[V1:%.+]]) : (i1) -> i1 backward static slice:
+ // BWD: [[V1]] = "slicing-test-op"([[V0]]) : (i1) -> i1
+ // BWD: matched: [[V1]] = "slicing-test-op"([[V0]]) : (i1) -> i1 backward static slice:
+ // BWD: [[V0]] = "slicing-test-op"([[V1]]) : (i1) -> i1
+ %0 = "slicing-test-op"(%1) : (i1) -> i1
+ %1 = "slicing-test-op"(%0) : (i1) -> i1
+ }
+
+ return
+}
|
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This analysis currently just crashes when applied to a graph region that has a use-def cycle. This PR fixes that by keeping track of the operations the DFS has already visited when following use-def edges and stopping once we visit an operation again.