[mlir][affine][Analysis] Add conservative bounds for semi-affine mods

This path 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:
```

%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 does not
+  /// cover all semi-affine expressions, so `addBound()` still may fail with
+  /// this set. 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 does not
+  /// cover all semi-affine expressions, so `addBound()` still may fail with
+  /// this set. 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/Analysis/Presburger/IntegerRelation.h

@@ -454,6 +454,20 @@ class IntegerRelation {
     addLocalFloorDiv(getMPIntVec(dividend), MPInt(divisor));
   }
 
+  /// Adds a new local variable as the mod of an affine function of other
+  /// variables. The coefficients of the operands of the mod are provided in
+  /// `lhs` and `rhs` respectively. Three constraints are added to provide a
+  /// conservative bound for the mod:
+  ///  1. rhs > 0 (assumption/precondition)
+  ///  2. lhs % rhs < rhs
+  ///  3. lhs % rhs >= 0
+  /// We ensure the rhs is positive so we can assume the result is positive.
+  void addLocalModConservativeBounds(ArrayRef<MPInt> lhs, ArrayRef<MPInt> rhs);
+  void addLocalModConservativeBounds(ArrayRef<int64_t> lhs,
+                                     ArrayRef<int64_t> rhs) {
+    addLocalModConservativeBounds(getMPIntVec(lhs), getMPIntVec(rhs));
+  }
+
   /// Projects out (aka eliminates) `num` variables starting at position
   /// `pos`. The resulting constraint system is the shadow along the dimensions
   /// that still exist. This method may not always be integer exact.

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,7 +413,8 @@ 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 void addLocalIdSemiAffine(AffineExpr localExpr, ArrayRef<int64_t> lhs,
+                                    ArrayRef<int64_t> rhs);
 
 private:
   /// Adds `expr`, which may be mod, ceildiv, floordiv or mod expression
@@ -422,7 +423,8 @@ class SimpleAffineExprFlattener
   /// 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,
+  void addLocalVariableSemiAffine(AffineExpr expr, ArrayRef<int64_t> lhs,
+                                  ArrayRef<int64_t> rhs,
                                   SmallVectorImpl<int64_t> &result,
                                   unsigned long resultSize);
 

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

mlir/lib/Analysis/FlatLinearValueConstraints.cpp

@@ -46,9 +46,15 @@ struct AffineExprFlattener : public SimpleAffineExprFlattener {
   // inequalities.
   IntegerPolyhedron localVarCst;
 
-  AffineExprFlattener(unsigned nDims, unsigned nSymbols)
+  AffineExprFlattener(unsigned nDims, unsigned nSymbols,
+                      bool addConservativeSemiAffineBounds = false)
       : SimpleAffineExprFlattener(nDims, nSymbols),
-        localVarCst(PresburgerSpace::getSetSpace(nDims, nSymbols)) {}
+        localVarCst(PresburgerSpace::getSetSpace(nDims, nSymbols)),
+        addConservativeSemiAffineBounds(addConservativeSemiAffineBounds) {}
+
+  bool hasUnhandledSemiAffineExpressions() const {
+    return unhandledSemiAffineExpressions;
+  }
 
 private:
   // Add a local variable (needed to flatten a mod, floordiv, ceildiv expr).
@@ -63,35 +69,61 @@ struct AffineExprFlattener : public SimpleAffineExprFlattener {
     // Update localVarCst.
     localVarCst.addLocalFloorDiv(dividend, divisor);
   }
+
+  // Add a local identifier (needed to flatten a mod, floordiv, ceildiv, mul
+  // expr) when the rhs is a symbolic expression. The local identifier added
+  // may be a floordiv, ceildiv, mul or mod of a pure affine/semi-affine
+  // function of other identifiers, coefficients of which are specified in the
+  // 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.
+  void addLocalIdSemiAffine(AffineExpr localExpr, ArrayRef<int64_t> lhs,
+                            ArrayRef<int64_t> rhs) override {
+    SimpleAffineExprFlattener::addLocalIdSemiAffine(localExpr, lhs, rhs);
+    if (!addConservativeSemiAffineBounds) {
+      unhandledSemiAffineExpressions = true;
+      return;
+    }
+    if (localExpr.getKind() == AffineExprKind::Mod) {
+      localVarCst.addLocalModConservativeBounds(lhs, rhs);
+      return;
+    }
+    // TODO: Support other semi-affine expressions.
+    unhandledSemiAffineExpressions = true;
+  }
+
+  bool addConservativeSemiAffineBounds = false;
+  bool unhandledSemiAffineExpressions = false;
 };
 
 } // namespace
 
 // Flattens the expressions in map. Returns failure if 'expr' was unable to be
 // flattened. For example two specific cases:
-// 1. semi-affine expressions not handled yet.
+// 1. an unhandled semi-affine expressions is found.
 // 2. has poison expression (i.e., division by zero).
 static LogicalResult
 getFlattenedAffineExprs(ArrayRef<AffineExpr> exprs, unsigned numDims,
                         unsigned numSymbols,
                         std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
-                        FlatLinearConstraints *localVarCst) {
+                        FlatLinearConstraints *localVarCst,
+                        bool addConservativeSemiAffineBounds = false) {
   if (exprs.empty()) {
     if (localVarCst)
       *localVarCst = FlatLinearConstraints(numDims, numSymbols);
     return success();
   }
 
-  AffineExprFlattener flattener(numDims, numSymbols);
+  AffineExprFlattener flattener(numDims, numSymbols,
+                                addConservativeSemiAffineBounds);
   // Use the same flattener to simplify each expression successively. This way
   // local variables / expressions are shared.
   for (auto expr : exprs) {
-    if (!expr.isPureAffine())
-      return failure();
-    // has poison expression
     auto flattenResult = flattener.walkPostOrder(expr);
     if (failed(flattenResult))
       return failure();
+    if (flattener.hasUnhandledSemiAffineExpressions())
+      return failure();
   }
 
   assert(flattener.operandExprStack.size() == exprs.size());
@@ -106,33 +138,33 @@ getFlattenedAffineExprs(ArrayRef<AffineExpr> exprs, unsigned numDims,
 }
 
 // Flattens 'expr' into 'flattenedExpr'. Returns failure if 'expr' was unable to
-// be flattened (semi-affine expressions not handled yet).
-LogicalResult
-mlir::getFlattenedAffineExpr(AffineExpr expr, unsigned numDims,
-                             unsigned numSymbols,
-                             SmallVectorImpl<int64_t> *flattenedExpr,
-                             FlatLinearConstraints *localVarCst) {
+// be flattened (an unhandled semi-affine was found).
+LogicalResult mlir::getFlattenedAffineExpr(
+    AffineExpr expr, unsigned numDims, unsigned numSymbols,
+    SmallVectorImpl<int64_t> *flattenedExpr, FlatLinearConstraints *localVarCst,
+    bool addConservativeSemiAffineBounds) {
   std::vector<SmallVector<int64_t, 8>> flattenedExprs;
-  LogicalResult ret = ::getFlattenedAffineExprs({expr}, numDims, numSymbols,
-                                                &flattenedExprs, localVarCst);
+  LogicalResult ret =
+      ::getFlattenedAffineExprs({expr}, numDims, numSymbols, &flattenedExprs,
+                                localVarCst, addConservativeSemiAffineBounds);
   *flattenedExpr = flattenedExprs[0];
   return ret;
 }
 
 /// Flattens the expressions in map. Returns failure if 'expr' was unable to be
-/// flattened (i.e., semi-affine expressions not handled yet).
+/// flattened (i.e., an unhandled semi-affine was found).
 LogicalResult mlir::getFlattenedAffineExprs(
     AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
-    FlatLinearConstraints *localVarCst) {
+    FlatLinearConstraints *localVarCst, bool addConservativeSemiAffineBounds) {
   if (map.getNumResults() == 0) {
     if (localVarCst)
       *localVarCst =
           FlatLinearConstraints(map.getNumDims(), map.getNumSymbols());
     return success();
   }
-  return ::getFlattenedAffineExprs(map.getResults(), map.getNumDims(),
-                                   map.getNumSymbols(), flattenedExprs,
-                                   localVarCst);
+  return ::getFlattenedAffineExprs(
+      map.getResults(), map.getNumDims(), map.getNumSymbols(), flattenedExprs,
+      localVarCst, addConservativeSemiAffineBounds);
 }
 
 LogicalResult mlir::getFlattenedAffineExprs(
@@ -641,9 +673,11 @@ void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
 }
 
 LogicalResult FlatLinearConstraints::flattenAlignedMapAndMergeLocals(
-    AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs) {
+    AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
+    bool addConservativeSemiAffineBounds) {
   FlatLinearConstraints localCst;
-  if (failed(getFlattenedAffineExprs(map, flattenedExprs, &localCst))) {
+  if (failed(getFlattenedAffineExprs(map, flattenedExprs, &localCst,
+                                     addConservativeSemiAffineBounds))) {
     LLVM_DEBUG(llvm::dbgs()
                << "composition unimplemented for semi-affine maps\n");
     return failure();
@@ -664,9 +698,9 @@ LogicalResult FlatLinearConstraints::flattenAlignedMapAndMergeLocals(
   return success();
 }
 
-LogicalResult FlatLinearConstraints::addBound(BoundType type, unsigned pos,
-                                              AffineMap boundMap,
-                                              bool isClosedBound) {
+LogicalResult FlatLinearConstraints::addBound(
+    BoundType type, unsigned pos, AffineMap boundMap, bool isClosedBound,
+    AddConservativeSemiAffineBounds addSemiAffineBounds) {
   assert(boundMap.getNumDims() == getNumDimVars() && "dim mismatch");
   assert(boundMap.getNumSymbols() == getNumSymbolVars() && "symbol mismatch");
   assert(pos < getNumDimAndSymbolVars() && "invalid position");
@@ -680,7 +714,9 @@ LogicalResult FlatLinearConstraints::addBound(BoundType type, unsigned pos,
   bool lower = type == BoundType::LB || type == BoundType::EQ;
 
   std::vector<SmallVector<int64_t, 8>> flatExprs;
-  if (failed(flattenAlignedMapAndMergeLocals(boundMap, &flatExprs)))
+  if (failed(flattenAlignedMapAndMergeLocals(
+          boundMap, &flatExprs,
+          addSemiAffineBounds == AddConservativeSemiAffineBounds::Yes)))
     return failure();
   assert(flatExprs.size() == boundMap.getNumResults());
 
@@ -716,9 +752,11 @@ LogicalResult FlatLinearConstraints::addBound(BoundType type, unsigned pos,
   return success();
 }
 
-LogicalResult FlatLinearConstraints::addBound(BoundType type, unsigned pos,
-                                              AffineMap boundMap) {
-  return addBound(type, pos, boundMap, /*isClosedBound=*/type != BoundType::UB);
+LogicalResult FlatLinearConstraints::addBound(
+    BoundType type, unsigned pos, AffineMap boundMap,
+    AddConservativeSemiAffineBounds addSemiAffineBounds) {
+  return addBound(type, pos, boundMap,
+                  /*isClosedBound=*/type != BoundType::UB, addSemiAffineBounds);
 }
 
 /// Compute an explicit representation for local vars. For all systems coming

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

@@ -1521,6 +1521,37 @@ void IntegerRelation::addLocalFloorDiv(ArrayRef<MPInt> dividend,
       getDivUpperBound(dividendCopy, divisor, dividendCopy.size() - 2));
 }
 
+/// Adds a new local variable as the mod of an affine function of other
+/// variables. The coefficients of the operands of the mod are provided in `lhs`
+/// and `rhs` respectively. Three constraints are added to provide a
+/// conservative bound for the mod:
+///  1. rhs > 0 (assumption/precondition)
+///  2. lhs % rhs < rhs
+///  3. lhs % rhs >= 0
+/// We ensure the rhs is positive so we can assume the result is positive.
+void IntegerRelation::addLocalModConservativeBounds(ArrayRef<MPInt> lhs,
+                                                    ArrayRef<MPInt> rhs) {
+  appendVar(VarKind::Local);
+
+  // Ensure the rhs is > 0 (most likely case).
+  // If this constraint does not hold the following bounds are incorrect.
+  SmallVector<MPInt, 8> rhsCopy(rhs);
+  rhsCopy.insert(rhsCopy.end() - 1, MPInt(0));
+  rhsCopy.back() -= MPInt(1);
+  addInequality(rhsCopy);
+
+  // rhs - (lhs % rhs) - 1 >= 0 i.e. lhs % rhs < rhs
+  SmallVector<MPInt, 8> resultUpperBound(rhs);
+  resultUpperBound.insert(resultUpperBound.end() - 1, MPInt(-1));
+  resultUpperBound.back() -= MPInt(1);
+  addInequality(resultUpperBound);
+
+  // lhs % rhs >= 0
+  SmallVector<MPInt, 8> resultLowerBound(rhs.size());
+  resultLowerBound.insert(resultLowerBound.end() - 1, MPInt(1));
+  addInequality(resultLowerBound);
+}
+
 /// Finds an equality that equates the specified variable to a constant.
 /// Returns the position of the equality row. If 'symbolic' is set to true,
 /// symbols are also treated like a constant, i.e., an affine function of the