| //===- BufferizationToMemRef.cpp - Bufferization to MemRef conversion -----===// |
| // |
| // 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 patterns to convert Bufferization dialect to MemRef |
| // dialect. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "../PassDetail.h" |
| #include "mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h" |
| #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" |
| #include "mlir/Dialect/Bufferization/IR/Bufferization.h" |
| #include "mlir/Dialect/MemRef/IR/MemRef.h" |
| #include "mlir/IR/BuiltinTypes.h" |
| #include "mlir/Support/LogicalResult.h" |
| #include "mlir/Transforms/DialectConversion.h" |
| |
| using namespace mlir; |
| |
| namespace { |
| /// The CloneOpConversion transforms all bufferization clone operations into |
| /// memref alloc and memref copy operations. In the dynamic-shape case, it also |
| /// emits additional dim and constant operations to determine the shape. This |
| /// conversion does not resolve memory leaks if it is used alone. |
| struct CloneOpConversion : public OpConversionPattern<bufferization::CloneOp> { |
| using OpConversionPattern<bufferization::CloneOp>::OpConversionPattern; |
| |
| LogicalResult |
| matchAndRewrite(bufferization::CloneOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| // Check for unranked memref types which are currently not supported. |
| Type type = op.getType(); |
| if (type.isa<UnrankedMemRefType>()) { |
| return rewriter.notifyMatchFailure( |
| op, "UnrankedMemRefType is not supported."); |
| } |
| |
| // Transform a clone operation into alloc + copy operation and pay |
| // attention to the shape dimensions. |
| MemRefType memrefType = type.cast<MemRefType>(); |
| Location loc = op->getLoc(); |
| SmallVector<Value, 4> dynamicOperands; |
| for (int i = 0; i < memrefType.getRank(); ++i) { |
| if (!memrefType.isDynamicDim(i)) |
| continue; |
| Value size = rewriter.createOrFold<arith::ConstantIndexOp>(loc, i); |
| Value dim = rewriter.createOrFold<memref::DimOp>(loc, op.input(), size); |
| dynamicOperands.push_back(dim); |
| } |
| Value alloc = rewriter.replaceOpWithNewOp<memref::AllocOp>(op, memrefType, |
| dynamicOperands); |
| rewriter.create<memref::CopyOp>(loc, op.input(), alloc); |
| return success(); |
| } |
| }; |
| } // namespace |
| |
| void mlir::populateBufferizationToMemRefConversionPatterns( |
| RewritePatternSet &patterns) { |
| patterns.add<CloneOpConversion>(patterns.getContext()); |
| } |
| |
| namespace { |
| struct BufferizationToMemRefPass |
| : public ConvertBufferizationToMemRefBase<BufferizationToMemRefPass> { |
| BufferizationToMemRefPass() = default; |
| |
| void runOnOperation() override { |
| RewritePatternSet patterns(&getContext()); |
| populateBufferizationToMemRefConversionPatterns(patterns); |
| |
| ConversionTarget target(getContext()); |
| target.addLegalDialect<memref::MemRefDialect>(); |
| target.addLegalOp<arith::ConstantOp>(); |
| target.addIllegalDialect<bufferization::BufferizationDialect>(); |
| |
| if (failed(applyPartialConversion(getOperation(), target, |
| std::move(patterns)))) |
| signalPassFailure(); |
| } |
| }; |
| } // namespace |
| |
| std::unique_ptr<Pass> mlir::createBufferizationToMemRefPass() { |
| return std::make_unique<BufferizationToMemRefPass>(); |
| } |