[mlir][linalg][bufferize] Limited support for scf.execute_region

Add support for analysis only.

Differential Revision: https://reviews.llvm.org/D114055
This commit is contained in:
Matthias Springer 2021-11-23 12:19:53 +09:00
parent 26c0dd83ab
commit fb99686bfd
2 changed files with 43 additions and 1 deletions

View File

@ -1516,6 +1516,46 @@ struct ConstantOpInterface
namespace scf_ext {
struct ExecuteRegionOpInterface
: public BufferizableOpInterface::ExternalModel<ExecuteRegionOpInterface,
scf::ExecuteRegionOp> {
SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
OpResult opResult) const {
// ExecuteRegionOps 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 ExecuteRegionOps in the analysis, the corresponding yield value
// is considered to be aliasing with the result.
auto executeRegionOp = cast<scf::ExecuteRegionOp>(op);
size_t resultNum = std::distance(op->getOpResults().begin(),
llvm::find(op->getOpResults(), opResult));
assert(executeRegionOp.region().getBlocks().size() == 1 &&
"expected exactly 1 block");
auto yieldOp = dyn_cast<scf::YieldOp>(
executeRegionOp.region().front().getTerminator());
assert(yieldOp && "expected scf.yield terminator in scf.execute_region");
return {&yieldOp->getOpOperand(resultNum)};
}
bool mustBufferizeInPlace(Operation *op, OpResult opResult) const {
// ExecuteRegionOp results always bufferize in-place. Since they have no
// OpOperands, they are mostly ignored by the analysis once alias sets are
// set up.
return true;
}
LogicalResult bufferize(Operation *op, OpBuilder &b,
BufferizationState &state) const {
// TODO: Add bufferization support when needed. scf.execute_region should be
// bufferized similar to scf.if.
bool hasTensorReturnType = any_of(
op->getResultTypes(), [](Type t) { return t.isa<TensorType>(); });
if (hasTensorReturnType)
return op->emitError(
"scf.execute_region with tensor result not supported");
return success();
}
};
struct IfOpInterface
: public BufferizableOpInterface::ExternalModel<IfOpInterface, scf::IfOp> {
SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
@ -2490,6 +2530,8 @@ struct TransferWriteOpInterface
void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
registry.addOpInterface<arith::ConstantOp, arith_ext::ConstantOpInterface>();
registry.addOpInterface<scf::ExecuteRegionOp,
scf_ext::ExecuteRegionOpInterface>();
registry.addOpInterface<scf::ForOp, scf_ext::ForOpInterface>();
registry.addOpInterface<scf::IfOp, scf_ext::IfOpInterface>();
registry.addOpInterface<scf::YieldOp, scf_ext::YieldOpInterface>();

View File

@ -160,7 +160,7 @@ func @mini_test_case1() -> tensor<10x20xf32> {
// -----
func @main() -> tensor<4xi32> {
// expected-error @+1 {{unsupported op with tensors}}
// expected-error @+1 {{scf.execute_region with tensor result not supported}}
%r = scf.execute_region -> tensor<4xi32> {
%A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
scf.yield %A: tensor<4xi32>