[all-commits] [llvm/llvm-project] 673999: [mlir][Linalg] Cleanup the drop unit dims pass in ...

[MaheshRavishankar via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 67399932c767f0a64c83a500dc6f7806c09d9401
      https://github.com/llvm/llvm-project/commit/67399932c767f0a64c83a500dc6f7806c09d9401
  Author: Mahesh Ravishankar <ravishankarm at google.com>
  Date:   2023-07-19 (Wed, 19 Jul 2023)

  Changed paths:
    M mlir/include/mlir/Dialect/Linalg/Passes.td
    M mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
    M mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
    M mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
    M mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
    R mlir/test/Dialect/Linalg/fold-unit-trip-loops.mlir
    A mlir/test/Dialect/Linalg/test-drop-unit-dims.mlir
    M mlir/test/lib/Dialect/Linalg/CMakeLists.txt
    A mlir/test/lib/Dialect/Linalg/TestLinalgDropUnitDims.cpp
    M mlir/tools/mlir-opt/mlir-opt.cpp

  Log Message:
  -----------
  [mlir][Linalg] Cleanup the drop unit dims pass in Linalg.

TL;DR the following API functions have been merged

```
void populateFoldUnitExtentDimsViaReshapesPatterns(RewritePatternSet &patterns);
void populateFoldUnitExtentDimsViaSlicesPatterns(RewritePatternSet &patterns);
```

into

```
void populateFoldUnitExtentDimsPatterns(RewritePatternSet &patterns,
                                        ControlDropUnitDims &options);
```

To use the previous functionality use

```
ControlDropUnitDims options;
// By default options.rankReductionStrategy is
// ControlDropUnitDims::RankReductionStrategy::ReassociativeReshape.
populateFoldUnitExtentDimsPatterns(patterns, options);
```

and

```
ControlDropUnitDims options;
options.rankReductionStrategy = ControlDropUnitDims::RankReductionStrategy::ExtractInsertSlice
populateFoldUnitExtentDimsPatterns(patterns, options);

```

This pass is quite old and needed to be updated based on the current
approach to transformations in Linalg

- Instead of two patterns, one to just remove loop dimensions that are
  unit extent (and using 0 in the indexing maps), and another to drop
  the unit-extents in the operand shapes, combine into a single
  transformation. This avoid creating an intermediate step with
  indexing maps having 0's in the domains exp ressions.

- Expose the core transformation as a utility function and add a
  pattern that calls this transformation.

This is a mostly NFC change, apart from the API change and dropping
the patterns/test that only dropped the loops that are unit extents.

Differential Revision: https://reviews.llvm.org/D155518