Page MenuHomePhabricator

[mlir][tensor] Add pattern to fold ExtractSliceOp, PadOp chains.
ClosedPublic

Authored by gysit on Mar 30 2022, 4:43 AM.

Details

Summary

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>

Diff Detail

Event Timeline

gysit created this revision.Mar 30 2022, 4:43 AM
Herald added a project: Restricted Project. · View Herald TranscriptMar 30 2022, 4:43 AM
gysit requested review of this revision.Mar 30 2022, 4:43 AM
gysit updated this revision to Diff 419131.Mar 30 2022, 6:54 AM

Make offset computation more robust.

hanchung accepted this revision.Mar 31 2022, 3:27 PM

Overall looks good to me, just few nits. :)

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
1962

maybe say they have the same constant padding value. I think constant is also a requirement in this pattern?

1963

do not have common padding dims?

2097

nit: s/PadOP/PadOp

This revision is now accepted and ready to land.Mar 31 2022, 3:27 PM
gysit updated this revision to Diff 419653.Apr 1 2022, 12:41 AM
gysit marked 2 inline comments as done.

Address comments.

Is there an opportunity to factor out folding logic into helpers similarly to what we have here ?

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
2050

Note, combining offsets also involves strides, see the link I shared from IREE.

gysit added inline comments.Apr 11 2022, 3:38 AM
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
2050

It is a bit specific here since one extract slice op / pad op pair always needs to have a zero offset and zero padding. I could separate this into a check and then use a method similar to the IREE one though? It would be more code overall but there may be reuse opportunities.

Also note that I cannot really perform actual computation since the tensor dialect shall not depend on the affine dialect.

gysit updated this revision to Diff 421879.Apr 11 2022, 4:56 AM

Rebase.

nicolasvasilache accepted this revision.Apr 11 2022, 6:33 AM
nicolasvasilache added inline comments.
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
2050

Ah yes, I had overlooked the dialect dependence issues.
Ok, in that case, and given that you also check explicitly for stride-1, I think this is fine.

This revision was landed with ongoing or failed builds.Apr 11 2022, 7:30 AM
This revision was automatically updated to reflect the committed changes.