[PATCH] D79329: [MLIR] Update the FunctionAndBlockSignatureConverter and NonVoidToVoidReturnOpConverter of Buffer Assignment

[Stephan Herhut via Phabricator via llvm-commits]

herhut added a comment.

Making it fail would be good enough for landing.



================
Comment at: mlir/test/Transforms/buffer-placement-prepration.mlir:34
+// CHECK-SAME: (i1, f16)
+//      CHECK: linalg.copy(%[[ARG2]], %[[RESULT]])
+//      CHECK: return %[[ARG0]], %[[ARG1]]
----------------
dfki-ehna wrote:
> herhut wrote:
> > If this returned a memref before, why would it not return a memref now? We should only rewrite the return values that got rewritten from `tensor` to `memref`. Not the ones that already were a `memref`. I'd assume those would be left alone.
> The idea of removing all Memref types from function results is proposed to avoid placing a DeallocOp after the return operation using buffer-placement pass. This case would happen if there is a defining operation for a return operand.
> 
> ```
> func @memref_in_function_results(%arg0: i1, %arg1: f16, %arg2:memref<2xf32>) -> (i1, memref<2xf32>, f16){
>   %0 = alloc() : memref<2xf32>,
>   lhlo.exp(%arg2, %0)
>   return %arg0, %0, %arg1  : i1, memref<2xf32>, f16
> }
> ```
I think it is better to just not allow this case. If one has a memref that escapes in the original program, it should still escape in the rewritten program. We cannot support this use case at the moment it seems, so it would be better to fail. 

I wonder what it would take, though, to handle this case. Would it suffice to know that the last use is a return and then treat the return like a `dealloc` from the local scope? This could be an alternative version of this rewriting that turns tensor typed returns into memref typed returns (and not arguments) and lets allocations escape. Would the main logic of placement still work in this setting? What would it take to support both modes?

I see a use for the latter mode in scenarios where the buffers for results are not provided from the context, which is particularly useful in the presence of dynamic shapes where the context might not know sizes. 

Such a rewrite would not be globally correct (the caller would need to free the buffer after all and that cannot be decided statically) but at least locally per function it would be. For global correctness we could think in the direction of reference counting but that would be a next step.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D79329/new/

https://reviews.llvm.org/D79329