2b2ce50fe8
This commit adds an API (`tileAndFuseConsumerOfSlice`) to fuse consumer to a producer within scf.for/scf.forall loop. To support this two new methods are added to the `TilingInterface` - `getIterationDomainTileFromOperandTile` - `getTiledImplementationFromOperandTile`. Consumer operations that implement this method can be used to be fused with tiled producer operands in a manner similar to (but essentially the inverse of) the fusion of an untiled producer with a tiled consumer. Note that this only does one `tiled producer` -> `consumer` fusion. This could be called repeatedly for fusing multiple consumers. The current implementation also is conservative in when this kicks in (like single use of the value returned by the inter-tile loops that surround the tiled producer, etc.) These can be relaxed over time. Signed-off-by: Abhishek Varma <abhvarma@amd.com> --------- Signed-off-by: Abhishek Varma <abhvarma@amd.com> Signed-off-by: Abhishek Varma <avarma094@gmail.com> Co-authored-by: cxy <chenxunyu1993@gmail.com>
66 lines
2.4 KiB
C++
66 lines
2.4 KiB
C++
//===- SwapExtractSliceWithProducerPatterns.cpp ---------------------------===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// Swap a `tensor.extract_slice` with the producer of the source if the producer
|
|
// implements the `TilingInterface`. When used in conjunction with tiling this
|
|
// effectively tiles + fuses the producer with its consumer.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
|
|
#include "mlir/Dialect/Utils/StaticValueUtils.h"
|
|
#include "mlir/Interfaces/TilingInterface.h"
|
|
|
|
using namespace mlir;
|
|
|
|
FailureOr<TilingResult> tensor::replaceExtractSliceWithTiledProducer(
|
|
OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producer) {
|
|
auto producerOp = dyn_cast<TilingInterface>(producer.getOwner());
|
|
if (!producerOp)
|
|
return failure();
|
|
|
|
// `TilingInterface` currently only supports strides being 1.
|
|
if (llvm::any_of(sliceOp.getMixedStrides(), [](OpFoldResult ofr) {
|
|
return !isConstantIntValue(ofr, 1);
|
|
}))
|
|
return failure();
|
|
|
|
FailureOr<TilingResult> tiledResult = producerOp.generateResultTileValue(
|
|
builder, producer.getResultNumber(), sliceOp.getMixedOffsets(),
|
|
sliceOp.getMixedSizes());
|
|
if (failed(tiledResult))
|
|
return failure();
|
|
|
|
return *tiledResult;
|
|
}
|
|
|
|
FailureOr<TilingResult> tensor::replaceInsertSliceWithTiledConsumer(
|
|
OpBuilder &builder, OffsetSizeAndStrideOpInterface sliceOp,
|
|
OpOperand &consumer) {
|
|
auto consumerOp = dyn_cast<TilingInterface>(consumer.getOwner());
|
|
if (!consumerOp)
|
|
return failure();
|
|
|
|
// `TilingInterface` currently only supports strides being 1.
|
|
if (llvm::any_of(sliceOp.getMixedStrides(), [](OpFoldResult ofr) {
|
|
return !isConstantIntValue(ofr, 1);
|
|
}))
|
|
return failure();
|
|
|
|
FailureOr<TilingResult> tiledResult =
|
|
consumerOp.getTiledImplementationFromOperandTile(
|
|
builder, consumer.getOperandNumber(), sliceOp.getMixedOffsets(),
|
|
sliceOp.getMixedSizes());
|
|
if (failed(tiledResult))
|
|
return failure();
|
|
|
|
return *tiledResult;
|
|
}
|