[mlir][affine][Analysis] Add conservative bounds for semi-affine mods (llvm#93576)

This patch adds support for computing bounds for semi-affine mod
expression to FlatLinearConstraints. This is then enabled within the
ScalableValueBoundsConstraintSet to allow computing the bounds of
scalable remainder loops.

E.g. computing the bound of something like:
```
// `1000 mod s0` is a semi-affine.
#remainder_start_index = affine_map<()[s0] -> (-(1000 mod s0) + 1000)>
#remaining_iterations = affine_map<(d0) -> (-d0 + 1000)>

%0 = affine.apply #remainder_start_index()[%c8_vscale]
scf.for %i = %0 to %c1000 step %c8_vscale {
  %remaining_iterations = affine.apply #remaining_iterations(%i)
  // The upper bound for the remainder loop iterations should be:
  // %c8_vscale - 1  (expressed as an affine map,
  // affine_map<()[s0] -> (s0 * 8 - 1)>, where s0 is vscale)
  %bound = "test.reify_bound"(%remaining_iterations) <{scalable, ...}>
}
```

There are caveats to this implementation. To be able to add a bound for
a `mod` we need to assume the rhs is positive (> 0). This may not be
known when adding the bounds for the `mod` expression. So to handle this
a constraint is added for `rhs > 0`, this may later be found not to hold
(in which case the constraints set becomes empty/invalid).

This is not a problem for computing scalable bounds where it's safe to
assume `s0` is vscale (or some positive multiple of it). But this may
need to be considered when enabling this feature elsewhere (to ensure
correctness).

Author: MacDue

mlir/include/mlir/Analysis/FlatLinearValueConstraints.h

@@ -66,6 +66,10 @@ class FlatLinearConstraints : public presburger::IntegerPolyhedron {
   /// Return the kind of this object.
   Kind getKind() const override { return Kind::FlatLinearConstraints; }
 
+  /// Flag to control if conservative semi-affine bounds should be added in
+  /// `addBound()`.
+  enum class AddConservativeSemiAffineBounds { No = 0, Yes };
+
   /// Adds a bound for the variable at the specified position with constraints
   /// being drawn from the specified bound map. In case of an EQ bound, the
   /// bound map is expected to have exactly one result. In case of a LB/UB, the
@@ -77,21 +81,39 @@ class FlatLinearConstraints : public presburger::IntegerPolyhedron {
   /// as a closed bound by +1/-1 respectively. In case of an EQ bound, it can
   /// only be added as a closed bound.
   ///
+  /// Conservative bounds for semi-affine expressions will be added if
+  /// `AddConservativeSemiAffineBounds` is set to `Yes`. This currently only
+  /// covers semi-affine `mod` expressions, so `addBound()` will still fail if
+  /// it encounters a semi-affine `floordiv`, `ceildiv`, or `mul`. Note: If
+  /// enabled it is possible for the resulting constraint set to become empty if
+  /// a precondition of a conservative bound is found not to hold.
+  ///
   /// Note: The dimensions/symbols of this FlatLinearConstraints must match the
   /// dimensions/symbols of the affine map.
-  LogicalResult addBound(presburger::BoundType type, unsigned pos,
-                         AffineMap boundMap, bool isClosedBound);
+  LogicalResult addBound(
+      presburger::BoundType type, unsigned pos, AffineMap boundMap,
+      bool isClosedBound,
+      AddConservativeSemiAffineBounds = AddConservativeSemiAffineBounds::No);
 
   /// Adds a bound for the variable at the specified position with constraints
   /// being drawn from the specified bound map. In case of an EQ bound, the
   /// bound map is expected to have exactly one result. In case of a LB/UB, the
   /// bound map may have more than one result, for each of which an inequality
   /// is added.
+  ///
+  /// Conservative bounds for semi-affine expressions will be added if
+  /// `AddConservativeSemiAffineBounds` is set to `Yes`. This currently only
+  /// covers semi-affine `mod` expressions, so `addBound()` will still fail if
+  /// it encounters a semi-affine `floordiv`, `ceildiv`, or `mul`. Note: If
+  /// enabled it is possible for the resulting constraint set to become empty if
+  /// a precondition of a conservative bound is found not to hold.
+  ///
   /// Note: The dimensions/symbols of this FlatLinearConstraints must match the
   /// dimensions/symbols of the affine map. By default the lower bound is closed
   /// and the upper bound is open.
-  LogicalResult addBound(presburger::BoundType type, unsigned pos,
-                         AffineMap boundMap);
+  LogicalResult addBound(
+      presburger::BoundType type, unsigned pos, AffineMap boundMap,
+      AddConservativeSemiAffineBounds = AddConservativeSemiAffineBounds::No);
 
   /// The `addBound` overload above hides the inherited overloads by default, so
   /// we explicitly introduce them here.
@@ -193,7 +215,8 @@ class FlatLinearConstraints : public presburger::IntegerPolyhedron {
   /// Note: This is a shared helper function of `addLowerOrUpperBound` and
   ///       `composeMatchingMap`.
   LogicalResult flattenAlignedMapAndMergeLocals(
-      AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs);
+      AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
+      bool addConservativeSemiAffineBounds = false);
 
   /// Prints the number of constraints, dimensions, symbols and locals in the
   /// FlatLinearConstraints. Also, prints for each variable whether there is
@@ -468,18 +491,19 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
 /// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the
 /// dimensions, symbols, and additional variables that represent floor divisions
 /// of dimensions, symbols, and in turn other floor divisions.  Returns failure
-/// if 'expr' could not be flattened (i.e., semi-affine is not yet handled).
+/// if 'expr' could not be flattened (i.e., an unhandled semi-affine was found).
 /// 'cst' contains constraints that connect newly introduced local variables
 /// to existing dimensional and symbolic variables. See documentation for
 /// AffineExprFlattener on how mod's and div's are flattened.
-LogicalResult getFlattenedAffineExpr(AffineExpr expr, unsigned numDims,
-                                     unsigned numSymbols,
-                                     SmallVectorImpl<int64_t> *flattenedExpr,
-                                     FlatLinearConstraints *cst = nullptr);
+LogicalResult
+getFlattenedAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols,
+                       SmallVectorImpl<int64_t> *flattenedExpr,
+                       FlatLinearConstraints *cst = nullptr,
+                       bool addConservativeSemiAffineBounds = false);
 
 /// Flattens the result expressions of the map to their corresponding flattened
 /// forms and set in 'flattenedExprs'. Returns failure if any expression in the
-/// map could not be flattened (i.e., semi-affine is not yet handled). 'cst'
+/// map could not be flattened (i.e., an unhandled semi-affine was found). 'cst'
 /// contains constraints that connect newly introduced local variables to
 /// existing dimensional and / symbolic variables. See documentation for
 /// AffineExprFlattener on how mod's and div's are flattened. For all affine
@@ -490,7 +514,8 @@ LogicalResult getFlattenedAffineExpr(AffineExpr expr, unsigned numDims,
 LogicalResult
 getFlattenedAffineExprs(AffineMap map,
                         std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
-                        FlatLinearConstraints *cst = nullptr);
+                        FlatLinearConstraints *cst = nullptr,
+                        bool addConservativeSemiAffineBounds = false);
 LogicalResult
 getFlattenedAffineExprs(IntegerSet set,
                         std::vector<SmallVector<int64_t, 8>> *flattenedExprs,

mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h

@@ -33,8 +33,9 @@ struct ScalableValueBoundsConstraintSet
       MLIRContext *context,
       ValueBoundsConstraintSet::StopConditionFn stopCondition,
       unsigned vscaleMin, unsigned vscaleMax)
-      : RTTIExtends(context, stopCondition), vscaleMin(vscaleMin),
-        vscaleMax(vscaleMax) {};
+      : RTTIExtends(context, stopCondition,
+                    /*addConservativeSemiAffineBounds=*/true),
+        vscaleMin(vscaleMin), vscaleMax(vscaleMax) {};
 
   using RTTIExtends::bound;
   using RTTIExtends::StopConditionFn;

mlir/include/mlir/IR/AffineExprVisitor.h

@@ -413,18 +413,22 @@ class SimpleAffineExprFlattener
   /// lhs of the mod, floordiv, ceildiv or mul expression and with respect to a
   /// symbolic rhs expression. `localExpr` is the simplified tree expression
   /// (AffineExpr) corresponding to the quantifier.
-  virtual void addLocalIdSemiAffine(AffineExpr localExpr);
+  virtual LogicalResult addLocalIdSemiAffine(ArrayRef<int64_t> lhs,
+                                             ArrayRef<int64_t> rhs,
+                                             AffineExpr localExpr);
 
 private:
-  /// Adds `expr`, which may be mod, ceildiv, floordiv or mod expression
+  /// Adds `localExpr`, which may be mod, ceildiv, floordiv or mod expression
   /// representing the affine expression corresponding to the quantifier
-  /// introduced as the local variable corresponding to `expr`. If the
+  /// introduced as the local variable corresponding to `localExpr`. If the
   /// quantifier is already present, we put the coefficient in the proper index
   /// of `result`, otherwise we add a new local variable and put the coefficient
   /// there.
-  void addLocalVariableSemiAffine(AffineExpr expr,
-                                  SmallVectorImpl<int64_t> &result,
-                                  unsigned long resultSize);
+  LogicalResult addLocalVariableSemiAffine(ArrayRef<int64_t> lhs,
+                                           ArrayRef<int64_t> rhs,
+                                           AffineExpr localExpr,
+                                           SmallVectorImpl<int64_t> &result,
+                                           unsigned long resultSize);
 
   // t = expr floordiv c   <=> t = q, c * q <= expr <= c * q + c - 1
   // A floordiv is thus flattened by introducing a new local variable q, and

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

@@ -313,7 +313,8 @@ class ValueBoundsConstraintSet
   /// An index-typed value or the dimension of a shaped-type value.
   using ValueDim = std::pair<Value, int64_t>;
 
-  ValueBoundsConstraintSet(MLIRContext *ctx, StopConditionFn stopCondition);
+  ValueBoundsConstraintSet(MLIRContext *ctx, StopConditionFn stopCondition,
+                           bool addConservativeSemiAffineBounds = false);
 
   /// 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
@@ -404,6 +405,9 @@ class ValueBoundsConstraintSet
 
   /// The current stop condition function.
   StopConditionFn stopCondition = nullptr;
+
+  /// Should conservative bounds be added for semi-affine expressions.
+  bool addConservativeSemiAffineBounds = false;
 };
 
 } // namespace mlir