MLIR  21.0.0git
ValueBoundsOpInterfaceImpl.cpp
Go to the documentation of this file.
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 
10 
13 
14 using namespace mlir;
15 using namespace mlir::affine;
16 
17 namespace mlir {
18 namespace {
19 
20 struct 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(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 
52 struct 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 
73 struct 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 struct AffineDelinearizeIndexOpInterface
95  : 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 
123 struct AffineLinearizeIndexOpInterface
124  : 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 } // namespace
153 } // namespace mlir
154 
156  DialectRegistry &registry) {
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 
168 FailureOr<int64_t>
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 logic
179  // can see opportunities for simplifying the affine map that
180  // `FlatLinearConstraints` can currently not see.
181  SmallVector<Value> mapOperands;
182  mapOperands.push_back(value1);
183  mapOperands.push_back(value2);
184  affine::fullyComposeAffineMapAndOperands(&map, &mapOperands);
187  ValueBoundsConstraintSet::Variable(map, mapOperands));
188 }
Base type for affine expression.
Definition: AffineExpr.h:68
AffineExpr replaceDimsAndSymbols(ArrayRef< AffineExpr > dimReplacements, ArrayRef< AffineExpr > symReplacements) const
This method substitutes any uses of dimensions and symbols (e.g.
Definition: AffineExpr.cpp:89
AffineExpr floorDiv(uint64_t v) const
Definition: AffineExpr.cpp:917
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:46
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
unsigned getNumSymbols() const
Definition: AffineMap.cpp:398
unsigned getNumDims() const
Definition: AffineMap.cpp:394
AffineExpr getResult(unsigned idx) const
Definition: AffineMap.cpp:411
This class is a general helper class for creating context-global objects like types,...
Definition: Builders.h:51
AffineExpr getAffineDimExpr(unsigned position)
Definition: Builders.cpp:360
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
bool addExtension(TypeID extensionID, std::unique_ptr< DialectExtensionBase > extension)
Add the given extension to the registry.
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
This class represents a single result from folding an operation.
Definition: OpDefinition.h:268
This class implements the operand iterators for the Operation class.
Definition: ValueRange.h:42
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
Definition: Operation.h:407
bool isIndex() const
Definition: Types.cpp:45
A variable that can be added to the constraint set as a "column".
A helper class to be used with ValueBoundsOpInterface.
AffineExpr getExpr(Value value, std::optional< int64_t > dim=std::nullopt)
Return an expression that represents the given index-typed value or shaped value dimension.
BoundBuilder bound(Value value)
Add a bound for the given index-typed value or shaped value.
static FailureOr< int64_t > computeConstantBound(presburger::BoundType type, const Variable &var, StopConditionFn stopCondition=nullptr, bool closedUB=false)
Compute a constant bound for the given variable.
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:132
Type getType() const
Return the type of this value.
Definition: Value.h:129
void fullyComposeAffineMapAndOperands(AffineMap *map, SmallVectorImpl< Value > *operands)
Given an affine map map and its input operands, this method composes into map, maps of AffineApplyOps...
Definition: AffineOps.cpp:1138
void registerValueBoundsOpInterfaceExternalModels(DialectRegistry &registry)
FailureOr< int64_t > fullyComposeAndComputeConstantDelta(Value value1, Value value2)
Compute a constant delta of the given two values.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:344
Include the generated interface declarations.
OpFoldResult getAsOpFoldResult(Value val)
Given a value, try to extract a constant Attribute.