fuse locations for hoisted and merged constants

Author: zyx-billy

mlir/lib/Transforms/Utils/FoldUtils.cpp

@@ -19,6 +19,38 @@
 
 using namespace mlir;
 
+// Fuse `foldedLocation` into the Location of `retainedOp`.
+// This will result in `retainedOp` having a FusedLoc with a StringAttr tag
+// "OpFold" 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.
+// Usage:
+// - When an op is deduplicated, fuse the location of the op to be removed into
+//   the op that is retained.
+// - When an op is hoisted to the front/back of a block, fuse the location of
+//   the parent region of the block into the hoisted op.
+static void appendFoldedLocation(Operation *retainedOp,
+                                 Location foldedLocation) {
+  constexpr std::string_view tag = "OpFold";
+  // Append into existing fused location if it has the same tag.
+  if (auto existingFusedLoc =
+          retainedOp->getLoc().dyn_cast<FusedLocWith<StringAttr>>()) {
+    StringAttr existingMetadata = existingFusedLoc.getMetadata();
+    if (existingMetadata.strref().equals(tag)) {
+      SmallVector<Location> locations(existingFusedLoc.getLocations());
+      locations.push_back(foldedLocation);
+      Location newFusedLoc =
+          FusedLoc::get(retainedOp->getContext(), locations, existingMetadata);
+      retainedOp->setLoc(newFusedLoc);
+      return;
+    }
+  }
+  // Create a new fusedloc with retainedOp's loc and foldedLocation.
+  Location newFusedLoc = FusedLoc::get(
+      retainedOp->getContext(), {retainedOp->getLoc(), foldedLocation},
+      StringAttr::get(retainedOp->getContext(), tag));
+  retainedOp->setLoc(newFusedLoc);
+}
+
 /// Given an operation, find the parent region that folded constants should be
 /// inserted into.
 static Region *
@@ -77,8 +109,10 @@ LogicalResult OperationFolder::tryToFold(Operation *op, bool *inPlaceUpdate) {
     // Check to see if we should rehoist, i.e. if a non-constant operation was
     // inserted before this one.
     Block *opBlock = op->getBlock();
-    if (&opBlock->front() != op && !isFolderOwnedConstant(op->getPrevNode()))
+    if (&opBlock->front() != op && !isFolderOwnedConstant(op->getPrevNode())) {
       op->moveBefore(&opBlock->front());
+      appendFoldedLocation(op, opBlock->getParent()->getLoc());
+    }
     return failure();
   }
 
@@ -112,8 +146,10 @@ bool OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
   // If this is a constant we unique'd, we don't need to insert, but we can
   // check to see if we should rehoist it.
   if (isFolderOwnedConstant(op)) {
-    if (&opBlock->front() != op && !isFolderOwnedConstant(op->getPrevNode()))
+    if (&opBlock->front() != op && !isFolderOwnedConstant(op->getPrevNode())) {
       op->moveBefore(&opBlock->front());
+      appendFoldedLocation(op, opBlock->getParent()->getLoc());
+    }
     return true;
   }
 
@@ -141,6 +177,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;
   }
@@ -151,8 +188,10 @@ bool OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
   // anything. Otherwise, we move the constant to the insertion block.
   Block *insertBlock = &insertRegion->front();
   if (opBlock != insertBlock || (&insertBlock->front() != op &&
-                                 !isFolderOwnedConstant(op->getPrevNode())))
+                                 !isFolderOwnedConstant(op->getPrevNode()))) {
     op->moveBefore(&insertBlock->front());
+    appendFoldedLocation(op, insertBlock->getParent()->getLoc());
+  }
 
   folderConstOp = op;
   referencedDialects[op].push_back(op->getDialect());
@@ -264,8 +303,10 @@ OperationFolder::processFoldResults(Operation *op,
       // with. This may not automatically happen if the operation being folded
       // was inserted before the constant within the insertion block.
       Block *opBlock = op->getBlock();
-      if (opBlock == constOp->getBlock() && &opBlock->front() != constOp)
+      if (opBlock == constOp->getBlock() && &opBlock->front() != constOp) {
         constOp->moveBefore(&opBlock->front());
+        appendFoldedLocation(constOp, opBlock->getParent()->getLoc());
+      }
 
       results.push_back(constOp->getResult(0));
       continue;
@@ -294,8 +335,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 +359,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;

mlir/test/Transforms/canonicalize-debuginfo.mlir

@@ -0,0 +1,42 @@
+// RUN: mlir-opt -allow-unregistered-dialect %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) {
+  // 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)
+  return %0, %1, %2: 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<"OpFold">[
+// CHECK-SAME: #[[LocConst0]]
+// CHECK-SAME: #[[LocConst1]]
+// CHECK-SAME: #[[LocConst2]]
+
+// -----
+
+// CHECK-LABEL: func @hoist_constant
+func.func @hoist_constant(%arg0: memref<8xi32>) {
+  // CHECK-NEXT: arith.constant 42 : i32 loc(#[[FusedWithFunction:.*]])
+  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
+} loc("hoist_constant":2:0)
+
+// CHECK-DAG: #[[LocConst0:.*]] = loc("hoist_constant":0:0)
+// CHECK-DAG: #[[LocConst1:.*]] = loc("hoist_constant":1:0)
+// CHECK-DAG: #[[LocFunc:.*]] = loc("hoist_constant":2:0)
+
+// CHECK: #[[FusedWithFunction]] = loc(fused<"OpFold">[
+// CHECK-SAME: #[[LocConst0]]
+// CHECK-SAME: #[[LocFunc]]
+// CHECK-SAME: #[[LocConst1]]