[Mlir-commits] [mlir] [mlir] Track known constant insertion effects (PR #232574)

[Hideto Ueno]

https://github.com/uenoku created https://github.com/llvm/llvm-project/pull/232574

Report whether inserting a known constant leaves the IR unchanged, hoists
the operation, or replaces it with an existing constant. Use this result to
accurately track IR changes in the greedy rewrite driver.

Assisted-by: gpt 6.1 Sol

>From fe846c92020abfd8518879b7b7a7f4ff3329d139 Mon Sep 17 00:00:00 2001
From: Hideto Ueno <uenoku.tokotoko at gmail.com>
Date: Thu, 15 Oct 2026 19:04:55 -0700
Subject: [PATCH] [mlir] Track known constant insertion effects

---
 mlir/include/mlir/Transforms/FoldUtils.h      | 19 +++++++++++++++----
 mlir/lib/Transforms/Utils/FoldUtils.cpp       | 13 +++++++++----
 .../Utils/GreedyPatternRewriteDriver.cpp      | 14 +++++++++-----
 3 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/mlir/include/mlir/Transforms/FoldUtils.h b/mlir/include/mlir/Transforms/FoldUtils.h
index ee89e8d0e7c3f7..53fb3e95d0dd27 100644
--- a/mlir/include/mlir/Transforms/FoldUtils.h
+++ b/mlir/include/mlir/Transforms/FoldUtils.h
@@ -45,15 +45,26 @@ class OperationFolder {
   LogicalResult tryToFold(Operation *op, bool *inPlaceUpdate = nullptr,
                           int maxIterations = INT_MAX);
 
+  /// The effect of inserting a known constant on the IR.
+  enum class ConstantInsertionResult {
+    /// The operation was retained without modifying the IR.
+    Unchanged,
+    /// The operation was retained but moved to the insertion point.
+    Hoisted,
+    /// The operation was replaced with an existing constant and erased.
+    Replaced
+  };
+
   /// Tries to fold a pre-existing constant operation. `constValue` represents
   /// the value of the constant, and can be optionally passed if the value is
   /// already known (e.g. if the constant was discovered by m_Constant). This is
   /// purely an optimization opportunity for callers that already know the value
-  /// of the constant. Returns false if an existing constant for `op` already
+  /// of the constant. Returns Replaced if an existing constant for `op` already
   /// exists in the folder, in which case `op` is replaced and erased.
-  /// Otherwise, returns true and `op` is inserted into the folder (and
-  /// hoisted if necessary).
-  bool insertKnownConstant(Operation *op, Attribute constValue = {});
+  /// Otherwise, `op` is inserted into the folder (and hoisted if necessary),
+  /// returning Hoisted if it was moved or Unchanged if the IR was not modified.
+  ConstantInsertionResult insertKnownConstant(Operation *op,
+                                              Attribute constValue = {});
 
   /// Notifies that the given constant `op` should be remove from this
   /// OperationFolder's internal bookkeeping.
diff --git a/mlir/lib/Transforms/Utils/FoldUtils.cpp b/mlir/lib/Transforms/Utils/FoldUtils.cpp
index 68ad3acf295c8f..90e30b6282d7ef 100644
--- a/mlir/lib/Transforms/Utils/FoldUtils.cpp
+++ b/mlir/lib/Transforms/Utils/FoldUtils.cpp
@@ -110,7 +110,8 @@ LogicalResult OperationFolder::tryToFold(Operation *op, bool *inPlaceUpdate,
   return success();
 }
 
-bool OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
+OperationFolder::ConstantInsertionResult
+OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
   Block *opBlock = op->getBlock();
 
   // If this is a constant we unique'd, we don't need to insert, but we can
@@ -119,8 +120,9 @@ bool OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
     if (&opBlock->front() != op && !isFolderOwnedConstant(op->getPrevNode())) {
       op->moveBefore(&opBlock->front());
       op->setLoc(erasedFoldedLocation);
+      return ConstantInsertionResult::Hoisted;
     }
-    return true;
+    return ConstantInsertionResult::Unchanged;
   }
 
   // Get the constant value of the op if necessary.
@@ -149,7 +151,7 @@ bool OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
     notifyRemoval(op);
     rewriter.replaceOp(op, folderConstOp->getResults());
     folderConstOp->setLoc(erasedFoldedLocation);
-    return false;
+    return ConstantInsertionResult::Replaced;
   }
 
   // Otherwise, we insert `op`. If `op` is in the insertion block and is either
@@ -157,18 +159,21 @@ bool OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
   // constant we unique'd (i.e. one we inserted), then we don't need to do
   // anything. Otherwise, we move the constant to the insertion block.
   // The location info is erased if the constant is moved to a different block.
+  ConstantInsertionResult result = ConstantInsertionResult::Unchanged;
   Block *insertBlock = &insertRegion->front();
   if (opBlock != insertBlock) {
     op->moveBefore(&insertBlock->front());
     op->setLoc(erasedFoldedLocation);
+    result = ConstantInsertionResult::Hoisted;
   } else if (&insertBlock->front() != op &&
              !isFolderOwnedConstant(op->getPrevNode())) {
     op->moveBefore(&insertBlock->front());
+    result = ConstantInsertionResult::Hoisted;
   }
 
   folderConstOp = op;
   referencedDialects[op].push_back(op->getDialect());
-  return true;
+  return result;
 }
 
 /// Notifies that the given constant `op` should be remove from this
diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index b71194e783ba06..de800cb14af9c2 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -841,6 +841,7 @@ class GreedyPatternRewriteIteration
 
 LogicalResult RegionPatternRewriteDriver::simplify(bool *changed) && {
   bool continueRewrites = false;
+  bool anyChanged = false;
   int64_t iteration = 0;
   MLIRContext *ctx = rewriter.getContext();
   do {
@@ -859,10 +860,12 @@ LogicalResult RegionPatternRewriteDriver::simplify(bool *changed) && {
       // Check for existing constants when populating the worklist. This avoids
       // accidentally reversing the constant order during processing.
       Attribute constValue;
-      if (matchPattern(op, m_Constant(&constValue)))
-        if (!folder.insertKnownConstant(op, constValue))
-          return true;
-      return false;
+      if (!matchPattern(op, m_Constant(&constValue)))
+        return false;
+      auto result = folder.insertKnownConstant(op, constValue);
+      anyChanged |=
+          result != OperationFolder::ConstantInsertionResult::Unchanged;
+      return result == OperationFolder::ConstantInsertionResult::Replaced;
     };
 
     if (!config.getUseTopDownTraversal()) {
@@ -921,10 +924,11 @@ LogicalResult RegionPatternRewriteDriver::simplify(bool *changed) && {
           }
         },
         {&region}, iteration);
+    anyChanged |= continueRewrites;
   } while (continueRewrites);
 
   if (changed)
-    *changed = iteration > 1;
+    *changed = anyChanged;
 
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return success(!continueRewrites);