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-exception
6//
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 AffineApplyOpInterface
21 : 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 logic
31 // can see opportunities for simplifying the affine map that
32 // `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(idx: 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 AffineMinOpInterface
53 : 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 AffineMaxOpInterface
74 : 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
94} // namespace
95} // namespace mlir
96
97void mlir::affine::registerValueBoundsOpInterfaceExternalModels(
98 DialectRegistry &registry) {
99 registry.addExtension(extensionFn: +[](MLIRContext *ctx, AffineDialect *dialect) {
100 AffineApplyOp::attachInterface<AffineApplyOpInterface>(*ctx);
101 AffineMaxOp::attachInterface<AffineMaxOpInterface>(*ctx);
102 AffineMinOp::attachInterface<AffineMinOpInterface>(*ctx);
103 });
104}
105
106FailureOr<int64_t>
107mlir::affine::fullyComposeAndComputeConstantDelta(Value value1, Value value2) {
108 assert(value1.getType().isIndex() && "expected index type");
109 assert(value2.getType().isIndex() && "expected index type");
110
111 // Subtract the two values/dimensions from each other. If the result is 0,
112 // both are equal.
113 Builder b(value1.getContext());
114 AffineMap map = AffineMap::get(/*dimCount=*/2, /*symbolCount=*/0,
115 result: b.getAffineDimExpr(position: 0) - b.getAffineDimExpr(position: 1));
116 // Fully compose the affine map with other ops because the folding logic
117 // can see opportunities for simplifying the affine map that
118 // `FlatLinearConstraints` can currently not see.
119 SmallVector<Value> mapOperands;
120 mapOperands.push_back(Elt: value1);
121 mapOperands.push_back(Elt: value2);
122 affine::fullyComposeAffineMapAndOperands(map: &map, operands: &mapOperands);
123 return ValueBoundsConstraintSet::computeConstantBound(
124 type: presburger::BoundType::EQ,
125 var: ValueBoundsConstraintSet::Variable(map, mapOperands));
126}
127

source code of mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp