[Mlir-commits] [mlir] [mlir][bufferization] Ownership dealloc: support `IsolatedFromAbove` (PR #97669)

[Nikhil Kalra]

nikalra wrote:

> What's the reason that you'd like IsolatedFromAbove ops to be treated like functions?

I guess it's less important that IsolatedFromAbove ops be treated like functions, but rather that the existing pass was crashing + producing invalid IR.

For some psuedo-ish input IR:
```
%0 = memref.alloc() : memref<64xf32>

%1 = isolated_from_above(%0) ({
^bb0(%a0: memref<64xf32>):
  %2 = do_something(%a0) : memref<64xf32> -> memref<32xf32>
  yield %2 : memref<32xf32> -> ()
}) : (memref<64xf32>) -> memref<32xf32>

return %1
```

The first problem I ran into was that the buffer-flow analysis pass was crashing because it assumed `%a0` was captured from the enclosing region, but as it traversed the use-def chain, couldn't find the source allocation. That's fixed by: https://github.com/llvm/llvm-project/blob/103d76069cdc7d92a6fdf10954c3624d9c539109/mlir/lib/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.cpp#L177

The second issue I ran into was that the pass assumed it could take ownership of the block argument and deallocate it after use, leading to something like this:

```
%0 = memref.alloc() : memref<64xf32>

isolated_from_above(%0) ({
^bb0(%a0: memref<64xf32>):
  do_something(%a0) : memref<64xf32> -> ()
  memref.dealloc(%a0) : memref<64xf32>  // last use of %a0 but it aliases %0
}) : memref<64xf32> -> ()

memref.dealloc(%0) : memref<64xf32>  // last use of %0
```

Treating IsolatedFromAbove ops like functions seemed like the most expedient way to produce valid IR, since the function contract prohibits taking ownership of the arguments.

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