[Mlir-commits] [mlir] [mlir][bufferization] handle exit aliases in buffer-loop-hoisting (PR #226951)

[Federico Bruzzone]

================
@@ -236,3 +236,64 @@ func.func @empty_to_tensor_alloc() -> tensor<2x2xf32> {
   %0 = bufferization.alloc_tensor() : tensor<2x2xf32>
   return %0 : tensor<2x2xf32>
 }
+
+// -----
+
+// Hoisting must stay within the operation targeted by the transform.
+// CHECK-LABEL: func @buffer_loop_hoisting_while_scope(
+//  CHECK-NOT: memref.alloc
+//      CHECK: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+//      CHECK: scf.condition{{.*}} %[[ALLOC]]
+//      CHECK: memref.load %[[LAST]]
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+    %loop = transform.structured.match ops{["scf.while"]} in %root : (!transform.any_op) -> !transform.any_op
+    transform.bufferization.buffer_loop_hoisting %loop : !transform.any_op
+    transform.yield
+  }
+  func.func @buffer_loop_hoisting_while_scope(%condition: i1) -> index {
+    %c0 = arith.constant 0 : index
+    %last = scf.while () : () -> memref<1xindex> {
+      %buffer = memref.alloc() : memref<1xindex>
+      memref.store %c0, %buffer[%c0] : memref<1xindex>
+      scf.condition(%condition) %buffer : memref<1xindex>
+    } do {
+    ^bb0(%current: memref<1xindex>):
+      scf.yield
+    }
+    %result = memref.load %last[%c0] : memref<1xindex>
+    return %result : index
+  }
+}
+
+// -----
+
+// Targeting the enclosing function permits hoisting before the loop.
+// CHECK-LABEL: func @buffer_loop_hoisting_function_scope(
+// CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: memref.load %[[LAST]]
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+    %func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+    transform.bufferization.buffer_loop_hoisting %func : !transform.any_op
+    transform.bufferization.buffer_loop_hoisting %func : !transform.any_op
----------------
FedericoBruzzone wrote:

Why is the transform applied twice on `%func`? If it is meant to show idempotency, a short comment would help. Otherwise the second call can go. Right?

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