MLIR  15.0.0git
AffineCanonicalizationUtils.cpp
Go to the documentation of this file.
1 //===- AffineCanonicalizationUtils.cpp - Affine Canonicalization in SCF ---===//
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 // Utility functions to canonicalize affine ops within SCF op regions.
10 //
11 //===----------------------------------------------------------------------===//
12 
18 #include "mlir/IR/AffineMap.h"
19 #include "mlir/IR/Matchers.h"
20 #include "mlir/IR/PatternMatch.h"
21 #include "llvm/Support/Debug.h"
22 
23 #define DEBUG_TYPE "mlir-scf-affine-utils"
24 
25 using namespace mlir;
26 using namespace presburger;
27 
29  SmallVector<Value> &target) {
30  target = llvm::to_vector<4>(llvm::map_range(source, [](Optional<Value> val) {
31  return val.hasValue() ? *val : Value();
32  }));
33 }
34 
35 /// Bound an identifier `pos` in a given FlatAffineValueConstraints with
36 /// constraints drawn from an affine map. Before adding the constraint, the
37 /// dimensions/symbols of the affine map are aligned with `constraints`.
38 /// `operands` are the SSA Value operands used with the affine map.
39 /// Note: This function adds a new symbol column to the `constraints` for each
40 /// dimension/symbol that exists in the affine map but not in `constraints`.
43  unsigned pos, AffineMap map,
44  ValueRange operands) {
45  SmallVector<Value> dims, syms, newSyms;
48 
49  AffineMap alignedMap =
50  alignAffineMapWithValues(map, operands, dims, syms, &newSyms);
51  for (unsigned i = syms.size(); i < newSyms.size(); ++i)
52  constraints.appendSymbolVar(newSyms[i]);
53  return constraints.addBound(type, pos, alignedMap);
54 }
55 
56 /// Add `val` to each result of `map`.
57 static AffineMap addConstToResults(AffineMap map, int64_t val) {
58  SmallVector<AffineExpr> newResults;
59  for (AffineExpr r : map.getResults())
60  newResults.push_back(r + val);
61  return AffineMap::get(map.getNumDims(), map.getNumSymbols(), newResults,
62  map.getContext());
63 }
64 
65 /// This function tries to canonicalize min/max operations by proving that their
66 /// value is bounded by the same lower and upper bound. In that case, the
67 /// operation can be folded away.
68 ///
69 /// Bounds are computed by FlatAffineValueConstraints. Invariants required for
70 /// finding/proving bounds should be supplied via `constraints`.
71 ///
72 /// 1. Add dimensions for `op` and `opBound` (lower or upper bound of `op`).
73 /// 2. Compute an upper bound of `op` (in case of `isMin`) or a lower bound (in
74 /// case of `!isMin`) and bind it to `opBound`. SSA values that are used in
75 /// `op` but are not part of `constraints`, are added as extra symbols.
76 /// 3. For each result of `op`: Add result as a dimension `r_i`. Prove that:
77 /// * If `isMin`: r_i >= opBound
78 /// * If `isMax`: r_i <= opBound
79 /// If this is the case, ub(op) == lb(op).
80 /// 4. Replace `op` with `opBound`.
81 ///
82 /// In summary, the following constraints are added throughout this function.
83 /// Note: `invar` are dimensions added by the caller to express the invariants.
84 /// (Showing only the case where `isMin`.)
85 ///
86 /// invar | op | opBound | r_i | extra syms... | const | eq/ineq
87 /// ------+-------+---------+-----+---------------+-------+-------------------
88 /// (various eq./ineq. constraining `invar`, added by the caller)
89 /// ... | 0 | 0 | 0 | 0 | ... | ...
90 /// ------+-------+---------+-----+---------------+-------+-------------------
91 /// (various ineq. constraining `op` in terms of `op` operands (`invar` and
92 /// extra `op` operands "extra syms" that are not in `invar`)).
93 /// ... | -1 | 0 | 0 | ... | ... | >= 0
94 /// ------+-------+---------+-----+---------------+-------+-------------------
95 /// (set `opBound` to `op` upper bound in terms of `invar` and "extra syms")
96 /// ... | 0 | -1 | 0 | ... | ... | = 0
97 /// ------+-------+---------+-----+---------------+-------+-------------------
98 /// (for each `op` map result r_i: set r_i to corresponding map result,
99 /// prove that r_i >= minOpUb via contradiction)
100 /// ... | 0 | 0 | -1 | ... | ... | = 0
101 /// 0 | 0 | 1 | -1 | 0 | -1 | >= 0
102 ///
103 static LogicalResult
105  ValueRange operands, bool isMin,
106  FlatAffineValueConstraints constraints) {
107  RewriterBase::InsertionGuard guard(rewriter);
108  unsigned numResults = map.getNumResults();
109 
110  // Add a few extra dimensions.
111  unsigned dimOp = constraints.appendDimVar(); // `op`
112  unsigned dimOpBound = constraints.appendDimVar(); // `op` lower/upper bound
113  unsigned resultDimStart = constraints.appendDimVar(/*num=*/numResults);
114 
115  // Add an inequality for each result expr_i of map:
116  // isMin: op <= expr_i, !isMin: op >= expr_i
117  auto boundType = isMin ? IntegerPolyhedron::UB : IntegerPolyhedron::LB;
118  // Upper bounds are exclusive, so add 1. (`affine.min` ops are inclusive.)
119  AffineMap mapLbUb = isMin ? addConstToResults(map, 1) : map;
120  if (failed(
121  alignAndAddBound(constraints, boundType, dimOp, mapLbUb, operands)))
122  return failure();
123 
124  // Try to compute a lower/upper bound for op, expressed in terms of the other
125  // `dims` and extra symbols.
126  SmallVector<AffineMap> opLb(1), opUb(1);
127  constraints.getSliceBounds(dimOp, 1, rewriter.getContext(), &opLb, &opUb);
128  AffineMap sliceBound = isMin ? opUb[0] : opLb[0];
129  // TODO: `getSliceBounds` may return multiple bounds at the moment. This is
130  // a TODO of `getSliceBounds` and not handled here.
131  if (!sliceBound || sliceBound.getNumResults() != 1)
132  return failure(); // No or multiple bounds found.
133  // Recover the inclusive UB in the case of an `affine.min`.
134  AffineMap boundMap = isMin ? addConstToResults(sliceBound, -1) : sliceBound;
135 
136  // Add an equality: Set dimOpBound to computed bound.
137  // Add back dimension for op. (Was removed by `getSliceBounds`.)
138  AffineMap alignedBoundMap = boundMap.shiftDims(/*shift=*/1, /*offset=*/dimOp);
139  if (failed(constraints.addBound(IntegerPolyhedron::EQ, dimOpBound,
140  alignedBoundMap)))
141  return failure();
142 
143  // If the constraint system is empty, there is an inconsistency. (E.g., this
144  // can happen if loop lb > ub.)
145  if (constraints.isEmpty())
146  return failure();
147 
148  // In the case of `isMin` (`!isMin` is inversed):
149  // Prove that each result of `map` has a lower bound that is equal to (or
150  // greater than) the upper bound of `op` (`dimOpBound`). In that case, `op`
151  // can be replaced with the bound. I.e., prove that for each result
152  // expr_i (represented by dimension r_i):
153  //
154  // r_i >= opBound
155  //
156  // To prove this inequality, add its negation to the constraint set and prove
157  // that the constraint set is empty.
158  for (unsigned i = resultDimStart; i < resultDimStart + numResults; ++i) {
159  FlatAffineValueConstraints newConstr(constraints);
160 
161  // Add an equality: r_i = expr_i
162  // Note: These equalities could have been added earlier and used to express
163  // minOp <= expr_i. However, then we run the risk that `getSliceBounds`
164  // computes minOpUb in terms of r_i dims, which is not desired.
166  map.getSubMap({i - resultDimStart}), operands)))
167  return failure();
168 
169  // If `isMin`: Add inequality: r_i < opBound
170  // equiv.: opBound - r_i - 1 >= 0
171  // If `!isMin`: Add inequality: r_i > opBound
172  // equiv.: -opBound + r_i - 1 >= 0
173  SmallVector<int64_t> ineq(newConstr.getNumCols(), 0);
174  ineq[dimOpBound] = isMin ? 1 : -1;
175  ineq[i] = isMin ? -1 : 1;
176  ineq[newConstr.getNumCols() - 1] = -1;
177  newConstr.addInequality(ineq);
178  if (!newConstr.isEmpty())
179  return failure();
180  }
181 
182  // Lower and upper bound of `op` are equal. Replace `minOp` with its bound.
183  AffineMap newMap = alignedBoundMap;
184  SmallVector<Value> newOperands;
185  unpackOptionalValues(constraints.getMaybeValues(), newOperands);
186  // If dims/symbols have known constant values, use those in order to simplify
187  // the affine map further.
188  for (int64_t i = 0, e = constraints.getNumVars(); i < e; ++i) {
189  // Skip unused operands and operands that are already constants.
190  if (!newOperands[i] || getConstantIntValue(newOperands[i]))
191  continue;
192  if (auto bound = constraints.getConstantBound(IntegerPolyhedron::EQ, i))
193  newOperands[i] =
194  rewriter.create<arith::ConstantIndexOp>(op->getLoc(), *bound);
195  }
196  mlir::canonicalizeMapAndOperands(&newMap, &newOperands);
197  rewriter.setInsertionPoint(op);
198  rewriter.replaceOpWithNewOp<AffineApplyOp>(op, newMap, newOperands);
199  return success();
200 }
201 
202 static LogicalResult
205  RewriterBase &rewriter) {
206  // IntegerPolyhedron does not support semi-affine expressions.
207  // Therefore, only constant step values are supported.
208  auto stepInt = getConstantIntValue(step);
209  if (!stepInt)
210  return failure();
211 
212  unsigned dimIv = constraints.appendDimVar(iv);
213  auto lbv = lb.dyn_cast<Value>();
214  unsigned dimLb =
215  lbv ? constraints.appendDimVar(lbv) : constraints.appendDimVar(/*num=*/1);
216  auto ubv = ub.dyn_cast<Value>();
217  unsigned dimUb =
218  ubv ? constraints.appendDimVar(ubv) : constraints.appendDimVar(/*num=*/1);
219 
220  // If loop lower/upper bounds are constant: Add EQ constraint.
223  if (lbInt)
224  constraints.addBound(IntegerPolyhedron::EQ, dimLb, *lbInt);
225  if (ubInt)
226  constraints.addBound(IntegerPolyhedron::EQ, dimUb, *ubInt);
227 
228  // Lower bound: iv >= lb (equiv.: iv - lb >= 0)
229  SmallVector<int64_t> ineqLb(constraints.getNumCols(), 0);
230  ineqLb[dimIv] = 1;
231  ineqLb[dimLb] = -1;
232  constraints.addInequality(ineqLb);
233 
234  // Upper bound
235  AffineExpr ivUb;
236  if (lbInt && ubInt && (*lbInt + *stepInt >= *ubInt)) {
237  // The loop has at most one iteration.
238  // iv < lb + 1
239  // TODO: Try to derive this constraint by simplifying the expression in
240  // the else-branch.
241  ivUb = rewriter.getAffineDimExpr(dimLb) + 1;
242  } else {
243  // The loop may have more than one iteration.
244  // iv < lb + step * ((ub - lb - 1) floorDiv step) + 1
245  AffineExpr exprLb = lbInt ? rewriter.getAffineConstantExpr(*lbInt)
246  : rewriter.getAffineDimExpr(dimLb);
247  AffineExpr exprUb = ubInt ? rewriter.getAffineConstantExpr(*ubInt)
248  : rewriter.getAffineDimExpr(dimUb);
249  ivUb = exprLb + 1 + (*stepInt * ((exprUb - exprLb - 1).floorDiv(*stepInt)));
250  }
251  auto map = AffineMap::get(
252  /*dimCount=*/constraints.getNumDimVars(),
253  /*symbolCount=*/constraints.getNumSymbolVars(), /*result=*/ivUb);
254 
255  return constraints.addBound(IntegerPolyhedron::UB, dimIv, map);
256 }
257 
258 /// Canonicalize min/max operations in the context of for loops with a known
259 /// range. Call `canonicalizeMinMaxOp` and add the following constraints to
260 /// the constraint system (along with the missing dimensions):
261 ///
262 /// * iv >= lb
263 /// * iv < lb + step * ((ub - lb - 1) floorDiv step) + 1
264 ///
265 /// Note: Due to limitations of IntegerPolyhedron, only constant step sizes
266 /// are currently supported.
268  Operation *op, AffineMap map,
269  ValueRange operands, bool isMin,
270  LoopMatcherFn loopMatcher) {
271  FlatAffineValueConstraints constraints;
272  DenseSet<Value> allIvs;
273 
274  // Find all iteration variables among `minOp`'s operands add constrain them.
275  for (Value operand : operands) {
276  // Skip duplicate ivs.
277  if (llvm::is_contained(allIvs, operand))
278  continue;
279 
280  // If `operand` is an iteration variable: Find corresponding loop
281  // bounds and step.
282  Value iv = operand;
283  OpFoldResult lb, ub, step;
284  if (failed(loopMatcher(operand, lb, ub, step)))
285  continue;
286  allIvs.insert(iv);
287 
288  if (failed(
289  addLoopRangeConstraints(constraints, iv, lb, ub, step, rewriter)))
290  return failure();
291  }
292 
293  return canonicalizeMinMaxOp(rewriter, op, map, operands, isMin, constraints);
294 }
295 
296 /// Try to simplify a min/max operation `op` after loop peeling. This function
297 /// can simplify min/max operations such as (ub is the previous upper bound of
298 /// the unpeeled loop):
299 /// ```
300 /// #map = affine_map<(d0)[s0, s1] -> (s0, -d0 + s1)>
301 /// %r = affine.min #affine.min #map(%iv)[%step, %ub]
302 /// ```
303 /// and rewrites them into (in the case the peeled loop):
304 /// ```
305 /// %r = %step
306 /// ```
307 /// min/max operations inside the partial iteration are rewritten in a similar
308 /// way.
309 ///
310 /// This function builds up a set of constraints, capable of proving that:
311 /// * Inside the peeled loop: min(step, ub - iv) == step
312 /// * Inside the partial iteration: min(step, ub - iv) == ub - iv
313 ///
314 /// Returns `success` if the given operation was replaced by a new operation;
315 /// `failure` otherwise.
316 ///
317 /// Note: `ub` is the previous upper bound of the loop (before peeling).
318 /// `insideLoop` must be true for min/max ops inside the loop and false for
319 /// affine.min ops inside the partial iteration. For an explanation of the other
320 /// parameters, see comment of `canonicalizeMinMaxOpInLoop`.
322  AffineMap map, ValueRange operands,
323  bool isMin, Value iv, Value ub,
324  Value step, bool insideLoop) {
325  FlatAffineValueConstraints constraints;
326  constraints.appendDimVar({iv, ub, step});
327  if (auto constUb = getConstantIntValue(ub))
328  constraints.addBound(IntegerPolyhedron::EQ, 1, *constUb);
329  if (auto constStep = getConstantIntValue(step))
330  constraints.addBound(IntegerPolyhedron::EQ, 2, *constStep);
331 
332  // Add loop peeling invariant. This is the main piece of knowledge that
333  // enables AffineMinOp simplification.
334  if (insideLoop) {
335  // ub - iv >= step (equiv.: -iv + ub - step + 0 >= 0)
336  // Intuitively: Inside the peeled loop, every iteration is a "full"
337  // iteration, i.e., step divides the iteration space `ub - lb` evenly.
338  constraints.addInequality({-1, 1, -1, 0});
339  } else {
340  // ub - iv < step (equiv.: iv + -ub + step - 1 >= 0)
341  // Intuitively: `iv` is the split bound here, i.e., the iteration variable
342  // value of the very last iteration (in the unpeeled loop). At that point,
343  // there are less than `step` elements remaining. (Otherwise, the peeled
344  // loop would run for at least one more iteration.)
345  constraints.addInequality({1, -1, 1, -1});
346  }
347 
348  return canonicalizeMinMaxOp(rewriter, op, map, operands, isMin, constraints);
349 }
TODO: Remove this file when SCCP and integer range analysis have been ported to the new framework...
static LogicalResult alignAndAddBound(FlatAffineValueConstraints &constraints, IntegerPolyhedron::BoundType type, unsigned pos, AffineMap map, ValueRange operands)
Bound an identifier pos in a given FlatAffineValueConstraints with constraints drawn from an affine m...
MLIRContext * getContext() const
Definition: Builders.h:54
unsigned appendSymbolVar(ValueRange vals)
Operation is a basic unit of execution within MLIR.
Definition: Operation.h:28
AffineMap getSubMap(ArrayRef< unsigned > resultPos) const
Returns the map consisting of the resultPos subset.
Definition: AffineMap.cpp:514
unsigned getNumSymbols() const
Definition: AffineMap.cpp:298
unsigned getNumDims() const
Definition: AffineMap.cpp:294
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:336
This class represents a single result from folding an operation.
Definition: OpDefinition.h:229
LogicalResult canonicalizeMinMaxOpInLoop(RewriterBase &rewriter, Operation *op, AffineMap map, ValueRange operands, bool isMin, LoopMatcherFn loopMatcher)
Try to canonicalize an min/max operations in the context of for loops with a known range...
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value...
Definition: LogicalResult.h:72
AffineMap alignAffineMapWithValues(AffineMap map, ValueRange operands, ValueRange dims, ValueRange syms, SmallVector< Value > *newSyms=nullptr)
Re-indexes the dimensions and symbols of an affine map with given operands values to align with dims ...
unsigned appendDimVar(ValueRange vals)
Append variables of the specified kind after the last variable of that kind.
AffineMap shiftDims(unsigned shift, unsigned offset=0) const
Replace dims[offset ...
Definition: AffineMap.h:220
BoundType
The type of bound: equal, lower bound or upper bound.
static AffineMap addConstToResults(AffineMap map, int64_t val)
Add val to each result of map.
void addInequality(ArrayRef< int64_t > inEq)
Adds an inequality (>= 0) from the coefficients specified in inEq.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:380
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
ArrayRef< Optional< Value > > getMaybeValues() const
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
int64_t floorDiv(int64_t lhs, int64_t rhs)
Returns the result of MLIR&#39;s floordiv operation on constants.
Definition: MathExtras.h:33
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
static LogicalResult addLoopRangeConstraints(FlatAffineValueConstraints &constraints, Value iv, OpFoldResult lb, OpFoldResult ub, OpFoldResult step, RewriterBase &rewriter)
Base type for affine expression.
Definition: AffineExpr.h:68
void canonicalizeMapAndOperands(AffineMap *map, SmallVectorImpl< Value > *operands)
Modifies both map and operands in-place so as to:
Definition: AffineOps.cpp:873
unsigned getNumResults() const
Definition: AffineMap.cpp:302
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:161
A multi-dimensional affine map Affine map&#39;s are immutable like Type&#39;s, and they are uniqued...
Definition: AffineMap.h:42
ArrayRef< AffineExpr > getResults() const
Definition: AffineMap.cpp:307
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:85
AffineExpr getAffineConstantExpr(int64_t constant)
Definition: Builders.cpp:293
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replaces the result op with a new op that is created without verification.
Definition: PatternMatch.h:451
FlatAffineValueConstraints represents an extension of IntegerPolyhedron where each non-local variable...
Specialization of arith.constant op that returns an integer of index type.
Definition: Arithmetic.h:80
bool isEmpty() const
Checks for emptiness by performing variable elimination on all variables, running the GCD test on eac...
unsigned getNumCols() const
Returns the number of columns in the constraint system.
AffineExpr getAffineDimExpr(unsigned position)
Definition: Builders.cpp:285
LogicalResult rewritePeeledMinMaxOp(RewriterBase &rewriter, Operation *op, AffineMap map, ValueRange operands, bool isMin, Value iv, Value ub, Value step, bool insideLoop)
Try to simplify a min/max operation op after loop peeling.
static void unpackOptionalValues(ArrayRef< Optional< Value >> source, SmallVector< Value > &target)
MLIRContext * getContext() const
Definition: AffineMap.cpp:253
static LogicalResult canonicalizeMinMaxOp(RewriterBase &rewriter, Operation *op, AffineMap map, ValueRange operands, bool isMin, FlatAffineValueConstraints constraints)
This function tries to canonicalize min/max operations by proving that their value is bounded by the ...
Optional< int64_t > getConstantIntValue(OpFoldResult ofr)
If ofr is a constant integer or an IntegerAttr, return the integer.
This class provides an abstraction over the different types of ranges over Values.
Optional< int64_t > getConstantBound(BoundType type, unsigned pos) const
Returns the constant bound for the pos^th variable if there is one; None otherwise.
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:398
void getSliceBounds(unsigned offset, unsigned num, MLIRContext *context, SmallVectorImpl< AffineMap > *lbMaps, SmallVectorImpl< AffineMap > *ubMaps, bool getClosedUB=false)
Computes the lower and upper bounds of the first num dimensional variables (starting at offset) as an...
LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap, bool isClosedBound)
Adds a bound for the variable at the specified position with constraints being drawn from the specifi...