[Mlir-commits] [mlir] [MLIR] Fix infinite loop in tryFold in Graph Regions (PR #189232)

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/189232

>From e1bd1208367d27c44ec37628db862632fb95e3fc Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Fri, 11 Sep 2026 10:06:11 +0200
Subject: [PATCH] [MLIR] Fix infinite loop in tryFold in Graph Regions

---
 mlir/lib/IR/Builders.cpp                 | 12 +++++++++++-
 mlir/test/Transforms/test-legalizer.mlir | 11 +++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp
index 98fdfb5dfd042..47e09faf0e04f 100644
--- a/mlir/lib/IR/Builders.cpp
+++ b/mlir/lib/IR/Builders.cpp
@@ -501,11 +501,21 @@ OpBuilder::tryFold(Operation *op, SmallVectorImpl<Value> &results,
   if (failed(op->fold(foldResults)))
     return cleanupFailure();
 
+  // Bound the number of in-place fold iterations. Legitimate chains are very
+  // short (e.g. foldCommutative swaps once, then the op folds to a value).
+  // Without a bound, circular SSA uses in graph regions can cause an infinite
+  // loop (e.g. addi(x, 0) where x is the op's own result).
+  constexpr int kMaxInPlaceFolds = 64;
   int count = 0;
   do {
     LDBG() << "Folded in place #" << count
            << " times: " << OpWithFlags(op, OpPrintingFlags().skipRegions());
-    count++;
+    if (++count >= kMaxInPlaceFolds) {
+      LDBG() << "Aborting after " << kMaxInPlaceFolds
+             << " in-place fold iterations: "
+             << OpWithFlags(op, OpPrintingFlags().skipRegions());
+      return cleanupFailure();
+    }
   } while (foldResults.empty() && succeeded(op->fold(foldResults)));
 
   // An in-place fold does not require generation of any constants.
diff --git a/mlir/test/Transforms/test-legalizer.mlir b/mlir/test/Transforms/test-legalizer.mlir
index 96aeccfb7c161..1b20422298f74 100644
--- a/mlir/test/Transforms/test-legalizer.mlir
+++ b/mlir/test/Transforms/test-legalizer.mlir
@@ -496,3 +496,14 @@ func.func @test_preorder_legalization() {
   // expected-remark @+1 {{'func.return' is not legalizable}}
   return
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/159675.
+// In a graph region, SSA values may be used before their definitions. Folding
+// the addi below first swaps its operands, then repeatedly attempts an
+// in-place self-fold. The legalizer must stop rather than looping forever.
+// expected-remark at +1 {{op 'arith.addi' is not legalizable}}
+%0 = "arith.addi"(%1, %0) <{overflowFlags = #arith.overflow<none>}> : (index, index) -> index
+// expected-remark at +1 {{op 'arith.constant' is not legalizable}}
+%1 = "arith.constant"() <{value = 0 : index}> : () -> index