[Mlir-commits] [mlir] [mlir] [tensor] Add patterns to remove whole slicing of tensors (PR #107046)

[llvmlistbot at llvm.org]

Menooker wrote:

The long story why we need this pattern.

In our downstream project, we have the IR for **tiling**, like

```mlir
%a = extract_slice %out[...] : tensor<...>
%result = scf.for(...) iter_args(%a1 = %a) {
     %slice = extract_slice %a1[...]: tensor<...>
     %0 = elementwise_compute(%slice) outs(%a) : tensor<...>
     %inserted = insert_slice %0 into %a1[...] : tensor<...>
     scf.yield %inserted
}
%final = insert_slice %result into %out[...]
```

If the tile size for the loop happens to be the same size of `%out`, the `scf.for` has trip-count as 1, and it can be canonicalized to

```mlir
%a = extract_slice %out[...] : tensor<...>  // slice a sub-tensor from out
%slice = extract_slice %a[...]: tensor<...>  // whole-slicing over %a. %slice has the same size of %a
%0 = elementwise_compute(%slice) outs(%a) : tensor<...>
%inserted = insert_slice %0 into %a[...] : tensor<...> //  overwriting %a as a whole. %0 has the same size of %a
%final = insert_slice %inserted into %out[...]
```

And we have `populateMergeConsecutiveInsertExtractSlicePatterns` to remove the redundant `extract_slice` and `insert_slice`. However, the shape of `%a = extract_slice %out[...] : tensor<...>` (as well as `%slice`) is dynamic, and  `populateMergeConsecutiveInsertExtractSlicePatterns` seems not handling that. And it finally generates IR:

```mlir
%a = extract_slice %out[...] : tensor<...>  // slice a sub-tensor from out
%0 = elementwise_compute(%a) outs(%a) : tensor<...>
%inserted = insert_slice %0 into %a[...] : tensor<...> //  overwriting %a as a whole. %0 has the same size of %a
%final = insert_slice %inserted into %out[...]
```

This will cause failure of in-placing buffer for `%a` when bufferization. Because both `elementwise_compute(%a) outs(%a)` and `insert_slice %0 into %a[...] ` writes to `%a`.

Hence we introduce this pattern to remove the redundant `insert_slice`. 

The pattern have some overlapping with `populateMergeConsecutiveInsertExtractSlicePatterns`. However, it targets at different patterns in the IR.


https://github.com/llvm/llvm-project/pull/107046