121 lines · cpp
1//===- AffineValueMap.cpp - MLIR Affine Value Map Class -------------------===//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/Dialect/Affine/IR/AffineValueMap.h"10#include "mlir/Dialect/Affine/IR/AffineOps.h"11 12using namespace mlir;13using namespace mlir::affine;14 15AffineValueMap::AffineValueMap(AffineMap map, ValueRange operands,16 ValueRange results)17 : map(map), operands(operands.begin(), operands.end()),18 results(results.begin(), results.end()) {}19 20void AffineValueMap::reset(AffineMap map, ValueRange operands,21 ValueRange results) {22 this->map.reset(map);23 this->operands.assign(operands.begin(), operands.end());24 this->results.assign(results.begin(), results.end());25}26 27void AffineValueMap::composeSimplifyAndCanonicalize() {28 AffineMap sMap = getAffineMap();29 fullyComposeAffineMapAndOperands(&sMap, &operands);30 // Full composition also canonicalizes and simplifies before returning. We31 // need to canonicalize once more to drop unused operands.32 canonicalizeMapAndOperands(&sMap, &operands);33 this->map.reset(sMap);34}35 36void AffineValueMap::difference(const AffineValueMap &a,37 const AffineValueMap &b, AffineValueMap *res) {38 assert(a.getNumResults() == b.getNumResults() && "invalid inputs");39 40 SmallVector<Value, 4> allOperands;41 allOperands.reserve(a.getNumOperands() + b.getNumOperands());42 auto aDims = a.getOperands().take_front(a.getNumDims());43 auto bDims = b.getOperands().take_front(b.getNumDims());44 auto aSyms = a.getOperands().take_back(a.getNumSymbols());45 auto bSyms = b.getOperands().take_back(b.getNumSymbols());46 allOperands.append(aDims.begin(), aDims.end());47 allOperands.append(bDims.begin(), bDims.end());48 allOperands.append(aSyms.begin(), aSyms.end());49 allOperands.append(bSyms.begin(), bSyms.end());50 51 // Shift dims and symbols of b's map.52 auto bMap = b.getAffineMap()53 .shiftDims(a.getNumDims())54 .shiftSymbols(a.getNumSymbols());55 56 // Construct the difference expressions.57 auto aMap = a.getAffineMap();58 SmallVector<AffineExpr, 4> diffExprs;59 diffExprs.reserve(a.getNumResults());60 for (unsigned i = 0, e = bMap.getNumResults(); i < e; ++i)61 diffExprs.push_back(aMap.getResult(i) - bMap.getResult(i));62 63 auto diffMap = AffineMap::get(bMap.getNumDims(), bMap.getNumSymbols(),64 diffExprs, bMap.getContext());65 fullyComposeAffineMapAndOperands(&diffMap, &allOperands);66 canonicalizeMapAndOperands(&diffMap, &allOperands);67 diffMap = simplifyAffineMap(diffMap);68 res->reset(diffMap, allOperands);69}70 71// Returns true and sets 'indexOfMatch' if 'valueToMatch' is found in72// 'valuesToSearch' beginning at 'indexStart'. Returns false otherwise.73static bool findIndex(Value valueToMatch, ArrayRef<Value> valuesToSearch,74 unsigned indexStart, unsigned *indexOfMatch) {75 unsigned size = valuesToSearch.size();76 for (unsigned i = indexStart; i < size; ++i) {77 if (valueToMatch == valuesToSearch[i]) {78 *indexOfMatch = i;79 return true;80 }81 }82 return false;83}84 85bool AffineValueMap::isMultipleOf(unsigned idx, int64_t factor) const {86 return map.isMultipleOf(idx, factor);87}88 89/// This method uses the invariant that operands are always positionally aligned90/// with the AffineDimExpr in the underlying AffineMap.91bool AffineValueMap::isFunctionOf(unsigned idx, Value value) const {92 unsigned index;93 if (!findIndex(value, operands, /*indexStart=*/0, &index)) {94 return false;95 }96 auto expr = const_cast<AffineValueMap *>(this)->getAffineMap().getResult(idx);97 // TODO: this is better implemented on a flattened representation.98 // At least for now it is conservative.99 return expr.isFunctionOfDim(index);100}101 102Value AffineValueMap::getOperand(unsigned i) const {103 return static_cast<Value>(operands[i]);104}105 106ArrayRef<Value> AffineValueMap::getOperands() const {107 return ArrayRef<Value>(operands);108}109 110AffineMap AffineValueMap::getAffineMap() const { return map.getAffineMap(); }111 112bool AffineValueMap::operator==(const AffineValueMap &other) const {113 AffineValueMap diff;114 AffineValueMap::difference(*this, other, &diff);115 return llvm::all_of(diff.getAffineMap().getResults(), [](AffineExpr e) {116 return e == getAffineConstantExpr(0, e.getContext());117 });118}119 120AffineValueMap::~AffineValueMap() = default;121