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
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(in_start: operands.begin(), in_end: operands.end());
24 this->results.assign(in_start: results.begin(), in_end: results.end());
25}
26
27void AffineValueMap::composeSimplifyAndCanonicalize() {
28 AffineMap sMap = getAffineMap();
29 fullyComposeAffineMapAndOperands(map: &sMap, operands: &operands);
30 // Full composition also canonicalizes and simplifies before returning. We
31 // need to canonicalize once more to drop unused operands.
32 canonicalizeMapAndOperands(map: &sMap, operands: &operands);
33 this->map.reset(map: 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(N: a.getNumOperands() + b.getNumOperands());
42 auto aDims = a.getOperands().take_front(N: a.getNumDims());
43 auto bDims = b.getOperands().take_front(N: b.getNumDims());
44 auto aSyms = a.getOperands().take_back(N: a.getNumSymbols());
45 auto bSyms = b.getOperands().take_back(N: b.getNumSymbols());
46 allOperands.append(in_start: aDims.begin(), in_end: aDims.end());
47 allOperands.append(in_start: bDims.begin(), in_end: bDims.end());
48 allOperands.append(in_start: aSyms.begin(), in_end: aSyms.end());
49 allOperands.append(in_start: bSyms.begin(), in_end: bSyms.end());
50
51 // Shift dims and symbols of b's map.
52 auto bMap = b.getAffineMap()
53 .shiftDims(shift: a.getNumDims())
54 .shiftSymbols(shift: a.getNumSymbols());
55
56 // Construct the difference expressions.
57 auto aMap = a.getAffineMap();
58 SmallVector<AffineExpr, 4> diffExprs;
59 diffExprs.reserve(N: a.getNumResults());
60 for (unsigned i = 0, e = bMap.getNumResults(); i < e; ++i)
61 diffExprs.push_back(Elt: aMap.getResult(idx: i) - bMap.getResult(idx: i));
62
63 auto diffMap = AffineMap::get(dimCount: bMap.getNumDims(), symbolCount: bMap.getNumSymbols(),
64 results: diffExprs, context: bMap.getContext());
65 fullyComposeAffineMapAndOperands(map: &diffMap, operands: &allOperands);
66 canonicalizeMapAndOperands(map: &diffMap, operands: &allOperands);
67 diffMap = simplifyAffineMap(map: diffMap);
68 res->reset(map: diffMap, operands: allOperands);
69}
70
71// Returns true and sets 'indexOfMatch' if 'valueToMatch' is found in
72// '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 aligned
90/// with the AffineDimExpr in the underlying AffineMap.
91bool AffineValueMap::isFunctionOf(unsigned idx, Value value) const {
92 unsigned index;
93 if (!findIndex(valueToMatch: value, valuesToSearch: operands, /*indexStart=*/0, indexOfMatch: &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(position: 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(a: *this, b: other, res: &diff);
115 return llvm::all_of(Range: diff.getAffineMap().getResults(), P: [](AffineExpr e) {
116 return e == getAffineConstantExpr(constant: 0, context: e.getContext());
117 });
118}
119
120AffineValueMap::~AffineValueMap() = default;
121

source code of mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp