15 #include "llvm/ADT/SmallBitVector.h"
16 #include "llvm/ADT/SmallSet.h"
17 #include "llvm/ADT/StringRef.h"
18 #include "llvm/Support/raw_ostream.h"
30 class AffineExprConstantFolder {
33 : numDims(numDims), operandConsts(operandConsts) {}
38 if (
auto result = constantFoldImpl(expr))
39 return IntegerAttr::get(IndexType::get(expr.
getContext()), *result);
44 std::optional<int64_t> constantFoldImpl(
AffineExpr expr) {
47 return constantFoldBinExpr(
48 expr, [](int64_t lhs, int64_t rhs) {
return lhs + rhs; });
50 return constantFoldBinExpr(
51 expr, [](int64_t lhs, int64_t rhs) {
return lhs * rhs; });
53 return constantFoldBinExpr(
54 expr, [](int64_t lhs, int64_t rhs) {
return mod(lhs, rhs); });
56 return constantFoldBinExpr(
57 expr, [](int64_t lhs, int64_t rhs) {
return floorDiv(lhs, rhs); });
59 return constantFoldBinExpr(
60 expr, [](int64_t lhs, int64_t rhs) {
return ceilDiv(lhs, rhs); });
65 .dyn_cast_or_null<IntegerAttr>())
69 if (
auto attr = operandConsts[numDims +
71 .dyn_cast_or_null<IntegerAttr>())
75 llvm_unreachable(
"Unknown AffineExpr");
79 std::optional<int64_t> constantFoldBinExpr(
AffineExpr expr,
80 int64_t (*op)(int64_t, int64_t)) {
82 if (
auto lhs = constantFoldImpl(binOpExpr.getLHS()))
83 if (
auto rhs = constantFoldImpl(binOpExpr.getRHS()))
84 return op(*lhs, *rhs);
106 assert(dims >= results &&
"Dimension mismatch");
122 broadcastedDims->clear();
127 unsigned resIdx = idxAndExpr.index();
131 if (constExpr.getValue() != 0)
134 broadcastedDims->push_back(resIdx);
137 if (dimExpr.getPosition() != suffixStart + resIdx)
161 unsigned projectionStart =
163 permutedDims.clear();
169 unsigned leadingBroadcast =
174 unsigned resIdx = idxAndExpr.index();
179 if (constExpr.getValue() != 0)
181 broadcastDims.push_back(resIdx);
183 if (dimExpr.getPosition() < projectionStart)
185 unsigned newPosition =
186 dimExpr.getPosition() - projectionStart + leadingBroadcast;
187 permutedDims[resIdx] = newPosition;
188 dimFound[newPosition] =
true;
197 for (
auto dim : broadcastDims) {
198 while (pos < dimFound.size() && dimFound[pos]) {
201 permutedDims[dim] = pos++;
209 assert(!permutation.empty() &&
210 "Cannot create permutation map from empty permutation vector");
212 for (
auto index : permutation)
214 const auto *m = std::max_element(permutation.begin(), permutation.end());
215 auto permutationMap =
AffineMap::get(*m + 1, 0, affExprs, context);
216 assert(permutationMap.isPermutation() &&
"Invalid permutation vector");
217 return permutationMap;
220 template <
typename AffineExprContainer>
223 assert(!exprsList.empty());
224 assert(!exprsList[0].empty());
225 auto context = exprsList[0][0].getContext();
226 int64_t maxDim = -1, maxSym = -1;
229 maps.reserve(exprsList.size());
230 for (
const auto &exprs : exprsList)
232 maxSym + 1, exprs, context));
249 uint64_t thisGcd = resultExpr.getLargestKnownDivisor();
260 dimExprs.reserve(numDims);
261 for (
unsigned i = 0; i < numDims; ++i)
263 return get(numDims, 0, dimExprs, context);
272 for (
unsigned i = 0, numDims =
getNumDims(); i < numDims; ++i) {
274 if (!expr || expr.getPosition() != i)
284 for (
unsigned i = 0, numSymbols =
getNumSymbols(); i < numSymbols; ++i) {
286 if (!expr || expr.getPosition() != i)
312 assert(
isConstant() &&
"map must have only constant results");
320 assert(map &&
"uninitialized map storage");
324 assert(map &&
"uninitialized map storage");
329 assert(map &&
"uninitialized map storage");
333 assert(map &&
"uninitialized map storage");
348 for (
unsigned i = 0, numResults =
getNumResults(); i < numResults; i++) {
368 if (integers.empty())
371 auto range = llvm::map_range(integers, [
this](int64_t i) {
372 return IntegerAttr::get(IndexType::get(
getContext()), i);
374 results.append(range.begin(), range.end());
384 AffineExprConstantFolder exprFolder(
getNumDims(), operandConstants);
389 auto folded = exprFolder.constantFold(expr);
396 results->push_back(folded.getInt());
398 exprs.push_back(expr);
423 unsigned numResultDims,
424 unsigned numResultSyms)
const {
430 return get(numResultDims, numResultSyms, results,
getContext());
437 unsigned numResultDims,
438 unsigned numResultSyms)
const {
442 newResults.push_back(e.replace(expr, replacement));
450 unsigned numResultDims,
451 unsigned numResultSyms)
const {
455 newResults.push_back(e.replace(map));
464 newResults.push_back(e.replace(map));
473 unsigned numSymbols = numSymbolsThisMap + map.
getNumSymbols();
475 for (
unsigned idx = 0; idx < numDims; ++idx) {
479 for (
unsigned idx = numSymbolsThisMap; idx < numSymbols; ++idx) {
480 newSymbols[idx - numSymbolsThisMap] =
488 exprs.push_back(expr.
compose(newMap));
495 exprs.reserve(values.size());
497 for (
auto v : values)
501 res.reserve(resMap.getNumResults());
502 for (
auto e : resMap.getResults())
523 if (seen[dim.getPosition()])
525 seen[dim.getPosition()] =
true;
528 if (!allowZeroInResults || !constExpr || constExpr.getValue() != 0)
545 exprs.reserve(resultPos.size());
546 for (
auto idx : resultPos)
573 const llvm::SmallBitVector &unusedDims) {
574 unsigned numDims = 0;
578 for (
unsigned dim = 0, e = map.
getNumDims(); dim < e; ++dim) {
579 if (unusedDims.test(dim))
587 resultExprs.push_back(e.replaceDims(dimReplacements));
601 allExprs.reserve(maps.size() * maps.front().getNumResults());
602 unsigned numDims = maps.front().getNumDims(),
603 numSymbols = maps.front().getNumSymbols();
604 for (
auto m : maps) {
605 assert(numDims == m.getNumDims() && numSymbols == m.getNumSymbols() &&
606 "expected maps with same num dims and symbols");
607 llvm::append_range(allExprs, m.getResults());
610 AffineMap::get(numDims, numSymbols, allExprs, maps.front().getContext()));
611 unsigned unifiedNumDims = unifiedMap.
getNumDims(),
615 res.reserve(maps.size());
616 for (
auto m : maps) {
618 unifiedResults.take_front(m.getNumResults()),
620 unifiedResults = unifiedResults.drop_front(m.getNumResults());
631 const llvm::SmallBitVector &unusedSymbols) {
632 unsigned numSymbols = 0;
636 for (
unsigned sym = 0, e = map.
getNumSymbols(); sym < e; ++sym) {
637 if (unusedSymbols.test(sym))
645 resultExprs.push_back(e.replaceSymbols(symReplacements));
650 llvm::SmallBitVector unusedSymbols(map.
getNumSymbols(),
true);
653 unusedSymbols.reset(symExpr.getPosition());
676 uniqueExprs.erase(std::unique(uniqueExprs.begin(), uniqueExprs.end()),
685 assert(map.
getNumSymbols() == 0 &&
"expected map without symbols");
688 auto expr = en.value();
691 if (exprs[d.getPosition()])
698 for (
auto expr : exprs)
700 seenExprs.push_back(expr);
712 for (
unsigned i : llvm::seq(
unsigned(0), map.
getNumResults())) {
715 assert(constExpr.getValue() == 0 &&
716 "Unexpected constant in projected permutation");
728 unsigned numResults = 0, numDims = 0, numSymbols = 0;
732 results.reserve(numResults);
733 for (
auto m : maps) {
734 for (
auto res : m.getResults())
735 results.push_back(res.shiftSymbols(m.getNumSymbols(), numSymbols));
737 numSymbols += m.getNumSymbols();
738 numDims =
std::max(m.getNumDims(), numDims);
741 maps.front().getContext());
745 const llvm::SmallBitVector &unusedDims) {
750 unsigned numDims = maps[0].getNumDims();
751 llvm::SmallBitVector numDimsBitVector(numDims,
true);
752 for (
const auto &m : maps) {
753 for (
unsigned i = 0; i < numDims; ++i) {
754 if (m.isFunctionOfDim(i))
755 numDimsBitVector.reset(i);
758 return numDimsBitVector;
766 : results(map.getResults().begin(), map.getResults().end()),
767 numDims(map.getNumDims()), numSymbols(map.getNumSymbols()),
768 context(map.getContext()) {}
775 llvm::append_range(results, map.
getResults());
static SmallVector< AffineMap, 4 > inferFromExprList(ArrayRef< AffineExprContainer > exprsList)
static SmallVector< AffineMap > compressUnusedImpl(ArrayRef< AffineMap > maps, llvm::function_ref< AffineMap(AffineMap)> compressionFun)
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
Affine binary operation expression.
An integer constant appearing in affine expression.
A dimensional identifier appearing in an affine expression.
unsigned getPosition() const
Base type for affine expression.
AffineExpr replaceDimsAndSymbols(ArrayRef< AffineExpr > dimReplacements, ArrayRef< AffineExpr > symReplacements) const
This method substitutes any uses of dimensions and symbols (e.g.
void walk(std::function< void(AffineExpr)> callback) const
Walk all of the AffineExpr's in this expression in postorder.
AffineExprKind getKind() const
Return the classification for this type.
AffineExpr compose(AffineMap map) const
Compose with an AffineMap.
constexpr bool isa() const
MLIRContext * getContext() const
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
int64_t getSingleConstantResult() const
Returns the constant result of this map.
static AffineMap getMinorIdentityMap(unsigned dims, unsigned results, MLIRContext *context)
Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most minor dimensions.
AffineMap getSliceMap(unsigned start, unsigned length) const
Returns the map consisting of length expressions starting from start.
AffineMap getMajorSubMap(unsigned numResults) const
Returns the map consisting of the most major numResults results.
MLIRContext * getContext() const
AffineMap partialConstantFold(ArrayRef< Attribute > operandConstants, SmallVectorImpl< int64_t > *results=nullptr) const
Propagates the constant operands into this affine map.
bool isMinorIdentity() const
Returns true if this affine map is a minor identity, i.e.
unsigned getDimPosition(unsigned idx) const
Extracts the position of the dimensional expression at the given result, when the caller knows it is ...
bool isConstant() const
Returns true if this affine map has only constant results.
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: () -> ().
bool isSingleConstant() const
Returns true if this affine map is a single result constant function.
bool isProjectedPermutation(bool allowZeroInResults=false) const
Returns true if the AffineMap represents a subset (i.e.
AffineMap getMinorSubMap(unsigned numResults) const
Returns the map consisting of the most minor numResults results.
uint64_t getLargestKnownDivisorOfMapExprs()
Get the largest known divisor of all map expressions.
constexpr AffineMap()=default
bool isEmpty() const
Returns true if this affine map is an empty map, i.e., () -> ().
std::optional< unsigned > getResultPosition(AffineExpr input) const
Extracts the first result position where input dimension resides.
unsigned getNumSymbols() const
bool isMinorIdentityWithBroadcasting(SmallVectorImpl< unsigned > *broadcastedDims=nullptr) const
Returns true if this affine map is a minor identity up to broadcasted dimensions which are indicated ...
unsigned getNumDims() const
ArrayRef< AffineExpr > getResults() const
SmallVector< int64_t > getConstantResults() const
Returns the constant results of this map.
bool isPermutationOfMinorIdentityWithBroadcasting(SmallVectorImpl< unsigned > &permutedDims) const
Return true if this affine map can be converted to a minor identity with broadcast by doing a permute...
LogicalResult constantFold(ArrayRef< Attribute > operandConstants, SmallVectorImpl< Attribute > &results) const
Folds the results of the application of an affine map on the provided operands to a constant if possi...
bool isSymbolIdentity() const
Returns true if this affine map is an identity affine map on the symbol identifiers.
unsigned getNumResults() const
AffineMap replaceDimsAndSymbols(ArrayRef< AffineExpr > dimReplacements, ArrayRef< AffineExpr > symReplacements, unsigned numResultDims, unsigned numResultSyms) const
This method substitutes any uses of dimensions and symbols (e.g.
unsigned getNumInputs() const
static SmallVector< AffineMap, 4 > inferFromExprList(ArrayRef< ArrayRef< AffineExpr >> exprsList)
Returns a vector of AffineMaps; each with as many results as exprs.size(), as many dims as the larges...
AffineExpr getResult(unsigned idx) const
AffineMap replace(AffineExpr expr, AffineExpr replacement, unsigned numResultDims, unsigned numResultSyms) const
Sparse replace method.
static AffineMap getPermutationMap(ArrayRef< unsigned > permutation, MLIRContext *context)
Returns an AffineMap representing a permutation.
void walkExprs(llvm::function_ref< void(AffineExpr)> callback) const
Walk all of the AffineExpr's in this mapping.
static AffineMap getConstantMap(int64_t val, MLIRContext *context)
Returns a single constant result affine map.
AffineMap getSubMap(ArrayRef< unsigned > resultPos) const
Returns the map consisting of the resultPos subset.
AffineMap compose(AffineMap map) const
Returns the AffineMap resulting from composing this with map.
bool isIdentity() const
Returns true if this affine map is an identity affine map.
bool isPermutation() const
Returns true if the AffineMap represents a symbol-less permutation map.
A symbolic identifier appearing in an affine expression.
unsigned getPosition() const
MLIRContext is the top-level object for a collection of MLIR operations.
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.
AffineMap simplifyAffineMap(AffineMap map)
Simplifies an affine map by simplifying its underlying AffineExpr results.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
AffineMap removeDuplicateExprs(AffineMap map)
Returns a map with the same dimension and symbol count as map, but whose results are the unique affin...
AffineMap inverseAndBroadcastProjectedPermutation(AffineMap map)
Return the reverse map of a projected permutation where the projected dimensions are transformed into...
int64_t floorDiv(int64_t lhs, int64_t rhs)
Returns the result of MLIR's floordiv operation on constants.
int64_t ceilDiv(int64_t lhs, int64_t rhs)
Returns the result of MLIR's ceildiv operation on constants.
AffineMap inversePermutation(AffineMap map)
Returns a map of codomain to domain dimensions such that the first codomain dimension for a particula...
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
AffineMap concatAffineMaps(ArrayRef< AffineMap > maps)
Concatenates a list of maps into a single AffineMap, stepping over potentially empty maps.
@ CeilDiv
RHS of ceildiv is always a constant or a symbolic expression.
@ Mul
RHS of mul is always a constant or a symbolic expression.
@ Mod
RHS of mod is always a constant or a symbolic expression with a positive value.
@ DimId
Dimensional identifier.
@ FloorDiv
RHS of floordiv is always a constant or a symbolic expression.
@ Constant
Constant integer.
@ SymbolId
Symbolic identifier.
AffineMap compressSymbols(AffineMap map, const llvm::SmallBitVector &unusedSymbols)
Drop the symbols that are not listed in unusedSymbols.
static void getMaxDimAndSymbol(ArrayRef< AffineExprContainer > exprsList, int64_t &maxDim, int64_t &maxSym)
Calculates maxmimum dimension and symbol positions from the expressions in exprsLists and stores them...
AffineMap compressUnusedDims(AffineMap map)
Drop the dims that are not used.
AffineMap getProjectedMap(AffineMap map, const llvm::SmallBitVector &projectedDimensions)
Returns the map that results from projecting out the dimensions specified in projectedDimensions.
AffineMap compressDims(AffineMap map, const llvm::SmallBitVector &unusedDims)
Drop the dims that are not listed in unusedDims.
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)
llvm::SmallBitVector getUnusedDimsBitVector(ArrayRef< AffineMap > maps)
AffineExpr simplifyAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols)
Simplify an affine expression by flattening and some amount of simple analysis.
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)
These free functions allow clients of the API to not use classes in detail.
AffineMap compressUnusedSymbols(AffineMap map)
Drop the symbols that are not used.
AffineExpr getAffineSymbolExpr(unsigned position, MLIRContext *context)
int64_t mod(int64_t lhs, int64_t rhs)
Returns MLIR's mod operation on constants.
This class represents an efficient way to signal success or failure.
void reset(AffineMap map)
Resets this MutableAffineMap with 'map'.
MutableAffineMap()=default
AffineMap getAffineMap() const
Get the AffineMap corresponding to this MutableAffineMap.
AffineExpr getResult(unsigned idx) const
bool isMultipleOf(unsigned idx, int64_t factor) const
Returns true if the idx'th result expression is a multiple of factor.
unsigned getNumResults() const
void simplify()
Simplify the (result) expressions in this map using analysis (used by.
ArrayRef< AffineExpr > results() const
The affine expressions for this (multi-dimensional) map.