[Mlir-commits] [mlir] [mlir] Notify rewriter listeners when hoisting constants (PR #233776)
Ivan Butygin
llvmlistbot at llvm.org
Wed Oct 21 10:51:26 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/2] [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/2] [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",
More information about the Mlir-commits
mailing list