604 lines · cpp
1//===- LoopAnalysis.cpp - Misc loop analysis routines //-------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements miscellaneous loop analysis routines.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"14 15#include "mlir/Analysis/SliceAnalysis.h"16#include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"17#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"18#include "mlir/Dialect/Affine/Analysis/NestedMatcher.h"19#include "mlir/Dialect/Affine/Analysis/Utils.h"20#include "mlir/Dialect/Affine/IR/AffineValueMap.h"21#include "llvm/Support/MathExtras.h"22 23#include "llvm/Support/Debug.h"24#include "llvm/Support/DebugLog.h"25#include <numeric>26#include <optional>27 28#define DEBUG_TYPE "affine-loop-analysis"29 30using namespace mlir;31using namespace mlir::affine;32 33namespace {34 35/// A directed graph to model relationships between MLIR Operations.36class DirectedOpGraph {37public:38 /// Add a node to the graph.39 void addNode(Operation *op) {40 assert(!hasNode(op) && "node already added");41 nodes.emplace_back(op);42 edges[op] = {};43 }44 45 /// Add an edge from `src` to `dest`.46 void addEdge(Operation *src, Operation *dest) {47 // This is a multi-graph.48 assert(hasNode(src) && "src node does not exist in graph");49 assert(hasNode(dest) && "dest node does not exist in graph");50 edges[src].push_back(getNode(dest));51 }52 53 /// Returns true if there is a (directed) cycle in the graph.54 bool hasCycle() { return dfs(/*cycleCheck=*/true); }55 56 void printEdges() {57 for (auto &en : edges) {58 llvm::dbgs() << *en.first << " (" << en.first << ")"59 << " has " << en.second.size() << " edges:\n";60 for (auto *node : en.second) {61 llvm::dbgs() << '\t' << *node->op << '\n';62 }63 }64 }65 66private:67 /// A node of a directed graph between MLIR Operations to model various68 /// relationships. This is meant to be used internally.69 struct DGNode {70 DGNode(Operation *op) : op(op) {};71 Operation *op;72 73 // Start and finish visit numbers are standard in DFS to implement things74 // like finding strongly connected components. These numbers are modified75 // during analyses on the graph and so seemingly const API methods will be76 // non-const.77 78 /// Start visit number.79 int vn = -1;80 81 /// Finish visit number.82 int fn = -1;83 };84 85 /// Get internal node corresponding to `op`.86 DGNode *getNode(Operation *op) {87 auto *value =88 llvm::find_if(nodes, [&](const DGNode &node) { return node.op == op; });89 assert(value != nodes.end() && "node doesn't exist in graph");90 return &*value;91 }92 93 /// Returns true if `key` is in the graph.94 bool hasNode(Operation *key) const {95 return llvm::find_if(nodes, [&](const DGNode &node) {96 return node.op == key;97 }) != nodes.end();98 }99 100 /// Perform a depth-first traversal of the graph setting visited and finished101 /// numbers. If `cycleCheck` is set, detects cycles and returns true as soon102 /// as the first cycle is detected, and false if there are no cycles. If103 /// `cycleCheck` is not set, completes the DFS and the `return` value doesn't104 /// have a meaning.105 bool dfs(bool cycleCheck = false) {106 for (DGNode &node : nodes) {107 node.vn = 0;108 node.fn = -1;109 }110 111 unsigned time = 0;112 for (DGNode &node : nodes) {113 if (node.vn == 0) {114 bool ret = dfsNode(node, cycleCheck, time);115 // Check if a cycle was already found.116 if (cycleCheck && ret)117 return true;118 } else if (cycleCheck && node.fn == -1) {119 // We have encountered a node whose visit has started but it's not120 // finished. So we have a cycle.121 return true;122 }123 }124 return false;125 }126 127 /// Perform depth-first traversal starting at `node`. Return true128 /// as soon as a cycle is found if `cycleCheck` was set. Update `time`.129 bool dfsNode(DGNode &node, bool cycleCheck, unsigned &time) const {130 auto nodeEdges = edges.find(node.op);131 assert(nodeEdges != edges.end() && "missing node in graph");132 node.vn = ++time;133 134 for (auto &neighbour : nodeEdges->second) {135 if (neighbour->vn == 0) {136 bool ret = dfsNode(*neighbour, cycleCheck, time);137 if (cycleCheck && ret)138 return true;139 } else if (cycleCheck && neighbour->fn == -1) {140 // We have encountered a node whose visit has started but it's not141 // finished. So we have a cycle.142 return true;143 }144 }145 146 // Update finish time.147 node.fn = ++time;148 149 return false;150 }151 152 // The list of nodes. The storage is owned by this class.153 SmallVector<DGNode> nodes;154 155 // Edges as an adjacency list.156 DenseMap<Operation *, SmallVector<DGNode *>> edges;157};158 159} // namespace160 161/// Returns the trip count of the loop as an affine expression if the latter is162/// expressible as an affine expression, and nullptr otherwise. The trip count163/// expression is simplified before returning. This method only utilizes map164/// composition to construct lower and upper bounds before computing the trip165/// count expressions.166void mlir::affine::getTripCountMapAndOperands(167 AffineForOp forOp, AffineMap *tripCountMap,168 SmallVectorImpl<Value> *tripCountOperands) {169 MLIRContext *context = forOp.getContext();170 int64_t step = forOp.getStepAsInt();171 int64_t loopSpan;172 if (forOp.hasConstantBounds()) {173 int64_t lb = forOp.getConstantLowerBound();174 int64_t ub = forOp.getConstantUpperBound();175 loopSpan = ub - lb;176 if (loopSpan < 0)177 loopSpan = 0;178 *tripCountMap = AffineMap::getConstantMap(179 llvm::divideCeilSigned(loopSpan, step), context);180 tripCountOperands->clear();181 return;182 }183 auto lbMap = forOp.getLowerBoundMap();184 auto ubMap = forOp.getUpperBoundMap();185 if (lbMap.getNumResults() != 1) {186 *tripCountMap = AffineMap();187 return;188 }189 190 // Difference of each upper bound expression from the single lower bound191 // expression (divided by the step) provides the expressions for the trip192 // count map.193 AffineValueMap ubValueMap(ubMap, forOp.getUpperBoundOperands());194 195 SmallVector<AffineExpr, 4> lbSplatExpr(ubValueMap.getNumResults(),196 lbMap.getResult(0));197 auto lbMapSplat = AffineMap::get(lbMap.getNumDims(), lbMap.getNumSymbols(),198 lbSplatExpr, context);199 AffineValueMap lbSplatValueMap(lbMapSplat, forOp.getLowerBoundOperands());200 201 AffineValueMap tripCountValueMap;202 AffineValueMap::difference(ubValueMap, lbSplatValueMap, &tripCountValueMap);203 for (unsigned i = 0, e = tripCountValueMap.getNumResults(); i < e; ++i)204 tripCountValueMap.setResult(i,205 tripCountValueMap.getResult(i).ceilDiv(step));206 207 *tripCountMap = tripCountValueMap.getAffineMap();208 tripCountOperands->assign(tripCountValueMap.getOperands().begin(),209 tripCountValueMap.getOperands().end());210}211 212/// Returns the trip count of the loop if it's a constant, std::nullopt213/// otherwise. This method uses affine expression analysis (in turn using214/// getTripCount) and is able to determine constant trip count in non-trivial215/// cases.216std::optional<uint64_t> mlir::affine::getConstantTripCount(AffineForOp forOp) {217 SmallVector<Value, 4> operands;218 AffineMap map;219 getTripCountMapAndOperands(forOp, &map, &operands);220 221 if (!map)222 return std::nullopt;223 224 // Take the min if all trip counts are constant.225 std::optional<uint64_t> tripCount;226 for (auto resultExpr : map.getResults()) {227 if (auto constExpr = dyn_cast<AffineConstantExpr>(resultExpr)) {228 if (tripCount.has_value())229 tripCount =230 std::min(*tripCount, static_cast<uint64_t>(constExpr.getValue()));231 else232 tripCount = constExpr.getValue();233 } else {234 return std::nullopt;235 }236 }237 return tripCount;238}239 240/// Returns the greatest known integral divisor of the trip count. Affine241/// expression analysis is used (indirectly through getTripCount), and242/// this method is thus able to determine non-trivial divisors.243uint64_t mlir::affine::getLargestDivisorOfTripCount(AffineForOp forOp) {244 SmallVector<Value, 4> operands;245 AffineMap map;246 getTripCountMapAndOperands(forOp, &map, &operands);247 248 if (!map)249 return 1;250 251 // The largest divisor of the trip count is the GCD of the individual largest252 // divisors.253 assert(map.getNumResults() >= 1 && "expected one or more results");254 std::optional<uint64_t> gcd;255 for (auto resultExpr : map.getResults()) {256 uint64_t thisGcd;257 if (auto constExpr = dyn_cast<AffineConstantExpr>(resultExpr)) {258 uint64_t tripCount = constExpr.getValue();259 // 0 iteration loops (greatest divisor is 2^64 - 1).260 if (tripCount == 0)261 thisGcd = std::numeric_limits<uint64_t>::max();262 else263 // The greatest divisor is the trip count.264 thisGcd = tripCount;265 } else {266 // Trip count is not a known constant; return its largest known divisor.267 thisGcd = resultExpr.getLargestKnownDivisor();268 }269 if (gcd.has_value())270 gcd = std::gcd(*gcd, thisGcd);271 else272 gcd = thisGcd;273 }274 assert(gcd.has_value() && "value expected per above logic");275 return *gcd;276}277 278/// Given an affine.for `iv` and an access `index` of type index, returns `true`279/// if `index` is independent of `iv` and false otherwise.280///281/// Prerequisites: `iv` and `index` of the proper type;282static bool isAccessIndexInvariant(Value iv, Value index) {283 assert(isAffineForInductionVar(iv) && "iv must be an affine.for iv");284 assert(isa<IndexType>(index.getType()) && "index must be of 'index' type");285 auto map = AffineMap::getMultiDimIdentityMap(/*numDims=*/1, iv.getContext());286 SmallVector<Value> operands = {index};287 AffineValueMap avm(map, operands);288 avm.composeSimplifyAndCanonicalize();289 return !avm.isFunctionOf(0, iv);290}291 292// Pre-requisite: Loop bounds should be in canonical form.293template <typename LoadOrStoreOp>294bool mlir::affine::isInvariantAccess(LoadOrStoreOp memOp, AffineForOp forOp) {295 AffineValueMap avm(memOp.getAffineMap(), memOp.getMapOperands());296 avm.composeSimplifyAndCanonicalize();297 return !llvm::is_contained(avm.getOperands(), forOp.getInductionVar());298}299 300// Explicitly instantiate the template so that the compiler knows we need them.301template bool mlir::affine::isInvariantAccess(AffineReadOpInterface,302 AffineForOp);303template bool mlir::affine::isInvariantAccess(AffineWriteOpInterface,304 AffineForOp);305template bool mlir::affine::isInvariantAccess(AffineLoadOp, AffineForOp);306template bool mlir::affine::isInvariantAccess(AffineStoreOp, AffineForOp);307 308DenseSet<Value> mlir::affine::getInvariantAccesses(Value iv,309 ArrayRef<Value> indices) {310 DenseSet<Value> res;311 for (Value index : indices) {312 if (isAccessIndexInvariant(iv, index))313 res.insert(index);314 }315 return res;316}317 318// TODO: check access stride.319template <typename LoadOrStoreOp>320bool mlir::affine::isContiguousAccess(Value iv, LoadOrStoreOp memoryOp,321 int *memRefDim) {322 static_assert(llvm::is_one_of<LoadOrStoreOp, AffineReadOpInterface,323 AffineWriteOpInterface>::value,324 "Must be called on either an affine read or write op");325 assert(memRefDim && "memRefDim == nullptr");326 auto memRefType = memoryOp.getMemRefType();327 328 if (!memRefType.getLayout().isIdentity())329 return memoryOp.emitError("NYI: non-trivial layout map"), false;330 331 int uniqueVaryingIndexAlongIv = -1;332 auto accessMap = memoryOp.getAffineMap();333 SmallVector<Value, 4> mapOperands(memoryOp.getMapOperands());334 unsigned numDims = accessMap.getNumDims();335 for (unsigned i = 0, e = memRefType.getRank(); i < e; ++i) {336 // Gather map operands used in result expr 'i' in 'exprOperands'.337 SmallVector<Value, 4> exprOperands;338 auto resultExpr = accessMap.getResult(i);339 resultExpr.walk([&](AffineExpr expr) {340 if (auto dimExpr = dyn_cast<AffineDimExpr>(expr))341 exprOperands.push_back(mapOperands[dimExpr.getPosition()]);342 else if (auto symExpr = dyn_cast<AffineSymbolExpr>(expr))343 exprOperands.push_back(mapOperands[numDims + symExpr.getPosition()]);344 });345 // Check access invariance of each operand in 'exprOperands'.346 for (Value exprOperand : exprOperands) {347 if (!isAccessIndexInvariant(iv, exprOperand)) {348 if (uniqueVaryingIndexAlongIv != -1) {349 // 2+ varying indices -> do not vectorize along iv.350 return false;351 }352 uniqueVaryingIndexAlongIv = i;353 }354 }355 }356 357 if (uniqueVaryingIndexAlongIv == -1)358 *memRefDim = -1;359 else360 *memRefDim = memRefType.getRank() - (uniqueVaryingIndexAlongIv + 1);361 return true;362}363 364template bool mlir::affine::isContiguousAccess(Value iv,365 AffineReadOpInterface loadOp,366 int *memRefDim);367template bool mlir::affine::isContiguousAccess(Value iv,368 AffineWriteOpInterface loadOp,369 int *memRefDim);370 371template <typename LoadOrStoreOp>372static bool isVectorElement(LoadOrStoreOp memoryOp) {373 auto memRefType = memoryOp.getMemRefType();374 return isa<VectorType>(memRefType.getElementType());375}376 377using VectorizableOpFun = std::function<bool(AffineForOp, Operation &)>;378 379static bool380isVectorizableLoopBodyWithOpCond(AffineForOp loop,381 const VectorizableOpFun &isVectorizableOp,382 NestedPattern &vectorTransferMatcher) {383 auto *forOp = loop.getOperation();384 385 // No vectorization across conditionals for now.386 auto conditionals = matcher::If();387 SmallVector<NestedMatch, 8> conditionalsMatched;388 conditionals.match(forOp, &conditionalsMatched);389 if (!conditionalsMatched.empty()) {390 return false;391 }392 393 // No vectorization for ops with operand or result types that are not394 // vectorizable.395 auto types = matcher::Op([](Operation &op) -> bool {396 if (llvm::any_of(op.getOperandTypes(), [](Type type) {397 if (MemRefType t = dyn_cast<MemRefType>(type))398 return !VectorType::isValidElementType(t.getElementType());399 return !VectorType::isValidElementType(type);400 }))401 return true;402 return !llvm::all_of(op.getResultTypes(), VectorType::isValidElementType);403 });404 SmallVector<NestedMatch, 8> opsMatched;405 types.match(forOp, &opsMatched);406 if (!opsMatched.empty()) {407 return false;408 }409 410 // No vectorization across unknown regions.411 auto regions = matcher::Op([](Operation &op) -> bool {412 return op.getNumRegions() != 0 && !isa<AffineIfOp, AffineForOp>(op);413 });414 SmallVector<NestedMatch, 8> regionsMatched;415 regions.match(forOp, ®ionsMatched);416 if (!regionsMatched.empty()) {417 return false;418 }419 420 SmallVector<NestedMatch, 8> vectorTransfersMatched;421 vectorTransferMatcher.match(forOp, &vectorTransfersMatched);422 if (!vectorTransfersMatched.empty()) {423 return false;424 }425 426 auto loadAndStores = matcher::Op(matcher::isLoadOrStore);427 SmallVector<NestedMatch, 8> loadAndStoresMatched;428 loadAndStores.match(forOp, &loadAndStoresMatched);429 for (auto ls : loadAndStoresMatched) {430 auto *op = ls.getMatchedOperation();431 auto load = dyn_cast<AffineLoadOp>(op);432 auto store = dyn_cast<AffineStoreOp>(op);433 // Only scalar types are considered vectorizable, all load/store must be434 // vectorizable for a loop to qualify as vectorizable.435 // TODO: ponder whether we want to be more general here.436 bool vector = load ? isVectorElement(load) : isVectorElement(store);437 if (vector) {438 return false;439 }440 if (isVectorizableOp && !isVectorizableOp(loop, *op)) {441 return false;442 }443 }444 return true;445}446 447bool mlir::affine::isVectorizableLoopBody(448 AffineForOp loop, int *memRefDim, NestedPattern &vectorTransferMatcher) {449 *memRefDim = -1;450 VectorizableOpFun fun([memRefDim](AffineForOp loop, Operation &op) {451 auto load = dyn_cast<AffineLoadOp>(op);452 auto store = dyn_cast<AffineStoreOp>(op);453 int thisOpMemRefDim = -1;454 bool isContiguous =455 load ? isContiguousAccess(loop.getInductionVar(),456 cast<AffineReadOpInterface>(*load),457 &thisOpMemRefDim)458 : isContiguousAccess(loop.getInductionVar(),459 cast<AffineWriteOpInterface>(*store),460 &thisOpMemRefDim);461 if (thisOpMemRefDim != -1) {462 // If memory accesses vary across different dimensions then the loop is463 // not vectorizable.464 if (*memRefDim != -1 && *memRefDim != thisOpMemRefDim)465 return false;466 *memRefDim = thisOpMemRefDim;467 }468 return isContiguous;469 });470 return isVectorizableLoopBodyWithOpCond(loop, fun, vectorTransferMatcher);471}472 473bool mlir::affine::isVectorizableLoopBody(474 AffineForOp loop, NestedPattern &vectorTransferMatcher) {475 return isVectorizableLoopBodyWithOpCond(loop, nullptr, vectorTransferMatcher);476}477 478/// Checks whether SSA dominance would be violated if a for op's body479/// operations are shifted by the specified shifts. This method checks if a480/// 'def' and all its uses have the same shift factor.481// TODO: extend this to check for memory-based dependence violation when we have482// the support.483bool mlir::affine::isOpwiseShiftValid(AffineForOp forOp,484 ArrayRef<uint64_t> shifts) {485 auto *forBody = forOp.getBody();486 assert(shifts.size() == forBody->getOperations().size());487 488 // Work backwards over the body of the block so that the shift of a use's489 // ancestor operation in the block gets recorded before it's looked up.490 DenseMap<Operation *, uint64_t> forBodyShift;491 for (const auto &it :492 llvm::enumerate(llvm::reverse(forBody->getOperations()))) {493 auto &op = it.value();494 495 // Get the index of the current operation, note that we are iterating in496 // reverse so we need to fix it up.497 size_t index = shifts.size() - it.index() - 1;498 499 // Remember the shift of this operation.500 uint64_t shift = shifts[index];501 forBodyShift.try_emplace(&op, shift);502 503 // Validate the results of this operation if it were to be shifted.504 for (unsigned i = 0, e = op.getNumResults(); i < e; ++i) {505 Value result = op.getResult(i);506 for (auto *user : result.getUsers()) {507 // If an ancestor operation doesn't lie in the block of forOp,508 // there is no shift to check.509 if (auto *ancOp = forBody->findAncestorOpInBlock(*user)) {510 assert(forBodyShift.count(ancOp) > 0 && "ancestor expected in map");511 if (shift != forBodyShift[ancOp])512 return false;513 }514 }515 }516 }517 return true;518}519 520bool mlir::affine::isTilingValid(ArrayRef<AffineForOp> loops) {521 assert(!loops.empty() && "no original loops provided");522 523 // We first find out all dependences we intend to check.524 SmallVector<Operation *, 8> loadAndStoreOps;525 loops[0]->walk([&](Operation *op) {526 if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op))527 loadAndStoreOps.push_back(op);528 });529 530 unsigned numOps = loadAndStoreOps.size();531 unsigned numLoops = loops.size();532 for (unsigned d = 1; d <= numLoops + 1; ++d) {533 for (unsigned i = 0; i < numOps; ++i) {534 Operation *srcOp = loadAndStoreOps[i];535 MemRefAccess srcAccess(srcOp);536 for (unsigned j = 0; j < numOps; ++j) {537 Operation *dstOp = loadAndStoreOps[j];538 MemRefAccess dstAccess(dstOp);539 540 SmallVector<DependenceComponent, 2> depComps;541 DependenceResult result = checkMemrefAccessDependence(542 srcAccess, dstAccess, d, /*dependenceConstraints=*/nullptr,543 &depComps);544 545 // Skip if there is no dependence in this case.546 if (!hasDependence(result))547 continue;548 549 // Check whether there is any negative direction vector in the550 // dependence components found above, which means that dependence is551 // violated by the default hyper-rect tiling method.552 LDBG() << "Checking whether tiling legality violated "553 << "for dependence at depth: " << Twine(d) << " between:"554 << OpWithFlags(srcAccess.opInst, OpPrintingFlags().skipRegions())555 << "\nand:\n"556 << OpWithFlags(dstAccess.opInst,557 OpPrintingFlags().skipRegions());558 for (const DependenceComponent &depComp : depComps) {559 if (depComp.lb.has_value() && depComp.ub.has_value() &&560 *depComp.lb < *depComp.ub && *depComp.ub < 0) {561 LDBG() << "Dependence component lb = " << Twine(*depComp.lb)562 << " ub = " << Twine(*depComp.ub)563 << " is negative at depth: " << Twine(d)564 << " and thus violates the legality rule.";565 return false;566 }567 }568 }569 }570 }571 572 return true;573}574 575bool mlir::affine::hasCyclicDependence(AffineForOp root) {576 // Collect all the memory accesses in the source nest grouped by their577 // immediate parent block.578 DirectedOpGraph graph;579 SmallVector<MemRefAccess> accesses;580 root->walk([&](Operation *op) {581 if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op)) {582 accesses.emplace_back(op);583 graph.addNode(op);584 }585 });586 587 // Construct the dependence graph for all the collected acccesses.588 unsigned rootDepth = getNestingDepth(root);589 for (const auto &accA : accesses) {590 for (const auto &accB : accesses) {591 if (accA.memref != accB.memref)592 continue;593 // Perform the dependence on all surrounding loops + the body.594 unsigned numCommonLoops =595 getNumCommonSurroundingLoops(*accA.opInst, *accB.opInst);596 for (unsigned d = rootDepth + 1; d <= numCommonLoops + 1; ++d) {597 if (!noDependence(checkMemrefAccessDependence(accA, accB, d)))598 graph.addEdge(accA.opInst, accB.opInst);599 }600 }601 }602 return graph.hasCycle();603}604