MLIR  22.0.0git
ScalableValueBoundsConstraintSet.cpp
Go to the documentation of this file.
1 //===- ScalableValueBoundsConstraintSet.cpp - Scalable Value Bounds -------===//
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 
11 namespace mlir::vector {
12 
13 FailureOr<ConstantOrScalableBound::BoundSize>
15  if (map.isSingleConstant())
16  return BoundSize{map.getSingleConstantResult(), /*scalable=*/false};
17  if (map.getNumResults() != 1 || map.getNumInputs() != 1)
18  return failure();
19  auto binop = dyn_cast<AffineBinaryOpExpr>(map.getResult(0));
20  if (!binop || binop.getKind() != AffineExprKind::Mul)
21  return failure();
22  auto matchConstant = [&](AffineExpr expr, int64_t &constant) -> bool {
23  if (auto cst = dyn_cast<AffineConstantExpr>(expr)) {
24  constant = cst.getValue();
25  return true;
26  }
27  return false;
28  };
29  // Match `s0 * cst` or `cst * s0`:
30  int64_t cst = 0;
31  auto lhs = binop.getLHS();
32  auto rhs = binop.getRHS();
33  if ((matchConstant(lhs, cst) && isa<AffineSymbolExpr>(rhs)) ||
34  (matchConstant(rhs, cst) && isa<AffineSymbolExpr>(lhs))) {
35  return BoundSize{cst, /*scalable=*/true};
36  }
37  return failure();
38 }
39 
41 
42 FailureOr<ConstantOrScalableBound>
44  Value value, std::optional<int64_t> dim, unsigned vscaleMin,
45  unsigned vscaleMax, presburger::BoundType boundType, bool closedUB,
46  StopConditionFn stopCondition) {
47  using namespace presburger;
48  assert(vscaleMin <= vscaleMax);
49 
50  // No stop condition specified: Keep adding constraints until the worklist
51  // is empty.
52  auto defaultStopCondition = [&](Value v, std::optional<int64_t> dim,
54  return false;
55  };
56 
58  value.getContext(), stopCondition ? stopCondition : defaultStopCondition,
59  vscaleMin, vscaleMax);
60  int64_t pos = scalableCstr.insert(value, dim, /*isSymbol=*/false);
61  scalableCstr.processWorklist();
62 
63  // Check the resulting constraints set is valid.
64  if (scalableCstr.cstr.isEmpty()) {
65  return failure();
66  }
67 
68  // Project out all columns apart from vscale and the starting point
69  // (value/dim). This should result in constraints in terms of vscale only.
70  auto projectOutFn = [&](ValueDim p) {
71  bool isStartingPoint =
72  p.first == value &&
73  p.second == dim.value_or(ValueBoundsConstraintSet::kIndexValue);
74  return p.first != scalableCstr.getVscaleValue() && !isStartingPoint;
75  };
76  scalableCstr.projectOut(projectOutFn);
77  scalableCstr.projectOutAnonymous(/*except=*/pos);
78  // Also project out local variables (these are not tracked by the
79  // ValueBoundsConstraintSet).
80  for (unsigned i = 0, e = scalableCstr.cstr.getNumLocalVars(); i < e; ++i) {
81  scalableCstr.cstr.projectOut(scalableCstr.cstr.getNumDimAndSymbolVars());
82  }
83 
84  assert(scalableCstr.cstr.getNumDimAndSymbolVars() ==
85  scalableCstr.positionToValueDim.size() &&
86  "inconsistent mapping state");
87 
88  // Check that the only columns left are vscale and the starting point.
89  for (int64_t i = 0; i < scalableCstr.cstr.getNumDimAndSymbolVars(); ++i) {
90  if (i == pos)
91  continue;
92  if (scalableCstr.positionToValueDim[i] !=
93  ValueDim(scalableCstr.getVscaleValue(),
95  return failure();
96  }
97  }
98 
99  SmallVector<AffineMap, 1> lowerBound(1), upperBound(1);
100  scalableCstr.cstr.getSliceBounds(pos, 1, value.getContext(), &lowerBound,
101  &upperBound, closedUB);
102 
103  auto invalidBound = [](auto &bound) {
104  return !bound[0] || bound[0].getNumResults() != 1;
105  };
106 
107  AffineMap bound = [&] {
108  if (boundType == BoundType::EQ && !invalidBound(lowerBound) &&
109  lowerBound[0] == upperBound[0]) {
110  return lowerBound[0];
111  } else if (boundType == BoundType::LB && !invalidBound(lowerBound)) {
112  return lowerBound[0];
113  } else if (boundType == BoundType::UB && !invalidBound(upperBound)) {
114  return upperBound[0];
115  }
116  return AffineMap{};
117  }();
118 
119  if (!bound)
120  return failure();
121 
122  return ConstantOrScalableBound{bound};
123 }
124 
125 } // namespace mlir::vector
function_ref< bool(Region *, ArrayRef< bool > visited)> StopConditionFn
Stop condition for traverseRegionGraph.
Base type for affine expression.
Definition: AffineExpr.h:68
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:46
int64_t getSingleConstantResult() const
Returns the constant result of this map.
Definition: AffineMap.cpp:377
bool isSingleConstant() const
Returns true if this affine map is a single result constant function.
Definition: AffineMap.cpp:369
unsigned getNumResults() const
Definition: AffineMap.cpp:398
unsigned getNumInputs() const
Definition: AffineMap.cpp:399
AffineExpr getResult(unsigned idx) const
Definition: AffineMap.cpp:407
A helper class to be used with ValueBoundsOpInterface.
static constexpr int64_t kIndexValue
Dimension identifier to indicate a value is index-typed.
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:108
BoundType
The type of bound: equal, lower bound or upper bound.
@ Mul
RHS of mul is always a constant or a symbolic expression.
A thin wrapper over an AffineMap which can represent a constant bound, or a scalable bound (in terms ...
FailureOr< BoundSize > getSize() const
Get the (possibly) scalable size of the bound, returns failure if the bound cannot be represented as ...
A version of ValueBoundsConstraintSet that can solve for scalable bounds.
static FailureOr< ConstantOrScalableBound > computeScalableBound(Value value, std::optional< int64_t > dim, unsigned vscaleMin, unsigned vscaleMax, presburger::BoundType boundType, bool closedUB=true, StopConditionFn stopCondition=nullptr)
Computes a (possibly) scalable bound for a given value.
Value getVscaleValue() const
Get the value of vscale. Returns nullptr vscale as not been encountered.