557 lines · cpp
1//===- AffineStructures.cpp - MLIR Affine Structures Class-----------------===//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// Structures for affine/polyhedral analysis of affine dialect ops.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"14#include "mlir/Analysis/Presburger/IntegerRelation.h"15#include "mlir/Analysis/Presburger/Utils.h"16#include "mlir/Dialect/Affine/IR/AffineOps.h"17#include "mlir/Dialect/Affine/IR/AffineValueMap.h"18#include "mlir/Dialect/Utils/StaticValueUtils.h"19#include "mlir/IR/IntegerSet.h"20#include "mlir/Support/LLVM.h"21#include "llvm/ADT/SmallVector.h"22#include "llvm/Support/Debug.h"23#include "llvm/Support/raw_ostream.h"24#include <optional>25 26#define DEBUG_TYPE "affine-structures"27 28using namespace mlir;29using namespace affine;30using namespace presburger;31 32LogicalResult33FlatAffineValueConstraints::addInductionVarOrTerminalSymbol(Value val) {34 if (containsVar(val))35 return success();36 37 // Caller is expected to fully compose map/operands if necessary.38 if (val.getDefiningOp<affine::AffineApplyOp>() ||39 (!isValidSymbol(val) && !isAffineInductionVar(val))) {40 LLVM_DEBUG(llvm::dbgs()41 << "only valid terminal symbols and affine IVs supported\n");42 return failure();43 }44 // Outer loop IVs could be used in forOp's bounds.45 if (auto loop = getForInductionVarOwner(val)) {46 appendDimVar(val);47 if (failed(this->addAffineForOpDomain(loop))) {48 LLVM_DEBUG(49 loop.emitWarning("failed to add domain info to constraint system"));50 return failure();51 }52 return success();53 }54 55 if (auto parallel = getAffineParallelInductionVarOwner(val)) {56 appendDimVar(parallel.getIVs());57 if (failed(this->addAffineParallelOpDomain(parallel))) {58 LLVM_DEBUG(parallel.emitWarning(59 "failed to add domain info to constraint system"));60 return failure();61 }62 return success();63 }64 65 // Add top level symbol.66 appendSymbolVar(val);67 // Check if the symbol is a constant.68 if (std::optional<int64_t> constOp = getConstantIntValue(val))69 addBound(BoundType::EQ, val, constOp.value());70 return success();71}72 73LogicalResult74FlatAffineValueConstraints::addAffineForOpDomain(AffineForOp forOp) {75 unsigned pos;76 // Pre-condition for this method.77 if (!findVar(forOp.getInductionVar(), &pos)) {78 assert(false && "Value not found");79 return failure();80 }81 82 int64_t step = forOp.getStepAsInt();83 if (step != 1) {84 if (!forOp.hasConstantLowerBound())85 LLVM_DEBUG(forOp.emitWarning("domain conservatively approximated"));86 else {87 // Add constraints for the stride.88 // (iv - lb) % step = 0 can be written as:89 // (iv - lb) - step * q = 0 where q = (iv - lb) / step.90 // Add local variable 'q' and add the above equality.91 // The first constraint is q = (iv - lb) floordiv step92 SmallVector<int64_t, 8> dividend(getNumCols(), 0);93 int64_t lb = forOp.getConstantLowerBound();94 dividend[pos] = 1;95 dividend.back() -= lb;96 unsigned qPos = addLocalFloorDiv(dividend, step);97 // Second constraint: (iv - lb) - step * q = 0.98 SmallVector<int64_t, 8> eq(getNumCols(), 0);99 eq[pos] = 1;100 eq.back() -= lb;101 // For the local var just added above.102 eq[qPos] = -step;103 addEquality(eq);104 }105 }106 107 if (forOp.hasConstantLowerBound()) {108 addBound(BoundType::LB, pos, forOp.getConstantLowerBound());109 } else {110 // Non-constant lower bound case.111 if (failed(addBound(BoundType::LB, pos, forOp.getLowerBoundMap(),112 forOp.getLowerBoundOperands())))113 return failure();114 }115 116 if (forOp.hasConstantUpperBound()) {117 addBound(BoundType::UB, pos, forOp.getConstantUpperBound() - 1);118 return success();119 }120 // Non-constant upper bound case.121 return addBound(BoundType::UB, pos, forOp.getUpperBoundMap(),122 forOp.getUpperBoundOperands());123}124 125LogicalResult FlatAffineValueConstraints::addAffineParallelOpDomain(126 AffineParallelOp parallelOp) {127 size_t ivPos = 0;128 for (Value iv : parallelOp.getIVs()) {129 unsigned pos;130 if (!findVar(iv, &pos)) {131 assert(false && "variable expected for the IV value");132 return failure();133 }134 135 AffineMap lowerBound = parallelOp.getLowerBoundMap(ivPos);136 if (lowerBound.isConstant())137 addBound(BoundType::LB, pos, lowerBound.getSingleConstantResult());138 else if (failed(addBound(BoundType::LB, pos, lowerBound,139 parallelOp.getLowerBoundsOperands())))140 return failure();141 142 auto upperBound = parallelOp.getUpperBoundMap(ivPos);143 if (upperBound.isConstant())144 addBound(BoundType::UB, pos, upperBound.getSingleConstantResult() - 1);145 else if (failed(addBound(BoundType::UB, pos, upperBound,146 parallelOp.getUpperBoundsOperands())))147 return failure();148 ++ivPos;149 }150 return success();151}152 153LogicalResult154FlatAffineValueConstraints::addDomainFromSliceMaps(ArrayRef<AffineMap> lbMaps,155 ArrayRef<AffineMap> ubMaps,156 ArrayRef<Value> operands) {157 assert(lbMaps.size() == ubMaps.size());158 assert(lbMaps.size() <= getNumDimVars());159 160 for (unsigned i = 0, e = lbMaps.size(); i < e; ++i) {161 AffineMap lbMap = lbMaps[i];162 AffineMap ubMap = ubMaps[i];163 assert(!lbMap || lbMap.getNumInputs() == operands.size());164 assert(!ubMap || ubMap.getNumInputs() == operands.size());165 166 // Check if this slice is just an equality along this dimension. If so,167 // retrieve the existing loop it equates to and add it to the system.168 if (lbMap && ubMap && lbMap.getNumResults() == 1 &&169 ubMap.getNumResults() == 1 &&170 lbMap.getResult(0) + 1 == ubMap.getResult(0) &&171 // The condition above will be true for maps describing a single172 // iteration (e.g., lbMap.getResult(0) = 0, ubMap.getResult(0) = 1).173 // Make sure we skip those cases by checking that the lb result is not174 // just a constant.175 !isa<AffineConstantExpr>(lbMap.getResult(0))) {176 // Limited support: we expect the lb result to be just a loop dimension.177 // Not supported otherwise for now.178 AffineDimExpr result = dyn_cast<AffineDimExpr>(lbMap.getResult(0));179 if (!result)180 return failure();181 182 AffineForOp loop =183 getForInductionVarOwner(operands[result.getPosition()]);184 if (!loop)185 return failure();186 187 if (failed(addAffineForOpDomain(loop)))188 return failure();189 continue;190 }191 192 // This slice refers to a loop that doesn't exist in the IR yet. Add its193 // bounds to the system assuming its dimension variable position is the194 // same as the position of the loop in the loop nest.195 if (lbMap && failed(addBound(BoundType::LB, i, lbMap, operands)))196 return failure();197 if (ubMap && failed(addBound(BoundType::UB, i, ubMap, operands)))198 return failure();199 }200 return success();201}202 203void FlatAffineValueConstraints::addAffineIfOpDomain(AffineIfOp ifOp) {204 IntegerSet set = ifOp.getIntegerSet();205 // Canonicalize set and operands to ensure unique values for206 // FlatAffineValueConstraints below and for early simplification.207 SmallVector<Value> operands(ifOp.getOperands());208 canonicalizeSetAndOperands(&set, &operands);209 210 // Create the base constraints from the integer set attached to ifOp.211 FlatAffineValueConstraints cst(set, operands);212 213 // Merge the constraints from ifOp to the current domain. We need first merge214 // and align the IDs from both constraints, and then append the constraints215 // from the ifOp into the current one.216 mergeAndAlignVarsWithOther(0, &cst);217 append(cst);218}219 220LogicalResult FlatAffineValueConstraints::addBound(BoundType type, unsigned pos,221 AffineMap boundMap,222 ValueRange boundOperands) {223 // Fully compose map and operands; canonicalize and simplify so that we224 // transitively get to terminal symbols or loop IVs.225 auto map = boundMap;226 SmallVector<Value, 4> operands(boundOperands.begin(), boundOperands.end());227 fullyComposeAffineMapAndOperands(&map, &operands);228 map = simplifyAffineMap(map);229 canonicalizeMapAndOperands(&map, &operands);230 for (Value operand : operands) {231 if (failed(addInductionVarOrTerminalSymbol(operand)))232 return failure();233 }234 return addBound(type, pos, computeAlignedMap(map, operands));235}236 237// Adds slice lower bounds represented by lower bounds in 'lbMaps' and upper238// bounds in 'ubMaps' to each value in `values' that appears in the constraint239// system. Note that both lower/upper bounds share the same operand list240// 'operands'.241// This function assumes 'values.size' == 'lbMaps.size' == 'ubMaps.size', and242// skips any null AffineMaps in 'lbMaps' or 'ubMaps'.243// Note that both lower/upper bounds use operands from 'operands'.244// Returns failure for unimplemented cases such as semi-affine expressions or245// expressions with mod/floordiv.246LogicalResult FlatAffineValueConstraints::addSliceBounds(247 ArrayRef<Value> values, ArrayRef<AffineMap> lbMaps,248 ArrayRef<AffineMap> ubMaps, ArrayRef<Value> operands) {249 assert(values.size() == lbMaps.size());250 assert(lbMaps.size() == ubMaps.size());251 252 for (unsigned i = 0, e = lbMaps.size(); i < e; ++i) {253 unsigned pos;254 if (!findVar(values[i], &pos))255 continue;256 257 AffineMap lbMap = lbMaps[i];258 AffineMap ubMap = ubMaps[i];259 assert(!lbMap || lbMap.getNumInputs() == operands.size());260 assert(!ubMap || ubMap.getNumInputs() == operands.size());261 262 // Check if this slice is just an equality along this dimension.263 if (lbMap && ubMap && lbMap.getNumResults() == 1 &&264 ubMap.getNumResults() == 1 &&265 lbMap.getResult(0) + 1 == ubMap.getResult(0)) {266 if (failed(addBound(BoundType::EQ, pos, lbMap, operands)))267 return failure();268 continue;269 }270 271 // If lower or upper bound maps are null or provide no results, it implies272 // that the source loop was not at all sliced, and the entire loop will be a273 // part of the slice.274 if (lbMap && lbMap.getNumResults() != 0 && ubMap &&275 ubMap.getNumResults() != 0) {276 if (failed(addBound(BoundType::LB, pos, lbMap, operands)))277 return failure();278 if (failed(addBound(BoundType::UB, pos, ubMap, operands)))279 return failure();280 } else {281 auto loop = getForInductionVarOwner(values[i]);282 if (failed(this->addAffineForOpDomain(loop)))283 return failure();284 }285 }286 return success();287}288 289LogicalResult290FlatAffineValueConstraints::composeMap(const AffineValueMap *vMap) {291 return composeMatchingMap(292 computeAlignedMap(vMap->getAffineMap(), vMap->getOperands()));293}294 295// Turn a symbol into a dimension.296static void turnSymbolIntoDim(FlatAffineValueConstraints *cst, Value value) {297 unsigned pos;298 if (cst->findVar(value, &pos) && pos >= cst->getNumDimVars() &&299 pos < cst->getNumDimAndSymbolVars()) {300 cst->swapVar(pos, cst->getNumDimVars());301 cst->setDimSymbolSeparation(cst->getNumSymbolVars() - 1);302 }303}304 305// Changes all symbol variables which are loop IVs to dim variables.306void FlatAffineValueConstraints::convertLoopIVSymbolsToDims() {307 // Gather all symbols which are loop IVs.308 SmallVector<Value, 4> loopIVs;309 for (unsigned i = getNumDimVars(), e = getNumDimAndSymbolVars(); i < e; i++) {310 if (hasValue(i) && getForInductionVarOwner(getValue(i)))311 loopIVs.push_back(getValue(i));312 }313 // Turn each symbol in 'loopIVs' into a dim variable.314 for (auto iv : loopIVs) {315 turnSymbolIntoDim(this, iv);316 }317}318 319void FlatAffineValueConstraints::getIneqAsAffineValueMap(320 unsigned pos, unsigned ineqPos, AffineValueMap &vmap,321 MLIRContext *context) const {322 unsigned numDims = getNumDimVars();323 unsigned numSyms = getNumSymbolVars();324 325 assert(pos < numDims && "invalid position");326 assert(ineqPos < getNumInequalities() && "invalid inequality position");327 328 // Get expressions for local vars.329 SmallVector<AffineExpr, 8> memo(getNumVars(), AffineExpr());330 if (failed(computeLocalVars(memo, context)))331 assert(false &&332 "one or more local exprs do not have an explicit representation");333 auto localExprs = ArrayRef<AffineExpr>(memo).take_back(getNumLocalVars());334 335 // Compute the AffineExpr lower/upper bound for this inequality.336 SmallVector<int64_t, 8> inequality = getInequality64(ineqPos);337 SmallVector<int64_t, 8> bound;338 bound.reserve(getNumCols() - 1);339 // Everything other than the coefficient at `pos`.340 bound.append(inequality.begin(), inequality.begin() + pos);341 bound.append(inequality.begin() + pos + 1, inequality.end());342 343 if (inequality[pos] > 0)344 // Lower bound.345 llvm::transform(bound, bound.begin(), std::negate<int64_t>());346 else347 // Upper bound (which is exclusive).348 bound.back() += 1;349 350 // Convert to AffineExpr (tree) form.351 auto boundExpr = getAffineExprFromFlatForm(bound, numDims - 1, numSyms,352 localExprs, context);353 354 // Get the values to bind to this affine expr (all dims and symbols).355 SmallVector<Value, 4> operands;356 getValues(0, pos, &operands);357 SmallVector<Value, 4> trailingOperands;358 getValues(pos + 1, getNumDimAndSymbolVars(), &trailingOperands);359 operands.append(trailingOperands.begin(), trailingOperands.end());360 vmap.reset(AffineMap::get(numDims - 1, numSyms, boundExpr), operands);361}362 363FlatAffineValueConstraints FlatAffineRelation::getDomainSet() const {364 FlatAffineValueConstraints domain = *this;365 // Convert all range variables to local variables.366 domain.convertToLocal(VarKind::SetDim, getNumDomainDims(),367 getNumDomainDims() + getNumRangeDims());368 return domain;369}370 371FlatAffineValueConstraints FlatAffineRelation::getRangeSet() const {372 FlatAffineValueConstraints range = *this;373 // Convert all domain variables to local variables.374 range.convertToLocal(VarKind::SetDim, 0, getNumDomainDims());375 return range;376}377 378void FlatAffineRelation::compose(const FlatAffineRelation &other) {379 assert(getNumDomainDims() == other.getNumRangeDims() &&380 "Domain of this and range of other do not match");381 assert(space.getDomainSpace().isAligned(other.getSpace().getRangeSpace()) &&382 "Values of domain of this and range of other do not match");383 384 FlatAffineRelation rel = other;385 386 // Convert `rel` from387 // [otherDomain] -> [otherRange]388 // to389 // [otherDomain] -> [otherRange thisRange]390 // and `this` from391 // [thisDomain] -> [thisRange]392 // to393 // [otherDomain thisDomain] -> [thisRange].394 unsigned removeDims = rel.getNumRangeDims();395 insertDomainVar(0, rel.getNumDomainDims());396 rel.appendRangeVar(getNumRangeDims());397 398 // Merge symbol and local variables.399 mergeSymbolVars(rel);400 mergeLocalVars(rel);401 402 // Convert `rel` from [otherDomain] -> [otherRange thisRange] to403 // [otherDomain] -> [thisRange] by converting first otherRange range vars404 // to local vars.405 rel.convertToLocal(VarKind::SetDim, rel.getNumDomainDims(),406 rel.getNumDomainDims() + removeDims);407 // Convert `this` from [otherDomain thisDomain] -> [thisRange] to408 // [otherDomain] -> [thisRange] by converting last thisDomain domain vars409 // to local vars.410 convertToLocal(VarKind::SetDim, getNumDomainDims() - removeDims,411 getNumDomainDims());412 413 auto thisMaybeValues = getMaybeValues(VarKind::SetDim);414 auto relMaybeValues = rel.getMaybeValues(VarKind::SetDim);415 416 // Add and match domain of `rel` to domain of `this`.417 for (unsigned i = 0, e = rel.getNumDomainDims(); i < e; ++i)418 if (relMaybeValues[i].has_value())419 setValue(i, *relMaybeValues[i]);420 // Add and match range of `this` to range of `rel`.421 for (unsigned i = 0, e = getNumRangeDims(); i < e; ++i) {422 unsigned rangeIdx = rel.getNumDomainDims() + i;423 if (thisMaybeValues[rangeIdx].has_value())424 rel.setValue(rangeIdx, *thisMaybeValues[rangeIdx]);425 }426 427 // Append `this` to `rel` and simplify constraints.428 rel.append(*this);429 rel.removeRedundantLocalVars();430 431 *this = rel;432}433 434void FlatAffineRelation::inverse() {435 unsigned oldDomain = getNumDomainDims();436 unsigned oldRange = getNumRangeDims();437 // Add new range vars.438 appendRangeVar(oldDomain);439 // Swap new vars with domain.440 for (unsigned i = 0; i < oldDomain; ++i)441 swapVar(i, oldDomain + oldRange + i);442 // Remove the swapped domain.443 removeVarRange(0, oldDomain);444 // Set domain and range as inverse.445 numDomainDims = oldRange;446 numRangeDims = oldDomain;447}448 449void FlatAffineRelation::insertDomainVar(unsigned pos, unsigned num) {450 assert(pos <= getNumDomainDims() &&451 "Var cannot be inserted at invalid position");452 insertDimVar(pos, num);453 numDomainDims += num;454}455 456void FlatAffineRelation::insertRangeVar(unsigned pos, unsigned num) {457 assert(pos <= getNumRangeDims() &&458 "Var cannot be inserted at invalid position");459 insertDimVar(getNumDomainDims() + pos, num);460 numRangeDims += num;461}462 463void FlatAffineRelation::appendDomainVar(unsigned num) {464 insertDimVar(getNumDomainDims(), num);465 numDomainDims += num;466}467 468void FlatAffineRelation::appendRangeVar(unsigned num) {469 insertDimVar(getNumDimVars(), num);470 numRangeDims += num;471}472 473void FlatAffineRelation::removeVarRange(VarKind kind, unsigned varStart,474 unsigned varLimit) {475 assert(varLimit <= getNumVarKind(kind));476 if (varStart >= varLimit)477 return;478 479 FlatAffineValueConstraints::removeVarRange(kind, varStart, varLimit);480 481 // If kind is not SetDim, domain and range don't need to be updated.482 if (kind != VarKind::SetDim)483 return;484 485 // Compute number of domain and range variables to remove. This is done by486 // intersecting the range of domain/range vars with range of vars to remove.487 unsigned intersectDomainLHS = std::min(varLimit, getNumDomainDims());488 unsigned intersectDomainRHS = varStart;489 unsigned intersectRangeLHS = std::min(varLimit, getNumDimVars());490 unsigned intersectRangeRHS = std::max(varStart, getNumDomainDims());491 492 if (intersectDomainLHS > intersectDomainRHS)493 numDomainDims -= intersectDomainLHS - intersectDomainRHS;494 if (intersectRangeLHS > intersectRangeRHS)495 numRangeDims -= intersectRangeLHS - intersectRangeRHS;496}497 498LogicalResult mlir::affine::getRelationFromMap(AffineMap &map,499 IntegerRelation &rel) {500 // Get flattened affine expressions.501 std::vector<SmallVector<int64_t, 8>> flatExprs;502 FlatAffineValueConstraints localVarCst;503 if (failed(getFlattenedAffineExprs(map, &flatExprs, &localVarCst)))504 return failure();505 506 const unsigned oldDimNum = localVarCst.getNumDimVars();507 const unsigned oldCols = localVarCst.getNumCols();508 const unsigned numRangeVars = map.getNumResults();509 const unsigned numDomainVars = map.getNumDims();510 511 // Add range as the new expressions.512 localVarCst.appendDimVar(numRangeVars);513 514 // Add identifiers to the local constraints as getFlattenedAffineExprs creates515 // a FlatLinearConstraints with no identifiers.516 for (unsigned i = 0, e = localVarCst.getNumDimAndSymbolVars(); i < e; ++i)517 localVarCst.setValue(i, Value());518 519 // Add equalities between source and range.520 SmallVector<int64_t, 8> eq(localVarCst.getNumCols());521 for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) {522 // Zero fill.523 llvm::fill(eq, 0);524 // Fill equality.525 for (unsigned j = 0, f = oldDimNum; j < f; ++j)526 eq[j] = flatExprs[i][j];527 for (unsigned j = oldDimNum, f = oldCols; j < f; ++j)528 eq[j + numRangeVars] = flatExprs[i][j];529 // Set this dimension to -1 to equate lhs and rhs and add equality.530 eq[numDomainVars + i] = -1;531 localVarCst.addEquality(eq);532 }533 534 rel = localVarCst;535 return success();536}537 538LogicalResult mlir::affine::getRelationFromMap(const AffineValueMap &map,539 IntegerRelation &rel) {540 541 AffineMap affineMap = map.getAffineMap();542 if (failed(getRelationFromMap(affineMap, rel)))543 return failure();544 545 // Set identifiers for domain and symbol variables.546 for (unsigned i = 0, e = affineMap.getNumDims(); i < e; ++i)547 rel.setId(VarKind::SetDim, i, Identifier(map.getOperand(i)));548 549 const unsigned mapNumResults = affineMap.getNumResults();550 for (unsigned i = 0, e = rel.getNumSymbolVars(); i < e; ++i)551 rel.setId(552 VarKind::Symbol, i,553 Identifier(map.getOperand(rel.getNumDimVars() + i - mapNumResults)));554 555 return success();556}557