[mlir] fix maybeReplaceWithConstant in IntRangeOptimizations

[makslevental]

If a dialect is caching/reusing constants when materializing then such constants might already have IntegerValueRangeLattices associated with them and the range endpoint bit widths might not match the new replacement (amongst other possible wackiness).

I observed this with %true = arith.constant true which was materialized but had an existing IntegerValueRangeLattice (i.e., solver.getOrCreateState<dataflow::IntegerValueRangeLattice> was not uninitalized) with range endpoint bit widths:

umin bit width: 32
umax bit width: 32
smin bit width: 32
smax bit width: 32

while the widths of the range end points for something like %20 = arith.cmpi slt, %19, %c1_i32 (a replacement candidate) would be

umin bit width: 1
umax bit width: 1
smin bit width: 1
smax bit width: 1

Thus, we should be clearing the analysis state each time a constant is reused.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

If a dialect is caching/reusing constants when materializing then such constants might already have IntegerValueRangeLattice associated with them and the range endpoint bit widths might not match the new replacement; I observed this with %true = arith.constant true which was materialized but had the range endpoint bit widths:

umin: 32
umax: 32
smin: 32
smax: 32

Thus, we should be clearing the analysis state each time a replacement is made.

---

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

1 Files Affected:

• (modified) mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp (+5-2)

diff --git a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
index 1cb9453ccf3c9..602d80a45993e 100644
--- a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
@@ -91,8 +91,11 @@ LogicalResult maybeReplaceWithConstant(DataFlowSolver &solver,
   if (!constOp)
     return failure();
 
-  copyIntegerRange(solver, value, constOp->getResult(0));
-  rewriter.replaceAllUsesWith(value, constOp->getResult(0));
+  OpResult res = constOp->getResult(0);
+  if (solver.lookupState<dataflow::IntegerValueRangeLattice>(res))
+    solver.eraseState(res);
+  copyIntegerRange(solver, value, res);
+  rewriter.replaceAllUsesWith(value, res);
   return success();
 }
 } // namespace mlir::dataflow

[llvmbot]

@llvm/pr-subscribers-mlir-arith

Author: Maksim Levental (makslevental)

Changes

If a dialect is caching/reusing constants when materializing then such constants might already have IntegerValueRangeLattice associated with them and the range endpoint bit widths might not match the new replacement; I observed this with %true = arith.constant true which was materialized but had the range endpoint bit widths:

umin: 32
umax: 32
smin: 32
smax: 32

Thus, we should be clearing the analysis state each time a replacement is made.

---

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

1 Files Affected:

• (modified) mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp (+5-2)

diff --git a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
index 1cb9453ccf3c9..602d80a45993e 100644
--- a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
@@ -91,8 +91,11 @@ LogicalResult maybeReplaceWithConstant(DataFlowSolver &solver,
   if (!constOp)
     return failure();
 
-  copyIntegerRange(solver, value, constOp->getResult(0));
-  rewriter.replaceAllUsesWith(value, constOp->getResult(0));
+  OpResult res = constOp->getResult(0);
+  if (solver.lookupState<dataflow::IntegerValueRangeLattice>(res))
+    solver.eraseState(res);
+  copyIntegerRange(solver, value, res);
+  rewriter.replaceAllUsesWith(value, res);
   return success();
 }
 } // namespace mlir::dataflow