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
|
//===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
//
// 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/Shape/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Shape/IR/Shape.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/PatternMatch.h"
using namespace mlir;
using namespace mlir::bufferization;
using namespace mlir::shape;
namespace mlir {
namespace shape {
namespace {
/// Bufferization of shape.assuming.
struct AssumingOpInterface
: public BufferizableOpInterface::ExternalModel<AssumingOpInterface,
shape::AssumingOp> {
AliasingOpOperandList
getAliasingOpOperands(Operation *op, OpResult opResult,
const AnalysisState &state) const {
// AssumingOps do not have tensor OpOperands. The yielded value can be any
// SSA value that is in scope. To allow for use-def chain traversal through
// AssumingOps in the analysis, the corresponding yield value is considered
// to be aliasing with the result.
auto assumingOp = cast<shape::AssumingOp>(op);
size_t resultNum = std::distance(op->getOpResults().begin(),
llvm::find(op->getOpResults(), opResult));
// TODO: Support multiple blocks.
assert(assumingOp.getDoRegion().getBlocks().size() == 1 &&
"expected exactly 1 block");
auto yieldOp = dyn_cast<shape::AssumingYieldOp>(
assumingOp.getDoRegion().front().getTerminator());
assert(yieldOp && "expected shape.assuming_yield terminator");
return {{&yieldOp->getOpOperand(resultNum), BufferRelation::Equivalent}};
}
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
const BufferizationOptions &options) const {
auto assumingOp = cast<shape::AssumingOp>(op);
assert(assumingOp.getDoRegion().getBlocks().size() == 1 &&
"only 1 block supported");
auto yieldOp = cast<shape::AssumingYieldOp>(
assumingOp.getDoRegion().front().getTerminator());
// Create new op and move over region.
TypeRange newResultTypes(yieldOp.getOperands());
auto newOp = rewriter.create<shape::AssumingOp>(
op->getLoc(), newResultTypes, assumingOp.getWitness());
newOp.getDoRegion().takeBody(assumingOp.getRegion());
// Update all uses of the old op.
rewriter.setInsertionPointAfter(newOp);
SmallVector<Value> newResults;
for (const auto &it : llvm::enumerate(assumingOp->getResultTypes())) {
if (isa<TensorType>(it.value())) {
newResults.push_back(rewriter.create<bufferization::ToTensorOp>(
assumingOp.getLoc(), newOp->getResult(it.index())));
} else {
newResults.push_back(newOp->getResult(it.index()));
}
}
// Replace old op.
rewriter.replaceOp(assumingOp, newResults);
return success();
}
};
/// Bufferization of shape.assuming_yield. Bufferized as part of their enclosing
/// ops, so this is for analysis only.
struct AssumingYieldOpInterface
: public BufferizableOpInterface::ExternalModel<AssumingYieldOpInterface,
shape::AssumingYieldOp> {
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
return true;
}
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
return false;
}
AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
assert(isa<shape::AssumingOp>(op->getParentOp()) &&
"expected that parent is an AssumingOp");
OpResult opResult =
op->getParentOp()->getResult(opOperand.getOperandNumber());
return {{opResult, BufferRelation::Equivalent}};
}
bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
// Yield operands always bufferize inplace. Otherwise, an alloc + copy
// may be generated inside the block. We should not return/yield allocations
// when possible.
return true;
}
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
const BufferizationOptions &options) const {
auto yieldOp = cast<shape::AssumingYieldOp>(op);
SmallVector<Value> newResults;
for (Value value : yieldOp.getOperands()) {
if (isa<TensorType>(value.getType())) {
FailureOr<Value> buffer = getBuffer(rewriter, value, options);
if (failed(buffer))
return failure();
newResults.push_back(*buffer);
} else {
newResults.push_back(value);
}
}
replaceOpWithNewBufferizedOp<shape::AssumingYieldOp>(rewriter, op,
newResults);
return success();
}
};
} // namespace
} // namespace shape
} // namespace mlir
void mlir::shape::registerBufferizableOpInterfaceExternalModels(
DialectRegistry ®istry) {
registry.addExtension(+[](MLIRContext *ctx, shape::ShapeDialect *dialect) {
shape::AssumingOp::attachInterface<AssumingOpInterface>(*ctx);
shape::AssumingYieldOp::attachInterface<AssumingYieldOpInterface>(*ctx);
});
}
|
