[mlir][bufferize] Fix tensor copy insertion for dynamic tensors
authorMatthias Springer <springerm@google.com>
Wed, 21 Dec 2022 11:39:19 +0000 (12:39 +0100)
committerMatthias Springer <springerm@google.com>
Wed, 21 Dec 2022 11:42:02 +0000 (12:42 +0100)
commit5f6d5ca0f80c79813ab3821dc87ae01c86146398
tree308febc1253aae787c41fc63ff539b3828014942
parent701890164d567866900f3087ffd2ad4da963111c
[mlir][bufferize] Fix tensor copy insertion for dynamic tensors

TensorCopyInsertion inserts bufferization.alloc_tensor ops in case of RaW conflicts. If such a tensor is dynamically shaped, tensor.dim ops are inserted. There is an optimization for ops such as tensor.extract_slice: A copy of the result is created instead of the operand. Afterwards, all uses of the result are updated. E.g.:

```
%0 = tensor.extract_slice ... : tensor<?xf32> to tensor<?xf32>
%1 = tensor.dim %0, %c0 : tensor<?xf32>
%2 = bufferization.alloc_tensor(%dim) : tensor<?xf32>
```

All uses of %0, except for tensor.dim and bufferization.alloc_tensor (if any), should be replaced. Before this change, the use in tensor.dim was also replaced, resulting in IR that had a dominance error.

Note: There is no test case for this fix because the bug cannot be triggered with tensor.extract_slice, which implements an interface to reify result shapes. This bug appeared in an external project with a tensor.extract_slice-like op that does not implement that interface, in which case tensor.dim ops must be created. We do not have such an op in MLIR to trigger this bug.

Differential Revision: https://reviews.llvm.org/D140471
mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp