llvm / llvm/llvm-project

Is `isGuaranteedCastCompatible` not guaranteeing compatibility sufficiently?

Open
#194,679 6 comments 0 reactions 0 assignees View on GitHub
mlir:bufferization
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.