[mlir][memref] annotate operand and result of realloc with proper memory attributes

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D143599

Author: aartbik

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

@@ -246,12 +246,16 @@ def MemRef_ReallocOp : MemRef_Op<"realloc"> {
     ```
   }];
 
-  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,

mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp

@@ -639,6 +639,16 @@ struct DefaultAllocationInterface
   }
 };
 
+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 @@ void bufferization::registerAllocationOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
     memref::AllocOp::attachInterface<DefaultAllocationInterface>(*ctx);
+    memref::ReallocOp::attachInterface<DefaultReallocationInterface>(*ctx);
   });
 }
 

mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir

@@ -1418,3 +1418,24 @@ func.func @test_affine_if_4(%arg0 : memref<10xf32>) -> memref<10xf32> {
 // CHECK-NEXT:    %[[ALLOC:.*]] = memref.alloc() : memref<10xf32>
 // CHECK-NEXT:    memref.dealloc %[[ALLOC]] : memref<10xf32>
 // CHECK-NEXT:    affine.if
+
+// -----
+
+// Ensure we free the realloc, not the alloc.
+
+// CHECK-LABEL: func @auto_dealloc()
+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
+}
+// CHECK-DAG:   %[[C10:.*]] = arith.constant 10 : index
+// CHECK-DAG:   %[[C100:.*]] = arith.constant 100 : index
+// CHECK-NEXT:  %[[A:.*]] = memref.alloc(%[[C10]]) : memref<?xi32>
+// CHECK-NEXT:  %[[R:.*]] = memref.realloc %alloc(%[[C100]]) : memref<?xi32> to memref<?xi32>
+// CHECK-NEXT:  memref.dealloc %[[R]] : memref<?xi32>
+// CHECK-NEXT:  return
+
+