[mlir][TransformDialect] Simplify the lowering of pack/unpack when these are just...
authorQuentin Colombet <quentin.colombet@gmail.com>
Wed, 5 Apr 2023 14:14:18 +0000 (16:14 +0200)
committerQuentin Colombet <quentin.colombet@gmail.com>
Thu, 13 Apr 2023 10:46:36 +0000 (10:46 +0000)
commit0bfbecf52e8fa04785cf0b5c38b25c13442fb53d
tree679c747af97f9cf6c644b7b4af7305945904bbb9
parent2031d7d66dc8bf7ce2168edf6eef6ba568c16d4f
[mlir][TransformDialect] Simplify the lowering of pack/unpack when these are just pad/unpad

This patch recognizes when tensor.pack/unpack operations are simple
tensor.pad/unpad (a.k.a. tensor.extract_slice) and lowers them in a simpler
sequence of instruction.

For pack, instead of doing:
```
pad
expand_shape
transpose
```
we do
```
pad
insert_slice
```

For unpack, instead of doing:
```
transpose
collapse_shape
extract_slice
```
we do
```
extract_slice
```

Note: returning nullptr for the transform dialect is fine. The related
handles are just ignored by the following transformation.

Differential Revision: https://reviews.llvm.org/D148159
mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
mlir/include/mlir/Dialect/Utils/IndexingUtils.h
mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
mlir/lib/Dialect/Utils/IndexingUtils.cpp
mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
mlir/test/Dialect/Linalg/transform-lower-pack.mlir