2861856baf
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` that will always take at most one parameter, vscale, and returns a single result, which is the bound of `value`. The API is defined as follows: ```c++ FailureOr<ConstantOrScalableBound> vector::ScalableValueBoundsConstraintSet::computeScalableBound( Value value, std::optional<int64_t> dim, unsigned vscaleMin, unsigned vscaleMax, presburger::BoundType boundType, bool closedUB = true, StopConditionFn stopCondition = nullptr); ``` 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).
52 lines
1.9 KiB
C++
52 lines
1.9 KiB
C++
//===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.h"
|
|
|
|
#include "mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h"
|
|
#include "mlir/Dialect/Vector/IR/VectorOps.h"
|
|
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace mlir::vector {
|
|
namespace {
|
|
|
|
struct VectorScaleOpInterface
|
|
: public ValueBoundsOpInterface::ExternalModel<VectorScaleOpInterface,
|
|
VectorScaleOp> {
|
|
void populateBoundsForIndexValue(Operation *op, Value value,
|
|
ValueBoundsConstraintSet &cstr) const {
|
|
auto *scalableCstr = dyn_cast<ScalableValueBoundsConstraintSet>(&cstr);
|
|
if (!scalableCstr)
|
|
return;
|
|
auto vscaleOp = cast<VectorScaleOp>(op);
|
|
assert(value == vscaleOp.getResult() && "invalid value");
|
|
if (auto vscale = scalableCstr->getVscaleValue()) {
|
|
// All copies of vscale are equivalent.
|
|
scalableCstr->bound(value) == cstr.getExpr(vscale);
|
|
} else {
|
|
// We know vscale is confined to [vscaleMin, vscaleMax].
|
|
scalableCstr->bound(value) >= scalableCstr->getVscaleMin();
|
|
scalableCstr->bound(value) <= scalableCstr->getVscaleMax();
|
|
scalableCstr->setVscale(vscaleOp);
|
|
}
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
} // namespace mlir::vector
|
|
|
|
void mlir::vector::registerValueBoundsOpInterfaceExternalModels(
|
|
DialectRegistry ®istry) {
|
|
registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
|
|
vector::VectorScaleOp::attachInterface<vector::VectorScaleOpInterface>(
|
|
*ctx);
|
|
});
|
|
}
|