[Mlir-commits] [mlir] [MLIR] Erase location of folded constants (PR #75415)
Billy Zhu
llvmlistbot at llvm.org
Thu Dec 14 09:54:32 PST 2023
https://github.com/zyx-billy updated https://github.com/llvm/llvm-project/pull/75415
>From d9045e45e07466d34fa0cce45ee07c1020fcf504 Mon Sep 17 00:00:00 2001
From: Billy Zhu <billyzhu at modular.com>
Date: Wed, 13 Dec 2023 15:51:44 -0800
Subject: [PATCH 1/2] erase loc when folding ops
---
mlir/include/mlir/Transforms/FoldUtils.h | 8 +++-
mlir/lib/Transforms/SCCP.cpp | 2 +-
mlir/lib/Transforms/Utils/FoldUtils.cpp | 13 +++---
.../Transform/test-pattern-application.mlir | 2 -
.../Transforms/canonicalize-debuginfo.mlir | 42 +++++++++++++++++++
.../Transforms/constant-fold-debuginfo.mlir | 42 +++++++++++++++++++
.../lib/Transforms/TestIntRangeInference.cpp | 5 +--
7 files changed, 100 insertions(+), 14 deletions(-)
create mode 100644 mlir/test/Transforms/canonicalize-debuginfo.mlir
create mode 100644 mlir/test/Transforms/constant-fold-debuginfo.mlir
diff --git a/mlir/include/mlir/Transforms/FoldUtils.h b/mlir/include/mlir/Transforms/FoldUtils.h
index 2600da361496cd..2e7a6fe3e362cb 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) {}
+ : erasedFoldedLocation(UnknownLoc::get(ctx)), interfaces(ctx),
+ rewriter(ctx, listener) {}
/// Tries to perform folding on the given `op`, including unifying
/// deduplicated constants. If successful, replaces `op`'s uses with
@@ -65,7 +66,7 @@ class OperationFolder {
/// be created in a parent block. On success this returns the constant
/// operation, nullptr otherwise.
Value getOrCreateConstant(Block *block, Dialect *dialect, Attribute value,
- Type type, Location loc);
+ Type type);
private:
/// This map keeps track of uniqued constants by dialect, attribute, and type.
@@ -95,6 +96,9 @@ class OperationFolder {
Dialect *dialect, Attribute value,
Type type, Location loc);
+ /// The location to overwrite with for folder-owned constants.
+ UnknownLoc erasedFoldedLocation;
+
/// 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/SCCP.cpp b/mlir/lib/Transforms/SCCP.cpp
index 14435b37acc914..b2d3929b045968 100644
--- a/mlir/lib/Transforms/SCCP.cpp
+++ b/mlir/lib/Transforms/SCCP.cpp
@@ -53,7 +53,7 @@ static LogicalResult replaceWithConstant(DataFlowSolver &solver,
Dialect *dialect = latticeValue.getConstantDialect();
Value constant = folder.getOrCreateConstant(
builder.getInsertionBlock(), dialect, latticeValue.getConstantValue(),
- value.getType(), value.getLoc());
+ value.getType());
if (!constant)
return failure();
diff --git a/mlir/lib/Transforms/Utils/FoldUtils.cpp b/mlir/lib/Transforms/Utils/FoldUtils.cpp
index 90ee5ba51de3ad..8221c12fd6b2dd 100644
--- a/mlir/lib/Transforms/Utils/FoldUtils.cpp
+++ b/mlir/lib/Transforms/Utils/FoldUtils.cpp
@@ -154,6 +154,7 @@ bool OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
!isFolderOwnedConstant(op->getPrevNode())))
op->moveBefore(&insertBlock->front());
+ op->setLoc(erasedFoldedLocation);
folderConstOp = op;
referencedDialects[op].push_back(op->getDialect());
return true;
@@ -193,8 +194,7 @@ void OperationFolder::clear() {
/// Get or create a constant using the given builder. On success this returns
/// the constant operation, nullptr otherwise.
Value OperationFolder::getOrCreateConstant(Block *block, Dialect *dialect,
- Attribute value, Type type,
- Location loc) {
+ Attribute value, Type type) {
// Find an insertion point for the constant.
auto *insertRegion = getInsertionRegion(interfaces, block);
auto &entry = insertRegion->front();
@@ -202,8 +202,8 @@ Value OperationFolder::getOrCreateConstant(Block *block, Dialect *dialect,
// Get the constant map for the insertion region of this operation.
auto &uniquedConstants = foldScopes[insertRegion];
- Operation *constOp =
- tryGetOrCreateConstant(uniquedConstants, dialect, value, type, loc);
+ Operation *constOp = tryGetOrCreateConstant(uniquedConstants, dialect, value,
+ type, erasedFoldedLocation);
return constOp ? constOp->getResult(0) : Value();
}
@@ -258,8 +258,9 @@ OperationFolder::processFoldResults(Operation *op,
// Check to see if there is a canonicalized version of this constant.
auto res = op->getResult(i);
Attribute attrRepl = foldResults[i].get<Attribute>();
- if (auto *constOp = tryGetOrCreateConstant(
- uniquedConstants, dialect, attrRepl, res.getType(), op->getLoc())) {
+ if (auto *constOp =
+ tryGetOrCreateConstant(uniquedConstants, dialect, attrRepl,
+ res.getType(), erasedFoldedLocation)) {
// Ensure that this constant dominates the operation we are replacing it
// with. This may not automatically happen if the operation being folded
// was inserted before the constant within the insertion block.
diff --git a/mlir/test/Dialect/Transform/test-pattern-application.mlir b/mlir/test/Dialect/Transform/test-pattern-application.mlir
index 2fd47c6bae396a..ff9a535c838432 100644
--- a/mlir/test/Dialect/Transform/test-pattern-application.mlir
+++ b/mlir/test/Dialect/Transform/test-pattern-application.mlir
@@ -179,7 +179,6 @@ module {
// CHECK: return %[[c5]]
func.func @canonicalization(%t: tensor<5xf32>) -> index {
%c0 = arith.constant 0 : index
- // expected-remark @below {{op was replaced}}
%dim = tensor.dim %t, %c0 : tensor<5xf32>
return %dim : index
}
@@ -191,7 +190,6 @@ transform.sequence failures(propagate) {
transform.apply_patterns to %1 {
transform.apply_patterns.canonicalization
} : !transform.any_op
- transform.test_print_remark_at_operand %0, "op was replaced" : !transform.any_op
}
// -----
diff --git a/mlir/test/Transforms/canonicalize-debuginfo.mlir b/mlir/test/Transforms/canonicalize-debuginfo.mlir
new file mode 100644
index 00000000000000..b80da2bbdbba70
--- /dev/null
+++ b/mlir/test/Transforms/canonicalize-debuginfo.mlir
@@ -0,0 +1,42 @@
+// 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(#[[UnknownLoc:.*]])
+ %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)
+ return %0, %1, %2, %3 : index, index, index, index
+}
+// CHECK: #[[UnknownLoc]] = loc(unknown)
+
+// -----
+
+// CHECK-LABEL: func @simple_hoist
+func.func @simple_hoist(%arg0: memref<8xi32>) -> i32 {
+ // CHECK: arith.constant 88 : i32 loc(#[[UnknownLoc:.*]])
+ %0 = arith.constant 42 : i32 loc("simple_hoist":0:0)
+ %1 = arith.constant 0 : index loc("simple_hoist":1:0)
+ memref.store %0, %arg0[%1] : memref<8xi32>
+
+ %2 = arith.constant 88 : i32 loc("simple_hoist":2:0)
+
+ return %2 : i32
+}
+// CHECK: #[[UnknownLoc]] = loc(unknown)
+
+// -----
+
+// CHECK-LABEL: func @hoist_and_merge
+func.func @hoist_and_merge(%arg0: memref<8xi32>) {
+ // CHECK-NEXT: arith.constant 42 : i32 loc(#[[UnknownLoc:.*]])
+ affine.for %arg1 = 0 to 8 {
+ %0 = arith.constant 42 : i32 loc("hoist_and_merge":0:0)
+ %1 = arith.constant 42 : i32 loc("hoist_and_merge":1:0)
+ memref.store %0, %arg0[%arg1] : memref<8xi32>
+ memref.store %1, %arg0[%arg1] : memref<8xi32>
+ }
+ return
+} loc("hoist_and_merge":2:0)
+// CHECK: #[[UnknownLoc]] = loc(unknown)
diff --git a/mlir/test/Transforms/constant-fold-debuginfo.mlir b/mlir/test/Transforms/constant-fold-debuginfo.mlir
new file mode 100644
index 00000000000000..c308bc477bee44
--- /dev/null
+++ b/mlir/test/Transforms/constant-fold-debuginfo.mlir
@@ -0,0 +1,42 @@
+// 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) {
+ // CHECK-NEXT: [[C:%.+]] = arith.constant 6 : i32 loc(#[[UnknownLoc:.*]])
+ %0 = arith.constant 1 : i32 loc("fold_and_merge":0:0)
+ %1 = arith.constant 5 : i32 loc("fold_and_merge":1:0)
+ %2 = arith.addi %0, %1 : i32 loc("fold_and_merge":2:0)
+
+ %3 = arith.constant 6 : i32 loc("fold_and_merge":3:0)
+
+ return %2, %3: i32, i32
+}
+// CHECK: #[[UnknownLoc]] = loc(unknown)
+
+// -----
+
+// CHECK-LABEL: func @materialize_different_dialect
+func.func @materialize_different_dialect() -> (f32, f32) {
+ // CHECK: arith.constant 1.{{0*}}e+00 : f32 loc(#[[UnknownLoc:.*]])
+ %0 = arith.constant -1.0 : f32 loc("materialize_different_dialect":0:0)
+ %1 = math.absf %0 : f32 loc("materialize_different_dialect":1:0)
+ %2 = arith.constant 1.0 : f32 loc("materialize_different_dialect":2:0)
+
+ return %1, %2: f32, f32
+}
+// CHECK: #[[UnknownLoc]] = loc(unknown)
+
+// -----
+
+// CHECK-LABEL: func @materialize_in_front
+func.func @materialize_in_front(%arg0: memref<8xi32>) {
+ // CHECK-NEXT: arith.constant 6 : i32 loc(#[[UnknownLoc:.*]])
+ affine.for %arg1 = 0 to 8 {
+ %1 = arith.constant 1 : i32 loc("materialize_in_front":0:0)
+ %2 = arith.constant 5 : i32 loc("materialize_in_front":1:0)
+ %3 = arith.addi %1, %2 : i32 loc("materialize_in_front":2:0)
+ memref.store %3, %arg0[%arg1] : memref<8xi32>
+ }
+ return
+} loc("materialize_in_front":3:0)
+// CHECK: #[[UnknownLoc]] = loc(unknown)
diff --git a/mlir/test/lib/Transforms/TestIntRangeInference.cpp b/mlir/test/lib/Transforms/TestIntRangeInference.cpp
index 2f6dd5b8095dfa..5758f6acf2f0ff 100644
--- a/mlir/test/lib/Transforms/TestIntRangeInference.cpp
+++ b/mlir/test/lib/Transforms/TestIntRangeInference.cpp
@@ -40,9 +40,8 @@ static LogicalResult replaceWithConstant(DataFlowSolver &solver, OpBuilder &b,
maybeDefiningOp ? maybeDefiningOp->getDialect()
: value.getParentRegion()->getParentOp()->getDialect();
Attribute constAttr = b.getIntegerAttr(value.getType(), *maybeConstValue);
- Value constant =
- folder.getOrCreateConstant(b.getInsertionBlock(), valueDialect, constAttr,
- value.getType(), value.getLoc());
+ Value constant = folder.getOrCreateConstant(
+ b.getInsertionBlock(), valueDialect, constAttr, value.getType());
if (!constant)
return failure();
>From 8ec47a9d0aeffcb11b3558b33b0f69e8d328ea8a Mon Sep 17 00:00:00 2001
From: Billy Zhu <billyzhu at modular.com>
Date: Thu, 14 Dec 2023 09:50:54 -0800
Subject: [PATCH 2/2] late erase
---
mlir/include/mlir/Transforms/FoldUtils.h | 5 +--
mlir/lib/Transforms/Utils/FoldUtils.cpp | 33 ++++++++++++-------
.../Transforms/canonicalize-debuginfo.mlir | 6 +++-
3 files changed, 29 insertions(+), 15 deletions(-)
diff --git a/mlir/include/mlir/Transforms/FoldUtils.h b/mlir/include/mlir/Transforms/FoldUtils.h
index 2e7a6fe3e362cb..96a796d5bf7746 100644
--- a/mlir/include/mlir/Transforms/FoldUtils.h
+++ b/mlir/include/mlir/Transforms/FoldUtils.h
@@ -91,10 +91,11 @@ class OperationFolder {
ArrayRef<OpFoldResult> foldResults);
/// Try to get or create a new constant entry. On success this returns the
- /// constant operation, nullptr otherwise.
+ /// constant operation, nullptr otherwise. The constant operation will have
+ /// an `erasedFoldedLocation`.
Operation *tryGetOrCreateConstant(ConstantMap &uniquedConstants,
Dialect *dialect, Attribute value,
- Type type, Location loc);
+ Type type);
/// The location to overwrite with for folder-owned constants.
UnknownLoc erasedFoldedLocation;
diff --git a/mlir/lib/Transforms/Utils/FoldUtils.cpp b/mlir/lib/Transforms/Utils/FoldUtils.cpp
index 8221c12fd6b2dd..8b316b5e1097e9 100644
--- a/mlir/lib/Transforms/Utils/FoldUtils.cpp
+++ b/mlir/lib/Transforms/Utils/FoldUtils.cpp
@@ -77,8 +77,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());
+ op->setLoc(erasedFoldedLocation);
+ }
return failure();
}
@@ -112,8 +114,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());
+ op->setLoc(erasedFoldedLocation);
+ }
return true;
}
@@ -142,6 +146,7 @@ bool OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
if (folderConstOp) {
notifyRemoval(op);
rewriter.replaceOp(op, folderConstOp->getResults());
+ folderConstOp->setLoc(erasedFoldedLocation);
return false;
}
@@ -151,10 +156,11 @@ 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());
+ op->setLoc(erasedFoldedLocation);
+ }
- op->setLoc(erasedFoldedLocation);
folderConstOp = op;
referencedDialects[op].push_back(op->getDialect());
return true;
@@ -202,8 +208,8 @@ Value OperationFolder::getOrCreateConstant(Block *block, Dialect *dialect,
// Get the constant map for the insertion region of this operation.
auto &uniquedConstants = foldScopes[insertRegion];
- Operation *constOp = tryGetOrCreateConstant(uniquedConstants, dialect, value,
- type, erasedFoldedLocation);
+ Operation *constOp =
+ tryGetOrCreateConstant(uniquedConstants, dialect, value, type);
return constOp ? constOp->getResult(0) : Value();
}
@@ -258,9 +264,8 @@ OperationFolder::processFoldResults(Operation *op,
// Check to see if there is a canonicalized version of this constant.
auto res = op->getResult(i);
Attribute attrRepl = foldResults[i].get<Attribute>();
- if (auto *constOp =
- tryGetOrCreateConstant(uniquedConstants, dialect, attrRepl,
- res.getType(), erasedFoldedLocation)) {
+ if (auto *constOp = tryGetOrCreateConstant(uniquedConstants, dialect,
+ attrRepl, res.getType())) {
// Ensure that this constant dominates the operation we are replacing it
// with. This may not automatically happen if the operation being folded
// was inserted before the constant within the insertion block.
@@ -291,15 +296,18 @@ OperationFolder::processFoldResults(Operation *op,
Operation *
OperationFolder::tryGetOrCreateConstant(ConstantMap &uniquedConstants,
Dialect *dialect, Attribute value,
- Type type, Location loc) {
+ Type type) {
// Check if an existing mapping already exists.
auto constKey = std::make_tuple(dialect, value, type);
Operation *&constOp = uniquedConstants[constKey];
- if (constOp)
+ if (constOp) {
+ constOp->setLoc(erasedFoldedLocation);
return constOp;
+ }
// If one doesn't exist, try to materialize one.
- if (!(constOp = materializeConstant(dialect, rewriter, value, type, loc)))
+ if (!(constOp = materializeConstant(dialect, rewriter, value, type,
+ erasedFoldedLocation)))
return nullptr;
// Check to see if the generated constant is in the expected dialect.
@@ -319,6 +327,7 @@ OperationFolder::tryGetOrCreateConstant(ConstantMap &uniquedConstants,
notifyRemoval(constOp);
rewriter.eraseOp(constOp);
referencedDialects[existingOp].push_back(dialect);
+ existingOp->setLoc(erasedFoldedLocation);
return constOp = existingOp;
}
diff --git a/mlir/test/Transforms/canonicalize-debuginfo.mlir b/mlir/test/Transforms/canonicalize-debuginfo.mlir
index b80da2bbdbba70..30c8022daa76b7 100644
--- a/mlir/test/Transforms/canonicalize-debuginfo.mlir
+++ b/mlir/test/Transforms/canonicalize-debuginfo.mlir
@@ -16,6 +16,8 @@ func.func @merge_constants() -> (index, index, index, index) {
// CHECK-LABEL: func @simple_hoist
func.func @simple_hoist(%arg0: memref<8xi32>) -> i32 {
// CHECK: arith.constant 88 : i32 loc(#[[UnknownLoc:.*]])
+ // CHECK: arith.constant 42 : i32 loc(#[[ConstLoc0:.*]])
+ // CHECK: arith.constant 0 : index loc(#[[ConstLoc1:.*]])
%0 = arith.constant 42 : i32 loc("simple_hoist":0:0)
%1 = arith.constant 0 : index loc("simple_hoist":1:0)
memref.store %0, %arg0[%1] : memref<8xi32>
@@ -24,7 +26,9 @@ func.func @simple_hoist(%arg0: memref<8xi32>) -> i32 {
return %2 : i32
}
-// CHECK: #[[UnknownLoc]] = loc(unknown)
+// CHECK-DAG: #[[UnknownLoc]] = loc(unknown)
+// CHECK-DAG: #[[ConstLoc0]] = loc("simple_hoist":0:0)
+// CHECK-DAG: #[[ConstLoc1]] = loc("simple_hoist":1:0)
// -----
More information about the Mlir-commits
mailing list