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

Author: makslevental

mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h

@@ -71,6 +71,18 @@ 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.
+/// Note, the results of this query may not be accurate for `index` if you plan
+/// to use a non-64-bit index.
+LogicalResult staticallyNonNegative(DataFlowSolver &solver, Value v);
+
 } // end namespace dataflow
 } // end namespace mlir
 

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);
 

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.
@@ -103,6 +77,10 @@ class DataFlowListener : public RewriterBase::Listener {
   DataFlowSolver &s;
 };
 
+// TODO: IntegerRangeAnalysis internally assumes index is 64bit and this pattern
+// (via staticallyNonNegative) relies on this. These transformations may not be
+// valid for 32bit index, need more investigation.
+
 template <typename Signed, typename Unsigned>
 struct ConvertOpToUnsigned final : OpRewritePattern<Signed> {
   ConvertOpToUnsigned(MLIRContext *context, DataFlowSolver &s)

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) {