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

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

Differential Revision: https://reviews.llvm.org/D145796

diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h
index 0282d2e..0f2ed99 100644 (file)
--- a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h
@@ -225,11 +225,11 @@ public:
   /// 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
@@ -241,13 +241,15 @@ public:
   /// 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

diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
index 3e7d811..22a98f0 100644 (file)
--- a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
@@ -920,7 +920,8 @@ static bool detectAsFloorDiv(const FlatAffineValueConstraints &cst,
 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() &&
@@ -970,8 +971,8 @@ FlatAffineValueConstraints::getLowerAndUpperBound(
     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.
@@ -1009,7 +1010,7 @@ FlatAffineValueConstraints::getLowerAndUpperBound(
 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.
@@ -1108,7 +1109,7 @@ void FlatAffineValueConstraints::getSliceBounds(
     // 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
@@ -1140,7 +1141,8 @@ void FlatAffineValueConstraints::getSliceBounds(
           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