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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
|
//===- SubsetInsertionOpInterfaceImpl.cpp - Tensor subsets ----------------===//
//
// 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/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Interfaces/SubsetOpInterface.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
using namespace mlir;
using namespace mlir::tensor;
namespace {
struct ExtractSliceOpSubsetOpInterface
: public SubsetOpInterface::ExternalModel<ExtractSliceOpSubsetOpInterface,
tensor::ExtractSliceOp> {
FailureOr<HyperrectangularSlice>
getAccessedHyperrectangularSlice(Operation *op) const {
return HyperrectangularSlice(cast<OffsetSizeAndStrideOpInterface>(op));
}
};
struct ExtractSliceOpSubsetExtractionOpInterface
: public SubsetExtractionOpInterface::ExternalModel<
ExtractSliceOpSubsetExtractionOpInterface, tensor::ExtractSliceOp> {
OpOperand &getSourceOperand(Operation *op) const {
return cast<tensor::ExtractSliceOp>(op).getSourceMutable();
}
};
template <typename OpTy>
struct InsertSliceLikeOpSubsetOpInterface
: public SubsetOpInterface::ExternalModel<
InsertSliceLikeOpSubsetOpInterface<OpTy>, OpTy> {
FailureOr<HyperrectangularSlice>
getAccessedHyperrectangularSlice(Operation *op) const {
return HyperrectangularSlice(cast<OffsetSizeAndStrideOpInterface>(op));
}
};
template <typename OpTy>
struct InsertSliceLikeOpSubsetInsertionOpInterface
: public SubsetInsertionOpInterface::ExternalModel<
InsertSliceLikeOpSubsetInsertionOpInterface<OpTy>, OpTy> {
OpOperand &getSourceOperand(Operation *op) const {
return cast<OpTy>(op).getSourceMutable();
}
OpOperand &getDestinationOperand(Operation *op) const {
return cast<OpTy>(op).getDestMutable();
}
Value buildSubsetExtraction(Operation *op, OpBuilder &builder,
Location loc) const {
auto insertSliceOp = cast<OpTy>(op);
auto extractOp = builder.create<tensor::ExtractSliceOp>(
loc, insertSliceOp.getSourceType(), insertSliceOp.getDest(),
insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
insertSliceOp.getMixedStrides());
return extractOp.getResult();
}
SmallVector<Value>
getValuesNeededToBuildSubsetExtraction(Operation *op) const {
auto insertSliceOp = cast<OpTy>(op);
SmallVector<Value> neededValues;
// Collect all values that are needed to construct the replacement op.
neededValues.append(insertSliceOp.getOffsets().begin(),
insertSliceOp.getOffsets().end());
neededValues.append(insertSliceOp.getSizes().begin(),
insertSliceOp.getSizes().end());
neededValues.append(insertSliceOp.getStrides().begin(),
insertSliceOp.getStrides().end());
neededValues.push_back(insertSliceOp.getDest());
return neededValues;
}
};
} // namespace
void mlir::tensor::registerSubsetOpInterfaceExternalModels(
DialectRegistry ®istry) {
registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
// Note: `SubsetExtractionOpInterface` and `SubsetInsertionOpInterface`
// require `SubsetOpInterface`.
ExtractSliceOp::attachInterface<ExtractSliceOpSubsetOpInterface>(*ctx);
ExtractSliceOp::attachInterface<ExtractSliceOpSubsetExtractionOpInterface>(
*ctx);
InsertSliceOp::attachInterface<
InsertSliceLikeOpSubsetOpInterface<InsertSliceOp>>(*ctx);
InsertSliceOp::attachInterface<
InsertSliceLikeOpSubsetInsertionOpInterface<InsertSliceOp>>(*ctx);
ParallelInsertSliceOp::attachInterface<
InsertSliceLikeOpSubsetOpInterface<ParallelInsertSliceOp>>(*ctx);
ParallelInsertSliceOp::attachInterface<
InsertSliceLikeOpSubsetInsertionOpInterface<ParallelInsertSliceOp>>(
*ctx);
});
}
|
