File: | build/source/mlir/lib/IR/AffineExpr.cpp |
Warning: | line 1072, column 7 Value stored to 'rhsPos' is never read |
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 == AffineExprKind::Add) |
50 | return lhs + rhs; |
51 | if (kind == AffineExprKind::Mul) |
52 | return lhs * rhs; |
53 | if (kind == AffineExprKind::FloorDiv) |
54 | return lhs.floorDiv(rhs); |
55 | if (kind == AffineExprKind::CeilDiv) |
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(); |
Value stored to 'rhsPos' is never read | |
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 | } |