[Mlir-commits] [mlir] [mlir][integer-range-analysis] expose helpers in header and fix ConstantIntRange print (PR #127888)

[Maksim Levental]

https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/127888

>From 39782228fc87272677f68229e5ef0850c8a88c76 Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental at gmail.com>
Date: Wed, 19 Feb 2025 14:44:42 -0500
Subject: [PATCH] [mlir][integer-range-analysis] expose helpers in header and
 fix ConstantIntRange print

---
 .../Analysis/DataFlow/IntegerRangeAnalysis.h  | 13 ++++++++++
 .../DataFlow/IntegerRangeAnalysis.cpp         | 18 +++++++++++++
 .../Transforms/UnsignedWhenEquivalent.cpp     | 26 -------------------
 .../lib/Interfaces/InferIntRangeInterface.cpp |  7 +++--
 4 files changed, 36 insertions(+), 28 deletions(-)

diff --git a/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
index f99eae379596b..6dc5a28a23b95 100644
--- a/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
@@ -71,6 +71,19 @@ class IntegerRangeAnalysis
                                unsigned firstIndex) override;
 };
 
+/// Succeeds if an op can be converted to its unsigned equivalent without
+/// changing its semantics. This is the case when none of its openands or
+/// results can be below 0 when analyzed from a signed perspective.
+LogicalResult staticallyNonNegative(DataFlowSolver &solver, Operation *op);
+
+/// Succeeds when a value is statically non-negative in that it has a lower
+/// bound on its value (if it is treated as signed) and that bound is
+/// non-negative.
+// TODO: IntegerRangeAnalysis internally assumes index is 64bit and this pattern
+// relies on this. These transformations may not be valid for 32bit index,
+// need more investigation.
+LogicalResult staticallyNonNegative(DataFlowSolver &solver, Value v);
+
 } // end namespace dataflow
 } // end namespace mlir
 
diff --git a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
index 722f4df18e981..c7a950d9a8871 100644
--- a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
@@ -37,6 +37,24 @@
 using namespace mlir;
 using namespace mlir::dataflow;
 
+namespace mlir::dataflow {
+LogicalResult staticallyNonNegative(DataFlowSolver &solver, Value v) {
+  auto *result = solver.lookupState<IntegerValueRangeLattice>(v);
+  if (!result || result->getValue().isUninitialized())
+    return failure();
+  const ConstantIntRanges &range = result->getValue().getValue();
+  return success(range.smin().isNonNegative());
+}
+
+LogicalResult staticallyNonNegative(DataFlowSolver &solver, Operation *op) {
+  auto nonNegativePred = [&solver](Value v) -> bool {
+    return succeeded(staticallyNonNegative(solver, v));
+  };
+  return success(llvm::all_of(op->getOperands(), nonNegativePred) &&
+                 llvm::all_of(op->getResults(), nonNegativePred));
+}
+} // namespace mlir::dataflow
+
 void IntegerValueRangeLattice::onUpdate(DataFlowSolver *solver) const {
   Lattice::onUpdate(solver);
 
diff --git a/mlir/lib/Dialect/Arith/Transforms/UnsignedWhenEquivalent.cpp b/mlir/lib/Dialect/Arith/Transforms/UnsignedWhenEquivalent.cpp
index 8922e93e399f9..9906028c21720 100644
--- a/mlir/lib/Dialect/Arith/Transforms/UnsignedWhenEquivalent.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/UnsignedWhenEquivalent.cpp
@@ -27,32 +27,6 @@ using namespace mlir;
 using namespace mlir::arith;
 using namespace mlir::dataflow;
 
-/// Succeeds when a value is statically non-negative in that it has a lower
-/// bound on its value (if it is treated as signed) and that bound is
-/// non-negative.
-// TODO: IntegerRangeAnalysis internally assumes index is 64bit and this pattern
-// relies on this. These transformations may not be valid for 32bit index,
-// need more investigation.
-static LogicalResult staticallyNonNegative(DataFlowSolver &solver, Value v) {
-  auto *result = solver.lookupState<IntegerValueRangeLattice>(v);
-  if (!result || result->getValue().isUninitialized())
-    return failure();
-  const ConstantIntRanges &range = result->getValue().getValue();
-  return success(range.smin().isNonNegative());
-}
-
-/// Succeeds if an op can be converted to its unsigned equivalent without
-/// changing its semantics. This is the case when none of its openands or
-/// results can be below 0 when analyzed from a signed perspective.
-static LogicalResult staticallyNonNegative(DataFlowSolver &solver,
-                                           Operation *op) {
-  auto nonNegativePred = [&solver](Value v) -> bool {
-    return succeeded(staticallyNonNegative(solver, v));
-  };
-  return success(llvm::all_of(op->getOperands(), nonNegativePred) &&
-                 llvm::all_of(op->getResults(), nonNegativePred));
-}
-
 /// Succeeds when the comparison predicate is a signed operation and all the
 /// operands are non-negative, indicating that the cmpi operation `op` can have
 /// its predicate changed to an unsigned equivalent.
diff --git a/mlir/lib/Interfaces/InferIntRangeInterface.cpp b/mlir/lib/Interfaces/InferIntRangeInterface.cpp
index 63658518dd4a3..1801e3f7c52fd 100644
--- a/mlir/lib/Interfaces/InferIntRangeInterface.cpp
+++ b/mlir/lib/Interfaces/InferIntRangeInterface.cpp
@@ -125,8 +125,11 @@ std::optional<APInt> ConstantIntRanges::getConstantValue() const {
 }
 
 raw_ostream &mlir::operator<<(raw_ostream &os, const ConstantIntRanges &range) {
-  return os << "unsigned : [" << range.umin() << ", " << range.umax()
-            << "] signed : [" << range.smin() << ", " << range.smax() << "]";
+  os << "unsigned : [";
+  range.umin().print(os, /*isSigned*/ false);
+  os << ", ";
+  range.umax().print(os, /*isSigned*/ false);
+  return os << "] signed : [" << range.smin() << ", " << range.smax() << "]";
 }
 
 IntegerValueRange IntegerValueRange::getMaxRange(Value value) {