1*075e3fddSMatthias Springer //===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
2*075e3fddSMatthias Springer //
3*075e3fddSMatthias Springer // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4*075e3fddSMatthias Springer // See https://llvm.org/LICENSE.txt for license information.
5*075e3fddSMatthias Springer // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6*075e3fddSMatthias Springer //
7*075e3fddSMatthias Springer //===----------------------------------------------------------------------===//
8*075e3fddSMatthias Springer 
9*075e3fddSMatthias Springer #include "mlir/Dialect/Arithmetic/Transforms/BufferizableOpInterfaceImpl.h"
10*075e3fddSMatthias Springer #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
11*075e3fddSMatthias Springer #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
12*075e3fddSMatthias Springer #include "mlir/Dialect/Bufferization/Transforms/BufferUtils.h"
13*075e3fddSMatthias Springer #include "mlir/Dialect/MemRef/IR/MemRef.h"
14*075e3fddSMatthias Springer #include "mlir/IR/Dialect.h"
15*075e3fddSMatthias Springer #include "mlir/IR/Operation.h"
16*075e3fddSMatthias Springer 
17*075e3fddSMatthias Springer using namespace mlir::bufferization;
18*075e3fddSMatthias Springer 
19*075e3fddSMatthias Springer namespace mlir {
20*075e3fddSMatthias Springer namespace arith {
21*075e3fddSMatthias Springer namespace {
22*075e3fddSMatthias Springer 
23*075e3fddSMatthias Springer /// Bufferization of arith.constant. Replace with memref.get_global.
24*075e3fddSMatthias Springer struct ConstantOpInterface
25*075e3fddSMatthias Springer     : public BufferizableOpInterface::ExternalModel<ConstantOpInterface,
26*075e3fddSMatthias Springer                                                     arith::ConstantOp> {
27*075e3fddSMatthias Springer   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
28*075e3fddSMatthias Springer                           const BufferizationState &state) const {
29*075e3fddSMatthias Springer     auto constantOp = cast<arith::ConstantOp>(op);
30*075e3fddSMatthias Springer 
31*075e3fddSMatthias Springer     // Only ranked tensors are supported.
32*075e3fddSMatthias Springer     if (!constantOp.getType().isa<RankedTensorType>())
33*075e3fddSMatthias Springer       return failure();
34*075e3fddSMatthias Springer 
35*075e3fddSMatthias Springer     // Only constants inside a module are supported.
36*075e3fddSMatthias Springer     auto moduleOp = constantOp->getParentOfType<ModuleOp>();
37*075e3fddSMatthias Springer     if (!moduleOp)
38*075e3fddSMatthias Springer       return failure();
39*075e3fddSMatthias Springer 
40*075e3fddSMatthias Springer     // Create global memory segment and replace tensor with memref pointing to
41*075e3fddSMatthias Springer     // that memory segment.
42*075e3fddSMatthias Springer     GlobalCreator globalCreator(moduleOp);
43*075e3fddSMatthias Springer     auto globalMemref = globalCreator.getGlobalFor(constantOp);
44*075e3fddSMatthias Springer     replaceOpWithNewBufferizedOp<memref::GetGlobalOp>(
45*075e3fddSMatthias Springer         rewriter, op, globalMemref.type(), globalMemref.getName());
46*075e3fddSMatthias Springer 
47*075e3fddSMatthias Springer     return success();
48*075e3fddSMatthias Springer   }
49*075e3fddSMatthias Springer 
50*075e3fddSMatthias Springer   bool isWritable(Operation *op, Value value,
51*075e3fddSMatthias Springer                   const BufferizationState &state) const {
52*075e3fddSMatthias Springer     // Memory locations returned by memref::GetGlobalOp may not be written to.
53*075e3fddSMatthias Springer     assert(value.isa<OpResult>());
54*075e3fddSMatthias Springer     return false;
55*075e3fddSMatthias Springer   }
56*075e3fddSMatthias Springer };
57*075e3fddSMatthias Springer 
58*075e3fddSMatthias Springer struct IndexCastOpInterface
59*075e3fddSMatthias Springer     : public BufferizableOpInterface::ExternalModel<IndexCastOpInterface,
60*075e3fddSMatthias Springer                                                     arith::IndexCastOp> {
61*075e3fddSMatthias Springer   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
62*075e3fddSMatthias Springer                               const BufferizationState &state) const {
63*075e3fddSMatthias Springer     return false;
64*075e3fddSMatthias Springer   }
65*075e3fddSMatthias Springer 
66*075e3fddSMatthias Springer   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
67*075e3fddSMatthias Springer                                const BufferizationState &state) const {
68*075e3fddSMatthias Springer     return false;
69*075e3fddSMatthias Springer   }
70*075e3fddSMatthias Springer 
71*075e3fddSMatthias Springer   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
72*075e3fddSMatthias Springer                                const BufferizationState &state) const {
73*075e3fddSMatthias Springer     return op->getResult(0);
74*075e3fddSMatthias Springer   }
75*075e3fddSMatthias Springer 
76*075e3fddSMatthias Springer   BufferRelation bufferRelation(Operation *op, OpResult opResult,
77*075e3fddSMatthias Springer                                 const BufferizationState &state) const {
78*075e3fddSMatthias Springer     return BufferRelation::Equivalent;
79*075e3fddSMatthias Springer   }
80*075e3fddSMatthias Springer 
81*075e3fddSMatthias Springer   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
82*075e3fddSMatthias Springer                           const BufferizationState &state) const {
83*075e3fddSMatthias Springer     auto castOp = cast<arith::IndexCastOp>(op);
84*075e3fddSMatthias Springer 
85*075e3fddSMatthias Springer     Value source = *state.getBuffer(rewriter, op->getOpOperand(0) /*in*/);
86*075e3fddSMatthias Springer     auto sourceType = source.getType().cast<BaseMemRefType>();
87*075e3fddSMatthias Springer 
88*075e3fddSMatthias Springer     // Result type should have same layout and address space as the source type.
89*075e3fddSMatthias Springer     MemRefLayoutAttrInterface layout = {};
90*075e3fddSMatthias Springer     if (auto rankedMemRefType = sourceType.dyn_cast<MemRefType>())
91*075e3fddSMatthias Springer       layout = rankedMemRefType.getLayout();
92*075e3fddSMatthias Springer     Type resultType =
93*075e3fddSMatthias Springer         getMemRefType(castOp.getType().cast<TensorType>(), state.getOptions(),
94*075e3fddSMatthias Springer                       layout, sourceType.getMemorySpace());
95*075e3fddSMatthias Springer 
96*075e3fddSMatthias Springer     replaceOpWithNewBufferizedOp<arith::IndexCastOp>(rewriter, op, source,
97*075e3fddSMatthias Springer                                                      resultType);
98*075e3fddSMatthias Springer     return success();
99*075e3fddSMatthias Springer   }
100*075e3fddSMatthias Springer };
101*075e3fddSMatthias Springer 
102*075e3fddSMatthias Springer } // namespace
103*075e3fddSMatthias Springer } // namespace arith
104*075e3fddSMatthias Springer } // namespace mlir
105*075e3fddSMatthias Springer 
106*075e3fddSMatthias Springer void mlir::arith::registerBufferizableOpInterfaceExternalModels(
107*075e3fddSMatthias Springer     DialectRegistry &registry) {
108*075e3fddSMatthias Springer   registry.addOpInterface<ConstantOp, ConstantOpInterface>();
109*075e3fddSMatthias Springer   registry.addOpInterface<IndexCastOp, IndexCastOpInterface>();
110*075e3fddSMatthias Springer }
111