MLIR  19.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 
13 #include <utility>
14 
22 #include "mlir/IR/AffineMap.h"
23 #include "mlir/IR/Matchers.h"
24 #include "mlir/IR/PatternMatch.h"
25 #include "llvm/Support/Debug.h"
26 
27 #define DEBUG_TYPE "mlir-scf-affine-utils"
28 
29 using namespace mlir;
30 using namespace affine;
31 using namespace presburger;
32 
34  OpFoldResult &ub, OpFoldResult &step) {
35  if (scf::ForOp forOp = scf::getForInductionVarOwner(iv)) {
36  lb = forOp.getLowerBound();
37  ub = forOp.getUpperBound();
38  step = forOp.getStep();
39  return success();
40  }
41  if (scf::ParallelOp parOp = scf::getParallelForInductionVarOwner(iv)) {
42  for (unsigned idx = 0; idx < parOp.getNumLoops(); ++idx) {
43  if (parOp.getInductionVars()[idx] == iv) {
44  lb = parOp.getLowerBound()[idx];
45  ub = parOp.getUpperBound()[idx];
46  step = parOp.getStep()[idx];
47  return success();
48  }
49  }
50  return failure();
51  }
52  if (scf::ForallOp forallOp = scf::getForallOpThreadIndexOwner(iv)) {
53  for (int64_t idx = 0; idx < forallOp.getRank(); ++idx) {
54  if (forallOp.getInductionVar(idx) == iv) {
55  lb = forallOp.getMixedLowerBound()[idx];
56  ub = forallOp.getMixedUpperBound()[idx];
57  step = forallOp.getMixedStep()[idx];
58  return success();
59  }
60  }
61  return failure();
62  }
63  return failure();
64 }
65 
68  FlatAffineValueConstraints constraints) {
69  RewriterBase::InsertionGuard guard(rewriter);
70  rewriter.setInsertionPoint(op);
71  FailureOr<AffineValueMap> simplified =
72  affine::simplifyConstrainedMinMaxOp(op, std::move(constraints));
73  if (failed(simplified))
74  return failure();
75  return rewriter.replaceOpWithNewOp<AffineApplyOp>(
76  op, simplified->getAffineMap(), simplified->getOperands());
77 }
78 
80  Value iv, OpFoldResult lb,
81  OpFoldResult ub, OpFoldResult step) {
82  Builder b(iv.getContext());
83 
84  // IntegerPolyhedron does not support semi-affine expressions.
85  // Therefore, only constant step values are supported.
86  auto stepInt = getConstantIntValue(step);
87  if (!stepInt)
88  return failure();
89 
90  unsigned dimIv = cstr.appendDimVar(iv);
91  auto lbv = llvm::dyn_cast_if_present<Value>(lb);
92  unsigned symLb =
93  lbv ? cstr.appendSymbolVar(lbv) : cstr.appendSymbolVar(/*num=*/1);
94  auto ubv = llvm::dyn_cast_if_present<Value>(ub);
95  unsigned symUb =
96  ubv ? cstr.appendSymbolVar(ubv) : cstr.appendSymbolVar(/*num=*/1);
97 
98  // If loop lower/upper bounds are constant: Add EQ constraint.
99  std::optional<int64_t> lbInt = getConstantIntValue(lb);
100  std::optional<int64_t> ubInt = getConstantIntValue(ub);
101  if (lbInt)
102  cstr.addBound(BoundType::EQ, symLb, *lbInt);
103  if (ubInt)
104  cstr.addBound(BoundType::EQ, symUb, *ubInt);
105 
106  // Lower bound: iv >= lb (equiv.: iv - lb >= 0)
107  SmallVector<int64_t> ineqLb(cstr.getNumCols(), 0);
108  ineqLb[dimIv] = 1;
109  ineqLb[symLb] = -1;
110  cstr.addInequality(ineqLb);
111 
112  // Upper bound
113  AffineExpr ivUb;
114  if (lbInt && ubInt && (*lbInt + *stepInt >= *ubInt)) {
115  // The loop has at most one iteration.
116  // iv < lb + 1
117  // TODO: Try to derive this constraint by simplifying the expression in
118  // the else-branch.
119  ivUb = b.getAffineSymbolExpr(symLb - cstr.getNumDimVars()) + 1;
120  } else {
121  // The loop may have more than one iteration.
122  // iv < lb + step * ((ub - lb - 1) floorDiv step) + 1
123  AffineExpr exprLb =
124  lbInt ? b.getAffineConstantExpr(*lbInt)
125  : b.getAffineSymbolExpr(symLb - cstr.getNumDimVars());
126  AffineExpr exprUb =
127  ubInt ? b.getAffineConstantExpr(*ubInt)
128  : b.getAffineSymbolExpr(symUb - cstr.getNumDimVars());
129  ivUb = exprLb + 1 + (*stepInt * ((exprUb - exprLb - 1).floorDiv(*stepInt)));
130  }
131  auto map = AffineMap::get(
132  /*dimCount=*/cstr.getNumDimVars(),
133  /*symbolCount=*/cstr.getNumSymbolVars(), /*result=*/ivUb);
134 
135  return cstr.addBound(BoundType::UB, dimIv, map);
136 }
137 
138 /// Canonicalize min/max operations in the context of for loops with a known
139 /// range. Call `canonicalizeMinMaxOp` and add the following constraints to
140 /// the constraint system (along with the missing dimensions):
141 ///
142 /// * iv >= lb
143 /// * iv < lb + step * ((ub - lb - 1) floorDiv step) + 1
144 ///
145 /// Note: Due to limitations of IntegerPolyhedron, only constant step sizes
146 /// are currently supported.
148  Operation *op,
149  LoopMatcherFn loopMatcher) {
150  FlatAffineValueConstraints constraints;
151  DenseSet<Value> allIvs;
152 
153  // Find all iteration variables among `minOp`'s operands add constrain them.
154  for (Value operand : op->getOperands()) {
155  // Skip duplicate ivs.
156  if (allIvs.contains(operand))
157  continue;
158 
159  // If `operand` is an iteration variable: Find corresponding loop
160  // bounds and step.
161  Value iv = operand;
162  OpFoldResult lb, ub, step;
163  if (failed(loopMatcher(operand, lb, ub, step)))
164  continue;
165  allIvs.insert(iv);
166 
167  if (failed(addLoopRangeConstraints(constraints, iv, lb, ub, step)))
168  return failure();
169  }
170 
171  return canonicalizeMinMaxOp(rewriter, op, constraints);
172 }
173 
174 /// Try to simplify the given affine.min/max operation `op` after loop peeling.
175 /// This function can simplify min/max operations such as (ub is the previous
176 /// upper bound of the unpeeled loop):
177 /// ```
178 /// #map = affine_map<(d0)[s0, s1] -> (s0, -d0 + s1)>
179 /// %r = affine.min #affine.min #map(%iv)[%step, %ub]
180 /// ```
181 /// and rewrites them into (in the case the peeled loop):
182 /// ```
183 /// %r = %step
184 /// ```
185 /// min/max operations inside the partial iteration are rewritten in a similar
186 /// way.
187 ///
188 /// This function builds up a set of constraints, capable of proving that:
189 /// * Inside the peeled loop: min(step, ub - iv) == step
190 /// * Inside the partial iteration: min(step, ub - iv) == ub - iv
191 ///
192 /// Returns `success` if the given operation was replaced by a new operation;
193 /// `failure` otherwise.
194 ///
195 /// Note: `ub` is the previous upper bound of the loop (before peeling).
196 /// `insideLoop` must be true for min/max ops inside the loop and false for
197 /// affine.min ops inside the partial iteration. For an explanation of the other
198 /// parameters, see comment of `canonicalizeMinMaxOpInLoop`.
200  Value iv, Value ub, Value step,
201  bool insideLoop) {
202  FlatAffineValueConstraints constraints;
203  constraints.appendDimVar({iv});
204  constraints.appendSymbolVar({ub, step});
205  if (auto constUb = getConstantIntValue(ub))
206  constraints.addBound(BoundType::EQ, 1, *constUb);
207  if (auto constStep = getConstantIntValue(step))
208  constraints.addBound(BoundType::EQ, 2, *constStep);
209 
210  // Add loop peeling invariant. This is the main piece of knowledge that
211  // enables AffineMinOp simplification.
212  if (insideLoop) {
213  // ub - iv >= step (equiv.: -iv + ub - step + 0 >= 0)
214  // Intuitively: Inside the peeled loop, every iteration is a "full"
215  // iteration, i.e., step divides the iteration space `ub - lb` evenly.
216  constraints.addInequality({-1, 1, -1, 0});
217  } else {
218  // ub - iv < step (equiv.: iv + -ub + step - 1 >= 0)
219  // Intuitively: `iv` is the split bound here, i.e., the iteration variable
220  // value of the very last iteration (in the unpeeled loop). At that point,
221  // there are less than `step` elements remaining. (Otherwise, the peeled
222  // loop would run for at least one more iteration.)
223  constraints.addInequality({1, -1, 1, -1});
224  }
225 
226  return canonicalizeMinMaxOp(rewriter, op, constraints);
227 }
static FailureOr< AffineApplyOp > canonicalizeMinMaxOp(RewriterBase &rewriter, Operation *op, FlatAffineValueConstraints constraints)
Base type for affine expression.
Definition: AffineExpr.h:69
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
This class is a general helper class for creating context-global objects like types,...
Definition: Builders.h:50
AffineExpr getAffineSymbolExpr(unsigned position)
Definition: Builders.cpp:375
AffineExpr getAffineConstantExpr(int64_t constant)
Definition: Builders.cpp:379
This class provides support for representing a failure result, or a valid value of type T.
Definition: LogicalResult.h:78
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:400
This class represents a single result from folding an operation.
Definition: OpDefinition.h:268
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
operand_range getOperands()
Returns an iterator on the underlying Value's.
Definition: Operation.h:373
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:400
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Definition: PatternMatch.h:536
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
MLIRContext * getContext() const
Utility to get the associated MLIRContext that this value is defined in.
Definition: Value.h:132
FlatAffineValueConstraints is an extension of FlatLinearValueConstraints with helper functions for Af...
LogicalResult addBound(presburger::BoundType type, unsigned pos, AffineMap boundMap, ValueRange operands)
Adds a bound for the variable at the specified position with constraints being drawn from the specifi...
void addInequality(ArrayRef< MPInt > inEq)
Adds an inequality (>= 0) from the coefficients specified in inEq.
unsigned getNumCols() const
Returns the number of columns in the constraint system.
FailureOr< AffineValueMap > simplifyConstrainedMinMaxOp(Operation *op, FlatAffineValueConstraints constraints)
Try to simplify the given affine.min or affine.max op to an affine map with a single result and opera...
Definition: Utils.cpp:2066
ParallelOp getParallelForInductionVarOwner(Value val)
Returns the parallel loop parent of an induction variable.
Definition: SCF.cpp:3035
LogicalResult matchForLikeLoop(Value iv, OpFoldResult &lb, OpFoldResult &ub, OpFoldResult &step)
Match "for loop"-like operations from the SCF dialect.
LogicalResult canonicalizeMinMaxOpInLoop(RewriterBase &rewriter, Operation *op, LoopMatcherFn loopMatcher)
Try to canonicalize the given affine.min/max operation in the context of for loops with a known range...
LogicalResult rewritePeeledMinMaxOp(RewriterBase &rewriter, Operation *op, Value iv, Value ub, Value step, bool insideLoop)
Try to simplify the given affine.min/max operation op after loop peeling.
LogicalResult addLoopRangeConstraints(affine::FlatAffineValueConstraints &cstr, Value iv, OpFoldResult lb, OpFoldResult ub, OpFoldResult step)
Populate the given constraint set with induction variable constraints of a "for" loop with the given ...
ForOp getForInductionVarOwner(Value val)
Returns the loop parent of an induction variable.
Definition: SCF.cpp:597
ForallOp getForallOpThreadIndexOwner(Value val)
Returns the ForallOp parent of an thread index variable.
Definition: SCF.cpp:1455
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
std::optional< int64_t > getConstantIntValue(OpFoldResult ofr)
If ofr is a constant integer or an IntegerAttr, return the integer.
int64_t floorDiv(int64_t lhs, int64_t rhs)
Returns the result of MLIR's floordiv operation on constants.
Definition: MathExtras.h:33
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
Definition: LogicalResult.h:72
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26