[Mlir-commits] [mlir] [mlir] [memref] Elevate `AllocOp`s to `GlobalOp`s pass (PR #211141)

[ioana ghiban]

ioghiban wrote:

This transformation is not safe when the allocated memref escapes the containing function.
See this example:
```mlir
/// BEFORE
func.func @allocs(%v0 : f32, %v1 : f32, %i : index) -> f32 {
  %a = call @init() : () -> memref<10xf32>
  %b = call @init() : () -> memref<10xf32>

  memref.store %v0, %a[%i] : memref<10xf32>
  memref.store %v1, %b[%i] : memref<10xf32>

  %result = memref.load %a[%i] : memref<10xf32>

  memref.dealloc %a : memref<10xf32>
  memref.dealloc %b : memref<10xf32>
  return %result : f32
}

func.func @init() -> memref<10xf32> {
  %alloc = memref.alloc() : memref<10xf32>
  return %alloc : memref<10xf32>
}

/// AFTER
func.func @allocs(%arg0: f32, %arg1: f32, %arg2: index) -> f32 {
  %0 = call @init() : () -> memref<10xf32>
  %1 = call @init() : () -> memref<10xf32>
  memref.store %arg0, %0[%arg2] : memref<10xf32>
  memref.store %arg1, %1[%arg2] : memref<10xf32>
  %2 = memref.load %0[%arg2] : memref<10xf32>
  memref.dealloc %0 : memref<10xf32>
  memref.dealloc %1 : memref<10xf32>
  return %2 : f32
}
func.func @init() -> memref<10xf32> {
  %0 = memref.get_global @global_alloc : memref<10xf32>
  return %0 : memref<10xf32>
}
memref.global "private" @global_alloc : memref<10xf32> = uninitialized
```

For example, if @init returns an allocation and is called twice, the original program returns two distinct memrefs. After elevation, both calls return the same `memref.get_global`, so stores through one result are observable through the other.

The caller-side `memref.deallocs` also remain, because the pattern only removes deallocations that directly use the original `memref.alloc` result. The transformed program consequently attempts to deallocate the same global twice.

It seems the pass must either reject escaping allocations or perform interprocedural escape/deallocation analysis, in addition to proving that the containing function is invoked at most once.

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