[mlir][affine][analysis] Fix `getClosedUB` handling in `getSliceBounds`

[springerm]

Summary

There were cases in which ubAdjustment was not applied to the resulting UB.

Depends On: D145696

[bondhugula]

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

228

The option can be simply named closedUb.

248

If closedUb is true, the upper bound is closed.

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

1145

get suffix can be dropped; it's implied and should be omitted to be consistent with such options throughout the codebase.

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

public:
   /// Computes the lower and upper bounds of the first `num` dimensional
   /// variables (starting at `offset`) as an affine map of the remaining
   /// variables (dimensional and symbolic). This method is able to detect
   /// variables as floordiv's and mod's of affine expressions of other
   /// variables with respect to (positive) constants. Sets bound map to a
   /// null AffineMap if such a bound can't be found (or yet unimplemented).
   ///
   /// By default the returned lower bounds are closed and upper bounds are open.
-  /// This can be changed by getClosedUB.
+  /// If `closedUb` is true, the upper bound is closed.
   void getSliceBounds(unsigned offset, unsigned num, MLIRContext *context,
                       SmallVectorImpl<AffineMap> *lbMaps,
                       SmallVectorImpl<AffineMap> *ubMaps,
-                      bool getClosedUB = false);
+                      bool closedUB = false);
 
   /// Composes an affine map whose dimensions and symbols match one to one with
   /// the dimensions and symbols of this FlatAffineConstraints. The results of
   /// the map `other` are added as the leading dimensions of this constraint
   /// system. Returns failure if `other` is a semi-affine map.
   LogicalResult composeMatchingMap(AffineMap other);
 
   /// Gets the lower and upper bound of the `offset` + `pos`th variable
   /// treating [0, offset) U [offset + num, symStartPos) as dimensions and
   /// [symStartPos, getNumDimAndSymbolVars) as symbols, and `pos` lies in
   /// [0, num). The multi-dimensional maps in the returned pair represent the
