Is `isGuaranteedCastCompatible` not guaranteeing compatibility sufficiently?
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
I'm getting some seemingly incorrect code generated from one shot bufferization in some cases, in particular with bufferizing `tensor.extract_slice`. Consider the following input
```mlir
func.func @dynamic_2D_extract_slice_static_zero_offset(%t: tensor, %s0: index, %s1: index) -> (tensor) {
%extracted_slice = tensor.extract_slice %t[0, 0] [%s0, %s1] [1, 1] : tensor to tensor
return %extracted_slice : tensor
}
```
Applying `-one-shot-bufferize="bufferize-function-boundaries unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map"` produces the following output
```mlir
func.func @dynamic_2D_extract_slice_static_zero_offset(%arg0: memref, %arg1: index, %arg2: index) -> memref {
%subview = memref.subview %arg0[0, 0] [%arg1, %arg2] [1, 1] : memref to memref>
%cast = memref.cast %subview : memref> to memref
return %cast : memref
}
```
However, this cast seems unsafe since it is stripping the memref type of the strided layout. As far as I know `memref` and `memref>` are not equivalent nor are they treated the same in llvm lowerings. What will happen here is that `%cast` will take the strides from `%arg0`, which could be wrong in some cases. Since `strides[0]` is dynamic we can't know that it should be the same for the source and destination. Even if they were they same, some llvm lowerings will have different behavior depending on the value of `type.getLayout().isIdentity()`. In the above examples, the source type produces `true` while the dest type produces `false`. This may be the wrong way to handle the lowering anyway, but that's a whole separate topic.
Moreover, note that the closest resemblance in the static case produces expected code. With input
```mlir
func.func @static_2D_extract_slice_static_zero_offset(%t: tensor<40x64xf32>) -> (tensor<39x63xf32>) {
%extracted_slice = tensor.extract_slice %t[0, 0] [39, 63] [1, 1] : tensor<40x64xf32> to tensor<39x63xf32>
return %extracted_slice : tensor<39x63xf32>
}
```
and same pass configuration, the result is
```mlir
func.func @static_2D_extract_slice_static_zero_offset(%arg0: memref<40x64xf32>) -> memref<39x63xf32> {
%subview = memref.subview %arg0[0, 0] [39, 63] [1, 1] : memref<40x64xf32> to memref<39x63xf32, strided<[64, 1]>>
%alloc = memref.alloc() {alignment = 64 : i64} : memref<39x63xf32>
memref.copy %subview, %alloc : memref<39x63xf32, strided<[64, 1]>> to memref<39x63xf32>
return %alloc : memref<39x63xf32>
}
```
In this case, a copy is generated instead of a cast.
Finally, the non-zero offset case for dynamic also produces expected code. Input
```mlir
func.func @dynamic_2D_extract_slice_static_non_zero_offset(%t: tensor, %s0: index, %s1: index) -> (tensor) {
%extracted_slice = tensor.extract_slice %t[10, 0] [%s0, %s1] [1, 1] : tensor to tensor
return %extracted_slice : tensor
}
```
and output
```mlir
func.func @dynamic_2D_extract_slice_static_non_zero_offset(%arg0: memref, %arg1: index, %arg2: index) -> memref {
%subview = memref.subview %arg0[10, 0] [%arg1, %arg2] [1, 1] : memref to memref>
%c0 = arith.constant 0 : index
%dim = memref.dim %subview, %c0 : memref>
%c1 = arith.constant 1 : index
%dim_0 = memref.dim %subview, %c1 : memref>
%alloc = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref
memref.copy %subview, %alloc : memref> to memref
return %alloc : memref
}
```
I've narrowed down the source of this behavior to the [isGuaranteedCastCompatible](https://github.com/llvm/llvm-project/blob/3c14034c55a296306ad0ea4990f0f1b34e9e5d6e/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp#L41) function. I think it needs a little modification to handle the `memref to memref>` correctly. Is there anything wrong in my understanding or analysis and this is actually correct expected behavior? If I'm correct, I'd be happy to submit a PR to fix.
Thanks in advance.
@matthias-springer
Contributor guide
Research direction
Start in mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp at isGuaranteedCastCompatible and inspect how the dynamic memref types in the reported tensor.extract_slice examples are classified. Reproduce the one-shot-bufferize command for the zero-offset and non-zero-offset cases, then verify that the resulting cast or copy preserves layout compatibility and does not assume an unknown stride is safe.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100