ee491ac91e
The opposite of tensor_to_memref is tensor_load. - Add some basic tensor_load/tensor_to_memref folding. - Add source/target materializations to BufferizeTypeConverter. - Add an example std bufferization pattern/pass that shows how the materialiations work together (more std bufferization patterns to come in subsequent commits). - In coming commits, I'll document how to write composable bufferization passes/patterns and update the other in-tree bufferization passes to match this convention. The populate* functions will of course continue to be exposed for power users. The naming on tensor_load/tensor_to_memref and their pretty forms are not very intuitive. I'm open to any suggestions here. One key observation is that the memref type must always be the one specified in the pretty form, since the tensor type can be inferred from the memref type but not vice-versa. With this, I've been able to replace all my custom bufferization type converters in npcomp with BufferizeTypeConverter! Part of the plan discussed in: https://llvm.discourse.group/t/what-is-the-strategy-for-tensor-memref-conversion-bufferization/1938/17 Differential Revision: https://reviews.llvm.org/D89437
63 lines
2.1 KiB
C++
63 lines
2.1 KiB
C++
//===- Bufferize.cpp - Bufferization for std ops --------------------------===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file implements bufferization of std ops.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Transforms/Bufferize.h"
|
|
#include "PassDetail.h"
|
|
#include "mlir/Dialect/StandardOps/IR/Ops.h"
|
|
#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
|
|
#include "mlir/Transforms/DialectConversion.h"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
class BufferizeTensorCastOp : public OpConversionPattern<TensorCastOp> {
|
|
public:
|
|
using OpConversionPattern::OpConversionPattern;
|
|
LogicalResult
|
|
matchAndRewrite(TensorCastOp op, ArrayRef<Value> operands,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
auto resultType = getTypeConverter()->convertType(op.getType());
|
|
rewriter.replaceOpWithNewOp<MemRefCastOp>(op, resultType, operands[0]);
|
|
return success();
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
void mlir::populateStdBufferizePatterns(MLIRContext *context,
|
|
BufferizeTypeConverter &typeConverter,
|
|
OwningRewritePatternList &patterns) {
|
|
patterns.insert<BufferizeTensorCastOp>(typeConverter, context);
|
|
}
|
|
|
|
namespace {
|
|
struct StdBufferizePass : public StdBufferizeBase<StdBufferizePass> {
|
|
void runOnFunction() override {
|
|
auto *context = &getContext();
|
|
BufferizeTypeConverter typeConverter;
|
|
OwningRewritePatternList patterns;
|
|
ConversionTarget target(*context);
|
|
|
|
target.addLegalDialect<StandardOpsDialect>();
|
|
|
|
populateStdBufferizePatterns(context, typeConverter, patterns);
|
|
target.addIllegalOp<TensorCastOp>();
|
|
|
|
if (failed(mlir::applyPartialConversion(getFunction(), target, patterns)))
|
|
signalPassFailure();
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
std::unique_ptr<Pass> mlir::createStdBufferizePass() {
|
|
return std::make_unique<StdBufferizePass>();
|
|
}
|