[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 f0f17ebbe6005cbf67ce5627558316607bea19ab 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. A later patch adds a fingerprint check of each
fold to the greedy driver, under
`MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS`. With that check and without
this fix, 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.