diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h --- a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h +++ b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h @@ -225,11 +225,11 @@ /// 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 *lbMaps, SmallVectorImpl *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 @@ /// 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 getLowerAndUpperBound(unsigned pos, unsigned offset, unsigned num, unsigned symStartPos, ArrayRef 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 --- a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp +++ b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp @@ -920,7 +920,8 @@ std::pair FlatAffineValueConstraints::getLowerAndUpperBound( unsigned pos, unsigned offset, unsigned num, unsigned symStartPos, - ArrayRef localExprs, MLIRContext *context) const { + ArrayRef 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 @@ 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 @@ void FlatAffineValueConstraints::getSliceBounds( unsigned offset, unsigned num, MLIRContext *context, SmallVectorImpl *lbMaps, SmallVectorImpl *ubMaps, - bool getClosedUB) { + bool closedUB) { assert(num < getNumDimVars() && "invalid range"); // Basic simplification. @@ -1108,7 +1109,7 @@ // 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 @@ 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