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

When merging constants by the operation folder, the location of the op
that remains should be updated to track the new meaning of this op. This
way we do not lose track of all possible source locations that the
constant op came from, and the final location of the op is less reliant
on the order of folding. This will also help debuggers understand how to
step these instructions.

This PR introduces a helper for operation folder to fuse another
location into the location of an op. When an op is deduplicated, fuse
the location of the op to be removed into the op that is retained. The
retained op now represents both original ops.

The FusedLoc will have a string metadata to help understand the reason
for the location fusion (motivated by the
[example](https://github.com/llvm/llvm-project/blob/71be8f3c23497e28c86f1135f564b16106d8d6fb/mlir/include/mlir/IR/BuiltinLocationAttributes.td#L130)
in the docstring of FusedLoc).

Author: zyx-billy

mlir/include/mlir/Transforms/FoldUtils.h

@@ -33,7 +33,8 @@ class Value;
 class OperationFolder {
 public:
   OperationFolder(MLIRContext *ctx, OpBuilder::Listener *listener = nullptr)
-      : interfaces(ctx), rewriter(ctx, listener) {}
+      : fusedLocationTag(StringAttr::get(ctx, "CSE")), interfaces(ctx),
+        rewriter(ctx, listener) {}
 
   /// Tries to perform folding on the given `op`, including unifying
   /// deduplicated constants. If successful, replaces `op`'s uses with
@@ -95,6 +96,15 @@ 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,6 +141,7 @@ 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;
   }
@@ -294,8 +295,10 @@ 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)
+  if (constOp) {
+    appendFoldedLocation(constOp, loc);
     return constOp;
+  }
 
   // If one doesn't exist, try to materialize one.
   if (!(constOp = materializeConstant(dialect, rewriter, value, type, loc)))
@@ -316,6 +319,7 @@ 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;
@@ -326,3 +330,32 @@ OperationFolder::tryGetOrCreateConstant(ConstantMap &uniquedConstants,
   auto newIt = uniquedConstants.insert({newKey, constOp});
   return newIt.first->second;
 }
+
+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(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(newFusedLoc);
+}

mlir/test/Transforms/canonicalize-debuginfo.mlir

@@ -0,0 +1,34 @@
+// 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) {
+  // 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
+}
+
+// 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-LABEL: func @hoist_constant
+func.func @hoist_constant(%arg0: memref<8xi32>) {
+  // CHECK-NEXT: arith.constant 42 : i32 loc(#[[FusedLoc:.*]])
+  affine.for %arg1 = 0 to 8 {
+    %0 = arith.constant 42 : i32 loc("hoist_constant":0:0)
+    %1 = arith.constant 42 : i32 loc("hoist_constant":1:0)
+    memref.store %0, %arg0[%arg1] : memref<8xi32>
+    memref.store %1, %arg0[%arg1] : memref<8xi32>
+  }
+  return
+}
+
+// CHECK-DAG: #[[LocConst0:.*]] = loc("hoist_constant":0:0)
+// CHECK-DAG: #[[LocConst1:.*]] = loc("hoist_constant":1:0)
+// CHECK: #[[FusedLoc]] = loc(fused<"CSE">[#[[LocConst0]], #[[LocConst1]]])

mlir/test/Transforms/constant-fold-debuginfo.mlir

@@ -0,0 +1,34 @@
+// RUN: mlir-opt %s -split-input-file -test-constant-fold -mlir-print-debuginfo | FileCheck %s
+
+// CHECK-LABEL: func @fold_and_merge
+func.func @fold_and_merge() -> (i32, i32) {
+  %0 = arith.constant 1 : i32
+  %1 = arith.constant 5 : i32
+
+  // CHECK-NEXT: [[C:%.+]] = arith.constant 6 : i32 loc(#[[FusedLoc:.*]])
+  %2 = arith.addi %0, %1 : i32 loc("fold_and_merge":0:0)
+
+  %3 = arith.constant 6 : i32 loc("fold_and_merge":1:0)
+
+  return %2, %3: i32, i32
+}
+
+// CHECK-DAG: #[[LocConst0:.*]] = loc("fold_and_merge":0:0)
+// CHECK-DAG: #[[LocConst1:.*]] = loc("fold_and_merge":1:0)
+// CHECK: #[[FusedLoc]] = loc(fused<"CSE">[#[[LocConst1]], #[[LocConst0]]])
+
+// -----
+
+// CHECK-LABEL: func @materialize_different_dialect
+func.func @materialize_different_dialect() -> (f32, f32) {
+  // CHECK: arith.constant 1.{{0*}}e+00 : f32 loc(#[[FusedLoc:.*]])
+  %0 = arith.constant -1.0 : f32
+  %1 = math.absf %0 : f32 loc("materialize_different_dialect":0:0)
+  %2 = arith.constant 1.0 : f32 loc("materialize_different_dialect":1:0)
+
+  return %1, %2: f32, f32
+}
+
+// CHECK-DAG: #[[LocConst0:.*]] = loc("materialize_different_dialect":0:0)
+// CHECK-DAG: #[[LocConst1:.*]] = loc("materialize_different_dialect":1:0)
+// CHECK: #[[FusedLoc]] = loc(fused<"CSE">[#[[LocConst1]], #[[LocConst0]]])