1//===- ValueBoundsOpInterface.h - Value Bounds ------------------*- 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#ifndef MLIR_INTERFACES_VALUEBOUNDSOPINTERFACE_H_
10#define MLIR_INTERFACES_VALUEBOUNDSOPINTERFACE_H_
11
12#include "mlir/Analysis/FlatLinearValueConstraints.h"
13#include "mlir/IR/Builders.h"
14#include "mlir/IR/OpDefinition.h"
15#include "mlir/IR/Value.h"
16#include "mlir/Interfaces/DestinationStyleOpInterface.h"
17#include "llvm/ADT/SetVector.h"
18#include "llvm/ADT/SmallVector.h"
19#include "llvm/Support/ExtensibleRTTI.h"
20
21#include <queue>
22
23namespace mlir {
24class OffsetSizeAndStrideOpInterface;
25
26/// A hyperrectangular slice, represented as a list of offsets, sizes and
27/// strides.
28class HyperrectangularSlice {
29public:
30 HyperrectangularSlice(ArrayRef<OpFoldResult> offsets,
31 ArrayRef<OpFoldResult> sizes,
32 ArrayRef<OpFoldResult> strides);
33
34 /// Create a hyperrectangular slice with unit strides.
35 HyperrectangularSlice(ArrayRef<OpFoldResult> offsets,
36 ArrayRef<OpFoldResult> sizes);
37
38 /// Infer a hyperrectangular slice from `OffsetSizeAndStrideOpInterface`.
39 HyperrectangularSlice(OffsetSizeAndStrideOpInterface op);
40
41 ArrayRef<OpFoldResult> getMixedOffsets() const { return mixedOffsets; }
42 ArrayRef<OpFoldResult> getMixedSizes() const { return mixedSizes; }
43 ArrayRef<OpFoldResult> getMixedStrides() const { return mixedStrides; }
44
45private:
46 SmallVector<OpFoldResult> mixedOffsets;
47 SmallVector<OpFoldResult> mixedSizes;
48 SmallVector<OpFoldResult> mixedStrides;
49};
50
51using ValueDimList = SmallVector<std::pair<Value, std::optional<int64_t>>>;
52
53/// A helper class to be used with `ValueBoundsOpInterface`. This class stores a
54/// constraint system and mapping of constrained variables to index-typed
55/// values or dimension sizes of shaped values.
56///
57/// Interface implementations of `ValueBoundsOpInterface` use `addBounds` to
58/// insert constraints about their results and/or region block arguments into
59/// the constraint set in the form of an AffineExpr. When a bound should be
60/// expressed in terms of another value/dimension, `getExpr` can be used to
61/// retrieve an AffineExpr that represents the specified value/dimension.
62///
63/// When a value/dimension is retrieved for the first time through `getExpr`,
64/// it is added to an internal worklist. See `computeBound` for more details.
65///
66/// Note: Any modification of existing IR invalides the data stored in this
67/// class. Adding new operations is allowed.
68class ValueBoundsConstraintSet
69 : public llvm::RTTIExtends<ValueBoundsConstraintSet, llvm::RTTIRoot> {
70protected:
71 /// Helper class that builds a bound for a shaped value dimension or
72 /// index-typed value.
73 class BoundBuilder {
74 public:
75 /// Specify a dimension, assuming that the underlying value is a shaped
76 /// value.
77 BoundBuilder &operator[](int64_t dim);
78
79 // These overloaded operators add lower/upper/equality bounds.
80 void operator<(AffineExpr expr);
81 void operator<=(AffineExpr expr);
82 void operator>(AffineExpr expr);
83 void operator>=(AffineExpr expr);
84 void operator==(AffineExpr expr);
85 void operator<(OpFoldResult ofr);
86 void operator<=(OpFoldResult ofr);
87 void operator>(OpFoldResult ofr);
88 void operator>=(OpFoldResult ofr);
89 void operator==(OpFoldResult ofr);
90 void operator<(int64_t i);
91 void operator<=(int64_t i);
92 void operator>(int64_t i);
93 void operator>=(int64_t i);
94 void operator==(int64_t i);
95
96 protected:
97 friend class ValueBoundsConstraintSet;
98 BoundBuilder(ValueBoundsConstraintSet &cstr, Value value)
99 : cstr(cstr), value(value) {}
100
101 private:
102 BoundBuilder(const BoundBuilder &) = delete;
103 BoundBuilder &operator=(const BoundBuilder &) = delete;
104 bool operator==(const BoundBuilder &) = delete;
105 bool operator!=(const BoundBuilder &) = delete;
106
107 ValueBoundsConstraintSet &cstr;
108 Value value;
109 std::optional<int64_t> dim;
110 };
111
112public:
113 static char ID;
114
115 /// A variable that can be added to the constraint set as a "column". The
116 /// value bounds infrastructure can compute bounds for variables and compare
117 /// two variables.
118 ///
119 /// Internally, a variable is represented as an affine map and operands.
120 class Variable {
121 public:
122 /// Construct a variable for an index-typed attribute or SSA value.
123 Variable(OpFoldResult ofr);
124
125 /// Construct a variable for an index-typed SSA value.
126 Variable(Value indexValue);
127
128 /// Construct a variable for a dimension of a shaped value.
129 Variable(Value shapedValue, int64_t dim);
130
131 /// Construct a variable for an index-typed attribute/SSA value or for a
132 /// dimension of a shaped value. A non-null dimension must be provided if
133 /// and only if `ofr` is a shaped value.
134 Variable(OpFoldResult ofr, std::optional<int64_t> dim);
135
136 /// Construct a variable for a map and its operands.
137 Variable(AffineMap map, ArrayRef<Variable> mapOperands);
138 Variable(AffineMap map, ArrayRef<Value> mapOperands);
139
140 MLIRContext *getContext() const { return map.getContext(); }
141
142 private:
143 friend class ValueBoundsConstraintSet;
144 AffineMap map;
145 ValueDimList mapOperands;
146 };
147
148 /// The stop condition when traversing the backward slice of a shaped value/
149 /// index-type value. The traversal continues until the stop condition
150 /// evaluates to "true" for a value.
151 ///
152 /// The first parameter of the function is the shaped value/index-typed
153 /// value. The second parameter is the dimension in case of a shaped value.
154 /// The third parameter is this constraint set.
155 using StopConditionFn = std::function<bool(
156 Value, std::optional<int64_t> /*dim*/, ValueBoundsConstraintSet &cstr)>;
157
158 /// Compute a bound for the given variable. The computed bound is stored in
159 /// `resultMap`. The operands of the bound are stored in `mapOperands`. An
160 /// operand is either an index-type SSA value or a shaped value and a
161 /// dimension.
162 ///
163 /// The bound is computed in terms of values/dimensions for which
164 /// `stopCondition` evaluates to "true". To that end, the backward slice
165 /// (reverse use-def chain) of the given value is visited in a worklist-driven
166 /// manner and the constraint set is populated according to
167 /// `ValueBoundsOpInterface` for each visited value.
168 ///
169 /// By default, lower/equal bounds are closed and upper bounds are open. If
170 /// `closedUB` is set to "true", upper bounds are also closed.
171 static LogicalResult
172 computeBound(AffineMap &resultMap, ValueDimList &mapOperands,
173 presburger::BoundType type, const Variable &var,
174 StopConditionFn stopCondition, bool closedUB = false);
175
176 /// Compute a bound in terms of the values/dimensions in `dependencies`. The
177 /// computed bound consists of only constant terms and dependent values (or
178 /// dimension sizes thereof).
179 static LogicalResult
180 computeDependentBound(AffineMap &resultMap, ValueDimList &mapOperands,
181 presburger::BoundType type, const Variable &var,
182 ValueDimList dependencies, bool closedUB = false);
183
184 /// Compute a bound in that is independent of all values in `independencies`.
185 ///
186 /// Independencies are the opposite of dependencies. The computed bound does
187 /// not contain any SSA values that are part of `independencies`. E.g., this
188 /// function can be used to make ops hoistable from loops. To that end, ops
189 /// must be made independent of loop induction variables (in the case of "for"
190 /// loops). Loop induction variables are the independencies; they may not
191 /// appear in the computed bound.
192 static LogicalResult
193 computeIndependentBound(AffineMap &resultMap, ValueDimList &mapOperands,
194 presburger::BoundType type, const Variable &var,
195 ValueRange independencies, bool closedUB = false);
196
197 /// Compute a constant bound for the given variable.
198 ///
199 /// This function traverses the backward slice of the given operands in a
200 /// worklist-driven manner until `stopCondition` evaluates to "true". The
201 /// constraint set is populated according to `ValueBoundsOpInterface` for each
202 /// visited value. (No constraints are added for values for which the stop
203 /// condition evaluates to "true".)
204 ///
205 /// The stop condition is optional: If none is specified, the backward slice
206 /// is traversed in a breadth-first manner until a constant bound could be
207 /// computed.
208 ///
209 /// By default, lower/equal bounds are closed and upper bounds are open. If
210 /// `closedUB` is set to "true", upper bounds are also closed.
211 static FailureOr<int64_t>
212 computeConstantBound(presburger::BoundType type, const Variable &var,
213 StopConditionFn stopCondition = nullptr,
214 bool closedUB = false);
215
216 /// Compute a constant delta between the given two values. Return "failure"
217 /// if a constant delta could not be determined.
218 ///
219 /// `dim1`/`dim2` must be `nullopt` if and only if `value1`/`value2` are
220 /// index-typed.
221 static FailureOr<int64_t>
222 computeConstantDelta(Value value1, Value value2,
223 std::optional<int64_t> dim1 = std::nullopt,
224 std::optional<int64_t> dim2 = std::nullopt);
225
226 /// Traverse the IR starting from the given value/dim and populate constraints
227 /// as long as the stop condition holds. Also process all values/dims that are
228 /// already on the worklist.
229 void populateConstraints(Value value, std::optional<int64_t> dim);
230
231 /// Comparison operator for `ValueBoundsConstraintSet::compare`.
232 enum ComparisonOperator { LT, LE, EQ, GT, GE };
233
234 /// Populate constraints for lhs/rhs (until the stop condition is met). Then,
235 /// try to prove that, based on the current state of this constraint set
236 /// (i.e., without analyzing additional IR or adding new constraints), the
237 /// "lhs" value/dim is LE/LT/EQ/GT/GE than the "rhs" value/dim.
238 ///
239 /// Return "true" if the specified relation between the two values/dims was
240 /// proven to hold. Return "false" if the specified relation could not be
241 /// proven. This could be because the specified relation does in fact not hold
242 /// or because there is not enough information in the constraint set. In other
243 /// words, if we do not know for sure, this function returns "false".
244 bool populateAndCompare(const Variable &lhs, ComparisonOperator cmp,
245 const Variable &rhs);
246
247 /// Return "true" if "lhs cmp rhs" was proven to hold. Return "false" if the
248 /// specified relation could not be proven. This could be because the
249 /// specified relation does in fact not hold or because there is not enough
250 /// information in the constraint set. In other words, if we do not know for
251 /// sure, this function returns "false".
252 ///
253 /// This function keeps traversing the backward slice of lhs/rhs until could
254 /// prove the relation or until it ran out of IR.
255 static bool compare(const Variable &lhs, ComparisonOperator cmp,
256 const Variable &rhs);
257
258 /// Compute whether the given variables are equal. Return "failure" if
259 /// equality could not be determined.
260 static FailureOr<bool> areEqual(const Variable &var1, const Variable &var2);
261
262 /// Return "true" if the given slices are guaranteed to be overlapping.
263 /// Return "false" if the given slices are guaranteed to be non-overlapping.
264 /// Return "failure" if unknown.
265 ///
266 /// Slices are overlapping if for all dimensions:
267 /// * offset1 + size1 * stride1 <= offset2
268 /// * and offset2 + size2 * stride2 <= offset1
269 ///
270 /// Slice are non-overlapping if the above constraint is not satisfied for
271 /// at least one dimension.
272 static FailureOr<bool> areOverlappingSlices(MLIRContext *ctx,
273 HyperrectangularSlice slice1,
274 HyperrectangularSlice slice2);
275
276 /// Return "true" if the given slices are guaranteed to be equivalent.
277 /// Return "false" if the given slices are guaranteed to be non-equivalent.
278 /// Return "failure" if unknown.
279 ///
280 /// Slices are equivalent if their offsets, sizes and strices are equal.
281 static FailureOr<bool> areEquivalentSlices(MLIRContext *ctx,
282 HyperrectangularSlice slice1,
283 HyperrectangularSlice slice2);
284
285 /// Add a bound for the given index-typed value or shaped value. This function
286 /// returns a builder that adds the bound.
287 BoundBuilder bound(Value value) { return BoundBuilder(*this, value); }
288
289 /// Return an expression that represents the given index-typed value or shaped
290 /// value dimension. If this value/dimension was not used so far, it is added
291 /// to the worklist.
292 ///
293 /// `dim` must be `nullopt` if and only if the given value is of index type.
294 AffineExpr getExpr(Value value, std::optional<int64_t> dim = std::nullopt);
295
296 /// Return an expression that represents a constant or index-typed SSA value.
297 /// In case of a value, if this value was not used so far, it is added to the
298 /// worklist.
299 AffineExpr getExpr(OpFoldResult ofr);
300
301 /// Return an expression that represents a constant.
302 AffineExpr getExpr(int64_t constant);
303
304 /// Debugging only: Dump the constraint set and the column-to-value/dim
305 /// mapping to llvm::errs.
306 void dump() const;
307
308protected:
309 /// Dimension identifier to indicate a value is index-typed. This is used for
310 /// internal data structures/API only.
311 static constexpr int64_t kIndexValue = -1;
312
313 /// An index-typed value or the dimension of a shaped-type value.
314 using ValueDim = std::pair<Value, int64_t>;
315
316 ValueBoundsConstraintSet(MLIRContext *ctx, StopConditionFn stopCondition);
317
318 /// Return "true" if, based on the current state of the constraint system,
319 /// "lhs cmp rhs" was proven to hold. Return "false" if the specified relation
320 /// could not be proven. This could be because the specified relation does in
321 /// fact not hold or because there is not enough information in the constraint
322 /// set. In other words, if we do not know for sure, this function returns
323 /// "false".
324 ///
325 /// This function does not analyze any IR and does not populate any additional
326 /// constraints.
327 bool comparePos(int64_t lhsPos, ComparisonOperator cmp, int64_t rhsPos);
328
329 /// Given an affine map with a single result (and map operands), add a new
330 /// column to the constraint set that represents the result of the map.
331 /// Traverse additional IR starting from the map operands as needed (as long
332 /// as the stop condition is not satisfied). Also process all values/dims that
333 /// are already on the worklist. Return the position of the newly added
334 /// column.
335 int64_t populateConstraints(AffineMap map, ValueDimList mapOperands);
336
337 /// Iteratively process all elements on the worklist until an index-typed
338 /// value or shaped value meets `stopCondition`. Such values are not processed
339 /// any further.
340 void processWorklist();
341
342 /// Bound the given column in the underlying constraint set by the given
343 /// expression.
344 void addBound(presburger::BoundType type, int64_t pos, AffineExpr expr);
345
346 /// Return the column position of the given value/dimension. Asserts that the
347 /// value/dimension exists in the constraint set.
348 int64_t getPos(Value value, std::optional<int64_t> dim = std::nullopt) const;
349
350 /// Return an affine expression that represents column `pos` in the constraint
351 /// set.
352 AffineExpr getPosExpr(int64_t pos);
353
354 /// Return "true" if the given value/dim is mapped (i.e., has a corresponding
355 /// column in the constraint system).
356 bool isMapped(Value value, std::optional<int64_t> dim = std::nullopt) const;
357
358 /// Insert a value/dimension into the constraint set. If `isSymbol` is set to
359 /// "false", a dimension is added. The value/dimension is added to the
360 /// worklist if `addToWorklist` is set.
361 ///
362 /// Note: There are certain affine restrictions wrt. dimensions. E.g., they
363 /// cannot be multiplied. Furthermore, bounds can only be queried for
364 /// dimensions but not for symbols.
365 int64_t insert(Value value, std::optional<int64_t> dim, bool isSymbol = true,
366 bool addToWorklist = true);
367
368 /// Insert an anonymous column into the constraint set. The column is not
369 /// bound to any value/dimension. If `isSymbol` is set to "false", a dimension
370 /// is added.
371 ///
372 /// Note: There are certain affine restrictions wrt. dimensions. E.g., they
373 /// cannot be multiplied. Furthermore, bounds can only be queried for
374 /// dimensions but not for symbols.
375 int64_t insert(bool isSymbol = true);
376
377 /// Insert the given affine map and its bound operands as a new column in the
378 /// constraint system. Return the position of the new column. Any operands
379 /// that were not analyzed yet are put on the worklist.
380 int64_t insert(AffineMap map, ValueDimList operands, bool isSymbol = true);
381 int64_t insert(const Variable &var, bool isSymbol = true);
382
383 /// Project out the given column in the constraint set.
384 void projectOut(int64_t pos);
385
386 /// Project out all columns for which the condition holds.
387 void projectOut(function_ref<bool(ValueDim)> condition);
388
389 void projectOutAnonymous(std::optional<int64_t> except = std::nullopt);
390
391 /// Mapping of columns to values/shape dimensions.
392 SmallVector<std::optional<ValueDim>> positionToValueDim;
393 /// Reverse mapping of values/shape dimensions to columns.
394 DenseMap<ValueDim, int64_t> valueDimToPosition;
395
396 /// Worklist of values/shape dimensions that have not been processed yet.
397 std::queue<int64_t> worklist;
398
399 /// Constraint system of equalities and inequalities.
400 FlatLinearConstraints cstr;
401
402 /// Builder for constructing affine expressions.
403 Builder builder;
404
405 /// The current stop condition function.
406 StopConditionFn stopCondition = nullptr;
407};
408
409} // namespace mlir
410
411#include "mlir/Interfaces/ValueBoundsOpInterface.h.inc"
412
413#endif // MLIR_INTERFACES_VALUEBOUNDSOPINTERFACE_H_
414

source code of mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h