[mlir][Interfaces] ValueBoundsOpInterface: Add API to compare values (#86915)

This commit adds a new public API to `ValueBoundsOpInterface` to compare
values/dims. Supported comparison operators are: LT, LE, EQ, GE, GT.

The new `ValueBoundsOpInterface::compare` API replaces and generalizes
`ValueBoundsOpInterface::areEqual`. Not only does it provide additional
comparison operators, it also works in cases where the difference
between the two values/dims is non-constant. The previous implementation
of `areEqual` used to compute a constant bound of `val1 - val2` (check
if it `== 0` or `!= 0`).

Note: This commit refactors, generalizes and adds a public API for
value/dim comparison. The comparison functionality itself was introduced
in #85895 and is already in use for analyzing `scf.if`.

In the long term, this improvement will allow for a more powerful
analysis of subset ops. A future commit will update
`areOverlappingSlices` to use the new comparison API.
(`areEquivalentSlices` is already using the new API.) This will improve
subset equivalence/disjointness checks with non-constant
offsets/sizes/strides.

Author: matthias-springer

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

@@ -211,7 +211,8 @@ class ValueBoundsConstraintSet
   /// Comparison operator for `ValueBoundsConstraintSet::compare`.
   enum ComparisonOperator { LT, LE, EQ, GT, GE };
 
-  /// Try to prove that, based on the current state of this constraint set
+  /// Populate constraints for lhs/rhs (until the stop condition is met). Then,
+  /// try to prove that, based on the current state of this constraint set
   /// (i.e., without analyzing additional IR or adding new constraints), the
   /// "lhs" value/dim is LE/LT/EQ/GT/GE than the "rhs" value/dim.
   ///
@@ -220,24 +221,37 @@ class ValueBoundsConstraintSet
   /// proven. This could be because the specified relation does in fact not hold
   /// or because there is not enough information in the constraint set. In other
   /// words, if we do not know for sure, this function returns "false".
-  bool compare(Value lhs, std::optional<int64_t> lhsDim, ComparisonOperator cmp,
-               Value rhs, std::optional<int64_t> rhsDim);
+  bool populateAndCompare(OpFoldResult lhs, std::optional<int64_t> lhsDim,
+                          ComparisonOperator cmp, OpFoldResult rhs,
+                          std::optional<int64_t> rhsDim);
+
+  /// Return "true" if "lhs cmp rhs" was proven to hold. Return "false" if the
+  /// specified relation could not be proven. This could be because the
+  /// specified relation does in fact not hold or because there is not enough
+  /// information in the constraint set. In other words, if we do not know for
+  /// sure, this function returns "false".
+  ///
+  /// This function keeps traversing the backward slice of lhs/rhs until could
+  /// prove the relation or until it ran out of IR.
+  static bool compare(OpFoldResult lhs, std::optional<int64_t> lhsDim,
+                      ComparisonOperator cmp, OpFoldResult rhs,
+                      std::optional<int64_t> rhsDim);
+  static bool compare(AffineMap lhs, ValueDimList lhsOperands,
+                      ComparisonOperator cmp, AffineMap rhs,
+                      ValueDimList rhsOperands);
+  static bool compare(AffineMap lhs, ArrayRef<Value> lhsOperands,
+                      ComparisonOperator cmp, AffineMap rhs,
+                      ArrayRef<Value> rhsOperands);
 
   /// Compute whether the given values/dimensions are equal. Return "failure" if
   /// equality could not be determined.
   ///
   /// `dim1`/`dim2` must be `nullopt` if and only if `value1`/`value2` are
   /// index-typed.
-  static FailureOr<bool> areEqual(Value value1, Value value2,
+  static FailureOr<bool> areEqual(OpFoldResult value1, OpFoldResult value2,
                                   std::optional<int64_t> dim1 = std::nullopt,
                                   std::optional<int64_t> dim2 = std::nullopt);
 
-  /// Compute whether the given values/attributes are equal. Return "failure" if
-  /// equality could not be determined.
-  ///
-  /// `ofr1`/`ofr2` must be of index type.
-  static FailureOr<bool> areEqual(OpFoldResult ofr1, OpFoldResult ofr2);
-
   /// Return "true" if the given slices are guaranteed to be overlapping.
   /// Return "false" if the given slices are guaranteed to be non-overlapping.
   /// Return "failure" if unknown.
@@ -294,6 +308,20 @@ class ValueBoundsConstraintSet
 
   ValueBoundsConstraintSet(MLIRContext *ctx, StopConditionFn stopCondition);
 
+  /// Return "true" if, based on the current state of the constraint system,
+  /// "lhs cmp rhs" was proven to hold. Return "false" if the specified relation
+  /// could not be proven. This could be because the specified relation does in
+  /// fact not hold or because there is not enough information in the constraint
+  /// set. In other words, if we do not know for sure, this function returns
+  /// "false".
+  ///
+  /// This function does not analyze any IR and does not populate any additional
+  /// constraints.
+  bool compareValueDims(OpFoldResult lhs, std::optional<int64_t> lhsDim,
+                        ComparisonOperator cmp, OpFoldResult rhs,
+                        std::optional<int64_t> rhsDim);
+  bool comparePos(int64_t lhsPos, ComparisonOperator cmp, int64_t rhsPos);
+
   /// Given an affine map with a single result (and map operands), add a new
   /// column to the constraint set that represents the result of the map.
   /// Traverse additional IR starting from the map operands as needed (as long
@@ -319,6 +347,10 @@ class ValueBoundsConstraintSet
   /// set.
   AffineExpr getPosExpr(int64_t pos);
 
+  /// Return "true" if the given value/dim is mapped (i.e., has a corresponding
+  /// column in the constraint system).
+  bool isMapped(Value value, std::optional<int64_t> dim = std::nullopt) const;
+
   /// Insert a value/dimension into the constraint set. If `isSymbol` is set to
   /// "false", a dimension is added. The value/dimension is added to the
   /// worklist if `addToWorklist` is set.
@@ -338,6 +370,11 @@ class ValueBoundsConstraintSet
   /// dimensions but not for symbols.
   int64_t insert(bool isSymbol = true);
 
+  /// Insert the given affine map and its bound operands as a new column in the
+  /// constraint system. Return the position of the new column. Any operands
+  /// that were not analyzed yet are put on the worklist.
+  int64_t insert(AffineMap map, ValueDimList operands, bool isSymbol = true);
+
   /// Project out the given column in the constraint set.
   void projectOut(int64_t pos);
 

mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp

@@ -58,20 +58,11 @@ struct ForOpInterface
     Value iterArg = forOp.getRegionIterArg(iterArgIdx);
     Value initArg = forOp.getInitArgs()[iterArgIdx];
 
-    // Populate constraints for the yielded value.
-    cstr.populateConstraints(yieldedValue, dim);
-    // Populate constraints for the iter_arg. This is just to ensure that the
-    // iter_arg is mapped in the constraint set, which is a prerequisite for
-    // `compare`. It may lead to a recursive call to this function in case the
-    // iter_arg was not visited when the constraints for the yielded value were
-    // populated, but no additional work is done.
-    cstr.populateConstraints(iterArg, dim);
-
     // An EQ constraint can be added if the yielded value (dimension size)
     // equals the corresponding block argument (dimension size).
-    if (cstr.compare(yieldedValue, dim,
-                     ValueBoundsConstraintSet::ComparisonOperator::EQ, iterArg,
-                     dim)) {
+    if (cstr.populateAndCompare(
+            yieldedValue, dim, ValueBoundsConstraintSet::ComparisonOperator::EQ,
+            iterArg, dim)) {
       if (dim.has_value()) {
         cstr.bound(value)[*dim] == cstr.getExpr(initArg, dim);
       } else {
@@ -113,10 +104,6 @@ struct IfOpInterface
     Value thenValue = ifOp.thenYield().getResults()[resultNum];
     Value elseValue = ifOp.elseYield().getResults()[resultNum];
 
-    // Populate constraints for the yielded value (and all values on the
-    // backward slice, as long as the current stop condition is not satisfied).
-    cstr.populateConstraints(thenValue, dim);
-    cstr.populateConstraints(elseValue, dim);
     auto boundsBuilder = cstr.bound(value);
     if (dim)
       boundsBuilder[*dim];
@@ -125,9 +112,9 @@ struct IfOpInterface
     // If thenValue <= elseValue:
     // * result <= elseValue
     // * result >= thenValue
-    if (cstr.compare(thenValue, dim,
-                     ValueBoundsConstraintSet::ComparisonOperator::LE,
-                     elseValue, dim)) {
+    if (cstr.populateAndCompare(
+            thenValue, dim, ValueBoundsConstraintSet::ComparisonOperator::LE,
+            elseValue, dim)) {
       if (dim) {
         cstr.bound(value)[*dim] >= cstr.getExpr(thenValue, dim);
         cstr.bound(value)[*dim] <= cstr.getExpr(elseValue, dim);
@@ -139,9 +126,9 @@ struct IfOpInterface
     // If elseValue <= thenValue:
     // * result <= thenValue
     // * result >= elseValue
-    if (cstr.compare(elseValue, dim,
-                     ValueBoundsConstraintSet::ComparisonOperator::LE,
-                     thenValue, dim)) {
+    if (cstr.populateAndCompare(
+            elseValue, dim, ValueBoundsConstraintSet::ComparisonOperator::LE,
+            thenValue, dim)) {
       if (dim) {
         cstr.bound(value)[*dim] >= cstr.getExpr(elseValue, dim);
         cstr.bound(value)[*dim] <= cstr.getExpr(thenValue, dim);