/
BufferizationTransformOps.cpp
96 lines (84 loc) · 3.8 KB
/
BufferizationTransformOps.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
//===- BufferizationTransformOps.h - Bufferization transform ops ----------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
#include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/PDL/IR/PDL.h"
#include "mlir/Dialect/PDL/IR/PDLTypes.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
using namespace mlir;
using namespace mlir::bufferization;
using namespace mlir::transform;
//===----------------------------------------------------------------------===//
// OneShotBufferizeOp
//===----------------------------------------------------------------------===//
LogicalResult
transform::OneShotBufferizeOp::apply(TransformResults &transformResults,
TransformState &state) {
OneShotBufferizationOptions options;
options.allowReturnAllocs = getAllowReturnAllocs();
options.allowUnknownOps = getAllowUnknownOps();
options.bufferizeFunctionBoundaries = getBufferizeFunctionBoundaries();
options.createDeallocs = getCreateDeallocs();
options.testAnalysisOnly = getTestAnalysisOnly();
options.printConflicts = getPrintConflicts();
ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
for (Operation *target : payloadOps) {
auto moduleOp = dyn_cast<ModuleOp>(target);
if (getTargetIsModule() && !moduleOp)
return emitError("expected ModuleOp target");
if (options.bufferizeFunctionBoundaries) {
if (!moduleOp)
return emitError("expected ModuleOp target");
if (failed(bufferization::runOneShotModuleBufferize(moduleOp, options)))
return emitError("bufferization failed");
} else {
if (failed(bufferization::runOneShotBufferize(target, options)))
return emitError("bufferization failed");
}
}
return success();
}
void transform::OneShotBufferizeOp::getEffects(
SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
effects.emplace_back(MemoryEffects::Read::get(), getTarget(),
TransformMappingResource::get());
// Handles that are not modules are not longer usable.
if (!getTargetIsModule())
effects.emplace_back(MemoryEffects::Free::get(), getTarget(),
TransformMappingResource::get());
}
//===----------------------------------------------------------------------===//
// Transform op registration
//===----------------------------------------------------------------------===//
namespace {
/// Registers new ops and declares PDL as dependent dialect since the additional
/// ops are using PDL types for operands and results.
class BufferizationTransformDialectExtension
: public transform::TransformDialectExtension<
BufferizationTransformDialectExtension> {
public:
BufferizationTransformDialectExtension() {
declareDependentDialect<bufferization::BufferizationDialect>();
declareDependentDialect<pdl::PDLDialect>();
declareDependentDialect<memref::MemRefDialect>();
registerTransformOps<
#define GET_OP_LIST
#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp.inc"
>();
}
};
} // namespace
#define GET_OP_CLASSES
#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp.inc"
void mlir::bufferization::registerTransformDialectExtension(
DialectRegistry ®istry) {
registry.addExtensions<BufferizationTransformDialectExtension>();
}