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 "../PassDetail.h"
15 #include "mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.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
34   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 
43     // Transform a clone operation into alloc + copy operation and pay
44     // attention to the shape dimensions.
45     MemRefType memrefType = type.cast<MemRefType>();
46     Location loc = op->getLoc();
47     SmallVector<Value, 4> dynamicOperands;
48     for (int i = 0; i < memrefType.getRank(); ++i) {
49       if (!memrefType.isDynamicDim(i))
50         continue;
51       Value size = rewriter.createOrFold<arith::ConstantIndexOp>(loc, i);
52       Value dim = rewriter.createOrFold<memref::DimOp>(loc, op.input(), size);
53       dynamicOperands.push_back(dim);
54     }
55     Value alloc = rewriter.replaceOpWithNewOp<memref::AllocOp>(op, memrefType,
56                                                                dynamicOperands);
57     rewriter.create<memref::CopyOp>(loc, op.input(), alloc);
58     return success();
59   }
60 };
61 } // namespace
62 
63 void mlir::populateBufferizationToMemRefConversionPatterns(
64     RewritePatternSet &patterns) {
65   patterns.add<CloneOpConversion>(patterns.getContext());
66 }
67 
68 namespace {
69 struct BufferizationToMemRefPass
70     : public ConvertBufferizationToMemRefBase<BufferizationToMemRefPass> {
71   BufferizationToMemRefPass() = default;
72 
73   void runOnOperation() override {
74     RewritePatternSet patterns(&getContext());
75     populateBufferizationToMemRefConversionPatterns(patterns);
76 
77     ConversionTarget target(getContext());
78     target.addLegalDialect<memref::MemRefDialect>();
79     target.addLegalOp<arith::ConstantOp>();
80     target.addIllegalDialect<bufferization::BufferizationDialect>();
81 
82     if (failed(applyPartialConversion(getOperation(), target,
83                                       std::move(patterns))))
84       signalPassFailure();
85   }
86 };
87 } // namespace
88 
89 std::unique_ptr<Pass> mlir::createBufferizationToMemRefPass() {
90   return std::make_unique<BufferizationToMemRefPass>();
91 }
92