Revert "[MLIR] Fuse locations of merged constants (#74670)"

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.

Author: MaskRay

mlir/include/mlir/Transforms/FoldUtils.h

@@ -33,8 +33,7 @@ class Value;
 class OperationFolder {
 public:
   OperationFolder(MLIRContext *ctx, OpBuilder::Listener *listener = nullptr)
-      : fusedLocationTag(StringAttr::get(ctx, "CSE")), interfaces(ctx),
-        rewriter(ctx, listener) {}
+      : interfaces(ctx), rewriter(ctx, listener) {}
 
   /// Tries to perform folding on the given `op`, including unifying
   /// deduplicated constants. If successful, replaces `op`'s uses with
@@ -96,15 +95,6 @@ class OperationFolder {
                                     Dialect *dialect, Attribute value,
                                     Type type, Location loc);
 
-  // Fuse `foldedLocation` into the Location of `retainedOp`. This will result
-  // in `retainedOp` having a FusedLoc with `fusedLocationTag` to help trace the
-  // source of the fusion. If `retainedOp` already had a FusedLoc with the same
-  // tag, `foldedLocation` will simply be appended to it.
-  void appendFoldedLocation(Operation *retainedOp, Location foldedLocation);
-
-  /// Tag for annotating fused locations as a result of merging constants.
-  StringAttr fusedLocationTag;
-
   /// A mapping between an insertion region and the constants that have been
   /// created within it.
   DenseMap<Region *, ConstantMap> foldScopes;

mlir/lib/Transforms/Utils/FoldUtils.cpp

@@ -141,7 +141,6 @@ bool OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
   // If there is an existing constant, replace `op`.
   if (folderConstOp) {
     notifyRemoval(op);
-    appendFoldedLocation(folderConstOp, op->getLoc());
     rewriter.replaceOp(op, folderConstOp->getResults());
     return false;
   }
@@ -295,10 +294,8 @@ OperationFolder::tryGetOrCreateConstant(ConstantMap &uniquedConstants,
   // Check if an existing mapping already exists.
   auto constKey = std::make_tuple(dialect, value, type);
   Operation *&constOp = uniquedConstants[constKey];
-  if (constOp) {
-    appendFoldedLocation(constOp, loc);
+  if (constOp)
     return constOp;
-  }
 
   // If one doesn't exist, try to materialize one.
   if (!(constOp = materializeConstant(dialect, rewriter, value, type, loc)))
@@ -319,7 +316,6 @@ OperationFolder::tryGetOrCreateConstant(ConstantMap &uniquedConstants,
   // materialized operation in favor of the existing one.
   if (auto *existingOp = uniquedConstants.lookup(newKey)) {
     notifyRemoval(constOp);
-    appendFoldedLocation(existingOp, constOp->getLoc());
     rewriter.eraseOp(constOp);
     referencedDialects[existingOp].push_back(dialect);
     return constOp = existingOp;
@@ -330,76 +326,3 @@ OperationFolder::tryGetOrCreateConstant(ConstantMap &uniquedConstants,
   auto newIt = uniquedConstants.insert({newKey, constOp});
   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) {
-  auto fusedLoc = dyn_cast<FusedLoc>(loc);
-  if (!fusedLoc)
-    return loc;
-
-  SetVector<Location> flattenedLocs;
-  Attribute metadata = fusedLoc.getMetadata();
-  ArrayRef<Location> unflattenedLocs = fusedLoc.getLocations();
-  bool hasAnyNestedLocChanged = false;
-
-  for (const Location &unflattenedLoc : unflattenedLocs) {
-    Location flattenedLoc = FlattenFusedLocationRecursively(unflattenedLoc);
-
-    auto flattenedFusedLoc = dyn_cast<FusedLoc>(flattenedLoc);
-    if (flattenedFusedLoc && (!flattenedFusedLoc.getMetadata() ||
-                              flattenedFusedLoc.getMetadata() == metadata)) {
-      hasAnyNestedLocChanged = true;
-      ArrayRef<Location> nestedLocations = flattenedFusedLoc.getLocations();
-      flattenedLocs.insert(nestedLocations.begin(), nestedLocations.end());
-    } else {
-      if (flattenedLoc != unflattenedLoc)
-        hasAnyNestedLocChanged = true;
-
-      flattenedLocs.insert(flattenedLoc);
-    }
-  }
-
-  if (!hasAnyNestedLocChanged &&
-      unflattenedLocs.size() == flattenedLocs.size()) {
-    return loc;
-  }
-
-  return FusedLoc::get(loc->getContext(), flattenedLocs.takeVector(),
-                       fusedLoc.getMetadata());
-}
-
-void OperationFolder::appendFoldedLocation(Operation *retainedOp,
-                                           Location foldedLocation) {
-  // Append into existing fused location if it has the same tag.
-  if (auto existingFusedLoc =
-          dyn_cast<FusedLocWith<StringAttr>>(retainedOp->getLoc())) {
-    StringAttr existingMetadata = existingFusedLoc.getMetadata();
-    if (existingMetadata == fusedLocationTag) {
-      ArrayRef<Location> existingLocations = existingFusedLoc.getLocations();
-      SetVector<Location> locations(existingLocations.begin(),
-                                    existingLocations.end());
-      locations.insert(foldedLocation);
-      Location newFusedLoc = FusedLoc::get(
-          retainedOp->getContext(), locations.takeVector(), existingMetadata);
-      retainedOp->setLoc(FlattenFusedLocationRecursively(newFusedLoc));
-      return;
-    }
-  }
-
-  // Create a new fusedloc with retainedOp's loc and foldedLocation.
-  // If they're already equal, no need to fuse.
-  if (retainedOp->getLoc() == foldedLocation)
-    return;
-
-  Location newFusedLoc =
-      FusedLoc::get(retainedOp->getContext(),
-                    {retainedOp->getLoc(), foldedLocation}, fusedLocationTag);
-  retainedOp->setLoc(FlattenFusedLocationRecursively(newFusedLoc));
-}