[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
Loading
Please sign in to comment