The only benefit of FunctionPass is that it filters out function declarations. This isn't enough to justify carrying it around, as we can simplify filter out declarations when necessary within the pass. We can also explore with better scheduling primitives to filter out declarations at the pipeline level in the future. The definition of FunctionPass is left intact for now to allow time for downstream users to migrate. Differential Revision: https://reviews.llvm.org/D117182
79 lines
2.9 KiB
C++
79 lines
2.9 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/Dialect/Bufferization/Transforms/Bufferize.h"
|
|
#include "PassDetail.h"
|
|
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
|
|
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
|
#include "mlir/Dialect/SCF/SCF.h"
|
|
#include "mlir/Dialect/StandardOps/IR/Ops.h"
|
|
#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/IR/BlockAndValueMapping.h"
|
|
#include "mlir/Transforms/DialectConversion.h"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
class BufferizeSelectOp : public OpConversionPattern<SelectOp> {
|
|
public:
|
|
using OpConversionPattern::OpConversionPattern;
|
|
LogicalResult
|
|
matchAndRewrite(SelectOp op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
if (!op.getCondition().getType().isa<IntegerType>())
|
|
return rewriter.notifyMatchFailure(op, "requires scalar condition");
|
|
|
|
rewriter.replaceOpWithNewOp<SelectOp>(op, adaptor.getCondition(),
|
|
adaptor.getTrueValue(),
|
|
adaptor.getFalseValue());
|
|
return success();
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
void mlir::populateStdBufferizePatterns(
|
|
bufferization::BufferizeTypeConverter &typeConverter,
|
|
RewritePatternSet &patterns) {
|
|
patterns.add<BufferizeSelectOp>(typeConverter, patterns.getContext());
|
|
}
|
|
|
|
namespace {
|
|
struct StdBufferizePass : public StdBufferizeBase<StdBufferizePass> {
|
|
void runOnOperation() override {
|
|
auto *context = &getContext();
|
|
bufferization::BufferizeTypeConverter typeConverter;
|
|
RewritePatternSet patterns(context);
|
|
ConversionTarget target(*context);
|
|
|
|
target.addLegalDialect<scf::SCFDialect, StandardOpsDialect,
|
|
memref::MemRefDialect>();
|
|
|
|
populateStdBufferizePatterns(typeConverter, patterns);
|
|
// We only bufferize the case of tensor selected type and scalar condition,
|
|
// as that boils down to a select over memref descriptors (don't need to
|
|
// touch the data).
|
|
target.addDynamicallyLegalOp<SelectOp>([&](SelectOp op) {
|
|
return typeConverter.isLegal(op.getType()) ||
|
|
!op.getCondition().getType().isa<IntegerType>();
|
|
});
|
|
if (failed(applyPartialConversion(getOperation(), target,
|
|
std::move(patterns))))
|
|
signalPassFailure();
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
std::unique_ptr<Pass> mlir::createStdBufferizePass() {
|
|
return std::make_unique<StdBufferizePass>();
|
|
}
|