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

Reply via email to