[Mlir-commits] [mlir] 3edfb37 - [mlir] Notify rewriter listeners when hoisting constants (#233776)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 21 23:13:13 PDT 2026
Author: Ivan Butygin
Date: 2026-10-22T09:13:08+03:00
New Revision: 3edfb3731ae243906a1fc7b46f6b87c5e1957667
URL: https://github.com/llvm/llvm-project/commit/3edfb3731ae243906a1fc7b46f6b87c5e1957667
DIFF: https://github.com/llvm/llvm-project/commit/3edfb3731ae243906a1fc7b46f6b87c5e1957667.diff
LOG: [mlir] Notify rewriter listeners when hoisting constants (#233776)
`OperationFolder` was calling `op->moveBefore` without notifying the
rewriter.
Also, fixed use-after-free exposed in `TestSingleFold`,
`notifyOperationInserted` called both when op is materialized and when
it moved in the block, causing same op to be added multiple times,
`notifyOperationErased` would only remove one instance leaving remaining
pointers dangling.
Added:
mlir/test/Transforms/constant-fold-notifications.mlir
mlir/test/Transforms/operation-folder-notifications.mlir
mlir/test/lib/IR/DumpNotifications.h
Modified:
mlir/lib/Transforms/Utils/FoldUtils.cpp
mlir/test/lib/Dialect/Test/TestPatterns.cpp
mlir/test/lib/Transforms/TestSingleFold.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Transforms/Utils/FoldUtils.cpp b/mlir/lib/Transforms/Utils/FoldUtils.cpp
index 90e30b6282d7ef..66eecd1b80164d 100644
--- a/mlir/lib/Transforms/Utils/FoldUtils.cpp
+++ b/mlir/lib/Transforms/Utils/FoldUtils.cpp
@@ -80,7 +80,7 @@ LogicalResult OperationFolder::tryToFold(Operation *op, bool *inPlaceUpdate,
// inserted before this one.
Block *opBlock = op->getBlock();
if (&opBlock->front() != op && !isFolderOwnedConstant(op->getPrevNode())) {
- op->moveBefore(&opBlock->front());
+ rewriter.moveOpBefore(op, &opBlock->front());
op->setLoc(erasedFoldedLocation);
}
return failure();
@@ -118,7 +118,7 @@ OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
// check to see if we should rehoist it.
if (isFolderOwnedConstant(op)) {
if (&opBlock->front() != op && !isFolderOwnedConstant(op->getPrevNode())) {
- op->moveBefore(&opBlock->front());
+ rewriter.moveOpBefore(op, &opBlock->front());
op->setLoc(erasedFoldedLocation);
return ConstantInsertionResult::Hoisted;
}
@@ -162,12 +162,12 @@ OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
ConstantInsertionResult result = ConstantInsertionResult::Unchanged;
Block *insertBlock = &insertRegion->front();
if (opBlock != insertBlock) {
- op->moveBefore(&insertBlock->front());
+ rewriter.moveOpBefore(op, &insertBlock->front());
op->setLoc(erasedFoldedLocation);
result = ConstantInsertionResult::Hoisted;
} else if (&insertBlock->front() != op &&
!isFolderOwnedConstant(op->getPrevNode())) {
- op->moveBefore(&insertBlock->front());
+ rewriter.moveOpBefore(op, &insertBlock->front());
result = ConstantInsertionResult::Hoisted;
}
@@ -288,7 +288,7 @@ OperationFolder::processFoldResults(Operation *op,
// was inserted before the constant within the insertion block.
Block *opBlock = op->getBlock();
if (opBlock == constOp->getBlock() && &opBlock->front() != constOp)
- constOp->moveBefore(&opBlock->front());
+ rewriter.moveOpBefore(constOp, &opBlock->front());
results.push_back(constOp->getResult(0));
continue;
diff --git a/mlir/test/Transforms/constant-fold-notifications.mlir b/mlir/test/Transforms/constant-fold-notifications.mlir
new file mode 100644
index 00000000000000..645ba9a978959c
--- /dev/null
+++ b/mlir/test/Transforms/constant-fold-notifications.mlir
@@ -0,0 +1,19 @@
+// RUN: mlir-opt %s -test-single-fold="dump-notifications" | FileCheck %s
+
+// Folding the comparisons and their operands reuses the constant 2. Moving
+// that constant before the earlier definition must notify the listener.
+// CHECK: notifyOperationInserted: arith.constant, previous = arith.constant
+// CHECK-NEXT: notifyOperationErased: arith.constant
+// CHECK: func.func @fold_and_reuse()
+// CHECK-NEXT: %[[TRUE:.*]] = arith.constant true
+// CHECK-NEXT: %[[FALSE:.*]] = arith.constant false
+// CHECK-NEXT: return %[[TRUE]], %[[FALSE]]
+func.func @fold_and_reuse() -> (i1, i1) {
+ %a = arith.constant 4 : i32
+ %b = arith.constant 2 : i32
+ %c = arith.constant 1 : i32
+ %d = arith.constant 2 : i32
+ %true = arith.cmpi sgt, %a, %b : i32
+ %false = arith.cmpi sgt, %c, %d : i32
+ return %true, %false : i1, i1
+}
diff --git a/mlir/test/Transforms/operation-folder-notifications.mlir b/mlir/test/Transforms/operation-folder-notifications.mlir
new file mode 100644
index 00000000000000..ff6c59617838fb
--- /dev/null
+++ b/mlir/test/Transforms/operation-folder-notifications.mlir
@@ -0,0 +1,35 @@
+// RUN: mlir-opt %s -test-greedy-patterns="dump-notifications top-down=true fold=false" -split-input-file | FileCheck %s
+
+// Hoisting across regions must report the old position. Unifying the second
+// constant must also report replacement, user modification, and erasure.
+// CHECK-LABEL: notifyOperationInserted: test.constant, previous = test.constant
+// CHECK-NEXT: notifyOperationReplaced: test.constant
+// CHECK-NEXT: notifyOperationModified: test.use
+// CHECK-NEXT: notifyOperationErased: test.constant
+// CHECK: func.func @hoist_and_unique()
+// CHECK-NEXT: %[[C:.*]] = "test.constant"
+// CHECK: "test.use"(%[[C]], %[[C]])
+func.func @hoist_and_unique() {
+ %outer = "test.one_region_with_recursive_memory_effects"() ({
+ %a = "test.constant"() {value = 1 : i32} : () -> i32
+ %b = "test.constant"() {value = 1 : i32} : () -> i32
+ "test.use"(%a, %b) : (i32, i32) -> ()
+ "test.return"() : () -> ()
+ }) : () -> i32
+ return
+}
+
+// -----
+
+// A constant behind an ordinary operation must report its move within the
+// same block, even on its first insertion into the folder.
+// CHECK-LABEL: notifyOperationInserted: test.constant, previous = func.return
+// CHECK: func.func @hoist_in_block()
+// CHECK-NEXT: %[[C:.*]] = "test.constant"
+// CHECK-NEXT: {{.*}} = "test.side_effect_op"
+// CHECK-NEXT: return %[[C]]
+func.func @hoist_in_block() -> i32 {
+ %unused = "test.side_effect_op"() {effects = [{effect = "write"}]} : () -> i32
+ %c = "test.constant"() {value = 1 : i32} : () -> i32
+ return %c : i32
+}
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 0c86304eef0140..7072977b14de26 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -6,6 +6,7 @@
//
//===----------------------------------------------------------------------===//
+#include "../../IR/DumpNotifications.h"
#include "TestDialect.h"
#include "TestOps.h"
#include "TestTypes.h"
@@ -492,7 +493,9 @@ struct TestGreedyPatternDriver
// Additional patterns for testing the GreedyPatternRewriteDriver.
patterns.insert<IncrementIntAttribute<3>>(&getContext());
+ DumpNotifications listener;
GreedyRewriteConfig config;
+ config.setListener(dumpNotifications ? &listener : nullptr);
config.setUseTopDownTraversal(useTopDownTraversal)
.setMaxIterations(this->maxIterations)
.enableFolding(this->fold)
@@ -500,6 +503,10 @@ struct TestGreedyPatternDriver
(void)applyPatternsGreedily(getOperation(), std::move(patterns), config);
}
+ Option<bool> dumpNotifications{
+ *this, "dump-notifications",
+ llvm::cl::desc("Print rewrite listener notifications"),
+ llvm::cl::init(false)};
Option<bool> useTopDownTraversal{
*this, "top-down",
llvm::cl::desc("Seed the worklist in general top-down order"),
@@ -515,51 +522,6 @@ struct TestGreedyPatternDriver
llvm::cl::init(GreedyRewriteConfig().isConstantCSEEnabled())};
};
-struct DumpNotifications : public RewriterBase::Listener {
- void notifyBlockInserted(Block *block, Region *previous,
- Region::iterator previousIt) override {
- llvm::outs() << "notifyBlockInserted";
- if (block->getParentOp()) {
- llvm::outs() << " into " << block->getParentOp()->getName() << ": ";
- } else {
- llvm::outs() << " into unknown op: ";
- }
- if (previous == nullptr) {
- llvm::outs() << "was unlinked\n";
- } else {
- llvm::outs() << "was linked\n";
- }
- }
- void notifyOperationInserted(Operation *op,
- OpBuilder::InsertPoint previous) override {
- llvm::outs() << "notifyOperationInserted: " << op->getName();
- if (!previous.isSet()) {
- llvm::outs() << ", was unlinked\n";
- } else {
- if (!previous.getPoint().getNodePtr()) {
- llvm::outs() << ", was linked, exact position unknown\n";
- } else if (previous.getPoint() == previous.getBlock()->end()) {
- llvm::outs() << ", was last in block\n";
- } else {
- llvm::outs() << ", previous = " << previous.getPoint()->getName()
- << "\n";
- }
- }
- }
- void notifyBlockErased(Block *block) override {
- llvm::outs() << "notifyBlockErased\n";
- }
- void notifyOperationErased(Operation *op) override {
- llvm::outs() << "notifyOperationErased: " << op->getName() << "\n";
- }
- void notifyOperationModified(Operation *op) override {
- llvm::outs() << "notifyOperationModified: " << op->getName() << "\n";
- }
- void notifyOperationReplaced(Operation *op, ValueRange values) override {
- llvm::outs() << "notifyOperationReplaced: " << op->getName() << "\n";
- }
-};
-
struct TestStrictPatternDriver
: public PassWrapper<TestStrictPatternDriver, OperationPass<func::FuncOp>> {
public:
diff --git a/mlir/test/lib/IR/DumpNotifications.h b/mlir/test/lib/IR/DumpNotifications.h
new file mode 100644
index 00000000000000..6ec54f68caa6d3
--- /dev/null
+++ b/mlir/test/lib/IR/DumpNotifications.h
@@ -0,0 +1,62 @@
+//===- DumpNotifications.h - Test notifications -----------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TEST_LIB_IR_DUMPNOTIFICATIONS_H
+#define MLIR_TEST_LIB_IR_DUMPNOTIFICATIONS_H
+
+#include "mlir/IR/PatternMatch.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace mlir {
+struct DumpNotifications : public RewriterBase::Listener {
+ void notifyBlockInserted(Block *block, Region *previous,
+ Region::iterator previousIt) override {
+ llvm::outs() << "notifyBlockInserted";
+ if (block->getParentOp()) {
+ llvm::outs() << " into " << block->getParentOp()->getName() << ": ";
+ } else {
+ llvm::outs() << " into unknown op: ";
+ }
+ if (previous == nullptr) {
+ llvm::outs() << "was unlinked\n";
+ } else {
+ llvm::outs() << "was linked\n";
+ }
+ }
+ void notifyOperationInserted(Operation *op,
+ OpBuilder::InsertPoint previous) override {
+ llvm::outs() << "notifyOperationInserted: " << op->getName();
+ if (!previous.isSet()) {
+ llvm::outs() << ", was unlinked\n";
+ } else {
+ if (!previous.getPoint().getNodePtr()) {
+ llvm::outs() << ", was linked, exact position unknown\n";
+ } else if (previous.getPoint() == previous.getBlock()->end()) {
+ llvm::outs() << ", was last in block\n";
+ } else {
+ llvm::outs() << ", previous = " << previous.getPoint()->getName()
+ << "\n";
+ }
+ }
+ }
+ void notifyBlockErased(Block *block) override {
+ llvm::outs() << "notifyBlockErased\n";
+ }
+ void notifyOperationErased(Operation *op) override {
+ llvm::outs() << "notifyOperationErased: " << op->getName() << "\n";
+ }
+ void notifyOperationModified(Operation *op) override {
+ llvm::outs() << "notifyOperationModified: " << op->getName() << "\n";
+ }
+ void notifyOperationReplaced(Operation *op, ValueRange values) override {
+ llvm::outs() << "notifyOperationReplaced: " << op->getName() << "\n";
+ }
+};
+} // namespace mlir
+
+#endif // MLIR_TEST_LIB_IR_DUMPNOTIFICATIONS_H
diff --git a/mlir/test/lib/Transforms/TestSingleFold.cpp b/mlir/test/lib/Transforms/TestSingleFold.cpp
index e55f36aea0a7cd..169315f1e5f710 100644
--- a/mlir/test/lib/Transforms/TestSingleFold.cpp
+++ b/mlir/test/lib/Transforms/TestSingleFold.cpp
@@ -6,8 +6,10 @@
//
//===----------------------------------------------------------------------===//
+#include "../IR/DumpNotifications.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/FoldUtils.h"
+#include "llvm/ADT/SetVector.h"
using namespace mlir;
@@ -33,22 +35,30 @@ struct TestSingleFold : public PassWrapper<TestSingleFold, OperationPass<>>,
StringRef getDescription() const final {
return "Test single-pass operation folding and dead constant elimination";
}
- // All constants in the operation post folding.
- SmallVector<Operation *> existingConstants;
+ // Moves can report the same constant more than once. Keep each constant
+ // once so that cleanup cannot visit an operation after erasing it.
+ llvm::SmallSetVector<Operation *, 8> existingConstants;
void foldOperation(Operation *op, OperationFolder &helper);
void runOnOperation() override;
void notifyOperationInserted(Operation *op,
OpBuilder::InsertPoint previous) override {
- existingConstants.push_back(op);
+ if (dumpNotifications)
+ notificationPrinter.notifyOperationInserted(op, previous);
+ existingConstants.insert(op);
}
void notifyOperationErased(Operation *op) override {
- auto *it = llvm::find(existingConstants, op);
- if (it != existingConstants.end())
- existingConstants.erase(it);
+ if (dumpNotifications)
+ notificationPrinter.notifyOperationErased(op);
+ existingConstants.remove(op);
}
+ DumpNotifications notificationPrinter;
+ Option<bool> dumpNotifications{
+ *this, "dump-notifications",
+ llvm::cl::desc("Print insertion and erasure notifications"),
+ llvm::cl::init(false)};
Option<int> maxIterations{*this, "max-iterations",
llvm::cl::desc("Max iterations in the tryToFold"),
llvm::cl::init(1)};
More information about the Mlir-commits
mailing list