a02ad6c177
This revision is needed to support bufferization of `cf.br`/`cf.cond_br`. It will also be useful for better analysis of loop ops. This revision generalizes `getAliasingOpResults` to `getAliasingValues`. An OpOperand can now not only alias with OpResults but also with BlockArguments. In the case of `cf.br` (will be added in a later revision): a `cf.br` operand will alias with the corresponding argument of the destination block. If an op does not implement the `BufferizableOpInterface`, the analysis in conservative. It previously assumed that an OpOperand may alias with each OpResult. It now assumes that an OpOperand may alias with each OpResult and each BlockArgument of the entry block. Differential Revision: https://reviews.llvm.org/D157957
64 lines
2.4 KiB
C++
64 lines
2.4 KiB
C++
//===- 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/MemRef/Transforms/BufferizableOpInterfaceImpl.h"
|
|
|
|
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
|
|
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
|
#include "mlir/IR/Dialect.h"
|
|
#include "mlir/IR/Operation.h"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::bufferization;
|
|
|
|
namespace {
|
|
/// Bufferization of memref.tensor_store. Replace with memref.copy.
|
|
struct TensorStoreOpInterface
|
|
: public BufferizableOpInterface::ExternalModel<TensorStoreOpInterface,
|
|
memref::TensorStoreOp> {
|
|
AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
|
|
const AnalysisState &state) const {
|
|
return {};
|
|
}
|
|
|
|
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
|
|
const AnalysisState &state) const {
|
|
assert(opOperand.getOperandNumber() == 0 && "expected src operand");
|
|
return true;
|
|
}
|
|
|
|
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
|
|
const AnalysisState &state) const {
|
|
// The memref operand is written but not the tensor operand.
|
|
assert(opOperand.getOperandNumber() == 0 && "expected src operand");
|
|
return false;
|
|
}
|
|
|
|
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
|
|
const BufferizationOptions &options) const {
|
|
auto tensorStoreOp = cast<memref::TensorStoreOp>(op);
|
|
auto srcBuffer = getBuffer(rewriter, tensorStoreOp.getTensor(), options);
|
|
if (failed(srcBuffer))
|
|
return failure();
|
|
if (failed(options.createMemCpy(rewriter, op->getLoc(), *srcBuffer,
|
|
tensorStoreOp.getMemref())))
|
|
return failure();
|
|
rewriter.eraseOp(tensorStoreOp);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
void mlir::memref::registerBufferizableOpInterfaceExternalModels(
|
|
DialectRegistry ®istry) {
|
|
registry.addExtension(+[](MLIRContext *ctx, MemRefDialect *dialect) {
|
|
TensorStoreOp::attachInterface<TensorStoreOpInterface>(*ctx);
|
|
});
|
|
}
|