[mlir][linalg] Refine how contiguous loads are identified
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
Differential Revision: https://reviews.llvm.org/D145385
parent
c7fcae52
Please register or sign in to comment