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