Page MenuHomePhabricator

[mlir][Linalg] Generalize linalg vectorization
ClosedPublic

Authored by nicolasvasilache on Apr 23 2021, 7:28 AM.

Details

Summary

This revision adds support for vectorizing more general linalg operations with projected permutation maps.

This is achieved by eagerly broadcasting the intermediate vector to the common size
of the iteration domain of the linalg op. This allows a much more natural expression of
generalized vectorization but may introduce additional computations until all the
proper canonicalizations are implemented.

This generalization modifies the vector.transfer_read/write permutation logic and
exposes the fact that the logic employed in vector.contract was too ad-hoc.

As a consequence, changes occur in the permutation / transposition logic for contraction. In turn this prompts supporting more cases in the lowering of contract
to matrix intrinsics, which is required to make the corresponding tests pass.

Diff Detail

Event Timeline

nicolasvasilache requested review of this revision.Apr 23 2021, 7:28 AM
ThomasRaoux added inline comments.Apr 23 2021, 8:51 AM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
184

You need to check that vecType != nullptr

202–203

Do we have a way to test for this in pre-condition so that we don't reach this part if running vectorization on linalg op with this case? That would break things like IREE where vectorization is ran on all generic ops.

mlir/lib/Dialect/Vector/VectorOps.cpp
2202

It doesn't depend on the TransferReadOp? Should this be static?

Fix incorrect order of rershape / transpose.

asaadaldien added inline comments.Apr 28 2021, 2:14 PM
mlir/include/mlir/Dialect/Vector/VectorOps.td
295

I think we should also have an accumulation argument Variadic<AnyType>:$acc

nicolasvasilache marked 3 inline comments as done.Apr 28 2021, 2:50 PM
nicolasvasilache added inline comments.
mlir/lib/Dialect/Vector/VectorOps.cpp
2202

It is a static method in TransferReadOp, see the .td decl.

nicolasvasilache marked an inline comment as done.

Address review.

nicolasvasilache marked an inline comment as done.Apr 28 2021, 2:57 PM
nicolasvasilache added inline comments.
mlir/include/mlir/Dialect/Vector/VectorOps.td
295

We can but let's add it in a separate CL when we have a concrete use case: this revision does not need it and is already complex enough to my taste.
When you do the lowering, if you see an opportunity to add the acc please do so.

nicolasvasilache marked an inline comment as done.Apr 28 2021, 3:06 PM

LGTM! I benchmarked this on IREE's CPU backend and its causing no regression.

My only concern is other downstream systems (cc: @ThomasRaoux for SPIR-V) might fail with illegal op because of not lowering vectro.multi_reduction, so how about optionally enable/disable the vectorization if any of iterators is a reduction ?

LGTM! I benchmarked this on IREE's CPU backend and its causing no regression.

My only concern is other downstream systems (cc: @ThomasRaoux for SPIR-V) might fail with illegal op because of not lowering vectro.multi_reduction, so how about optionally enable/disable the vectorization if any of iterators is a reduction ?

I think it is easy enough to do this filtering on the IREE side so I don't think we have to add it here.

ThomasRaoux accepted this revision.Apr 28 2021, 10:54 PM
This revision is now accepted and ready to land.Apr 28 2021, 10:54 PM

also + @sgrechanik re. reduction detection that we'll want to unify

This revision was automatically updated to reflect the committed changes.