226 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// Utility functions to canonicalize affine ops within SCF op regions.10//11//===----------------------------------------------------------------------===//12 13#include <utility>14 15#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"16#include "mlir/Dialect/Affine/Analysis/Utils.h"17#include "mlir/Dialect/Affine/IR/AffineOps.h"18#include "mlir/Dialect/Affine/IR/AffineValueMap.h"19#include "mlir/Dialect/SCF/IR/SCF.h"20#include "mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h"21#include "mlir/Dialect/Utils/StaticValueUtils.h"22#include "mlir/IR/AffineMap.h"23#include "mlir/IR/PatternMatch.h"24 25#define DEBUG_TYPE "mlir-scf-affine-utils"26 27using namespace mlir;28using namespace affine;29using namespace presburger;30 31LogicalResult scf::matchForLikeLoop(Value iv, OpFoldResult &lb,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 64static FailureOr<AffineApplyOp>65canonicalizeMinMaxOp(RewriterBase &rewriter, Operation *op,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 77LogicalResult scf::addLoopRangeConstraints(FlatAffineValueConstraints &cstr,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 bound111 AffineExpr ivUb;112 if (lbInt && ubInt && (*lbInt + *stepInt >= *ubInt)) {113 // The loop has at most one iteration.114 // iv < lb + 1115 // TODO: Try to derive this constraint by simplifying the expression in116 // 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) + 1121 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 known137/// range. Call `canonicalizeMinMaxOp` and add the following constraints to138/// the constraint system (along with the missing dimensions):139///140/// * iv >= lb141/// * iv < lb + step * ((ub - lb - 1) floorDiv step) + 1142///143/// Note: Due to limitations of IntegerPolyhedron, only constant step sizes144/// are currently supported.145LogicalResult scf::canonicalizeMinMaxOpInLoop(RewriterBase &rewriter,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 loop158 // 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 previous174/// 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 = %step182/// ```183/// min/max operations inside the partial iteration are rewritten in a similar184/// way.185///186/// This function builds up a set of constraints, capable of proving that:187/// * Inside the peeled loop: min(step, ub - iv) == step188/// * Inside the partial iteration: min(step, ub - iv) == ub - iv189///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 for195/// affine.min ops inside the partial iteration. For an explanation of the other196/// parameters, see comment of `canonicalizeMinMaxOpInLoop`.197LogicalResult 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 that209 // 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 variable218 // value of the very last iteration (in the unpeeled loop). At that point,219 // there are less than `step` elements remaining. (Otherwise, the peeled220 // 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}226