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 #include "../PassDetail.h"
16 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
17 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
18 #include "mlir/Dialect/MemRef/IR/MemRef.h"
19 #include "mlir/IR/BuiltinTypes.h"
20 #include "mlir/Support/LogicalResult.h"
21 #include "mlir/Transforms/DialectConversion.h"
22
23 using namespace mlir;
24
25 namespace {
26 /// The CloneOpConversion transforms all bufferization clone operations into
27 /// memref alloc and memref copy operations. In the dynamic-shape case, it also
28 /// emits additional dim and constant operations to determine the shape. This
29 /// conversion does not resolve memory leaks if it is used alone.
30 struct CloneOpConversion : public OpConversionPattern<bufferization::CloneOp> {
31 using OpConversionPattern<bufferization::CloneOp>::OpConversionPattern;
32
33 LogicalResult
matchAndRewrite__anonf46417ab0111::CloneOpConversion34 matchAndRewrite(bufferization::CloneOp op, OpAdaptor adaptor,
35 ConversionPatternRewriter &rewriter) const override {
36 // Check for unranked memref types which are currently not supported.
37 Type type = op.getType();
38 if (type.isa<UnrankedMemRefType>()) {
39 return rewriter.notifyMatchFailure(
40 op, "UnrankedMemRefType is not supported.");
41 }
42 MemRefType memrefType = type.cast<MemRefType>();
43 MemRefLayoutAttrInterface layout;
44 auto allocType =
45 MemRefType::get(memrefType.getShape(), memrefType.getElementType(),
46 layout, memrefType.getMemorySpace());
47 // Since this implementation always allocates, certain result types of the
48 // clone op cannot be lowered.
49 if (!memref::CastOp::areCastCompatible({allocType}, {memrefType}))
50 return failure();
51
52 // Transform a clone operation into alloc + copy operation and pay
53 // attention to the shape dimensions.
54 Location loc = op->getLoc();
55 SmallVector<Value, 4> dynamicOperands;
56 for (int i = 0; i < memrefType.getRank(); ++i) {
57 if (!memrefType.isDynamicDim(i))
58 continue;
59 Value size = rewriter.createOrFold<arith::ConstantIndexOp>(loc, i);
60 Value dim =
61 rewriter.createOrFold<memref::DimOp>(loc, op.getInput(), size);
62 dynamicOperands.push_back(dim);
63 }
64
65 // Allocate a memref with identity layout.
66 Value alloc = rewriter.create<memref::AllocOp>(op->getLoc(), allocType,
67 dynamicOperands);
68 // Cast the allocation to the specified type if needed.
69 if (memrefType != allocType)
70 alloc = rewriter.create<memref::CastOp>(op->getLoc(), memrefType, alloc);
71 rewriter.replaceOp(op, alloc);
72 rewriter.create<memref::CopyOp>(loc, op.getInput(), alloc);
73 return success();
74 }
75 };
76 } // namespace
77
populateBufferizationToMemRefConversionPatterns(RewritePatternSet & patterns)78 void mlir::populateBufferizationToMemRefConversionPatterns(
79 RewritePatternSet &patterns) {
80 patterns.add<CloneOpConversion>(patterns.getContext());
81 }
82
83 namespace {
84 struct BufferizationToMemRefPass
85 : public ConvertBufferizationToMemRefBase<BufferizationToMemRefPass> {
86 BufferizationToMemRefPass() = default;
87
runOnOperation__anonf46417ab0211::BufferizationToMemRefPass88 void runOnOperation() override {
89 RewritePatternSet patterns(&getContext());
90 populateBufferizationToMemRefConversionPatterns(patterns);
91
92 ConversionTarget target(getContext());
93 target.addLegalDialect<memref::MemRefDialect>();
94 target.addLegalOp<arith::ConstantOp>();
95 target.addIllegalDialect<bufferization::BufferizationDialect>();
96
97 if (failed(applyPartialConversion(getOperation(), target,
98 std::move(patterns))))
99 signalPassFailure();
100 }
101 };
102 } // namespace
103
createBufferizationToMemRefPass()104 std::unique_ptr<Pass> mlir::createBufferizationToMemRefPass() {
105 return std::make_unique<BufferizationToMemRefPass>();
106 }
107