[Mlir-commits] [mlir] [MLIR][Bufferization] Bail on automatic deallocation to enable reentrant behaviour (PR #72289)

[Matthias Springer]

https://github.com/matthias-springer requested changes to this pull request.

This change can lead to double deallocs, e.g.:
```
func.func private @foo(%c: i1) {
  %0 = memref.alloc() : memref<5xf32>
  %1 = memref.alloc() : memref<5xf32>
  %2 = arith.select %c, %0, %1 : memref<5xf32>
  memref.dealloc %2 : memref<5xf32>
  // additional deallocs will be added for %0 and %1
  return
}
```

In the above case, we may be able to add a workaround in the buffer deallocation pass to fix the double dealloc. The real problem are deallocs within nested regions. In such cases, the block that contains the dealloc **must** take ownership of the memref. Otherwise, another block (the one that has ownership instead) will deallocate the memref. There is currently no way to force ownership without redesigning parts of the algorithm. (I think it gets even trickier with loops where we do not know if it has 0 or more iterations.)

Ideally, the buffer deallocation pass should run a single time towards the end of the compilation process. (Note, we have a similar situation with One-Shot Bufferize: if you do partial bufferization, you cannot run One-Shot Bufferize a second time and bufferize the rest. There are certain `to_tensor` ops that it will reject.)

If you have pre-existing deallocs in your IR, there are two options:
1. Remove all deallocs and let the buffer deallocation pass re-insert them. We have a pass that does that for `memref.realloc`: `--expand-realloc=emit-deallocs=false`. Deallocs may not be inserted at the same position. E.g., deallocs will always be inserted in the same block as the alloc.
2. Annotate `memref.alloc`+`memref.dealloc` with `bufferization.manual_deallocation`. (You could do that after running the buffer deallocation pass.) Such ops will be ignored by the buffer deallocation pass. Note, this would not work with the example shown above if only one alloc has that attribute because the dealloc may receive a buffer that is not under manual deallocation.


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