| Issue |
176739
|
| Summary |
[MLIR][Bufferization] one-shot-bufferize with bufferize-function-boundaries drops ranked return (memref.cast ignored)
|
| Labels |
mlir
|
| Assignees |
|
| Reporter |
tridhapuku
|
### Description
one-shot-bufferize with bufferize-function-boundaries changes the function result type to memref<*xf32> and returns the unranked value, even though a ranked memref.cast is created. This drops static shape information at the function boundary.
### Reproducer
```
module {
func.func @foo(%arg0: tensor<64x20x40xf32>) -> tensor<64x20x40xf32> {
%u = tensor.cast %arg0 : tensor<64x20x40xf32> to tensor<*xf32>
%r = call @relu(%u) : (tensor<*xf32>) -> tensor<*xf32>
%b = tensor.cast %r : tensor<*xf32> to tensor<64x20x40xf32>
return %b : tensor<64x20x40xf32>
}
func.func private @relu(tensor<*xf32>) -> tensor<*xf32>
}
```
### Commands
#### Failing (default)
`mlir-opt repro.mlir -one-shot-bufferize="bufferize-function-boundaries" -verify-each -o -`
```
module {
func.func @foo(%arg0: memref<64x20x40xf32, strided<[?, ?, ?], offset: ?>>) -> memref<*xf32> {
%cast = memref.cast %arg0 : memref<64x20x40xf32, strided<[?, ?, ?], offset: ?>> to memref<*xf32>
%0 = call @relu(%cast) : (memref<*xf32>) -> memref<*xf32>
%cast_0 = memref.cast %0 : memref<*xf32> to memref<64x20x40xf32, strided<[?, ?, ?], offset: ?>>
return %0 : memref<*xf32>
}
func.func private @relu(memref<*xf32>) -> memref<*xf32>
}
```
return %0 : memref<*xf32> is expected to be **memref<64x20x40xf32, strided<[?, ?, ?], offset: ?>>**.
#### Works with fully-dynamic-layout-map
`mlir-opt repro.mlir -one-shot-bufferize="bufferize-function-boundaries function-boundary-type-conversion=fully-dynamic-layout-map" -verify-each -o -`
### Observed output (key lines)
With default bufferize-function-boundaries:
- func.func @foo(...) -> memref<*xf32>
- return %0 : memref<*xf32>
- %cast_0 = memref.cast %0 : memref<*xf32> to memref<64x20x40xf32, ...> is created but not returned
With fully-dynamic-layout-map:
- func.func @foo(...) -> memref<64x20x40xf32, ...>
- return %cast_0 : memref<64x20x40xf32, ...>
### Expected behavior
When the original function returns a ranked tensor (here: tensor<64x20x40xf32>), the bufferized function boundary should return the corresponding ranked memref type and return the casted value (or otherwise preserve the ranked ABI).
### Environment
- llvm-project commit: 38be580aeef2fe3adf50b24854f39dad4a45c283
- mlir-opt --version: LLVM version 23.0.0git
### Notes
I have a minimal reproducer and I’m planning to work on a fix; I can open a draft PR with a regression test shortly.
_______________________________________________
llvm-bugs mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-bugs