1 | //===----------------------------------------------------------------------===// |
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/Affine/IR/AffineOps.h" |
10 | #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h" |
11 | #include "mlir/Dialect/Bufferization/IR/Bufferization.h" |
12 | #include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h" |
13 | #include "mlir/Dialect/MemRef/IR/MemRef.h" |
14 | #include "mlir/Dialect/Tensor/IR/Tensor.h" |
15 | #include "mlir/IR/BuiltinTypes.h" |
16 | #include "mlir/Interfaces/FunctionInterfaces.h" |
17 | #include "mlir/Transforms/InliningUtils.h" |
18 | |
19 | using namespace mlir; |
20 | using namespace mlir::bufferization; |
21 | |
22 | #include "mlir/Dialect/Bufferization/IR/BufferizationOpsDialect.cpp.inc" |
23 | |
24 | /// Attribute name used to mark function arguments who's buffers can be written |
25 | /// to during One-Shot Module Bufferize. |
26 | constexpr const ::llvm::StringLiteral BufferizationDialect::kWritableAttrName; |
27 | |
28 | /// Attribute name used to mark the bufferization layout for region arguments |
29 | /// during One-Shot Module Bufferize. |
30 | constexpr const ::llvm::StringLiteral |
31 | BufferizationDialect::kBufferLayoutAttrName; |
32 | |
33 | /// An attribute that can be attached to ops with an allocation and/or |
34 | /// deallocation side effect. It indicates that the op is under a "manual |
35 | /// deallocation" scheme. In the case of an allocation op, the returned |
36 | /// value is *not* an automatically managed allocation and assigned an |
37 | /// ownership of "false". Furthermore, only deallocation ops that are |
38 | /// guaranteed to deallocate a buffer under "manual deallocation" are |
39 | /// allowed to have this attribute. (Deallocation ops without this |
40 | /// attribute are rejected by the ownership-based buffer deallocation pass.) |
41 | constexpr const ::llvm::StringLiteral BufferizationDialect::kManualDeallocation; |
42 | |
43 | //===----------------------------------------------------------------------===// |
44 | // Bufferization Dialect Interfaces |
45 | //===----------------------------------------------------------------------===// |
46 | |
47 | namespace { |
48 | struct BufferizationInlinerInterface : public DialectInlinerInterface { |
49 | using DialectInlinerInterface::DialectInlinerInterface; |
50 | |
51 | /// Operations in Bufferization dialect are always legal to inline. |
52 | bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final { |
53 | return true; |
54 | } |
55 | }; |
56 | |
57 | template <typename Tensor> |
58 | struct BuiltinTensorExternalModel |
59 | : TensorLikeType::ExternalModel<BuiltinTensorExternalModel<Tensor>, |
60 | Tensor> {}; |
61 | |
62 | template <typename MemRef> |
63 | struct BuiltinMemRefExternalModel |
64 | : BufferLikeType::ExternalModel<BuiltinMemRefExternalModel<MemRef>, |
65 | MemRef> {}; |
66 | } // namespace |
67 | |
68 | //===----------------------------------------------------------------------===// |
69 | // Bufferization Dialect |
70 | //===----------------------------------------------------------------------===// |
71 | |
72 | void mlir::bufferization::BufferizationDialect::initialize() { |
73 | addOperations< |
74 | #define GET_OP_LIST |
75 | #include "mlir/Dialect/Bufferization/IR/BufferizationOps.cpp.inc" |
76 | >(); |
77 | addInterfaces<BufferizationInlinerInterface>(); |
78 | |
79 | // Note: Unlike with other external models, declaring bufferization's |
80 | // "promised interfaces" in builtins for TensorLike and BufferLike type |
81 | // interfaces is not possible (due to builtins being independent of |
82 | // bufferization). Thus, the compromise is to attach these interfaces directly |
83 | // during dialect initialization. |
84 | RankedTensorType::attachInterface< |
85 | BuiltinTensorExternalModel<RankedTensorType>>(*getContext()); |
86 | UnrankedTensorType::attachInterface< |
87 | BuiltinTensorExternalModel<UnrankedTensorType>>(*getContext()); |
88 | MemRefType::attachInterface<BuiltinMemRefExternalModel<MemRefType>>( |
89 | *getContext()); |
90 | UnrankedMemRefType::attachInterface< |
91 | BuiltinMemRefExternalModel<UnrankedMemRefType>>(*getContext()); |
92 | } |
93 | |
94 | LogicalResult BufferizationDialect::verifyRegionArgAttribute( |
95 | Operation *op, unsigned /*regionIndex*/, unsigned argIndex, |
96 | NamedAttribute attr) { |
97 | if (attr.getName() == kWritableAttrName) { |
98 | if (!llvm::isa<BoolAttr>(attr.getValue())) { |
99 | return op->emitError() << "'" << kWritableAttrName |
100 | << "' is expected to be a boolean attribute" ; |
101 | } |
102 | if (!isa<FunctionOpInterface>(op)) |
103 | return op->emitError() << "expected '" << kWritableAttrName |
104 | << "' to be used on function-like operations" ; |
105 | if (cast<FunctionOpInterface>(op).isExternal()) |
106 | return op->emitError() << "'" << kWritableAttrName |
107 | << "' is invalid on external functions" ; |
108 | return success(); |
109 | } |
110 | if (attr.getName() == kBufferAccessAttrName) { |
111 | if (!llvm::isa<StringAttr>(attr.getValue())) { |
112 | return op->emitError() << "'" << kBufferAccessAttrName |
113 | << "' is expected to be a string attribute" ; |
114 | } |
115 | StringRef str = llvm::cast<StringAttr>(attr.getValue()).getValue(); |
116 | if (str != "none" && str != "read" && str != "write" && str != "read-write" ) |
117 | return op->emitError() |
118 | << "invalid value for '" << kBufferAccessAttrName << "'" ; |
119 | if (!isa<FunctionOpInterface>(op)) |
120 | return op->emitError() << "expected '" << kBufferAccessAttrName |
121 | << "' to be used on function-like operations" ; |
122 | return success(); |
123 | } |
124 | if (attr.getName() == kBufferLayoutAttrName) { |
125 | if (!llvm::isa<MemRefLayoutAttrInterface>(attr.getValue())) { |
126 | return op->emitError() << "'" << kBufferLayoutAttrName |
127 | << "' is expected to be a memref layout attribute" ; |
128 | } |
129 | if (!isa<FunctionOpInterface>(op)) |
130 | return op->emitError() << "expected '" << kBufferLayoutAttrName |
131 | << "' to be used on function-like operations" ; |
132 | return success(); |
133 | } |
134 | return op->emitError() << "attribute '" << kBufferLayoutAttrName |
135 | << "' not supported as a region arg attribute by the " |
136 | "bufferization dialect" ; |
137 | } |
138 | |
139 | LogicalResult |
140 | BufferizationDialect::verifyOperationAttribute(Operation *op, |
141 | NamedAttribute attr) { |
142 | using bufferization::BufferizableOpInterface; |
143 | |
144 | if (attr.getName() == kManualDeallocation) { |
145 | if (!mlir::hasEffect<MemoryEffects::Allocate>(op) && |
146 | !mlir::hasEffect<MemoryEffects::Free>(op)) |
147 | return op->emitOpError("attribute '" ) |
148 | << kManualDeallocation |
149 | << "' can be used only on ops that have an allocation and/or free " |
150 | "side effect" ; |
151 | return success(); |
152 | } |
153 | |
154 | return op->emitError() |
155 | << "attribute '" << attr.getName() |
156 | << "' not supported as an op attribute by the bufferization dialect" ; |
157 | } |
158 | |