-  /// max and min of potentially multiple affine expressions. The upper bound is
-  /// exclusive. `localExprs` holds pre-computed AffineExpr's for all local
-  /// variables in the system.
+  /// max and min of potentially multiple affine expressions. `localExprs` holds
+  /// pre-computed AffineExpr's for all local variables in the system.
+  ///
+  /// By default the returned lower bounds are closed and upper bounds are open.
+  /// If `closedUb` is true, the upper bound is closed.
   std::pair<AffineMap, AffineMap>
   getLowerAndUpperBound(unsigned pos, unsigned offset, unsigned num,
                         unsigned symStartPos, ArrayRef<AffineExpr> localExprs,
-                        MLIRContext *context) const;
+                        MLIRContext *context, bool closedUB = false) const;
 
   /// Returns the bound for the variable at `pos` from the inequality at
   /// `ineqPos` as a 1-d affine value map (affine map + operands). The returned
   /// affine value map can either be a lower bound or an upper bound depending
   /// on the sign of atIneq(ineqPos, pos). Asserts if the row at `ineqPos` does
   /// not involve the `pos`th variable.
   void getIneqAsAffineValueMap(unsigned pos, unsigned ineqPos,
                                AffineValueMap &vmap,

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

static bool detectAsFloorDiv(const FlatAffineValueConstraints &cst,
   // Successfully detected the floordiv.
   exprs[pos] = dividendExpr.floorDiv(divisor);
   return true;
 }
 
 std::pair<AffineMap, AffineMap>
 FlatAffineValueConstraints::getLowerAndUpperBound(
     unsigned pos, unsigned offset, unsigned num, unsigned symStartPos,
-    ArrayRef<AffineExpr> localExprs, MLIRContext *context) const {
+    ArrayRef<AffineExpr> localExprs, MLIRContext *context,
+    bool closedUB) const {
   assert(pos + offset < getNumDimVars() && "invalid dim start pos");
   assert(symStartPos >= (pos + offset) && "invalid sym start pos");
   assert(getNumLocalVars() == localExprs.size() &&
          "incorrect local exprs count");
 
   SmallVector<unsigned, 4> lbIndices, ubIndices, eqIndices;
   getLowerAndUpperBoundIndices(pos + offset, &lbIndices, &ubIndices, &eqIndices,
                                offset, num);
FlatAffineValueConstraints::getLowerAndUpperBound(
   // Upper bound expressions.
   for (auto idx : ubIndices) {
     auto ineq = getInequality64(idx);
     // Extract the upper bound (in terms of other coeff's + const).
     addCoeffs(ineq, ub);
     auto expr =
         getAffineExprFromFlatForm(ub, dimCount, symCount, localExprs, context);
     expr = expr.floorDiv(std::abs(ineq[pos + offset]));
-    // Upper bound is exclusive.
-    ubExprs.push_back(expr + 1);
+    int64_t ubAdjustment = closedUB ? 0 : 1;
+    ubExprs.push_back(expr + ubAdjustment);
   }
 
   // Equalities. It's both a lower and a upper bound.
   SmallVector<int64_t, 4> b;
   for (auto idx : eqIndices) {
     auto eq = getEquality64(idx);
     addCoeffs(eq, b);
     if (eq[pos + offset] > 0)

 /// Computes the lower and upper bounds of the first 'num' dimensional
 /// variables (starting at 'offset') as affine maps of the remaining
 /// variables (dimensional and symbolic variables). Local variables are
 /// themselves explicitly computed as affine functions of other variables in
 /// this process if needed.
 void FlatAffineValueConstraints::getSliceBounds(
     unsigned offset, unsigned num, MLIRContext *context,
     SmallVectorImpl<AffineMap> *lbMaps, SmallVectorImpl<AffineMap> *ubMaps,
-    bool getClosedUB) {
+    bool closedUB) {
   assert(num < getNumDimVars() && "invalid range");
 
   // Basic simplification.
   normalizeConstraintsByGCD();
 
   LLVM_DEBUG(llvm::dbgs() << "getSliceBounds for first " << num
                           << " variables\n");
   LLVM_DEBUG(dump());
for (unsigned pos = 0; pos < getNumVars(); pos++) {
       memo[pos] = expr;
       changed = true;
     }
     // This loop is guaranteed to reach a fixed point - since once an
     // variable's explicit form is computed (in memo[pos]), it's not updated
     // again.
   } while (changed);
 
-  int64_t ubAdjustment = getClosedUB ? 0 : 1;
+  int64_t ubAdjustment = closedUB ? 0 : 1;
 
   // Set the lower and upper bound maps for all the variables that were
   // computed as affine expressions of the rest as the "detected expr" and
   // "detected expr + 1" respectively; set the undetected ones to null.
   std::optional<FlatAffineValueConstraints> tmpClone;
   for (unsigned pos = 0; pos < num; pos++) {
     unsigned numMapDims = getNumDimVars() - num;
     unsigned numMapSymbols = getNumSymbolVars();
if (expr) {
         // Work on a copy so that we don't update this constraint system.
         if (!tmpClone) {
           tmpClone.emplace(FlatAffineValueConstraints(*this));
           // Removing redundant inequalities is necessary so that we don't get
           // redundant loop bounds.
           tmpClone->removeRedundantInequalities();
         }
         std::tie(lbMap, ubMap) = tmpClone->getLowerAndUpperBound(
-            pos, offset, num, getNumDimVars(), /*localExprs=*/{}, context);
+            pos, offset, num, getNumDimVars(), /*localExprs=*/{}, context,
+            closedUB);
       }
 
       // If the above fails, we'll just use the constant lower bound and the
       // constant upper bound (if they exist) as the slice bounds.
       // TODO: being conservative for the moment in cases that
       // lead to multiple bounds - until getConstDifference in LoopFusion.cpp is
       // fixed (b/126426796).
       if (!lbMap || lbMap.getNumResults() > 1) {