1 | //===- AffineExpr.h - MLIR Affine Expr Class --------------------*- C++ -*-===// |
---|---|
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 | // An affine expression is an affine combination of dimension identifiers and |
10 | // symbols, including ceildiv/floordiv/mod by a constant integer. |
11 | // |
12 | //===----------------------------------------------------------------------===// |
13 | |
14 | #ifndef MLIR_IR_AFFINEEXPR_H |
15 | #define MLIR_IR_AFFINEEXPR_H |
16 | |
17 | #include "mlir/IR/Visitors.h" |
18 | #include "mlir/Support/LLVM.h" |
19 | #include "llvm/ADT/DenseMapInfo.h" |
20 | #include "llvm/ADT/Hashing.h" |
21 | #include "llvm/ADT/SmallVector.h" |
22 | #include "llvm/Support/Casting.h" |
23 | #include <type_traits> |
24 | |
25 | namespace mlir { |
26 | |
27 | class MLIRContext; |
28 | class AffineMap; |
29 | class IntegerSet; |
30 | |
31 | namespace detail { |
32 | |
33 | struct AffineExprStorage; |
34 | struct AffineBinaryOpExprStorage; |
35 | struct AffineDimExprStorage; |
36 | struct AffineConstantExprStorage; |
37 | |
38 | } // namespace detail |
39 | |
40 | enum class AffineExprKind { |
41 | Add, |
42 | /// RHS of mul is always a constant or a symbolic expression. |
43 | Mul, |
44 | /// RHS of mod is always a constant or a symbolic expression with a positive |
45 | /// value. |
46 | Mod, |
47 | /// RHS of floordiv is always a constant or a symbolic expression. |
48 | FloorDiv, |
49 | /// RHS of ceildiv is always a constant or a symbolic expression. |
50 | CeilDiv, |
51 | |
52 | /// This is a marker for the last affine binary op. The range of binary |
53 | /// op's is expected to be this element and earlier. |
54 | LAST_AFFINE_BINARY_OP = CeilDiv, |
55 | |
56 | /// Constant integer. |
57 | Constant, |
58 | /// Dimensional identifier. |
59 | DimId, |
60 | /// Symbolic identifier. |
61 | SymbolId, |
62 | }; |
63 | |
64 | /// Base type for affine expression. |
65 | /// AffineExpr's are immutable value types with intuitive operators to |
66 | /// operate on chainable, lightweight compositions. |
67 | /// An AffineExpr is an interface to the underlying storage type pointer. |
68 | class AffineExpr { |
69 | public: |
70 | using ImplType = detail::AffineExprStorage; |
71 | |
72 | constexpr AffineExpr() {} |
73 | /* implicit */ AffineExpr(const ImplType *expr) |
74 | : expr(const_cast<ImplType *>(expr)) {} |
75 | |
76 | bool operator==(AffineExpr other) const { return expr == other.expr; } |
77 | bool operator!=(AffineExpr other) const { return !(*this == other); } |
78 | bool operator==(int64_t v) const; |
79 | bool operator!=(int64_t v) const { return !(*this == v); } |
80 | explicit operator bool() const { return expr; } |
81 | |
82 | bool operator!() const { return expr == nullptr; } |
83 | |
84 | MLIRContext *getContext() const; |
85 | |
86 | /// Return the classification for this type. |
87 | AffineExprKind getKind() const; |
88 | |
89 | void print(raw_ostream &os) const; |
90 | void dump() const; |
91 | |
92 | /// Returns true if this expression is made out of only symbols and |
93 | /// constants, i.e., it does not involve dimensional identifiers. |
94 | bool isSymbolicOrConstant() const; |
95 | |
96 | /// Returns true if this is a pure affine expression, i.e., multiplication, |
97 | /// floordiv, ceildiv, and mod is only allowed w.r.t constants. |
98 | bool isPureAffine() const; |
99 | |
100 | /// Returns the greatest known integral divisor of this affine expression. The |
101 | /// result is always positive. |
102 | int64_t getLargestKnownDivisor() const; |
103 | |
104 | /// Return true if the affine expression is a multiple of 'factor'. |
105 | bool isMultipleOf(int64_t factor) const; |
106 | |
107 | /// Return true if the affine expression involves AffineDimExpr `position`. |
108 | bool isFunctionOfDim(unsigned position) const; |
109 | |
110 | /// Return true if the affine expression involves AffineSymbolExpr `position`. |
111 | bool isFunctionOfSymbol(unsigned position) const; |
112 | |
113 | /// Walk all of the AffineExpr's in this expression in postorder. This allows |
114 | /// a lambda walk function that can either return `void` or a WalkResult. With |
115 | /// a WalkResult, interrupting is supported. |
116 | template <typename FnT, typename RetT = detail::walkResultType<FnT>> |
117 | RetT walk(FnT &&callback) const { |
118 | return walk<RetT>(*this, callback); |
119 | } |
120 | |
121 | /// This method substitutes any uses of dimensions and symbols (e.g. |
122 | /// dim#0 with dimReplacements[0]) and returns the modified expression tree. |
123 | /// This is a dense replacement method: a replacement must be specified for |
124 | /// every single dim and symbol. |
125 | AffineExpr replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements, |
126 | ArrayRef<AffineExpr> symReplacements) const; |
127 | |
128 | /// Dim-only version of replaceDimsAndSymbols. |
129 | AffineExpr replaceDims(ArrayRef<AffineExpr> dimReplacements) const; |
130 | |
131 | /// Symbol-only version of replaceDimsAndSymbols. |
132 | AffineExpr replaceSymbols(ArrayRef<AffineExpr> symReplacements) const; |
133 | |
134 | /// Sparse replace method. Replace `expr` by `replacement` and return the |
135 | /// modified expression tree. |
136 | AffineExpr replace(AffineExpr expr, AffineExpr replacement) const; |
137 | |
138 | /// Sparse replace method. If `*this` appears in `map` replaces it by |
139 | /// `map[*this]` and return the modified expression tree. Otherwise traverse |
140 | /// `*this` and apply replace with `map` on its subexpressions. |
141 | AffineExpr replace(const DenseMap<AffineExpr, AffineExpr> &map) const; |
142 | |
143 | /// Replace dims[offset ... numDims) |
144 | /// by dims[offset + shift ... shift + numDims). |
145 | AffineExpr shiftDims(unsigned numDims, unsigned shift, |
146 | unsigned offset = 0) const; |
147 | |
148 | /// Replace symbols[offset ... numSymbols) |
149 | /// by symbols[offset + shift ... shift + numSymbols). |
150 | AffineExpr shiftSymbols(unsigned numSymbols, unsigned shift, |
151 | unsigned offset = 0) const; |
152 | |
153 | AffineExpr operator+(int64_t v) const; |
154 | AffineExpr operator+(AffineExpr other) const; |
155 | AffineExpr operator-() const; |
156 | AffineExpr operator-(int64_t v) const; |
157 | AffineExpr operator-(AffineExpr other) const; |
158 | AffineExpr operator*(int64_t v) const; |
159 | AffineExpr operator*(AffineExpr other) const; |
160 | AffineExpr floorDiv(uint64_t v) const; |
161 | AffineExpr floorDiv(AffineExpr other) const; |
162 | AffineExpr ceilDiv(uint64_t v) const; |
163 | AffineExpr ceilDiv(AffineExpr other) const; |
164 | AffineExpr operator%(uint64_t v) const; |
165 | AffineExpr operator%(AffineExpr other) const; |
166 | |
167 | /// Compose with an AffineMap. |
168 | /// Returns the composition of this AffineExpr with `map`. |
169 | /// |
170 | /// Prerequisites: |
171 | /// `this` and `map` are composable, i.e. that the number of AffineDimExpr of |
172 | /// `this` is smaller than the number of results of `map`. If a result of a |
173 | /// map does not have a corresponding AffineDimExpr, that result simply does |
174 | /// not appear in the produced AffineExpr. |
175 | /// |
176 | /// Example: |
177 | /// expr: `d0 + d2` |
178 | /// map: `(d0, d1, d2)[s0, s1] -> (d0 + s1, d1 + s0, d0 + d1 + d2)` |
179 | /// returned expr: `d0 * 2 + d1 + d2 + s1` |
180 | AffineExpr compose(AffineMap map) const; |
181 | |
182 | friend ::llvm::hash_code hash_value(AffineExpr arg); |
183 | |
184 | /// Methods supporting C API. |
185 | const void *getAsOpaquePointer() const { |
186 | return static_cast<const void *>(expr); |
187 | } |
188 | static AffineExpr getFromOpaquePointer(const void *pointer) { |
189 | return AffineExpr( |
190 | reinterpret_cast<ImplType *>(const_cast<void *>(pointer))); |
191 | } |
192 | |
193 | ImplType *getImpl() const { return expr; } |
194 | |
195 | protected: |
196 | ImplType *expr{nullptr}; |
197 | |
198 | private: |
199 | /// A trampoline for the templated non-static AffineExpr::walk method to |
200 | /// dispatch lambda `callback`'s of either a void result type or a |
201 | /// WalkResult type. Walk all of the AffineExprs in `e` in postorder. Users |
202 | /// should use the regular (non-static) `walk` method. |
203 | template <typename WalkRetTy> |
204 | static WalkRetTy walk(AffineExpr e, |
205 | function_ref<WalkRetTy(AffineExpr)> callback); |
206 | }; |
207 | |
208 | /// Affine binary operation expression. An affine binary operation could be an |
209 | /// add, mul, floordiv, ceildiv, or a modulo operation. (Subtraction is |
210 | /// represented through a multiply by -1 and add.) These expressions are always |
211 | /// constructed in a simplified form. For eg., the LHS and RHS operands can't |
212 | /// both be constants. There are additional canonicalizing rules depending on |
213 | /// the op type: see checks in the constructor. |
214 | class AffineBinaryOpExpr : public AffineExpr { |
215 | public: |
216 | using ImplType = detail::AffineBinaryOpExprStorage; |
217 | /* implicit */ AffineBinaryOpExpr(AffineExpr::ImplType *ptr); |
218 | AffineExpr getLHS() const; |
219 | AffineExpr getRHS() const; |
220 | }; |
221 | |
222 | /// A dimensional identifier appearing in an affine expression. |
223 | class AffineDimExpr : public AffineExpr { |
224 | public: |
225 | using ImplType = detail::AffineDimExprStorage; |
226 | /* implicit */ AffineDimExpr(AffineExpr::ImplType *ptr); |
227 | unsigned getPosition() const; |
228 | }; |
229 | |
230 | /// A symbolic identifier appearing in an affine expression. |
231 | class AffineSymbolExpr : public AffineExpr { |
232 | public: |
233 | using ImplType = detail::AffineDimExprStorage; |
234 | /* implicit */ AffineSymbolExpr(AffineExpr::ImplType *ptr); |
235 | unsigned getPosition() const; |
236 | }; |
237 | |
238 | /// An integer constant appearing in affine expression. |
239 | class AffineConstantExpr : public AffineExpr { |
240 | public: |
241 | using ImplType = detail::AffineConstantExprStorage; |
242 | /* implicit */ AffineConstantExpr(AffineExpr::ImplType *ptr = nullptr); |
243 | int64_t getValue() const; |
244 | }; |
245 | |
246 | /// Make AffineExpr hashable. |
247 | inline ::llvm::hash_code hash_value(AffineExpr arg) { |
248 | return ::llvm::hash_value(ptr: arg.expr); |
249 | } |
250 | |
251 | inline AffineExpr operator+(int64_t val, AffineExpr expr) { return expr + val; } |
252 | inline AffineExpr operator*(int64_t val, AffineExpr expr) { return expr * val; } |
253 | inline AffineExpr operator-(int64_t val, AffineExpr expr) { |
254 | return expr * (-1) + val; |
255 | } |
256 | |
257 | /// These free functions allow clients of the API to not use classes in detail. |
258 | AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context); |
259 | AffineExpr getAffineSymbolExpr(unsigned position, MLIRContext *context); |
260 | AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context); |
261 | SmallVector<AffineExpr> getAffineConstantExprs(ArrayRef<int64_t> constants, |
262 | MLIRContext *context); |
263 | AffineExpr getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs, |
264 | AffineExpr rhs); |
265 | |
266 | /// Constructs an affine expression from a flat ArrayRef. If there are local |
267 | /// identifiers (neither dimensional nor symbolic) that appear in the sum of |
268 | /// products expression, 'localExprs' is expected to have the AffineExpr |
269 | /// for it, and is substituted into. The ArrayRef 'eq' is expected to be in the |
270 | /// format [dims, symbols, locals, constant term]. |
271 | AffineExpr getAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs, |
272 | unsigned numDims, unsigned numSymbols, |
273 | ArrayRef<AffineExpr> localExprs, |
274 | MLIRContext *context); |
275 | |
276 | raw_ostream &operator<<(raw_ostream &os, AffineExpr expr); |
277 | |
278 | /// Simplify an affine expression by flattening and some amount of simple |
279 | /// analysis. This has complexity linear in the number of nodes in 'expr'. |
280 | /// Returns the simplified expression, which is the same as the input expression |
281 | /// if it can't be simplified. When `expr` is semi-affine, a simplified |
282 | /// semi-affine expression is constructed in the sorted order of dimension and |
283 | /// symbol positions. |
284 | AffineExpr simplifyAffineExpr(AffineExpr expr, unsigned numDims, |
285 | unsigned numSymbols); |
286 | |
287 | namespace detail { |
288 | template <int N> |
289 | void bindDims(MLIRContext *ctx) {} |
290 | |
291 | template <int N, typename AffineExprTy, typename... AffineExprTy2> |
292 | void bindDims(MLIRContext *ctx, AffineExprTy &e, AffineExprTy2 &...exprs) { |
293 | e = getAffineDimExpr(position: N, context: ctx); |
294 | bindDims<N + 1, AffineExprTy2 &...>(ctx, exprs...); |
295 | } |
296 | |
297 | template <int N> |
298 | void bindSymbols(MLIRContext *ctx) {} |
299 | |
300 | template <int N, typename AffineExprTy, typename... AffineExprTy2> |
301 | void bindSymbols(MLIRContext *ctx, AffineExprTy &e, AffineExprTy2 &...exprs) { |
302 | e = getAffineSymbolExpr(position: N, context: ctx); |
303 | bindSymbols<N + 1, AffineExprTy2 &...>(ctx, exprs...); |
304 | } |
305 | |
306 | } // namespace detail |
307 | |
308 | /// Bind a list of AffineExpr references to DimExpr at positions: |
309 | /// [0 .. sizeof...(exprs)] |
310 | template <typename... AffineExprTy> |
311 | void bindDims(MLIRContext *ctx, AffineExprTy &...exprs) { |
312 | detail::bindDims<0>(ctx, exprs...); |
313 | } |
314 | |
315 | template <typename AffineExprTy> |
316 | void bindDimsList(MLIRContext *ctx, MutableArrayRef<AffineExprTy> exprs) { |
317 | int idx = 0; |
318 | for (AffineExprTy &e : exprs) |
319 | e = getAffineDimExpr(position: idx++, context: ctx); |
320 | } |
321 | |
322 | /// Bind a list of AffineExpr references to SymbolExpr at positions: |
323 | /// [0 .. sizeof...(exprs)] |
324 | template <typename... AffineExprTy> |
325 | void bindSymbols(MLIRContext *ctx, AffineExprTy &...exprs) { |
326 | detail::bindSymbols<0>(ctx, exprs...); |
327 | } |
328 | |
329 | template <typename AffineExprTy> |
330 | void bindSymbolsList(MLIRContext *ctx, MutableArrayRef<AffineExprTy> exprs) { |
331 | int idx = 0; |
332 | for (AffineExprTy &e : exprs) |
333 | e = getAffineSymbolExpr(position: idx++, context: ctx); |
334 | } |
335 | |
336 | /// Get a lower or upper (depending on `isUpper`) bound for `expr` while using |
337 | /// the constant lower and upper bounds for its inputs provided in |
338 | /// `constLowerBounds` and `constUpperBounds`. Return std::nullopt if such a |
339 | /// bound can't be computed. This method only handles simple sum of product |
340 | /// expressions (w.r.t constant coefficients) so as to not depend on anything |
341 | /// heavyweight in `Analysis`. Expressions of the form: c0*d0 + c1*d1 + c2*s0 + |
342 | /// ... + c_n are handled. Expressions involving floordiv, ceildiv, mod or |
343 | /// semi-affine ones will lead a none being returned. |
344 | std::optional<int64_t> |
345 | getBoundForAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols, |
346 | ArrayRef<std::optional<int64_t>> constLowerBounds, |
347 | ArrayRef<std::optional<int64_t>> constUpperBounds, |
348 | bool isUpper); |
349 | |
350 | } // namespace mlir |
351 | |
352 | namespace llvm { |
353 | |
354 | // AffineExpr hash just like pointers |
355 | template <> |
356 | struct DenseMapInfo<mlir::AffineExpr> { |
357 | static mlir::AffineExpr getEmptyKey() { |
358 | auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey(); |
359 | return mlir::AffineExpr(static_cast<mlir::AffineExpr::ImplType *>(pointer)); |
360 | } |
361 | static mlir::AffineExpr getTombstoneKey() { |
362 | auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey(); |
363 | return mlir::AffineExpr(static_cast<mlir::AffineExpr::ImplType *>(pointer)); |
364 | } |
365 | static unsigned getHashValue(mlir::AffineExpr val) { |
366 | return mlir::hash_value(arg: val); |
367 | } |
368 | static bool isEqual(mlir::AffineExpr LHS, mlir::AffineExpr RHS) { |
369 | return LHS == RHS; |
370 | } |
371 | }; |
372 | |
373 | /// Add support for llvm style casts. We provide a cast between To and From if |
374 | /// From is mlir::AffineExpr or derives from it. |
375 | template <typename To, typename From> |
376 | struct CastInfo<To, From, |
377 | std::enable_if_t<std::is_same_v<mlir::AffineExpr, |
378 | std::remove_const_t<From>> || |
379 | std::is_base_of_v<mlir::AffineExpr, From>>> |
380 | : NullableValueCastFailed<To>, |
381 | DefaultDoCastIfPossible<To, From, CastInfo<To, From>> { |
382 | |
383 | static inline bool isPossible(mlir::AffineExpr expr) { |
384 | /// Return a constant true instead of a dynamic true when casting to self or |
385 | /// up the hierarchy. |
386 | if constexpr (std::is_base_of_v<To, From>) { |
387 | return true; |
388 | } else { |
389 | if constexpr (std::is_same_v<To, ::mlir::AffineBinaryOpExpr>) |
390 | return expr.getKind() <= ::mlir::AffineExprKind::LAST_AFFINE_BINARY_OP; |
391 | if constexpr (std::is_same_v<To, ::mlir::AffineDimExpr>) |
392 | return expr.getKind() == ::mlir::AffineExprKind::DimId; |
393 | if constexpr (std::is_same_v<To, ::mlir::AffineSymbolExpr>) |
394 | return expr.getKind() == ::mlir::AffineExprKind::SymbolId; |
395 | if constexpr (std::is_same_v<To, ::mlir::AffineConstantExpr>) |
396 | return expr.getKind() == ::mlir::AffineExprKind::Constant; |
397 | } |
398 | } |
399 | static inline To doCast(mlir::AffineExpr expr) { return To(expr.getImpl()); } |
400 | }; |
401 | |
402 | } // namespace llvm |
403 | |
404 | #endif // MLIR_IR_AFFINEEXPR_H |
405 |
Definitions
- AffineExprKind
- AffineExpr
- AffineExpr
- AffineExpr
- operator==
- operator!=
- operator!=
- operator bool
- operator!
- walk
- getAsOpaquePointer
- getFromOpaquePointer
- getImpl
- AffineBinaryOpExpr
- AffineDimExpr
- AffineSymbolExpr
- AffineConstantExpr
- hash_value
- operator+
- operator*
- operator-
- bindDims
- bindDims
- bindSymbols
- bindSymbols
- bindDims
- bindDimsList
- bindSymbols
- bindSymbolsList
- DenseMapInfo
- getEmptyKey
- getTombstoneKey
- getHashValue
- isEqual
- CastInfo
- isPossible
Learn to use CMake with our Intro Training
Find out more