This is an archive of the discontinued LLVM Phabricator instance.

[mlir][Vector] Add support for masked vector gather ops
ClosedPublic

Authored by dcaballe on Feb 13 2023, 12:15 PM.

Details

Summary

This patch adds support for masked vector.gather ops using the
vector.mask representation. It includes the implementation of the
MaskableOpInterface, Linalg vectorizer support and lowering to LLVM.

Diff Detail

Event Timeline

dcaballe created this revision.Feb 13 2023, 12:15 PM
Herald added a project: Restricted Project. · View Herald Transcript
dcaballe requested review of this revision.Feb 13 2023, 12:15 PM
dcaballe added inline comments.
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
628

I'll need https://reviews.llvm.org/D142634 to be able to add a test to Linalg vectorizer but I can do that when that diff lands. I've tested this end to end for now.
(FYI, @awarzynski)

ThomasRaoux accepted this revision.Feb 14 2023, 11:27 AM

Looks good to me. Would be better to have a test to at least make sure it doesn't crash or emit wrong code until everything is ready

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
628

What happens if we try to vectorize a gather right now? Do we skip vectorization? Could be a negative test for now?

This revision is now accepted and ready to land.Feb 14 2023, 11:27 AM
dcaballe added inline comments.Feb 14 2023, 11:50 AM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
628

We generate vector.gather ops when we vectorize some tensor.extract ops but that is temporarily disabled by default due to performance reason. We have a flag to enable it but that flag is missing in the transform dialect op masked_vectorize so I can write a test but it won't be vectorized if we don't have a way to enable the flag from the transform dialect op. We need https://reviews.llvm.org/D142634. I can write a test where we don't vectorize and then add the flag as part of https://reviews.llvm.org/D142634. Let me do that

dcaballe added inline comments.Feb 14 2023, 9:49 PM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
628

Actually the test is already in https://reviews.llvm.org/D142634 :) so we can land this.

This revision was automatically updated to reflect the committed changes.