diff --git a/mlir/include/mlir/Analysis/LoopAnalysis.h b/mlir/include/mlir/Analysis/LoopAnalysis.h --- a/mlir/include/mlir/Analysis/LoopAnalysis.h +++ b/mlir/include/mlir/Analysis/LoopAnalysis.h @@ -34,10 +34,8 @@ /// multi-result map. The trip count expression is simplified before returning. /// This method only utilizes map composition to construct lower and upper /// bounds before computing the trip count expressions -// TODO: this should be moved into 'Transforms/' and be replaced by a pure -// analysis method relying on FlatAffineConstraints -void buildTripCountMapAndOperands(AffineForOp forOp, AffineMap *map, - SmallVectorImpl *operands); +void getTripCountMapAndOperands(AffineForOp forOp, AffineMap *map, + SmallVectorImpl *operands); /// Returns the trip count of the loop if it's a constant, None otherwise. This /// uses affine expression analysis and is able to determine constant trip count diff --git a/mlir/lib/Analysis/LoopAnalysis.cpp b/mlir/lib/Analysis/LoopAnalysis.cpp --- a/mlir/lib/Analysis/LoopAnalysis.cpp +++ b/mlir/lib/Analysis/LoopAnalysis.cpp @@ -32,21 +32,19 @@ /// expression is simplified before returning. This method only utilizes map /// composition to construct lower and upper bounds before computing the trip /// count expressions. -void mlir::buildTripCountMapAndOperands( +void mlir::getTripCountMapAndOperands( AffineForOp forOp, AffineMap *tripCountMap, SmallVectorImpl *tripCountOperands) { - int64_t loopSpan; - + MLIRContext *context = forOp.getContext(); int64_t step = forOp.getStep(); - OpBuilder b(forOp.getOperation()); - + int64_t loopSpan; if (forOp.hasConstantBounds()) { int64_t lb = forOp.getConstantLowerBound(); int64_t ub = forOp.getConstantUpperBound(); loopSpan = ub - lb; if (loopSpan < 0) loopSpan = 0; - *tripCountMap = b.getConstantAffineMap(ceilDiv(loopSpan, step)); + *tripCountMap = AffineMap::getConstantMap(ceilDiv(loopSpan, step), context); tripCountOperands->clear(); return; } @@ -65,7 +63,7 @@ SmallVector lbSplatExpr(ubValueMap.getNumResults(), lbMap.getResult(0)); auto lbMapSplat = AffineMap::get(lbMap.getNumDims(), lbMap.getNumSymbols(), - lbSplatExpr, b.getContext()); + lbSplatExpr, context); AffineValueMap lbSplatValueMap(lbMapSplat, forOp.getLowerBoundOperands()); AffineValueMap tripCountValueMap; @@ -82,14 +80,10 @@ /// Returns the trip count of the loop if it's a constant, None otherwise. This /// method uses affine expression analysis (in turn using getTripCount) and is /// able to determine constant trip count in non-trivial cases. -// FIXME(mlir-team): this is really relying on buildTripCountMapAndOperands; -// being an analysis utility, it shouldn't. Replace with a version that just -// works with analysis structures (FlatAffineConstraints) and thus doesn't -// update the IR. Optional mlir::getConstantTripCount(AffineForOp forOp) { SmallVector operands; AffineMap map; - buildTripCountMapAndOperands(forOp, &map, &operands); + getTripCountMapAndOperands(forOp, &map, &operands); if (!map) return None; @@ -115,7 +109,7 @@ uint64_t mlir::getLargestDivisorOfTripCount(AffineForOp forOp) { SmallVector operands; AffineMap map; - buildTripCountMapAndOperands(forOp, &map, &operands); + getTripCountMapAndOperands(forOp, &map, &operands); if (!map) return 1; diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp --- a/mlir/lib/Transforms/Utils/LoopUtils.cpp +++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp @@ -68,7 +68,7 @@ AffineMap tripCountMap; SmallVector tripCountOperands; - buildTripCountMapAndOperands(forOp, &tripCountMap, &tripCountOperands); + getTripCountMapAndOperands(forOp, &tripCountMap, &tripCountOperands); // Sometimes the trip count cannot be expressed as an affine expression. if (!tripCountMap) {