1//===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
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/Arith/Transforms/BufferizableOpInterfaceImpl.h"
10#include "mlir/Dialect/Arith/IR/Arith.h"
11#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
12#include "mlir/Dialect/Bufferization/Transforms/BufferUtils.h"
13#include "mlir/Dialect/MemRef/IR/MemRef.h"
14#include "mlir/IR/Attributes.h"
15#include "mlir/IR/Dialect.h"
16#include "mlir/IR/Operation.h"
17
18using namespace mlir;
19using namespace mlir::bufferization;
20
21namespace {
22/// Bufferization of arith.constant. Replace with memref.get_global.
23struct ConstantOpInterface
24 : public BufferizableOpInterface::ExternalModel<ConstantOpInterface,
25 arith::ConstantOp> {
26 LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
27 const BufferizationOptions &options,
28 BufferizationState &state) const {
29 auto constantOp = cast<arith::ConstantOp>(Val: op);
30 auto type = dyn_cast<RankedTensorType>(Val: constantOp.getType());
31
32 // Only ranked tensors are supported.
33 if (!type)
34 return failure();
35
36 Attribute memorySpace;
37 if (auto memSpace = options.defaultMemorySpaceFn(type))
38 memorySpace = *memSpace;
39 else
40 return constantOp->emitError(message: "could not infer memory space");
41
42 // Only constants inside a module are supported.
43 auto moduleOp = constantOp->getParentOfType<ModuleOp>();
44 if (!moduleOp)
45 return failure();
46
47 // Create global memory segment and replace tensor with memref pointing to
48 // that memory segment.
49 FailureOr<memref::GlobalOp> globalOp =
50 getGlobalFor(constantOp, symbolTables&: state.getSymbolTables(),
51 alignment: options.bufferAlignment, memorySpace);
52 if (failed(Result: globalOp))
53 return failure();
54 memref::GlobalOp globalMemref = *globalOp;
55 replaceOpWithNewBufferizedOp<memref::GetGlobalOp>(
56 rewriter, op, args: globalMemref.getType(), args: globalMemref.getName());
57
58 return success();
59 }
60
61 bool isWritable(Operation *op, Value value,
62 const AnalysisState &state) const {
63 // Memory locations returned by memref::GetGlobalOp may not be written to.
64 assert(isa<OpResult>(value));
65 return false;
66 }
67};
68
69struct IndexCastOpInterface
70 : public BufferizableOpInterface::ExternalModel<IndexCastOpInterface,
71 arith::IndexCastOp> {
72 bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
73 const AnalysisState &state) const {
74 return false;
75 }
76
77 bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
78 const AnalysisState &state) const {
79 return false;
80 }
81
82 AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
83 const AnalysisState &state) const {
84 return {{op->getResult(idx: 0), BufferRelation::Equivalent}};
85 }
86
87 LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
88 const BufferizationOptions &options,
89 BufferizationState &state) const {
90 auto castOp = cast<arith::IndexCastOp>(Val: op);
91 auto resultTensorType = cast<TensorType>(Val: castOp.getType());
92
93 FailureOr<Value> source =
94 getBuffer(rewriter, value: castOp.getIn(), options, state);
95 if (failed(Result: source))
96 return failure();
97 auto sourceType = cast<BaseMemRefType>(Val: source->getType());
98
99 // Result type should have same layout and address space as the source type.
100 BaseMemRefType resultType;
101 if (auto rankedMemRefType = dyn_cast<MemRefType>(Val&: sourceType)) {
102 resultType = MemRefType::get(
103 shape: rankedMemRefType.getShape(), elementType: resultTensorType.getElementType(),
104 layout: rankedMemRefType.getLayout(), memorySpace: rankedMemRefType.getMemorySpace());
105 } else {
106 auto unrankedMemrefType = cast<UnrankedMemRefType>(Val&: sourceType);
107 resultType = UnrankedMemRefType::get(elementType: resultTensorType.getElementType(),
108 memorySpace: unrankedMemrefType.getMemorySpace());
109 }
110
111 replaceOpWithNewBufferizedOp<arith::IndexCastOp>(rewriter, op, args&: resultType,
112 args&: *source);
113 return success();
114 }
115};
116
117/// Bufferization of arith.select. Just replace the operands.
118struct SelectOpInterface
119 : public BufferizableOpInterface::ExternalModel<SelectOpInterface,
120 arith::SelectOp> {
121 bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
122 const AnalysisState &state) const {
123 return false;
124 }
125
126 bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
127 const AnalysisState &state) const {
128 return false;
129 }
130
131 AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
132 const AnalysisState &state) const {
133 return {{op->getOpResult(idx: 0) /*result*/, BufferRelation::Equivalent,
134 /*isDefinite=*/false}};
135 }
136
137 LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
138 const BufferizationOptions &options,
139 BufferizationState &state) const {
140 auto selectOp = cast<arith::SelectOp>(Val: op);
141 Location loc = selectOp.getLoc();
142
143 // Elementwise conditions are not supported yet. To bufferize such an op,
144 // it could be lowered to an elementwise "linalg.generic" with a new
145 // "tensor.empty" out tensor, followed by "empty tensor elimination". Such
146 // IR will bufferize.
147 if (!selectOp.getCondition().getType().isInteger(width: 1))
148 return op->emitOpError(message: "only i1 condition values are supported");
149
150 // TODO: It would be more efficient to copy the result of the `select` op
151 // instead of its OpOperands. In the worst case, 2 copies are inserted at
152 // the moment (one for each tensor). When copying the op result, only one
153 // copy would be needed.
154 FailureOr<Value> maybeTrueBuffer =
155 getBuffer(rewriter, value: selectOp.getTrueValue(), options, state);
156 FailureOr<Value> maybeFalseBuffer =
157 getBuffer(rewriter, value: selectOp.getFalseValue(), options, state);
158 if (failed(Result: maybeTrueBuffer) || failed(Result: maybeFalseBuffer))
159 return failure();
160 Value trueBuffer = *maybeTrueBuffer;
161 Value falseBuffer = *maybeFalseBuffer;
162
163 // The "true" and the "false" operands must have the same type. If the
164 // buffers have different types, they differ only in their layout map. Cast
165 // both of them to the most dynamic MemRef type.
166 if (trueBuffer.getType() != falseBuffer.getType()) {
167 auto targetType = bufferization::detail::asMemRefType(
168 bufferType: bufferization::getBufferType(value: selectOp.getResult(), options, state));
169 if (failed(Result: targetType))
170 return failure();
171 if (trueBuffer.getType() != *targetType)
172 trueBuffer =
173 rewriter.create<memref::CastOp>(location: loc, args&: *targetType, args&: trueBuffer);
174 if (falseBuffer.getType() != *targetType)
175 falseBuffer =
176 rewriter.create<memref::CastOp>(location: loc, args&: *targetType, args&: falseBuffer);
177 }
178
179 replaceOpWithNewBufferizedOp<arith::SelectOp>(
180 rewriter, op, args: selectOp.getCondition(), args&: trueBuffer, args&: falseBuffer);
181 return success();
182 }
183
184 FailureOr<BufferLikeType>
185 getBufferType(Operation *op, Value value, const BufferizationOptions &options,
186 const BufferizationState &state,
187 SmallVector<Value> &invocationStack) const {
188 auto selectOp = cast<arith::SelectOp>(Val: op);
189 assert(value == selectOp.getResult() && "invalid value");
190 auto trueType =
191 bufferization::detail::asMemRefType(bufferType: bufferization::getBufferType(
192 value: selectOp.getTrueValue(), options, state, invocationStack));
193 auto falseType =
194 bufferization::detail::asMemRefType(bufferType: bufferization::getBufferType(
195 value: selectOp.getFalseValue(), options, state, invocationStack));
196 if (failed(Result: trueType) || failed(Result: falseType))
197 return failure();
198 if (*trueType == *falseType)
199 return cast<BufferLikeType>(Val&: *trueType);
200 if (trueType->getMemorySpace() != falseType->getMemorySpace())
201 return op->emitError(message: "inconsistent memory space on true/false operands");
202
203 // If the buffers have different types, they differ only in their layout
204 // map.
205 auto memrefType = llvm::cast<MemRefType>(Val&: *trueType);
206 return cast<BufferLikeType>(Val: getMemRefTypeWithFullyDynamicLayout(
207 tensorType: RankedTensorType::get(shape: memrefType.getShape(),
208 elementType: memrefType.getElementType()),
209 memorySpace: memrefType.getMemorySpace()));
210 }
211};
212
213} // namespace
214
215void mlir::arith::registerBufferizableOpInterfaceExternalModels(
216 DialectRegistry &registry) {
217 registry.addExtension(extensionFn: +[](MLIRContext *ctx, ArithDialect *dialect) {
218 ConstantOp::attachInterface<ConstantOpInterface>(context&: *ctx);
219 IndexCastOp::attachInterface<IndexCastOpInterface>(context&: *ctx);
220 SelectOp::attachInterface<SelectOpInterface>(context&: *ctx);
221 });
222}
223

source code of mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp