[llvm-branch-commits] [mlir] [mlir][shape] Report the in-place change of the `assuming_all` fold (PR #228771)

[Victor Perez via llvm-branch-commits]

https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/228771

>From 6a977e0356a3761171626771a1fdd33c4045d138 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20P=C3=A9rez=20Carrasco?=
 <victor.pc.upm at gmail.com>
Date: Fri, 2 Oct 2026 18:48:55 -0700
Subject: [PATCH] [mlir][shape] Report the in-place change of the assuming_all
 fold
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The `shape.assuming_all` fold visits its inputs in reverse order. It
erases each constant input from the operands, and it returns null at
the first input that is not constant. If the fold already erased an
input at that point, it changed the op but reported a failure, so no
driver knew about the change. The fold now returns the result of the op
in this case, which is the in-place signal of the single-result fold
form.

No test fails on main. The later patch "[mlir] Apply partial folds in
the greedy pattern rewrite driver" adds a fingerprint check of each fold
to the greedy driver. The check runs only in a build with
`-DMLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON`. In such a build of that
patch, with this fix reverted, the existing test
`Dialect/Shape/canonicalize.mlir` fails:

  $ mlir-opt -split-input-file -allow-unregistered-dialect \
      -canonicalize="test-convergence" \
      mlir/test/Dialect/Shape/canonicalize.mlir
  LLVM ERROR: fold returned failure but changed the operation

Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
 mlir/lib/Dialect/Shape/IR/Shape.cpp | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index e01b705e5610a..4487c2e652276 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -615,11 +615,16 @@ void AssumingAllOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
 OpFoldResult AssumingAllOp::fold(FoldAdaptor adaptor) {
   // Iterate in reverse to first handle all constant operands. They are
   // guaranteed to be the tail of the inputs because this is commutative.
-  for (int idx = adaptor.getInputs().size() - 1; idx >= 0; idx--) {
+  int numInputs = adaptor.getInputs().size();
+  for (int idx = numInputs - 1; idx >= 0; idx--) {
     Attribute a = adaptor.getInputs()[idx];
-    // Cannot fold if any inputs are not constant;
-    if (!a)
-      return nullptr;
+    // Cannot fold if any inputs are not constant; if the loop already erased
+    // constant inputs, the op's own result reports an in-place fold.
+    if (!a) {
+      if (idx == numInputs - 1)
+        return nullptr;
+      return getResult();
+    }
 
     // We do not need to keep statically known values after handling them in
     // this method.