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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/IR/AffineMap.h"
10 #include "AffineMapDetail.h"
11 #include "mlir/IR/Attributes.h"
12 #include "mlir/IR/StandardTypes.h"
13 #include "mlir/Support/LogicalResult.h"
14 #include "mlir/Support/MathExtras.h"
15 #include "llvm/ADT/StringRef.h"
16 #include "llvm/Support/raw_ostream.h"
17 
18 using namespace mlir;
19 
20 namespace {
21 
22 // AffineExprConstantFolder evaluates an affine expression using constant
23 // operands passed in 'operandConsts'. Returns an IntegerAttr attribute
24 // representing the constant value of the affine expression evaluated on
25 // constant 'operandConsts', or nullptr if it can't be folded.
26 class AffineExprConstantFolder {
27 public:
28   AffineExprConstantFolder(unsigned numDims, ArrayRef<Attribute> operandConsts)
29       : numDims(numDims), operandConsts(operandConsts) {}
30 
31   /// Attempt to constant fold the specified affine expr, or return null on
32   /// failure.
33   IntegerAttr constantFold(AffineExpr expr) {
34     if (auto result = constantFoldImpl(expr))
35       return IntegerAttr::get(IndexType::get(expr.getContext()), *result);
36     return nullptr;
37   }
38 
39 private:
40   Optional<int64_t> constantFoldImpl(AffineExpr expr) {
41     switch (expr.getKind()) {
42     case AffineExprKind::Add:
43       return constantFoldBinExpr(
44           expr, [](int64_t lhs, int64_t rhs) { return lhs + rhs; });
45     case AffineExprKind::Mul:
46       return constantFoldBinExpr(
47           expr, [](int64_t lhs, int64_t rhs) { return lhs * rhs; });
48     case AffineExprKind::Mod:
49       return constantFoldBinExpr(
50           expr, [](int64_t lhs, int64_t rhs) { return mod(lhs, rhs); });
51     case AffineExprKind::FloorDiv:
52       return constantFoldBinExpr(
53           expr, [](int64_t lhs, int64_t rhs) { return floorDiv(lhs, rhs); });
54     case AffineExprKind::CeilDiv:
55       return constantFoldBinExpr(
56           expr, [](int64_t lhs, int64_t rhs) { return ceilDiv(lhs, rhs); });
57     case AffineExprKind::Constant:
58       return expr.cast<AffineConstantExpr>().getValue();
59     case AffineExprKind::DimId:
60       if (auto attr = operandConsts[expr.cast<AffineDimExpr>().getPosition()]
61                           .dyn_cast_or_null<IntegerAttr>())
62         return attr.getInt();
63       return llvm::None;
64     case AffineExprKind::SymbolId:
65       if (auto attr = operandConsts[numDims +
66                                     expr.cast<AffineSymbolExpr>().getPosition()]
67                           .dyn_cast_or_null<IntegerAttr>())
68         return attr.getInt();
69       return llvm::None;
70     }
71     llvm_unreachable("Unknown AffineExpr");
72   }
73 
74   // TODO: Change these to operate on APInts too.
75   Optional<int64_t> constantFoldBinExpr(AffineExpr expr,
76                                         int64_t (*op)(int64_t, int64_t)) {
77     auto binOpExpr = expr.cast<AffineBinaryOpExpr>();
78     if (auto lhs = constantFoldImpl(binOpExpr.getLHS()))
79       if (auto rhs = constantFoldImpl(binOpExpr.getRHS()))
80         return op(*lhs, *rhs);
81     return llvm::None;
82   }
83 
84   // The number of dimension operands in AffineMap containing this expression.
85   unsigned numDims;
86   // The constant valued operands used to evaluate this AffineExpr.
87   ArrayRef<Attribute> operandConsts;
88 };
89 
90 } // end anonymous namespace
91 
92 /// Returns a single constant result affine map.
93 AffineMap AffineMap::getConstantMap(int64_t val, MLIRContext *context) {
94   return get(/*dimCount=*/0, /*symbolCount=*/0,
95              {getAffineConstantExpr(val, context)});
96 }
97 
98 /// Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most
99 /// minor dimensions.
100 AffineMap AffineMap::getMinorIdentityMap(unsigned dims, unsigned results,
101                                          MLIRContext *context) {
102   assert(dims >= results && "Dimension mismatch");
103   auto id = AffineMap::getMultiDimIdentityMap(dims, context);
104   return AffineMap::get(dims, 0, id.getResults().take_back(results), context);
105 }
106 
107 bool AffineMap::isMinorIdentity(AffineMap map) {
108   if (!map)
109     return false;
110   return map == getMinorIdentityMap(map.getNumDims(), map.getNumResults(),
111                                     map.getContext());
112 };
113 
114 /// Returns an AffineMap representing a permutation.
115 AffineMap AffineMap::getPermutationMap(ArrayRef<unsigned> permutation,
116                                        MLIRContext *context) {
117   assert(!permutation.empty() &&
118          "Cannot create permutation map from empty permutation vector");
119   SmallVector<AffineExpr, 4> affExprs;
120   for (auto index : permutation)
121     affExprs.push_back(getAffineDimExpr(index, context));
122   auto m = std::max_element(permutation.begin(), permutation.end());
123   auto permutationMap = AffineMap::get(*m + 1, 0, affExprs, context);
124   assert(permutationMap.isPermutation() && "Invalid permutation vector");
125   return permutationMap;
126 }
127 
128 template <typename AffineExprContainer>
129 static void getMaxDimAndSymbol(ArrayRef<AffineExprContainer> exprsList,
130                                int64_t &maxDim, int64_t &maxSym) {
131   for (const auto &exprs : exprsList) {
132     for (auto expr : exprs) {
133       expr.walk([&maxDim, &maxSym](AffineExpr e) {
134         if (auto d = e.dyn_cast<AffineDimExpr>())
135           maxDim = std::max(maxDim, static_cast<int64_t>(d.getPosition()));
136         if (auto s = e.dyn_cast<AffineSymbolExpr>())
137           maxSym = std::max(maxSym, static_cast<int64_t>(s.getPosition()));
138       });
139     }
140   }
141 }
142 
143 template <typename AffineExprContainer>
144 static SmallVector<AffineMap, 4>
145 inferFromExprList(ArrayRef<AffineExprContainer> exprsList) {
146   assert(!exprsList.empty());
147   assert(!exprsList[0].empty());
148   auto context = exprsList[0][0].getContext();
149   int64_t maxDim = -1, maxSym = -1;
150   getMaxDimAndSymbol(exprsList, maxDim, maxSym);
151   SmallVector<AffineMap, 4> maps;
152   maps.reserve(exprsList.size());
153   for (const auto &exprs : exprsList)
154     maps.push_back(AffineMap::get(/*dimCount=*/maxDim + 1,
155                                   /*symbolCount=*/maxSym + 1, exprs, context));
156   return maps;
157 }
158 
159 SmallVector<AffineMap, 4>
160 AffineMap::inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList) {
161   return ::inferFromExprList(exprsList);
162 }
163 
164 SmallVector<AffineMap, 4>
165 AffineMap::inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList) {
166   return ::inferFromExprList(exprsList);
167 }
168 
169 AffineMap AffineMap::getMultiDimIdentityMap(unsigned numDims,
170                                             MLIRContext *context) {
171   SmallVector<AffineExpr, 4> dimExprs;
172   dimExprs.reserve(numDims);
173   for (unsigned i = 0; i < numDims; ++i)
174     dimExprs.push_back(mlir::getAffineDimExpr(i, context));
175   return get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs, context);
176 }
177 
178 MLIRContext *AffineMap::getContext() const { return map->context; }
179 
180 bool AffineMap::isIdentity() const {
181   if (getNumDims() != getNumResults())
182     return false;
183   ArrayRef<AffineExpr> results = getResults();
184   for (unsigned i = 0, numDims = getNumDims(); i < numDims; ++i) {
185     auto expr = results[i].dyn_cast<AffineDimExpr>();
186     if (!expr || expr.getPosition() != i)
187       return false;
188   }
189   return true;
190 }
191 
192 bool AffineMap::isEmpty() const {
193   return getNumDims() == 0 && getNumSymbols() == 0 && getNumResults() == 0;
194 }
195 
196 bool AffineMap::isSingleConstant() const {
197   return getNumResults() == 1 && getResult(0).isa<AffineConstantExpr>();
198 }
199 
200 int64_t AffineMap::getSingleConstantResult() const {
201   assert(isSingleConstant() && "map must have a single constant result");
202   return getResult(0).cast<AffineConstantExpr>().getValue();
203 }
204 
205 unsigned AffineMap::getNumDims() const {
206   assert(map && "uninitialized map storage");
207   return map->numDims;
208 }
209 unsigned AffineMap::getNumSymbols() const {
210   assert(map && "uninitialized map storage");
211   return map->numSymbols;
212 }
213 unsigned AffineMap::getNumResults() const {
214   assert(map && "uninitialized map storage");
215   return map->results.size();
216 }
217 unsigned AffineMap::getNumInputs() const {
218   assert(map && "uninitialized map storage");
219   return map->numDims + map->numSymbols;
220 }
221 
222 ArrayRef<AffineExpr> AffineMap::getResults() const {
223   assert(map && "uninitialized map storage");
224   return map->results;
225 }
226 AffineExpr AffineMap::getResult(unsigned idx) const {
227   assert(map && "uninitialized map storage");
228   return map->results[idx];
229 }
230 
231 /// Folds the results of the application of an affine map on the provided
232 /// operands to a constant if possible. Returns false if the folding happens,
233 /// true otherwise.
234 LogicalResult
235 AffineMap::constantFold(ArrayRef<Attribute> operandConstants,
236                         SmallVectorImpl<Attribute> &results) const {
237   assert(getNumInputs() == operandConstants.size());
238 
239   // Fold each of the result expressions.
240   AffineExprConstantFolder exprFolder(getNumDims(), operandConstants);
241   // Constant fold each AffineExpr in AffineMap and add to 'results'.
242   for (auto expr : getResults()) {
243     auto folded = exprFolder.constantFold(expr);
244     // If we didn't fold to a constant, then folding fails.
245     if (!folded)
246       return failure();
247 
248     results.push_back(folded);
249   }
250   assert(results.size() == getNumResults() &&
251          "constant folding produced the wrong number of results");
252   return success();
253 }
254 
255 /// Walk all of the AffineExpr's in this mapping. Each node in an expression
256 /// tree is visited in postorder.
257 void AffineMap::walkExprs(std::function<void(AffineExpr)> callback) const {
258   for (auto expr : getResults())
259     expr.walk(callback);
260 }
261 
262 /// This method substitutes any uses of dimensions and symbols (e.g.
263 /// dim#0 with dimReplacements[0]) in subexpressions and returns the modified
264 /// expression mapping.  Because this can be used to eliminate dims and
265 /// symbols, the client needs to specify the number of dims and symbols in
266 /// the result.  The returned map always has the same number of results.
267 AffineMap AffineMap::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
268                                            ArrayRef<AffineExpr> symReplacements,
269                                            unsigned numResultDims,
270                                            unsigned numResultSyms) {
271   SmallVector<AffineExpr, 8> results;
272   results.reserve(getNumResults());
273   for (auto expr : getResults())
274     results.push_back(
275         expr.replaceDimsAndSymbols(dimReplacements, symReplacements));
276 
277   return get(numResultDims, numResultSyms, results, getContext());
278 }
279 
280 AffineMap AffineMap::compose(AffineMap map) {
281   assert(getNumDims() == map.getNumResults() && "Number of results mismatch");
282   // Prepare `map` by concatenating the symbols and rewriting its exprs.
283   unsigned numDims = map.getNumDims();
284   unsigned numSymbolsThisMap = getNumSymbols();
285   unsigned numSymbols = numSymbolsThisMap + map.getNumSymbols();
286   SmallVector<AffineExpr, 8> newDims(numDims);
287   for (unsigned idx = 0; idx < numDims; ++idx) {
288     newDims[idx] = getAffineDimExpr(idx, getContext());
289   }
290   SmallVector<AffineExpr, 8> newSymbols(numSymbols);
291   for (unsigned idx = numSymbolsThisMap; idx < numSymbols; ++idx) {
292     newSymbols[idx - numSymbolsThisMap] =
293         getAffineSymbolExpr(idx, getContext());
294   }
295   auto newMap =
296       map.replaceDimsAndSymbols(newDims, newSymbols, numDims, numSymbols);
297   SmallVector<AffineExpr, 8> exprs;
298   exprs.reserve(getResults().size());
299   for (auto expr : getResults())
300     exprs.push_back(expr.compose(newMap));
301   return AffineMap::get(numDims, numSymbols, exprs, map.getContext());
302 }
303 
304 bool AffineMap::isProjectedPermutation() {
305   if (getNumSymbols() > 0)
306     return false;
307   SmallVector<bool, 8> seen(getNumInputs(), false);
308   for (auto expr : getResults()) {
309     if (auto dim = expr.dyn_cast<AffineDimExpr>()) {
310       if (seen[dim.getPosition()])
311         return false;
312       seen[dim.getPosition()] = true;
313       continue;
314     }
315     return false;
316   }
317   return true;
318 }
319 
320 bool AffineMap::isPermutation() {
321   if (getNumDims() != getNumResults())
322     return false;
323   return isProjectedPermutation();
324 }
325 
326 AffineMap AffineMap::getSubMap(ArrayRef<unsigned> resultPos) {
327   SmallVector<AffineExpr, 4> exprs;
328   exprs.reserve(resultPos.size());
329   for (auto idx : resultPos) {
330     exprs.push_back(getResult(idx));
331   }
332   return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());
333 }
334 
335 AffineMap mlir::simplifyAffineMap(AffineMap map) {
336   SmallVector<AffineExpr, 8> exprs;
337   for (auto e : map.getResults()) {
338     exprs.push_back(
339         simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols()));
340   }
341   return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs,
342                         map.getContext());
343 }
344 
345 AffineMap mlir::removeDuplicateExprs(AffineMap map) {
346   auto results = map.getResults();
347   SmallVector<AffineExpr, 4> uniqueExprs(results.begin(), results.end());
348   uniqueExprs.erase(std::unique(uniqueExprs.begin(), uniqueExprs.end()),
349                     uniqueExprs.end());
350   return AffineMap::get(map.getNumDims(), map.getNumSymbols(), uniqueExprs,
351                         map.getContext());
352 }
353 
354 AffineMap mlir::inversePermutation(AffineMap map) {
355   if (map.isEmpty())
356     return map;
357   assert(map.getNumSymbols() == 0 && "expected map without symbols");
358   SmallVector<AffineExpr, 4> exprs(map.getNumDims());
359   for (auto en : llvm::enumerate(map.getResults())) {
360     auto expr = en.value();
361     // Skip non-permutations.
362     if (auto d = expr.dyn_cast<AffineDimExpr>()) {
363       if (exprs[d.getPosition()])
364         continue;
365       exprs[d.getPosition()] = getAffineDimExpr(en.index(), d.getContext());
366     }
367   }
368   SmallVector<AffineExpr, 4> seenExprs;
369   seenExprs.reserve(map.getNumDims());
370   for (auto expr : exprs)
371     if (expr)
372       seenExprs.push_back(expr);
373   if (seenExprs.size() != map.getNumInputs())
374     return AffineMap();
375   return AffineMap::get(map.getNumResults(), 0, seenExprs, map.getContext());
376 }
377 
378 AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
379   unsigned numResults = 0;
380   for (auto m : maps)
381     numResults += m.getNumResults();
382   unsigned numDims = 0;
383   SmallVector<AffineExpr, 8> results;
384   results.reserve(numResults);
385   for (auto m : maps) {
386     assert(m.getNumSymbols() == 0 && "expected map without symbols");
387     results.append(m.getResults().begin(), m.getResults().end());
388     numDims = std::max(m.getNumDims(), numDims);
389   }
390   return AffineMap::get(numDims, /*numSymbols=*/0, results,
391                         maps.front().getContext());
392 }
393 
394 //===----------------------------------------------------------------------===//
395 // MutableAffineMap.
396 //===----------------------------------------------------------------------===//
397 
398 MutableAffineMap::MutableAffineMap(AffineMap map)
399     : numDims(map.getNumDims()), numSymbols(map.getNumSymbols()),
400       context(map.getContext()) {
401   for (auto result : map.getResults())
402     results.push_back(result);
403 }
404 
405 void MutableAffineMap::reset(AffineMap map) {
406   results.clear();
407   numDims = map.getNumDims();
408   numSymbols = map.getNumSymbols();
409   context = map.getContext();
410   for (auto result : map.getResults())
411     results.push_back(result);
412 }
413 
414 bool MutableAffineMap::isMultipleOf(unsigned idx, int64_t factor) const {
415   if (results[idx].isMultipleOf(factor))
416     return true;
417 
418   // TODO(bondhugula): use simplifyAffineExpr and FlatAffineConstraints to
419   // complete this (for a more powerful analysis).
420   return false;
421 }
422 
423 // Simplifies the result affine expressions of this map. The expressions have to
424 // be pure for the simplification implemented.
425 void MutableAffineMap::simplify() {
426   // Simplify each of the results if possible.
427   // TODO(ntv): functional-style map
428   for (unsigned i = 0, e = getNumResults(); i < e; i++) {
429     results[i] = simplifyAffineExpr(getResult(i), numDims, numSymbols);
430   }
431 }
432 
433 AffineMap MutableAffineMap::getAffineMap() const {
434   return AffineMap::get(numDims, numSymbols, results, context);
435 }
436