988 lines · cpp
1//===- AffineMap.cpp - MLIR Affine Map Classes ----------------------------===//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#include "mlir/IR/AffineMap.h"10#include "AffineMapDetail.h"11#include "mlir/IR/AffineExpr.h"12#include "mlir/IR/Builders.h"13#include "mlir/IR/BuiltinAttributes.h"14#include "mlir/IR/BuiltinTypes.h"15#include "llvm/ADT/STLExtras.h"16#include "llvm/ADT/SmallBitVector.h"17#include "llvm/ADT/SmallVector.h"18#include "llvm/Support/MathExtras.h"19#include <numeric>20#include <optional>21#include <type_traits>22 23using namespace mlir;24 25using llvm::divideCeilSigned;26using llvm::divideFloorSigned;27using llvm::mod;28 29namespace {30 31// AffineExprConstantFolder evaluates an affine expression using constant32// operands passed in 'operandConsts'. Returns an IntegerAttr attribute33// representing the constant value of the affine expression evaluated on34// constant 'operandConsts', or nullptr if it can't be folded.35class AffineExprConstantFolder {36public:37 AffineExprConstantFolder(unsigned numDims, ArrayRef<Attribute> operandConsts)38 : numDims(numDims), operandConsts(operandConsts) {}39 40 /// Attempt to constant fold the specified affine expr, or return null on41 /// failure.42 IntegerAttr constantFold(AffineExpr expr) {43 if (auto result = constantFoldImpl(expr))44 return IntegerAttr::get(IndexType::get(expr.getContext()), *result);45 return nullptr;46 }47 48 bool hasPoison() const { return hasPoison_; }49 50private:51 std::optional<int64_t> constantFoldImpl(AffineExpr expr) {52 switch (expr.getKind()) {53 case AffineExprKind::Add:54 return constantFoldBinExpr(55 expr, [](int64_t lhs, int64_t rhs) { return lhs + rhs; });56 case AffineExprKind::Mul:57 return constantFoldBinExpr(58 expr, [](int64_t lhs, int64_t rhs) { return lhs * rhs; });59 case AffineExprKind::Mod:60 return constantFoldBinExpr(61 expr, [this](int64_t lhs, int64_t rhs) -> std::optional<int64_t> {62 if (rhs < 1) {63 hasPoison_ = true;64 return std::nullopt;65 }66 return mod(lhs, rhs);67 });68 case AffineExprKind::FloorDiv:69 return constantFoldBinExpr(70 expr, [this](int64_t lhs, int64_t rhs) -> std::optional<int64_t> {71 if (rhs == 0) {72 hasPoison_ = true;73 return std::nullopt;74 }75 return divideFloorSigned(lhs, rhs);76 });77 case AffineExprKind::CeilDiv:78 return constantFoldBinExpr(79 expr, [this](int64_t lhs, int64_t rhs) -> std::optional<int64_t> {80 if (rhs == 0) {81 hasPoison_ = true;82 return std::nullopt;83 }84 return divideCeilSigned(lhs, rhs);85 });86 case AffineExprKind::Constant:87 return cast<AffineConstantExpr>(expr).getValue();88 case AffineExprKind::DimId:89 if (auto attr = llvm::dyn_cast_or_null<IntegerAttr>(90 operandConsts[cast<AffineDimExpr>(expr).getPosition()]))91 return attr.getInt();92 return std::nullopt;93 case AffineExprKind::SymbolId:94 if (auto attr = llvm::dyn_cast_or_null<IntegerAttr>(95 operandConsts[numDims +96 cast<AffineSymbolExpr>(expr).getPosition()]))97 return attr.getInt();98 return std::nullopt;99 }100 llvm_unreachable("Unknown AffineExpr");101 }102 103 // TODO: Change these to operate on APInts too.104 std::optional<int64_t> constantFoldBinExpr(105 AffineExpr expr,106 llvm::function_ref<std::optional<int64_t>(int64_t, int64_t)> op) {107 auto binOpExpr = cast<AffineBinaryOpExpr>(expr);108 if (auto lhs = constantFoldImpl(binOpExpr.getLHS()))109 if (auto rhs = constantFoldImpl(binOpExpr.getRHS()))110 return op(*lhs, *rhs);111 return std::nullopt;112 }113 114 // The number of dimension operands in AffineMap containing this expression.115 unsigned numDims;116 // The constant valued operands used to evaluate this AffineExpr.117 ArrayRef<Attribute> operandConsts;118 bool hasPoison_{false};119};120 121} // namespace122 123/// Returns a single constant result affine map.124AffineMap AffineMap::getConstantMap(int64_t val, MLIRContext *context) {125 return get(/*dimCount=*/0, /*symbolCount=*/0,126 {getAffineConstantExpr(val, context)});127}128 129/// Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most130/// minor dimensions.131AffineMap AffineMap::getMinorIdentityMap(unsigned dims, unsigned results,132 MLIRContext *context) {133 assert(dims >= results && "Dimension mismatch");134 auto id = AffineMap::getMultiDimIdentityMap(dims, context);135 return AffineMap::get(dims, 0, id.getResults().take_back(results), context);136}137 138AffineMap AffineMap::getFilteredIdentityMap(139 MLIRContext *ctx, unsigned numDims,140 llvm::function_ref<bool(AffineDimExpr)> keepDimFilter) {141 auto identityMap = getMultiDimIdentityMap(numDims, ctx);142 143 // Apply filter to results.144 llvm::SmallBitVector dropDimResults(numDims);145 for (auto [idx, resultExpr] : llvm::enumerate(identityMap.getResults()))146 dropDimResults[idx] = !keepDimFilter(cast<AffineDimExpr>(resultExpr));147 148 return identityMap.dropResults(dropDimResults);149}150 151bool AffineMap::isMinorIdentity() const {152 return getNumDims() >= getNumResults() &&153 *this ==154 getMinorIdentityMap(getNumDims(), getNumResults(), getContext());155}156 157SmallVector<unsigned> AffineMap::getBroadcastDims() const {158 SmallVector<unsigned> broadcastedDims;159 for (const auto &[resIdx, expr] : llvm::enumerate(getResults())) {160 if (auto constExpr = dyn_cast<AffineConstantExpr>(expr)) {161 if (constExpr.getValue() != 0)162 continue;163 broadcastedDims.push_back(resIdx);164 }165 }166 167 return broadcastedDims;168}169 170/// Returns true if this affine map is a minor identity up to broadcasted171/// dimensions which are indicated by value 0 in the result.172bool AffineMap::isMinorIdentityWithBroadcasting(173 SmallVectorImpl<unsigned> *broadcastedDims) const {174 if (broadcastedDims)175 broadcastedDims->clear();176 if (getNumDims() < getNumResults())177 return false;178 unsigned suffixStart = getNumDims() - getNumResults();179 for (const auto &idxAndExpr : llvm::enumerate(getResults())) {180 unsigned resIdx = idxAndExpr.index();181 AffineExpr expr = idxAndExpr.value();182 if (auto constExpr = dyn_cast<AffineConstantExpr>(expr)) {183 // Each result may be either a constant 0 (broadcasted dimension).184 if (constExpr.getValue() != 0)185 return false;186 if (broadcastedDims)187 broadcastedDims->push_back(resIdx);188 } else if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {189 // Or it may be the input dimension corresponding to this result position.190 if (dimExpr.getPosition() != suffixStart + resIdx)191 return false;192 } else {193 return false;194 }195 }196 return true;197}198 199/// Return true if this affine map can be converted to a minor identity with200/// broadcast by doing a permute. Return a permutation (there may be201/// several) to apply to get to a minor identity with broadcasts.202/// Ex:203/// * (d0, d1, d2) -> (0, d1) maps to minor identity (d1, 0 = d2) with204/// perm = [1, 0] and broadcast d2205/// * (d0, d1, d2) -> (d0, 0) cannot be mapped to a minor identity by206/// permutation + broadcast207/// * (d0, d1, d2, d3) -> (0, d1, d3) maps to minor identity (d1, 0 = d2, d3)208/// with perm = [1, 0, 2] and broadcast d2209/// * (d0, d1) -> (d1, 0, 0, d0) maps to minor identity (d0, d1) with extra210/// leading broadcat dimensions. The map returned would be (0, 0, d0, d1) with211/// perm = [3, 0, 1, 2]212bool AffineMap::isPermutationOfMinorIdentityWithBroadcasting(213 SmallVectorImpl<unsigned> &permutedDims) const {214 unsigned projectionStart =215 getNumResults() < getNumInputs() ? getNumInputs() - getNumResults() : 0;216 permutedDims.clear();217 SmallVector<unsigned> broadcastDims;218 permutedDims.resize(getNumResults(), 0);219 // If there are more results than input dimensions we want the new map to220 // start with broadcast dimensions in order to be a minor identity with221 // broadcasting.222 unsigned leadingBroadcast =223 getNumResults() > getNumInputs() ? getNumResults() - getNumInputs() : 0;224 llvm::SmallBitVector dimFound(std::max(getNumInputs(), getNumResults()),225 false);226 for (const auto &idxAndExpr : llvm::enumerate(getResults())) {227 unsigned resIdx = idxAndExpr.index();228 AffineExpr expr = idxAndExpr.value();229 // Each result may be either a constant 0 (broadcast dimension) or a230 // dimension.231 if (auto constExpr = dyn_cast<AffineConstantExpr>(expr)) {232 if (constExpr.getValue() != 0)233 return false;234 broadcastDims.push_back(resIdx);235 } else if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {236 if (dimExpr.getPosition() < projectionStart)237 return false;238 unsigned newPosition =239 dimExpr.getPosition() - projectionStart + leadingBroadcast;240 permutedDims[resIdx] = newPosition;241 dimFound[newPosition] = true;242 } else {243 return false;244 }245 }246 // Find a permuation for the broadcast dimension. Since they are broadcasted247 // any valid permutation is acceptable. We just permute the dim into a slot248 // without an existing dimension.249 unsigned pos = 0;250 for (auto dim : broadcastDims) {251 while (pos < dimFound.size() && dimFound[pos]) {252 pos++;253 }254 permutedDims[dim] = pos++;255 }256 return true;257}258 259/// Returns an AffineMap representing a permutation.260AffineMap AffineMap::getPermutationMap(ArrayRef<unsigned> permutation,261 MLIRContext *context) {262 assert(!permutation.empty() &&263 "Cannot create permutation map from empty permutation vector");264 const auto *m = llvm::max_element(permutation);265 auto permutationMap = getMultiDimMapWithTargets(*m + 1, permutation, context);266 assert(permutationMap.isPermutation() && "Invalid permutation vector");267 return permutationMap;268}269AffineMap AffineMap::getPermutationMap(ArrayRef<int64_t> permutation,270 MLIRContext *context) {271 SmallVector<unsigned> perm = llvm::map_to_vector(272 permutation, [](int64_t i) { return static_cast<unsigned>(i); });273 return AffineMap::getPermutationMap(perm, context);274}275 276AffineMap AffineMap::getMultiDimMapWithTargets(unsigned numDims,277 ArrayRef<unsigned> targets,278 MLIRContext *context) {279 SmallVector<AffineExpr, 4> affExprs;280 for (unsigned t : targets)281 affExprs.push_back(getAffineDimExpr(t, context));282 AffineMap result = AffineMap::get(/*dimCount=*/numDims, /*symbolCount=*/0,283 affExprs, context);284 return result;285}286 287/// Creates an affine map each for each list of AffineExpr's in `exprsList`288/// while inferring the right number of dimensional and symbolic inputs needed289/// based on the maximum dimensional and symbolic identifier appearing in the290/// expressions.291template <typename AffineExprContainer>292static SmallVector<AffineMap, 4>293inferFromExprList(ArrayRef<AffineExprContainer> exprsList,294 MLIRContext *context) {295 if (exprsList.empty())296 return {};297 int64_t maxDim = -1, maxSym = -1;298 getMaxDimAndSymbol(exprsList, maxDim, maxSym);299 SmallVector<AffineMap, 4> maps;300 maps.reserve(exprsList.size());301 for (const auto &exprs : exprsList)302 maps.push_back(AffineMap::get(/*dimCount=*/maxDim + 1,303 /*symbolCount=*/maxSym + 1, exprs, context));304 return maps;305}306 307SmallVector<AffineMap, 4>308AffineMap::inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList,309 MLIRContext *context) {310 return ::inferFromExprList(exprsList, context);311}312 313SmallVector<AffineMap, 4>314AffineMap::inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList,315 MLIRContext *context) {316 return ::inferFromExprList(exprsList, context);317}318 319uint64_t AffineMap::getLargestKnownDivisorOfMapExprs() {320 uint64_t gcd = 0;321 for (AffineExpr resultExpr : getResults()) {322 uint64_t thisGcd = resultExpr.getLargestKnownDivisor();323 gcd = std::gcd(gcd, thisGcd);324 }325 if (gcd == 0)326 gcd = std::numeric_limits<uint64_t>::max();327 return gcd;328}329 330AffineMap AffineMap::getMultiDimIdentityMap(unsigned numDims,331 MLIRContext *context) {332 SmallVector<AffineExpr, 4> dimExprs;333 dimExprs.reserve(numDims);334 for (unsigned i = 0; i < numDims; ++i)335 dimExprs.push_back(mlir::getAffineDimExpr(i, context));336 return get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs, context);337}338 339MLIRContext *AffineMap::getContext() const { return map->context; }340 341bool AffineMap::isIdentity() const {342 if (getNumDims() != getNumResults())343 return false;344 ArrayRef<AffineExpr> results = getResults();345 for (unsigned i = 0, numDims = getNumDims(); i < numDims; ++i) {346 auto expr = dyn_cast<AffineDimExpr>(results[i]);347 if (!expr || expr.getPosition() != i)348 return false;349 }350 return true;351}352 353bool AffineMap::isSymbolIdentity() const {354 if (getNumSymbols() != getNumResults())355 return false;356 ArrayRef<AffineExpr> results = getResults();357 for (unsigned i = 0, numSymbols = getNumSymbols(); i < numSymbols; ++i) {358 auto expr = dyn_cast<AffineDimExpr>(results[i]);359 if (!expr || expr.getPosition() != i)360 return false;361 }362 return true;363}364 365bool AffineMap::isEmpty() const {366 return getNumDims() == 0 && getNumSymbols() == 0 && getNumResults() == 0;367}368 369bool AffineMap::isSingleConstant() const {370 return getNumResults() == 1 && isa<AffineConstantExpr>(getResult(0));371}372 373bool AffineMap::isConstant() const {374 return llvm::all_of(getResults(), llvm::IsaPred<AffineConstantExpr>);375}376 377int64_t AffineMap::getSingleConstantResult() const {378 assert(isSingleConstant() && "map must have a single constant result");379 return cast<AffineConstantExpr>(getResult(0)).getValue();380}381 382SmallVector<int64_t> AffineMap::getConstantResults() const {383 assert(isConstant() && "map must have only constant results");384 SmallVector<int64_t> result;385 for (auto expr : getResults())386 result.emplace_back(cast<AffineConstantExpr>(expr).getValue());387 return result;388}389 390unsigned AffineMap::getNumDims() const {391 assert(map && "uninitialized map storage");392 return map->numDims;393}394unsigned AffineMap::getNumSymbols() const {395 assert(map && "uninitialized map storage");396 return map->numSymbols;397}398unsigned AffineMap::getNumResults() const { return getResults().size(); }399unsigned AffineMap::getNumInputs() const {400 assert(map && "uninitialized map storage");401 return map->numDims + map->numSymbols;402}403ArrayRef<AffineExpr> AffineMap::getResults() const {404 assert(map && "uninitialized map storage");405 return map->results();406}407AffineExpr AffineMap::getResult(unsigned idx) const {408 return getResults()[idx];409}410 411unsigned AffineMap::getDimPosition(unsigned idx) const {412 return cast<AffineDimExpr>(getResult(idx)).getPosition();413}414 415std::optional<unsigned> AffineMap::getResultPosition(AffineExpr input) const {416 if (!isa<AffineDimExpr>(input))417 return std::nullopt;418 419 for (unsigned i = 0, numResults = getNumResults(); i < numResults; i++) {420 if (getResult(i) == input)421 return i;422 }423 424 return std::nullopt;425}426 427/// Folds the results of the application of an affine map on the provided428/// operands to a constant if possible. Returns false if the folding happens,429/// true otherwise.430LogicalResult AffineMap::constantFold(ArrayRef<Attribute> operandConstants,431 SmallVectorImpl<Attribute> &results,432 bool *hasPoison) const {433 // Attempt partial folding.434 SmallVector<int64_t, 2> integers;435 partialConstantFold(operandConstants, &integers, hasPoison);436 437 // If all expressions folded to a constant, populate results with attributes438 // containing those constants.439 if (integers.empty())440 return failure();441 442 auto range = llvm::map_range(integers, [this](int64_t i) {443 return IntegerAttr::get(IndexType::get(getContext()), i);444 });445 results.append(range.begin(), range.end());446 return success();447}448 449AffineMap AffineMap::partialConstantFold(ArrayRef<Attribute> operandConstants,450 SmallVectorImpl<int64_t> *results,451 bool *hasPoison) const {452 assert(getNumInputs() == operandConstants.size());453 454 // Fold each of the result expressions.455 AffineExprConstantFolder exprFolder(getNumDims(), operandConstants);456 SmallVector<AffineExpr, 4> exprs;457 exprs.reserve(getNumResults());458 459 for (auto expr : getResults()) {460 auto folded = exprFolder.constantFold(expr);461 if (exprFolder.hasPoison() && hasPoison) {462 *hasPoison = true;463 return {};464 }465 // If did not fold to a constant, keep the original expression, and clear466 // the integer results vector.467 if (folded) {468 exprs.push_back(469 getAffineConstantExpr(folded.getInt(), folded.getContext()));470 if (results)471 results->push_back(folded.getInt());472 } else {473 exprs.push_back(expr);474 if (results) {475 results->clear();476 results = nullptr;477 }478 }479 }480 481 return get(getNumDims(), getNumSymbols(), exprs, getContext());482}483 484/// Walk all of the AffineExpr's in this mapping. Each node in an expression485/// tree is visited in postorder.486void AffineMap::walkExprs(llvm::function_ref<void(AffineExpr)> callback) const {487 for (auto expr : getResults())488 expr.walk(callback);489}490 491/// This method substitutes any uses of dimensions and symbols (e.g.492/// dim#0 with dimReplacements[0]) in subexpressions and returns the modified493/// expression mapping. Because this can be used to eliminate dims and494/// symbols, the client needs to specify the number of dims and symbols in495/// the result. The returned map always has the same number of results.496AffineMap AffineMap::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,497 ArrayRef<AffineExpr> symReplacements,498 unsigned numResultDims,499 unsigned numResultSyms) const {500 SmallVector<AffineExpr, 8> results;501 results.reserve(getNumResults());502 for (auto expr : getResults())503 results.push_back(504 expr.replaceDimsAndSymbols(dimReplacements, symReplacements));505 return get(numResultDims, numResultSyms, results, getContext());506}507 508/// Sparse replace method. Apply AffineExpr::replace(`expr`, `replacement`) to509/// each of the results and return a new AffineMap with the new results and510/// with the specified number of dims and symbols.511AffineMap AffineMap::replace(AffineExpr expr, AffineExpr replacement,512 unsigned numResultDims,513 unsigned numResultSyms) const {514 SmallVector<AffineExpr, 4> newResults;515 newResults.reserve(getNumResults());516 for (AffineExpr e : getResults())517 newResults.push_back(e.replace(expr, replacement));518 return AffineMap::get(numResultDims, numResultSyms, newResults, getContext());519}520 521/// Sparse replace method. Apply AffineExpr::replace(`map`) to each of the522/// results and return a new AffineMap with the new results and with the523/// specified number of dims and symbols.524AffineMap AffineMap::replace(const DenseMap<AffineExpr, AffineExpr> &map,525 unsigned numResultDims,526 unsigned numResultSyms) const {527 SmallVector<AffineExpr, 4> newResults;528 newResults.reserve(getNumResults());529 for (AffineExpr e : getResults())530 newResults.push_back(e.replace(map));531 return AffineMap::get(numResultDims, numResultSyms, newResults, getContext());532}533 534AffineMap535AffineMap::replace(const DenseMap<AffineExpr, AffineExpr> &map) const {536 SmallVector<AffineExpr, 4> newResults;537 newResults.reserve(getNumResults());538 for (AffineExpr e : getResults())539 newResults.push_back(e.replace(map));540 return AffineMap::inferFromExprList(newResults, getContext()).front();541}542 543AffineMap AffineMap::dropResults(const llvm::SmallBitVector &positions) const {544 auto exprs = llvm::to_vector<4>(getResults());545 // TODO: this is a pretty terrible API .. is there anything better?546 for (auto pos = positions.find_last(); pos != -1;547 pos = positions.find_prev(pos))548 exprs.erase(exprs.begin() + pos);549 return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());550}551 552AffineMap AffineMap::compose(AffineMap map) const {553 assert(getNumDims() == map.getNumResults() && "Number of results mismatch");554 // Prepare `map` by concatenating the symbols and rewriting its exprs.555 unsigned numDims = map.getNumDims();556 unsigned numSymbolsThisMap = getNumSymbols();557 unsigned numSymbols = numSymbolsThisMap + map.getNumSymbols();558 SmallVector<AffineExpr, 8> newDims(numDims);559 for (unsigned idx = 0; idx < numDims; ++idx) {560 newDims[idx] = getAffineDimExpr(idx, getContext());561 }562 SmallVector<AffineExpr, 8> newSymbols(numSymbols - numSymbolsThisMap);563 for (unsigned idx = numSymbolsThisMap; idx < numSymbols; ++idx) {564 newSymbols[idx - numSymbolsThisMap] =565 getAffineSymbolExpr(idx, getContext());566 }567 auto newMap =568 map.replaceDimsAndSymbols(newDims, newSymbols, numDims, numSymbols);569 SmallVector<AffineExpr, 8> exprs;570 exprs.reserve(getResults().size());571 for (auto expr : getResults())572 exprs.push_back(expr.compose(newMap));573 return AffineMap::get(numDims, numSymbols, exprs, map.getContext());574}575 576SmallVector<int64_t, 4> AffineMap::compose(ArrayRef<int64_t> values) const {577 assert(getNumSymbols() == 0 && "Expected symbol-less map");578 SmallVector<AffineExpr, 4> exprs;579 MLIRContext *ctx = getContext();580 for (int64_t value : values)581 exprs.push_back(getAffineConstantExpr(value, ctx));582 SmallVector<int64_t, 4> res;583 res.reserve(getNumResults());584 for (auto e : getResults())585 res.push_back(cast<AffineConstantExpr>(e.replaceDims(exprs)).getValue());586 return res;587}588 589size_t AffineMap::getNumOfZeroResults() const {590 size_t res = 0;591 for (auto expr : getResults()) {592 auto constExpr = dyn_cast<AffineConstantExpr>(expr);593 if (constExpr && constExpr.getValue() == 0)594 res++;595 }596 597 return res;598}599 600AffineMap AffineMap::dropZeroResults() {601 SmallVector<AffineExpr> newExprs;602 603 for (auto expr : getResults()) {604 auto constExpr = dyn_cast<AffineConstantExpr>(expr);605 if (!constExpr || constExpr.getValue() != 0)606 newExprs.push_back(expr);607 }608 return AffineMap::get(getNumDims(), getNumSymbols(), newExprs, getContext());609}610 611bool AffineMap::isProjectedPermutation(bool allowZeroInResults) const {612 if (getNumSymbols() > 0)613 return false;614 615 // Having more results than inputs means that results have duplicated dims or616 // zeros that can't be mapped to input dims.617 if (getNumResults() > getNumInputs())618 return false;619 620 SmallVector<bool, 8> seen(getNumInputs(), false);621 // A projected permutation can have, at most, only one instance of each input622 // dimension in the result expressions. Zeros are allowed as long as the623 // number of result expressions is lower or equal than the number of input624 // expressions.625 for (auto expr : getResults()) {626 if (auto dim = dyn_cast<AffineDimExpr>(expr)) {627 if (seen[dim.getPosition()])628 return false;629 seen[dim.getPosition()] = true;630 } else {631 auto constExpr = dyn_cast<AffineConstantExpr>(expr);632 if (!allowZeroInResults || !constExpr || constExpr.getValue() != 0)633 return false;634 }635 }636 637 // Results are either dims or zeros and zeros can be mapped to input dims.638 return true;639}640 641bool AffineMap::isPermutation() const {642 if (getNumDims() != getNumResults())643 return false;644 return isProjectedPermutation();645}646 647AffineMap AffineMap::getSubMap(ArrayRef<unsigned> resultPos) const {648 SmallVector<AffineExpr, 4> exprs;649 exprs.reserve(resultPos.size());650 for (auto idx : resultPos)651 exprs.push_back(getResult(idx));652 return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());653}654 655AffineMap AffineMap::getSliceMap(unsigned start, unsigned length) const {656 return AffineMap::get(getNumDims(), getNumSymbols(),657 getResults().slice(start, length), getContext());658}659 660AffineMap AffineMap::getMajorSubMap(unsigned numResults) const {661 if (numResults == 0)662 return AffineMap();663 if (numResults > getNumResults())664 return *this;665 return getSliceMap(0, numResults);666}667 668AffineMap AffineMap::getMinorSubMap(unsigned numResults) const {669 if (numResults == 0)670 return AffineMap();671 if (numResults > getNumResults())672 return *this;673 return getSliceMap(getNumResults() - numResults, numResults);674}675 676/// Implementation detail to compress multiple affine maps with a compressionFun677/// that is expected to be either compressUnusedDims or compressUnusedSymbols.678/// The implementation keeps track of num dims and symbols across the different679/// affine maps.680static SmallVector<AffineMap> compressUnusedListImpl(681 ArrayRef<AffineMap> maps,682 llvm::function_ref<AffineMap(AffineMap)> compressionFun) {683 if (maps.empty())684 return SmallVector<AffineMap>();685 SmallVector<AffineExpr> allExprs;686 allExprs.reserve(maps.size() * maps.front().getNumResults());687 unsigned numDims = maps.front().getNumDims(),688 numSymbols = maps.front().getNumSymbols();689 for (auto m : maps) {690 assert(numDims == m.getNumDims() && numSymbols == m.getNumSymbols() &&691 "expected maps with same num dims and symbols");692 llvm::append_range(allExprs, m.getResults());693 }694 AffineMap unifiedMap = compressionFun(695 AffineMap::get(numDims, numSymbols, allExprs, maps.front().getContext()));696 unsigned unifiedNumDims = unifiedMap.getNumDims(),697 unifiedNumSymbols = unifiedMap.getNumSymbols();698 ArrayRef<AffineExpr> unifiedResults = unifiedMap.getResults();699 SmallVector<AffineMap> res;700 res.reserve(maps.size());701 for (auto m : maps) {702 res.push_back(AffineMap::get(unifiedNumDims, unifiedNumSymbols,703 unifiedResults.take_front(m.getNumResults()),704 m.getContext()));705 unifiedResults = unifiedResults.drop_front(m.getNumResults());706 }707 return res;708}709 710AffineMap mlir::compressDims(AffineMap map,711 const llvm::SmallBitVector &unusedDims) {712 return projectDims(map, unusedDims, /*compressDimsFlag=*/true);713}714 715AffineMap mlir::compressUnusedDims(AffineMap map) {716 return compressDims(map, getUnusedDimsBitVector({map}));717}718 719SmallVector<AffineMap> mlir::compressUnusedDims(ArrayRef<AffineMap> maps) {720 return compressUnusedListImpl(721 maps, [](AffineMap m) { return compressUnusedDims(m); });722}723 724AffineMap mlir::compressSymbols(AffineMap map,725 const llvm::SmallBitVector &unusedSymbols) {726 return projectSymbols(map, unusedSymbols, /*compressSymbolsFlag=*/true);727}728 729AffineMap mlir::compressUnusedSymbols(AffineMap map) {730 return compressSymbols(map, getUnusedSymbolsBitVector({map}));731}732 733SmallVector<AffineMap> mlir::compressUnusedSymbols(ArrayRef<AffineMap> maps) {734 return compressUnusedListImpl(735 maps, [](AffineMap m) { return compressUnusedSymbols(m); });736}737 738AffineMap mlir::foldAttributesIntoMap(Builder &b, AffineMap map,739 ArrayRef<OpFoldResult> operands,740 SmallVector<Value> &remainingValues) {741 SmallVector<AffineExpr> dimReplacements, symReplacements;742 int64_t numDims = 0;743 for (int64_t i = 0; i < map.getNumDims(); ++i) {744 if (auto attr = dyn_cast<Attribute>(operands[i])) {745 dimReplacements.push_back(746 b.getAffineConstantExpr(cast<IntegerAttr>(attr).getInt()));747 } else {748 dimReplacements.push_back(b.getAffineDimExpr(numDims++));749 remainingValues.push_back(cast<Value>(operands[i]));750 }751 }752 int64_t numSymbols = 0;753 for (int64_t i = 0; i < map.getNumSymbols(); ++i) {754 if (auto attr = dyn_cast<Attribute>(operands[i + map.getNumDims()])) {755 symReplacements.push_back(756 b.getAffineConstantExpr(cast<IntegerAttr>(attr).getInt()));757 } else {758 symReplacements.push_back(b.getAffineSymbolExpr(numSymbols++));759 remainingValues.push_back(cast<Value>(operands[i + map.getNumDims()]));760 }761 }762 return map.replaceDimsAndSymbols(dimReplacements, symReplacements, numDims,763 numSymbols);764}765 766AffineMap mlir::simplifyAffineMap(AffineMap map) {767 SmallVector<AffineExpr, 8> exprs;768 for (auto e : map.getResults()) {769 exprs.push_back(770 simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols()));771 }772 return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs,773 map.getContext());774}775 776AffineMap mlir::removeDuplicateExprs(AffineMap map) {777 auto results = map.getResults();778 SmallVector<AffineExpr, 4> uniqueExprs(results);779 uniqueExprs.erase(llvm::unique(uniqueExprs), uniqueExprs.end());780 return AffineMap::get(map.getNumDims(), map.getNumSymbols(), uniqueExprs,781 map.getContext());782}783 784AffineMap mlir::inversePermutation(AffineMap map) {785 if (map.isEmpty())786 return map;787 assert(map.getNumSymbols() == 0 && "expected map without symbols");788 SmallVector<AffineExpr, 4> exprs(map.getNumDims());789 for (const auto &en : llvm::enumerate(map.getResults())) {790 auto expr = en.value();791 // Skip non-permutations.792 if (auto d = dyn_cast<AffineDimExpr>(expr)) {793 if (exprs[d.getPosition()])794 continue;795 exprs[d.getPosition()] = getAffineDimExpr(en.index(), d.getContext());796 }797 }798 SmallVector<AffineExpr, 4> seenExprs;799 seenExprs.reserve(map.getNumDims());800 for (auto expr : exprs)801 if (expr)802 seenExprs.push_back(expr);803 if (seenExprs.size() != map.getNumInputs())804 return AffineMap();805 return AffineMap::get(map.getNumResults(), 0, seenExprs, map.getContext());806}807 808AffineMap mlir::inverseAndBroadcastProjectedPermutation(AffineMap map) {809 assert(map.isProjectedPermutation(/*allowZeroInResults=*/true));810 MLIRContext *context = map.getContext();811 AffineExpr zero = mlir::getAffineConstantExpr(0, context);812 // Start with all the results as 0.813 SmallVector<AffineExpr, 4> exprs(map.getNumInputs(), zero);814 for (unsigned i : llvm::seq(unsigned(0), map.getNumResults())) {815 // Skip zeros from input map. 'exprs' is already initialized to zero.816 if (auto constExpr = dyn_cast<AffineConstantExpr>(map.getResult(i))) {817 assert(constExpr.getValue() == 0 &&818 "Unexpected constant in projected permutation");819 (void)constExpr;820 continue;821 }822 823 // Reverse each dimension existing in the original map result.824 exprs[map.getDimPosition(i)] = getAffineDimExpr(i, context);825 }826 return AffineMap::get(map.getNumResults(), /*symbolCount=*/0, exprs, context);827}828 829AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps,830 MLIRContext *context) {831 if (maps.empty())832 return AffineMap::get(context);833 unsigned numResults = 0, numDims = 0, numSymbols = 0;834 for (auto m : maps)835 numResults += m.getNumResults();836 SmallVector<AffineExpr, 8> results;837 results.reserve(numResults);838 for (auto m : maps) {839 for (auto res : m.getResults())840 results.push_back(res.shiftSymbols(m.getNumSymbols(), numSymbols));841 842 numSymbols += m.getNumSymbols();843 numDims = std::max(m.getNumDims(), numDims);844 }845 return AffineMap::get(numDims, numSymbols, results, context);846}847 848/// Common implementation to project out dimensions or symbols from an affine849/// map based on the template type.850/// Additionally, if 'compress' is true, the projected out dimensions or symbols851/// are also dropped from the resulting map.852template <typename AffineDimOrSymExpr>853static AffineMap projectCommonImpl(AffineMap map,854 const llvm::SmallBitVector &toProject,855 bool compress) {856 static_assert(llvm::is_one_of<AffineDimOrSymExpr, AffineDimExpr,857 AffineSymbolExpr>::value,858 "expected AffineDimExpr or AffineSymbolExpr");859 860 constexpr bool isDim = std::is_same<AffineDimOrSymExpr, AffineDimExpr>::value;861 int64_t numDimOrSym = (isDim) ? map.getNumDims() : map.getNumSymbols();862 SmallVector<AffineExpr> replacements;863 replacements.reserve(numDimOrSym);864 865 auto createNewDimOrSym = (isDim) ? getAffineDimExpr : getAffineSymbolExpr;866 867 using replace_fn_ty =868 std::function<AffineExpr(AffineExpr, ArrayRef<AffineExpr>)>;869 replace_fn_ty replaceDims = [](AffineExpr e,870 ArrayRef<AffineExpr> replacements) {871 return e.replaceDims(replacements);872 };873 replace_fn_ty replaceSymbols = [](AffineExpr e,874 ArrayRef<AffineExpr> replacements) {875 return e.replaceSymbols(replacements);876 };877 replace_fn_ty replaceNewDimOrSym = (isDim) ? replaceDims : replaceSymbols;878 879 MLIRContext *context = map.getContext();880 int64_t newNumDimOrSym = 0;881 for (unsigned dimOrSym = 0; dimOrSym < numDimOrSym; ++dimOrSym) {882 if (toProject.test(dimOrSym)) {883 replacements.push_back(getAffineConstantExpr(0, context));884 continue;885 }886 int64_t newPos = compress ? newNumDimOrSym++ : dimOrSym;887 replacements.push_back(createNewDimOrSym(newPos, context));888 }889 SmallVector<AffineExpr> resultExprs;890 resultExprs.reserve(map.getNumResults());891 for (auto e : map.getResults())892 resultExprs.push_back(replaceNewDimOrSym(e, replacements));893 894 int64_t numDims = (compress && isDim) ? newNumDimOrSym : map.getNumDims();895 int64_t numSyms = (compress && !isDim) ? newNumDimOrSym : map.getNumSymbols();896 return AffineMap::get(numDims, numSyms, resultExprs, context);897}898 899AffineMap mlir::projectDims(AffineMap map,900 const llvm::SmallBitVector &projectedDimensions,901 bool compressDimsFlag) {902 return projectCommonImpl<AffineDimExpr>(map, projectedDimensions,903 compressDimsFlag);904}905 906AffineMap mlir::projectSymbols(AffineMap map,907 const llvm::SmallBitVector &projectedSymbols,908 bool compressSymbolsFlag) {909 return projectCommonImpl<AffineSymbolExpr>(map, projectedSymbols,910 compressSymbolsFlag);911}912 913AffineMap mlir::getProjectedMap(AffineMap map,914 const llvm::SmallBitVector &projectedDimensions,915 bool compressDimsFlag,916 bool compressSymbolsFlag) {917 map = projectDims(map, projectedDimensions, compressDimsFlag);918 if (compressSymbolsFlag)919 map = compressUnusedSymbols(map);920 return map;921}922 923llvm::SmallBitVector mlir::getUnusedDimsBitVector(ArrayRef<AffineMap> maps) {924 unsigned numDims = maps[0].getNumDims();925 llvm::SmallBitVector numDimsBitVector(numDims, true);926 for (AffineMap m : maps) {927 for (unsigned i = 0; i < numDims; ++i) {928 if (m.isFunctionOfDim(i))929 numDimsBitVector.reset(i);930 }931 }932 return numDimsBitVector;933}934 935llvm::SmallBitVector mlir::getUnusedSymbolsBitVector(ArrayRef<AffineMap> maps) {936 unsigned numSymbols = maps[0].getNumSymbols();937 llvm::SmallBitVector numSymbolsBitVector(numSymbols, true);938 for (AffineMap m : maps) {939 for (unsigned i = 0; i < numSymbols; ++i) {940 if (m.isFunctionOfSymbol(i))941 numSymbolsBitVector.reset(i);942 }943 }944 return numSymbolsBitVector;945}946 947AffineMap948mlir::expandDimsToRank(AffineMap map, int64_t rank,949 const llvm::SmallBitVector &projectedDimensions) {950 auto id = AffineMap::getMultiDimIdentityMap(rank, map.getContext());951 AffineMap proj = id.dropResults(projectedDimensions);952 return map.compose(proj);953}954 955//===----------------------------------------------------------------------===//956// MutableAffineMap.957//===----------------------------------------------------------------------===//958 959MutableAffineMap::MutableAffineMap(AffineMap map)960 : results(map.getResults()), numDims(map.getNumDims()),961 numSymbols(map.getNumSymbols()), context(map.getContext()) {}962 963void MutableAffineMap::reset(AffineMap map) {964 results.clear();965 numDims = map.getNumDims();966 numSymbols = map.getNumSymbols();967 context = map.getContext();968 llvm::append_range(results, map.getResults());969}970 971bool MutableAffineMap::isMultipleOf(unsigned idx, int64_t factor) const {972 return results[idx].isMultipleOf(factor);973}974 975// Simplifies the result affine expressions of this map. The expressions976// have to be pure for the simplification implemented.977void MutableAffineMap::simplify() {978 // Simplify each of the results if possible.979 // TODO: functional-style map980 for (unsigned i = 0, e = getNumResults(); i < e; i++) {981 results[i] = simplifyAffineExpr(getResult(i), numDims, numSymbols);982 }983}984 985AffineMap MutableAffineMap::getAffineMap() const {986 return AffineMap::get(numDims, numSymbols, results, context);987}988