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