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 
16 namespace mlir {
17 namespace arith {
18 namespace {
19 
20 struct AddIOpInterface
21  : public ValueBoundsOpInterface::ExternalModel<AddIOpInterface, AddIOp> {
22  void populateBoundsForIndexValue(Operation *op, Value value,
23  ValueBoundsConstraintSet &cstr) const {
24  auto addIOp = cast<AddIOp>(op);
25  assert(value == addIOp.getResult() && "invalid value");
26 
27  // Note: `getExpr` has a side effect: it may add a new column to the
28  // constraint system. The evaluation order of addition operands is
29  // unspecified in C++. To make sure that all compilers produce the exact
30  // same results (that can be FileCheck'd), it is important that `getExpr`
31  // is called first and assigned to temporary variables, and the addition
32  // is performed afterwards.
33  AffineExpr lhs = cstr.getExpr(addIOp.getLhs());
34  AffineExpr rhs = cstr.getExpr(addIOp.getRhs());
35  cstr.bound(value) == lhs + rhs;
36  }
37 };
38 
39 struct ConstantOpInterface
40  : public ValueBoundsOpInterface::ExternalModel<ConstantOpInterface,
41  ConstantOp> {
42  void populateBoundsForIndexValue(Operation *op, Value value,
43  ValueBoundsConstraintSet &cstr) const {
44  auto constantOp = cast<ConstantOp>(op);
45  assert(value == constantOp.getResult() && "invalid value");
46 
47  if (auto attr = llvm::dyn_cast<IntegerAttr>(constantOp.getValue()))
48  cstr.bound(value) == attr.getInt();
49  }
50 };
51 
52 struct SubIOpInterface
53  : public ValueBoundsOpInterface::ExternalModel<SubIOpInterface, SubIOp> {
54  void populateBoundsForIndexValue(Operation *op, Value value,
55  ValueBoundsConstraintSet &cstr) const {
56  auto subIOp = cast<SubIOp>(op);
57  assert(value == subIOp.getResult() && "invalid value");
58 
59  AffineExpr lhs = cstr.getExpr(subIOp.getLhs());
60  AffineExpr rhs = cstr.getExpr(subIOp.getRhs());
61  cstr.bound(value) == lhs - rhs;
62  }
63 };
64 
65 struct MulIOpInterface
66  : public ValueBoundsOpInterface::ExternalModel<MulIOpInterface, MulIOp> {
67  void populateBoundsForIndexValue(Operation *op, Value value,
68  ValueBoundsConstraintSet &cstr) const {
69  auto mulIOp = cast<MulIOp>(op);
70  assert(value == mulIOp.getResult() && "invalid value");
71 
72  AffineExpr lhs = cstr.getExpr(mulIOp.getLhs());
73  AffineExpr rhs = cstr.getExpr(mulIOp.getRhs());
74  cstr.bound(value) == lhs *rhs;
75  }
76 };
77 
78 struct FloorDivSIOpInterface
79  : public ValueBoundsOpInterface::ExternalModel<FloorDivSIOpInterface,
80  FloorDivSIOp> {
81  void populateBoundsForIndexValue(Operation *op, Value value,
82  ValueBoundsConstraintSet &cstr) const {
83  auto divSIOp = cast<FloorDivSIOp>(op);
84  assert(value == divSIOp.getResult() && "invalid value");
85 
86  AffineExpr lhs = cstr.getExpr(divSIOp.getLhs());
87  AffineExpr rhs = cstr.getExpr(divSIOp.getRhs());
88  cstr.bound(value) == lhs.floorDiv(rhs);
89  }
90 };
91 
92 struct SelectOpInterface
93  : public ValueBoundsOpInterface::ExternalModel<SelectOpInterface,
94  SelectOp> {
95 
96  static void populateBounds(SelectOp selectOp, std::optional<int64_t> dim,
98  Value value = selectOp.getResult();
99  Value condition = selectOp.getCondition();
100  Value trueValue = selectOp.getTrueValue();
101  Value falseValue = selectOp.getFalseValue();
102 
103  if (isa<ShapedType>(condition.getType())) {
104  // If the condition is a shaped type, the condition is applied
105  // element-wise. All three operands must have the same shape.
106  cstr.bound(value)[*dim] == cstr.getExpr(trueValue, dim);
107  cstr.bound(value)[*dim] == cstr.getExpr(falseValue, dim);
108  cstr.bound(value)[*dim] == cstr.getExpr(condition, dim);
109  return;
110  }
111 
112  // Populate constraints for the true/false values (and all values on the
113  // backward slice, as long as the current stop condition is not satisfied).
114  cstr.populateConstraints(trueValue, dim);
115  cstr.populateConstraints(falseValue, dim);
116  auto boundsBuilder = cstr.bound(value);
117  if (dim)
118  boundsBuilder[*dim];
119 
120  // Compare yielded values.
121  // If trueValue <= falseValue:
122  // * result <= falseValue
123  // * result >= trueValue
124  if (cstr.populateAndCompare(
125  /*lhs=*/{trueValue, dim},
126  ValueBoundsConstraintSet::ComparisonOperator::LE,
127  /*rhs=*/{falseValue, dim})) {
128  if (dim) {
129  cstr.bound(value)[*dim] >= cstr.getExpr(trueValue, dim);
130  cstr.bound(value)[*dim] <= cstr.getExpr(falseValue, dim);
131  } else {
132  cstr.bound(value) >= trueValue;
133  cstr.bound(value) <= falseValue;
134  }
135  }
136  // If falseValue <= trueValue:
137  // * result <= trueValue
138  // * result >= falseValue
139  if (cstr.populateAndCompare(
140  /*lhs=*/{falseValue, dim},
141  ValueBoundsConstraintSet::ComparisonOperator::LE,
142  /*rhs=*/{trueValue, dim})) {
143  if (dim) {
144  cstr.bound(value)[*dim] >= cstr.getExpr(falseValue, dim);
145  cstr.bound(value)[*dim] <= cstr.getExpr(trueValue, dim);
146  } else {
147  cstr.bound(value) >= falseValue;
148  cstr.bound(value) <= trueValue;
149  }
150  }
151  }
152 
153  void populateBoundsForIndexValue(Operation *op, Value value,
154  ValueBoundsConstraintSet &cstr) const {
155  populateBounds(cast<SelectOp>(op), /*dim=*/std::nullopt, cstr);
156  }
157 
158  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
159  ValueBoundsConstraintSet &cstr) const {
160  populateBounds(cast<SelectOp>(op), dim, cstr);
161  }
162 };
163 } // namespace
164 } // namespace arith
165 } // namespace mlir
166 
168  DialectRegistry &registry) {
169  registry.addExtension(+[](MLIRContext *ctx, arith::ArithDialect *dialect) {
170  arith::AddIOp::attachInterface<arith::AddIOpInterface>(*ctx);
171  arith::ConstantOp::attachInterface<arith::ConstantOpInterface>(*ctx);
172  arith::SubIOp::attachInterface<arith::SubIOpInterface>(*ctx);
173  arith::MulIOp::attachInterface<arith::MulIOpInterface>(*ctx);
174  arith::FloorDivSIOp::attachInterface<arith::FloorDivSIOpInterface>(*ctx);
175  arith::SelectOp::attachInterface<arith::SelectOpInterface>(*ctx);
176  });
177 }
Base type for affine expression.
Definition: AffineExpr.h:68
AffineExpr floorDiv(uint64_t v) const
Definition: AffineExpr.cpp:921
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
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
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.
void populateConstraints(Value value, std::optional< int64_t > dim)
Traverse the IR starting from the given value/dim and populate constraints as long as the stop condit...
bool populateAndCompare(const Variable &lhs, ComparisonOperator cmp, const Variable &rhs)
Populate constraints for lhs/rhs (until the stop condition is met).
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Type getType() const
Return the type of this value.
Definition: Value.h:105
void registerValueBoundsOpInterfaceExternalModels(DialectRegistry &registry)
Include the generated interface declarations.