[MLIR] Nested regions test for Buffer Assginment

[dfki-ehna]

Summary

Due to the extension of Liveness, Buffer Assignment can now work on nested regions. This PR provides a test case to show that existing functionally of BA works properly.

[mehdi_amini]

mlir/test/Transforms/buffer-placement.mlir

416

Can you expand the doc on what this test is doing? Either here or in line with the checks themselves.

452

Nit: can you indent it to represent the expected nesting? It'll make it more readable.

[dfki-ehna]

Add description and indent.

[herhut]

Thanks, I am happy to see this land. Just some nits.

mlir/test/Transforms/buffer-placement.mlir

416

Could you reflow the comment to match 80 chars?

455

Could you add where the BB3 starts? That makes it easier to understand what is going on here.

mlir/test/Transforms/buffer-placement.mlir

^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
     linalg.yield %tmp1 : f32
   }: memref<2xf32>, memref<2xf32>
   dealloc %0 : memref<2xf32>
   "linalg.copy"(%0, %arg1) : (memref<2xf32>, memref<2xf32>) -> ()
   return
 }
 
 //      CHECK: linalg.copy
-// CHECK-NEXT: dealloc
+// CHECK-NEXT: dealloc
+
+// -----
+
+// Test Case: Nested regions
+
+#map0 = affine_map<(d0) -> (d0)>
+
+// CHECK-LABEL: func @nested_regions_and_cond_branch
+func @nested_regions_and_cond_branch(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
+  cond_br %arg0, ^bb1, ^bb2
+^bb1:
+  br ^bb3(%arg1 : memref<2xf32>)
+^bb2:
+  %0 = alloc() : memref<2xf32>
+  linalg.generic {args_in = 1 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0], iterator_types = ["parallel"]} %arg1, %0 {
+  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
+    %1 = alloc() : memref<2xf32>
+    linalg.generic {args_in = 1 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0], iterator_types = ["parallel"]} %arg1, %1 {
+    ^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
+      %tmp2 = exp %gen2_arg0 : f32
+      linalg.yield %tmp2 : f32
+    }: memref<2xf32>, memref<2xf32>
+    %tmp1 = exp %gen1_arg0 : f32
+    linalg.yield %tmp1 : f32
+  }: memref<2xf32>, memref<2xf32>
+  br ^bb3(%0 : memref<2xf32>)
+^bb3(%1: memref<2xf32>):
+  "linalg.copy"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
+  return
+}
+//      CHECK: (%[[cond:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %{{.*}}: {{.*}})
+// CHECK-NEXT: %[[GENERIC1_ALLOC:.*]] = alloc()
+// CHECK-NEXT: cond_br %[[cond]], ^[[BB1:.*]], ^[[BB2:.*]]
+// CHECK-NEXT: ^[[BB1]]:
+// CHECK-NEXT: br ^
+// CHECK-NEXT: ^[[BB2]]:
+// CHECK-NEXT: linalg.generic {{{.*}}} %[[ARG1]], %[[GENERIC1_ALLOC]]
+// CHECK-NEXT: ^{{.*}}:
+// CHECK-NEXT: %[[GENERIC2_ALLOC:.*]] = alloc()
+// CHECK-NEXT: linalg.generic {{{.*}}} %[[ARG1]], %[[GENERIC2_ALLOC]]
+//      CHECK: dealloc %[[GENERIC2_ALLOC]]
+// CHECK-NEXT: %{{.*}} = exp
+//      CHECK: linalg.copy
+// CHECK-NEXT: dealloc %[[GENERIC1_ALLOC]]
 No newline at end of file