[Mlir-commits] [mlir] [MLIR] Fuse locations of hoisted / merged constants (PR #74670)
Billy Zhu
llvmlistbot at llvm.org
Wed Dec 6 14:43:34 PST 2023
https://github.com/zyx-billy updated https://github.com/llvm/llvm-project/pull/74670
>From 90c1e7b4c9a65cff7724775fc7fc4fbcac0a3e2c Mon Sep 17 00:00:00 2001
From: Billy Zhu <billyzhu at modular.com>
Date: Wed, 6 Dec 2023 13:10:57 -0800
Subject: [PATCH] fuse locations for hoisted and merged constants
---
mlir/lib/Transforms/Utils/FoldUtils.cpp | 54 +++++++++++++++++--
.../Transforms/canonicalize-debuginfo.mlir | 42 +++++++++++++++
2 files changed, 91 insertions(+), 5 deletions(-)
create mode 100644 mlir/test/Transforms/canonicalize-debuginfo.mlir
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]]
More information about the Mlir-commits
mailing list