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
25namespace mlir {
26
27class MLIRContext;
28class AffineMap;
29class IntegerSet;
30
31namespace detail {
32
33struct AffineExprStorage;
34struct AffineBinaryOpExprStorage;
35struct AffineDimExprStorage;
36struct AffineConstantExprStorage;
37
38} // namespace detail
39
40enum 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.
68class AffineExpr {
69public:
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
195protected:
196 ImplType *expr{nullptr};
197
198private:
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.
214class AffineBinaryOpExpr : public AffineExpr {
215public:
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.
223class AffineDimExpr : public AffineExpr {
224public:
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.
231class AffineSymbolExpr : public AffineExpr {
232public:
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.
239class AffineConstantExpr : public AffineExpr {
240public:
241 using ImplType = detail::AffineConstantExprStorage;
242 /* implicit */ AffineConstantExpr(AffineExpr::ImplType *ptr = nullptr);
243 int64_t getValue() const;
244};
245
246/// Make AffineExpr hashable.
247inline ::llvm::hash_code hash_value(AffineExpr arg) {
248 return ::llvm::hash_value(ptr: arg.expr);
249}
250
251inline AffineExpr operator+(int64_t val, AffineExpr expr) { return expr + val; }
252inline AffineExpr operator*(int64_t val, AffineExpr expr) { return expr * val; }
253inline 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.
258AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context);
259AffineExpr getAffineSymbolExpr(unsigned position, MLIRContext *context);
260AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context);
261SmallVector<AffineExpr> getAffineConstantExprs(ArrayRef<int64_t> constants,
262 MLIRContext *context);
263AffineExpr 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].
271AffineExpr getAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs,
272 unsigned numDims, unsigned numSymbols,
273 ArrayRef<AffineExpr> localExprs,
274 MLIRContext *context);
275
276raw_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.
284AffineExpr simplifyAffineExpr(AffineExpr expr, unsigned numDims,
285 unsigned numSymbols);
286
287namespace detail {
288template <int N>
289void bindDims(MLIRContext *ctx) {}
290
291template <int N, typename AffineExprTy, typename... AffineExprTy2>
292void bindDims(MLIRContext *ctx, AffineExprTy &e, AffineExprTy2 &...exprs) {
293 e = getAffineDimExpr(position: N, context: ctx);
294 bindDims<N + 1, AffineExprTy2 &...>(ctx, exprs...);
295}
296
297template <int N>
298void bindSymbols(MLIRContext *ctx) {}
299
300template <int N, typename AffineExprTy, typename... AffineExprTy2>
301void 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)]
310template <typename... AffineExprTy>
311void bindDims(MLIRContext *ctx, AffineExprTy &...exprs) {
312 detail::bindDims<0>(ctx, exprs...);
313}
314
315template <typename AffineExprTy>
316void 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)]
324template <typename... AffineExprTy>
325void bindSymbols(MLIRContext *ctx, AffineExprTy &...exprs) {
326 detail::bindSymbols<0>(ctx, exprs...);
327}
328
329template <typename AffineExprTy>
330void 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.
344std::optional<int64_t>
345getBoundForAffineExpr(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
352namespace llvm {
353
354// AffineExpr hash just like pointers
355template <>
356struct 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.
375template <typename To, typename From>
376struct 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

source code of mlir/include/mlir/IR/AffineExpr.h