[Mlir-commits] [mlir] [mlir][bufferization] Add scf.for test cases to static memory planner (PR #215221)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Krish Gupta (KrxGu)

<details>
<summary>Changes</summary>

Extends test coverage from #<!-- -->213634 to document how the static memory planner handles scf.for patterns.

Four new tests (15–18):

- Test 15: alloc nested inside an scf.for body is skipped , same rule as test 13 for scf.if, only entry-block allocs get planned.
- Test 16: entry-block alloc passed as an iter_arg, freed and replaced each iteration. Skipped because the iter_arg dealloc also potentially frees the per-iteration nested alloc, which the arena doesn't own.
- Test 17: entry-block alloc as iter_arg, fresh buffer yielded each iteration without freeing the previous one (potential memory leak). Skipped for the same reason as test 16.
- Test 18: entry-block alloc freed directly inside the loop body, fresh buffer allocated and yielded each iteration (potential double-free / memory leak). Skipped for the same reason.

Also updates the test 13 comment to mention scf.for alongside scf.if since the same rule applies to both.

No pass changes , this is tests only.

---
Full diff: https://github.com/llvm/llvm-project/pull/215221.diff


1 Files Affected:

- (modified) mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir (+88-1) 


``````````diff
diff --git a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
index 5bf7b27133504..8bcfb023b2a7b 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -285,7 +285,7 @@ func.func @scf_if_result_aliases(%c: i1) {
 
 // -----
 
-// Test 13: Alloc nested inside an scf.if body is left untouched (not planned).
+// Test 13: Alloc nested inside a conditional/loop body is left untouched.
 // Only entry-block allocs are planned; the nested %b keeps its alloc/dealloc.
 // CHECK-LABEL: func @scf_if_nested_alloc_skipped
 func.func @scf_if_nested_alloc_skipped(%c: i1) {
@@ -322,3 +322,90 @@ func.func @scf_if_shared_nested_dealloc_skipped(%c: i1) {
   memref.dealloc %0 : memref<1024xf32>
   return
 }
+// -----
+
+// Test 15: Alloc nested inside an scf.for body is left untouched (same rule as
+// test 13 — only entry-block allocs are planned).
+// CHECK-LABEL: func @scf_for_nested_alloc_skipped
+func.func @scf_for_nested_alloc_skipped(%lb: index, %ub: index, %step: index) {
+  // CHECK-NOT: memref.view
+  // CHECK: scf.for
+  // CHECK: memref.alloc() : memref<1024xf32>
+  // CHECK: memref.dealloc
+  scf.for %iv = %lb to %ub step %step {
+    %b = memref.alloc() : memref<1024xf32>
+    memref.dealloc %b : memref<1024xf32>
+  }
+  return
+}
+
+// -----
+
+// Test 16: Entry-block alloc passed as scf.for iter_arg; each iteration frees
+// the current iter_arg and allocates a fresh buffer. The reverse-alias guard
+// conservatively skips %a because dealloc(%arg0) may also free the per-iteration
+// nested %b (which is not managed by the arena).
+// CHECK-LABEL: func @scf_for_iter_arg_nested_alloc
+func.func @scf_for_iter_arg_nested_alloc(%lb: index, %ub: index, %step: index) {
+  // CHECK: memref.alloc() : memref<1024xf32>
+  // CHECK-NOT: memref.view
+  // CHECK: scf.for
+  // CHECK: memref.dealloc
+  // CHECK: memref.alloc() : memref<1024xf32>
+  // CHECK: memref.dealloc
+  %a = memref.alloc() : memref<1024xf32>
+  %0 = scf.for %iv = %lb to %ub step %step iter_args(%arg0 = %a) -> memref<1024xf32> {
+    memref.dealloc %arg0 : memref<1024xf32>
+    %b = memref.alloc() : memref<1024xf32>
+    scf.yield %b : memref<1024xf32>
+  }
+  memref.dealloc %0 : memref<1024xf32>
+  return
+}
+
+// -----
+
+// Test 17: Entry-block alloc passed as scf.for iter_arg; each iteration
+// allocates a fresh buffer and yields it without freeing the previous iter_arg
+// (potential memory leak at runtime if the loop executes). The reverse-alias
+// guard skips %a because dealloc(%0) may also free the nested per-iteration %b.
+// CHECK-LABEL: func @scf_for_nested_alloc_yielded
+func.func @scf_for_nested_alloc_yielded(%lb: index, %ub: index, %step: index) {
+  // CHECK: memref.alloc() : memref<1024xf32>
+  // CHECK-NOT: memref.view
+  // CHECK: scf.for
+  // CHECK: memref.alloc() : memref<1024xf32>
+  // CHECK: memref.dealloc
+  %a = memref.alloc() : memref<1024xf32>
+  %0 = scf.for %iv = %lb to %ub step %step iter_args(%arg0 = %a) -> memref<1024xf32> {
+    %b = memref.alloc() : memref<1024xf32>
+    scf.yield %b : memref<1024xf32>
+  }
+  memref.dealloc %0 : memref<1024xf32>
+  return
+}
+
+// -----
+
+// Test 18: Entry-block alloc passed as scf.for iter_arg; the original %a is
+// freed directly inside the loop body (not via the iter_arg), and a fresh
+// buffer is allocated and yielded (potential double-free / memory leak at
+// runtime if the loop executes more than once). The reverse-alias guard skips
+// %a because dealloc(%0) may also free the nested per-iteration %b.
+// CHECK-LABEL: func @scf_for_orig_alloc_freed_in_body
+func.func @scf_for_orig_alloc_freed_in_body(%lb: index, %ub: index, %step: index) {
+  // CHECK: memref.alloc() : memref<1024xf32>
+  // CHECK-NOT: memref.view
+  // CHECK: scf.for
+  // CHECK: memref.dealloc
+  // CHECK: memref.alloc() : memref<1024xf32>
+  // CHECK: memref.dealloc
+  %a = memref.alloc() : memref<1024xf32>
+  %0 = scf.for %iv = %lb to %ub step %step iter_args(%arg0 = %a) -> memref<1024xf32> {
+    memref.dealloc %a : memref<1024xf32>
+    %b = memref.alloc() : memref<1024xf32>
+    scf.yield %b : memref<1024xf32>
+  }
+  memref.dealloc %0 : memref<1024xf32>
+  return
+}
\ No newline at end of file

``````````

</details>


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