514 lines
23 KiB
TableGen
514 lines
23 KiB
TableGen
//===- TilingInterface.td - Interface for tiling operations *- 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 file contains an interface to allow operations to generate a tiled
|
|
// implementation of themselves.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_TILINGINTERFACE
|
|
#define MLIR_TILINGINTERFACE
|
|
|
|
include "mlir/IR/OpBase.td"
|
|
|
|
def TilingInterface : OpInterface<"TilingInterface"> {
|
|
let description = [{
|
|
This interface allows operations to expose information needed to tile them.
|
|
|
|
The intent of this interface is to separate the generation of the loop
|
|
structure (and constructs used for it) from the information needed from
|
|
the operation to be able to tile them. As a result an implementation of
|
|
the tiling algorithm (like `scf::tileUsingSCF`) can generate the inter-tile
|
|
loop structure, and call into the methods of the interface to be able to
|
|
tile any operation that implements the interface.
|
|
|
|
This interface is also meant to help with "tile and fuse", i.e. the process
|
|
of fusing a producer with a consumer by
|
|
a) Tiling the consumer
|
|
b) Based on the tile of the producer used by the tiled consumer,
|
|
materialize the tiled implementation of a producer to generate that
|
|
tile (and use it immediately in the consumer)
|
|
You could also fuse a consumer with a producer by
|
|
a) Tiling the producer
|
|
b) Based on the tile produced, materialize the tiled implementation of
|
|
a consumer that uses this tile.
|
|
Note that the tile and fuse does not make any calculation on whether it
|
|
is "profitable to do this", but simply provides a mechansim to implement
|
|
the transformation when such a fusion is needed by the caller.
|
|
|
|
For any operation to be tilable, an operation has to implement the
|
|
following two methods (see description below)
|
|
- `getLoopIteratorTypes`
|
|
- `getIterationDomain`
|
|
- `getTiledImplementation`
|
|
- `getResultTilePosition`
|
|
|
|
For an operation to be "tiled and fused" with its (already tiled) consumer,
|
|
an operation has to implement the following additional method (see
|
|
description below):
|
|
- `generateResultTileValue`
|
|
- `getIterationDomainTileFromResultTile`
|
|
|
|
For an operation to be "tiled and fused" with its (already tiled) producer,
|
|
an operation has to implement the following additional methods (see
|
|
description below):
|
|
- `getTiledImplementationFromOperandTiles`
|
|
- `getIterationDomainTileFromOperandTiles`.
|
|
}];
|
|
let cppNamespace = "::mlir";
|
|
let methods = [
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Returns a list of iterator types that describe the number of loops.
|
|
}],
|
|
/*retType=*/"::mlir::SmallVector<::mlir::utils::IteratorType>",
|
|
/*methodName=*/"getLoopIteratorTypes",
|
|
/*args=*/(ins),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/"return {};"
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Returns a list of ranges that describe the loop bounds and
|
|
step for the loops of the operation.
|
|
}],
|
|
/*retTy=*/"::mlir::SmallVector<::mlir::Range>",
|
|
/*methodName=*/"getIterationDomain",
|
|
/*args=*/(ins "::mlir::OpBuilder &":$b),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/"return {};"
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Method to generate the tiled implementation of an operation.
|
|
|
|
Given a tile of the iteration space (as returned by
|
|
`getIterationDomain`), generate in-place the code that represents
|
|
the computation corresponding to that tile of the iteration space.
|
|
It is the responsibility of the implementation of this method in
|
|
the operation to generate the slices of the operands needed for the
|
|
tiled implementation.
|
|
- `offsets` provides the offset of the tile in the coordinate system
|
|
of the original iteration space, i.e., if an iteration space
|
|
dimension had non-zero offset, it will be included in the offset
|
|
provided here (as opposed to zero-based offset "relative" to the
|
|
iteration space).
|
|
- `sizes` provides the size of the tile.
|
|
|
|
The returned `TilingResult` must return for each result of the
|
|
untiled operation, a `Value` that is the result of the tiled
|
|
operation.
|
|
}],
|
|
/*retType=*/"::mlir::FailureOr<::mlir::TilingResult>",
|
|
/*methodName=*/"getTiledImplementation",
|
|
/*args=*/(ins
|
|
"::mlir::OpBuilder &":$b,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult> ":$offsets,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult> ":$sizes),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return {};
|
|
}]
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Method to return the position of the result tile computed by the
|
|
tiled operation.
|
|
|
|
For operations that return a value (typically a value of type
|
|
`RankedTensorType`), the generated tiled computation has to also
|
|
recompute a replacement for the results of the original operation.
|
|
The tiled implementation of the operation returns a tile of the
|
|
result(s). This methods returns information about what part of the
|
|
result tensor is computed by the tiled implementation. The manner in
|
|
which these tiles get put together to get the final result is upto
|
|
the surrounding loop construct. If an operation has no results, (for
|
|
example an operation that operates only on memrefs), then this method
|
|
need not be implemented by the operation.
|
|
- `resultNumber` is the result number of the original operation
|
|
being processed.
|
|
- `offsets` provides the offset of the tile in the coordinate system
|
|
of the original iteration space, i.e., if an iteration space
|
|
dimension had non-zero offset, it will be included in the offset
|
|
provided here (as opposed to zero-based offset "relative" to the
|
|
iteration space).
|
|
- `sizes` provides the size of the tile.
|
|
- `resultOffsets` is the offsets of the tile of the result generated
|
|
by the tiled implementation (returned by value).
|
|
- `resultSizes` is the size of the tile of the result generated
|
|
by the tiled implementation (returned by value).
|
|
|
|
Note: It is undefined behaviour if there is overlap between the
|
|
tiles of the result generated by the tiled implementation.
|
|
}],
|
|
/*retType=*/"::llvm::LogicalResult",
|
|
/*methodName=*/"getResultTilePosition",
|
|
/*args=*/(ins
|
|
"::mlir::OpBuilder &":$b,
|
|
"unsigned":$resultNumber,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult> ":$offsets,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult> ":$sizes,
|
|
"::mlir::SmallVector<::mlir::OpFoldResult> &":$resultOffsets,
|
|
"::mlir::SmallVector<::mlir::OpFoldResult> &":$resultSizes),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return failure();
|
|
}]
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Method to generate the code that produces a tile of the result.
|
|
|
|
This method is required to allow operations to be "tiled and fused"
|
|
with an (already tiled) consumer. Typically, for two operations with
|
|
producer -> consumer relation ship, to compute a tile of the
|
|
consumer a `slice` of the producer is needed. This method allows
|
|
computing that slice of the producer in-place, thereby "fusing"
|
|
the operations at tile-granularity. This method is different from
|
|
`getTiledImplementation`, which produces a tiled implementation
|
|
for a tile of the iteration space. This method produces a tiled
|
|
implementation based on the tile of producer required.
|
|
- `resultNumber` is the result of the producer used by the consumer.
|
|
- `offsets` is the offset of the slice of the producer result used by
|
|
the tiled implementation of the consumer.
|
|
- `sizes` is the size of the slice of the producer result used by the
|
|
consumer.
|
|
If fusion of the producer with the consumer is not legal for the
|
|
operation/result, this method should return failure.
|
|
|
|
Note: This method only deals with the mechanism of implementing the
|
|
fusion. In general the fusion might result in recomputation (based on
|
|
the way the result is produced by the producer and the access pattern
|
|
used in the consumer to access). This is upto the caller to handle
|
|
appropriately.
|
|
}],
|
|
/*retType=*/"::mlir::FailureOr<::mlir::TilingResult>",
|
|
/*methodName=*/"generateResultTileValue",
|
|
/*args=*/(ins
|
|
"::mlir::OpBuilder &":$b,
|
|
"unsigned":$resultNumber,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult>":$offsets,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult>":$sizes),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return failure();
|
|
}]
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Method to generate the tiled implementation of an operation that uses
|
|
the exact tiles of the given operands.
|
|
|
|
This method is required to allow operations to be "tiled and fused"
|
|
with an (already tiled) producer. Given tiles of the producer, this
|
|
method generates the tile of the consumer that uses exactly these
|
|
produced tiles. In some sense it is the "reverse" of
|
|
`generateResultTileValue`.
|
|
- `operandNumbers` is the list of operands whose tiles are "producers".
|
|
- `allOffsets` is the offset of the slice of the producer used by the
|
|
tiled implementation of the consumer.
|
|
- `allSizes` is the size of the slice of the producer used by the
|
|
consumer.
|
|
If it is illegal to fuse with a producer along the given operand tiles for
|
|
an operation, the implementation should return a failure.
|
|
}],
|
|
/*retType=*/"::mlir::FailureOr<::mlir::TilingResult>",
|
|
/*methodName=*/"getTiledImplementationFromOperandTiles",
|
|
/*args=*/(ins
|
|
"::mlir::OpBuilder &":$b,
|
|
"::mlir::ArrayRef<unsigned>":$operandNumbers,
|
|
"::mlir::ArrayRef<::mlir::SmallVector<::mlir::OpFoldResult>>":$allOffsets,
|
|
"::mlir::ArrayRef<::mlir::SmallVector<::mlir::OpFoldResult>>":$allSizes),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return failure();
|
|
}]
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Method to return the tile of the iteration domain that uses a given
|
|
tile of the operand.
|
|
|
|
This method is required to allow operations to be "tiled and fused"
|
|
with an (already tiled) producer. Given tiles of operands,
|
|
returns the tile of the iteration space that uses these tiles.
|
|
- `operandNumbers` is the list of operands whose tiles are "produced"
|
|
by the producer(s).
|
|
- `allOffsets` is the offset of the slice of the producers used by
|
|
the tiled implementation of the consumer.
|
|
- `allSizes` is the size of the slice of the producers used by the
|
|
consumer.
|
|
If it is illegal to fuse with the producer slices for an operation,
|
|
or if this mapping cannot be computed, the implementation should
|
|
return a failure.
|
|
|
|
Note that unlike the "tile consumer and fuse producer" case, the
|
|
"tile producer and fuse consumer" requires an additional method to get
|
|
the iteration tile space that encompasses all uses of the given operand
|
|
tile. The reason for this is, consider
|
|
```mlir
|
|
%1 = scf.for... {
|
|
%2 = <tiled_producer_op>
|
|
%3 = tensor.insert_slice %2 into ...
|
|
scf.yield %3
|
|
}
|
|
%4 = <consumer_op>)(... %1... )
|
|
... <some_op>(... %4 ...)
|
|
```
|
|
|
|
when fused this becomes
|
|
```
|
|
%1 = scf.for... {
|
|
%2 = <tiled_producer_op>
|
|
%3 = <tiled_consumer_op>(... %2...)
|
|
%4 = tensor.insert_slice %3 into ...
|
|
scf.yield %4
|
|
}
|
|
... <some_op>(... %1 ...)
|
|
```
|
|
|
|
i.e, when fusing the consumer, the replacement for the result of the
|
|
consumer needs to be returned to replace the uses of the consumer.
|
|
For the tile+fuse algorithm to do this it needs information about
|
|
which tile of the iteration space encompasses all uses of the tile
|
|
produced and use that to compute what are the results produced. Note
|
|
that this iteration space might be the entire iteration space of the
|
|
operation, or multiple operand tiles might map to intersecting
|
|
iteration spaces. It is upto the caller to make sure that it is still
|
|
fusable with producer in this scenario, or it must return a failure.
|
|
|
|
Note that this method is only used as a way to implement the
|
|
transformation. It does not provide guarantees on whether such a
|
|
transformation is profitable.
|
|
|
|
For most cases `getTiledImplementationFromOperandTiles` could be a
|
|
implemented using `getIterationDomainTileFromOperandTiles` +
|
|
`getTiledImplementation` methods.
|
|
}],
|
|
/*retType=*/"::llvm::LogicalResult",
|
|
/*methodName=*/"getIterationDomainTileFromOperandTiles",
|
|
/*args=*/(ins
|
|
"::mlir::OpBuilder &":$b,
|
|
"::mlir::ArrayRef<unsigned>":$operandNumbers,
|
|
"::mlir::ArrayRef<::mlir::SmallVector<::mlir::OpFoldResult>> ":$allOffsets,
|
|
"::mlir::ArrayRef<::mlir::SmallVector<::mlir::OpFoldResult>> ":$allSizes,
|
|
"::mlir::SmallVectorImpl<::mlir::OpFoldResult> &":$iterDomainOffsets,
|
|
"::mlir::SmallVectorImpl<::mlir::OpFoldResult> &":$iterDomainSizes),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return failure();
|
|
}]
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Method to return the tile of the iteration domain based
|
|
on the given tile of the certain result.
|
|
|
|
This method is required to allow operations to be "tiled and fused"
|
|
with an (already tiled) consumer. Given a tile of an result,
|
|
returns the tile of the iteration space that uses this tile.
|
|
- `resultNumber` is the result of the producer used by the consumer.
|
|
- `offsets` is the offset of the slice of the producer result used by
|
|
the tiled implementation of the consumer.
|
|
- `sizes` is the size of the slice of the producer result used by the
|
|
consumer.
|
|
If fusion of the producer with the consumer is not legal for the
|
|
result, or if this mapping cannot be computed, the implementation
|
|
should return a failure.
|
|
|
|
For most cases `generateResultTileValue` could be a implemented using
|
|
`getIterationDomainTileFromResultTile` + `getTiledImplementation`
|
|
methods.
|
|
}],
|
|
/*retType=*/"::llvm::LogicalResult",
|
|
/*methodName=*/"getIterationDomainTileFromResultTile",
|
|
/*args=*/(ins
|
|
"::mlir::OpBuilder &":$b,
|
|
"unsigned":$resultNumber,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult> ":$offsets,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult> ":$sizes,
|
|
"::mlir::SmallVectorImpl<::mlir::OpFoldResult> &":$iterDomainOffsets,
|
|
"::mlir::SmallVectorImpl<::mlir::OpFoldResult> &":$iterDomainSizes),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return failure();
|
|
}]
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Generates the scalar implementation of the operation.
|
|
|
|
Given the list `ivs` that represent points in the iteration space
|
|
(as specified by `getIterationDomain()`) returns the scalar operations
|
|
that represent the computation at that point in the iteration space.
|
|
This method is typically used as the "exit path", i.e. once all
|
|
transformations are done, this method can be used to lower to scalar
|
|
code that can then be lowered to LLVM or SPIR-V dialects.
|
|
}],
|
|
/*retType=*/"::llvm::LogicalResult",
|
|
/*methodName=*/"generateScalarImplementation",
|
|
/*args=*/(ins
|
|
"::mlir::OpBuilder &":$b,
|
|
"::mlir::Location ":$loc,
|
|
"::mlir::ValueRange ":$ivs),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return failure();
|
|
}]
|
|
>,
|
|
//===------------------------------------------------------------------===//
|
|
// Interface methods for querying fusability.
|
|
//===------------------------------------------------------------------===//
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Indicates whether it is possible to fuse this operation with the given
|
|
result slice. This method is not allowed to generate any IR.
|
|
}],
|
|
/*retTy=*/"bool",
|
|
/*methodName=*/"isOpFusableWithConsumerSlice",
|
|
/*args=*/(ins
|
|
"unsigned":$resultNumber,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult>":$offsets,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult>":$sizes
|
|
),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return false;
|
|
}]
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Indicates whether it is possible to fuse this operation with the given
|
|
list of operand slices. This method is not allowed to generate any IR.
|
|
}],
|
|
/*retTy=*/"bool",
|
|
/*methodName=*/"isOpFusableWithProducerSlices",
|
|
/*args=*/(ins
|
|
"::mlir::ArrayRef<unsigned>":$operandNumbers,
|
|
"::mlir::ArrayRef<::mlir::SmallVector<::mlir::OpFoldResult>>":$allOffsets,
|
|
"::mlir::ArrayRef<::mlir::SmallVector<::mlir::OpFoldResult>>":$allSizes
|
|
),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return false;
|
|
}]
|
|
>
|
|
];
|
|
}
|
|
|
|
def PartialReductionOpInterface :
|
|
OpInterface<"PartialReductionOpInterface", [TilingInterface]> {
|
|
let description = [{
|
|
Interface for allowing operations to expose information needed to
|
|
tile reductions using partial reduction followed by merge. This
|
|
extends the `TilingInterface` to allow splitting a reduction
|
|
dimension into a parallel dimension and reduction dimension.
|
|
The materialized inter-tile loop could either be the reduction dimension
|
|
(i.e. `ReductionTilingStrategy::PartialReductionOuterReduction`) or
|
|
the parallel dimension (i.e
|
|
`ReductionTilingStrategy::PartialReductionOuterReduction`).
|
|
}];
|
|
let cppNamespace = "::mlir";
|
|
let methods = [
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Method to generate a tensor initalized with the identity value of the
|
|
reduction operator. The tensor shape is equal to operation result
|
|
shape with new dimension for each non zero tile size.
|
|
}],
|
|
/*retType=*/"::mlir::FailureOr<SmallVector<Value>>",
|
|
/*methodName=*/"generateInitialTensorForPartialReduction",
|
|
/*args=*/(ins
|
|
"::mlir::OpBuilder &":$b,
|
|
"Location":$loc,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult>":$tileSizes,
|
|
"const ::mlir::SetVector<unsigned> &":$reductionDims),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return failure();
|
|
}]
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Method to generate a tiled version of the operation where the tiled
|
|
reduction dimension are converted to parallel dimensions with a size
|
|
less or equal to the tile size. This is meant to be used with
|
|
`mergeReductions` method which will combine the partial reductions.
|
|
The method receives the `offset` and `sizes` for all iteration space
|
|
dimensions, as well as the iteration number of the tiled reduction
|
|
dimensions (which is the induction variable of the inter-tile loop
|
|
for the reduction dimension divided by the step of the loop) in
|
|
`splitReductionIvs`.
|
|
}],
|
|
/*retType=*/"::mlir::FailureOr<TilingResult>",
|
|
/*methodName=*/"tileToPartialReduction",
|
|
/*args=*/(ins
|
|
"::mlir::OpBuilder &":$b,
|
|
"Location ":$loc,
|
|
"::mlir::ReductionTilingStrategy":$tilingStrategy,
|
|
"ValueRange":$init,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult>":$offsets,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult>":$sizes,
|
|
"const ::llvm::SetVector<unsigned> &":$reductionDims,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult>":$splitReductionIvs),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return failure();
|
|
}]
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Method to merge partial reductions for an operation that has been
|
|
tiled along the reduction dimensions. This will only apply the
|
|
reduction the operation.
|
|
}],
|
|
/*retType=*/"::mlir::FailureOr<MergeResult>",
|
|
/*methodName=*/"mergeReductions",
|
|
/*args=*/(ins
|
|
"::mlir::OpBuilder &":$b,
|
|
"Location ":$loc,
|
|
"ValueRange":$partialReduce,
|
|
"const ::mlir::SetVector<unsigned> &":$reductionDims),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return failure();
|
|
}]
|
|
>,
|
|
InterfaceMethod<
|
|
/*desc=*/[{
|
|
Method to return the position of the partial result tile computed by
|
|
the tiled operation. This is same as
|
|
TilingInterface:::getResultTilePosition, but determines the result
|
|
tile position for partial reduction.
|
|
The method receives the `offset` and `sizes` for all iteration space
|
|
dimensions, as well as the iteration number of the tiled reduction
|
|
dimensions (which is the induction variable of the inter-tile loop
|
|
for the reduction dimension divided by the tile size specified) in
|
|
`splitReductionIvs`.
|
|
}],
|
|
/*retType=*/"::llvm::LogicalResult",
|
|
/*methodName=*/"getPartialResultTilePosition",
|
|
/*args=*/(ins
|
|
"::mlir::OpBuilder &":$b,
|
|
"unsigned":$resultNumber,
|
|
"ReductionTilingStrategy":$tilingStrategy,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult> ":$offsets,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult> ":$sizes,
|
|
"const ::mlir::SetVector<unsigned> &":$reductionDims,
|
|
"::mlir::ArrayRef<::mlir::OpFoldResult>":$splitReductionIvs,
|
|
"::mlir::SmallVector<::mlir::OpFoldResult> &":$resultOffsets,
|
|
"::mlir::SmallVector<::mlir::OpFoldResult> &":$resultSizes),
|
|
/*methodBody=*/"",
|
|
/*defaultImplementation=*/[{
|
|
return failure();
|
|
}]
|
|
>
|
|
];
|
|
}
|
|
#endif // MLIR_TILINGINTERFACE
|