[Mlir-commits] [mlir] 87e2e89 - [MLIR] Fuse locations of merged constants (#74670)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Dec 11 19:31:58 PST 2023
Author: Billy Zhu
Date: 2023-12-11T19:31:54-08:00
New Revision: 87e2e89019ec4405fa47c3b4585be4e67473b590
URL: https://github.com/llvm/llvm-project/commit/87e2e89019ec4405fa47c3b4585be4e67473b590
DIFF: https://github.com/llvm/llvm-project/commit/87e2e89019ec4405fa47c3b4585be4e67473b590.diff
LOG: [MLIR] Fuse locations of merged constants (#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).
Added:
mlir/test/Transforms/canonicalize-debuginfo.mlir
mlir/test/Transforms/constant-fold-debuginfo.mlir
Modified:
mlir/include/mlir/Transforms/FoldUtils.h
mlir/lib/Transforms/Utils/FoldUtils.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Transforms/FoldUtils.h b/mlir/include/mlir/Transforms/FoldUtils.h
index 2600da361496c..28fa18cf942de 100644
--- a/mlir/include/mlir/Transforms/FoldUtils.h
+++ b/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;
diff --git a/mlir/lib/Transforms/Utils/FoldUtils.cpp b/mlir/lib/Transforms/Utils/FoldUtils.cpp
index 90ee5ba51de3a..dfc63ed6c4a54 100644
--- a/mlir/lib/Transforms/Utils/FoldUtils.cpp
+++ b/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);
+}
diff --git a/mlir/test/Transforms/canonicalize-debuginfo.mlir b/mlir/test/Transforms/canonicalize-debuginfo.mlir
new file mode 100644
index 0000000000000..034c9163a8059
--- /dev/null
+++ b/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]]])
diff --git a/mlir/test/Transforms/constant-fold-debuginfo.mlir b/mlir/test/Transforms/constant-fold-debuginfo.mlir
new file mode 100644
index 0000000000000..79a25f860a484
--- /dev/null
+++ b/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_
diff erent_dialect
+func.func @materialize_
diff erent_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_
diff erent_dialect":0:0)
+ %2 = arith.constant 1.0 : f32 loc("materialize_
diff erent_dialect":1:0)
+
+ return %1, %2: f32, f32
+}
+
+// CHECK-DAG: #[[LocConst0:.*]] = loc("materialize_
diff erent_dialect":0:0)
+// CHECK-DAG: #[[LocConst1:.*]] = loc("materialize_
diff erent_dialect":1:0)
+// CHECK: #[[FusedLoc]] = loc(fused<"CSE">[#[[LocConst1]], #[[LocConst0]]])
More information about the Mlir-commits
mailing list