23 #include "llvm/ADT/DenseSet.h"
24 #include "llvm/ADT/SmallPtrSet.h"
25 #include "llvm/ADT/SmallString.h"
28 #include <type_traits>
39 AffineForOp forOp,
AffineMap *tripCountMap,
42 int64_t step = forOp.getStep();
44 if (forOp.hasConstantBounds()) {
45 int64_t lb = forOp.getConstantLowerBound();
46 int64_t ub = forOp.getConstantUpperBound();
51 tripCountOperands->clear();
54 auto lbMap = forOp.getLowerBoundMap();
55 auto ubMap = forOp.getUpperBoundMap();
56 if (lbMap.getNumResults() != 1) {
68 auto lbMapSplat =
AffineMap::get(lbMap.getNumDims(), lbMap.getNumSymbols(),
69 lbSplatExpr, context);
70 AffineValueMap lbSplatValueMap(lbMapSplat, forOp.getLowerBoundOperands());
74 for (
unsigned i = 0, e = tripCountValueMap.
getNumResults(); i < e; ++i)
79 tripCountOperands->assign(tripCountValueMap.
getOperands().begin(),
96 std::optional<uint64_t> tripCount;
99 if (tripCount.has_value())
101 std::min(*tripCount,
static_cast<uint64_t
>(constExpr.getValue()));
103 tripCount = constExpr.getValue();
123 assert(map.
getNumResults() >= 1 &&
"expected one or more results");
124 std::optional<uint64_t>
gcd;
128 uint64_t tripCount = constExpr.getValue();
137 thisGcd = resultExpr.getLargestKnownDivisor();
144 assert(
gcd.has_value() &&
"value expected per above logic");
165 assert(isa<IndexType>(index.
getType()) &&
"index must be of IndexType");
169 if (affineApplyOps.empty()) {
174 if (affineApplyOps.size() > 1) {
175 affineApplyOps[0]->emitRemark(
176 "CompositionAffineMapsPass must have been run: there should be at most "
177 "one AffineApplyOp, returning false conservatively.");
181 auto composeOp = cast<AffineApplyOp>(affineApplyOps[0]);
184 return !composeOp.getAffineValueMap().isFunctionOf(0, iv);
190 for (
auto val : indices) {
217 template <
typename LoadOrStoreOp>
221 llvm::is_one_of<LoadOrStoreOp, AffineLoadOp, AffineStoreOp>::value,
222 "Must be called on either LoadOp or StoreOp");
223 assert(memRefDim &&
"memRefDim == nullptr");
224 auto memRefType = memoryOp.getMemRefType();
226 if (!memRefType.getLayout().isIdentity())
227 return memoryOp.emitError(
"NYI: non-trivial layoutMap"),
false;
229 int uniqueVaryingIndexAlongIv = -1;
230 auto accessMap = memoryOp.getAffineMap();
232 unsigned numDims = accessMap.getNumDims();
233 for (
unsigned i = 0, e = memRefType.getRank(); i < e; ++i) {
236 auto resultExpr = accessMap.getResult(i);
239 exprOperands.push_back(mapOperands[dimExpr.getPosition()]);
241 exprOperands.push_back(mapOperands[numDims + symExpr.getPosition()]);
244 for (
auto exprOperand : exprOperands) {
246 if (uniqueVaryingIndexAlongIv != -1) {
250 uniqueVaryingIndexAlongIv = i;
255 if (uniqueVaryingIndexAlongIv == -1)
258 *memRefDim = memRefType.getRank() - (uniqueVaryingIndexAlongIv + 1);
262 template <
typename LoadOrStoreOp>
264 auto memRefType = memoryOp.getMemRefType();
265 return isa<VectorType>(memRefType.getElementType());
274 auto *forOp = loop.getOperation();
279 conditionals.match(forOp, &conditionalsMatched);
280 if (!conditionalsMatched.empty()) {
288 if (MemRefType t = dyn_cast<MemRefType>(type))
289 return !VectorType::isValidElementType(t.getElementType());
290 return !VectorType::isValidElementType(type);
294 return !VectorType::isValidElementType(type);
298 types.match(forOp, &opsMatched);
299 if (!opsMatched.empty()) {
305 return op.
getNumRegions() != 0 && !isa<AffineIfOp, AffineForOp>(op);
308 regions.match(forOp, ®ionsMatched);
309 if (!regionsMatched.empty()) {
314 vectorTransferMatcher.
match(forOp, &vectorTransfersMatched);
315 if (!vectorTransfersMatched.empty()) {
321 loadAndStores.match(forOp, &loadAndStoresMatched);
322 for (
auto ls : loadAndStoresMatched) {
323 auto *op = ls.getMatchedOperation();
324 auto load = dyn_cast<AffineLoadOp>(op);
325 auto store = dyn_cast<AffineStoreOp>(op);
333 if (isVectorizableOp && !isVectorizableOp(loop, *op)) {
341 AffineForOp loop,
int *memRefDim,
NestedPattern &vectorTransferMatcher) {
344 auto load = dyn_cast<AffineLoadOp>(op);
345 auto store = dyn_cast<AffineStoreOp>(op);
346 int thisOpMemRefDim = -1;
351 if (thisOpMemRefDim != -1) {
354 if (*memRefDim != -1 && *memRefDim != thisOpMemRefDim)
356 *memRefDim = thisOpMemRefDim;
375 auto *forBody = forOp.getBody();
376 assert(shifts.size() == forBody->getOperations().size());
381 for (
const auto &it :
383 auto &op = it.value();
387 size_t index = shifts.size() - it.index() - 1;
390 uint64_t shift = shifts[index];
391 forBodyShift.try_emplace(&op, shift);
396 for (
auto *user : result.
getUsers()) {
399 if (
auto *ancOp = forBody->findAncestorOpInBlock(*user)) {
400 assert(forBodyShift.count(ancOp) > 0 &&
"ancestor expected in map");
401 if (shift != forBodyShift[ancOp])
static bool isVectorizableLoopBodyWithOpCond(AffineForOp loop, const VectorizableOpFun &isVectorizableOp, NestedPattern &vectorTransferMatcher)
std::function< bool(AffineForOp, Operation &)> VectorizableOpFun
static bool isContiguousAccess(Value iv, LoadOrStoreOp memoryOp, int *memRefDim)
Given:
static bool isAccessIndexInvariant(Value iv, Value index)
Given an induction variable iv of type AffineForOp and an access index of type index,...
static bool isVectorElement(LoadOrStoreOp memoryOp)
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
static Value min(ImplicitLocOpBuilder &builder, Value value, Value bound)
An integer constant appearing in affine expression.
A dimensional identifier appearing in an affine expression.
Base type for affine expression.
AffineExpr ceilDiv(uint64_t v) const
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
ArrayRef< AffineExpr > getResults() const
unsigned getNumResults() const
static AffineMap getConstantMap(int64_t val, MLIRContext *context)
Returns a single constant result affine map.
A symbolic identifier appearing in an affine expression.
MLIRContext is the top-level object for a collection of MLIR operations.
Operation is the basic unit of execution within MLIR.
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
unsigned getNumRegions()
Returns the number of regions held by this operation.
operand_type_range getOperandTypes()
result_type_range getResultTypes()
unsigned getNumResults()
Return the number of results held by this operation.
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Type getType() const
Return the type of this value.
user_range getUsers() const
An AffineValueMap is an affine map plus its ML value operands and results for analysis purposes.
ArrayRef< Value > getOperands() const
AffineExpr getResult(unsigned i)
AffineMap getAffineMap() const
void setResult(unsigned i, AffineExpr e)
unsigned getNumResults() const
static void difference(const AffineValueMap &a, const AffineValueMap &b, AffineValueMap *res)
Return the value map that is the difference of value maps 'a' and 'b', represented as an affine map a...
void match(Operation *op, SmallVectorImpl< NestedMatch > *matches)
Returns all the top-level matches in op.
NestedPattern If(const NestedPattern &child)
bool isLoadOrStore(Operation &op)
NestedPattern Op(FilterFunctionType filter=defaultFilterFunction)
std::optional< uint64_t > getConstantTripCount(AffineForOp forOp)
Returns the trip count of the loop if it's a constant, std::nullopt otherwise.
bool isVectorizableLoopBody(AffineForOp loop, NestedPattern &vectorTransferMatcher)
Checks whether the loop is structurally vectorizable; i.e.
DenseSet< Value, DenseMapInfo< Value > > getInvariantAccesses(Value iv, ArrayRef< Value > indices)
Given an induction variable iv of type AffineForOp and indices of type IndexType, returns the set of ...
void getTripCountMapAndOperands(AffineForOp forOp, AffineMap *map, SmallVectorImpl< Value > *operands)
Returns the trip count of the loop as an affine map with its corresponding operands if the latter is ...
void getReachableAffineApplyOps(ArrayRef< Value > operands, SmallVectorImpl< Operation * > &affineApplyOps)
Returns in affineApplyOps, the sequence of those AffineApplyOp Operations that are reachable via a se...
bool isAffineForInductionVar(Value val)
Returns true if the provided value is the induction variable of an AffineForOp.
uint64_t getLargestDivisorOfTripCount(AffineForOp forOp)
Returns the greatest known integral divisor of the trip count.
bool isOpwiseShiftValid(AffineForOp forOp, ArrayRef< uint64_t > shifts)
Checks where SSA dominance would be violated if a for op's body operations are shifted by the specifi...
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
LLVM_ATTRIBUTE_ALWAYS_INLINE MPInt gcd(const MPInt &a, const MPInt &b)
This header declares functions that assit transformations in the MemRef dialect.
int64_t ceilDiv(int64_t lhs, int64_t rhs)
Returns the result of MLIR's ceildiv operation on constants.