[Mlir-commits] [mlir] [mlir] Notify rewriter listeners when hoisting constants (PR #233776)
Ivan Butygin
llvmlistbot at llvm.org
Wed Oct 21 10:49:49 PDT 2026
https://github.com/Hardcode84 updated https://github.com/llvm/llvm-project/pull/233776
>From 0dd25d65833546e231bc175eaa3f508ec5b0449a Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Tue, 20 Oct 2026 12:11:20 +0200
Subject: [PATCH 1/3] [mlir] Notify rewriter listeners when hoisting constants
---
mlir/lib/Transforms/Utils/FoldUtils.cpp | 10 +-
mlir/unittests/Transforms/CMakeLists.txt | 4 +
mlir/unittests/Transforms/FoldUtils.cpp | 114 +++++++++++++++++++++++
3 files changed, 123 insertions(+), 5 deletions(-)
create mode 100644 mlir/unittests/Transforms/FoldUtils.cpp
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/unittests/Transforms/CMakeLists.txt b/mlir/unittests/Transforms/CMakeLists.txt
index dc5920087b505e..53287cfefa719f 100644
--- a/mlir/unittests/Transforms/CMakeLists.txt
+++ b/mlir/unittests/Transforms/CMakeLists.txt
@@ -1,8 +1,12 @@
add_mlir_unittest(MLIRTransformsTests
Canonicalizer.cpp
DialectConversion.cpp
+ FoldUtils.cpp
)
mlir_target_link_libraries(MLIRTransformsTests
PRIVATE
MLIRParser
MLIRTransforms)
+
+target_include_directories(MLIRTransformsTests PRIVATE "${MLIR_BINARY_DIR}/test/lib/Dialect/Test")
+target_link_libraries(MLIRTransformsTests PRIVATE MLIRTestDialect)
diff --git a/mlir/unittests/Transforms/FoldUtils.cpp b/mlir/unittests/Transforms/FoldUtils.cpp
new file mode 100644
index 00000000000000..33faceb1cefae8
--- /dev/null
+++ b/mlir/unittests/Transforms/FoldUtils.cpp
@@ -0,0 +1,114 @@
+//===- FoldUtils.cpp - Operation folding tests --------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/FoldUtils.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Verifier.h"
+#include "mlir/Parser/Parser.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "gtest/gtest.h"
+
+#include "../../test/lib/Dialect/Test/TestDialect.h"
+
+using namespace mlir;
+
+namespace {
+static std::string region(StringRef body) {
+ return ("%r = \"test.one_region_with_recursive_memory_effects\"() ({\n" +
+ body + "\n\"test.return\"() : () -> ()\n}) : () -> i32\n")
+ .str();
+}
+
+class OperationFolderTest : public testing::Test {
+protected:
+ OperationFolderTest() { context.loadDialect<test::TestDialect>(); }
+ MLIRContext context;
+};
+
+TEST_F(OperationFolderTest, ConstantHoistingNotifiesListener) {
+ // Constant CSE must report moves and replacements through the driver's
+ // listener chain, including changes to users of the replaced constant.
+ auto moduleOp = parseSourceString<ModuleOp>(region(R"mlir(
+ %a = "test.constant"() {value = 1 : i32} : () -> i32
+ %b = "test.constant"() {value = 1 : i32} : () -> i32
+ "test.use"(%a, %b) : (i32, i32) -> ()
+ )mlir"),
+ &context);
+ ASSERT_TRUE(moduleOp);
+ struct Listener : RewriterBase::Listener {
+ unsigned moved = 0, erased = 0, modified = 0;
+ void notifyOperationInserted(Operation *op,
+ OpBuilder::InsertPoint previous) override {
+ if (op->hasTrait<OpTrait::ConstantLike>()) {
+ EXPECT_TRUE(previous.isSet());
+ EXPECT_NE(previous.getBlock(), op->getBlock());
+ ++moved;
+ }
+ }
+ void notifyOperationErased(Operation *op) override {
+ if (op->hasTrait<OpTrait::ConstantLike>())
+ ++erased;
+ }
+ void notifyOperationModified(Operation *op) override {
+ if (op->getName().getStringRef() == "test.use")
+ ++modified;
+ }
+ } listener;
+ GreedyRewriteConfig config;
+ config.setListener(&listener)
+ .enableFolding(false)
+ .enableConstantCSE(true)
+ .setRegionSimplificationLevel(GreedySimplifyRegionLevel::Disabled);
+ ASSERT_TRUE(succeeded(
+ applyPatternsGreedily(*moduleOp, RewritePatternSet(&context), config)));
+ EXPECT_EQ(listener.moved, 1u);
+ EXPECT_EQ(listener.erased, 1u);
+ EXPECT_EQ(listener.modified, 1u);
+ EXPECT_TRUE(succeeded(verify(*moduleOp)));
+}
+
+TEST_F(OperationFolderTest, ConstantRehoistingNotifiesListener) {
+ // A constant already owned by the folder must report each move, whether
+ // rehoisting is requested by constant insertion or by folding.
+ auto moduleOp = parseSourceString<ModuleOp>(R"mlir(
+ %a = "test.constant"() {value = 1 : i32} : () -> i32
+ "test.use"(%a) : (i32) -> ()
+ )mlir",
+ &context);
+ ASSERT_TRUE(moduleOp);
+ struct Listener : RewriterBase::Listener {
+ unsigned moved = 0;
+ void notifyOperationInserted(Operation *op,
+ OpBuilder::InsertPoint previous) override {
+ EXPECT_TRUE(previous.isSet());
+ EXPECT_EQ(previous.getBlock(), op->getBlock());
+ ++moved;
+ }
+ } listener;
+ Operation *constant = &moduleOp->getBody()->front();
+ Operation *user = constant->getNextNode();
+ OperationFolder folder(&context, &listener);
+ EXPECT_EQ(folder.insertKnownConstant(constant),
+ OperationFolder::ConstantInsertionResult::Unchanged);
+ // Inserting an operation ahead of the constant requires rehoisting
+ // without changing the constant's users or value.
+ OpBuilder builder(&context);
+ builder.setInsertionPoint(constant);
+ OperationState state(builder.getUnknownLoc(), "test.side_effect_op");
+ state.addTypes(builder.getI32Type());
+ Operation *inserted = builder.create(state);
+ EXPECT_EQ(folder.insertKnownConstant(constant),
+ OperationFolder::ConstantInsertionResult::Hoisted);
+ EXPECT_EQ(listener.moved, 1u);
+ inserted->moveBefore(constant);
+ EXPECT_TRUE(failed(folder.tryToFold(constant)));
+ EXPECT_EQ(listener.moved, 2u);
+ EXPECT_EQ(user->getOperand(0), constant->getResult(0));
+ EXPECT_TRUE(succeeded(verify(*moduleOp)));
+}
+} // namespace
>From d4ce41394e3ded56a06e6deb915a88eaf20236ee Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Tue, 20 Oct 2026 16:05:45 +0200
Subject: [PATCH 2/3] [mlir] Track each constant once in the single-fold test
pass
---
mlir/test/lib/Transforms/TestSingleFold.cpp | 12 ++++++------
1 file changed, 6 insertions(+), 6 deletions(-)
diff --git a/mlir/test/lib/Transforms/TestSingleFold.cpp b/mlir/test/lib/Transforms/TestSingleFold.cpp
index e55f36aea0a7cd..3369a3780ef2b2 100644
--- a/mlir/test/lib/Transforms/TestSingleFold.cpp
+++ b/mlir/test/lib/Transforms/TestSingleFold.cpp
@@ -8,6 +8,7 @@
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/FoldUtils.h"
+#include "llvm/ADT/SetVector.h"
using namespace mlir;
@@ -33,20 +34,19 @@ 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);
+ existingConstants.insert(op);
}
void notifyOperationErased(Operation *op) override {
- auto *it = llvm::find(existingConstants, op);
- if (it != existingConstants.end())
- existingConstants.erase(it);
+ existingConstants.remove(op);
}
Option<int> maxIterations{*this, "max-iterations",
>From 4638c186b85bfcb9d5e1dcb8513b7842dfd84fbb Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Wed, 21 Oct 2026 19:02:59 +0200
Subject: [PATCH 3/3] [mlir] Test constant move notifications with lit
---
.../operation-folder-notifications.mlir | 56 +++++++++
mlir/test/lib/Dialect/Test/TestPatterns.cpp | 46 +++++++
mlir/unittests/Transforms/CMakeLists.txt | 4 -
mlir/unittests/Transforms/FoldUtils.cpp | 114 ------------------
4 files changed, 102 insertions(+), 118 deletions(-)
create mode 100644 mlir/test/Transforms/operation-folder-notifications.mlir
delete mode 100644 mlir/unittests/Transforms/FoldUtils.cpp
diff --git a/mlir/test/Transforms/operation-folder-notifications.mlir b/mlir/test/Transforms/operation-folder-notifications.mlir
new file mode 100644
index 00000000000000..52a72d5673f795
--- /dev/null
+++ b/mlir/test/Transforms/operation-folder-notifications.mlir
@@ -0,0 +1,56 @@
+// RUN: mlir-opt %s -test-operation-folder -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"() : () -> i32
+ %c = "test.constant"() {value = 1 : i32} : () -> i32
+ return %c : i32
+}
+
+// -----
+
+// The test pass inserts an operation ahead of this folder-owned constant.
+// Both APIs must report rehoisting it, even though tryToFold returns failure.
+// CHECK-LABEL: Rehoist with insertKnownConstant
+// CHECK-NEXT: notifyOperationInserted: test.constant, previous = func.return
+// CHECK-NEXT: Rehoist with tryToFold
+// CHECK-NEXT: notifyOperationInserted: test.constant, previous = func.return
+// CHECK-NEXT: Fold before owned constant
+// CHECK-NEXT: notifyOperationInserted: test.constant, previous = test.side_effect_op
+// CHECK-NEXT: notifyOperationReplaced: test.constant
+// CHECK-NEXT: notifyOperationErased: test.constant
+// CHECK: func.func @rehoist()
+// CHECK-NEXT: %[[C:.*]] = "test.constant"
+// CHECK-NEXT: {{.*}} = "test.side_effect_op"
+// CHECK-NEXT: return %[[C]]
+func.func @rehoist() -> i32 {
+ %c = "test.constant"() {value = 1 : i32, test.rehoist} : () -> i32
+ return %c : i32
+}
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 0c86304eef0140..d63575a79b134c 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -560,6 +560,51 @@ struct DumpNotifications : public RewriterBase::Listener {
}
};
+struct TestOperationFolder
+ : public PassWrapper<TestOperationFolder, OperationPass<func::FuncOp>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOperationFolder)
+
+ StringRef getArgument() const final { return "test-operation-folder"; }
+ StringRef getDescription() const final {
+ return "Test constant hoisting and rehoisting notifications";
+ }
+
+ void runOnOperation() override {
+ DumpNotifications listener;
+ OperationFolder folder(&getContext(), &listener);
+ SmallVector<Operation *> constants;
+ getOperation()->walk([&](Operation *op) {
+ if (matchPattern(op, m_Constant()))
+ constants.push_back(op);
+ });
+ for (Operation *op : constants) {
+ if (folder.insertKnownConstant(op) ==
+ OperationFolder::ConstantInsertionResult::Replaced)
+ continue;
+ if (!op->hasAttr("test.rehoist"))
+ continue;
+
+ // A rewrite can insert an operation ahead of a folder-owned constant.
+ // Both entry points must report moving the constant back to the front.
+ OpBuilder builder(op);
+ OperationState state(op->getLoc(), "test.side_effect_op");
+ state.addTypes(builder.getI32Type());
+ Operation *inserted = builder.create(state);
+ llvm::outs() << "Rehoist with insertKnownConstant\n";
+ (void)folder.insertKnownConstant(op);
+ inserted->moveBefore(op);
+ llvm::outs() << "Rehoist with tryToFold\n";
+ (void)folder.tryToFold(op);
+
+ // Reusing a constant for an earlier operation must restore dominance.
+ builder.setInsertionPoint(op);
+ Operation *duplicate = builder.clone(*op);
+ llvm::outs() << "Fold before owned constant\n";
+ (void)folder.tryToFold(duplicate);
+ }
+ }
+};
+
struct TestStrictPatternDriver
: public PassWrapper<TestStrictPatternDriver, OperationPass<func::FuncOp>> {
public:
@@ -2551,6 +2596,7 @@ void registerPatternsTestPass() {
PassRegistration<TestDerivedAttributeDriver>();
PassRegistration<TestGreedyPatternDriver>();
+ PassRegistration<TestOperationFolder>();
PassRegistration<TestStrictPatternDriver>();
PassRegistration<TestWalkPatternDriver>();
diff --git a/mlir/unittests/Transforms/CMakeLists.txt b/mlir/unittests/Transforms/CMakeLists.txt
index 53287cfefa719f..dc5920087b505e 100644
--- a/mlir/unittests/Transforms/CMakeLists.txt
+++ b/mlir/unittests/Transforms/CMakeLists.txt
@@ -1,12 +1,8 @@
add_mlir_unittest(MLIRTransformsTests
Canonicalizer.cpp
DialectConversion.cpp
- FoldUtils.cpp
)
mlir_target_link_libraries(MLIRTransformsTests
PRIVATE
MLIRParser
MLIRTransforms)
-
-target_include_directories(MLIRTransformsTests PRIVATE "${MLIR_BINARY_DIR}/test/lib/Dialect/Test")
-target_link_libraries(MLIRTransformsTests PRIVATE MLIRTestDialect)
diff --git a/mlir/unittests/Transforms/FoldUtils.cpp b/mlir/unittests/Transforms/FoldUtils.cpp
deleted file mode 100644
index 33faceb1cefae8..00000000000000
--- a/mlir/unittests/Transforms/FoldUtils.cpp
+++ /dev/null
@@ -1,114 +0,0 @@
-//===- FoldUtils.cpp - Operation folding tests --------------------------===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Transforms/FoldUtils.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/IR/Verifier.h"
-#include "mlir/Parser/Parser.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "gtest/gtest.h"
-
-#include "../../test/lib/Dialect/Test/TestDialect.h"
-
-using namespace mlir;
-
-namespace {
-static std::string region(StringRef body) {
- return ("%r = \"test.one_region_with_recursive_memory_effects\"() ({\n" +
- body + "\n\"test.return\"() : () -> ()\n}) : () -> i32\n")
- .str();
-}
-
-class OperationFolderTest : public testing::Test {
-protected:
- OperationFolderTest() { context.loadDialect<test::TestDialect>(); }
- MLIRContext context;
-};
-
-TEST_F(OperationFolderTest, ConstantHoistingNotifiesListener) {
- // Constant CSE must report moves and replacements through the driver's
- // listener chain, including changes to users of the replaced constant.
- auto moduleOp = parseSourceString<ModuleOp>(region(R"mlir(
- %a = "test.constant"() {value = 1 : i32} : () -> i32
- %b = "test.constant"() {value = 1 : i32} : () -> i32
- "test.use"(%a, %b) : (i32, i32) -> ()
- )mlir"),
- &context);
- ASSERT_TRUE(moduleOp);
- struct Listener : RewriterBase::Listener {
- unsigned moved = 0, erased = 0, modified = 0;
- void notifyOperationInserted(Operation *op,
- OpBuilder::InsertPoint previous) override {
- if (op->hasTrait<OpTrait::ConstantLike>()) {
- EXPECT_TRUE(previous.isSet());
- EXPECT_NE(previous.getBlock(), op->getBlock());
- ++moved;
- }
- }
- void notifyOperationErased(Operation *op) override {
- if (op->hasTrait<OpTrait::ConstantLike>())
- ++erased;
- }
- void notifyOperationModified(Operation *op) override {
- if (op->getName().getStringRef() == "test.use")
- ++modified;
- }
- } listener;
- GreedyRewriteConfig config;
- config.setListener(&listener)
- .enableFolding(false)
- .enableConstantCSE(true)
- .setRegionSimplificationLevel(GreedySimplifyRegionLevel::Disabled);
- ASSERT_TRUE(succeeded(
- applyPatternsGreedily(*moduleOp, RewritePatternSet(&context), config)));
- EXPECT_EQ(listener.moved, 1u);
- EXPECT_EQ(listener.erased, 1u);
- EXPECT_EQ(listener.modified, 1u);
- EXPECT_TRUE(succeeded(verify(*moduleOp)));
-}
-
-TEST_F(OperationFolderTest, ConstantRehoistingNotifiesListener) {
- // A constant already owned by the folder must report each move, whether
- // rehoisting is requested by constant insertion or by folding.
- auto moduleOp = parseSourceString<ModuleOp>(R"mlir(
- %a = "test.constant"() {value = 1 : i32} : () -> i32
- "test.use"(%a) : (i32) -> ()
- )mlir",
- &context);
- ASSERT_TRUE(moduleOp);
- struct Listener : RewriterBase::Listener {
- unsigned moved = 0;
- void notifyOperationInserted(Operation *op,
- OpBuilder::InsertPoint previous) override {
- EXPECT_TRUE(previous.isSet());
- EXPECT_EQ(previous.getBlock(), op->getBlock());
- ++moved;
- }
- } listener;
- Operation *constant = &moduleOp->getBody()->front();
- Operation *user = constant->getNextNode();
- OperationFolder folder(&context, &listener);
- EXPECT_EQ(folder.insertKnownConstant(constant),
- OperationFolder::ConstantInsertionResult::Unchanged);
- // Inserting an operation ahead of the constant requires rehoisting
- // without changing the constant's users or value.
- OpBuilder builder(&context);
- builder.setInsertionPoint(constant);
- OperationState state(builder.getUnknownLoc(), "test.side_effect_op");
- state.addTypes(builder.getI32Type());
- Operation *inserted = builder.create(state);
- EXPECT_EQ(folder.insertKnownConstant(constant),
- OperationFolder::ConstantInsertionResult::Hoisted);
- EXPECT_EQ(listener.moved, 1u);
- inserted->moveBefore(constant);
- EXPECT_TRUE(failed(folder.tryToFold(constant)));
- EXPECT_EQ(listener.moved, 2u);
- EXPECT_EQ(user->getOperand(0), constant->getResult(0));
- EXPECT_TRUE(succeeded(verify(*moduleOp)));
-}
-} // namespace
More information about the Mlir-commits
mailing list