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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
|
//===- BufferResultsToOutParams.cpp - Calling convention conversion -------===//
//
// 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/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Operation.h"
#include "mlir/Pass/Pass.h"
namespace mlir {
namespace bufferization {
#define GEN_PASS_DEF_BUFFERRESULTSTOOUTPARAMS
#include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
} // namespace bufferization
} // namespace mlir
using namespace mlir;
/// Return `true` if the given MemRef type has a fully dynamic layout.
static bool hasFullyDynamicLayoutMap(MemRefType type) {
int64_t offset;
SmallVector<int64_t, 4> strides;
if (failed(getStridesAndOffset(type, strides, offset)))
return false;
if (!llvm::all_of(strides, ShapedType::isDynamicStrideOrOffset))
return false;
if (!ShapedType::isDynamicStrideOrOffset(offset))
return false;
return true;
}
/// Return `true` if the given MemRef type has a static identity layout (i.e.,
/// no layout).
static bool hasStaticIdentityLayout(MemRefType type) {
return type.getLayout().isIdentity();
}
// Updates the func op and entry block.
//
// Any args appended to the entry block are added to `appendedEntryArgs`.
static LogicalResult
updateFuncOp(func::FuncOp func,
SmallVectorImpl<BlockArgument> &appendedEntryArgs) {
auto functionType = func.getFunctionType();
// Collect information about the results will become appended arguments.
SmallVector<Type, 6> erasedResultTypes;
BitVector erasedResultIndices(functionType.getNumResults());
for (const auto &resultType : llvm::enumerate(functionType.getResults())) {
if (auto memrefType = resultType.value().dyn_cast<MemRefType>()) {
if (!hasStaticIdentityLayout(memrefType) &&
!hasFullyDynamicLayoutMap(memrefType)) {
// Only buffers with static identity layout can be allocated. These can
// be casted to memrefs with fully dynamic layout map. Other layout maps
// are not supported.
return func->emitError()
<< "cannot create out param for result with unsupported layout";
}
erasedResultIndices.set(resultType.index());
erasedResultTypes.push_back(memrefType);
}
}
// Add the new arguments to the function type.
auto newArgTypes = llvm::to_vector<6>(
llvm::concat<const Type>(functionType.getInputs(), erasedResultTypes));
auto newFunctionType = FunctionType::get(func.getContext(), newArgTypes,
functionType.getResults());
func.setType(newFunctionType);
// Transfer the result attributes to arg attributes.
auto erasedIndicesIt = erasedResultIndices.set_bits_begin();
for (int i = 0, e = erasedResultTypes.size(); i < e; ++i, ++erasedIndicesIt) {
func.setArgAttrs(functionType.getNumInputs() + i,
func.getResultAttrs(*erasedIndicesIt));
}
// Erase the results.
func.eraseResults(erasedResultIndices);
// Add the new arguments to the entry block if the function is not external.
if (func.isExternal())
return success();
Location loc = func.getLoc();
for (Type type : erasedResultTypes)
appendedEntryArgs.push_back(func.front().addArgument(type, loc));
return success();
}
// Updates all ReturnOps in the scope of the given func::FuncOp by either
// keeping them as return values or copying the associated buffer contents into
// the given out-params.
static void updateReturnOps(func::FuncOp func,
ArrayRef<BlockArgument> appendedEntryArgs) {
func.walk([&](func::ReturnOp op) {
SmallVector<Value, 6> copyIntoOutParams;
SmallVector<Value, 6> keepAsReturnOperands;
for (Value operand : op.getOperands()) {
if (operand.getType().isa<MemRefType>())
copyIntoOutParams.push_back(operand);
else
keepAsReturnOperands.push_back(operand);
}
OpBuilder builder(op);
for (auto t : llvm::zip(copyIntoOutParams, appendedEntryArgs))
builder.create<memref::CopyOp>(op.getLoc(), std::get<0>(t),
std::get<1>(t));
builder.create<func::ReturnOp>(op.getLoc(), keepAsReturnOperands);
op.erase();
});
}
// Updates all CallOps in the scope of the given ModuleOp by allocating
// temporary buffers for newly introduced out params.
static LogicalResult updateCalls(ModuleOp module) {
bool didFail = false;
module.walk([&](func::CallOp op) {
SmallVector<Value, 6> replaceWithNewCallResults;
SmallVector<Value, 6> replaceWithOutParams;
for (OpResult result : op.getResults()) {
if (result.getType().isa<MemRefType>())
replaceWithOutParams.push_back(result);
else
replaceWithNewCallResults.push_back(result);
}
SmallVector<Value, 6> outParams;
OpBuilder builder(op);
for (Value memref : replaceWithOutParams) {
if (!memref.getType().cast<MemRefType>().hasStaticShape()) {
op.emitError()
<< "cannot create out param for dynamically shaped result";
didFail = true;
return;
}
auto memrefType = memref.getType().cast<MemRefType>();
auto allocType =
MemRefType::get(memrefType.getShape(), memrefType.getElementType(),
AffineMap(), memrefType.getMemorySpaceAsInt());
Value outParam = builder.create<memref::AllocOp>(op.getLoc(), allocType);
if (!hasStaticIdentityLayout(memrefType)) {
// Layout maps are already checked in `updateFuncOp`.
assert(hasFullyDynamicLayoutMap(memrefType) &&
"layout map not supported");
outParam =
builder.create<memref::CastOp>(op.getLoc(), memrefType, outParam);
}
memref.replaceAllUsesWith(outParam);
outParams.push_back(outParam);
}
auto newOperands = llvm::to_vector<6>(op.getOperands());
newOperands.append(outParams.begin(), outParams.end());
auto newResultTypes = llvm::to_vector<6>(llvm::map_range(
replaceWithNewCallResults, [](Value v) { return v.getType(); }));
auto newCall = builder.create<func::CallOp>(op.getLoc(), op.getCalleeAttr(),
newResultTypes, newOperands);
for (auto t : llvm::zip(replaceWithNewCallResults, newCall.getResults()))
std::get<0>(t).replaceAllUsesWith(std::get<1>(t));
op.erase();
});
return failure(didFail);
}
LogicalResult
mlir::bufferization::promoteBufferResultsToOutParams(ModuleOp module) {
for (auto func : module.getOps<func::FuncOp>()) {
SmallVector<BlockArgument, 6> appendedEntryArgs;
if (failed(updateFuncOp(func, appendedEntryArgs)))
return failure();
if (func.isExternal())
continue;
updateReturnOps(func, appendedEntryArgs);
}
if (failed(updateCalls(module)))
return failure();
return success();
}
namespace {
struct BufferResultsToOutParamsPass
: bufferization::impl::BufferResultsToOutParamsBase<
BufferResultsToOutParamsPass> {
void runOnOperation() override {
if (failed(bufferization::promoteBufferResultsToOutParams(getOperation())))
return signalPassFailure();
}
};
} // namespace
std::unique_ptr<Pass>
mlir::bufferization::createBufferResultsToOutParamsPass() {
return std::make_unique<BufferResultsToOutParamsPass>();
}
|
