[Mlir-commits] [mlir] [MLIR] Erase location of folded constants (PR #75415)
Billy Zhu
llvmlistbot at llvm.org
Wed Dec 13 17:13:56 PST 2023
https://github.com/zyx-billy created https://github.com/llvm/llvm-project/pull/75415
Follow up to the discussion from #75258, and serves as an alternate solution for #74670.
Set the location to Unknown for deduplicated / moved / materialized constants by OperationFolder. This makes sure that the folded constants don't end up with an arbitrary location of one of the original ops that became it, and that hoisted ops don't confuse the stepping order.
One foreseeable side effect is that since the canonicalizer always hoists / materializes constants to the front, any IR going through the canonicalizer will have its constants' locations erased. If it makes a difference, we can consider making OperationFolder more lazy by only hoisting when deduplicating to partially avoid this issue.
>From 24954d593255487f3b5668e82d6a15835cdaedf8 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] 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 | 54 +++++++++++++++++++
.../Transforms/constant-fold-debuginfo.mlir | 42 +++++++++++++++
.../lib/Transforms/TestIntRangeInference.cpp | 5 +-
7 files changed, 112 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..2442b1ebbf6987
--- /dev/null
+++ b/mlir/test/Transforms/canonicalize-debuginfo.mlir
@@ -0,0 +1,54 @@
+// 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-DAG: #[[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)
+
+// -----
+
+// CHECK-LABEL: func @canonicalization(
+// CHECK: %[[c5:.*]] = arith.constant 5 : index
+// 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
+}
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();
More information about the Mlir-commits
mailing list