File: | build/source/mlir/lib/IR/AffineExpr.cpp |
Warning: | line 25, column 54 Access to field 'context' results in a dereference of a null pointer (loaded from field 'expr') |
Press '?' to see keyboard shortcuts
Keyboard shortcuts:
1 | //===- AffineExpr.cpp - MLIR Affine Expr 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 <utility> | |||
10 | ||||
11 | #include "AffineExprDetail.h" | |||
12 | #include "mlir/IR/AffineExpr.h" | |||
13 | #include "mlir/IR/AffineExprVisitor.h" | |||
14 | #include "mlir/IR/AffineMap.h" | |||
15 | #include "mlir/IR/IntegerSet.h" | |||
16 | #include "mlir/Support/MathExtras.h" | |||
17 | #include "mlir/Support/TypeID.h" | |||
18 | #include "llvm/ADT/STLExtras.h" | |||
19 | #include <numeric> | |||
20 | #include <optional> | |||
21 | ||||
22 | using namespace mlir; | |||
23 | using namespace mlir::detail; | |||
24 | ||||
25 | MLIRContext *AffineExpr::getContext() const { return expr->context; } | |||
| ||||
26 | ||||
27 | AffineExprKind AffineExpr::getKind() const { return expr->kind; } | |||
28 | ||||
29 | /// Walk all of the AffineExprs in this subgraph in postorder. | |||
30 | void AffineExpr::walk(std::function<void(AffineExpr)> callback) const { | |||
31 | struct AffineExprWalker : public AffineExprVisitor<AffineExprWalker> { | |||
32 | std::function<void(AffineExpr)> callback; | |||
33 | ||||
34 | AffineExprWalker(std::function<void(AffineExpr)> callback) | |||
35 | : callback(std::move(callback)) {} | |||
36 | ||||
37 | void visitAffineBinaryOpExpr(AffineBinaryOpExpr expr) { callback(expr); } | |||
38 | void visitConstantExpr(AffineConstantExpr expr) { callback(expr); } | |||
39 | void visitDimExpr(AffineDimExpr expr) { callback(expr); } | |||
40 | void visitSymbolExpr(AffineSymbolExpr expr) { callback(expr); } | |||
41 | }; | |||
42 | ||||
43 | AffineExprWalker(std::move(callback)).walkPostOrder(*this); | |||
44 | } | |||
45 | ||||
46 | // Dispatch affine expression construction based on kind. | |||
47 | AffineExpr mlir::getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs, | |||
48 | AffineExpr rhs) { | |||
49 | if (kind
| |||
50 | return lhs + rhs; | |||
51 | if (kind
| |||
52 | return lhs * rhs; | |||
53 | if (kind
| |||
54 | return lhs.floorDiv(rhs); | |||
55 | if (kind
| |||
56 | return lhs.ceilDiv(rhs); | |||
57 | if (kind == AffineExprKind::Mod) | |||
58 | return lhs % rhs; | |||
59 | ||||
60 | llvm_unreachable("unknown binary operation on affine expressions")::llvm::llvm_unreachable_internal("unknown binary operation on affine expressions" , "mlir/lib/IR/AffineExpr.cpp", 60); | |||
61 | } | |||
62 | ||||
63 | /// This method substitutes any uses of dimensions and symbols (e.g. | |||
64 | /// dim#0 with dimReplacements[0]) and returns the modified expression tree. | |||
65 | AffineExpr | |||
66 | AffineExpr::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements, | |||
67 | ArrayRef<AffineExpr> symReplacements) const { | |||
68 | switch (getKind()) { | |||
69 | case AffineExprKind::Constant: | |||
70 | return *this; | |||
71 | case AffineExprKind::DimId: { | |||
72 | unsigned dimId = cast<AffineDimExpr>().getPosition(); | |||
73 | if (dimId >= dimReplacements.size()) | |||
74 | return *this; | |||
75 | return dimReplacements[dimId]; | |||
76 | } | |||
77 | case AffineExprKind::SymbolId: { | |||
78 | unsigned symId = cast<AffineSymbolExpr>().getPosition(); | |||
79 | if (symId >= symReplacements.size()) | |||
80 | return *this; | |||
81 | return symReplacements[symId]; | |||
82 | } | |||
83 | case AffineExprKind::Add: | |||
84 | case AffineExprKind::Mul: | |||
85 | case AffineExprKind::FloorDiv: | |||
86 | case AffineExprKind::CeilDiv: | |||
87 | case AffineExprKind::Mod: | |||
88 | auto binOp = cast<AffineBinaryOpExpr>(); | |||
89 | auto lhs = binOp.getLHS(), rhs = binOp.getRHS(); | |||
90 | auto newLHS = lhs.replaceDimsAndSymbols(dimReplacements, symReplacements); | |||
91 | auto newRHS = rhs.replaceDimsAndSymbols(dimReplacements, symReplacements); | |||
92 | if (newLHS == lhs && newRHS == rhs) | |||
93 | return *this; | |||
94 | return getAffineBinaryOpExpr(getKind(), newLHS, newRHS); | |||
95 | } | |||
96 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 96); | |||
97 | } | |||
98 | ||||
99 | AffineExpr AffineExpr::replaceDims(ArrayRef<AffineExpr> dimReplacements) const { | |||
100 | return replaceDimsAndSymbols(dimReplacements, {}); | |||
101 | } | |||
102 | ||||
103 | AffineExpr | |||
104 | AffineExpr::replaceSymbols(ArrayRef<AffineExpr> symReplacements) const { | |||
105 | return replaceDimsAndSymbols({}, symReplacements); | |||
106 | } | |||
107 | ||||
108 | /// Replace dims[offset ... numDims) | |||
109 | /// by dims[offset + shift ... shift + numDims). | |||
110 | AffineExpr AffineExpr::shiftDims(unsigned numDims, unsigned shift, | |||
111 | unsigned offset) const { | |||
112 | SmallVector<AffineExpr, 4> dims; | |||
113 | for (unsigned idx = 0; idx < offset; ++idx) | |||
114 | dims.push_back(getAffineDimExpr(idx, getContext())); | |||
115 | for (unsigned idx = offset; idx < numDims; ++idx) | |||
116 | dims.push_back(getAffineDimExpr(idx + shift, getContext())); | |||
117 | return replaceDimsAndSymbols(dims, {}); | |||
118 | } | |||
119 | ||||
120 | /// Replace symbols[offset ... numSymbols) | |||
121 | /// by symbols[offset + shift ... shift + numSymbols). | |||
122 | AffineExpr AffineExpr::shiftSymbols(unsigned numSymbols, unsigned shift, | |||
123 | unsigned offset) const { | |||
124 | SmallVector<AffineExpr, 4> symbols; | |||
125 | for (unsigned idx = 0; idx < offset; ++idx) | |||
126 | symbols.push_back(getAffineSymbolExpr(idx, getContext())); | |||
127 | for (unsigned idx = offset; idx < numSymbols; ++idx) | |||
128 | symbols.push_back(getAffineSymbolExpr(idx + shift, getContext())); | |||
129 | return replaceDimsAndSymbols({}, symbols); | |||
130 | } | |||
131 | ||||
132 | /// Sparse replace method. Return the modified expression tree. | |||
133 | AffineExpr | |||
134 | AffineExpr::replace(const DenseMap<AffineExpr, AffineExpr> &map) const { | |||
135 | auto it = map.find(*this); | |||
136 | if (it != map.end()) | |||
137 | return it->second; | |||
138 | switch (getKind()) { | |||
139 | default: | |||
140 | return *this; | |||
141 | case AffineExprKind::Add: | |||
142 | case AffineExprKind::Mul: | |||
143 | case AffineExprKind::FloorDiv: | |||
144 | case AffineExprKind::CeilDiv: | |||
145 | case AffineExprKind::Mod: | |||
146 | auto binOp = cast<AffineBinaryOpExpr>(); | |||
147 | auto lhs = binOp.getLHS(), rhs = binOp.getRHS(); | |||
148 | auto newLHS = lhs.replace(map); | |||
149 | auto newRHS = rhs.replace(map); | |||
150 | if (newLHS == lhs && newRHS == rhs) | |||
151 | return *this; | |||
152 | return getAffineBinaryOpExpr(getKind(), newLHS, newRHS); | |||
153 | } | |||
154 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 154); | |||
155 | } | |||
156 | ||||
157 | /// Sparse replace method. Return the modified expression tree. | |||
158 | AffineExpr AffineExpr::replace(AffineExpr expr, AffineExpr replacement) const { | |||
159 | DenseMap<AffineExpr, AffineExpr> map; | |||
160 | map.insert(std::make_pair(expr, replacement)); | |||
161 | return replace(map); | |||
| ||||
162 | } | |||
163 | /// Returns true if this expression is made out of only symbols and | |||
164 | /// constants (no dimensional identifiers). | |||
165 | bool AffineExpr::isSymbolicOrConstant() const { | |||
166 | switch (getKind()) { | |||
167 | case AffineExprKind::Constant: | |||
168 | return true; | |||
169 | case AffineExprKind::DimId: | |||
170 | return false; | |||
171 | case AffineExprKind::SymbolId: | |||
172 | return true; | |||
173 | ||||
174 | case AffineExprKind::Add: | |||
175 | case AffineExprKind::Mul: | |||
176 | case AffineExprKind::FloorDiv: | |||
177 | case AffineExprKind::CeilDiv: | |||
178 | case AffineExprKind::Mod: { | |||
179 | auto expr = this->cast<AffineBinaryOpExpr>(); | |||
180 | return expr.getLHS().isSymbolicOrConstant() && | |||
181 | expr.getRHS().isSymbolicOrConstant(); | |||
182 | } | |||
183 | } | |||
184 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 184); | |||
185 | } | |||
186 | ||||
187 | /// Returns true if this is a pure affine expression, i.e., multiplication, | |||
188 | /// floordiv, ceildiv, and mod is only allowed w.r.t constants. | |||
189 | bool AffineExpr::isPureAffine() const { | |||
190 | switch (getKind()) { | |||
191 | case AffineExprKind::SymbolId: | |||
192 | case AffineExprKind::DimId: | |||
193 | case AffineExprKind::Constant: | |||
194 | return true; | |||
195 | case AffineExprKind::Add: { | |||
196 | auto op = cast<AffineBinaryOpExpr>(); | |||
197 | return op.getLHS().isPureAffine() && op.getRHS().isPureAffine(); | |||
198 | } | |||
199 | ||||
200 | case AffineExprKind::Mul: { | |||
201 | // TODO: Canonicalize the constants in binary operators to the RHS when | |||
202 | // possible, allowing this to merge into the next case. | |||
203 | auto op = cast<AffineBinaryOpExpr>(); | |||
204 | return op.getLHS().isPureAffine() && op.getRHS().isPureAffine() && | |||
205 | (op.getLHS().template isa<AffineConstantExpr>() || | |||
206 | op.getRHS().template isa<AffineConstantExpr>()); | |||
207 | } | |||
208 | case AffineExprKind::FloorDiv: | |||
209 | case AffineExprKind::CeilDiv: | |||
210 | case AffineExprKind::Mod: { | |||
211 | auto op = cast<AffineBinaryOpExpr>(); | |||
212 | return op.getLHS().isPureAffine() && | |||
213 | op.getRHS().template isa<AffineConstantExpr>(); | |||
214 | } | |||
215 | } | |||
216 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 216); | |||
217 | } | |||
218 | ||||
219 | // Returns the greatest known integral divisor of this affine expression. | |||
220 | int64_t AffineExpr::getLargestKnownDivisor() const { | |||
221 | AffineBinaryOpExpr binExpr(nullptr); | |||
222 | switch (getKind()) { | |||
223 | case AffineExprKind::DimId: | |||
224 | [[fallthrough]]; | |||
225 | case AffineExprKind::SymbolId: | |||
226 | return 1; | |||
227 | case AffineExprKind::CeilDiv: | |||
228 | [[fallthrough]]; | |||
229 | case AffineExprKind::FloorDiv: { | |||
230 | // If the RHS is a constant and divides the known divisor on the LHS, the | |||
231 | // quotient is a known divisor of the expression. | |||
232 | binExpr = this->cast<AffineBinaryOpExpr>(); | |||
233 | auto rhs = binExpr.getRHS().dyn_cast<AffineConstantExpr>(); | |||
234 | // Leave alone undefined expressions. | |||
235 | if (rhs && rhs.getValue() != 0) { | |||
236 | int64_t lhsDiv = binExpr.getLHS().getLargestKnownDivisor(); | |||
237 | if (lhsDiv % rhs.getValue() == 0) | |||
238 | return lhsDiv / rhs.getValue(); | |||
239 | } | |||
240 | return 1; | |||
241 | } | |||
242 | case AffineExprKind::Constant: | |||
243 | return std::abs(this->cast<AffineConstantExpr>().getValue()); | |||
244 | case AffineExprKind::Mul: { | |||
245 | binExpr = this->cast<AffineBinaryOpExpr>(); | |||
246 | return binExpr.getLHS().getLargestKnownDivisor() * | |||
247 | binExpr.getRHS().getLargestKnownDivisor(); | |||
248 | } | |||
249 | case AffineExprKind::Add: | |||
250 | [[fallthrough]]; | |||
251 | case AffineExprKind::Mod: { | |||
252 | binExpr = cast<AffineBinaryOpExpr>(); | |||
253 | return std::gcd((uint64_t)binExpr.getLHS().getLargestKnownDivisor(), | |||
254 | (uint64_t)binExpr.getRHS().getLargestKnownDivisor()); | |||
255 | } | |||
256 | } | |||
257 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 257); | |||
258 | } | |||
259 | ||||
260 | bool AffineExpr::isMultipleOf(int64_t factor) const { | |||
261 | AffineBinaryOpExpr binExpr(nullptr); | |||
262 | uint64_t l, u; | |||
263 | switch (getKind()) { | |||
264 | case AffineExprKind::SymbolId: | |||
265 | [[fallthrough]]; | |||
266 | case AffineExprKind::DimId: | |||
267 | return factor * factor == 1; | |||
268 | case AffineExprKind::Constant: | |||
269 | return cast<AffineConstantExpr>().getValue() % factor == 0; | |||
270 | case AffineExprKind::Mul: { | |||
271 | binExpr = cast<AffineBinaryOpExpr>(); | |||
272 | // It's probably not worth optimizing this further (to not traverse the | |||
273 | // whole sub-tree under - it that would require a version of isMultipleOf | |||
274 | // that on a 'false' return also returns the largest known divisor). | |||
275 | return (l = binExpr.getLHS().getLargestKnownDivisor()) % factor == 0 || | |||
276 | (u = binExpr.getRHS().getLargestKnownDivisor()) % factor == 0 || | |||
277 | (l * u) % factor == 0; | |||
278 | } | |||
279 | case AffineExprKind::Add: | |||
280 | case AffineExprKind::FloorDiv: | |||
281 | case AffineExprKind::CeilDiv: | |||
282 | case AffineExprKind::Mod: { | |||
283 | binExpr = cast<AffineBinaryOpExpr>(); | |||
284 | return std::gcd((uint64_t)binExpr.getLHS().getLargestKnownDivisor(), | |||
285 | (uint64_t)binExpr.getRHS().getLargestKnownDivisor()) % | |||
286 | factor == | |||
287 | 0; | |||
288 | } | |||
289 | } | |||
290 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 290); | |||
291 | } | |||
292 | ||||
293 | bool AffineExpr::isFunctionOfDim(unsigned position) const { | |||
294 | if (getKind() == AffineExprKind::DimId) { | |||
295 | return *this == mlir::getAffineDimExpr(position, getContext()); | |||
296 | } | |||
297 | if (auto expr = this->dyn_cast<AffineBinaryOpExpr>()) { | |||
298 | return expr.getLHS().isFunctionOfDim(position) || | |||
299 | expr.getRHS().isFunctionOfDim(position); | |||
300 | } | |||
301 | return false; | |||
302 | } | |||
303 | ||||
304 | bool AffineExpr::isFunctionOfSymbol(unsigned position) const { | |||
305 | if (getKind() == AffineExprKind::SymbolId) { | |||
306 | return *this == mlir::getAffineSymbolExpr(position, getContext()); | |||
307 | } | |||
308 | if (auto expr = this->dyn_cast<AffineBinaryOpExpr>()) { | |||
309 | return expr.getLHS().isFunctionOfSymbol(position) || | |||
310 | expr.getRHS().isFunctionOfSymbol(position); | |||
311 | } | |||
312 | return false; | |||
313 | } | |||
314 | ||||
315 | AffineBinaryOpExpr::AffineBinaryOpExpr(AffineExpr::ImplType *ptr) | |||
316 | : AffineExpr(ptr) {} | |||
317 | AffineExpr AffineBinaryOpExpr::getLHS() const { | |||
318 | return static_cast<ImplType *>(expr)->lhs; | |||
319 | } | |||
320 | AffineExpr AffineBinaryOpExpr::getRHS() const { | |||
321 | return static_cast<ImplType *>(expr)->rhs; | |||
322 | } | |||
323 | ||||
324 | AffineDimExpr::AffineDimExpr(AffineExpr::ImplType *ptr) : AffineExpr(ptr) {} | |||
325 | unsigned AffineDimExpr::getPosition() const { | |||
326 | return static_cast<ImplType *>(expr)->position; | |||
327 | } | |||
328 | ||||
329 | /// Returns true if the expression is divisible by the given symbol with | |||
330 | /// position `symbolPos`. The argument `opKind` specifies here what kind of | |||
331 | /// division or mod operation called this division. It helps in implementing the | |||
332 | /// commutative property of the floordiv and ceildiv operations. If the argument | |||
333 | ///`exprKind` is floordiv and `expr` is also a binary expression of a floordiv | |||
334 | /// operation, then the commutative property can be used otherwise, the floordiv | |||
335 | /// operation is not divisible. The same argument holds for ceildiv operation. | |||
336 | static bool isDivisibleBySymbol(AffineExpr expr, unsigned symbolPos, | |||
337 | AffineExprKind opKind) { | |||
338 | // The argument `opKind` can either be Modulo, Floordiv or Ceildiv only. | |||
339 | assert((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv ||(static_cast <bool> ((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv ) && "unexpected opKind") ? void (0) : __assert_fail ( "(opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv) && \"unexpected opKind\"" , "mlir/lib/IR/AffineExpr.cpp", 341, __extension__ __PRETTY_FUNCTION__ )) | |||
340 | opKind == AffineExprKind::CeilDiv) &&(static_cast <bool> ((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv ) && "unexpected opKind") ? void (0) : __assert_fail ( "(opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv) && \"unexpected opKind\"" , "mlir/lib/IR/AffineExpr.cpp", 341, __extension__ __PRETTY_FUNCTION__ )) | |||
341 | "unexpected opKind")(static_cast <bool> ((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv ) && "unexpected opKind") ? void (0) : __assert_fail ( "(opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv) && \"unexpected opKind\"" , "mlir/lib/IR/AffineExpr.cpp", 341, __extension__ __PRETTY_FUNCTION__ )); | |||
342 | switch (expr.getKind()) { | |||
343 | case AffineExprKind::Constant: | |||
344 | return expr.cast<AffineConstantExpr>().getValue() == 0; | |||
345 | case AffineExprKind::DimId: | |||
346 | return false; | |||
347 | case AffineExprKind::SymbolId: | |||
348 | return (expr.cast<AffineSymbolExpr>().getPosition() == symbolPos); | |||
349 | // Checks divisibility by the given symbol for both operands. | |||
350 | case AffineExprKind::Add: { | |||
351 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); | |||
352 | return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind) && | |||
353 | isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, opKind); | |||
354 | } | |||
355 | // Checks divisibility by the given symbol for both operands. Consider the | |||
356 | // expression `(((s1*s0) floordiv w) mod ((s1 * s2) floordiv p)) floordiv s1`, | |||
357 | // this is a division by s1 and both the operands of modulo are divisible by | |||
358 | // s1 but it is not divisible by s1 always. The third argument is | |||
359 | // `AffineExprKind::Mod` for this reason. | |||
360 | case AffineExprKind::Mod: { | |||
361 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); | |||
362 | return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, | |||
363 | AffineExprKind::Mod) && | |||
364 | isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, | |||
365 | AffineExprKind::Mod); | |||
366 | } | |||
367 | // Checks if any of the operand divisible by the given symbol. | |||
368 | case AffineExprKind::Mul: { | |||
369 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); | |||
370 | return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind) || | |||
371 | isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, opKind); | |||
372 | } | |||
373 | // Floordiv and ceildiv are divisible by the given symbol when the first | |||
374 | // operand is divisible, and the affine expression kind of the argument expr | |||
375 | // is same as the argument `opKind`. This can be inferred from commutative | |||
376 | // property of floordiv and ceildiv operations and are as follow: | |||
377 | // (exp1 floordiv exp2) floordiv exp3 = (exp1 floordiv exp3) floordiv exp2 | |||
378 | // (exp1 ceildiv exp2) ceildiv exp3 = (exp1 ceildiv exp3) ceildiv expr2 | |||
379 | // It will fail if operations are not same. For example: | |||
380 | // (exps1 ceildiv exp2) floordiv exp3 can not be simplified. | |||
381 | case AffineExprKind::FloorDiv: | |||
382 | case AffineExprKind::CeilDiv: { | |||
383 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); | |||
384 | if (opKind != expr.getKind()) | |||
385 | return false; | |||
386 | return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, expr.getKind()); | |||
387 | } | |||
388 | } | |||
389 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 389); | |||
390 | } | |||
391 | ||||
392 | /// Divides the given expression by the given symbol at position `symbolPos`. It | |||
393 | /// considers the divisibility condition is checked before calling itself. A | |||
394 | /// null expression is returned whenever the divisibility condition fails. | |||
395 | static AffineExpr symbolicDivide(AffineExpr expr, unsigned symbolPos, | |||
396 | AffineExprKind opKind) { | |||
397 | // THe argument `opKind` can either be Modulo, Floordiv or Ceildiv only. | |||
398 | assert((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv ||(static_cast <bool> ((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv ) && "unexpected opKind") ? void (0) : __assert_fail ( "(opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv) && \"unexpected opKind\"" , "mlir/lib/IR/AffineExpr.cpp", 400, __extension__ __PRETTY_FUNCTION__ )) | |||
399 | opKind == AffineExprKind::CeilDiv) &&(static_cast <bool> ((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv ) && "unexpected opKind") ? void (0) : __assert_fail ( "(opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv) && \"unexpected opKind\"" , "mlir/lib/IR/AffineExpr.cpp", 400, __extension__ __PRETTY_FUNCTION__ )) | |||
400 | "unexpected opKind")(static_cast <bool> ((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv ) && "unexpected opKind") ? void (0) : __assert_fail ( "(opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv) && \"unexpected opKind\"" , "mlir/lib/IR/AffineExpr.cpp", 400, __extension__ __PRETTY_FUNCTION__ )); | |||
401 | switch (expr.getKind()) { | |||
402 | case AffineExprKind::Constant: | |||
403 | if (expr.cast<AffineConstantExpr>().getValue() != 0) | |||
404 | return nullptr; | |||
405 | return getAffineConstantExpr(0, expr.getContext()); | |||
406 | case AffineExprKind::DimId: | |||
407 | return nullptr; | |||
408 | case AffineExprKind::SymbolId: | |||
409 | return getAffineConstantExpr(1, expr.getContext()); | |||
410 | // Dividing both operands by the given symbol. | |||
411 | case AffineExprKind::Add: { | |||
412 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); | |||
413 | return getAffineBinaryOpExpr( | |||
414 | expr.getKind(), symbolicDivide(binaryExpr.getLHS(), symbolPos, opKind), | |||
415 | symbolicDivide(binaryExpr.getRHS(), symbolPos, opKind)); | |||
416 | } | |||
417 | // Dividing both operands by the given symbol. | |||
418 | case AffineExprKind::Mod: { | |||
419 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); | |||
420 | return getAffineBinaryOpExpr( | |||
421 | expr.getKind(), | |||
422 | symbolicDivide(binaryExpr.getLHS(), symbolPos, expr.getKind()), | |||
423 | symbolicDivide(binaryExpr.getRHS(), symbolPos, expr.getKind())); | |||
424 | } | |||
425 | // Dividing any of the operand by the given symbol. | |||
426 | case AffineExprKind::Mul: { | |||
427 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); | |||
428 | if (!isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind)) | |||
429 | return binaryExpr.getLHS() * | |||
430 | symbolicDivide(binaryExpr.getRHS(), symbolPos, opKind); | |||
431 | return symbolicDivide(binaryExpr.getLHS(), symbolPos, opKind) * | |||
432 | binaryExpr.getRHS(); | |||
433 | } | |||
434 | // Dividing first operand only by the given symbol. | |||
435 | case AffineExprKind::FloorDiv: | |||
436 | case AffineExprKind::CeilDiv: { | |||
437 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); | |||
438 | return getAffineBinaryOpExpr( | |||
439 | expr.getKind(), | |||
440 | symbolicDivide(binaryExpr.getLHS(), symbolPos, expr.getKind()), | |||
441 | binaryExpr.getRHS()); | |||
442 | } | |||
443 | } | |||
444 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 444); | |||
445 | } | |||
446 | ||||
447 | /// Simplify a semi-affine expression by handling modulo, floordiv, or ceildiv | |||
448 | /// operations when the second operand simplifies to a symbol and the first | |||
449 | /// operand is divisible by that symbol. It can be applied to any semi-affine | |||
450 | /// expression. Returned expression can either be a semi-affine or pure affine | |||
451 | /// expression. | |||
452 | static AffineExpr simplifySemiAffine(AffineExpr expr) { | |||
453 | switch (expr.getKind()) { | |||
454 | case AffineExprKind::Constant: | |||
455 | case AffineExprKind::DimId: | |||
456 | case AffineExprKind::SymbolId: | |||
457 | return expr; | |||
458 | case AffineExprKind::Add: | |||
459 | case AffineExprKind::Mul: { | |||
460 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); | |||
461 | return getAffineBinaryOpExpr(expr.getKind(), | |||
462 | simplifySemiAffine(binaryExpr.getLHS()), | |||
463 | simplifySemiAffine(binaryExpr.getRHS())); | |||
464 | } | |||
465 | // Check if the simplification of the second operand is a symbol, and the | |||
466 | // first operand is divisible by it. If the operation is a modulo, a constant | |||
467 | // zero expression is returned. In the case of floordiv and ceildiv, the | |||
468 | // symbol from the simplification of the second operand divides the first | |||
469 | // operand. Otherwise, simplification is not possible. | |||
470 | case AffineExprKind::FloorDiv: | |||
471 | case AffineExprKind::CeilDiv: | |||
472 | case AffineExprKind::Mod: { | |||
473 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); | |||
474 | AffineExpr sLHS = simplifySemiAffine(binaryExpr.getLHS()); | |||
475 | AffineExpr sRHS = simplifySemiAffine(binaryExpr.getRHS()); | |||
476 | AffineSymbolExpr symbolExpr = | |||
477 | simplifySemiAffine(binaryExpr.getRHS()).dyn_cast<AffineSymbolExpr>(); | |||
478 | if (!symbolExpr) | |||
479 | return getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS); | |||
480 | unsigned symbolPos = symbolExpr.getPosition(); | |||
481 | if (!isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, expr.getKind())) | |||
482 | return getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS); | |||
483 | if (expr.getKind() == AffineExprKind::Mod) | |||
484 | return getAffineConstantExpr(0, expr.getContext()); | |||
485 | return symbolicDivide(sLHS, symbolPos, expr.getKind()); | |||
486 | } | |||
487 | } | |||
488 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 488); | |||
489 | } | |||
490 | ||||
491 | static AffineExpr getAffineDimOrSymbol(AffineExprKind kind, unsigned position, | |||
492 | MLIRContext *context) { | |||
493 | auto assignCtx = [context](AffineDimExprStorage *storage) { | |||
494 | storage->context = context; | |||
495 | }; | |||
496 | ||||
497 | StorageUniquer &uniquer = context->getAffineUniquer(); | |||
498 | return uniquer.get<AffineDimExprStorage>( | |||
499 | assignCtx, static_cast<unsigned>(kind), position); | |||
500 | } | |||
501 | ||||
502 | AffineExpr mlir::getAffineDimExpr(unsigned position, MLIRContext *context) { | |||
503 | return getAffineDimOrSymbol(AffineExprKind::DimId, position, context); | |||
504 | } | |||
505 | ||||
506 | AffineSymbolExpr::AffineSymbolExpr(AffineExpr::ImplType *ptr) | |||
507 | : AffineExpr(ptr) {} | |||
508 | unsigned AffineSymbolExpr::getPosition() const { | |||
509 | return static_cast<ImplType *>(expr)->position; | |||
510 | } | |||
511 | ||||
512 | AffineExpr mlir::getAffineSymbolExpr(unsigned position, MLIRContext *context) { | |||
513 | return getAffineDimOrSymbol(AffineExprKind::SymbolId, position, context); | |||
514 | ; | |||
515 | } | |||
516 | ||||
517 | AffineConstantExpr::AffineConstantExpr(AffineExpr::ImplType *ptr) | |||
518 | : AffineExpr(ptr) {} | |||
519 | int64_t AffineConstantExpr::getValue() const { | |||
520 | return static_cast<ImplType *>(expr)->constant; | |||
521 | } | |||
522 | ||||
523 | bool AffineExpr::operator==(int64_t v) const { | |||
524 | return *this == getAffineConstantExpr(v, getContext()); | |||
525 | } | |||
526 | ||||
527 | AffineExpr mlir::getAffineConstantExpr(int64_t constant, MLIRContext *context) { | |||
528 | auto assignCtx = [context](AffineConstantExprStorage *storage) { | |||
529 | storage->context = context; | |||
530 | }; | |||
531 | ||||
532 | StorageUniquer &uniquer = context->getAffineUniquer(); | |||
533 | return uniquer.get<AffineConstantExprStorage>(assignCtx, constant); | |||
534 | } | |||
535 | ||||
536 | /// Simplify add expression. Return nullptr if it can't be simplified. | |||
537 | static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) { | |||
538 | auto lhsConst = lhs.dyn_cast<AffineConstantExpr>(); | |||
539 | auto rhsConst = rhs.dyn_cast<AffineConstantExpr>(); | |||
540 | // Fold if both LHS, RHS are a constant. | |||
541 | if (lhsConst && rhsConst) | |||
542 | return getAffineConstantExpr(lhsConst.getValue() + rhsConst.getValue(), | |||
543 | lhs.getContext()); | |||
544 | ||||
545 | // Canonicalize so that only the RHS is a constant. (4 + d0 becomes d0 + 4). | |||
546 | // If only one of them is a symbolic expressions, make it the RHS. | |||
547 | if (lhs.isa<AffineConstantExpr>() || | |||
548 | (lhs.isSymbolicOrConstant() && !rhs.isSymbolicOrConstant())) { | |||
549 | return rhs + lhs; | |||
550 | } | |||
551 | ||||
552 | // At this point, if there was a constant, it would be on the right. | |||
553 | ||||
554 | // Addition with a zero is a noop, return the other input. | |||
555 | if (rhsConst) { | |||
556 | if (rhsConst.getValue() == 0) | |||
557 | return lhs; | |||
558 | } | |||
559 | // Fold successive additions like (d0 + 2) + 3 into d0 + 5. | |||
560 | auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>(); | |||
561 | if (lBin && rhsConst && lBin.getKind() == AffineExprKind::Add) { | |||
562 | if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) | |||
563 | return lBin.getLHS() + (lrhs.getValue() + rhsConst.getValue()); | |||
564 | } | |||
565 | ||||
566 | // Detect "c1 * expr + c_2 * expr" as "(c1 + c2) * expr". | |||
567 | // c1 is rRhsConst, c2 is rLhsConst; firstExpr, secondExpr are their | |||
568 | // respective multiplicands. | |||
569 | std::optional<int64_t> rLhsConst, rRhsConst; | |||
570 | AffineExpr firstExpr, secondExpr; | |||
571 | AffineConstantExpr rLhsConstExpr; | |||
572 | auto lBinOpExpr = lhs.dyn_cast<AffineBinaryOpExpr>(); | |||
573 | if (lBinOpExpr && lBinOpExpr.getKind() == AffineExprKind::Mul && | |||
574 | (rLhsConstExpr = lBinOpExpr.getRHS().dyn_cast<AffineConstantExpr>())) { | |||
575 | rLhsConst = rLhsConstExpr.getValue(); | |||
576 | firstExpr = lBinOpExpr.getLHS(); | |||
577 | } else { | |||
578 | rLhsConst = 1; | |||
579 | firstExpr = lhs; | |||
580 | } | |||
581 | ||||
582 | auto rBinOpExpr = rhs.dyn_cast<AffineBinaryOpExpr>(); | |||
583 | AffineConstantExpr rRhsConstExpr; | |||
584 | if (rBinOpExpr && rBinOpExpr.getKind() == AffineExprKind::Mul && | |||
585 | (rRhsConstExpr = rBinOpExpr.getRHS().dyn_cast<AffineConstantExpr>())) { | |||
586 | rRhsConst = rRhsConstExpr.getValue(); | |||
587 | secondExpr = rBinOpExpr.getLHS(); | |||
588 | } else { | |||
589 | rRhsConst = 1; | |||
590 | secondExpr = rhs; | |||
591 | } | |||
592 | ||||
593 | if (rLhsConst && rRhsConst && firstExpr == secondExpr) | |||
594 | return getAffineBinaryOpExpr( | |||
595 | AffineExprKind::Mul, firstExpr, | |||
596 | getAffineConstantExpr(*rLhsConst + *rRhsConst, lhs.getContext())); | |||
597 | ||||
598 | // When doing successive additions, bring constant to the right: turn (d0 + 2) | |||
599 | // + d1 into (d0 + d1) + 2. | |||
600 | if (lBin && lBin.getKind() == AffineExprKind::Add) { | |||
601 | if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) { | |||
602 | return lBin.getLHS() + rhs + lrhs; | |||
603 | } | |||
604 | } | |||
605 | ||||
606 | // Detect and transform "expr - q * (expr floordiv q)" to "expr mod q", where | |||
607 | // q may be a constant or symbolic expression. This leads to a much more | |||
608 | // efficient form when 'c' is a power of two, and in general a more compact | |||
609 | // and readable form. | |||
610 | ||||
611 | // Process '(expr floordiv c) * (-c)'. | |||
612 | if (!rBinOpExpr) | |||
613 | return nullptr; | |||
614 | ||||
615 | auto lrhs = rBinOpExpr.getLHS(); | |||
616 | auto rrhs = rBinOpExpr.getRHS(); | |||
617 | ||||
618 | AffineExpr llrhs, rlrhs; | |||
619 | ||||
620 | // Check if lrhsBinOpExpr is of the form (expr floordiv q) * q, where q is a | |||
621 | // symbolic expression. | |||
622 | auto lrhsBinOpExpr = lrhs.dyn_cast<AffineBinaryOpExpr>(); | |||
623 | // Check rrhsConstOpExpr = -1. | |||
624 | auto rrhsConstOpExpr = rrhs.dyn_cast<AffineConstantExpr>(); | |||
625 | if (rrhsConstOpExpr && rrhsConstOpExpr.getValue() == -1 && lrhsBinOpExpr && | |||
626 | lrhsBinOpExpr.getKind() == AffineExprKind::Mul) { | |||
627 | // Check llrhs = expr floordiv q. | |||
628 | llrhs = lrhsBinOpExpr.getLHS(); | |||
629 | // Check rlrhs = q. | |||
630 | rlrhs = lrhsBinOpExpr.getRHS(); | |||
631 | auto llrhsBinOpExpr = llrhs.dyn_cast<AffineBinaryOpExpr>(); | |||
632 | if (!llrhsBinOpExpr || llrhsBinOpExpr.getKind() != AffineExprKind::FloorDiv) | |||
633 | return nullptr; | |||
634 | if (llrhsBinOpExpr.getRHS() == rlrhs && lhs == llrhsBinOpExpr.getLHS()) | |||
635 | return lhs % rlrhs; | |||
636 | } | |||
637 | ||||
638 | // Process lrhs, which is 'expr floordiv c'. | |||
639 | AffineBinaryOpExpr lrBinOpExpr = lrhs.dyn_cast<AffineBinaryOpExpr>(); | |||
640 | if (!lrBinOpExpr || lrBinOpExpr.getKind() != AffineExprKind::FloorDiv) | |||
641 | return nullptr; | |||
642 | ||||
643 | llrhs = lrBinOpExpr.getLHS(); | |||
644 | rlrhs = lrBinOpExpr.getRHS(); | |||
645 | ||||
646 | if (lhs == llrhs && rlrhs == -rrhs) { | |||
647 | return lhs % rlrhs; | |||
648 | } | |||
649 | return nullptr; | |||
650 | } | |||
651 | ||||
652 | AffineExpr AffineExpr::operator+(int64_t v) const { | |||
653 | return *this + getAffineConstantExpr(v, getContext()); | |||
654 | } | |||
655 | AffineExpr AffineExpr::operator+(AffineExpr other) const { | |||
656 | if (auto simplified = simplifyAdd(*this, other)) | |||
657 | return simplified; | |||
658 | ||||
659 | StorageUniquer &uniquer = getContext()->getAffineUniquer(); | |||
660 | return uniquer.get<AffineBinaryOpExprStorage>( | |||
661 | /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Add), *this, other); | |||
662 | } | |||
663 | ||||
664 | /// Simplify a multiply expression. Return nullptr if it can't be simplified. | |||
665 | static AffineExpr simplifyMul(AffineExpr lhs, AffineExpr rhs) { | |||
666 | auto lhsConst = lhs.dyn_cast<AffineConstantExpr>(); | |||
667 | auto rhsConst = rhs.dyn_cast<AffineConstantExpr>(); | |||
668 | ||||
669 | if (lhsConst && rhsConst) | |||
670 | return getAffineConstantExpr(lhsConst.getValue() * rhsConst.getValue(), | |||
671 | lhs.getContext()); | |||
672 | ||||
673 | assert(lhs.isSymbolicOrConstant() || rhs.isSymbolicOrConstant())(static_cast <bool> (lhs.isSymbolicOrConstant() || rhs. isSymbolicOrConstant()) ? void (0) : __assert_fail ("lhs.isSymbolicOrConstant() || rhs.isSymbolicOrConstant()" , "mlir/lib/IR/AffineExpr.cpp", 673, __extension__ __PRETTY_FUNCTION__ )); | |||
674 | ||||
675 | // Canonicalize the mul expression so that the constant/symbolic term is the | |||
676 | // RHS. If both the lhs and rhs are symbolic, swap them if the lhs is a | |||
677 | // constant. (Note that a constant is trivially symbolic). | |||
678 | if (!rhs.isSymbolicOrConstant() || lhs.isa<AffineConstantExpr>()) { | |||
679 | // At least one of them has to be symbolic. | |||
680 | return rhs * lhs; | |||
681 | } | |||
682 | ||||
683 | // At this point, if there was a constant, it would be on the right. | |||
684 | ||||
685 | // Multiplication with a one is a noop, return the other input. | |||
686 | if (rhsConst) { | |||
687 | if (rhsConst.getValue() == 1) | |||
688 | return lhs; | |||
689 | // Multiplication with zero. | |||
690 | if (rhsConst.getValue() == 0) | |||
691 | return rhsConst; | |||
692 | } | |||
693 | ||||
694 | // Fold successive multiplications: eg: (d0 * 2) * 3 into d0 * 6. | |||
695 | auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>(); | |||
696 | if (lBin && rhsConst && lBin.getKind() == AffineExprKind::Mul) { | |||
697 | if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) | |||
698 | return lBin.getLHS() * (lrhs.getValue() * rhsConst.getValue()); | |||
699 | } | |||
700 | ||||
701 | // When doing successive multiplication, bring constant to the right: turn (d0 | |||
702 | // * 2) * d1 into (d0 * d1) * 2. | |||
703 | if (lBin && lBin.getKind() == AffineExprKind::Mul) { | |||
704 | if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) { | |||
705 | return (lBin.getLHS() * rhs) * lrhs; | |||
706 | } | |||
707 | } | |||
708 | ||||
709 | return nullptr; | |||
710 | } | |||
711 | ||||
712 | AffineExpr AffineExpr::operator*(int64_t v) const { | |||
713 | return *this * getAffineConstantExpr(v, getContext()); | |||
714 | } | |||
715 | AffineExpr AffineExpr::operator*(AffineExpr other) const { | |||
716 | if (auto simplified = simplifyMul(*this, other)) | |||
717 | return simplified; | |||
718 | ||||
719 | StorageUniquer &uniquer = getContext()->getAffineUniquer(); | |||
720 | return uniquer.get<AffineBinaryOpExprStorage>( | |||
721 | /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mul), *this, other); | |||
722 | } | |||
723 | ||||
724 | // Unary minus, delegate to operator*. | |||
725 | AffineExpr AffineExpr::operator-() const { | |||
726 | return *this * getAffineConstantExpr(-1, getContext()); | |||
727 | } | |||
728 | ||||
729 | // Delegate to operator+. | |||
730 | AffineExpr AffineExpr::operator-(int64_t v) const { return *this + (-v); } | |||
731 | AffineExpr AffineExpr::operator-(AffineExpr other) const { | |||
732 | return *this + (-other); | |||
733 | } | |||
734 | ||||
735 | static AffineExpr simplifyFloorDiv(AffineExpr lhs, AffineExpr rhs) { | |||
736 | auto lhsConst = lhs.dyn_cast<AffineConstantExpr>(); | |||
737 | auto rhsConst = rhs.dyn_cast<AffineConstantExpr>(); | |||
738 | ||||
739 | // mlir floordiv by zero or negative numbers is undefined and preserved as is. | |||
740 | if (!rhsConst || rhsConst.getValue() < 1) | |||
741 | return nullptr; | |||
742 | ||||
743 | if (lhsConst) | |||
744 | return getAffineConstantExpr( | |||
745 | floorDiv(lhsConst.getValue(), rhsConst.getValue()), lhs.getContext()); | |||
746 | ||||
747 | // Fold floordiv of a multiply with a constant that is a multiple of the | |||
748 | // divisor. Eg: (i * 128) floordiv 64 = i * 2. | |||
749 | if (rhsConst == 1) | |||
750 | return lhs; | |||
751 | ||||
752 | // Simplify (expr * const) floordiv divConst when expr is known to be a | |||
753 | // multiple of divConst. | |||
754 | auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>(); | |||
755 | if (lBin && lBin.getKind() == AffineExprKind::Mul) { | |||
756 | if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) { | |||
757 | // rhsConst is known to be a positive constant. | |||
758 | if (lrhs.getValue() % rhsConst.getValue() == 0) | |||
759 | return lBin.getLHS() * (lrhs.getValue() / rhsConst.getValue()); | |||
760 | } | |||
761 | } | |||
762 | ||||
763 | // Simplify (expr1 + expr2) floordiv divConst when either expr1 or expr2 is | |||
764 | // known to be a multiple of divConst. | |||
765 | if (lBin && lBin.getKind() == AffineExprKind::Add) { | |||
766 | int64_t llhsDiv = lBin.getLHS().getLargestKnownDivisor(); | |||
767 | int64_t lrhsDiv = lBin.getRHS().getLargestKnownDivisor(); | |||
768 | // rhsConst is known to be a positive constant. | |||
769 | if (llhsDiv % rhsConst.getValue() == 0 || | |||
770 | lrhsDiv % rhsConst.getValue() == 0) | |||
771 | return lBin.getLHS().floorDiv(rhsConst.getValue()) + | |||
772 | lBin.getRHS().floorDiv(rhsConst.getValue()); | |||
773 | } | |||
774 | ||||
775 | return nullptr; | |||
776 | } | |||
777 | ||||
778 | AffineExpr AffineExpr::floorDiv(uint64_t v) const { | |||
779 | return floorDiv(getAffineConstantExpr(v, getContext())); | |||
780 | } | |||
781 | AffineExpr AffineExpr::floorDiv(AffineExpr other) const { | |||
782 | if (auto simplified = simplifyFloorDiv(*this, other)) | |||
783 | return simplified; | |||
784 | ||||
785 | StorageUniquer &uniquer = getContext()->getAffineUniquer(); | |||
786 | return uniquer.get<AffineBinaryOpExprStorage>( | |||
787 | /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::FloorDiv), *this, | |||
788 | other); | |||
789 | } | |||
790 | ||||
791 | static AffineExpr simplifyCeilDiv(AffineExpr lhs, AffineExpr rhs) { | |||
792 | auto lhsConst = lhs.dyn_cast<AffineConstantExpr>(); | |||
793 | auto rhsConst = rhs.dyn_cast<AffineConstantExpr>(); | |||
794 | ||||
795 | if (!rhsConst || rhsConst.getValue() < 1) | |||
796 | return nullptr; | |||
797 | ||||
798 | if (lhsConst) | |||
799 | return getAffineConstantExpr( | |||
800 | ceilDiv(lhsConst.getValue(), rhsConst.getValue()), lhs.getContext()); | |||
801 | ||||
802 | // Fold ceildiv of a multiply with a constant that is a multiple of the | |||
803 | // divisor. Eg: (i * 128) ceildiv 64 = i * 2. | |||
804 | if (rhsConst.getValue() == 1) | |||
805 | return lhs; | |||
806 | ||||
807 | // Simplify (expr * const) ceildiv divConst when const is known to be a | |||
808 | // multiple of divConst. | |||
809 | auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>(); | |||
810 | if (lBin && lBin.getKind() == AffineExprKind::Mul) { | |||
811 | if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) { | |||
812 | // rhsConst is known to be a positive constant. | |||
813 | if (lrhs.getValue() % rhsConst.getValue() == 0) | |||
814 | return lBin.getLHS() * (lrhs.getValue() / rhsConst.getValue()); | |||
815 | } | |||
816 | } | |||
817 | ||||
818 | return nullptr; | |||
819 | } | |||
820 | ||||
821 | AffineExpr AffineExpr::ceilDiv(uint64_t v) const { | |||
822 | return ceilDiv(getAffineConstantExpr(v, getContext())); | |||
823 | } | |||
824 | AffineExpr AffineExpr::ceilDiv(AffineExpr other) const { | |||
825 | if (auto simplified = simplifyCeilDiv(*this, other)) | |||
826 | return simplified; | |||
827 | ||||
828 | StorageUniquer &uniquer = getContext()->getAffineUniquer(); | |||
829 | return uniquer.get<AffineBinaryOpExprStorage>( | |||
830 | /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::CeilDiv), *this, | |||
831 | other); | |||
832 | } | |||
833 | ||||
834 | static AffineExpr simplifyMod(AffineExpr lhs, AffineExpr rhs) { | |||
835 | auto lhsConst = lhs.dyn_cast<AffineConstantExpr>(); | |||
836 | auto rhsConst = rhs.dyn_cast<AffineConstantExpr>(); | |||
837 | ||||
838 | // mod w.r.t zero or negative numbers is undefined and preserved as is. | |||
839 | if (!rhsConst || rhsConst.getValue() < 1) | |||
840 | return nullptr; | |||
841 | ||||
842 | if (lhsConst) | |||
843 | return getAffineConstantExpr(mod(lhsConst.getValue(), rhsConst.getValue()), | |||
844 | lhs.getContext()); | |||
845 | ||||
846 | // Fold modulo of an expression that is known to be a multiple of a constant | |||
847 | // to zero if that constant is a multiple of the modulo factor. Eg: (i * 128) | |||
848 | // mod 64 is folded to 0, and less trivially, (i*(j*4*(k*32))) mod 128 = 0. | |||
849 | if (lhs.getLargestKnownDivisor() % rhsConst.getValue() == 0) | |||
850 | return getAffineConstantExpr(0, lhs.getContext()); | |||
851 | ||||
852 | // Simplify (expr1 + expr2) mod divConst when either expr1 or expr2 is | |||
853 | // known to be a multiple of divConst. | |||
854 | auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>(); | |||
855 | if (lBin && lBin.getKind() == AffineExprKind::Add) { | |||
856 | int64_t llhsDiv = lBin.getLHS().getLargestKnownDivisor(); | |||
857 | int64_t lrhsDiv = lBin.getRHS().getLargestKnownDivisor(); | |||
858 | // rhsConst is known to be a positive constant. | |||
859 | if (llhsDiv % rhsConst.getValue() == 0) | |||
860 | return lBin.getRHS() % rhsConst.getValue(); | |||
861 | if (lrhsDiv % rhsConst.getValue() == 0) | |||
862 | return lBin.getLHS() % rhsConst.getValue(); | |||
863 | } | |||
864 | ||||
865 | // Simplify (e % a) % b to e % b when b evenly divides a | |||
866 | if (lBin && lBin.getKind() == AffineExprKind::Mod) { | |||
867 | auto intermediate = lBin.getRHS().dyn_cast<AffineConstantExpr>(); | |||
868 | if (intermediate && intermediate.getValue() >= 1 && | |||
869 | mod(intermediate.getValue(), rhsConst.getValue()) == 0) { | |||
870 | return lBin.getLHS() % rhsConst.getValue(); | |||
871 | } | |||
872 | } | |||
873 | ||||
874 | return nullptr; | |||
875 | } | |||
876 | ||||
877 | AffineExpr AffineExpr::operator%(uint64_t v) const { | |||
878 | return *this % getAffineConstantExpr(v, getContext()); | |||
879 | } | |||
880 | AffineExpr AffineExpr::operator%(AffineExpr other) const { | |||
881 | if (auto simplified = simplifyMod(*this, other)) | |||
882 | return simplified; | |||
883 | ||||
884 | StorageUniquer &uniquer = getContext()->getAffineUniquer(); | |||
885 | return uniquer.get<AffineBinaryOpExprStorage>( | |||
886 | /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mod), *this, other); | |||
887 | } | |||
888 | ||||
889 | AffineExpr AffineExpr::compose(AffineMap map) const { | |||
890 | SmallVector<AffineExpr, 8> dimReplacements(map.getResults().begin(), | |||
891 | map.getResults().end()); | |||
892 | return replaceDimsAndSymbols(dimReplacements, {}); | |||
893 | } | |||
894 | raw_ostream &mlir::operator<<(raw_ostream &os, AffineExpr expr) { | |||
895 | expr.print(os); | |||
896 | return os; | |||
897 | } | |||
898 | ||||
899 | /// Constructs an affine expression from a flat ArrayRef. If there are local | |||
900 | /// identifiers (neither dimensional nor symbolic) that appear in the sum of | |||
901 | /// products expression, `localExprs` is expected to have the AffineExpr | |||
902 | /// for it, and is substituted into. The ArrayRef `flatExprs` is expected to be | |||
903 | /// in the format [dims, symbols, locals, constant term]. | |||
904 | AffineExpr mlir::getAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs, | |||
905 | unsigned numDims, | |||
906 | unsigned numSymbols, | |||
907 | ArrayRef<AffineExpr> localExprs, | |||
908 | MLIRContext *context) { | |||
909 | // Assert expected numLocals = flatExprs.size() - numDims - numSymbols - 1. | |||
910 | assert(flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() &&(static_cast <bool> (flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && "unexpected number of local expressions" ) ? void (0) : __assert_fail ("flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && \"unexpected number of local expressions\"" , "mlir/lib/IR/AffineExpr.cpp", 911, __extension__ __PRETTY_FUNCTION__ )) | |||
911 | "unexpected number of local expressions")(static_cast <bool> (flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && "unexpected number of local expressions" ) ? void (0) : __assert_fail ("flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && \"unexpected number of local expressions\"" , "mlir/lib/IR/AffineExpr.cpp", 911, __extension__ __PRETTY_FUNCTION__ )); | |||
912 | ||||
913 | auto expr = getAffineConstantExpr(0, context); | |||
914 | // Dimensions and symbols. | |||
915 | for (unsigned j = 0; j < numDims + numSymbols; j++) { | |||
916 | if (flatExprs[j] == 0) | |||
917 | continue; | |||
918 | auto id = j < numDims ? getAffineDimExpr(j, context) | |||
919 | : getAffineSymbolExpr(j - numDims, context); | |||
920 | expr = expr + id * flatExprs[j]; | |||
921 | } | |||
922 | ||||
923 | // Local identifiers. | |||
924 | for (unsigned j = numDims + numSymbols, e = flatExprs.size() - 1; j < e; | |||
925 | j++) { | |||
926 | if (flatExprs[j] == 0) | |||
927 | continue; | |||
928 | auto term = localExprs[j - numDims - numSymbols] * flatExprs[j]; | |||
929 | expr = expr + term; | |||
930 | } | |||
931 | ||||
932 | // Constant term. | |||
933 | int64_t constTerm = flatExprs[flatExprs.size() - 1]; | |||
934 | if (constTerm != 0) | |||
935 | expr = expr + constTerm; | |||
936 | return expr; | |||
937 | } | |||
938 | ||||
939 | /// Constructs a semi-affine expression from a flat ArrayRef. If there are | |||
940 | /// local identifiers (neither dimensional nor symbolic) that appear in the sum | |||
941 | /// of products expression, `localExprs` is expected to have the AffineExprs for | |||
942 | /// it, and is substituted into. The ArrayRef `flatExprs` is expected to be in | |||
943 | /// the format [dims, symbols, locals, constant term]. The semi-affine | |||
944 | /// expression is constructed in the sorted order of dimension and symbol | |||
945 | /// position numbers. Note: local expressions/ids are used for mod, div as well | |||
946 | /// as symbolic RHS terms for terms that are not pure affine. | |||
947 | static AffineExpr getSemiAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs, | |||
948 | unsigned numDims, | |||
949 | unsigned numSymbols, | |||
950 | ArrayRef<AffineExpr> localExprs, | |||
951 | MLIRContext *context) { | |||
952 | assert(!flatExprs.empty() && "flatExprs cannot be empty")(static_cast <bool> (!flatExprs.empty() && "flatExprs cannot be empty" ) ? void (0) : __assert_fail ("!flatExprs.empty() && \"flatExprs cannot be empty\"" , "mlir/lib/IR/AffineExpr.cpp", 952, __extension__ __PRETTY_FUNCTION__ )); | |||
953 | ||||
954 | // Assert expected numLocals = flatExprs.size() - numDims - numSymbols - 1. | |||
955 | assert(flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() &&(static_cast <bool> (flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && "unexpected number of local expressions" ) ? void (0) : __assert_fail ("flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && \"unexpected number of local expressions\"" , "mlir/lib/IR/AffineExpr.cpp", 956, __extension__ __PRETTY_FUNCTION__ )) | |||
956 | "unexpected number of local expressions")(static_cast <bool> (flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && "unexpected number of local expressions" ) ? void (0) : __assert_fail ("flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && \"unexpected number of local expressions\"" , "mlir/lib/IR/AffineExpr.cpp", 956, __extension__ __PRETTY_FUNCTION__ )); | |||
957 | ||||
958 | AffineExpr expr = getAffineConstantExpr(0, context); | |||
959 | ||||
960 | // We design indices as a pair which help us present the semi-affine map as | |||
961 | // sum of product where terms are sorted based on dimension or symbol | |||
962 | // position: <keyA, keyB> for expressions of the form dimension * symbol, | |||
963 | // where keyA is the position number of the dimension and keyB is the | |||
964 | // position number of the symbol. For dimensional expressions we set the index | |||
965 | // as (position number of the dimension, -1), as we want dimensional | |||
966 | // expressions to appear before symbolic and product of dimensional and | |||
967 | // symbolic expressions having the dimension with the same position number. | |||
968 | // For symbolic expression set the index as (position number of the symbol, | |||
969 | // maximum of last dimension and symbol position) number. For example, we want | |||
970 | // the expression we are constructing to look something like: d0 + d0 * s0 + | |||
971 | // s0 + d1*s1 + s1. | |||
972 | ||||
973 | // Stores the affine expression corresponding to a given index. | |||
974 | DenseMap<std::pair<unsigned, signed>, AffineExpr> indexToExprMap; | |||
975 | // Stores the constant coefficient value corresponding to a given | |||
976 | // dimension, symbol or a non-pure affine expression stored in `localExprs`. | |||
977 | DenseMap<std::pair<unsigned, signed>, int64_t> coefficients; | |||
978 | // Stores the indices as defined above, and later sorted to produce | |||
979 | // the semi-affine expression in the desired form. | |||
980 | SmallVector<std::pair<unsigned, signed>, 8> indices; | |||
981 | ||||
982 | // Example: expression = d0 + d0 * s0 + 2 * s0. | |||
983 | // indices = [{0,-1}, {0, 0}, {0, 1}] | |||
984 | // coefficients = [{{0, -1}, 1}, {{0, 0}, 1}, {{0, 1}, 2}] | |||
985 | // indexToExprMap = [{{0, -1}, d0}, {{0, 0}, d0 * s0}, {{0, 1}, s0}] | |||
986 | ||||
987 | // Adds entries to `indexToExprMap`, `coefficients` and `indices`. | |||
988 | auto addEntry = [&](std::pair<unsigned, signed> index, int64_t coefficient, | |||
989 | AffineExpr expr) { | |||
990 | assert(!llvm::is_contained(indices, index) &&(static_cast <bool> (!llvm::is_contained(indices, index ) && "Key is already present in indices vector and overwriting will " "happen in `indexToExprMap` and `coefficients`!") ? void (0) : __assert_fail ("!llvm::is_contained(indices, index) && \"Key is already present in indices vector and overwriting will \" \"happen in `indexToExprMap` and `coefficients`!\"" , "mlir/lib/IR/AffineExpr.cpp", 992, __extension__ __PRETTY_FUNCTION__ )) | |||
991 | "Key is already present in indices vector and overwriting will "(static_cast <bool> (!llvm::is_contained(indices, index ) && "Key is already present in indices vector and overwriting will " "happen in `indexToExprMap` and `coefficients`!") ? void (0) : __assert_fail ("!llvm::is_contained(indices, index) && \"Key is already present in indices vector and overwriting will \" \"happen in `indexToExprMap` and `coefficients`!\"" , "mlir/lib/IR/AffineExpr.cpp", 992, __extension__ __PRETTY_FUNCTION__ )) | |||
992 | "happen in `indexToExprMap` and `coefficients`!")(static_cast <bool> (!llvm::is_contained(indices, index ) && "Key is already present in indices vector and overwriting will " "happen in `indexToExprMap` and `coefficients`!") ? void (0) : __assert_fail ("!llvm::is_contained(indices, index) && \"Key is already present in indices vector and overwriting will \" \"happen in `indexToExprMap` and `coefficients`!\"" , "mlir/lib/IR/AffineExpr.cpp", 992, __extension__ __PRETTY_FUNCTION__ )); | |||
993 | ||||
994 | indices.push_back(index); | |||
995 | coefficients.insert({index, coefficient}); | |||
996 | indexToExprMap.insert({index, expr}); | |||
997 | }; | |||
998 | ||||
999 | // Design indices for dimensional or symbolic terms, and store the indices, | |||
1000 | // constant coefficient corresponding to the indices in `coefficients` map, | |||
1001 | // and affine expression corresponding to indices in `indexToExprMap` map. | |||
1002 | ||||
1003 | // Ensure we do not have duplicate keys in `indexToExpr` map. | |||
1004 | unsigned offsetSym = 0; | |||
1005 | signed offsetDim = -1; | |||
1006 | for (unsigned j = numDims; j < numDims + numSymbols; ++j) { | |||
1007 | if (flatExprs[j] == 0) | |||
1008 | continue; | |||
1009 | // For symbolic expression set the index as <position number | |||
1010 | // of the symbol, max(dimCount, symCount)> number, | |||
1011 | // as we want symbolic expressions with the same positional number to | |||
1012 | // appear after dimensional expressions having the same positional number. | |||
1013 | std::pair<unsigned, signed> indexEntry( | |||
1014 | j - numDims, std::max(numDims, numSymbols) + offsetSym++); | |||
1015 | addEntry(indexEntry, flatExprs[j], | |||
1016 | getAffineSymbolExpr(j - numDims, context)); | |||
1017 | } | |||
1018 | ||||
1019 | // Denotes semi-affine product, modulo or division terms, which has been added | |||
1020 | // to the `indexToExpr` map. | |||
1021 | SmallVector<bool, 4> addedToMap(flatExprs.size() - numDims - numSymbols - 1, | |||
1022 | false); | |||
1023 | unsigned lhsPos, rhsPos; | |||
1024 | // Construct indices for product terms involving dimension, symbol or constant | |||
1025 | // as lhs/rhs, and store the indices, constant coefficient corresponding to | |||
1026 | // the indices in `coefficients` map, and affine expression corresponding to | |||
1027 | // in indices in `indexToExprMap` map. | |||
1028 | for (const auto &it : llvm::enumerate(localExprs)) { | |||
1029 | AffineExpr expr = it.value(); | |||
1030 | if (flatExprs[numDims + numSymbols + it.index()] == 0) | |||
1031 | continue; | |||
1032 | AffineExpr lhs = expr.cast<AffineBinaryOpExpr>().getLHS(); | |||
1033 | AffineExpr rhs = expr.cast<AffineBinaryOpExpr>().getRHS(); | |||
1034 | if (!((lhs.isa<AffineDimExpr>() || lhs.isa<AffineSymbolExpr>()) && | |||
1035 | (rhs.isa<AffineDimExpr>() || rhs.isa<AffineSymbolExpr>() || | |||
1036 | rhs.isa<AffineConstantExpr>()))) { | |||
1037 | continue; | |||
1038 | } | |||
1039 | if (rhs.isa<AffineConstantExpr>()) { | |||
1040 | // For product/modulo/division expressions, when rhs of modulo/division | |||
1041 | // expression is constant, we put 0 in place of keyB, because we want | |||
1042 | // them to appear earlier in the semi-affine expression we are | |||
1043 | // constructing. When rhs is constant, we place 0 in place of keyB. | |||
1044 | if (lhs.isa<AffineDimExpr>()) { | |||
1045 | lhsPos = lhs.cast<AffineDimExpr>().getPosition(); | |||
1046 | std::pair<unsigned, signed> indexEntry(lhsPos, offsetDim--); | |||
1047 | addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], | |||
1048 | expr); | |||
1049 | } else { | |||
1050 | lhsPos = lhs.cast<AffineSymbolExpr>().getPosition(); | |||
1051 | std::pair<unsigned, signed> indexEntry( | |||
1052 | lhsPos, std::max(numDims, numSymbols) + offsetSym++); | |||
1053 | addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], | |||
1054 | expr); | |||
1055 | } | |||
1056 | } else if (lhs.isa<AffineDimExpr>()) { | |||
1057 | // For product/modulo/division expressions having lhs as dimension and rhs | |||
1058 | // as symbol, we order the terms in the semi-affine expression based on | |||
1059 | // the pair: <keyA, keyB> for expressions of the form dimension * symbol, | |||
1060 | // where keyA is the position number of the dimension and keyB is the | |||
1061 | // position number of the symbol. | |||
1062 | lhsPos = lhs.cast<AffineDimExpr>().getPosition(); | |||
1063 | rhsPos = rhs.cast<AffineSymbolExpr>().getPosition(); | |||
1064 | std::pair<unsigned, signed> indexEntry(lhsPos, rhsPos); | |||
1065 | addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], expr); | |||
1066 | } else { | |||
1067 | // For product/modulo/division expressions having both lhs and rhs as | |||
1068 | // symbol, we design indices as a pair: <keyA, keyB> for expressions | |||
1069 | // of the form dimension * symbol, where keyA is the position number of | |||
1070 | // the dimension and keyB is the position number of the symbol. | |||
1071 | lhsPos = lhs.cast<AffineSymbolExpr>().getPosition(); | |||
1072 | rhsPos = rhs.cast<AffineSymbolExpr>().getPosition(); | |||
1073 | std::pair<unsigned, signed> indexEntry( | |||
1074 | lhsPos, std::max(numDims, numSymbols) + offsetSym++); | |||
1075 | addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], expr); | |||
1076 | } | |||
1077 | addedToMap[it.index()] = true; | |||
1078 | } | |||
1079 | ||||
1080 | for (unsigned j = 0; j < numDims; ++j) { | |||
1081 | if (flatExprs[j] == 0) | |||
1082 | continue; | |||
1083 | // For dimensional expressions we set the index as <position number of the | |||
1084 | // dimension, 0>, as we want dimensional expressions to appear before | |||
1085 | // symbolic ones and products of dimensional and symbolic expressions | |||
1086 | // having the dimension with the same position number. | |||
1087 | std::pair<unsigned, signed> indexEntry(j, offsetDim--); | |||
1088 | addEntry(indexEntry, flatExprs[j], getAffineDimExpr(j, context)); | |||
1089 | } | |||
1090 | ||||
1091 | // Constructing the simplified semi-affine sum of product/division/mod | |||
1092 | // expression from the flattened form in the desired sorted order of indices | |||
1093 | // of the various individual product/division/mod expressions. | |||
1094 | llvm::sort(indices); | |||
1095 | for (const std::pair<unsigned, unsigned> index : indices) { | |||
1096 | assert(indexToExprMap.lookup(index) &&(static_cast <bool> (indexToExprMap.lookup(index) && "cannot find key in `indexToExprMap` map") ? void (0) : __assert_fail ("indexToExprMap.lookup(index) && \"cannot find key in `indexToExprMap` map\"" , "mlir/lib/IR/AffineExpr.cpp", 1097, __extension__ __PRETTY_FUNCTION__ )) | |||
1097 | "cannot find key in `indexToExprMap` map")(static_cast <bool> (indexToExprMap.lookup(index) && "cannot find key in `indexToExprMap` map") ? void (0) : __assert_fail ("indexToExprMap.lookup(index) && \"cannot find key in `indexToExprMap` map\"" , "mlir/lib/IR/AffineExpr.cpp", 1097, __extension__ __PRETTY_FUNCTION__ )); | |||
1098 | expr = expr + indexToExprMap.lookup(index) * coefficients.lookup(index); | |||
1099 | } | |||
1100 | ||||
1101 | // Local identifiers. | |||
1102 | for (unsigned j = numDims + numSymbols, e = flatExprs.size() - 1; j < e; | |||
1103 | j++) { | |||
1104 | // If the coefficient of the local expression is 0, continue as we need not | |||
1105 | // add it in out final expression. | |||
1106 | if (flatExprs[j] == 0 || addedToMap[j - numDims - numSymbols]) | |||
1107 | continue; | |||
1108 | auto term = localExprs[j - numDims - numSymbols] * flatExprs[j]; | |||
1109 | expr = expr + term; | |||
1110 | } | |||
1111 | ||||
1112 | // Constant term. | |||
1113 | int64_t constTerm = flatExprs.back(); | |||
1114 | if (constTerm != 0) | |||
1115 | expr = expr + constTerm; | |||
1116 | return expr; | |||
1117 | } | |||
1118 | ||||
1119 | SimpleAffineExprFlattener::SimpleAffineExprFlattener(unsigned numDims, | |||
1120 | unsigned numSymbols) | |||
1121 | : numDims(numDims), numSymbols(numSymbols), numLocals(0) { | |||
1122 | operandExprStack.reserve(8); | |||
1123 | } | |||
1124 | ||||
1125 | // In pure affine t = expr * c, we multiply each coefficient of lhs with c. | |||
1126 | // | |||
1127 | // In case of semi affine multiplication expressions, t = expr * symbolic_expr, | |||
1128 | // introduce a local variable p (= expr * symbolic_expr), and the affine | |||
1129 | // expression expr * symbolic_expr is added to `localExprs`. | |||
1130 | void SimpleAffineExprFlattener::visitMulExpr(AffineBinaryOpExpr expr) { | |||
1131 | assert(operandExprStack.size() >= 2)(static_cast <bool> (operandExprStack.size() >= 2) ? void (0) : __assert_fail ("operandExprStack.size() >= 2", "mlir/lib/IR/AffineExpr.cpp", 1131, __extension__ __PRETTY_FUNCTION__ )); | |||
1132 | SmallVector<int64_t, 8> rhs = operandExprStack.back(); | |||
1133 | operandExprStack.pop_back(); | |||
1134 | SmallVector<int64_t, 8> &lhs = operandExprStack.back(); | |||
1135 | ||||
1136 | // Flatten semi-affine multiplication expressions by introducing a local | |||
1137 | // variable in place of the product; the affine expression | |||
1138 | // corresponding to the quantifier is added to `localExprs`. | |||
1139 | if (!expr.getRHS().isa<AffineConstantExpr>()) { | |||
1140 | MLIRContext *context = expr.getContext(); | |||
1141 | AffineExpr a = getAffineExprFromFlatForm(lhs, numDims, numSymbols, | |||
1142 | localExprs, context); | |||
1143 | AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols, | |||
1144 | localExprs, context); | |||
1145 | addLocalVariableSemiAffine(a * b, lhs, lhs.size()); | |||
1146 | return; | |||
1147 | } | |||
1148 | ||||
1149 | // Get the RHS constant. | |||
1150 | auto rhsConst = rhs[getConstantIndex()]; | |||
1151 | for (unsigned i = 0, e = lhs.size(); i < e; i++) { | |||
1152 | lhs[i] *= rhsConst; | |||
1153 | } | |||
1154 | } | |||
1155 | ||||
1156 | void SimpleAffineExprFlattener::visitAddExpr(AffineBinaryOpExpr expr) { | |||
1157 | assert(operandExprStack.size() >= 2)(static_cast <bool> (operandExprStack.size() >= 2) ? void (0) : __assert_fail ("operandExprStack.size() >= 2", "mlir/lib/IR/AffineExpr.cpp", 1157, __extension__ __PRETTY_FUNCTION__ )); | |||
1158 | const auto &rhs = operandExprStack.back(); | |||
1159 | auto &lhs = operandExprStack[operandExprStack.size() - 2]; | |||
1160 | assert(lhs.size() == rhs.size())(static_cast <bool> (lhs.size() == rhs.size()) ? void ( 0) : __assert_fail ("lhs.size() == rhs.size()", "mlir/lib/IR/AffineExpr.cpp" , 1160, __extension__ __PRETTY_FUNCTION__)); | |||
1161 | // Update the LHS in place. | |||
1162 | for (unsigned i = 0, e = rhs.size(); i < e; i++) { | |||
1163 | lhs[i] += rhs[i]; | |||
1164 | } | |||
1165 | // Pop off the RHS. | |||
1166 | operandExprStack.pop_back(); | |||
1167 | } | |||
1168 | ||||
1169 | // | |||
1170 | // t = expr mod c <=> t = expr - c*q and c*q <= expr <= c*q + c - 1 | |||
1171 | // | |||
1172 | // A mod expression "expr mod c" is thus flattened by introducing a new local | |||
1173 | // variable q (= expr floordiv c), such that expr mod c is replaced with | |||
1174 | // 'expr - c * q' and c * q <= expr <= c * q + c - 1 are added to localVarCst. | |||
1175 | // | |||
1176 | // In case of semi-affine modulo expressions, t = expr mod symbolic_expr, | |||
1177 | // introduce a local variable m (= expr mod symbolic_expr), and the affine | |||
1178 | // expression expr mod symbolic_expr is added to `localExprs`. | |||
1179 | void SimpleAffineExprFlattener::visitModExpr(AffineBinaryOpExpr expr) { | |||
1180 | assert(operandExprStack.size() >= 2)(static_cast <bool> (operandExprStack.size() >= 2) ? void (0) : __assert_fail ("operandExprStack.size() >= 2", "mlir/lib/IR/AffineExpr.cpp", 1180, __extension__ __PRETTY_FUNCTION__ )); | |||
1181 | ||||
1182 | SmallVector<int64_t, 8> rhs = operandExprStack.back(); | |||
1183 | operandExprStack.pop_back(); | |||
1184 | SmallVector<int64_t, 8> &lhs = operandExprStack.back(); | |||
1185 | MLIRContext *context = expr.getContext(); | |||
1186 | ||||
1187 | // Flatten semi affine modulo expressions by introducing a local | |||
1188 | // variable in place of the modulo value, and the affine expression | |||
1189 | // corresponding to the quantifier is added to `localExprs`. | |||
1190 | if (!expr.getRHS().isa<AffineConstantExpr>()) { | |||
1191 | AffineExpr dividendExpr = getAffineExprFromFlatForm( | |||
1192 | lhs, numDims, numSymbols, localExprs, context); | |||
1193 | AffineExpr divisorExpr = getAffineExprFromFlatForm(rhs, numDims, numSymbols, | |||
1194 | localExprs, context); | |||
1195 | AffineExpr modExpr = dividendExpr % divisorExpr; | |||
1196 | addLocalVariableSemiAffine(modExpr, lhs, lhs.size()); | |||
1197 | return; | |||
1198 | } | |||
1199 | ||||
1200 | int64_t rhsConst = rhs[getConstantIndex()]; | |||
1201 | // TODO: handle modulo by zero case when this issue is fixed | |||
1202 | // at the other places in the IR. | |||
1203 | assert(rhsConst > 0 && "RHS constant has to be positive")(static_cast <bool> (rhsConst > 0 && "RHS constant has to be positive" ) ? void (0) : __assert_fail ("rhsConst > 0 && \"RHS constant has to be positive\"" , "mlir/lib/IR/AffineExpr.cpp", 1203, __extension__ __PRETTY_FUNCTION__ )); | |||
1204 | ||||
1205 | // Check if the LHS expression is a multiple of modulo factor. | |||
1206 | unsigned i, e; | |||
1207 | for (i = 0, e = lhs.size(); i < e; i++) | |||
1208 | if (lhs[i] % rhsConst != 0) | |||
1209 | break; | |||
1210 | // If yes, modulo expression here simplifies to zero. | |||
1211 | if (i == lhs.size()) { | |||
1212 | std::fill(lhs.begin(), lhs.end(), 0); | |||
1213 | return; | |||
1214 | } | |||
1215 | ||||
1216 | // Add a local variable for the quotient, i.e., expr % c is replaced by | |||
1217 | // (expr - q * c) where q = expr floordiv c. Do this while canceling out | |||
1218 | // the GCD of expr and c. | |||
1219 | SmallVector<int64_t, 8> floorDividend(lhs); | |||
1220 | uint64_t gcd = rhsConst; | |||
1221 | for (unsigned i = 0, e = lhs.size(); i < e; i++) | |||
1222 | gcd = std::gcd(gcd, (uint64_t)std::abs(lhs[i])); | |||
1223 | // Simplify the numerator and the denominator. | |||
1224 | if (gcd != 1) { | |||
1225 | for (unsigned i = 0, e = floorDividend.size(); i < e; i++) | |||
1226 | floorDividend[i] = floorDividend[i] / static_cast<int64_t>(gcd); | |||
1227 | } | |||
1228 | int64_t floorDivisor = rhsConst / static_cast<int64_t>(gcd); | |||
1229 | ||||
1230 | // Construct the AffineExpr form of the floordiv to store in localExprs. | |||
1231 | ||||
1232 | AffineExpr dividendExpr = getAffineExprFromFlatForm( | |||
1233 | floorDividend, numDims, numSymbols, localExprs, context); | |||
1234 | AffineExpr divisorExpr = getAffineConstantExpr(floorDivisor, context); | |||
1235 | AffineExpr floorDivExpr = dividendExpr.floorDiv(divisorExpr); | |||
1236 | int loc; | |||
1237 | if ((loc = findLocalId(floorDivExpr)) == -1) { | |||
1238 | addLocalFloorDivId(floorDividend, floorDivisor, floorDivExpr); | |||
1239 | // Set result at top of stack to "lhs - rhsConst * q". | |||
1240 | lhs[getLocalVarStartIndex() + numLocals - 1] = -rhsConst; | |||
1241 | } else { | |||
1242 | // Reuse the existing local id. | |||
1243 | lhs[getLocalVarStartIndex() + loc] = -rhsConst; | |||
1244 | } | |||
1245 | } | |||
1246 | ||||
1247 | void SimpleAffineExprFlattener::visitCeilDivExpr(AffineBinaryOpExpr expr) { | |||
1248 | visitDivExpr(expr, /*isCeil=*/true); | |||
1249 | } | |||
1250 | void SimpleAffineExprFlattener::visitFloorDivExpr(AffineBinaryOpExpr expr) { | |||
1251 | visitDivExpr(expr, /*isCeil=*/false); | |||
1252 | } | |||
1253 | ||||
1254 | void SimpleAffineExprFlattener::visitDimExpr(AffineDimExpr expr) { | |||
1255 | operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0)); | |||
1256 | auto &eq = operandExprStack.back(); | |||
1257 | assert(expr.getPosition() < numDims && "Inconsistent number of dims")(static_cast <bool> (expr.getPosition() < numDims && "Inconsistent number of dims") ? void (0) : __assert_fail ("expr.getPosition() < numDims && \"Inconsistent number of dims\"" , "mlir/lib/IR/AffineExpr.cpp", 1257, __extension__ __PRETTY_FUNCTION__ )); | |||
1258 | eq[getDimStartIndex() + expr.getPosition()] = 1; | |||
1259 | } | |||
1260 | ||||
1261 | void SimpleAffineExprFlattener::visitSymbolExpr(AffineSymbolExpr expr) { | |||
1262 | operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0)); | |||
1263 | auto &eq = operandExprStack.back(); | |||
1264 | assert(expr.getPosition() < numSymbols && "inconsistent number of symbols")(static_cast <bool> (expr.getPosition() < numSymbols && "inconsistent number of symbols") ? void (0) : __assert_fail ("expr.getPosition() < numSymbols && \"inconsistent number of symbols\"" , "mlir/lib/IR/AffineExpr.cpp", 1264, __extension__ __PRETTY_FUNCTION__ )); | |||
1265 | eq[getSymbolStartIndex() + expr.getPosition()] = 1; | |||
1266 | } | |||
1267 | ||||
1268 | void SimpleAffineExprFlattener::visitConstantExpr(AffineConstantExpr expr) { | |||
1269 | operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0)); | |||
1270 | auto &eq = operandExprStack.back(); | |||
1271 | eq[getConstantIndex()] = expr.getValue(); | |||
1272 | } | |||
1273 | ||||
1274 | void SimpleAffineExprFlattener::addLocalVariableSemiAffine( | |||
1275 | AffineExpr expr, SmallVectorImpl<int64_t> &result, | |||
1276 | unsigned long resultSize) { | |||
1277 | assert(result.size() == resultSize &&(static_cast <bool> (result.size() == resultSize && "`result` vector passed is not of correct size") ? void (0) : __assert_fail ("result.size() == resultSize && \"`result` vector passed is not of correct size\"" , "mlir/lib/IR/AffineExpr.cpp", 1278, __extension__ __PRETTY_FUNCTION__ )) | |||
1278 | "`result` vector passed is not of correct size")(static_cast <bool> (result.size() == resultSize && "`result` vector passed is not of correct size") ? void (0) : __assert_fail ("result.size() == resultSize && \"`result` vector passed is not of correct size\"" , "mlir/lib/IR/AffineExpr.cpp", 1278, __extension__ __PRETTY_FUNCTION__ )); | |||
1279 | int loc; | |||
1280 | if ((loc = findLocalId(expr)) == -1) | |||
1281 | addLocalIdSemiAffine(expr); | |||
1282 | std::fill(result.begin(), result.end(), 0); | |||
1283 | if (loc == -1) | |||
1284 | result[getLocalVarStartIndex() + numLocals - 1] = 1; | |||
1285 | else | |||
1286 | result[getLocalVarStartIndex() + loc] = 1; | |||
1287 | } | |||
1288 | ||||
1289 | // t = expr floordiv c <=> t = q, c * q <= expr <= c * q + c - 1 | |||
1290 | // A floordiv is thus flattened by introducing a new local variable q, and | |||
1291 | // replacing that expression with 'q' while adding the constraints | |||
1292 | // c * q <= expr <= c * q + c - 1 to localVarCst (done by | |||
1293 | // IntegerRelation::addLocalFloorDiv). | |||
1294 | // | |||
1295 | // A ceildiv is similarly flattened: | |||
1296 | // t = expr ceildiv c <=> t = (expr + c - 1) floordiv c | |||
1297 | // | |||
1298 | // In case of semi affine division expressions, t = expr floordiv symbolic_expr | |||
1299 | // or t = expr ceildiv symbolic_expr, introduce a local variable q (= expr | |||
1300 | // floordiv/ceildiv symbolic_expr), and the affine floordiv/ceildiv is added to | |||
1301 | // `localExprs`. | |||
1302 | void SimpleAffineExprFlattener::visitDivExpr(AffineBinaryOpExpr expr, | |||
1303 | bool isCeil) { | |||
1304 | assert(operandExprStack.size() >= 2)(static_cast <bool> (operandExprStack.size() >= 2) ? void (0) : __assert_fail ("operandExprStack.size() >= 2", "mlir/lib/IR/AffineExpr.cpp", 1304, __extension__ __PRETTY_FUNCTION__ )); | |||
1305 | ||||
1306 | MLIRContext *context = expr.getContext(); | |||
1307 | SmallVector<int64_t, 8> rhs = operandExprStack.back(); | |||
1308 | operandExprStack.pop_back(); | |||
1309 | SmallVector<int64_t, 8> &lhs = operandExprStack.back(); | |||
1310 | ||||
1311 | // Flatten semi affine division expressions by introducing a local | |||
1312 | // variable in place of the quotient, and the affine expression corresponding | |||
1313 | // to the quantifier is added to `localExprs`. | |||
1314 | if (!expr.getRHS().isa<AffineConstantExpr>()) { | |||
1315 | AffineExpr a = getAffineExprFromFlatForm(lhs, numDims, numSymbols, | |||
1316 | localExprs, context); | |||
1317 | AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols, | |||
1318 | localExprs, context); | |||
1319 | AffineExpr divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b); | |||
1320 | addLocalVariableSemiAffine(divExpr, lhs, lhs.size()); | |||
1321 | return; | |||
1322 | } | |||
1323 | ||||
1324 | // This is a pure affine expr; the RHS is a positive constant. | |||
1325 | int64_t rhsConst = rhs[getConstantIndex()]; | |||
1326 | // TODO: handle division by zero at the same time the issue is | |||
1327 | // fixed at other places. | |||
1328 | assert(rhsConst > 0 && "RHS constant has to be positive")(static_cast <bool> (rhsConst > 0 && "RHS constant has to be positive" ) ? void (0) : __assert_fail ("rhsConst > 0 && \"RHS constant has to be positive\"" , "mlir/lib/IR/AffineExpr.cpp", 1328, __extension__ __PRETTY_FUNCTION__ )); | |||
1329 | ||||
1330 | // Simplify the floordiv, ceildiv if possible by canceling out the greatest | |||
1331 | // common divisors of the numerator and denominator. | |||
1332 | uint64_t gcd = std::abs(rhsConst); | |||
1333 | for (unsigned i = 0, e = lhs.size(); i < e; i++) | |||
1334 | gcd = std::gcd(gcd, (uint64_t)std::abs(lhs[i])); | |||
1335 | // Simplify the numerator and the denominator. | |||
1336 | if (gcd != 1) { | |||
1337 | for (unsigned i = 0, e = lhs.size(); i < e; i++) | |||
1338 | lhs[i] = lhs[i] / static_cast<int64_t>(gcd); | |||
1339 | } | |||
1340 | int64_t divisor = rhsConst / static_cast<int64_t>(gcd); | |||
1341 | // If the divisor becomes 1, the updated LHS is the result. (The | |||
1342 | // divisor can't be negative since rhsConst is positive). | |||
1343 | if (divisor == 1) | |||
1344 | return; | |||
1345 | ||||
1346 | // If the divisor cannot be simplified to one, we will have to retain | |||
1347 | // the ceil/floor expr (simplified up until here). Add an existential | |||
1348 | // quantifier to express its result, i.e., expr1 div expr2 is replaced | |||
1349 | // by a new identifier, q. | |||
1350 | AffineExpr a = | |||
1351 | getAffineExprFromFlatForm(lhs, numDims, numSymbols, localExprs, context); | |||
1352 | AffineExpr b = getAffineConstantExpr(divisor, context); | |||
1353 | ||||
1354 | int loc; | |||
1355 | AffineExpr divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b); | |||
1356 | if ((loc = findLocalId(divExpr)) == -1) { | |||
1357 | if (!isCeil) { | |||
1358 | SmallVector<int64_t, 8> dividend(lhs); | |||
1359 | addLocalFloorDivId(dividend, divisor, divExpr); | |||
1360 | } else { | |||
1361 | // lhs ceildiv c <=> (lhs + c - 1) floordiv c | |||
1362 | SmallVector<int64_t, 8> dividend(lhs); | |||
1363 | dividend.back() += divisor - 1; | |||
1364 | addLocalFloorDivId(dividend, divisor, divExpr); | |||
1365 | } | |||
1366 | } | |||
1367 | // Set the expression on stack to the local var introduced to capture the | |||
1368 | // result of the division (floor or ceil). | |||
1369 | std::fill(lhs.begin(), lhs.end(), 0); | |||
1370 | if (loc == -1) | |||
1371 | lhs[getLocalVarStartIndex() + numLocals - 1] = 1; | |||
1372 | else | |||
1373 | lhs[getLocalVarStartIndex() + loc] = 1; | |||
1374 | } | |||
1375 | ||||
1376 | // Add a local identifier (needed to flatten a mod, floordiv, ceildiv expr). | |||
1377 | // The local identifier added is always a floordiv of a pure add/mul affine | |||
1378 | // function of other identifiers, coefficients of which are specified in | |||
1379 | // dividend and with respect to a positive constant divisor. localExpr is the | |||
1380 | // simplified tree expression (AffineExpr) corresponding to the quantifier. | |||
1381 | void SimpleAffineExprFlattener::addLocalFloorDivId(ArrayRef<int64_t> dividend, | |||
1382 | int64_t divisor, | |||
1383 | AffineExpr localExpr) { | |||
1384 | assert(divisor > 0 && "positive constant divisor expected")(static_cast <bool> (divisor > 0 && "positive constant divisor expected" ) ? void (0) : __assert_fail ("divisor > 0 && \"positive constant divisor expected\"" , "mlir/lib/IR/AffineExpr.cpp", 1384, __extension__ __PRETTY_FUNCTION__ )); | |||
1385 | for (SmallVector<int64_t, 8> &subExpr : operandExprStack) | |||
1386 | subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0); | |||
1387 | localExprs.push_back(localExpr); | |||
1388 | numLocals++; | |||
1389 | // dividend and divisor are not used here; an override of this method uses it. | |||
1390 | } | |||
1391 | ||||
1392 | void SimpleAffineExprFlattener::addLocalIdSemiAffine(AffineExpr localExpr) { | |||
1393 | for (SmallVector<int64_t, 8> &subExpr : operandExprStack) | |||
1394 | subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0); | |||
1395 | localExprs.push_back(localExpr); | |||
1396 | ++numLocals; | |||
1397 | } | |||
1398 | ||||
1399 | int SimpleAffineExprFlattener::findLocalId(AffineExpr localExpr) { | |||
1400 | SmallVectorImpl<AffineExpr>::iterator it; | |||
1401 | if ((it = llvm::find(localExprs, localExpr)) == localExprs.end()) | |||
1402 | return -1; | |||
1403 | return it - localExprs.begin(); | |||
1404 | } | |||
1405 | ||||
1406 | /// Simplify the affine expression by flattening it and reconstructing it. | |||
1407 | AffineExpr mlir::simplifyAffineExpr(AffineExpr expr, unsigned numDims, | |||
1408 | unsigned numSymbols) { | |||
1409 | // Simplify semi-affine expressions separately. | |||
1410 | if (!expr.isPureAffine()) | |||
1411 | expr = simplifySemiAffine(expr); | |||
1412 | ||||
1413 | SimpleAffineExprFlattener flattener(numDims, numSymbols); | |||
1414 | flattener.walkPostOrder(expr); | |||
1415 | ArrayRef<int64_t> flattenedExpr = flattener.operandExprStack.back(); | |||
1416 | if (!expr.isPureAffine() && | |||
1417 | expr == getAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols, | |||
1418 | flattener.localExprs, | |||
1419 | expr.getContext())) | |||
1420 | return expr; | |||
1421 | AffineExpr simplifiedExpr = | |||
1422 | expr.isPureAffine() | |||
1423 | ? getAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols, | |||
1424 | flattener.localExprs, expr.getContext()) | |||
1425 | : getSemiAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols, | |||
1426 | flattener.localExprs, | |||
1427 | expr.getContext()); | |||
1428 | ||||
1429 | flattener.operandExprStack.pop_back(); | |||
1430 | assert(flattener.operandExprStack.empty())(static_cast <bool> (flattener.operandExprStack.empty() ) ? void (0) : __assert_fail ("flattener.operandExprStack.empty()" , "mlir/lib/IR/AffineExpr.cpp", 1430, __extension__ __PRETTY_FUNCTION__ )); | |||
1431 | return simplifiedExpr; | |||
1432 | } |