annotate mlir/lib/IR/AffineMap.cpp @ 173:0572611fdcc8 llvm10 llvm12

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