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

[springerm]

mlir/test/Dialect/Bufferization/Transforms/realloc.mlir

1 ↗

(On Diff #495908)

You can put the test case in this file: Dialect/Bufferization/Transforms/buffer-deallocation.mlir

[aartbik]

mlir/test/Dialect/Bufferization/Transforms/realloc.mlir

1 ↗

(On Diff #495908)

Ah, that makes sense yes. Done.

[aartbik]

merged test files

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

let description = [{
 
     ```mlir
     %new = memref.realloc %old : memref<64xf32> to memref<124xf32>
     %4 = memref.load %new[%index]   // ok
     %5 = memref.load %old[%index]   // undefined behavior
     ```
   }];
 
-  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,
                   "Value":$source,
                   CArg<"Value", "Value()">:$dynamicResultSize), [{
       return build($_builder, $_state, resultType, source, dynamicResultSize,
                    IntegerAttr());
     }]>];

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

return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)
         .getOperation();
   }
   static std::optional<Value> buildClone(OpBuilder &builder, Value alloc) {
     return builder.create<bufferization::CloneOp>(alloc.getLoc(), alloc)
         .getResult();
   }
 };
 
+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.
 struct BufferDeallocationPass
     : public bufferization::impl::BufferDeallocationBase<
           BufferDeallocationPass> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<bufferization::BufferizationDialect>();
LogicalResult bufferization::deallocateBuffers(Operation *op) {
 
   return success();
 }
 
 void bufferization::registerAllocationOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
     memref::AllocOp::attachInterface<DefaultAllocationInterface>(*ctx);
+    memref::ReallocOp::attachInterface<DefaultReallocationInterface>(*ctx);
   });
 }
 
 //===----------------------------------------------------------------------===//
 // BufferDeallocationPass construction
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<Pass> mlir::bufferization::createBufferDeallocationPass() {
   return std::make_unique<BufferDeallocationPass>();
 }

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

%0 = affine.if #set() -> memref<10xf32> {
     %alloc = memref.alloc() : memref<10xf32>
     affine.yield %alloc : memref<10xf32>
   }
   return %0 : 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
+
+