-
Notifications
You must be signed in to change notification settings - Fork 12.3k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[MLIR] Flatten fused locations when merging constants. #75218
Conversation
@llvm/pr-subscribers-mlir-core @llvm/pr-subscribers-mlir Author: Benjamin Chetioui (bchetioui) ChangesPR 74670 added support for merging locations at constant folding time. We have discovered that in some cases, the number of locations grows so big as to cause a compilation process to OOM. In that case, many of the locations end up appearing several times in nested fused locations. We add here a helper that always flattens fused locations in order to eliminate duplicates in the case of nested fused locations. Full diff: https://github.com/llvm/llvm-project/pull/75218.diff 2 Files Affected:
diff --git a/mlir/lib/Transforms/Utils/FoldUtils.cpp b/mlir/lib/Transforms/Utils/FoldUtils.cpp
index dfc63ed6c4a542..f414cca38f36ff 100644
--- a/mlir/lib/Transforms/Utils/FoldUtils.cpp
+++ b/mlir/lib/Transforms/Utils/FoldUtils.cpp
@@ -331,6 +331,34 @@ OperationFolder::tryGetOrCreateConstant(ConstantMap &uniquedConstants,
return newIt.first->second;
}
+namespace {
+
+/// Helper that flattens nested fused locations to a single fused location.
+/// Fused locations nested under non-fused locations are not flattened, and
+/// calling this on non-fused locations is a no-op as a result. The metadata
+/// of the outer fused location is retained in the result.
+Location FlattenFusedLocationRecursively(const Location loc) {
+ if (auto fusedLoc = dyn_cast<FusedLoc>(loc)) {
+ SetVector<Location> flattenedLocs;
+
+ for (const Location &unflattenedLoc : fusedLoc.getLocations()) {
+ Location flattenedLoc = FlattenFusedLocationRecursively(unflattenedLoc);
+ if (auto flattenedFusedLoc = dyn_cast<FusedLoc>(flattenedLoc)) {
+ ArrayRef<Location> nestedLocations = flattenedFusedLoc.getLocations();
+ flattenedLocs.insert(nestedLocations.begin(), nestedLocations.end());
+ } else {
+ flattenedLocs.insert(flattenedLoc);
+ }
+ }
+
+ return FusedLoc::get(loc->getContext(), flattenedLocs.takeVector(),
+ fusedLoc.getMetadata());
+ }
+
+ return loc;
+}
+} // anonymous namespace
+
void OperationFolder::appendFoldedLocation(Operation *retainedOp,
Location foldedLocation) {
// Append into existing fused location if it has the same tag.
@@ -344,7 +372,7 @@ void OperationFolder::appendFoldedLocation(Operation *retainedOp,
locations.insert(foldedLocation);
Location newFusedLoc = FusedLoc::get(
retainedOp->getContext(), locations.takeVector(), existingMetadata);
- retainedOp->setLoc(newFusedLoc);
+ retainedOp->setLoc(FlattenFusedLocationRecursively(newFusedLoc));
return;
}
}
@@ -357,5 +385,5 @@ void OperationFolder::appendFoldedLocation(Operation *retainedOp,
Location newFusedLoc =
FusedLoc::get(retainedOp->getContext(),
{retainedOp->getLoc(), foldedLocation}, fusedLocationTag);
- retainedOp->setLoc(newFusedLoc);
+ retainedOp->setLoc(FlattenFusedLocationRecursively(newFusedLoc));
}
diff --git a/mlir/test/Transforms/canonicalize-debuginfo.mlir b/mlir/test/Transforms/canonicalize-debuginfo.mlir
index 034c9163a8059f..3cf98900a7c54a 100644
--- a/mlir/test/Transforms/canonicalize-debuginfo.mlir
+++ b/mlir/test/Transforms/canonicalize-debuginfo.mlir
@@ -1,19 +1,21 @@
// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(canonicalize{test-convergence}))' -split-input-file -mlir-print-debuginfo | FileCheck %s
// CHECK-LABEL: func @merge_constants
-func.func @merge_constants() -> (index, index, index, index) {
+func.func @merge_constants() -> (index, index, index, index, index) {
// CHECK-NEXT: arith.constant 42 : index loc(#[[FusedLoc:.*]])
%0 = arith.constant 42 : index loc("merge_constants":0:0)
%1 = arith.constant 42 : index loc("merge_constants":1:0)
%2 = arith.constant 42 : index loc("merge_constants":2:0)
%3 = arith.constant 42 : index loc("merge_constants":2:0) // repeated loc
- return %0, %1, %2, %3: index, index, index, index
+ %4 = arith.constant 42 : index loc(fused<"some_label">["merge_constants":3:0])
+ return %0, %1, %2, %3, %4 : index, index, index, index, index
}
// CHECK-DAG: #[[LocConst0:.*]] = loc("merge_constants":0:0)
// CHECK-DAG: #[[LocConst1:.*]] = loc("merge_constants":1:0)
// CHECK-DAG: #[[LocConst2:.*]] = loc("merge_constants":2:0)
-// CHECK: #[[FusedLoc]] = loc(fused<"CSE">[#[[LocConst0]], #[[LocConst1]], #[[LocConst2]]])
+// CHECK-DAG: #[[LocConst3:.*]] = loc("merge_constants":3:0)
+// CHECK: #[[FusedLoc]] = loc(fused<"CSE">[#[[LocConst0]], #[[LocConst1]], #[[LocConst2]], #[[LocConst3]]])
// -----
|
|
930030e
to
e885e18
Compare
10b5964
to
a4ac415
Compare
[PR 74670](llvm#74670) added support for merging locations at constant folding time. We have discovered that in some cases, the number of locations grows so big as to cause a compilation process to OOM. In that case, many of the locations end up appearing several times in nested fused locations. We add here a helper that always flattens fused locations in order to eliminate duplicates in the case of nested fused locations. We only allow flattening nested fused locations when the inner fused location has no metadata, or has the same metadata as the outer fused location.
a4ac415
to
c78733f
Compare
This is a follow-up on [PR 75218](llvm#75218) that avoids reconstructing a fused loc in the `FlattenFusedLocationRecursively` helper when there has been no change.
This is a follow-up on [PR 75218](#75218) that avoids reconstructing a fused loc in the `FlattenFusedLocationRecursively` helper when there has been no change.
This reverts commit 87e2e89. and its follow-ups 0d1490f (llvm#75218) and 6fe3cd5 (llvm#75312). We observed significant OOM/timeout issues due to llvm#74670 to quite a few services including google-research/swirl-lm. The follow-up llvm#75218 and llvm#75312 do not address the issue. Perhaps this is worth more investigation.
This reverts commit 87e2e89. and its follow-ups 0d1490f (#75218) and 6fe3cd5 (#75312). We observed significant OOM/timeout issues due to #74670 to quite a few services including google-research/swirl-lm. The follow-up #75218 and #75312 do not address the issue. Perhaps this is worth more investigation.
PR 74670 added support for merging locations at constant folding time. We have discovered that in some cases, the number of locations grows so big as to cause a compilation process to OOM. In that case, many of the locations end up appearing several times in nested fused locations.
We add here a helper that always flattens fused locations in order to eliminate duplicates in the case of nested fused locations.