Page MenuHomePhabricator

[MLIR] Vectorize tensor.extract on 1-d tensor
ClosedPublic

Authored by pzread on Sep 13 2022, 9:11 AM.

Details

Summary

This patch implements the vectorization of tensor.extract for the basic 1-d lookup case.

It only vectorizes the tensor.extract to a vector.gather when the op extracts value from an 1-d tensor.

Related discussion: https://github.com/iree-org/iree/issues/9198

Diff Detail

Event Timeline

pzread created this revision.Sep 13 2022, 9:11 AM
Herald added a project: Restricted Project. · View Herald Transcript
pzread updated this revision to Diff 459783.Sep 13 2022, 9:34 AM

Refactor.

pzread retitled this revision from [MLIR]Vectorize tensor.extract on 1-d table to [MLIR] Vectorize tensor.extract on 1-d tensor.Sep 13 2022, 9:40 AM
pzread edited the summary of this revision. (Show Details)
pzread edited the summary of this revision. (Show Details)
pzread added a reviewer: dcaballe.
pzread published this revision for review.Sep 13 2022, 9:42 AM

Thanks, @pzread! LGTM! A few comments.

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

nit: /*name=*/

593

Have you tried to move this into vectorizeOneOp instead of using a custom hook? I might be wrong here but, IIRC, custom vectorization hooks are supposed to be exceptional cases with higher priority that may fall back to other vectorization approach if the hook pattern fails. Most of the regular operation should be vectorized directly in vectorizeOneOp.

@nicolasvasilache can probably correct me/clarify.

649–651

I don't quite get what we are trying to do here. Operands coming from outside would be broadcasted or kept scalar, as needed. What is the problem? Perhaps, elaborating a bit more the comment would help.

mlir/test/Dialect/Linalg/vectorization.mlir
1497–1499

Do we have plans to support this case soon? :)

1507

nit: remove last // -----. It will run an empty test.

pzread updated this revision to Diff 466951.Oct 11 2022, 3:33 PM

Address comments

pzread updated this revision to Diff 466957.Oct 11 2022, 3:47 PM
pzread marked an inline comment as done.

Add comments.

pzread marked an inline comment as done.Oct 11 2022, 3:49 PM
pzread added inline comments.
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
593

The vectorizeOneOp can only handle ElementwiseMappable ops, which can take either scalar or vector as their operands. tensor::extract needs to be transformed into vector.gather, so I use a custom hook to handle it.

The comment on vectorizeOneOp also mentioned custom hook is designed to handle all non-ElementwiseMappable ops.

649–651

So the problem here is that these checks in vectorizeStaticLinalgOpPrecondition is intented to check ElementwiseMappable ops (assume all operands should be converted from scalar to vector). For tensor.extract (which isn't ElementwiseMappable), its first operand is a tensor, but we have a custom hook to transform it into vector.gather.

I changed this part and added a CustomVectorizationPrecondition list to put those custom precondition checks for custom-handled ops. Let me know WDYT.

mlir/test/Dialect/Linalg/vectorization.mlir
1497–1499

I guess we need more discussions on these cases : )

dcaballe accepted this revision.Oct 17 2022, 3:20 PM
dcaballe added inline comments.
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
649–651

It looks like we are implementing our ad-hoc pattern rewriter here but I guess it's ok for now.
My main concern is how this custom preconditions should compose with the existing pre-conditions. I think it makes sense for them to be auto-contained pre-conditions as they have all the context about the operation being evaluated.

mlir/test/Dialect/Linalg/vectorization.mlir
1497–1499

I think @awarzynski might be looking into this :)

This revision is now accepted and ready to land.Oct 17 2022, 3:20 PM

Thanks for the review @dcaballe : ) If you are okay with this revision, could you help me commit the change? thx

This revision was automatically updated to reflect the committed changes.