MLIR  20.0.0git
ExpandRealloc.cpp
Go to the documentation of this file.
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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
11 
16 
17 namespace mlir {
18 namespace memref {
19 #define GEN_PASS_DEF_EXPANDREALLOC
20 #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
21 } // namespace memref
22 } // namespace mlir
23 
24 using namespace mlir;
25 
26 namespace {
27 
28 /// The `realloc` operation performs a conditional allocation and copy to
29 /// increase the size of a buffer if necessary. This pattern converts the
30 /// `realloc` operation into this sequence of simpler operations.
31 
32 /// Example of an expansion:
33 /// ```mlir
34 /// %realloc = memref.realloc %alloc (%size) : memref<?xf32> to memref<?xf32>
35 /// ```
36 /// is expanded to
37 /// ```mlir
38 /// %c0 = arith.constant 0 : index
39 /// %dim = memref.dim %alloc, %c0 : memref<?xf32>
40 /// %is_old_smaller = arith.cmpi ult, %dim, %arg1
41 /// %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, %subview
45 /// memref.dealloc %alloc
46 /// scf.yield %alloc_0 : memref<?xf32>
47 /// } else {
48 /// %reinterpret_cast = memref.reinterpret_cast %alloc to
49 /// offset: [0], sizes: [%size], strides: [1]
50 /// scf.yield %reinterpret_cast : memref<?xf32>
51 /// }
52 /// ```
53 struct 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 = rewriter.create<memref::DimOp>(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 if
88  // 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 = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
92  lhs, rhs);
93  auto ifOp = rewriter.create<scf::IfOp>(
94  loc, cond,
95  [&](OpBuilder &builder, Location loc) {
96  // Allocate the new buffer. If it is a dynamic memref we need to pass
97  // an additional operand for the size at runtime, otherwise the static
98  // 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 = builder.create<memref::AllocOp>(
104  loc, op.getResult().getType(), dynamicSizeOperands,
105  op.getAlignmentAttr());
106 
107  // Take a subview of the new (bigger) buffer such that we can copy the
108  // old values over (the copy operation requires both operands to have
109  // the same shape).
110  Value subview = builder.create<memref::SubViewOp>(
111  loc, newAlloc, ArrayRef<OpFoldResult>{rewriter.getIndexAttr(0)},
112  ArrayRef<OpFoldResult>{currSize},
113  ArrayRef<OpFoldResult>{rewriter.getIndexAttr(1)});
114  builder.create<memref::CopyOp>(loc, op.getSource(), subview);
115 
116  // Insert the deallocation of the old buffer only if requested
117  // (enabled by default).
118  if (emitDeallocs)
119  builder.create<memref::DeallocOp>(loc, op.getSource());
120 
121  builder.create<scf::YieldOp>(loc, newAlloc);
122  },
123  [&](OpBuilder &builder, Location loc) {
124  // We need to reinterpret-cast here because either the input or output
125  // type might be static, which means we need to cast from static to
126  // dynamic or vice-versa. If both are static and the original buffer
127  // is already bigger than the requested size, the cast represents a
128  // subview operation.
129  Value casted = builder.create<memref::ReinterpretCastOp>(
130  loc, cast<MemRefType>(op.getResult().getType()), op.getSource(),
131  rewriter.getIndexAttr(0), ArrayRef<OpFoldResult>{targetSize},
132  ArrayRef<OpFoldResult>{rewriter.getIndexAttr(1)});
133  builder.create<scf::YieldOp>(loc, casted);
134  });
135 
136  rewriter.replaceOp(op, ifOp.getResult(0));
137  return success();
138  }
139 
140 private:
141  const bool emitDeallocs;
142 };
143 
144 struct ExpandReallocPass
145  : public memref::impl::ExpandReallocBase<ExpandReallocPass> {
146  ExpandReallocPass(bool emitDeallocs)
147  : memref::impl::ExpandReallocBase<ExpandReallocPass>() {
148  this->emitDeallocs.setValue(emitDeallocs);
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 } // namespace
167 
169  bool emitDeallocs) {
170  patterns.add<ExpandReallocOpPattern>(patterns.getContext(), emitDeallocs);
171 }
172 
173 std::unique_ptr<Pass> mlir::memref::createExpandReallocPass(bool emitDeallocs) {
174  return std::make_unique<ExpandReallocPass>(emitDeallocs);
175 }
static MLIRContext * getContext(OpFoldResult val)
This class describes a specific conversion target.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:66
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
This class helps build Operations.
Definition: Builders.h:216
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:497
This class represents a single result from folding an operation.
Definition: OpDefinition.h:268
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:791
MLIRContext * getContext() const
Definition: PatternMatch.h:829
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
Definition: PatternMatch.h:853
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
std::unique_ptr< Pass > createExpandReallocPass(bool emitDeallocs=true)
Creates an operation pass to expand memref.realloc operations into their components.
void populateExpandReallocPatterns(RewritePatternSet &patterns, bool emitDeallocs=true)
Appends patterns for expanding memref.realloc operations.
Include the generated interface declarations.
Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc, OpFoldResult ofr)
Converts an OpFoldResult to a Value.
Definition: Utils.cpp:112
LogicalResult applyPartialConversion(ArrayRef< Operation * > ops, const ConversionTarget &target, const FrozenRewritePatternSet &patterns, ConversionConfig config=ConversionConfig())
Below we define several entry points for operation conversion.
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:358