2018-09-12 10:21:23 -07:00
|
|
|
//===- 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"
|
2018-10-17 18:01:44 -07:00
|
|
|
#include "mlir/Analysis/AffineStructures.h"
|
|
|
|
#include "mlir/Analysis/MLFunctionMatcher.h"
|
|
|
|
#include "mlir/IR/Builders.h"
|
|
|
|
#include "mlir/IR/BuiltinOps.h"
|
2018-09-12 10:21:23 -07:00
|
|
|
#include "mlir/IR/Statements.h"
|
2018-10-17 18:01:44 -07:00
|
|
|
#include "mlir/StandardOps/StandardOps.h"
|
2018-10-30 07:54:23 -07:00
|
|
|
#include "mlir/Support/Functional.h"
|
2018-10-03 10:07:54 -07:00
|
|
|
#include "mlir/Support/MathExtras.h"
|
2018-09-12 10:21:23 -07:00
|
|
|
|
2018-10-03 15:39:12 -07:00
|
|
|
using namespace mlir;
|
2018-09-12 10:21:23 -07:00
|
|
|
|
|
|
|
/// 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.
|
2018-10-08 13:47:18 -07:00
|
|
|
AffineExpr mlir::getTripCountExpr(const ForStmt &forStmt) {
|
2018-09-12 10:21:23 -07:00
|
|
|
// 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 {
|
2018-10-09 16:39:24 -07:00
|
|
|
auto lbMap = forStmt.getLowerBoundMap();
|
|
|
|
auto ubMap = forStmt.getUpperBoundMap();
|
2018-09-12 10:21:23 -07:00
|
|
|
// TODO(bondhugula): handle max/min of multiple expressions.
|
2018-10-09 16:39:24 -07:00
|
|
|
if (lbMap.getNumResults() != 1 || ubMap.getNumResults() != 1)
|
2018-09-12 10:21:23 -07:00
|
|
|
return nullptr;
|
|
|
|
|
|
|
|
// TODO(bondhugula): handle bounds with different operands.
|
|
|
|
// Bounds have different operands, unhandled for now.
|
2018-09-18 10:22:03 -07:00
|
|
|
if (!forStmt.matchingBoundOperandList())
|
2018-09-12 10:21:23 -07:00
|
|
|
return nullptr;
|
|
|
|
|
|
|
|
// ub_expr - lb_expr + 1
|
2018-10-09 16:39:24 -07:00
|
|
|
AffineExpr lbExpr(lbMap.getResult(0));
|
|
|
|
AffineExpr ubExpr(ubMap.getResult(0));
|
2018-10-03 15:39:12 -07:00
|
|
|
auto loopSpanExpr = simplifyAffineExpr(
|
2018-10-09 16:39:24 -07:00
|
|
|
ubExpr - lbExpr + 1, std::max(lbMap.getNumDims(), ubMap.getNumDims()),
|
|
|
|
std::max(lbMap.getNumSymbols(), ubMap.getNumSymbols()));
|
2018-10-08 13:47:18 -07:00
|
|
|
auto cExpr = loopSpanExpr.dyn_cast<AffineConstantExpr>();
|
2018-09-12 10:21:23 -07:00
|
|
|
if (!cExpr)
|
2018-10-08 10:20:25 -07:00
|
|
|
return loopSpanExpr.ceilDiv(step);
|
2018-10-09 10:59:27 -07:00
|
|
|
loopSpan = cExpr.getValue();
|
2018-09-12 10:21:23 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
// 0 iteration loops.
|
2018-09-27 18:03:27 -07:00
|
|
|
if (loopSpan < 0)
|
2018-09-12 10:21:23 -07:00
|
|
|
return 0;
|
|
|
|
|
2018-10-08 10:20:25 -07:00
|
|
|
return getAffineConstantExpr(static_cast<uint64_t>(ceilDiv(loopSpan, step)),
|
|
|
|
context);
|
2018-09-12 10:21:23 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
/// 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) {
|
2018-10-03 15:39:12 -07:00
|
|
|
auto tripCountExpr = getTripCountExpr(forStmt);
|
2018-09-12 10:21:23 -07:00
|
|
|
|
2018-10-08 08:09:50 -07:00
|
|
|
if (!tripCountExpr)
|
|
|
|
return None;
|
|
|
|
|
2018-10-08 13:47:18 -07:00
|
|
|
if (auto constExpr = tripCountExpr.dyn_cast<AffineConstantExpr>())
|
2018-10-09 10:59:27 -07:00
|
|
|
return constExpr.getValue();
|
2018-09-12 10:21:23 -07:00
|
|
|
|
|
|
|
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) {
|
2018-10-03 15:39:12 -07:00
|
|
|
auto tripCountExpr = getTripCountExpr(forStmt);
|
2018-09-12 10:21:23 -07:00
|
|
|
|
|
|
|
if (!tripCountExpr)
|
|
|
|
return 1;
|
|
|
|
|
2018-10-08 13:47:18 -07:00
|
|
|
if (auto constExpr = tripCountExpr.dyn_cast<AffineConstantExpr>()) {
|
2018-10-09 10:59:27 -07:00
|
|
|
uint64_t tripCount = constExpr.getValue();
|
2018-09-12 10:21:23 -07:00
|
|
|
|
|
|
|
// 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.
|
2018-10-09 10:59:27 -07:00
|
|
|
return tripCountExpr.getLargestKnownDivisor();
|
2018-09-12 10:21:23 -07:00
|
|
|
}
|
2018-10-17 18:01:44 -07:00
|
|
|
|
2018-10-30 14:59:22 -07:00
|
|
|
bool mlir::isAccessInvariant(const MLValue &input, MemRefType memRefType,
|
2018-10-30 07:54:23 -07:00
|
|
|
ArrayRef<MLValue *> indices, unsigned dim) {
|
2018-10-30 14:59:22 -07:00
|
|
|
assert(indices.size() == memRefType.getRank());
|
2018-10-17 18:01:44 -07:00
|
|
|
assert(dim < indices.size());
|
2018-10-30 14:59:22 -07:00
|
|
|
auto layoutMap = memRefType.getAffineMaps();
|
|
|
|
assert(memRefType.getAffineMaps().size() <= 1);
|
2018-10-17 18:01:44 -07:00
|
|
|
// TODO(ntv): remove dependency on Builder once we support non-identity
|
|
|
|
// layout map.
|
2018-10-30 14:59:22 -07:00
|
|
|
Builder b(memRefType.getContext());
|
2018-10-17 18:01:44 -07:00
|
|
|
assert(layoutMap.empty() ||
|
|
|
|
layoutMap[0] == b.getMultiDimIdentityMap(indices.size()));
|
2018-10-18 11:14:26 -07:00
|
|
|
(void)layoutMap;
|
2018-10-17 18:01:44 -07:00
|
|
|
|
|
|
|
SmallVector<OperationStmt *, 4> affineApplyOps;
|
|
|
|
getReachableAffineApplyOps({indices[dim]}, affineApplyOps);
|
|
|
|
|
|
|
|
if (affineApplyOps.empty()) {
|
|
|
|
// Pointer equality test because of MLValue pointer semantics.
|
2018-10-30 07:54:23 -07:00
|
|
|
return indices[dim] != &input;
|
2018-10-17 18:01:44 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
assert(affineApplyOps.size() == 1 &&
|
|
|
|
"CompositionAffineMapsPass must have "
|
|
|
|
"been run: there should be at most one AffineApplyOp");
|
2018-10-19 09:07:58 -07:00
|
|
|
auto composeOp = affineApplyOps[0]->cast<AffineApplyOp>();
|
2018-10-30 07:54:23 -07:00
|
|
|
// We need yet another level of indirection because the `dim` index of the
|
|
|
|
// access may not correspond to the `dim` index of composeOp.
|
|
|
|
unsigned idx = std::numeric_limits<unsigned>::max();
|
|
|
|
unsigned numResults = composeOp->getNumResults();
|
|
|
|
for (unsigned i = 0; i < numResults; ++i) {
|
|
|
|
if (indices[dim] == composeOp->getResult(i)) {
|
|
|
|
idx = i;
|
|
|
|
break;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
assert(idx < std::numeric_limits<unsigned>::max());
|
|
|
|
return !AffineValueMap(*composeOp)
|
|
|
|
.isFunctionOf(idx, &const_cast<MLValue &>(input));
|
2018-10-17 18:01:44 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
/// Determines whether a load or a store has a contiguous access along the
|
|
|
|
/// value `input`. Contiguous is defined as either invariant or varying only
|
|
|
|
/// along the fastest varying memory dimension.
|
|
|
|
// TODO(ntv): allow more advanced notions of contiguity (non-fastest varying,
|
|
|
|
// check strides, ...).
|
|
|
|
template <typename LoadOrStoreOpPointer>
|
2018-10-30 07:54:23 -07:00
|
|
|
static bool isContiguousAccess(const MLValue &input,
|
|
|
|
LoadOrStoreOpPointer memoryOp,
|
|
|
|
unsigned fastestVaryingDim) {
|
|
|
|
using namespace functional;
|
|
|
|
auto indices = map([](SSAValue *val) { return dyn_cast<MLValue>(val); },
|
|
|
|
memoryOp->getIndices());
|
2018-10-30 14:59:22 -07:00
|
|
|
auto memRefType = memoryOp->getMemRefType();
|
2018-10-30 07:54:23 -07:00
|
|
|
for (unsigned d = 0, numIndices = indices.size(); d < numIndices; ++d) {
|
|
|
|
if (fastestVaryingDim == (numIndices - 1) - d) {
|
|
|
|
continue;
|
|
|
|
}
|
2018-10-17 18:01:44 -07:00
|
|
|
if (!isAccessInvariant(input, memRefType, indices, d)) {
|
|
|
|
return false;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return true;
|
|
|
|
}
|
|
|
|
|
2018-10-30 07:54:23 -07:00
|
|
|
template <typename LoadOrStoreOpPointer>
|
|
|
|
static bool isVectorElement(LoadOrStoreOpPointer memoryOp) {
|
2018-10-30 14:59:22 -07:00
|
|
|
auto memRefType = memoryOp->getMemRefType();
|
|
|
|
return memRefType.getElementType().template isa<VectorType>();
|
2018-10-30 07:54:23 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
bool mlir::isVectorizableLoop(const ForStmt &loop, unsigned fastestVaryingDim) {
|
|
|
|
if (!matcher::isParallelLoop(loop) && !matcher::isReductionLoop(loop)) {
|
|
|
|
return false;
|
|
|
|
}
|
|
|
|
|
|
|
|
// No vectorization across conditionals for now.
|
|
|
|
auto conditionals = matcher::If();
|
|
|
|
auto *forStmt = const_cast<ForStmt *>(&loop);
|
|
|
|
auto conditionalsMatched = conditionals.match(forStmt);
|
|
|
|
if (!conditionalsMatched.empty()) {
|
|
|
|
return false;
|
|
|
|
}
|
|
|
|
|
|
|
|
auto loadAndStores = matcher::Op(matcher::isLoadOrStore);
|
|
|
|
auto loadAndStoresMatched = loadAndStores.match(forStmt);
|
|
|
|
for (auto ls : loadAndStoresMatched) {
|
2018-10-17 18:01:44 -07:00
|
|
|
auto *op = cast<OperationStmt>(ls.first);
|
2018-10-19 09:07:58 -07:00
|
|
|
auto load = op->dyn_cast<LoadOp>();
|
|
|
|
auto store = op->dyn_cast<StoreOp>();
|
2018-10-30 07:54:23 -07:00
|
|
|
// Only scalar types are considered vectorizable, all load/store must be
|
|
|
|
// vectorizable for a loop to qualify as vectorizable.
|
|
|
|
// TODO(ntv): ponder whether we want to be more general here.
|
|
|
|
bool vector = load ? isVectorElement(load) : isVectorElement(store);
|
|
|
|
if (vector) {
|
|
|
|
return false;
|
|
|
|
}
|
|
|
|
bool contiguous = load ? isContiguousAccess(loop, load, fastestVaryingDim)
|
|
|
|
: isContiguousAccess(loop, store, fastestVaryingDim);
|
2018-10-17 18:01:44 -07:00
|
|
|
if (!contiguous) {
|
|
|
|
return false;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return true;
|
|
|
|
}
|
2018-10-18 11:14:26 -07:00
|
|
|
|
|
|
|
/// Checks whether SSA dominance would be violated if a for stmt's body
|
|
|
|
/// statements are shifted by the specified shifts. This method checks if a
|
|
|
|
/// 'def' and all its uses have the same shift factor.
|
|
|
|
// TODO(mlir-team): extend this to check for memory-based dependence
|
|
|
|
// violation when we have the support.
|
|
|
|
bool mlir::isStmtwiseShiftValid(const ForStmt &forStmt,
|
|
|
|
ArrayRef<uint64_t> shifts) {
|
|
|
|
assert(shifts.size() == forStmt.getStatements().size());
|
|
|
|
unsigned s = 0;
|
|
|
|
for (const auto &stmt : forStmt) {
|
|
|
|
// A for or if stmt does not produce any def/results (that are used
|
|
|
|
// outside).
|
|
|
|
if (const auto *opStmt = dyn_cast<OperationStmt>(&stmt)) {
|
|
|
|
for (unsigned i = 0, e = opStmt->getNumResults(); i < e; ++i) {
|
|
|
|
const MLValue *result = opStmt->getResult(i);
|
|
|
|
for (const StmtOperand &use : result->getUses()) {
|
|
|
|
// If an ancestor statement doesn't lie in the block of forStmt, there
|
|
|
|
// is no shift to check.
|
|
|
|
// This is a naive way. If performance becomes an issue, a map can
|
|
|
|
// be used to store 'shifts' - to look up the shift for a statement in
|
|
|
|
// constant time.
|
|
|
|
if (auto *ancStmt = forStmt.findAncestorStmtInBlock(*use.getOwner()))
|
|
|
|
if (shifts[s] != shifts[forStmt.findStmtPosInBlock(*ancStmt)])
|
|
|
|
return false;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
s++;
|
|
|
|
}
|
|
|
|
return true;
|
|
|
|
}
|