1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
|
//===- 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.
// TODO: move to FoldTensorSubsetOps and unify APIs with FoldMemRefAliasOps.
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(affine::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.
// TODO: move to FoldTensorSubsetOps and unify APIs with FoldMemRefAliasOps.
template <typename OpTy>
struct MergeConsecutiveInsertSlice : public OpRewritePattern<OpTy> {
using OpRewritePattern<OpTy>::OpRewritePattern;
LogicalResult matchAndRewrite(OpTy nextOp,
PatternRewriter &rewriter) const override {
auto prevOp = nextOp.getSource().template 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<OpTy>(
nextOp, prevOp.getSource(), nextOp.getDest(), nextOp.getMixedOffsets(),
nextOp.getMixedSizes(), nextOp.getMixedStrides());
return success();
}
};
} // namespace
void mlir::tensor::populateMergeConsecutiveInsertExtractSlicePatterns(
RewritePatternSet &patterns) {
patterns.add<MergeConsecutiveExtractSlice,
MergeConsecutiveInsertSlice<InsertSliceOp>,
MergeConsecutiveInsertSlice<ParallelInsertSliceOp>>(
patterns.getContext());
}
|
