ab4797229c
- extend loop unroll-jam similar to loop unroll for affine bounds - extend both loop unroll/unroll-jam to deal with cleanup loop for non multiple of unroll factor. - extend promotion of single iteration loops to work with affine bounds - fix typo bugs in loop unroll - refactor common code b/w loop unroll and loop unroll-jam - move prototypes of non-pass transforms to LoopUtils.h - add additional builder methods. - introduce loopUnrollUpTo(factor) to unroll by either factor or trip count, whichever is less. - remove Statement::isInnermost (not used for now - will come back at the right place/in right form later) PiperOrigin-RevId: 213471227
119 lines
4.3 KiB
C++
119 lines
4.3 KiB
C++
//===- LoopAnalysis.cpp - Misc loop analysis routines //-------------------===//
|
|
//
|
|
// Copyright 2019 The MLIR Authors.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
// =============================================================================
|
|
//
|
|
// This file implements miscellaneous loop analysis routines.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Analysis/LoopAnalysis.h"
|
|
|
|
#include "mlir/Analysis/AffineAnalysis.h"
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/IR/AffineMap.h"
|
|
#include "mlir/IR/Statements.h"
|
|
|
|
using mlir::AffineExpr;
|
|
|
|
/// Returns the trip count of the loop as an affine expression if the latter is
|
|
/// expressible as an affine expression, and nullptr otherwise. The trip count
|
|
/// expression is simplified before returning.
|
|
AffineExpr *mlir::getTripCountExpr(const ForStmt &forStmt) {
|
|
// upper_bound - lower_bound + 1
|
|
int64_t loopSpan;
|
|
|
|
int64_t step = forStmt.getStep();
|
|
auto *context = forStmt.getContext();
|
|
|
|
if (forStmt.hasConstantBounds()) {
|
|
int64_t lb = forStmt.getConstantLowerBound();
|
|
int64_t ub = forStmt.getConstantUpperBound();
|
|
loopSpan = ub - lb + 1;
|
|
} else {
|
|
auto *lbMap = forStmt.getLowerBoundMap();
|
|
auto *ubMap = forStmt.getUpperBoundMap();
|
|
// TODO(bondhugula): handle max/min of multiple expressions.
|
|
if (lbMap->getNumResults() != 1 || ubMap->getNumResults() != 1)
|
|
return nullptr;
|
|
|
|
// TODO(bondhugula): handle bounds with different operands.
|
|
// Bounds have different operands, unhandled for now.
|
|
if (!forStmt.matchingBoundOperandList())
|
|
return nullptr;
|
|
|
|
// ub_expr - lb_expr + 1
|
|
auto *lbExpr = lbMap->getResult(0);
|
|
auto *ubExpr = ubMap->getResult(0);
|
|
auto *loopSpanExpr = AffineBinaryOpExpr::getAdd(
|
|
AffineBinaryOpExpr::getSub(ubExpr, lbExpr, context), 1, context);
|
|
|
|
if (auto *expr = simplifyAffineExpr(loopSpanExpr, lbMap->getNumDims(),
|
|
lbMap->getNumSymbols(), context))
|
|
loopSpanExpr = expr;
|
|
|
|
auto *cExpr = dyn_cast<AffineConstantExpr>(loopSpanExpr);
|
|
if (!cExpr)
|
|
return AffineBinaryOpExpr::getCeilDiv(loopSpanExpr, std::abs(step),
|
|
context);
|
|
loopSpan = cExpr->getValue();
|
|
}
|
|
|
|
// 0 iteration loops.
|
|
if ((loopSpan < 0 && step >= 1) || (loopSpan > 0 && step <= -1))
|
|
return 0;
|
|
|
|
return AffineConstantExpr::get(
|
|
static_cast<uint64_t>(loopSpan % step == 0 ? loopSpan / step
|
|
: loopSpan / step + 1),
|
|
context);
|
|
}
|
|
|
|
/// 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.
|
|
llvm::Optional<uint64_t> mlir::getConstantTripCount(const ForStmt &forStmt) {
|
|
AffineExpr *tripCountExpr = getTripCountExpr(forStmt);
|
|
|
|
if (auto *constExpr = dyn_cast_or_null<AffineConstantExpr>(tripCountExpr))
|
|
return constExpr->getValue();
|
|
|
|
return None;
|
|
}
|
|
|
|
/// Returns the greatest known integral divisor of the trip count. Affine
|
|
/// expression analysis is used (indirectly through getTripCount), and
|
|
/// this method is thus able to determine non-trivial divisors.
|
|
uint64_t mlir::getLargestDivisorOfTripCount(const ForStmt &forStmt) {
|
|
AffineExpr *tripCountExpr = getTripCountExpr(forStmt);
|
|
|
|
if (!tripCountExpr)
|
|
return 1;
|
|
|
|
if (auto *constExpr = dyn_cast<AffineConstantExpr>(tripCountExpr)) {
|
|
uint64_t tripCount = constExpr->getValue();
|
|
|
|
// 0 iteration loops (greatest divisor is 2^64 - 1).
|
|
if (tripCount == 0)
|
|
return ULONG_MAX;
|
|
|
|
// The greatest divisor is the trip count.
|
|
return tripCount;
|
|
}
|
|
|
|
// Trip count is not a known constant; return its largest known divisor.
|
|
return tripCountExpr->getLargestKnownDivisor();
|
|
}
|