[Mlir-commits] [mlir] [mlir][bufferization] Ownership dealloc: support `IsolatedFromAbove` (PR #97669)

[Matthias Springer]

================
@@ -0,0 +1,125 @@
+// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation \
+// RUN:   -buffer-deallocation-simplification -split-input-file %s | FileCheck %s
+
+func.func @function_call() {
+  %alloc = memref.alloc() : memref<f64>
+  %alloc2 = memref.alloc() : memref<f64>
+  %ret = test.isolated_one_region_with_recursive_memory_effects %alloc {
+  ^bb0(%arg1: memref<f64>):
+    test.region_yield %arg1 : memref<f64>
+  } : (memref<f64>) -> memref<f64>
+  test.copy(%ret, %alloc2) : (memref<f64>, memref<f64>)
+  return
+}
+
+// CHECK-LABEL: func @function_call()
+//       CHECK: [[ALLOC0:%.+]] = memref.alloc(
+//  CHECK-NEXT: [[ALLOC1:%.+]] = memref.alloc(
+//  CHECK-NEXT: [[RET:%.+]]:2 = test.isolated_one_region_with_recursive_memory_effects [[ALLOC0]]
+//  CHECK-NEXT: ^bb0([[ARG:%.+]]: memref<f64>)
+//       CHECK:   test.region_yield [[ARG]]
+//   CHECK-NOT:   bufferization.dealloc
+//       CHECK: }
+//  CHECK-NEXT: test.copy
+//  CHECK-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]#0
+//  CHECK-NEXT: bufferization.dealloc ([[ALLOC0]], [[ALLOC1]], [[BASE]] :{{.*}}) if (%true, %true, [[RET]]#1)
+
+// -----
+
+func.func @function_call_requries_merged_ownership_mid_block(%arg0: i1) {
+  %alloc = memref.alloc() : memref<f64>
+  %alloc2 = memref.alloca() : memref<f64>
+  %0 = arith.select %arg0, %alloc, %alloc2 : memref<f64>
+  %ret = test.isolated_one_region_with_recursive_memory_effects %0 {
+  ^bb0(%arg1: memref<f64>):
+    test.region_yield %arg1 : memref<f64>
+  } : (memref<f64>) -> (memref<f64>)
+  test.copy(%ret, %alloc) : (memref<f64>, memref<f64>)
+  return
+}
+
+// CHECK-LABEL: func @function_call_requries_merged_ownership_mid_block
+//       CHECK:   [[ALLOC0:%.+]] = memref.alloc(
+//  CHECK-NEXT:   [[ALLOC1:%.+]] = memref.alloca(
+//  CHECK-NEXT:   [[SELECT:%.+]] = arith.select{{.*}}[[ALLOC0]], [[ALLOC1]]
+//  CHECK-NEXT:   [[RET:%.+]]:2 = test.isolated_one_region_with_recursive_memory_effects [[SELECT]]
+//  CHECK-NEXT:   ^bb0([[ARG:%.+]]: memref<f64>)
+//       CHECK:     test.region_yield [[ARG]]
+//   CHECK-NOT:     bufferization.dealloc
+//       CHECK:   }
+//  CHECK-NEXT:   test.copy
+//  CHECK-NEXT:   [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]#0
+//  CHECK-NEXT:   bufferization.dealloc ([[ALLOC0]], [[BASE]] :
+//  CHECK-SAME:     if (%true, [[RET]]#1)
+//   CHECK-NOT:     retain
+//  CHECK-NEXT:   return
+
+// -----
+
+func.func @g(%arg0: memref<f32>) -> memref<f32> {
+  %0 = test.isolated_one_region_with_recursive_memory_effects %arg0 {
+  ^bb0(%arg1: memref<f32>):
+    test.region_yield %arg1 : memref<f32>
+  } : (memref<f32>) -> (memref<f32>)
+  return %0 : memref<f32>
+}
+
+// CHECK-LABEL:   func.func @g(
+// CHECK-SAME:                 %[[VAL_0:.*]]: memref<f32>) -> memref<f32> {
+// CHECK:           %[[BLOCK:.*]]:2 = test.isolated_one_region_with_recursive_memory_effects %[[VAL_0]] {
+// CHECK:           ^bb0(%[[ARG:.*]]: memref<f32>):
+// CHECK:             test.region_yield %[[ARG]], %false : memref<f32>, i1
+// CHECK:           } : (memref<f32>) -> (memref<f32>, i1)
+// CHECK:           %[[VAL_4:.*]] = scf.if %[[BLOCK]]#1 -> (memref<f32>) {
+// CHECK:             scf.yield %[[BLOCK]]#0 : memref<f32>
+// CHECK:           } else {
+// CHECK:             %[[VAL_6:.*]] = bufferization.clone %[[BLOCK]]#0 : memref<f32> to memref<f32>
+// CHECK:             scf.yield %[[VAL_6]] : memref<f32>
+// CHECK:           }
+// CHECK:           %[[BUF:.*]], %[[OFFSET:.*]] = memref.extract_strided_metadata %[[BLOCK]]#0 : memref<f32> -> memref<f32>, index
+// CHECK:           %[[VAL_11:.*]] = bufferization.dealloc (%[[BUF]] : memref<f32>) if (%[[BLOCK]]#1) retain (%[[VAL_4]] : memref<f32>)
+// CHECK:           return %[[VAL_4]] : memref<f32>
+// CHECK:         }
+
+// -----
+
+func.func @alloc_yielded_from_block() {
+  %alloc = memref.alloc() : memref<f64>
+  %alloc2 = memref.alloc() : memref<f64>
+  %ret = test.isolated_one_region_with_recursive_memory_effects %alloc {
+  ^bb0(%arg1: memref<f64>):
----------------
matthias-springer wrote:

What are the semantics of this op? Are there as many bbargs as memref operands? And are they the same buffer? Then the `RegionBranchOpInterface` should capture that.

In `getSuccessorRegions`: `regions.emplace_back(&getBody(), bbarg);`
And by adding `getEntrySuccessorOperands`.
Also the terminator must implement the `RegionBranchTerminatorOpInterface`.

https://github.com/llvm/llvm-project/pull/97669