[Mlir-commits] [mlir] b959831 - [MLIR][Arith] Fix int-range-optimizations miscompile from stale solver state (#188992)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-27T16:00:55+01:00
New Revision: b959831ed2e4ef9f819b6504085f70d1d14dbac0

URL: https://github.com/llvm/llvm-project/commit/b959831ed2e4ef9f819b6504085f70d1d14dbac0
DIFF: https://github.com/llvm/llvm-project/commit/b959831ed2e4ef9f819b6504085f70d1d14dbac0.diff

LOG: [MLIR][Arith] Fix int-range-optimizations miscompile from stale solver state (#188992)

The `--int-range-optimizations` pass runs the `DataFlowSolver` once,
then calls `applyPatternsGreedily` with a `DataFlowListener` that erases
solver state when ops are deleted. However, the greedy driver's
`simplifyRegions` step (which calls `runRegionDCE` between pattern
iterations) can remove block arguments without notifying the listener.
This frees the `BlockArgumentImpl` storage, which may be reused by a
subsequent allocation. The solver then finds stale lattice state keyed
at the reused address and incorrectly treats the new block argument as a
known constant, causing a miscompile.

The existing `enableFolding(false)` was added for the same class of bug
(folding can also remove block arguments). This patch extends the fix by
also disabling region simplification, preventing dead-arg elimination
from causing the same address-reuse problem.

Fixes #137281
Fixes #126195

Assisted-by: Claude Code

Added: 
    

Modified: 
    mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
    mlir/test/Dialect/Arith/int-range-opts.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
index 5813f3c3ea746..6acfc2c15af42 100644
--- a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
@@ -678,14 +678,19 @@ struct IntRangeOptimizationsPass final
     RewritePatternSet patterns(ctx);
     populateIntRangeOptimizationsPatterns(patterns, solver);
 
-    // Disable folding to avoid potential control-flow folding that would break
-    // the solver state: this happens for example when a block argument is
-    // folded and other block arguments inherit from the address of the folded
-    // block argument. Further queries to the remaining block arguments would
-    // then return the solver state associated to the original block argument.
-    if (failed(applyPatternsGreedily(
-            op, std::move(patterns),
-            GreedyRewriteConfig().enableFolding(false).setListener(&listener))))
+    // Disable folding and region simplification to avoid breaking the solver
+    // state. Both can remove block arguments (folding via control-flow
+    // simplification, region simplification via dead-arg elimination), which
+    // frees their underlying storage. A subsequent allocation may reuse the
+    // same address for a 
diff erent block argument, causing stale solver state
+    // to be associated with the new argument and producing incorrect constants.
+    if (failed(
+            applyPatternsGreedily(op, std::move(patterns),
+                                  GreedyRewriteConfig()
+                                      .enableFolding(false)
+                                      .setRegionSimplificationLevel(
+                                          GreedySimplifyRegionLevel::Disabled)
+                                      .setListener(&listener))))
       signalPassFailure();
   }
 };

diff  --git a/mlir/test/Dialect/Arith/int-range-opts.mlir b/mlir/test/Dialect/Arith/int-range-opts.mlir
index e6e48d30cece5..0cfbdf3e33c93 100644
--- a/mlir/test/Dialect/Arith/int-range-opts.mlir
+++ b/mlir/test/Dialect/Arith/int-range-opts.mlir
@@ -148,3 +148,40 @@ func.func @analysis_crash(%arg0: i32, %arg1: tensor<128xi1>) -> tensor<128xi64>
   %2 = arith.extsi %1 : tensor<128xi32> to tensor<128xi64>
   return %2 : tensor<128xi64>
 }
+
+// -----
+
+// Regression test: block args whose values come from multiple predecessors
+// (some of which are dead) must not be incorrectly replaced with a constant
+// from a 
diff erent block arg whose storage was freed by dead-arg elimination
+// between pattern applications. Here %result depends on %cond and must not be
+// replaced. (https://github.com/llvm/llvm-project/issues/137281)
+
+// CHECK-LABEL: func @no_miscompile_block_arg_address_reuse
+// CHECK: cf.cond_br %[[COND:.*]], ^bb5, ^bb6
+// CHECK: ^bb7([[RESULT:%.*]]: i64):
+// CHECK-NEXT: vector.print [[RESULT]] : i64
+func.func @no_miscompile_block_arg_address_reuse(%cond: i1) {
+  %false = arith.constant false
+  %c4 = arith.constant 4 : i64
+  %c99 = arith.constant 99 : i64
+  cf.cond_br %false, ^bb1, ^bb2
+^bb1:
+  cf.br ^bb3(%c4, %c99 : i64, i64)
+^bb2:
+  cf.br ^bb3(%c4, %c99 : i64, i64)
+^bb3(%a: i64, %b: i64):
+  "test.use"(%a) : (i64) -> ()
+  cf.br ^bb4
+^bb4:
+  cf.cond_br %cond, ^bb5, ^bb6
+^bb5:
+  %c333 = arith.constant 333 : i64
+  cf.br ^bb7(%c333 : i64)
+^bb6:
+  %c444 = arith.constant 444 : i64
+  cf.br ^bb7(%c444 : i64)
+^bb7(%result: i64):
+  vector.print %result : i64
+  return
+}