This is an archive of the discontinued LLVM Phabricator instance.

[mlir][linalg] Refine how contiguous loads are identified
ClosedPublic

Authored by awarzynski on Mar 6 2023, 7:22 AM.

Details

Summary

Vectorization of tensor.extract using contiguous loads
(vector.transfer_read) was introduced in [1]. This patch updates and
refines the existing logic (so that more cases of contiguous can be
identified), as well as adds more tests.

Specifically, contiguous load operations are identified by making sure
that:

  1. non-trailing indices for tensor.extract are loop invariant (so, e.g., there are no "jumps" from one row to the other between iterations),
  2. the trailing index for tensor.extract increments by 1 with every loop iteration (so that it's always adjacent elements that are loaded).

This patch introduces:

  • isLoopInvariantIdx for step 1., and
  • isContiguousLoadIdx for step 2.

These new methods replace:

  • isContiguousLoadIdx, and isBasedOnIndexOp.

Both approaches lead to similar end-result (none of the existing tests
required updating). However, with the updated approach, it's much easier
to treat the trailing and non-trailing indices separately and to add
more cases for which contiguous loads can be used.

[1] https://reviews.llvm.org/D141998

Diff Detail

Event Timeline

awarzynski created this revision.Mar 6 2023, 7:22 AM
Herald added a reviewer: hanchung. · View Herald Transcript
Herald added a project: Restricted Project. · View Herald Transcript
awarzynski requested review of this revision.Mar 6 2023, 7:22 AM
awarzynski updated this revision to Diff 502643.Mar 6 2023, 7:24 AM

Apply clang-format

dcaballe accepted this revision.Mar 6 2023, 10:49 PM

LGTM, thanks! I'd love to see this also working for 2D vectors :)

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

Just to confirm, preconditions won't let us hit these asserts, right? Otherwise, we should gracefully bail out without crashing.

675

dyn_cast -> isa

709

Same for these asserts

716

dyn_cast -> isa

This revision is now accepted and ready to land.Mar 6 2023, 10:49 PM

I'd love to see this also working for 2D vectors :)

Haha, me too, just need a good reference example :)

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

Correct. I've added these asserts primarily to document the limitations of this method (it's easier to do it with code than with words).

awarzynski updated this revision to Diff 503061.Mar 7 2023, 8:34 AM

Address PR comments

If there are no new comments, I will merge this later today.

This revision was landed with ongoing or failed builds.Mar 8 2023, 12:24 AM
This revision was automatically updated to reflect the committed changes.