[Mlir-commits] [mlir] [MLIR] Fuse locations of hoisted / merged constants (PR #74670)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Dec 6 16:47:14 PST 2023
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Billy Zhu (zyx-billy)
<details>
<summary>Changes</summary>
When hoisting & 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. There are two ways this is used now:
- 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.
- 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. The hoisted op is now considered part of the setup of the block.
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).
---
Full diff: https://github.com/llvm/llvm-project/pull/74670.diff
2 Files Affected:
- (modified) mlir/lib/Transforms/Utils/FoldUtils.cpp (+49-5)
- (added) mlir/test/Transforms/canonicalize-debuginfo.mlir (+42)
``````````diff
diff --git a/mlir/lib/Transforms/Utils/FoldUtils.cpp b/mlir/lib/Transforms/Utils/FoldUtils.cpp
index 90ee5ba51de3ad..23fcb4796892d6 100644
--- a/mlir/lib/Transforms/Utils/FoldUtils.cpp
+++ b/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;
diff --git a/mlir/test/Transforms/canonicalize-debuginfo.mlir b/mlir/test/Transforms/canonicalize-debuginfo.mlir
new file mode 100644
index 00000000000000..9cee2fec83bad6
--- /dev/null
+++ b/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]]
``````````
</details>
https://github.com/llvm/llvm-project/pull/74670
More information about the Mlir-commits
mailing list