diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td --- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td +++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td @@ -246,12 +246,16 @@ ``` }]; - let arguments = (ins MemRefRankOf<[AnyType], [1]>:$source, + // Note that we conceptually mark the operands as freeing the incoming + // memref and allocating the outcoming memref, even though this may not + // physically happen on each execution. + + let arguments = (ins Arg<MemRefRankOf<[AnyType], [1]>, "", [MemFree]>:$source, Optional<Index>:$dynamicResultSize, ConfinedAttr<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$alignment); - let results = (outs MemRefRankOf<[AnyType], [1]>); + let results = (outs Res<MemRefRankOf<[AnyType], [1]>, "", [MemAlloc<DefaultResource>]>); let builders = [ OpBuilder<(ins "MemRefType":$resultType, diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp @@ -639,6 +639,16 @@ } }; +struct DefaultReallocationInterface + : public bufferization::AllocationOpInterface::ExternalModel< + DefaultAllocationInterface, memref::ReallocOp> { + static std::optional<Operation *> buildDealloc(OpBuilder &builder, + Value realloc) { + return builder.create<memref::DeallocOp>(realloc.getLoc(), realloc) + .getOperation(); + } +}; + /// The actual buffer deallocation pass that inserts and moves dealloc nodes /// into the right positions. Furthermore, it inserts additional clones if /// necessary. It uses the algorithm described at the top of the file. @@ -703,6 +713,7 @@ DialectRegistry ®istry) { registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) { memref::AllocOp::attachInterface<DefaultAllocationInterface>(*ctx); + memref::ReallocOp::attachInterface<DefaultReallocationInterface>(*ctx); }); } diff --git a/mlir/test/Dialect/Bufferization/Transforms/realloc.mlir b/mlir/test/Dialect/Bufferization/Transforms/realloc.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Bufferization/Transforms/realloc.mlir @@ -0,0 +1,20 @@ +// RUN: mlir-opt --buffer-deallocation %s | FileCheck %s + +// Ensure we free the realloc, not the alloc + +// CHECK-LABEL: func @auto_dealloc() +// CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index +// CHECK-DAG: %[[C100:.*]] = arith.constant 100 : index +// CHECK: %[[A:.*]] = memref.alloc(%[[C10]]) : memref<?xi32> +// CHECK: %[[R:.*]] = memref.realloc %alloc(%[[C100]]) : memref<?xi32> to memref<?xi32> +// CHECK: memref.dealloc %[[R]] : memref<?xi32> +// CHECK: return +func.func @auto_dealloc() { + %c10 = arith.constant 10 : index + %c100 = arith.constant 100 : index + %alloc = memref.alloc(%c10) : memref<?xi32> + %realloc = memref.realloc %alloc(%c100) : memref<?xi32> to memref<?xi32> + return +} + +