[all-commits] [llvm/llvm-project] 973dbe: [mlir][tensor] Add pattern to fold ExtractSliceOp, ...

[Tobias Gysi via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 973dbe20f681ca885edd0f5e63cde62dbdb6c186
      https://github.com/llvm/llvm-project/commit/973dbe20f681ca885edd0f5e63cde62dbdb6c186
  Author: gysit <gysit at google.com>
  Date:   2022-04-11 (Mon, 11 Apr 2022)

  Changed paths:
    M mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
    M mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
    M mlir/test/Dialect/Tensor/canonicalize.mlir

  Log Message:
  -----------
  [mlir][tensor] Add pattern to fold ExtractSliceOp, PadOp chains.

The pattern folds chains of tensor::ExtractSliceOp, tensor::PadOp pairs if they pad different dimensions. Repeated tiling and padding of the tiled dimensions may introduce such chains. This canonicalization pattern folds these chains to a single tensor::ExtractSliceOp, tensor::PadOp pair that pads all dimensions at once, which simplifies vectorization and bufferization.

Example:
```mlir
   %0 = tensor.extract_slice %input[16, 0] [%sz0, 64] [1, 1]
       : tensor<64x64xf32> to tensor<?x64xf32>
   %1 = tensor.pad %0 low[0, 0] high[%pw0, 0] { ...
     } : tensor<?x64xf32> to tensor<8x64xf32>
   %2 = tensor.extract_slice %1[0, 4] [8, %sz1] [1, 1]
        : tensor<8x64xf32> to tensor<8x?xf32>
   %res = tensor.pad %2 nofold low[0, 0] high[0, %pw1] { ...
     } : tensor<8x?xf32> to tensor<8x4xf32>
```
folds into:
 ```mlir
   %0 = tensor.extract_slice %input[16, 4] [%sz0, %sz1] [1, 1]
        : tensor<64x64xf32> to tensor<?x?xf32>
   %res = tensor.pad %0 nofold low[0, 0] high[%pw0, %pw1] { ...
     } : tensor<?x?xf32> to tensor<8x4xf32>
 ```

Reviewed By: nicolasvasilache, hanchung

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