1 //===- Bufferize.cpp - Bufferization for std ops --------------------------===//
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 bufferization of std.func's and std.call's.
10 //
11 //===----------------------------------------------------------------------===//
12
13 #include "PassDetail.h"
14 #include "mlir/Dialect/StandardOps/IR/Ops.h"
15 #include "mlir/Dialect/StandardOps/Transforms/FuncConversions.h"
16 #include "mlir/Dialect/StandardOps/Transforms/Passes.h"
17 #include "mlir/Transforms/Bufferize.h"
18 #include "mlir/Transforms/DialectConversion.h"
19
20 using namespace mlir;
21
22 namespace {
23 struct FuncBufferizePass : public FuncBufferizeBase<FuncBufferizePass> {
24 using FuncBufferizeBase<FuncBufferizePass>::FuncBufferizeBase;
25
runOnOperation__anon082463c30111::FuncBufferizePass26 void runOnOperation() override {
27 auto module = getOperation();
28 auto *context = &getContext();
29
30 BufferizeTypeConverter typeConverter;
31 OwningRewritePatternList patterns;
32 ConversionTarget target(*context);
33
34 populateFuncOpTypeConversionPattern(patterns, context, typeConverter);
35 target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
36 return typeConverter.isSignatureLegal(op.getType()) &&
37 typeConverter.isLegal(&op.getBody());
38 });
39 populateCallOpTypeConversionPattern(patterns, context, typeConverter);
40 target.addDynamicallyLegalOp<CallOp>(
41 [&](CallOp op) { return typeConverter.isLegal(op); });
42
43 populateBranchOpInterfaceAndReturnOpTypeConversionPattern(patterns, context,
44 typeConverter);
45 target.addLegalOp<ModuleOp, ModuleTerminatorOp, TensorLoadOp,
46 TensorToMemrefOp>();
47 target.addDynamicallyLegalOp<ReturnOp>(
48 [&](ReturnOp op) { return typeConverter.isLegal(op); });
49 // Mark terminators as legal if they have the ReturnLike trait or
50 // implement the BranchOpInterface and have valid types. If they do not
51 // implement the trait or interface, mark them as illegal no matter what.
52 target.markUnknownOpDynamicallyLegal([&](Operation *op) {
53 // If it is not a terminator, ignore it.
54 if (op->isKnownNonTerminator())
55 return true;
56 // If it is not the last operation in the block, also ignore it. We do
57 // this to handle unknown operations, as well.
58 Block *block = op->getBlock();
59 if (!block || &block->back() != op)
60 return true;
61 // ReturnLike operations have to be legalized with their parent. For
62 // return this is handled, for other ops they remain as is.
63 if (op->hasTrait<OpTrait::ReturnLike>())
64 return true;
65 // All successor operands of branch like operations must be rewritten.
66 if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
67 for (int p = 0, e = op->getBlock()->getNumSuccessors(); p < e; ++p) {
68 auto successorOperands = branchOp.getSuccessorOperands(p);
69 if (successorOperands.hasValue() &&
70 !typeConverter.isLegal(successorOperands.getValue().getTypes()))
71 return false;
72 }
73 return true;
74 }
75 return false;
76 });
77
78 if (failed(applyFullConversion(module, target, std::move(patterns))))
79 signalPassFailure();
80 }
81 };
82 } // namespace
83
createFuncBufferizePass()84 std::unique_ptr<Pass> mlir::createFuncBufferizePass() {
85 return std::make_unique<FuncBufferizePass>();
86 }
87