shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

[mlir][tensor] Implement getBufferType for ReshapeOp.

Browse Source 
This function should be implemented for ops that work in one-shot
bufferization.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D151548
This commit is contained in:
Ingo Müller 2023-05-26 09:43:16 +00:00
parent 4b1eb4cf0e
commit 9dbb8eefd4
1 changed files with 19 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
View File

@ -992,13 +992,28 @@ struct ReshapeOpInterface
getBuffer(rewriter, reshapeOp.getShape(), options);
if (failed(srcBuffer) || failed(shapeBuffer))
return failure();
auto resultMemRefType = getMemRefTypeWithStaticIdentityLayout(
reshapeOp.getResult().getType(),
cast<BaseMemRefType>(srcBuffer->getType()).getMemorySpace());
auto maybeResultMemRefType =
bufferization::getBufferType(reshapeOp.getResult(), options);
if (failed(maybeResultMemRefType))
return failure();
replaceOpWithNewBufferizedOp<memref::ReshapeOp>(
rewriter, op, resultMemRefType, *srcBuffer, *shapeBuffer);
rewriter, op, maybeResultMemRefType.value(), *srcBuffer, *shapeBuffer);
return success();
}
FailureOr<BaseMemRefType>
getBufferType(Operation *op, Value value, const BufferizationOptions &options,
const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
auto reshapeOp = cast<tensor::ReshapeOp>(op);
assert(value == reshapeOp.getResult() && "unexpected value provided");
auto maybeSourceBufferType = bufferization::getBufferType(
reshapeOp.getSource(), options, fixedTypes);
if (failed(maybeSourceBufferType))
return failure();
return getMemRefTypeWithStaticIdentityLayout(
reshapeOp.getResult().getType(),
cast<BaseMemRefType>(maybeSourceBufferType.value()).getMemorySpace());
}
};
/// Analysis of ParallelInsertSliceOp.