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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/188992.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp (+13-8) 
- (modified) mlir/test/Dialect/Arith/int-range-opts.mlir (+37) 


``````````diff
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 different 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 different 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
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/188992