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