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.getStepAsInt();
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;
98 if (
auto constExpr = dyn_cast<AffineConstantExpr>(resultExpr)) {
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;
127 if (
auto constExpr = dyn_cast<AffineConstantExpr>(resultExpr)) {
128 uint64_t tripCount = constExpr.getValue();
137 thisGcd = resultExpr.getLargestKnownDivisor();
144 assert(
gcd.has_value() &&
"value expected per above logic");
154 assert(isa<IndexType>(index.
getType()) &&
"index must be of 'index' type");
163 template <
typename LoadOrStoreOp>
167 return !llvm::is_contained(avm.
getOperands(), forOp.getInductionVar());
181 for (
auto val : indices) {
190 template <
typename LoadOrStoreOp>
193 static_assert(llvm::is_one_of<LoadOrStoreOp, AffineReadOpInterface,
194 AffineWriteOpInterface>::value,
195 "Must be called on either an affine read or write op");
196 assert(memRefDim &&
"memRefDim == nullptr");
197 auto memRefType = memoryOp.getMemRefType();
199 if (!memRefType.getLayout().isIdentity())
200 return memoryOp.emitError(
"NYI: non-trivial layout map"),
false;
202 int uniqueVaryingIndexAlongIv = -1;
203 auto accessMap = memoryOp.getAffineMap();
205 unsigned numDims = accessMap.getNumDims();
206 for (
unsigned i = 0, e = memRefType.getRank(); i < e; ++i) {
209 auto resultExpr = accessMap.getResult(i);
211 if (
auto dimExpr = dyn_cast<AffineDimExpr>(expr))
212 exprOperands.push_back(mapOperands[dimExpr.getPosition()]);
213 else if (
auto symExpr = dyn_cast<AffineSymbolExpr>(expr))
214 exprOperands.push_back(mapOperands[numDims + symExpr.getPosition()]);
217 for (
Value exprOperand : exprOperands) {
219 if (uniqueVaryingIndexAlongIv != -1) {
223 uniqueVaryingIndexAlongIv = i;
228 if (uniqueVaryingIndexAlongIv == -1)
231 *memRefDim = memRefType.getRank() - (uniqueVaryingIndexAlongIv + 1);
236 AffineReadOpInterface loadOp,
239 AffineWriteOpInterface loadOp,
242 template <
typename LoadOrStoreOp>
244 auto memRefType = memoryOp.getMemRefType();
245 return isa<VectorType>(memRefType.getElementType());
254 auto *forOp = loop.getOperation();
259 conditionals.match(forOp, &conditionalsMatched);
260 if (!conditionalsMatched.empty()) {
268 if (MemRefType t = dyn_cast<MemRefType>(type))
269 return !VectorType::isValidElementType(t.getElementType());
270 return !VectorType::isValidElementType(type);
274 return !VectorType::isValidElementType(type);
278 types.match(forOp, &opsMatched);
279 if (!opsMatched.empty()) {
285 return op.
getNumRegions() != 0 && !isa<AffineIfOp, AffineForOp>(op);
288 regions.match(forOp, ®ionsMatched);
289 if (!regionsMatched.empty()) {
294 vectorTransferMatcher.
match(forOp, &vectorTransfersMatched);
295 if (!vectorTransfersMatched.empty()) {
301 loadAndStores.match(forOp, &loadAndStoresMatched);
302 for (
auto ls : loadAndStoresMatched) {
303 auto *op = ls.getMatchedOperation();
304 auto load = dyn_cast<AffineLoadOp>(op);
305 auto store = dyn_cast<AffineStoreOp>(op);
313 if (isVectorizableOp && !isVectorizableOp(loop, *op)) {
321 AffineForOp loop,
int *memRefDim,
NestedPattern &vectorTransferMatcher) {
324 auto load = dyn_cast<AffineLoadOp>(op);
325 auto store = dyn_cast<AffineStoreOp>(op);
326 int thisOpMemRefDim = -1;
329 cast<AffineReadOpInterface>(*load),
332 cast<AffineWriteOpInterface>(*store),
334 if (thisOpMemRefDim != -1) {
337 if (*memRefDim != -1 && *memRefDim != thisOpMemRefDim)
339 *memRefDim = thisOpMemRefDim;
358 auto *forBody = forOp.getBody();
359 assert(shifts.size() == forBody->getOperations().size());
364 for (
const auto &it :
366 auto &op = it.value();
370 size_t index = shifts.size() - it.index() - 1;
373 uint64_t shift = shifts[index];
374 forBodyShift.try_emplace(&op, shift);
379 for (
auto *user : result.
getUsers()) {
382 if (
auto *ancOp = forBody->findAncestorOpInBlock(*user)) {
383 assert(forBodyShift.count(ancOp) > 0 &&
"ancestor expected in map");
384 if (shift != forBodyShift[ancOp])
static bool isVectorizableLoopBodyWithOpCond(AffineForOp loop, const VectorizableOpFun &isVectorizableOp, NestedPattern &vectorTransferMatcher)
std::function< bool(AffineForOp, Operation &)> VectorizableOpFun
static bool isAccessIndexInvariant(Value iv, Value index)
Given an affine.for iv and an access index of type index, returns true if index is independent of iv ...
static bool isVectorElement(LoadOrStoreOp memoryOp)
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
static Value min(ImplicitLocOpBuilder &builder, Value value, Value bound)
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 getMultiDimIdentityMap(unsigned numDims, MLIRContext *context)
Returns an AffineMap with 'numDims' identity result dim exprs.
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.
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...
MLIRContext * getContext() const
Utility to get the associated MLIRContext that this value is defined in.
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.
void composeSimplifyAndCanonicalize()
Composes all incoming affine.apply ops and then simplifies and canonicalizes the map and operands.
ArrayRef< Value > getOperands() const
AffineExpr getResult(unsigned i)
AffineMap getAffineMap() const
bool isFunctionOf(unsigned idx, Value value) const
Return true if the idx^th result depends on 'value', false otherwise.
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 ...
bool isInvariantAccess(LoadOrStoreOp memOp, AffineForOp forOp)
Checks if an affine read or write operation depends on forOp's IV, i.e., if the memory access is inva...
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 isContiguousAccess(Value iv, LoadOrStoreOp memoryOp, int *memRefDim)
Given:
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)
Include the generated interface declarations.
int64_t ceilDiv(int64_t lhs, int64_t rhs)
Returns the result of MLIR's ceildiv operation on constants.