172 lines · cpp
1//===- ExpandRealloc.cpp - Expand memref.realloc ops into it's components -===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/MemRef/Transforms/Passes.h"10#include "mlir/Dialect/MemRef/Transforms/Transforms.h"11 12#include "mlir/Dialect/Arith/IR/Arith.h"13#include "mlir/Dialect/MemRef/IR/MemRef.h"14#include "mlir/Dialect/SCF/IR/SCF.h"15#include "mlir/Transforms/DialectConversion.h"16 17namespace mlir {18namespace memref {19#define GEN_PASS_DEF_EXPANDREALLOCPASS20#include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"21} // namespace memref22} // namespace mlir23 24using namespace mlir;25 26namespace {27 28/// The `realloc` operation performs a conditional allocation and copy to29/// increase the size of a buffer if necessary. This pattern converts the30/// `realloc` operation into this sequence of simpler operations.31 32/// Example of an expansion:33/// ```mlir34/// %realloc = memref.realloc %alloc (%size) : memref<?xf32> to memref<?xf32>35/// ```36/// is expanded to37/// ```mlir38/// %c0 = arith.constant 0 : index39/// %dim = memref.dim %alloc, %c0 : memref<?xf32>40/// %is_old_smaller = arith.cmpi ult, %dim, %arg141/// %realloc = scf.if %is_old_smaller -> (memref<?xf32>) {42/// %new_alloc = memref.alloc(%size) : memref<?xf32>43/// %subview = memref.subview %new_alloc[0] [%dim] [1]44/// memref.copy %alloc, %subview45/// memref.dealloc %alloc46/// scf.yield %alloc_0 : memref<?xf32>47/// } else {48/// %reinterpret_cast = memref.reinterpret_cast %alloc to49/// offset: [0], sizes: [%size], strides: [1]50/// scf.yield %reinterpret_cast : memref<?xf32>51/// }52/// ```53struct ExpandReallocOpPattern : public OpRewritePattern<memref::ReallocOp> {54 ExpandReallocOpPattern(MLIRContext *ctx, bool emitDeallocs)55 : OpRewritePattern(ctx), emitDeallocs(emitDeallocs) {}56 57 LogicalResult matchAndRewrite(memref::ReallocOp op,58 PatternRewriter &rewriter) const final {59 Location loc = op.getLoc();60 assert(op.getType().getRank() == 1 &&61 "result MemRef must have exactly one rank");62 assert(op.getSource().getType().getRank() == 1 &&63 "source MemRef must have exactly one rank");64 assert(op.getType().getLayout().isIdentity() &&65 "result MemRef must have identity layout (or none)");66 assert(op.getSource().getType().getLayout().isIdentity() &&67 "source MemRef must have identity layout (or none)");68 69 // Get the size of the original buffer.70 int64_t inputSize =71 cast<BaseMemRefType>(op.getSource().getType()).getDimSize(0);72 OpFoldResult currSize = rewriter.getIndexAttr(inputSize);73 if (ShapedType::isDynamic(inputSize)) {74 Value dimZero = getValueOrCreateConstantIndexOp(rewriter, loc,75 rewriter.getIndexAttr(0));76 currSize = memref::DimOp::create(rewriter, loc, op.getSource(), dimZero)77 .getResult();78 }79 80 // Get the requested size that the new buffer should have.81 int64_t outputSize =82 cast<BaseMemRefType>(op.getResult().getType()).getDimSize(0);83 OpFoldResult targetSize = ShapedType::isDynamic(outputSize)84 ? OpFoldResult{op.getDynamicResultSize()}85 : rewriter.getIndexAttr(outputSize);86 87 // Only allocate a new buffer and copy over the values in the old buffer if88 // the old buffer is smaller than the requested size.89 Value lhs = getValueOrCreateConstantIndexOp(rewriter, loc, currSize);90 Value rhs = getValueOrCreateConstantIndexOp(rewriter, loc, targetSize);91 Value cond = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::ult,92 lhs, rhs);93 auto ifOp = scf::IfOp::create(94 rewriter, loc, cond,95 [&](OpBuilder &builder, Location loc) {96 // Allocate the new buffer. If it is a dynamic memref we need to pass97 // an additional operand for the size at runtime, otherwise the static98 // size is encoded in the result type.99 SmallVector<Value> dynamicSizeOperands;100 if (op.getDynamicResultSize())101 dynamicSizeOperands.push_back(op.getDynamicResultSize());102 103 Value newAlloc = memref::AllocOp::create(104 builder, loc, op.getResult().getType(), dynamicSizeOperands,105 op.getAlignmentAttr());106 107 // Take a subview of the new (bigger) buffer such that we can copy the108 // old values over (the copy operation requires both operands to have109 // the same shape).110 Value subview = memref::SubViewOp::create(111 builder, loc, newAlloc,112 ArrayRef<OpFoldResult>{rewriter.getIndexAttr(0)},113 ArrayRef<OpFoldResult>{currSize},114 ArrayRef<OpFoldResult>{rewriter.getIndexAttr(1)});115 memref::CopyOp::create(builder, loc, op.getSource(), subview);116 117 // Insert the deallocation of the old buffer only if requested118 // (enabled by default).119 if (emitDeallocs)120 memref::DeallocOp::create(builder, loc, op.getSource());121 122 scf::YieldOp::create(builder, loc, newAlloc);123 },124 [&](OpBuilder &builder, Location loc) {125 // We need to reinterpret-cast here because either the input or output126 // type might be static, which means we need to cast from static to127 // dynamic or vice-versa. If both are static and the original buffer128 // is already bigger than the requested size, the cast represents a129 // subview operation.130 Value casted = memref::ReinterpretCastOp::create(131 builder, loc, cast<MemRefType>(op.getResult().getType()),132 op.getSource(), rewriter.getIndexAttr(0),133 ArrayRef<OpFoldResult>{targetSize},134 ArrayRef<OpFoldResult>{rewriter.getIndexAttr(1)});135 scf::YieldOp::create(builder, loc, casted);136 });137 138 rewriter.replaceOp(op, ifOp.getResult(0));139 return success();140 }141 142private:143 const bool emitDeallocs;144};145 146struct ExpandReallocPass147 : public memref::impl::ExpandReallocPassBase<ExpandReallocPass> {148 using Base::Base;149 150 void runOnOperation() override {151 MLIRContext &ctx = getContext();152 153 RewritePatternSet patterns(&ctx);154 memref::populateExpandReallocPatterns(patterns, emitDeallocs.getValue());155 ConversionTarget target(ctx);156 157 target.addLegalDialect<arith::ArithDialect, scf::SCFDialect,158 memref::MemRefDialect>();159 target.addIllegalOp<memref::ReallocOp>();160 if (failed(applyPartialConversion(getOperation(), target,161 std::move(patterns))))162 signalPassFailure();163 }164};165 166} // namespace167 168void mlir::memref::populateExpandReallocPatterns(RewritePatternSet &patterns,169 bool emitDeallocs) {170 patterns.add<ExpandReallocOpPattern>(patterns.getContext(), emitDeallocs);171}172