1//===- BufferizationToMemRef.cpp - Bufferization to MemRef conversion -----===//
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// This file implements patterns to convert Bufferization dialect to MemRef
10// dialect.
11//
12//===----------------------------------------------------------------------===//
13
14#include "mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h"
15
16#include "mlir/Dialect/Arith/IR/Arith.h"
17#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
18#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
19#include "mlir/Dialect/Func/IR/FuncOps.h"
20#include "mlir/Dialect/MemRef/IR/MemRef.h"
21#include "mlir/Dialect/SCF/IR/SCF.h"
22#include "mlir/IR/BuiltinTypes.h"
23#include "mlir/Pass/Pass.h"
24#include "mlir/Transforms/DialectConversion.h"
25
26namespace mlir {
27#define GEN_PASS_DEF_CONVERTBUFFERIZATIONTOMEMREFPASS
28#include "mlir/Conversion/Passes.h.inc"
29} // namespace mlir
30
31using namespace mlir;
32
33namespace {
34/// The CloneOpConversion transforms all bufferization clone operations into
35/// memref alloc and memref copy operations. In the dynamic-shape case, it also
36/// emits additional dim and constant operations to determine the shape. This
37/// conversion does not resolve memory leaks if it is used alone.
38struct CloneOpConversion : public OpConversionPattern<bufferization::CloneOp> {
39 using OpConversionPattern<bufferization::CloneOp>::OpConversionPattern;
40
41 LogicalResult
42 matchAndRewrite(bufferization::CloneOp op, OpAdaptor adaptor,
43 ConversionPatternRewriter &rewriter) const override {
44 Location loc = op->getLoc();
45
46 Type type = op.getType();
47 Value alloc;
48
49 if (auto unrankedType = dyn_cast<UnrankedMemRefType>(type)) {
50 // Constants
51 Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
52 Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
53
54 // Dynamically evaluate the size and shape of the unranked memref
55 Value rank = rewriter.create<memref::RankOp>(loc, op.getInput());
56 MemRefType allocType =
57 MemRefType::get({ShapedType::kDynamic}, rewriter.getIndexType());
58 Value shape = rewriter.create<memref::AllocaOp>(loc, allocType, rank);
59
60 // Create a loop to query dimension sizes, store them as a shape, and
61 // compute the total size of the memref
62 auto loopBody = [&](OpBuilder &builder, Location loc, Value i,
63 ValueRange args) {
64 auto acc = args.front();
65 auto dim = rewriter.create<memref::DimOp>(loc, op.getInput(), i);
66
67 rewriter.create<memref::StoreOp>(loc, dim, shape, i);
68 acc = rewriter.create<arith::MulIOp>(loc, acc, dim);
69
70 rewriter.create<scf::YieldOp>(loc, acc);
71 };
72 auto size = rewriter
73 .create<scf::ForOp>(loc, zero, rank, one, ValueRange(one),
74 loopBody)
75 .getResult(0);
76
77 MemRefType memrefType = MemRefType::get({ShapedType::kDynamic},
78 unrankedType.getElementType());
79
80 // Allocate new memref with 1D dynamic shape, then reshape into the
81 // shape of the original unranked memref
82 alloc = rewriter.create<memref::AllocOp>(loc, memrefType, size);
83 alloc =
84 rewriter.create<memref::ReshapeOp>(loc, unrankedType, alloc, shape);
85 } else {
86 MemRefType memrefType = cast<MemRefType>(type);
87 MemRefLayoutAttrInterface layout;
88 auto allocType =
89 MemRefType::get(memrefType.getShape(), memrefType.getElementType(),
90 layout, memrefType.getMemorySpace());
91 // Since this implementation always allocates, certain result types of
92 // the clone op cannot be lowered.
93 if (!memref::CastOp::areCastCompatible({allocType}, {memrefType}))
94 return failure();
95
96 // Transform a clone operation into alloc + copy operation and pay
97 // attention to the shape dimensions.
98 SmallVector<Value, 4> dynamicOperands;
99 for (int i = 0; i < memrefType.getRank(); ++i) {
100 if (!memrefType.isDynamicDim(i))
101 continue;
102 Value dim = rewriter.createOrFold<memref::DimOp>(loc, op.getInput(), i);
103 dynamicOperands.push_back(Elt: dim);
104 }
105
106 // Allocate a memref with identity layout.
107 alloc = rewriter.create<memref::AllocOp>(loc, allocType, dynamicOperands);
108 // Cast the allocation to the specified type if needed.
109 if (memrefType != allocType)
110 alloc =
111 rewriter.create<memref::CastOp>(op->getLoc(), memrefType, alloc);
112 }
113
114 rewriter.replaceOp(op, alloc);
115 rewriter.create<memref::CopyOp>(loc, op.getInput(), alloc);
116 return success();
117 }
118};
119
120} // namespace
121
122namespace {
123struct BufferizationToMemRefPass
124 : public impl::ConvertBufferizationToMemRefPassBase<
125 BufferizationToMemRefPass> {
126 BufferizationToMemRefPass() = default;
127
128 void runOnOperation() override {
129 if (!isa<ModuleOp, FunctionOpInterface>(getOperation())) {
130 emitError(getOperation()->getLoc(),
131 "root operation must be a builtin.module or a function");
132 signalPassFailure();
133 return;
134 }
135
136 bufferization::DeallocHelperMap deallocHelperFuncMap;
137 if (auto module = dyn_cast<ModuleOp>(getOperation())) {
138 OpBuilder builder = OpBuilder::atBlockBegin(block: module.getBody());
139
140 // Build dealloc helper function if there are deallocs.
141 getOperation()->walk([&](bufferization::DeallocOp deallocOp) {
142 Operation *symtableOp =
143 deallocOp->getParentWithTrait<OpTrait::SymbolTable>();
144 if (deallocOp.getMemrefs().size() > 1 &&
145 !deallocHelperFuncMap.contains(Val: symtableOp)) {
146 SymbolTable symbolTable(symtableOp);
147 func::FuncOp helperFuncOp =
148 bufferization::buildDeallocationLibraryFunction(
149 builder, getOperation()->getLoc(), symbolTable);
150 deallocHelperFuncMap[symtableOp] = helperFuncOp;
151 }
152 });
153 }
154
155 RewritePatternSet patterns(&getContext());
156 patterns.add<CloneOpConversion>(arg: patterns.getContext());
157 bufferization::populateBufferizationDeallocLoweringPattern(
158 patterns, deallocHelperFuncMap);
159
160 ConversionTarget target(getContext());
161 target.addLegalDialect<memref::MemRefDialect, arith::ArithDialect,
162 scf::SCFDialect, func::FuncDialect>();
163 target.addIllegalDialect<bufferization::BufferizationDialect>();
164
165 if (failed(applyPartialConversion(getOperation(), target,
166 std::move(patterns))))
167 signalPassFailure();
168 }
169};
170} // namespace
171

source code of mlir/lib/Conversion/BufferizationToMemRef/BufferizationToMemRef.cpp