[Mlir-commits] [mlir] [mlir] Notify rewriter listeners when hoisting constants (PR #233776)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 20 03:32:50 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-core
Author: Ivan Butygin (Hardcode84)
<details>
<summary>Changes</summary>
`OperationFolder` was calling `op->moveBefore` without notifying the rewriter.
---
Full diff: https://github.com/llvm/llvm-project/pull/233776.diff
3 Files Affected:
- (modified) mlir/lib/Transforms/Utils/FoldUtils.cpp (+5-5)
- (modified) mlir/unittests/Transforms/CMakeLists.txt (+4)
- (added) mlir/unittests/Transforms/FoldUtils.cpp (+114)
``````````diff
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
``````````
</details>
https://github.com/llvm/llvm-project/pull/233776
More information about the Mlir-commits
mailing list