[Mlir-commits] [mlir] [mlir][Vector] Add utility for computing scalable value bounds (PR #83876)
Benjamin Maxwell
llvmlistbot at llvm.org
Mon Mar 4 09:36:22 PST 2024
https://github.com/MacDue created https://github.com/llvm/llvm-project/pull/83876
This adds a new API built with the `ValueBoundsConstraintSet` to compute the bounds of possibly scalable quantities. It uses knowledge of the range of vscale (which is defined by the target architecture), to solve for the bound as either a constant or an expression in terms of vscale.
The result is an `AffineMap` will always take at most one parameter, vscale, and return a single result, which is the bound of `value`.
The API is defined as follows:
```c++
FailureOr<ConstantOrScalableBound>
vector::computeScalableBound(Value value, std::optional<int64_t> dim,
unsigned vscaleMin, unsigned vscaleMax,
presburger::BoundType boundType);
```
Note: `ConstantOrScalableBound` is a thin wrapper over the `AffineMap` with a utility for converting the bound to a single quantity (i.e. a size and scalable flag).
We believe this API could prove useful downstream in IREE (which uses a similar analysis to hoist allocas, which currently fails for scalable vectors).
>From ffb5404711c2f3ce1a7340a33bd8796bc6b083b9 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Fri, 1 Mar 2024 18:23:09 +0000
Subject: [PATCH] [mlir][Vector] Add utility for computing scalable value
bounds
This adds a new API built with the `ValueBoundsConstraintSet` to compute
the bounds of possibly scalable quantities. It uses knowledge of the
range of vscale (which is defined by the target architecture), to solve
for the bound as either a constant or an expression in terms of vscale.
The result is an `AffineMap` will always take at most one parameter,
vscale, and return a single result, which is the bound of `value`.
The API is defined as follows:
```c++
FailureOr<ConstantOrScalableBound>
vector::computeScalableBound(Value value, std::optional<int64_t> dim,
unsigned vscaleMin, unsigned vscaleMax,
presburger::BoundType boundType);
```
Note: `ConstantOrScalableBound` is a thin wrapper over the `AffineMap`
with a utility for converting the bound to a single quantity (i.e. a
size and scalable flag).
We believe this API could prove useful downstream in IREE (which uses
a similar analysis to hoist allocas, which currently fails for scalable
vectors).
---
.../mlir/Dialect/Vector/Utils/VectorUtils.h | 29 ++++
.../mlir/Interfaces/ValueBoundsOpInterface.h | 20 ++-
mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp | 130 +++++++++++++++++
.../lib/Interfaces/ValueBoundsOpInterface.cpp | 102 +++++++++----
.../Vector/test-scalable-upper-bound.mlir | 137 ++++++++++++++++++
.../Dialect/Affine/TestReifyValueBounds.cpp | 30 +++-
6 files changed, 417 insertions(+), 31 deletions(-)
create mode 100644 mlir/test/Dialect/Vector/test-scalable-upper-bound.mlir
diff --git a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
index f6b03a0f2c8007..635d609d0b3e71 100644
--- a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
+++ b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
@@ -9,6 +9,7 @@
#ifndef MLIR_DIALECT_VECTOR_UTILS_VECTORUTILS_H_
#define MLIR_DIALECT_VECTOR_UTILS_VECTORUTILS_H_
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/BuiltinAttributes.h"
@@ -98,6 +99,34 @@ bool isContiguousSlice(MemRefType memrefType, VectorType vectorType);
std::optional<StaticTileOffsetRange>
createUnrollIterator(VectorType vType, int64_t targetRank = 1);
+struct ConstantOrScalableBound {
+ AffineMap map;
+
+ struct BoundSize {
+ int64_t baseSize{0};
+ bool scalable{false};
+ };
+
+ /// Get the (possibly) scalable size of the bound, returns failure if the
+ /// bound cannot be represented as a single quantity.
+ FailureOr<BoundSize> getSize() const;
+};
+
+/// Computes a (possibly) scalable bound for a given value. This is similar to
+/// `ValueBoundsConstraintSet::computeConstantBound()`, but uses knowledge of
+/// the range of vscale to compute either a constant bound, an expression in
+/// terms of vscale, or failure if no bound can be computed.
+///
+/// The resulting `AffineMap` will always take at most one parameter, vscale,
+/// and return a single result, which is the bound of `value`.
+///
+/// Note: `vscaleMin` must be `<=` to `vscaleMax`. If `vscaleMin` ==
+/// `vscaleMax`, the resulting bound (if found), will be constant.
+FailureOr<ConstantOrScalableBound>
+computeScalableBound(Value value, std::optional<int64_t> dim,
+ unsigned vscaleMin, unsigned vscaleMax,
+ presburger::BoundType boundType);
+
} // namespace vector
/// Constructs a permutation map of invariant memref indices to vector
diff --git a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
index 28dadfb9ecf868..6d0e16bf215f8a 100644
--- a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
+++ b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
@@ -265,10 +265,28 @@ class ValueBoundsConstraintSet {
ValueBoundsConstraintSet(MLIRContext *ctx);
+ /// A callback to allow injecting custom value bounds constraints.
+ /// It takes the current value, the dim (or kIndexValue), and a reference to
+ /// the constraints set.
+ using PopulateCustomValueBoundsFn =
+ function_ref<void(Value, int64_t, ValueBoundsConstraintSet &)>;
+
+ /// Populates the constraint set for a value/map without actually computing
+ /// the bound.
+ int64_t populateConstraintsSet(
+ Value value, std::optional<int64_t> dim = std::nullopt,
+ PopulateCustomValueBoundsFn customValueBounds = nullptr,
+ StopConditionFn stopCondition = nullptr);
+ int64_t populateConstraintsSet(
+ AffineMap map, ValueDimList mapOperands,
+ PopulateCustomValueBoundsFn customValueBounds = nullptr,
+ StopConditionFn stopCondition = nullptr, int64_t *posOut = nullptr);
+
/// Iteratively process all elements on the worklist until an index-typed
/// value or shaped value meets `stopCondition`. Such values are not processed
/// any further.
- void processWorklist(StopConditionFn stopCondition);
+ void processWorklist(StopConditionFn stopCondition,
+ PopulateCustomValueBoundsFn customValueBounds = nullptr);
/// Bound the given column in the underlying constraint set by the given
/// expression.
diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index d613672608c3ad..77b2ad6ff540ce 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -24,6 +24,7 @@
#include "mlir/IR/IntegerSet.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/TypeUtilities.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/MathExtras.h"
@@ -300,3 +301,132 @@ vector::createUnrollIterator(VectorType vType, int64_t targetRank) {
shapeToUnroll = shapeToUnroll.slice(0, firstScalableDim);
return StaticTileOffsetRange(shapeToUnroll, /*unrollStep=*/1);
}
+
+FailureOr<vector::ConstantOrScalableBound::BoundSize>
+vector::ConstantOrScalableBound::getSize() const {
+ if (map.isSingleConstant())
+ return BoundSize{map.getSingleConstantResult(), /*scalable=*/false};
+ if (map.getNumResults() != 1 || map.getNumInputs() != 1)
+ return failure();
+ auto binop = dyn_cast<AffineBinaryOpExpr>(map.getResult(0));
+ if (!binop || binop.getKind() != AffineExprKind::Mul)
+ return failure();
+ auto matchConstant = [&](AffineExpr expr, int64_t &constant) -> bool {
+ if (auto cst = dyn_cast<AffineConstantExpr>(expr)) {
+ constant = cst.getValue();
+ return true;
+ }
+ return false;
+ };
+ // Match `s0 * cst` or `cst * s0`:
+ int64_t cst = 0;
+ auto lhs = binop.getLHS();
+ auto rhs = binop.getRHS();
+ if ((matchConstant(lhs, cst) && isa<AffineSymbolExpr>(rhs)) ||
+ (matchConstant(rhs, cst) && isa<AffineSymbolExpr>(lhs))) {
+ return BoundSize{cst, /*scalable=*/true};
+ }
+ return failure();
+}
+
+namespace {
+struct ScalableValueBoundsConstraintSet : public ValueBoundsConstraintSet {
+ using ValueBoundsConstraintSet::ValueBoundsConstraintSet;
+
+ static Operation *getOwnerOfValue(Value value) {
+ if (auto bbArg = dyn_cast<BlockArgument>(value))
+ return bbArg.getOwner()->getParentOp();
+ return value.getDefiningOp();
+ }
+
+ static FailureOr<AffineMap>
+ computeScalableBound(Value value, std::optional<int64_t> dim,
+ unsigned vscaleMin, unsigned vscaleMax,
+ presburger::BoundType boundType) {
+ using namespace presburger;
+
+ assert(vscaleMin <= vscaleMax);
+ ScalableValueBoundsConstraintSet cstr(value.getContext());
+
+ Value vscale;
+ int64_t pos = cstr.populateConstraintsSet(
+ value, dim,
+ /* Custom vscale value bounds */
+ [&vscale, vscaleMin, vscaleMax](Value value, int64_t dim,
+ ValueBoundsConstraintSet &cstr) {
+ if (dim != ValueBoundsConstraintSet::kIndexValue)
+ return;
+ if (isa_and_present<vector::VectorScaleOp>(getOwnerOfValue(value))) {
+ if (vscale) {
+ // All copies of vscale are equivalent.
+ cstr.bound(value) == cstr.getExpr(vscale);
+ } else {
+ // We know vscale is confined to [vscaleMin, vscaleMax].
+ cstr.bound(value) >= vscaleMin;
+ cstr.bound(value) <= vscaleMax;
+ vscale = value;
+ }
+ }
+ },
+ /* Stop condition */
+ [](auto, auto) {
+ // Keep adding constraints till the worklist is empty.
+ return false;
+ });
+
+ // Project out all variables apart from the first vscale.
+ cstr.projectOut([&](ValueDim p) { return p.first != vscale; });
+
+ assert(cstr.cstr.getNumDimAndSymbolVars() ==
+ cstr.positionToValueDim.size() &&
+ "inconsistent mapping state");
+
+ for (int64_t i = 0; i < cstr.cstr.getNumDimAndSymbolVars(); ++i) {
+ if (i == pos)
+ continue;
+ if (cstr.positionToValueDim[i] !=
+ ValueDim(vscale, ValueBoundsConstraintSet::kIndexValue)) {
+ return failure();
+ }
+ }
+
+ SmallVector<AffineMap, 1> lowerBound(1), upperBound(1);
+ cstr.cstr.getSliceBounds(pos, 1, value.getContext(), &lowerBound,
+ &upperBound,
+ /*closedUB=*/true);
+
+ auto invalidBound = [](auto &bound) {
+ return !bound[0] || bound[0].getNumResults() != 1;
+ };
+
+ AffineMap bound = [&] {
+ if (boundType == BoundType::EQ && !invalidBound(lowerBound) &&
+ lowerBound[0] == lowerBound[0]) {
+ return lowerBound[0];
+ } else if (boundType == BoundType::LB && !invalidBound(lowerBound)) {
+ return lowerBound[0];
+ } else if (boundType == BoundType::UB && !invalidBound(upperBound)) {
+ return upperBound[0];
+ }
+ return AffineMap{};
+ }();
+
+ if (!bound)
+ return failure();
+
+ return bound;
+ }
+};
+
+} // namespace
+
+FailureOr<vector::ConstantOrScalableBound>
+vector::computeScalableBound(Value value, std::optional<int64_t> dim,
+ unsigned vscaleMin, unsigned vscaleMax,
+ presburger::BoundType boundType) {
+ auto bound = ScalableValueBoundsConstraintSet::computeScalableBound(
+ value, dim, vscaleMin, vscaleMax, boundType);
+ if (failed(bound))
+ return failure();
+ return ConstantOrScalableBound{*bound};
+}
diff --git a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
index 85abc2df894797..ac4e3b935a0542 100644
--- a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
+++ b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
@@ -191,7 +191,9 @@ static Operation *getOwnerOfValue(Value value) {
return value.getDefiningOp();
}
-void ValueBoundsConstraintSet::processWorklist(StopConditionFn stopCondition) {
+void ValueBoundsConstraintSet::processWorklist(
+ StopConditionFn stopCondition,
+ PopulateCustomValueBoundsFn customValueBounds) {
while (!worklist.empty()) {
int64_t pos = worklist.front();
worklist.pop();
@@ -215,8 +217,11 @@ void ValueBoundsConstraintSet::processWorklist(StopConditionFn stopCondition) {
if (stopCondition(value, maybeDim))
continue;
- // Query `ValueBoundsOpInterface` for constraints. New items may be added to
- // the worklist.
+ // 1. Query `customValueBounds` for constraints (if provided).
+ if (customValueBounds)
+ customValueBounds(value, dim, *this);
+
+ // 2. Query `ValueBoundsOpInterface` for constraints.
auto valueBoundsOp =
dyn_cast<ValueBoundsOpInterface>(getOwnerOfValue(value));
if (valueBoundsOp) {
@@ -228,6 +233,8 @@ void ValueBoundsConstraintSet::processWorklist(StopConditionFn stopCondition) {
continue;
}
+ // Steps 1 and 2 above may add new items to the worklist.
+
// If the op does not implement `ValueBoundsOpInterface`, check if it
// implements the `DestinationStyleOpInterface`. OpResults of such ops are
// tied to OpOperands. Tied values have the same shape.
@@ -471,55 +478,92 @@ FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
closedUB);
}
+FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
+ presburger::BoundType type, AffineMap map, ArrayRef<Value> operands,
+ StopConditionFn stopCondition, bool closedUB) {
+ ValueDimList valueDims;
+ for (Value v : operands) {
+ assert(v.getType().isIndex() && "expected index type");
+ valueDims.emplace_back(v, std::nullopt);
+ }
+ return computeConstantBound(type, map, valueDims, stopCondition, closedUB);
+}
+
FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
presburger::BoundType type, AffineMap map, ValueDimList operands,
StopConditionFn stopCondition, bool closedUB) {
assert(map.getNumResults() == 1 && "expected affine map with one result");
ValueBoundsConstraintSet cstr(map.getContext());
- int64_t pos = cstr.insert(/*isSymbol=*/false);
+
+ int64_t pos = 0;
+ if (stopCondition) {
+ cstr.populateConstraintsSet(map, operands, nullptr, stopCondition, &pos);
+ } else {
+ // No stop condition specified: Keep adding constraints until a bound could
+ // be computed.
+ cstr.populateConstraintsSet(
+ map, operands, nullptr,
+ [&](Value v, std::optional<int64_t> dim) {
+ return cstr.cstr.getConstantBound64(type, pos).has_value();
+ },
+ &pos);
+ }
+ // Compute constant bound for `valueDim`.
+ int64_t ubAdjustment = closedUB ? 0 : 1;
+ if (auto bound = cstr.cstr.getConstantBound64(type, pos))
+ return type == BoundType::UB ? *bound + ubAdjustment : *bound;
+ return failure();
+}
+
+int64_t ValueBoundsConstraintSet::populateConstraintsSet(
+ Value value, std::optional<int64_t> dim,
+ PopulateCustomValueBoundsFn customValueBounds,
+ StopConditionFn stopCondition) {
+#ifndef NDEBUG
+ assertValidValueDim(value, dim);
+#endif // NDEBUG
+
+ AffineMap map =
+ AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0,
+ Builder(value.getContext()).getAffineDimExpr(0));
+ return populateConstraintsSet(map, {{value, dim}}, customValueBounds,
+ stopCondition);
+}
+
+int64_t ValueBoundsConstraintSet::populateConstraintsSet(
+ AffineMap map, ValueDimList operands,
+ PopulateCustomValueBoundsFn customValueBounds,
+ StopConditionFn stopCondition, int64_t *posOut) {
+ assert(map.getNumResults() == 1 && "expected affine map with one result");
+ int64_t pos = insert(/*isSymbol=*/false);
+ if (posOut)
+ *posOut = pos;
// Add map and operands to the constraint set. Dimensions are converted to
// symbols. All operands are added to the worklist.
auto mapper = [&](std::pair<Value, std::optional<int64_t>> v) {
- return cstr.getExpr(v.first, v.second);
+ return getExpr(v.first, v.second);
};
SmallVector<AffineExpr> dimReplacements = llvm::to_vector(
llvm::map_range(ArrayRef(operands).take_front(map.getNumDims()), mapper));
SmallVector<AffineExpr> symReplacements = llvm::to_vector(
llvm::map_range(ArrayRef(operands).drop_front(map.getNumDims()), mapper));
- cstr.addBound(
+ addBound(
presburger::BoundType::EQ, pos,
map.getResult(0).replaceDimsAndSymbols(dimReplacements, symReplacements));
// Process the backward slice of `operands` (i.e., reverse use-def chain)
// until `stopCondition` is met.
if (stopCondition) {
- cstr.processWorklist(stopCondition);
+ processWorklist(stopCondition, customValueBounds);
} else {
- // No stop condition specified: Keep adding constraints until a bound could
- // be computed.
- cstr.processWorklist(
- /*stopCondition=*/[&](Value v, std::optional<int64_t> dim) {
- return cstr.cstr.getConstantBound64(type, pos).has_value();
- });
+ // No stop condition specified: Keep adding constraints until the worklist
+ // is empty.
+ processWorklist([](Value v, std::optional<int64_t> dim) { return false; },
+ customValueBounds);
}
- // Compute constant bound for `valueDim`.
- int64_t ubAdjustment = closedUB ? 0 : 1;
- if (auto bound = cstr.cstr.getConstantBound64(type, pos))
- return type == BoundType::UB ? *bound + ubAdjustment : *bound;
- return failure();
-}
-
-FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
- presburger::BoundType type, AffineMap map, ArrayRef<Value> operands,
- StopConditionFn stopCondition, bool closedUB) {
- ValueDimList valueDims;
- for (Value v : operands) {
- assert(v.getType().isIndex() && "expected index type");
- valueDims.emplace_back(v, std::nullopt);
- }
- return computeConstantBound(type, map, valueDims, stopCondition, closedUB);
+ return pos;
}
FailureOr<int64_t>
diff --git a/mlir/test/Dialect/Vector/test-scalable-upper-bound.mlir b/mlir/test/Dialect/Vector/test-scalable-upper-bound.mlir
new file mode 100644
index 00000000000000..2afc4db874b73e
--- /dev/null
+++ b/mlir/test/Dialect/Vector/test-scalable-upper-bound.mlir
@@ -0,0 +1,137 @@
+// RUN: mlir-opt %s -test-affine-reify-value-bounds -cse -verify-diagnostics \
+// RUN: -verify-diagnostics -split-input-file | FileCheck %s
+
+#fixedDim0Map = affine_map<(d0)[s0] -> (-d0 + 32400, s0)>
+#fixedDim1Map = affine_map<(d0)[s0] -> (-d0 + 16, s0)>
+
+// Here the upper bound for min_i is 4 x vscale, as we know 4 x vscale is
+// always less than 32400. The bound for min_j is 16 as at vscale > 4,
+// 4 x vscale will be > 16, so the value will be clamped at 16.
+
+// CHECK: #[[$SCALABLE_BOUND_MAP_0:.*]] = affine_map<()[s0] -> (s0 * 4)>
+
+// CHECK-LABEL: @fixed_size_loop_nest
+// CHECK-DAG: %[[SCALABLE_BOUND:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_0]]()[%vscale]
+// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
+// CHECK: "test.some_use"(%[[SCALABLE_BOUND]], %[[C16]]) : (index, index) -> ()
+func.func @fixed_size_loop_nest() {
+ %c16 = arith.constant 16 : index
+ %c32400 = arith.constant 32400 : index
+ %c4 = arith.constant 4 : index
+ %c0 = arith.constant 0 : index
+ %vscale = vector.vscale
+ %c4_vscale = arith.muli %vscale, %c4 : index
+ scf.for %i = %c0 to %c32400 step %c4_vscale {
+ %min_i = affine.min #fixedDim0Map(%i)[%c4_vscale]
+ scf.for %j = %c0 to %c16 step %c4_vscale {
+ %min_j = affine.min #fixedDim1Map(%j)[%c4_vscale]
+ %bound_i = "test.reify_scalable_bound"(%min_i) {type = "UB"} : (index) -> index
+ %bound_j = "test.reify_scalable_bound"(%min_j) {type = "UB"} : (index) -> index
+ "test.some_use"(%bound_i, %bound_j) : (index, index) -> ()
+ }
+ }
+ return
+}
+
+// -----
+
+#dynamicDim0Map = affine_map<(d0, d1)[s0] -> (-d0 + d1, s0)>
+#dynamicDim1Map = affine_map<(d0, d1)[s0] -> (-d0 + d1, s0)>
+
+// Here upper bounds for both min_i and min_j are both 4 x vscale, as we know
+// that is always the largest value they could take. As if `dim < 4 x vscale`
+// then 4 x vscale is an overestimate, and if `dim > 4 x vscale` then the min
+// will be clamped to 4 x vscale.
+
+// CHECK: #[[$SCALABLE_BOUND_MAP_1:.*]] = affine_map<()[s0] -> (s0 * 4)>
+
+// CHECK-LABEL: @dynamic_size_loop_nest
+// CHECK: %[[SCALABLE_BOUND:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_1]]()[%vscale]
+// CHECK: "test.some_use"(%[[SCALABLE_BOUND]], %[[SCALABLE_BOUND]]) : (index, index) -> ()
+func.func @dynamic_size_loop_nest(%dim0: index, %dim1: index) {
+ %c4 = arith.constant 4 : index
+ %c0 = arith.constant 0 : index
+ %vscale = vector.vscale
+ %c4_vscale = arith.muli %vscale, %c4 : index
+ scf.for %i = %c0 to %dim0 step %c4_vscale {
+ %min_i = affine.min #dynamicDim0Map(%i)[%c4_vscale, %dim0]
+ scf.for %j = %c0 to %dim1 step %c4_vscale {
+ %min_j = affine.min #dynamicDim1Map(%j)[%c4_vscale, %dim1]
+ %bound_i = "test.reify_scalable_bound"(%min_i) {type = "UB"} : (index) -> index
+ %bound_j = "test.reify_scalable_bound"(%min_j) {type = "UB"} : (index) -> index
+ "test.some_use"(%bound_i, %bound_j) : (index, index) -> ()
+ }
+ }
+ return
+}
+
+// -----
+
+// Here the upper bound is just a value + a constant.
+
+// CHECK: #[[$SCALABLE_BOUND_MAP_2:.*]] = affine_map<()[s0] -> (s0 + 8)>
+
+// CHECK-LABEL: @add_to_vscale
+// CHECK: %[[SCALABLE_BOUND:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_2]]()[%vscale]
+// CHECK: "test.some_use"(%[[SCALABLE_BOUND]]) : (index) -> ()
+func.func @add_to_vscale() {
+ %vscale = vector.vscale
+ %c8 = arith.constant 8 : index
+ %vscale_plus_c8 = arith.addi %vscale, %c8 : index
+ %bound = "test.reify_scalable_bound"(%vscale_plus_c8) {type = "UB"} : (index) -> index
+ "test.some_use"(%bound) : (index) -> ()
+ return
+}
+
+// -----
+
+// Here we know vscale is always 2 so we get a constant upper bound.
+
+// CHECK-LABEL: @vscale_fixed_size
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: "test.some_use"(%[[C2]]) : (index) -> ()
+func.func @vscale_fixed_size() {
+ %vscale = vector.vscale
+ %bound = "test.reify_scalable_bound"(%vscale) {type = "UB", vscale_min = 2, vscale_max = 2} : (index) -> index
+ "test.some_use"(%bound) : (index) -> ()
+ return
+}
+
+// -----
+
+// Here we don't know the upper bound (%a is underspecified)
+
+func.func @unknown_bound(%a: index) {
+ %vscale = vector.vscale
+ %vscale_plus_a = arith.muli %vscale, %a : index
+ // expected-error @below{{could not reify bound}}
+ %bound = "test.reify_scalable_bound"(%vscale_plus_a) {type = "UB"} : (index) -> index
+ "test.some_use"(%bound) : (index) -> ()
+ return
+}
+
+// -----
+
+// Here we have two vscale values (that have not been CSE'd), but they should
+// still be treated as equivalent.
+
+// CHECK: #[[$SCALABLE_BOUND_MAP_3:.*]] = affine_map<()[s0] -> (s0 * 6)>
+
+// CHECK-LABEL: @duplicate_vscale_values
+// CHECK: %[[SCALABLE_BOUND:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_3]]()[%vscale]
+// CHECK: "test.some_use"(%[[SCALABLE_BOUND]]) : (index) -> ()
+func.func @duplicate_vscale_values() {
+ %c4 = arith.constant 4 : index
+ %vscale_0 = vector.vscale
+
+ %c2 = arith.constant 2 : index
+ %vscale_1 = vector.vscale
+
+ %c4_vscale = arith.muli %vscale_0, %c4 : index
+ %c2_vscale = arith.muli %vscale_1, %c2 : index
+ %add = arith.addi %c2_vscale, %c4_vscale : index
+
+ %bound = "test.reify_scalable_bound"(%add) {type = "UB"} : (index) -> index
+ "test.some_use"(%bound) : (index) -> ()
+ return
+}
diff --git a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
index 39671a930f2e21..a2fe4b3f6c34f7 100644
--- a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
@@ -13,6 +13,7 @@
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include "mlir/Pass/Pass.h"
@@ -75,7 +76,8 @@ static LogicalResult testReifyValueBounds(func::FuncOp funcOp,
WalkResult result = funcOp.walk([&](Operation *op) {
// Look for test.reify_bound ops.
if (op->getName().getStringRef() == "test.reify_bound" ||
- op->getName().getStringRef() == "test.reify_constant_bound") {
+ op->getName().getStringRef() == "test.reify_constant_bound" ||
+ op->getName().getStringRef() == "test.reify_scalable_bound") {
if (op->getNumOperands() != 1 || op->getNumResults() != 1 ||
!op->getResultTypes()[0].isIndex()) {
op->emitOpError("invalid op");
@@ -110,6 +112,9 @@ static LogicalResult testReifyValueBounds(func::FuncOp funcOp,
bool constant =
op->getName().getStringRef() == "test.reify_constant_bound";
+ bool scalable = !constant && op->getName().getStringRef() ==
+ "test.reify_scalable_bound";
+
// Prepare stop condition. By default, reify in terms of the op's
// operands. No stop condition is used when a constant was requested.
std::function<bool(Value, std::optional<int64_t>)> stopCondition =
@@ -137,6 +142,29 @@ static LogicalResult testReifyValueBounds(func::FuncOp funcOp,
if (succeeded(reifiedConst))
reified =
FailureOr<OpFoldResult>(rewriter.getIndexAttr(*reifiedConst));
+ } else if (scalable) {
+ unsigned vscaleMin = 1;
+ unsigned vscaleMax = 16;
+
+ if (auto attr = "vscale_min"; op->hasAttrOfType<IntegerAttr>(attr))
+ vscaleMin = unsigned(op->getAttrOfType<IntegerAttr>(attr).getInt());
+ if (auto attr = "vscale_max"; op->hasAttrOfType<IntegerAttr>(attr))
+ vscaleMax = unsigned(op->getAttrOfType<IntegerAttr>(attr).getInt());
+
+ auto loc = op->getLoc();
+ auto reifiedScalable = vector::computeScalableBound(
+ value, dim, vscaleMin, vscaleMax, *boundType);
+ if (succeeded(reifiedScalable)) {
+ SmallVector<std::pair<Value, std::optional<int64_t>>, 1>
+ vscaleOperand;
+ if (reifiedScalable->map.getNumInputs() == 1) {
+ // The only possible input to the bound is vscale.
+ vscaleOperand.push_back(std::make_pair(
+ rewriter.create<vector::VectorScaleOp>(loc), std::nullopt));
+ }
+ reified = affine::materializeComputedBound(
+ rewriter, loc, reifiedScalable->map, vscaleOperand);
+ }
} else {
if (dim) {
if (useArithOps) {
More information about the Mlir-commits
mailing list