1 | //===- SubsetInsertionOpInterfaceImpl.cpp - Tensor subsets ----------------===// |
2 | // |
3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
4 | // See https://llvm.org/LICENSE.txt for license information. |
5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | |
9 | #include "mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h" |
10 | |
11 | #include "mlir/Dialect/Linalg/IR/Linalg.h" |
12 | #include "mlir/Interfaces/SubsetOpInterface.h" |
13 | |
14 | using namespace mlir; |
15 | using namespace mlir::linalg; |
16 | |
17 | namespace { |
18 | struct LinalgCopyOpSubsetOpInterface |
19 | : public SubsetOpInterface::ExternalModel<LinalgCopyOpSubsetOpInterface, |
20 | linalg::CopyOp> { |
21 | bool operatesOnEquivalentSubset( |
22 | Operation *op, SubsetOpInterface candidate, |
23 | function_ref<bool(Value, Value)> equivalenceFn) const { |
24 | // linalg.copy operates on the entire destination tensor. |
25 | if (auto otherCopyOp = dyn_cast<linalg::CopyOp>(candidate.getOperation())) |
26 | return equivalenceFn(cast<linalg::CopyOp>(op).getOutputs()[0], |
27 | otherCopyOp.getOutputs()[0]); |
28 | // In the absence of an analysis, "false" is a conservative way to implement |
29 | // this interface. |
30 | return false; |
31 | } |
32 | |
33 | bool operatesOnDisjointSubset( |
34 | Operation *op, SubsetOpInterface candidate, |
35 | function_ref<bool(Value, Value)> equivalenceFn) const { |
36 | // In the absence of an analysis, "false" is a conservative way to implement |
37 | // this interface. |
38 | return false; |
39 | } |
40 | }; |
41 | |
42 | struct LinalgCopyOpInterface |
43 | : public SubsetInsertionOpInterface::ExternalModel<LinalgCopyOpInterface, |
44 | linalg::CopyOp> { |
45 | OpOperand &getSourceOperand(Operation *op) const { |
46 | auto copyOp = cast<CopyOp>(op); |
47 | return llvm::getSingleElement(copyOp.getInputsMutable()); |
48 | } |
49 | |
50 | bool |
51 | isEquivalentSubset(Operation *op, Value candidate, |
52 | function_ref<bool(Value, Value)> equivalenceFn) const { |
53 | auto copyOp = cast<CopyOp>(op); |
54 | return equivalenceFn(candidate, |
55 | llvm::getSingleElement(copyOp.getOutputs())); |
56 | } |
57 | |
58 | Value (Operation *op, OpBuilder &builder, |
59 | Location loc) const { |
60 | auto copyOp = cast<CopyOp>(op); |
61 | return llvm::getSingleElement(copyOp.getOutputs()); |
62 | } |
63 | |
64 | SmallVector<Value> |
65 | (Operation *op) const { |
66 | auto copyOp = cast<CopyOp>(op); |
67 | return {llvm::getSingleElement(copyOp.getOutputs())}; |
68 | } |
69 | }; |
70 | } // namespace |
71 | |
72 | void mlir::linalg::registerSubsetOpInterfaceExternalModels( |
73 | DialectRegistry ®istry) { |
74 | registry.addExtension(extensionFn: +[](MLIRContext *ctx, linalg::LinalgDialect *dialect) { |
75 | linalg::CopyOp::attachInterface<LinalgCopyOpSubsetOpInterface>(*ctx); |
76 | linalg::CopyOp::attachInterface<LinalgCopyOpInterface>(*ctx); |
77 | }); |
78 | } |
79 | |