359 lines
14 KiB
TableGen
359 lines
14 KiB
TableGen
//===- VectorInterfaces.td - Vector interfaces -------------*- tablegen -*-===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// Defines the interface for operations on vectors.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_INTERFACES_VECTORINTERFACES
|
|
#define MLIR_INTERFACES_VECTORINTERFACES
|
|
|
|
include "mlir/IR/OpBase.td"
|
|
|
|
def VectorUnrollOpInterface : OpInterface<"VectorUnrollOpInterface"> {
|
|
let description = [{
|
|
Encodes properties of an operation on vectors that can be unrolled.
|
|
}];
|
|
let cppNamespace = "::mlir";
|
|
|
|
let methods = [
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Return the shape ratio of unrolling to the target vector shape
|
|
`targetShape`. Return `std::nullopt` if the op cannot be unrolled to the
|
|
target vector shape.
|
|
}],
|
|
/*retTy=*/"::std::optional<::llvm::SmallVector<int64_t, 4>>",
|
|
/*methodName=*/"getShapeForUnroll",
|
|
/*args=*/(ins),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
assert($_op->getNumResults() == 1);
|
|
auto vt =
|
|
::llvm::dyn_cast<::mlir::VectorType>($_op.getResult().getType());
|
|
if (!vt)
|
|
return ::std::nullopt;
|
|
::llvm::SmallVector<int64_t, 4> res(
|
|
vt.getShape().begin(), vt.getShape().end());
|
|
return res;
|
|
}]
|
|
>,
|
|
];
|
|
}
|
|
|
|
def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
|
|
let description = [{
|
|
Encodes properties of a `vector.transfer_read` or `vector.transfer_write`
|
|
operation. Vector transfer ops have:
|
|
|
|
- A shaped value that the op reads from/writes to: a memref or a tensor.
|
|
- A vector, either as a result or as an operand.
|
|
- Indices that describe where the transfer from/to the shaped value starts.
|
|
- An optional mask.
|
|
- An optional in_bounds array to indicate transfer dimensions that are
|
|
guaranteed to be in-bounds.
|
|
- A permutation map to indicate transposes and broadcasts.
|
|
|
|
The "vector rank" is the rank of the vector type. E.g.:
|
|
```
|
|
// Transfer with shaped value rank 2 and vector (transfer) rank 1.
|
|
%0 = vector.transfer_read %arg0[%c3, %c3], %f0
|
|
{permutation_map = affine_map<(d0, d1) -> (d0)>}
|
|
: memref<?x?xf32>, vector<128xf32>
|
|
```
|
|
|
|
The "vector transfer rank" is the number of dimensions that participate in
|
|
the transfer and broadcasts, and matches the number of results in the
|
|
permutation map. In most cases, the vector rank matches the vector transfer
|
|
rank; the only exception is when a vector is flattened as part of the
|
|
transfer (see `getPermutationMap`).
|
|
}];
|
|
let cppNamespace = "::mlir";
|
|
|
|
let methods = [
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Return the `in_bounds` attribute name.
|
|
}],
|
|
/*retTy=*/"::mlir::StringRef",
|
|
/*methodName=*/"getInBoundsAttrName",
|
|
/*args=*/(ins)
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Return the `permutation_map` attribute name.
|
|
}],
|
|
/*retTy=*/"::mlir::StringRef",
|
|
/*methodName=*/"getPermutationMapAttrName",
|
|
/*args=*/(ins)
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Return the optional in_bounds attribute that specifies for each vector
|
|
dimension whether it is in-bounds or not. (Broadcast dimensions are
|
|
always in-bounds).
|
|
}],
|
|
/*retTy=*/"::mlir::ArrayAttr",
|
|
/*methodName=*/"getInBounds",
|
|
/*args=*/(ins)
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Return the memref or ranked tensor operand that this operation operates
|
|
on. In case of a "read" operation, that's the source from which the
|
|
operation reads. In case of a "write" operation, that's the destination
|
|
into which the operation writes.
|
|
TODO: Change name of operand, which is not accurate for xfer_write.
|
|
}],
|
|
/*retTy=*/"::mlir::Value",
|
|
/*methodName=*/"getBase",
|
|
/*args=*/(ins)
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Return the vector that this operation operates on. In case of a "read",
|
|
that's the vector OpResult. In case of a "write", that's the vector
|
|
operand value that is written by the op.
|
|
}],
|
|
/*retTy=*/"::mlir::Value",
|
|
/*methodName=*/"getVector",
|
|
/*args=*/(ins)
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Return the type of the vector that this operation operates on.
|
|
}],
|
|
/*retTy=*/"::mlir::VectorType",
|
|
/*methodName=*/"getVectorType",
|
|
/*args=*/(ins)
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Return the indices that specify the starting offsets into the source
|
|
operand. The starting offsets are guaranteed to be in-bounds except
|
|
when the transfer operation is masked or the `in_bounds` array doesn't
|
|
specify that that dimension is not necessarily in bounds. It is undefinied
|
|
behavior if these guarantees are violated.
|
|
}],
|
|
/*retTy=*/"::mlir::OperandRange",
|
|
/*methodName=*/"getIndices",
|
|
/*args=*/(ins)
|
|
>,
|
|
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Return the permutation map that describes the mapping of vector
|
|
dimensions to source dimensions, as well as broadcast dimensions.
|
|
|
|
The permutation result has one result per vector transfer dimension.
|
|
Each result is either a dim expression, indicating the corresponding
|
|
dimension in the source operand, or a constant "0" expression,
|
|
indicating a broadcast dimension.
|
|
|
|
Note: Nested vector dimensions that are flattened by this op are not
|
|
accounted for in the permutation map. E.g.:
|
|
```
|
|
// Vector type has rank 4, but permutation map has only 2 results. That
|
|
// is because there are only 2 transfer dimensions.
|
|
%0 = vector.transfer_read %arg1[%c3, %c3], %vf0
|
|
{permutation_map = affine_map<(d0, d1) -> (d0, d1)>}
|
|
: memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
|
|
```
|
|
}],
|
|
/*retTy=*/"::mlir::AffineMap",
|
|
/*methodName=*/"getPermutationMap",
|
|
/*args=*/(ins)
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Return the mask operand if the op has a mask. Otherwise, return an
|
|
empty value.
|
|
}],
|
|
/*retTy=*/"Value",
|
|
/*methodName=*/"getMask",
|
|
/*args=*/(ins)
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Returns whether replacing the base operand with a shaped value of
|
|
`newBaseType` and updating the permutation map according to
|
|
`newPermutationMap`, where the replacement would adhere to the
|
|
conditions specified by `updateStartingPosition`, passes any additional
|
|
op-specific constraints.
|
|
}],
|
|
/*retTy=*/"::mlir::LogicalResult",
|
|
/*methodName=*/"mayUpdateStartingPosition",
|
|
/*args=*/(ins "::mlir::ShapedType":$newBaseType,
|
|
"::mlir::AffineMap":$newPermutationMap),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return ::mlir::success();
|
|
}]
|
|
>,
|
|
// Note: The initial motivating factor for adding this interface method is to
|
|
// enable folding operations like `memref.subview` with transfer ops in a
|
|
// generic way. If future usecases require extending this method, it should be
|
|
// done.
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Updates the base of this transfer operation to `newBase`, the indices into
|
|
that base into `newIndices`, and the permutation map to `newPermutationMap`.
|
|
The new base, indices, and map must:
|
|
- Keep the same element type as the base (this condition could be relaxed
|
|
in the future)
|
|
- Preserve the number of transfer dimensions and avoid inserting new
|
|
dimensions unaccounted for in the permutation map - that is, if the transfer operation is
|
|
reading or writing a 4x2 vector to/from an Nx4x2 tensor or memref, it is
|
|
not valid to update that base to an Nx8 one, but updating in an N0xN1x4x2
|
|
base - or even an Nx16x32 one - is acceptable. Similarly, if the transfer
|
|
rank is 2 and the base is MxN, replacing it with a base that is Mx1xN or
|
|
MxKxN is only permitted if the permutation map is updated to skip the new
|
|
dimension. This restriction means that the updater does not need to reshape
|
|
or otherwise adjust the vector type of the transfer or other related attributes
|
|
like the inbounds attribute array.
|
|
- Must not shrink the size of a transfer dimension such that an `inbounds`
|
|
annotation becomes incorrect.
|
|
- Must pass any op-specific conditions in `mayUpdateStartingPosition`
|
|
|
|
The update is performed in-place. Implementations of the interface
|
|
should use `RewriterBase::modifyOpInPlace` to ensure proper pattern
|
|
rewriter operation.
|
|
}],
|
|
/*retTy=*/"void",
|
|
/*methodName=*/"updateStartingPosition",
|
|
/*args=*/(ins "::mlir::RewriterBase&":$rewriter, "::mlir::Value":$newBase,
|
|
"::mlir::ValueRange":$newIndices,
|
|
"::mlir::AffineMapAttr":$newPermutationMap),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
rewriter.modifyOpInPlace($_op, [&]() {
|
|
$_op.getBaseMutable().assign(newBase);
|
|
$_op.getIndicesMutable().assign(newIndices);
|
|
$_op.setPermutationMapAttr(newPermutationMap);
|
|
});
|
|
}]
|
|
>
|
|
];
|
|
|
|
let extraSharedClassDeclaration = [{
|
|
/// Return a vector of all in_bounds values as booleans (one per vector
|
|
/// transfer dimension).
|
|
::llvm::SmallVector<bool> getInBoundsValues() {
|
|
::llvm::SmallVector<bool> inBounds;
|
|
for (int64_t i = 0, e = $_op.getTransferRank(); i < e; ++i)
|
|
inBounds.push_back($_op.isDimInBounds(i));
|
|
return inBounds;
|
|
}
|
|
|
|
/// Return the number of leading shaped dimensions (of the "source" operand)
|
|
/// that do not participate in the permutation map.
|
|
unsigned getLeadingShapedRank() {
|
|
return $_op.getShapedType().getRank() - $_op.getTransferRank();
|
|
}
|
|
|
|
/// Return the mask type if the op has a mask. Otherwise, return an empty
|
|
/// VectorType.
|
|
::mlir::VectorType getMaskType() {
|
|
return $_op.getMask()
|
|
? ::llvm::cast<::mlir::VectorType>($_op.getMask().getType())
|
|
: ::mlir::VectorType();
|
|
}
|
|
|
|
/// Return the shaped type of the "source" operand value.
|
|
::mlir::ShapedType getShapedType() {
|
|
return ::llvm::cast<::mlir::ShapedType>($_op.getBase().getType());
|
|
}
|
|
|
|
/// Return the number of dimensions that participate in the permutation map.
|
|
unsigned getTransferRank() {
|
|
return $_op.getPermutationMap().getNumResults();
|
|
}
|
|
|
|
/// Return "true" if at least one of the vector dimensions is a broadcasted
|
|
/// dimension.
|
|
bool hasBroadcastDim() {
|
|
for (unsigned i = 0, e = $_op.getTransferRank(); i < e; ++i) {
|
|
if ($_op.isBroadcastDim(i))
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
/// Return "true" if at least one of the vector dimensions may be
|
|
/// out-of-bounds.
|
|
bool hasOutOfBoundsDim() {
|
|
for (unsigned idx = 0, e = $_op.getTransferRank(); idx < e; ++idx)
|
|
if (!$_op.isDimInBounds(idx))
|
|
return true;
|
|
return false;
|
|
}
|
|
|
|
/// Return "true" if the specified vector transfer dimension is a
|
|
/// broadcasted dimension.
|
|
bool isBroadcastDim(unsigned dim) {
|
|
auto expr = $_op.getPermutationMap().getResult(dim);
|
|
auto constExpr = ::llvm::dyn_cast<::mlir::AffineConstantExpr>(expr);
|
|
return constExpr && constExpr.getValue() == 0;
|
|
}
|
|
|
|
/// Return "true" if the vector transfer dimension `dim` is in-bounds.
|
|
/// Return "false" otherwise.
|
|
bool isDimInBounds(unsigned dim) {
|
|
auto inBounds = $_op.getInBounds();
|
|
return ::llvm::cast<::mlir::BoolAttr>(inBounds[dim]).getValue();
|
|
}
|
|
|
|
/// Helper function to account for the fact that `permutationMap` results
|
|
/// and `op.getIndices` sizes may not match and may not be aligned. The
|
|
/// first `getLeadingShapedRank()` indices may just be indexed and not
|
|
/// transferred from/into the vector.
|
|
/// For example:
|
|
/// ```
|
|
/// vector.transfer %0[%i, %j, %k, %c0]
|
|
/// : memref<?x?x?x?xf32>, vector<2x4xf32>
|
|
/// ```
|
|
/// with `permutation_map = (d0, d1, d2, d3) -> (d2, d3)`.
|
|
/// Provide a zip function to coiterate on 2 running indices: `resultIdx`
|
|
/// and `indicesIdx` which accounts for this misalignment.
|
|
void zipResultAndIndexing(
|
|
::llvm::function_ref<void(int64_t, int64_t)> fun) {
|
|
for (int64_t resultIdx = 0,
|
|
indicesIdx = $_op.getLeadingShapedRank(),
|
|
eResult = $_op.getTransferRank();
|
|
resultIdx < eResult;
|
|
++resultIdx, ++indicesIdx)
|
|
fun(resultIdx, indicesIdx);
|
|
}
|
|
|
|
/// Return the shape of the hyperrectangular slice within the tensor/memref
|
|
/// operand that is accessed by the transfer op.
|
|
/// For example:
|
|
/// ```
|
|
/// vector.transfer %w0[%i, %j, %k] {
|
|
/// permutation_map = affine_map<(d0, d1, d2) -> (d1, d0, 0)>} :
|
|
/// : tensor<?x?x?xf32>, vector<4x2x6xf32>
|
|
/// ```
|
|
/// returns a shape [2, 4, 1].
|
|
SmallVector<int64_t> getTransferChunkAccessed() {
|
|
SmallVector<int64_t> dimSizes($_op.getPermutationMap().getNumDims(), 1);
|
|
for (auto vecDims : llvm::zip($_op.getPermutationMap().getResults(),
|
|
$_op.getVectorType().getShape())) {
|
|
AffineExpr dim = std::get<0>(vecDims);
|
|
int64_t size = std::get<1>(vecDims);
|
|
// Skip broadcast.
|
|
if (::llvm::isa<AffineConstantExpr>(dim))
|
|
continue;
|
|
dimSizes[::llvm::cast<AffineDimExpr>(dim).getPosition()] = size;
|
|
}
|
|
return dimSizes;
|
|
}
|
|
}];
|
|
}
|
|
|
|
#endif // MLIR_INTERFACES_VECTORINTERFACES
|