189 lines · cpp
1//===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===//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#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"10 11#include "mlir/Dialect/Affine/IR/AffineOps.h"12#include "mlir/Interfaces/ValueBoundsOpInterface.h"13 14using namespace mlir;15using namespace mlir::affine;16 17namespace mlir {18namespace {19 20struct AffineApplyOpInterface21 : public ValueBoundsOpInterface::ExternalModel<AffineApplyOpInterface,22 AffineApplyOp> {23 void populateBoundsForIndexValue(Operation *op, Value value,24 ValueBoundsConstraintSet &cstr) const {25 auto applyOp = cast<AffineApplyOp>(op);26 assert(value == applyOp.getResult() && "invalid value");27 assert(applyOp.getAffineMap().getNumResults() == 1 &&28 "expected single result");29 30 // Fully compose this affine.apply with other ops because the folding logic31 // can see opportunities for simplifying the affine map that32 // `FlatLinearConstraints` can currently not see.33 AffineMap map = applyOp.getAffineMap();34 SmallVector<Value> operands = llvm::to_vector(applyOp.getOperands());35 fullyComposeAffineMapAndOperands(&map, &operands);36 37 // Align affine map result with dims/symbols in the constraint set.38 AffineExpr expr = map.getResult(0);39 SmallVector<AffineExpr> dimReplacements, symReplacements;40 for (int64_t i = 0, e = map.getNumDims(); i < e; ++i)41 dimReplacements.push_back(cstr.getExpr(operands[i]));42 for (int64_t i = map.getNumDims(),43 e = map.getNumDims() + map.getNumSymbols();44 i < e; ++i)45 symReplacements.push_back(cstr.getExpr(operands[i]));46 AffineExpr bound =47 expr.replaceDimsAndSymbols(dimReplacements, symReplacements);48 cstr.bound(value) == bound;49 }50};51 52struct AffineMinOpInterface53 : public ValueBoundsOpInterface::ExternalModel<AffineMinOpInterface,54 AffineMinOp> {55 void populateBoundsForIndexValue(Operation *op, Value value,56 ValueBoundsConstraintSet &cstr) const {57 auto minOp = cast<AffineMinOp>(op);58 assert(value == minOp.getResult() && "invalid value");59 60 // Align affine map results with dims/symbols in the constraint set.61 for (AffineExpr expr : minOp.getAffineMap().getResults()) {62 SmallVector<AffineExpr> dimReplacements = llvm::to_vector(llvm::map_range(63 minOp.getDimOperands(), [&](Value v) { return cstr.getExpr(v); }));64 SmallVector<AffineExpr> symReplacements = llvm::to_vector(llvm::map_range(65 minOp.getSymbolOperands(), [&](Value v) { return cstr.getExpr(v); }));66 AffineExpr bound =67 expr.replaceDimsAndSymbols(dimReplacements, symReplacements);68 cstr.bound(value) <= bound;69 }70 };71};72 73struct AffineMaxOpInterface74 : public ValueBoundsOpInterface::ExternalModel<AffineMaxOpInterface,75 AffineMaxOp> {76 void populateBoundsForIndexValue(Operation *op, Value value,77 ValueBoundsConstraintSet &cstr) const {78 auto maxOp = cast<AffineMaxOp>(op);79 assert(value == maxOp.getResult() && "invalid value");80 81 // Align affine map results with dims/symbols in the constraint set.82 for (AffineExpr expr : maxOp.getAffineMap().getResults()) {83 SmallVector<AffineExpr> dimReplacements = llvm::to_vector(llvm::map_range(84 maxOp.getDimOperands(), [&](Value v) { return cstr.getExpr(v); }));85 SmallVector<AffineExpr> symReplacements = llvm::to_vector(llvm::map_range(86 maxOp.getSymbolOperands(), [&](Value v) { return cstr.getExpr(v); }));87 AffineExpr bound =88 expr.replaceDimsAndSymbols(dimReplacements, symReplacements);89 cstr.bound(value) >= bound;90 }91 };92};93 94struct AffineDelinearizeIndexOpInterface95 : public ValueBoundsOpInterface::ExternalModel<96 AffineDelinearizeIndexOpInterface, AffineDelinearizeIndexOp> {97 void populateBoundsForIndexValue(Operation *rawOp, Value value,98 ValueBoundsConstraintSet &cstr) const {99 auto op = cast<AffineDelinearizeIndexOp>(rawOp);100 auto result = cast<OpResult>(value);101 assert(result.getOwner() == rawOp &&102 "bounded value isn't a result of this delinearize_index");103 unsigned resIdx = result.getResultNumber();104 105 AffineExpr linearIdx = cstr.getExpr(op.getLinearIndex());106 107 SmallVector<OpFoldResult> basis = op.getPaddedBasis();108 AffineExpr divisor = cstr.getExpr(1);109 for (OpFoldResult basisElem : llvm::drop_begin(basis, resIdx + 1))110 divisor = divisor * cstr.getExpr(basisElem);111 112 if (resIdx == 0) {113 cstr.bound(value) == linearIdx.floorDiv(divisor);114 if (!basis.front().isNull())115 cstr.bound(value) < cstr.getExpr(basis.front());116 return;117 }118 AffineExpr thisBasis = cstr.getExpr(basis[resIdx]);119 cstr.bound(value) == (linearIdx % (thisBasis * divisor)).floorDiv(divisor);120 }121};122 123struct AffineLinearizeIndexOpInterface124 : public ValueBoundsOpInterface::ExternalModel<125 AffineLinearizeIndexOpInterface, AffineLinearizeIndexOp> {126 void populateBoundsForIndexValue(Operation *rawOp, Value value,127 ValueBoundsConstraintSet &cstr) const {128 auto op = cast<AffineLinearizeIndexOp>(rawOp);129 assert(value == op.getResult() &&130 "value isn't the result of this linearize");131 132 AffineExpr bound = cstr.getExpr(0);133 AffineExpr stride = cstr.getExpr(1);134 SmallVector<OpFoldResult> basis = op.getPaddedBasis();135 OperandRange multiIndex = op.getMultiIndex();136 unsigned numArgs = multiIndex.size();137 for (auto [revArgNum, length] : llvm::enumerate(llvm::reverse(basis))) {138 unsigned argNum = numArgs - (revArgNum + 1);139 if (argNum == 0)140 break;141 OpFoldResult indexAsFoldRes = getAsOpFoldResult(multiIndex[argNum]);142 bound = bound + cstr.getExpr(indexAsFoldRes) * stride;143 stride = stride * cstr.getExpr(length);144 }145 bound = bound + cstr.getExpr(op.getMultiIndex().front()) * stride;146 cstr.bound(value) == bound;147 if (op.getDisjoint() && !basis.front().isNull()) {148 cstr.bound(value) < stride *cstr.getExpr(basis.front());149 }150 }151};152} // namespace153} // namespace mlir154 155void mlir::affine::registerValueBoundsOpInterfaceExternalModels(156 DialectRegistry ®istry) {157 registry.addExtension(+[](MLIRContext *ctx, AffineDialect *dialect) {158 AffineApplyOp::attachInterface<AffineApplyOpInterface>(*ctx);159 AffineMaxOp::attachInterface<AffineMaxOpInterface>(*ctx);160 AffineMinOp::attachInterface<AffineMinOpInterface>(*ctx);161 AffineDelinearizeIndexOp::attachInterface<162 AffineDelinearizeIndexOpInterface>(*ctx);163 AffineLinearizeIndexOp::attachInterface<AffineLinearizeIndexOpInterface>(164 *ctx);165 });166}167 168FailureOr<int64_t>169mlir::affine::fullyComposeAndComputeConstantDelta(Value value1, Value value2) {170 assert(value1.getType().isIndex() && "expected index type");171 assert(value2.getType().isIndex() && "expected index type");172 173 // Subtract the two values/dimensions from each other. If the result is 0,174 // both are equal.175 Builder b(value1.getContext());176 AffineMap map = AffineMap::get(/*dimCount=*/2, /*symbolCount=*/0,177 b.getAffineDimExpr(0) - b.getAffineDimExpr(1));178 // Fully compose the affine map with other ops because the folding logic179 // can see opportunities for simplifying the affine map that180 // `FlatLinearConstraints` can currently not see.181 SmallVector<Value> mapOperands;182 mapOperands.push_back(value1);183 mapOperands.push_back(value2);184 affine::fullyComposeAffineMapAndOperands(&map, &mapOperands);185 return ValueBoundsConstraintSet::computeConstantBound(186 presburger::BoundType::EQ,187 ValueBoundsConstraintSet::Variable(map, mapOperands));188}189