//===- LoopAnalysis.cpp - Misc loop analysis routines //-------------------===// // // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // //===----------------------------------------------------------------------===// // // This file implements miscellaneous loop analysis routines. // //===----------------------------------------------------------------------===// #include "mlir/Analysis/LoopAnalysis.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/NestedMatcher.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineValueMap.h" #include "mlir/Support/MathExtras.h" #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/SmallString.h" #include using namespace mlir; /// 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. This method only utilizes map /// composition to construct lower and upper bounds before computing the trip /// count expressions. void mlir::buildTripCountMapAndOperands( AffineForOp forOp, AffineMap *tripCountMap, SmallVectorImpl *tripCountOperands) { int64_t loopSpan; int64_t step = forOp.getStep(); OpBuilder b(forOp.getOperation()); 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)); tripCountOperands->clear(); return; } auto lbMap = forOp.getLowerBoundMap(); auto ubMap = forOp.getUpperBoundMap(); if (lbMap.getNumResults() != 1) { *tripCountMap = AffineMap(); return; } // Difference of each upper bound expression from the single lower bound // expression (divided by the step) provides the expressions for the trip // count map. AffineValueMap ubValueMap(ubMap, forOp.getUpperBoundOperands()); SmallVector lbSplatExpr(ubValueMap.getNumResults(), lbMap.getResult(0)); auto lbMapSplat = AffineMap::get(lbMap.getNumDims(), lbMap.getNumSymbols(), lbSplatExpr, b.getContext()); AffineValueMap lbSplatValueMap(lbMapSplat, forOp.getLowerBoundOperands()); AffineValueMap tripCountValueMap; AffineValueMap::difference(ubValueMap, lbSplatValueMap, &tripCountValueMap); for (unsigned i = 0, e = tripCountValueMap.getNumResults(); i < e; ++i) tripCountValueMap.setResult(i, tripCountValueMap.getResult(i).ceilDiv(step)); *tripCountMap = tripCountValueMap.getAffineMap(); tripCountOperands->assign(tripCountValueMap.getOperands().begin(), tripCountValueMap.getOperands().end()); } /// 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); if (!map) return None; // Take the min if all trip counts are constant. Optional tripCount; for (auto resultExpr : map.getResults()) { if (auto constExpr = resultExpr.dyn_cast()) { if (tripCount.hasValue()) tripCount = std::min(tripCount.getValue(), static_cast(constExpr.getValue())); else tripCount = constExpr.getValue(); } else return None; } return tripCount; } /// 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(AffineForOp forOp) { SmallVector operands; AffineMap map; buildTripCountMapAndOperands(forOp, &map, &operands); if (!map) return 1; // The largest divisor of the trip count is the GCD of the individual largest // divisors. assert(map.getNumResults() >= 1 && "expected one or more results"); Optional gcd; for (auto resultExpr : map.getResults()) { uint64_t thisGcd; if (auto constExpr = resultExpr.dyn_cast()) { uint64_t tripCount = constExpr.getValue(); // 0 iteration loops (greatest divisor is 2^64 - 1). if (tripCount == 0) thisGcd = std::numeric_limits::max(); else // The greatest divisor is the trip count. thisGcd = tripCount; } else { // Trip count is not a known constant; return its largest known divisor. thisGcd = resultExpr.getLargestKnownDivisor(); } if (gcd.hasValue()) gcd = llvm::GreatestCommonDivisor64(gcd.getValue(), thisGcd); else gcd = thisGcd; } assert(gcd.hasValue() && "value expected per above logic"); return gcd.getValue(); } /// Given an induction variable `iv` of type AffineForOp and an access `index` /// of type index, returns `true` if `index` is independent of `iv` and /// false otherwise. The determination supports composition with at most one /// AffineApplyOp. The 'at most one AffineApplyOp' comes from the fact that /// the composition of AffineApplyOp needs to be canonicalized by construction /// to avoid writing code that composes arbitrary numbers of AffineApplyOps /// everywhere. To achieve this, at the very least, the compose-affine-apply /// pass must have been run. /// /// Prerequisites: /// 1. `iv` and `index` of the proper type; /// 2. at most one reachable AffineApplyOp from index; /// /// Returns false in cases with more than one AffineApplyOp, this is /// conservative. static bool isAccessIndexInvariant(Value iv, Value index) { assert(isForInductionVar(iv) && "iv must be a AffineForOp"); assert(index.getType().isa() && "index must be of IndexType"); SmallVector affineApplyOps; getReachableAffineApplyOps({index}, affineApplyOps); if (affineApplyOps.empty()) { // Pointer equality test because of Value pointer semantics. return index != iv; } if (affineApplyOps.size() > 1) { affineApplyOps[0]->emitRemark( "CompositionAffineMapsPass must have been run: there should be at most " "one AffineApplyOp, returning false conservatively."); return false; } auto composeOp = cast(affineApplyOps[0]); // We need yet another level of indirection because the `dim` index of the // access may not correspond to the `dim` index of composeOp. return !composeOp.getAffineValueMap().isFunctionOf(0, iv); } DenseSet mlir::getInvariantAccesses(Value iv, ArrayRef indices) { DenseSet res; for (unsigned idx = 0, n = indices.size(); idx < n; ++idx) { auto val = indices[idx]; if (isAccessIndexInvariant(iv, val)) { res.insert(val); } } return res; } /// Given: /// 1. an induction variable `iv` of type AffineForOp; /// 2. a `memoryOp` of type const LoadOp& or const StoreOp&; /// determines whether `memoryOp` has a contiguous access along `iv`. Contiguous /// is defined as either invariant or varying only along a unique MemRef dim. /// Upon success, the unique MemRef dim is written in `memRefDim` (or -1 to /// convey the memRef access is invariant along `iv`). /// /// Prerequisites: /// 1. `memRefDim` ~= nullptr; /// 2. `iv` of the proper type; /// 3. the MemRef accessed by `memoryOp` has no layout map or at most an /// identity layout map. /// /// Currently only supports no layoutMap or identity layoutMap in the MemRef. /// Returns false if the MemRef has a non-identity layoutMap or more than 1 /// layoutMap. This is conservative. /// // TODO: check strides. template static bool isContiguousAccess(Value iv, LoadOrStoreOp memoryOp, int *memRefDim) { static_assert( llvm::is_one_of::value, "Must be called on either LoadOp or StoreOp"); assert(memRefDim && "memRefDim == nullptr"); auto memRefType = memoryOp.getMemRefType(); auto layoutMap = memRefType.getAffineMaps(); // TODO: remove dependence on Builder once we support non-identity layout map. Builder b(memoryOp.getContext()); if (layoutMap.size() >= 2 || (layoutMap.size() == 1 && !(layoutMap[0] == b.getMultiDimIdentityMap(layoutMap[0].getNumDims())))) { return memoryOp.emitError("NYI: non-trivial layoutMap"), false; } int uniqueVaryingIndexAlongIv = -1; auto accessMap = memoryOp.getAffineMap(); SmallVector mapOperands(memoryOp.getMapOperands()); unsigned numDims = accessMap.getNumDims(); for (unsigned i = 0, e = memRefType.getRank(); i < e; ++i) { // Gather map operands used result expr 'i' in 'exprOperands'. SmallVector exprOperands; auto resultExpr = accessMap.getResult(i); resultExpr.walk([&](AffineExpr expr) { if (auto dimExpr = expr.dyn_cast()) exprOperands.push_back(mapOperands[dimExpr.getPosition()]); else if (auto symExpr = expr.dyn_cast()) exprOperands.push_back(mapOperands[numDims + symExpr.getPosition()]); }); // Check access invariance of each operand in 'exprOperands'. for (auto exprOperand : exprOperands) { if (!isAccessIndexInvariant(iv, exprOperand)) { if (uniqueVaryingIndexAlongIv != -1) { // 2+ varying indices -> do not vectorize along iv. return false; } uniqueVaryingIndexAlongIv = i; } } } if (uniqueVaryingIndexAlongIv == -1) *memRefDim = -1; else *memRefDim = memRefType.getRank() - (uniqueVaryingIndexAlongIv + 1); return true; } template static bool isVectorElement(LoadOrStoreOp memoryOp) { auto memRefType = memoryOp.getMemRefType(); return memRefType.getElementType().template isa(); } using VectorizableOpFun = std::function; static bool isVectorizableLoopBodyWithOpCond(AffineForOp loop, VectorizableOpFun isVectorizableOp, NestedPattern &vectorTransferMatcher) { auto *forOp = loop.getOperation(); // No vectorization across conditionals for now. auto conditionals = matcher::If(); SmallVector conditionalsMatched; conditionals.match(forOp, &conditionalsMatched); if (!conditionalsMatched.empty()) { return false; } // No vectorization across unknown regions. auto regions = matcher::Op([](Operation &op) -> bool { return op.getNumRegions() != 0 && !isa(op); }); SmallVector regionsMatched; regions.match(forOp, ®ionsMatched); if (!regionsMatched.empty()) { return false; } SmallVector vectorTransfersMatched; vectorTransferMatcher.match(forOp, &vectorTransfersMatched); if (!vectorTransfersMatched.empty()) { return false; } auto loadAndStores = matcher::Op(matcher::isLoadOrStore); SmallVector loadAndStoresMatched; loadAndStores.match(forOp, &loadAndStoresMatched); for (auto ls : loadAndStoresMatched) { auto *op = ls.getMatchedOperation(); auto load = dyn_cast(op); auto store = dyn_cast(op); // Only scalar types are considered vectorizable, all load/store must be // vectorizable for a loop to qualify as vectorizable. // TODO: ponder whether we want to be more general here. bool vector = load ? isVectorElement(load) : isVectorElement(store); if (vector) { return false; } if (isVectorizableOp && !isVectorizableOp(loop, *op)) { return false; } } return true; } bool mlir::isVectorizableLoopBody(AffineForOp loop, int *memRefDim, NestedPattern &vectorTransferMatcher) { VectorizableOpFun fun([memRefDim](AffineForOp loop, Operation &op) { auto load = dyn_cast(op); auto store = dyn_cast(op); return load ? isContiguousAccess(loop.getInductionVar(), load, memRefDim) : isContiguousAccess(loop.getInductionVar(), store, memRefDim); }); return isVectorizableLoopBodyWithOpCond(loop, fun, vectorTransferMatcher); } bool mlir::isVectorizableLoopBody(AffineForOp loop, NestedPattern &vectorTransferMatcher) { return isVectorizableLoopBodyWithOpCond(loop, nullptr, vectorTransferMatcher); } /// Checks whether SSA dominance would be violated if a for op's body /// operations are shifted by the specified shifts. This method checks if a /// 'def' and all its uses have the same shift factor. // TODO: extend this to check for memory-based dependence violation when we have // the support. bool mlir::isOpwiseShiftValid(AffineForOp forOp, ArrayRef shifts) { auto *forBody = forOp.getBody(); assert(shifts.size() == forBody->getOperations().size()); // Work backwards over the body of the block so that the shift of a use's // ancestor operation in the block gets recorded before it's looked up. DenseMap forBodyShift; for (auto it : llvm::enumerate(llvm::reverse(forBody->getOperations()))) { auto &op = it.value(); // Get the index of the current operation, note that we are iterating in // reverse so we need to fix it up. size_t index = shifts.size() - it.index() - 1; // Remember the shift of this operation. uint64_t shift = shifts[index]; forBodyShift.try_emplace(&op, shift); // Validate the results of this operation if it were to be shifted. for (unsigned i = 0, e = op.getNumResults(); i < e; ++i) { Value result = op.getResult(i); for (auto *user : result.getUsers()) { // If an ancestor operation doesn't lie in the block of forOp, // there is no shift to check. if (auto *ancOp = forBody->findAncestorOpInBlock(*user)) { assert(forBodyShift.count(ancOp) > 0 && "ancestor expected in map"); if (shift != forBodyShift[ancOp]) return false; } } } } return true; }