[Mlir-commits] [mlir] [MLIR] Flatten fused locations when merging constants. (PR #75218)

[Benjamin Chetioui]

================
@@ -331,6 +331,39 @@ OperationFolder::tryGetOrCreateConstant(ConstantMap &uniquedConstants,
   return newIt.first->second;
 }
 
+/// 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.
+///
+/// Fused locations are only flattened into parent fused locations if the
+/// child fused location has no metadata, or if the metadata of the parent and
+/// child fused locations are the same---this to avoid breaking cases where
+/// metadata matter.
+static Location FlattenFusedLocationRecursively(const Location loc) {
+  if (auto fusedLoc = dyn_cast<FusedLoc>(loc)) {
+    SetVector<Location> flattenedLocs;
+    Attribute metadata = fusedLoc.getMetadata();
+
+    for (const Location &unflattenedLoc : fusedLoc.getLocations()) {
+      Location flattenedLoc = FlattenFusedLocationRecursively(unflattenedLoc);
+      auto flattenedFusedLoc = dyn_cast<FusedLoc>(flattenedLoc);
+
+      if (flattenedFusedLoc && (!flattenedFusedLoc.getMetadata() ||
+                                flattenedFusedLoc.getMetadata() == metadata)) {
+        ArrayRef<Location> nestedLocations = flattenedFusedLoc.getLocations();
+        flattenedLocs.insert(nestedLocations.begin(), nestedLocations.end());
+      } else {
+        flattenedLocs.insert(flattenedLoc);
+      }
+    }
+
+    return FusedLoc::get(loc->getContext(), flattenedLocs.takeVector(),
+                         fusedLoc.getMetadata());
----------------
bchetioui wrote:

I don't think there is a case where we can skip the iteration. `flattenedLoc == unflattenedLoc` is a case that we may still be able to simplify if `unflattenedLoc` is a `FusedLoc` that has the same metadata as its parent `FusedLoc` (we go from `FusedLoc(FusedLoc(...))` to `FusedLoc(...)`.

However, we can check if a change happened in the loop, and avoid reconstructing if nothing changed. I sent a [new PR](https://github.com/llvm/llvm-project/pull/75312) and requested your review.

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