Skip to content
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

Merged
merged 1 commit into from
Dec 12, 2023

Conversation

bchetioui
Copy link
Member

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.

@bchetioui bchetioui requested a review from d0k December 12, 2023 17:14
@llvmbot llvmbot added mlir:core MLIR Core Infrastructure mlir labels Dec 12, 2023
@llvmbot
Copy link
Member

llvmbot commented Dec 12, 2023

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Benjamin Chetioui (bchetioui)

Changes

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.


Full diff: https://github.com/llvm/llvm-project/pull/75218.diff

2 Files Affected:

  • (modified) mlir/lib/Transforms/Utils/FoldUtils.cpp (+30-2)
  • (modified) mlir/test/Transforms/canonicalize-debuginfo.mlir (+5-3)
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]]])
 
 // -----
 

@bchetioui bchetioui requested a review from zyx-billy December 12, 2023 17:15
Copy link

github-actions bot commented Dec 12, 2023

:white_check_mark: With the latest revision this PR passed the C/C++ code formatter.

@bchetioui bchetioui force-pushed the flatten-fused-locations branch from 930030e to e885e18 Compare December 12, 2023 17:36
@bchetioui bchetioui force-pushed the flatten-fused-locations branch 2 times, most recently from 10b5964 to a4ac415 Compare December 12, 2023 20:25
[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.
@bchetioui bchetioui force-pushed the flatten-fused-locations branch from a4ac415 to c78733f Compare December 12, 2023 21:00
@bchetioui bchetioui merged commit 0d1490f into llvm:main Dec 12, 2023
3 of 4 checks passed
@bchetioui bchetioui deleted the flatten-fused-locations branch December 12, 2023 21:03
bchetioui added a commit to bchetioui/llvm-project that referenced this pull request Dec 13, 2023
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.
bchetioui added a commit that referenced this pull request Dec 13, 2023
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.
MaskRay added a commit to MaskRay/llvm-project that referenced this pull request Dec 13, 2023
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.
MaskRay added a commit that referenced this pull request Dec 13, 2023
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
mlir:core MLIR Core Infrastructure mlir
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants