[mlir][linalg] Add option to generate rank-reducing slices in DropUnitDims
authorMatthias Springer <springerm@google.com>
Wed, 14 Dec 2022 13:02:13 +0000 (14:02 +0100)
committerMatthias Springer <springerm@google.com>
Wed, 14 Dec 2022 13:10:04 +0000 (14:10 +0100)
commite07149c91f4bb43c1413bc1fbe19dc6eff2fcde6
treef10ddb556a37ccfa2c5052b41da70a7a9ba0dd82
parentbefd167050ce46517a799cf0f997e2f75c20eb3b
[mlir][linalg] Add option to generate rank-reducing slices in DropUnitDims

This change extends the `ReplaceUnitExtents` pattern so that users can choose between of two strategies for generating rank reductions:
* CollapseShapeOp / ExpandShapeOp (was already implemented but code was cleaned up; default strategy)
* rank-reducing ExtractSliceOp / InsertSliceOp

Also add helper functions to the memref dialect that we already have on the tensor dialect: `getMixedSizes`, `createCanonicalRankReducingSubViewOp`, `rankReduceIfNeeded`.

We are using ReassociationIndices instead of ReassoicationExprs in many other places and this makes the code easier to read. Also adding a new test case (that also passed before).

Differential Revision: https://reviews.llvm.org/D139947
mlir/include/mlir/Dialect/Linalg/Passes.td
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
mlir/include/mlir/Dialect/MemRef/IR/MemRef.h
mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir