[Util] Fix intermittent crash in OptimizeIntArithmeticPass (#20637)

We're seeing intermittent failures in 'sharktank_model_tests ::
rocm_hip_w7900' tests on main:
-
https://github.com/iree-org/iree/actions/runs/14600283542/job/40957129060?pr=20601
-
https://github.com/iree-org/iree/actions/runs/14636457142/job/41069048559
-
https://github.com/iree-org/iree/actions/runs/14630167718/job/41051068642

The following issue contains more information on the crash and a
reproducer: #20636

I am not entirely sure whether this PR contains a fix or a workaround.
On one hand, it seems like correctly updating and invalidating all
states manually in the transformations should lead to the same result
without needing to recompute all states. On the other hand, the
`DataFlowSolver` mentions that all analysis states should be erased when
IR changes before a re-run
(https://github.com/llvm/llvm-project/blob/12a31658ea36cda74157c6b4e6b6c031e39a19c0/mlir/include/mlir/Analysis/DataFlowFramework.h#L320):

> /// Steps to re-run a data-flow analysis when IR changes:
/// 1. Erase all analysis states as they are no longer valid.
/// 2. Re-run the analysis using `initializeAndRun`.

Note: I didn't add a test because I don't have a deterministic
reproducer. As the crash is intermittent, the reproducer needs to be run
a large number of times and will fail once in 10-200 runs maybe.

Signed-off-by: Jorn Tuyls <[email protected]>

Author: jtuyls

compiler/src/iree/compiler/Dialect/Util/Transforms/OptimizeIntArithmetic.cpp

@@ -493,6 +493,12 @@ class OptimizeIntArithmeticPass
     FrozenRewritePatternSet frozenPatterns(std::move(patterns));
     for (int i = 0;; ++i) {
       LLVM_DEBUG(dbgs() << "  * Starting iteration: " << i << "\n");
+      // NOTE(jtuyls): Erasing all data flow solver states here should be safe.
+      // Manually updating and invalidating appropriate states in the
+      // inserted patterns might speed up the transformation, but this
+      // currently results in occasional crashes. See
+      // https://github.com/iree-org/iree/issues/20636.
+      solver.eraseAllStates();
       if (failed(solver.initializeAndRun(op))) {
         emitError(op->getLoc()) << "failed to perform int range analysis";
         return signalPassFailure();