83 lines
3.2 KiB
C++
83 lines
3.2 KiB
C++
//===- MergeConsecutiveInsertExtractSlicePatterns.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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Affine/ViewLikeInterfaceUtils.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/OpDefinition.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::tensor;
|
|
|
|
namespace {
|
|
/// Merges consecutive tensor.extract_slice ops into one.
|
|
struct MergeConsecutiveExtractSlice : public OpRewritePattern<ExtractSliceOp> {
|
|
using OpRewritePattern::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(ExtractSliceOp nextOp,
|
|
PatternRewriter &rewriter) const override {
|
|
auto prevOp = nextOp.getSource().getDefiningOp<ExtractSliceOp>();
|
|
if (!prevOp)
|
|
return failure();
|
|
|
|
SmallVector<OpFoldResult> newOffsets, newSizes, newStrides;
|
|
if (failed(mergeOffsetsSizesAndStrides(rewriter, nextOp.getLoc(), prevOp,
|
|
nextOp, prevOp.getDroppedDims(),
|
|
newOffsets, newSizes, newStrides)))
|
|
return failure();
|
|
|
|
rewriter.replaceOpWithNewOp<ExtractSliceOp>(nextOp, nextOp.getType(),
|
|
prevOp.getSource(), newOffsets,
|
|
newSizes, newStrides);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
/// Merges consecutive tensor.insert_slice ops into one.
|
|
struct MergeConsecutiveInsertSlice : public OpRewritePattern<InsertSliceOp> {
|
|
using OpRewritePattern::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(InsertSliceOp nextOp,
|
|
PatternRewriter &rewriter) const override {
|
|
auto prevOp = nextOp.getSource().getDefiningOp<InsertSliceOp>();
|
|
if (!prevOp)
|
|
return failure();
|
|
|
|
if (!prevOp.hasUnitStride() || !nextOp.hasUnitStride())
|
|
return failure();
|
|
|
|
// The first insert_slice op should be rank reducing to make sure we cover
|
|
// the full source tensor to be inserted in the second insert_slice op.
|
|
SliceVerificationResult result =
|
|
isRankReducedType(prevOp.getDestType(), prevOp.getSourceType());
|
|
if (result != SliceVerificationResult::Success)
|
|
return failure();
|
|
|
|
// Dynamic dimensions can pass rank reducing check in the above, e.g,
|
|
// inserting <?xf32> into <1x?x1xf32>. For such cases we cannot be certain
|
|
// the dynamic size covers the full tensor.
|
|
if (!prevOp.getSourceType().hasStaticShape() ||
|
|
!prevOp.getDestType().hasStaticShape())
|
|
return failure();
|
|
|
|
rewriter.replaceOpWithNewOp<InsertSliceOp>(
|
|
nextOp, prevOp.getSource(), nextOp.getDest(), nextOp.getMixedOffsets(),
|
|
nextOp.getMixedSizes(), nextOp.getMixedStrides());
|
|
return success();
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
void mlir::tensor::populateMergeConsecutiveInsertExtractSlicePatterns(
|
|
RewritePatternSet &patterns) {
|
|
patterns.add<MergeConsecutiveExtractSlice, MergeConsecutiveInsertSlice>(
|
|
patterns.getContext());
|
|
}
|