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
19using namespace mlir;
20using 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.
26constexpr const ::llvm::StringLiteral BufferizationDialect::kWritableAttrName;
27
28/// Attribute name used to mark the bufferization layout for region arguments
29/// during One-Shot Module Bufferize.
30constexpr 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.)
41constexpr const ::llvm::StringLiteral BufferizationDialect::kManualDeallocation;
42
43//===----------------------------------------------------------------------===//
44// Bufferization Dialect Interfaces
45//===----------------------------------------------------------------------===//
46
47namespace {
48struct 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
57template <typename Tensor>
58struct BuiltinTensorExternalModel
59 : TensorLikeType::ExternalModel<BuiltinTensorExternalModel<Tensor>,
60 Tensor> {
61 llvm::FailureOr<BufferLikeType> getBufferType(
62 mlir::Type tensor, const BufferizationOptions &options,
63 llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
64 auto tensorType = cast<TensorType>(Val&: tensor);
65 auto memSpace = options.defaultMemorySpaceFn(tensorType);
66 if (!memSpace.has_value())
67 return emitError() << "could not infer memory space";
68
69 return cast<BufferLikeType>(
70 Val: getMemRefType(tensorType, options, /*layout=*/{}, memorySpace: *memSpace));
71 }
72
73 mlir::LogicalResult verifyCompatibleBufferType(
74 mlir::Type tensor, BufferLikeType bufferType,
75 llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
76 assert(isa<TensorType>(tensor) && "expected tensor type");
77 assert(isa<BaseMemRefType>(bufferType) && "expected memref type");
78
79 auto tensorType = cast<ShapedType>(Val&: tensor);
80 auto memrefType = cast<ShapedType>(Val&: bufferType);
81
82 if (tensorType.getShape() != memrefType.getShape())
83 return emitError() << "shapes do not match";
84
85 if (tensorType.getElementType() != memrefType.getElementType())
86 return emitError() << "element types do not match";
87
88 return mlir::success();
89 }
90};
91
92template <typename MemRef>
93struct BuiltinMemRefExternalModel
94 : BufferLikeType::ExternalModel<BuiltinMemRefExternalModel<MemRef>,
95 MemRef> {};
96} // namespace
97
98//===----------------------------------------------------------------------===//
99// Bufferization Dialect
100//===----------------------------------------------------------------------===//
101
102void mlir::bufferization::BufferizationDialect::initialize() {
103 addOperations<
104#define GET_OP_LIST
105#include "mlir/Dialect/Bufferization/IR/BufferizationOps.cpp.inc"
106 >();
107 addInterfaces<BufferizationInlinerInterface>();
108
109 // Note: Unlike with other external models, declaring bufferization's
110 // "promised interfaces" in builtins for TensorLike and BufferLike type
111 // interfaces is not possible (due to builtins being independent of
112 // bufferization). Thus, the compromise is to attach these interfaces directly
113 // during dialect initialization.
114 RankedTensorType::attachInterface<
115 BuiltinTensorExternalModel<RankedTensorType>>(context&: *getContext());
116 UnrankedTensorType::attachInterface<
117 BuiltinTensorExternalModel<UnrankedTensorType>>(context&: *getContext());
118 MemRefType::attachInterface<BuiltinMemRefExternalModel<MemRefType>>(
119 context&: *getContext());
120 UnrankedMemRefType::attachInterface<
121 BuiltinMemRefExternalModel<UnrankedMemRefType>>(context&: *getContext());
122}
123
124LogicalResult BufferizationDialect::verifyRegionArgAttribute(
125 Operation *op, unsigned /*regionIndex*/, unsigned argIndex,
126 NamedAttribute attr) {
127 if (attr.getName() == kWritableAttrName) {
128 if (!llvm::isa<BoolAttr>(Val: attr.getValue())) {
129 return op->emitError() << "'" << kWritableAttrName
130 << "' is expected to be a boolean attribute";
131 }
132 if (!isa<FunctionOpInterface>(Val: op))
133 return op->emitError() << "expected '" << kWritableAttrName
134 << "' to be used on function-like operations";
135 if (cast<FunctionOpInterface>(Val: op).isExternal())
136 return op->emitError() << "'" << kWritableAttrName
137 << "' is invalid on external functions";
138 return success();
139 }
140 if (attr.getName() == kBufferAccessAttrName) {
141 if (!llvm::isa<StringAttr>(Val: attr.getValue())) {
142 return op->emitError() << "'" << kBufferAccessAttrName
143 << "' is expected to be a string attribute";
144 }
145 StringRef str = llvm::cast<StringAttr>(Val: attr.getValue()).getValue();
146 if (str != "none" && str != "read" && str != "write" && str != "read-write")
147 return op->emitError()
148 << "invalid value for '" << kBufferAccessAttrName << "'";
149 if (!isa<FunctionOpInterface>(Val: op))
150 return op->emitError() << "expected '" << kBufferAccessAttrName
151 << "' to be used on function-like operations";
152 return success();
153 }
154 if (attr.getName() == kBufferLayoutAttrName) {
155 if (!llvm::isa<MemRefLayoutAttrInterface>(Val: attr.getValue())) {
156 return op->emitError() << "'" << kBufferLayoutAttrName
157 << "' is expected to be a memref layout attribute";
158 }
159 if (!isa<FunctionOpInterface>(Val: op))
160 return op->emitError() << "expected '" << kBufferLayoutAttrName
161 << "' to be used on function-like operations";
162 return success();
163 }
164 return op->emitError() << "attribute '" << kBufferLayoutAttrName
165 << "' not supported as a region arg attribute by the "
166 "bufferization dialect";
167}
168
169LogicalResult
170BufferizationDialect::verifyOperationAttribute(Operation *op,
171 NamedAttribute attr) {
172 using bufferization::BufferizableOpInterface;
173
174 if (attr.getName() == kManualDeallocation) {
175 if (!mlir::hasEffect<MemoryEffects::Allocate>(op) &&
176 !mlir::hasEffect<MemoryEffects::Free>(op))
177 return op->emitOpError(message: "attribute '")
178 << kManualDeallocation
179 << "' can be used only on ops that have an allocation and/or free "
180 "side effect";
181 return success();
182 }
183
184 return op->emitError()
185 << "attribute '" << attr.getName()
186 << "' not supported as an op attribute by the bufferization dialect";
187}
188

source code of mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp