//===- LinalgInterfaces.td - Linalg Interfaces Declaration -*- 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
//
//===----------------------------------------------------------------------===//
//
// This is the definition file for the structured interface sfor Linalg ops.
//
//===----------------------------------------------------------------------===//
#ifndef LINALG_IR_LINALGINTERFACES
#define LINALG_IR_LINALGINTERFACES
include "mlir/Interfaces/DestinationStyleOpInterface.td"
include "mlir/IR/OpBase.td"
// The 'LinalgContractionOpInterface' provides access to the
// 'ContractionOpInterface'.
def LinalgContractionOpInterface : OpInterface<"ContractionOpInterface"> {
let description = [{
A Linalg contraction is defined in general terms:
1. Has 2 input and 1 output shapes.
2. Has at least one reduction dimension.
3. Has only projected permutation indexing maps.
4. its body computes `u5(u1(c) + u2(u3(a) * u4(b)))` on some field
(AddOpType, MulOpType), where u1, u2, u3, u4 and u5 represent scalar unary
operations that may change the type (e.g. for mixed-precision).
As a consequence, when vectorization of such an op occurs, the only special
behavior is that the (unique) MulOpType is vectorized into a
`vector.contract`. All other ops are handled in a generic fashion.
In the future, we may wish to allow more input arguments and elementwise and
constant operations that do not involve the reduction dimension(s).
}];
let cppNamespace = "::mlir::linalg";
let verify = [{ return detail::verifyContractionInterface($_op); }];
let verifyWithRegions = 1;
let methods = [
InterfaceMethod<
/*desc=*/"Returns the left-hand side operand.",
/*retTy=*/"Value",
/*methodName=*/"lhs",
/*args=*/(ins),
/*methodBody=*/[{
return $_op.getOperation()->getOperand(0);
}]>,
InterfaceMethod<
/*desc=*/"Returns the right-hand side operand.",
/*retTy=*/"Value",
/*methodName=*/"rhs",
/*args=*/(ins),
/*methodBody=*/[{
return $_op.getOperation()->getOperand(1);
}]>,
InterfaceMethod<
/*desc=*/[{
Returns whether the given op has indexing maps that correspond to a
row-major matmul operation.
}],
/*retTy=*/"bool",
/*methodName=*/"isRowMajorMatmul",
/*args=*/(ins),
/*methodBody=*/[{
return mlir::isRowMajorMatmul($_op.getIndexingMaps());
}]>,
InterfaceMethod<
/*desc=*/[{
Returns whether the given op has indexing maps that correspond to a
column-major matmul operation.
}],
/*retTy=*/"bool",
/*methodName=*/"isColumnMajorMatmul",
/*args=*/(ins),
/*methodBody=*/[{
return mlir::isColumnMajorMatmul($_op.getIndexingMaps());
}]>,
InterfaceMethod<
/*desc=*/[{
Returns whether the given op has indexing maps that correspond to a
row-major batch matmul operation.
}],
/*retTy=*/"bool",
/*methodName=*/"isRowMajorBatchMatmul",
/*args=*/(ins),
/*methodBody=*/[{
return mlir::isRowMajorBatchMatmul($_op.getIndexingMaps());
}]>,
InterfaceMethod<
/*desc=*/[{
Returns whether the given op has indexing maps that correspond to a
vector-matrix multiplication.
}],
/*retTy=*/"bool",
/*methodName=*/"isVecmat",
/*args=*/(ins),
/*methodBody=*/[{
return mlir::isVecmat($_op.getIndexingMaps());
}]>,
InterfaceMethod<
/*desc=*/[{
Returns whether the given op has indexing maps that correspond to a
batched vector-matrix multiplication.
}],
/*retTy=*/"bool",
/*methodName=*/"isBatchVecmat",
/*args=*/(ins),
/*methodBody=*/[{
return mlir::isBatchVecmat($_op.getIndexingMaps());
}]>,
InterfaceMethod<
/*desc=*/[{
Returns whether the given op has indexing maps that correspond to a
matrix-vector multiplication.
}],
/*retTy=*/"bool",
/*methodName=*/"isMatvec",
/*args=*/(ins),
/*methodBody=*/[{
return mlir::isMatvec($_op.getIndexingMaps());
}]>,
InterfaceMethod<
/*desc=*/[{
Returns whether the given op has indexing maps that correspond to a
batched matrix-vector multiplication.
}],
/*retTy=*/"bool",
/*methodName=*/"isBatchMatvec",
/*args=*/(ins),
/*methodBody=*/[{
return mlir::isBatchMatvec($_op.getIndexingMaps());
}]>,
];
}
def LinalgConvolutionOpInterface : OpInterface<"ConvolutionOpInterface"> {
let description = [{
A convolution is defined in general terms:
1. Has an `image` and a `filter` operand.
2. Has one `output` operand.
3. The indexing maps of the input have expressions that satisfy
```
AffineExpr ::== AffineDimExpr | ConvolvedExpr
ConvolvedExpr ::== MulExpr (`+` MulExpr)+
MulExpr ::== AffineDimExpr (`*` (AffineConstantExpr | AffineSymbolExpr))?
```
4. The filter and the output have projected permutation maps.
5. Each of the loops can be qualified as one of,
- Loop over batch dimension,
- Loop over output image dimensions,
- Loop over output channel dimensions,
- Loop over convolved filter dimensions,
- Loop over input channel dimension.
}];
let cppNamespace = "::mlir::linalg";
let verify = [{ return detail::verifyConvolutionInterface($_op); }];
let methods = [
InterfaceMethod<
/*desc=*/"Return the image operand.",
/*retTy=*/"Value",
/*methodName=*/"image",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return $_op.getOperation()->getOperand(0);
}]
>,
InterfaceMethod<
/*desc=*/"Return the filter operand.",
/*retTy=*/"Value",
/*methodName=*/"filter",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return $_op.getOperation()->getOperand(1);
}]
>,
];
}
def LinalgFillOpInterface : OpInterface<"FillOpInterface"> {
let description = [{
A fill operation is defined in general terms:
1. Has a scalar `value` operand.
2. Has one `output` operand.
}];
let cppNamespace = "::mlir::linalg";
let verify = [{ return detail::verifyFillInterface($_op); }];
let methods = [
InterfaceMethod<
/*desc=*/"Return the fill value.",
/*retTy=*/"Value",
/*methodName=*/"value",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return $_op.getOperation()->getOperand(0);
}]
>,
InterfaceMethod<
/*desc=*/"Return the output operand.",
/*retTy=*/"Value",
/*methodName=*/"output",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return $_op.getOperation()->getOperand(1);
}]
>,
InterfaceMethod<
/*desc=*/"Return the result.",
/*retTy=*/"Value",
/*methodName=*/"result",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
if ($_op.getOperation()->getResults().empty())
return nullptr;
return $_op.getOperation()->getResults().front();
}]
>,
];
}
// The 'LinalgStructuredInterface' provides access to the 'LinalgOp' interface.
def LinalgStructuredInterface
: OpInterface<"LinalgOp", [DestinationStyleOpInterface]> {
let cppNamespace = "::mlir::linalg";
let methods = [
//===------------------------------------------------------------------===//
// Loop types handling.
//===------------------------------------------------------------------===//
InterfaceMethod<
/*desc=*/[{
Return the number of parallel loops.
}],
/*retTy=*/"unsigned",
/*methodName=*/"getNumParallelLoops",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return llvm::count($_op.getIteratorTypesArray(),
utils::IteratorType::parallel);
}]
>,
InterfaceMethod<
/*desc=*/[{
Return true if all loops are parallel.
}],
/*retTy=*/"bool",
/*methodName=*/"isAllParallelLoops",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return getNumParallelLoops() == getNumParallelLoops();
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the dims that are parallel loops.
}],
/*retTy=*/"void",
/*methodName=*/"getParallelDims",
/*args=*/(ins "SmallVectorImpl<unsigned> &":$res),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return findPositionsOfType($_op.getIteratorTypesArray(),
utils::IteratorType::parallel, res);
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the number of reduction loops.
}],
/*retTy=*/"unsigned",
/*methodName=*/"getNumReductionLoops",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return llvm::count($_op.getIteratorTypesArray(),
utils::IteratorType::reduction);
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the dims that are reduction loops.
}],
/*retTy=*/"void",
/*methodName=*/"getReductionDims",
/*args=*/(ins "SmallVectorImpl<unsigned> &":$res),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return findPositionsOfType($_op.getIteratorTypesArray(),
utils::IteratorType::reduction, res);
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the total number of loops within the current operation.
}],
/*retTy=*/"unsigned",
/*methodName=*/"getNumLoops",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return $_op.getIteratorTypesArray().size();
}]
>,
InterfaceMethod<
/*desc=*/[{
Returns true if the current operation has only one loop and it's a
reduction loop.
}],
/*retTy=*/"bool",
/*methodName=*/"hasSingleReductionLoop",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
auto iters = $_op.getIteratorTypesArray();
return iters.size() == 1 &&
llvm::count(iters, utils::IteratorType::reduction) == 1;
}]>,
//===------------------------------------------------------------------===//
// Input and Init arguments handling.
//===------------------------------------------------------------------===//
InterfaceMethod<
/*desc=*/[{
Return true if the payload uses the value loaded from `opOperand`. This
is useful to avoid loading from "write-only" memory that may be
uninitialized, as well as properly cloning "read-write" operands.
}],
/*retTy=*/"bool",
/*methodName=*/"payloadUsesValueFromOperand",
/*args=*/(ins "OpOperand *":$opOperand),
/*methodBody=*/"",
/*defaultImplementation=*/[{
unsigned bbArgNumber = opOperand->getOperandNumber();
// Init tensors have uses.
return !getBlock()->getArgument(bbArgNumber).use_empty();
}]
>,
InterfaceMethod<
/*desc=*/[{
Returns true only if linalgOp takes one input and produces one result.
}],
/*retTy=*/"bool",
/*methodName=*/"isSingleInputOutput",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return $_op.getNumDpsInputs() == 1 && $_op.getNumDpsInits() == 1;
}]
>,
InterfaceMethod<
/*desc=*/[{
Return true if `opOperand` is an init tensor. This is true when it is
an output tensor operand whose value is used in the payload region.
}],
/*retTy=*/"bool",
/*methodName=*/"isInitTensor",
/*args=*/(ins "OpOperand *":$opOperand),
/*methodBody=*/"",
/*defaultImplementation=*/[{
if (!$_op.isDpsInit(opOperand))
return false;
return payloadUsesValueFromOperand(opOperand);
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the `opOperand` rank or zero for scalars or vectors not wrapped within a tensor or a memref.
}],
/*retTy=*/"int64_t",
/*methodName=*/"getRank",
/*args=*/(ins "OpOperand*":$opOperand),
/*methodBody=*/"",
/*defaultImplementation=*/[{
assert(opOperand->getOwner() == this->getOperation());
Type t = opOperand->get().getType();
// A VectorType is an elemental type, do not consider its rank for the operand.
if (isa<VectorType>(t))
return 0;
// Tensor and Memref container types have a rank.
if (auto shapedType = ::llvm::dyn_cast<ShapedType>(t)) {
// Failsafe.
assert((isa<MemRefType>(t) || isa<RankedTensorType>(t)) &&
"expected a ranked tensor or memref in LinalgInterface::getRank");
return shapedType.getRank();
}
return 0;
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the input block arguments of the region.
}],
/*retTy=*/"Block::BlockArgListType",
/*methodName=*/"getRegionInputArgs",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return getBlock()->getArguments().take_front($_op.getNumDpsInputs());
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the output block arguments of the region.
}],
/*retTy=*/"Block::BlockArgListType",
/*methodName=*/"getRegionOutputArgs",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return getBlock()->getArguments().take_back($_op.getNumDpsInits());
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the `opOperand` shape or an empty vector for scalars or vectors
not wrapped within a tensor or a memref.
}],
/*retTy=*/"ArrayRef<int64_t>",
/*methodName=*/"getShape",
/*args=*/(ins "OpOperand*":$opOperand),
/*methodBody=*/"",
/*defaultImplementation=*/[{
assert(opOperand->getOwner() == this->getOperation());
Type t = opOperand->get().getType();
// A VectorType is an elemental type, do not consider its rank for the operand.
if (isa<VectorType>(t))
return {};
if (auto shapedType = ::llvm::dyn_cast<ShapedType>(t)) {
// Failsafe.
assert((isa<MemRefType>(t) || isa<RankedTensorType>(t)) &&
"expected a ranked tensor or memref in LinalgInterface::getRank");
return shapedType.getShape();
}
return {};
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the block argument for an `opOperand`.
}],
/*retTy=*/"BlockArgument",
/*methodName=*/"getMatchingBlockArgument",
/*args=*/(ins "OpOperand *":$opOperand),
/*methodBody=*/"",
/*defaultImplementation=*/[{
assert(opOperand->getOwner() == this->getOperation());
return getBlock()->getArgument(opOperand->getOperandNumber());
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the operand for a `blockArgument`.
}],
/*retTy=*/"OpOperand *",
/*methodName=*/"getMatchingOpOperand",
/*args=*/(ins "BlockArgument":$blockArgument),
/*methodBody=*/"",
/*defaultImplementation=*/[{
assert(blockArgument.getOwner() == getBlock());
return &this->getOperation()->getOpOperand(
blockArgument.getArgNumber());
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the input or output indexing map for `opOperand`.
}],
/*retTy=*/"AffineMap",
/*methodName=*/"getMatchingIndexingMap",
/*args=*/(ins "OpOperand*":$opOperand),
/*methodBody=*/"",
/*defaultImplementation=*/[{
assert(opOperand->getOwner() == this->getOperation());
auto indexingMaps =
$_op.getIndexingMaps().template getAsValueRange<AffineMapAttr>();
return *(indexingMaps.begin() + opOperand->getOperandNumber());
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the indexing map for a `result`.
}],
/*retTy=*/"AffineMap",
/*methodName=*/"getIndexingMapMatchingResult",
/*args=*/(ins "OpResult":$result),
/*methodBody=*/"",
/*defaultImplementation=*/[{
assert(result.getOwner() == this->getOperation());
auto indexingMaps =
$_op.getIndexingMaps().template getAsValueRange<AffineMapAttr>();
return *(indexingMaps.begin() + $_op.getNumDpsInputs() +
result.getResultNumber());
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the value yielded by the region corresponding to an output
`opOperand`.
}],
/*retTy=*/"OpOperand *",
/*methodName=*/"getMatchingYieldValue",
/*args=*/(ins "OpOperand*":$opOperand),
/*methodBody=*/"",
/*defaultImplementation=*/[{
assert(opOperand->getOwner() == this->getOperation());
int64_t resultIndex =
opOperand->getOperandNumber() - $_op.getNumDpsInputs();
assert(resultIndex >= 0 &&
resultIndex < this->getOperation()->getNumResults());
Operation *yieldOp = getBlock()->getTerminator();
return &yieldOp->getOpOperand(resultIndex);
}]
>,
//===------------------------------------------------------------------===//
// Other interface methods.
//===------------------------------------------------------------------===//
InterfaceMethod<
/*desc=*/[{
Return the single block constituting the body of the operation by
calling the getBody method on the concrete operation.
}],
/*retTy=*/"Block*",
/*methodName=*/"getBlock",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
// Assume the concrete operation implements the
// SingleBlockImplicitTerminator trait.
return $_op.getBody();
}]
>,
InterfaceMethod<
/*desc=*/[{
Return iterator types in the current operation.
Default implementation assumes that the operation has an attribute
`iterator_types`, but it's not always the case. Sometimes iterator types
can be infered from other parameters and in such cases default
getIteratorTypesArray should be overriden.
}],
/*retTy=*/"SmallVector<utils::IteratorType>",
/*methodName=*/"getIteratorTypesArray",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
auto range = $_op.getIteratorTypes()
.template getAsValueRange<IteratorTypeAttr,
utils::IteratorType>();
return {range.begin(), range.end()};
}]
>,
InterfaceMethod<
/*desc=*/[{
Return true if the indexing map is depending on the current op instance.
This means that the indexing map is dynamically synthesized by using the
op instance's concrete attributes, instead of being static for all
instances of the same op kind.
}],
/*retTy=*/"bool",
/*methodName=*/"hasDynamicIndexingMaps",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{ return false; }]
>,
InterfaceMethod<
/*desc=*/[{
Verify all attributes used by indexing maps are valid.
}],
/*retTy=*/"LogicalResult",
/*methodName=*/"verifyIndexingMapRequiredAttributes",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{ return success(); }]
>,
InterfaceMethod<
/*desc=*/[{
Return the indexing maps attribute within the current operation.
}],
/*retTy=*/"ArrayAttr",
/*methodName=*/"getIndexingMaps"
>,
InterfaceMethod<
/*desc=*/[{
Return the indexing maps within the current operation.
}],
/*retTy=*/"SmallVector<AffineMap>",
/*methodName=*/"getIndexingMapsArray",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
auto range = $_op.getIndexingMaps()
.template getAsValueRange<AffineMapAttr>();
return {range.begin(), range.end()};
}]
>,
InterfaceMethod<
/*desc=*/[{
Return true if any of the operands has a dynamic shape.
}],
/*retTy=*/"bool",
/*methodName=*/"hasDynamicShape",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return llvm::any_of(getStaticShape(), ShapedType::isDynamic);
}]
>,
InterfaceMethod<
/*desc=*/[{
Return the name registered for this op when lowering to an external
library call.
}],
/*retTy=*/"std::string",
/*methodName=*/"getLibraryCallName",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return $_op.getLibraryCallName();
}]
>,
InterfaceMethod<
/*desc=*/[{
Return whether the op accesses the iteration indices.
}],
/*retTy=*/"bool",
/*methodName=*/"hasIndexSemantics",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/""
>,
InterfaceMethod<
/*desc=*/[{
Return op operands that have a corresponding argument in the basic block.
By default, the block should have an argument for each operand, but there
are expection. For example, in `map` output operand isn't used in
the block.
}],
/*retTy=*/"::llvm::SmallVector<OpOperand *>",
/*methodName=*/"getOpOperandsMatchingBBargs",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
::llvm::SmallVector<OpOperand *> result;
result.reserve($_op->getNumOperands());
llvm::transform(
this->getOperation()->getOpOperands(),
std::back_inserter(result),
[](OpOperand &opOperand) { return &opOperand; });
return result;
}]
>,
InterfaceMethod<
/*desc=*/[{
Given a dimension of the iteration space of a Linalg operation, finds an
operand in the operation that is defined on such dimension. Returns
whether such operand was found or not. If found, also returns the
operand value and the dimension position within the operand.
}],
/*retTy=*/"LogicalResult",
/*methodName=*/"mapIterationSpaceDimToOperandDim",
/*args=*/(ins "unsigned":$dimPos,
"::mlir::Value &":$operand,
"unsigned &":$operandDimPos),
/*methodBody=*/"",
/*defaultImplementation=*/[{
// Retrieve the operand and its dimension position from the first
// operand with a permutation map that is defined on such dimension.
for (auto [i, idxMap] : llvm::enumerate($_op.getIndexingMapsArray())) {
if (idxMap.isProjectedPermutation()) {
if (auto mayOperandDim = idxMap.getResultPosition(
getAffineDimExpr(dimPos, idxMap.getContext()))) {
operand = $_op->getOperand(i);
operandDimPos = *mayOperandDim;
return success();
}
}
}
return failure();
}]
>,
InterfaceMethod<
/*desc=*/[{
Given a dimension of the iteration space of a Linalg operation, finds
all the operands in the operation that are defined on such dimension.
Returns all the operand values found and their dimension positions in
`operandDimPairs`.
}],
/*retTy=*/"void",
/*methodName=*/"mapIterationSpaceDimToAllOperandDims",
/*args=*/(ins "unsigned":$dimPos,
"mlir::SmallVectorImpl<std::pair<Value, unsigned>>&":$operandDimPairs),
/*methodBody=*/"",
/*defaultImplementation=*/[{
for (auto [i, idxMap] : llvm::enumerate($_op.getIndexingMapsArray())) {
if (idxMap.isProjectedPermutation()) {
if (auto mayOperandDim = idxMap.getResultPosition(
getAffineDimExpr(dimPos, idxMap.getContext()))) {
operandDimPairs.push_back({$_op->getOperand(i), *mayOperandDim});
}
}
}
return;
}]
>,
//===------------------------------------------------------------------===//
// Linalg generalization hooks.
//===------------------------------------------------------------------===//
InterfaceMethod<
/*desc=*/[{
Hook to provide a custom AffineMap used to compute all the operand
subshapes given loop bounds. This is used to answer the question: "given
an iteration space over the codomain, what are the subshapes of the
operands involved in the computation".
The default behavior is to just concatenate all the indexing maps.
A custom AffineMap allows providing a map that can be used to
compute subshapes even in cases where the concatenation of indexing maps
(i.e. the data traversal order) is not a simple permutation of the loop
traversal order. It is then possible to define ops with skewed data
traversal order for which we can still easily compute hyperrectangular
loop bounds and subviews.
}],
/*retTy=*/"AffineMap",
/*methodName=*/"getLoopsToShapesMap",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
auto maps = $_op.getIndexingMapsArray();
return concatAffineMaps(maps);
}]
>,
InterfaceMethod<
/*desc=*/[{
Hook to provide a custom AffineMap used to construct the
hyperrectangular loop iteration space given all the operand subshapes.
This is used to answer the question:
"Given a list of operand ranges, what is the subportion of the iteration
space involved in the computation".
This is the inverse problem of `getLoopsToShapesMap`.
Return the empty AffineMap when such an AffineMap cannot be constructed.
The default behavior is based on a very simple inference procedure that
only works with permutation affine maps.
A more advanced Tensor-Comprehension like inference is possible but has
proven to be ambiguous in unfavorable case.
A safer and more robust alternative is to allow each op to define
its own AffineMap.
}],
/*retTy=*/"AffineMap",
/*methodName=*/"getShapesToLoopsMap",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return inversePermutation(getLoopsToShapesMap());
}]
>,
InterfaceMethod<
/*desc=*/[{
Checks if the given operands can be dropped, and the remaining
operands can still compute the bounds of the op.
}],
/*retTy=*/"bool",
/*methodName=*/"canOpOperandsBeDropped",
/*args=*/(ins "ArrayRef<OpOperand *>":$droppedOperands),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return detail::canOpOperandsBeDroppedImpl($_op, droppedOperands);
}]
>,
InterfaceMethod<
/*desc=*/[{
Like `getShape`, but only returns statically-known information, without
generating any new IR. For each shape dimension, returns >=0 if that
dimension is statically known, or ShapedType::kDynamic otherwise.
}],
/*retTy=*/"SmallVector<int64_t>",
/*methodName=*/"getStaticShape",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
SmallVector<int64_t> res;
for (OpOperand &opOperand : this->getOperation()->getOpOperands())
llvm::append_range(res, getShape(&opOperand));
return res;
}]
>,
InterfaceMethod<
/*desc=*/[{
Returns the statically-known loop ranges. Composes
`getShapesToLoopsMap()` with the result of `getStaticShape`.
Returns ShapedType::kDynamic for non-statically-known loop ranges.
This is expected to be called by a valid Linalg op
}],
/*retTy=*/"SmallVector<int64_t, 4>",
/*methodName=*/"getStaticLoopRanges",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
SmallVector<int64_t> viewSizes = getStaticShape();
AffineMap invertedMap = getShapesToLoopsMap();
assert(invertedMap && "expected a valid Linalg op to call the method");
return invertedMap.compose(viewSizes);
}]
>,
//===------------------------------------------------------------------===//
// Other static interface methods.
//===------------------------------------------------------------------===//
StaticInterfaceMethod<
/*desc=*/[{
Returns the region builder for constructing the body for linalg.generic.
Returns a null function if this named op does not define a region
builder.
}],
/*retTy=*/"std::function<void(ImplicitLocOpBuilder &, Block &, ArrayRef<NamedAttribute>)>",
/*methodName=*/"getRegionBuilder",
(ins),
[{ return ConcreteOp::getRegionBuilder(); }]
>,
InterfaceMethod<
/*desc=*/[{
Return true if all the indexing maps are projected permutations.
Otherwise return false.
}],
/*retTy=*/"bool",
/*methodName=*/"hasOnlyProjectedPermutations",
(ins),
[{
return llvm::all_of($_op.getIndexingMapsArray(),
[](AffineMap map) { return map.isProjectedPermutation(); });
}]
>
];
let extraClassDeclaration = [{
/// Return the flat list of all operand dimension sizes in the order they
/// appear in the operands.
SmallVector<OpFoldResult> createFlatListOfOperandDims(OpBuilder &, Location);
/// Return the flat list of all operands' static dimension sizes in the
/// order they appear in the operands. All operand dimension sizes have to
/// be statically known.
SmallVector<int64_t, 4> createFlatListOfOperandStaticDims();
/// Create the loop ranges to materialize the computation over the current
/// operands. This is done by applying `getShapesToLoopsMap` to
/// `createFlatListOfOperandDims`.
SmallVector<Range, 4> createLoopRanges(OpBuilder &b, Location loc);
/// Compute the static loop sizes necessary to vectorize the computation.
/// This is done by applying `getShapesToLoopsMap` to
/// `createFlatListOfOperandStaticDims`.
SmallVector<int64_t, 4> computeStaticLoopSizes();
/// Returns the value that expresses the shape of the output in terms of
/// shape of the input operands where possible
LogicalResult reifyResultShapes(OpBuilder &b,
ReifiedRankedShapedTypeDims &reifiedReturnShapes);
/// Return the index in the indexingMaps vector that corresponds to this `opOperand`
int64_t getIndexingMapIndex(OpOperand *opOperand);
}];
let verify = [{ return detail::verifyStructuredOpInterface($_op); }];
let verifyWithRegions = 1;
}
def AggregatedOpInterface : OpInterface<"AggregatedOpInterface"> {
let description = [{
Interface for decomposing aggregated operations into a sequence of simpler
ops.
}];
let cppNamespace = "::mlir::linalg";
let methods = [
InterfaceMethod<
/*desc=*/[{
Method to decompose the operation into simpler operations.
On success, this method returns one `Value` per result in the
original operation.
The order of the returned values must match the order of the
original values.
In other words, the returned vector can be used directly with
`RewriterBase::replaceOp(this, returnedValues)`.
}],
/*retType=*/"FailureOr<SmallVector<Value>>",
/*methodName=*/"decomposeOperation",
/*args=*/(ins
"OpBuilder &":$b),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return {};
}]
>
];
}
#endif // LINALG_IR_LINALGINTERFACES
Codebase Browser
llvm