[mlir][linalg] Refine how contiguous loads are identified
authorAndrzej Warzynski <andrzej.warzynski@gmail.com>
Sun, 5 Mar 2023 12:23:17 +0000 (12:23 +0000)
committerAndrzej Warzynski <andrzej.warzynski@gmail.com>
Wed, 8 Mar 2023 07:49:04 +0000 (07:49 +0000)
commit7a078b65fb880eb97bb47a26e938cdbd23531688
treed3c39fdc601162313e90d2798e62f6871dec9d6b
parentc7fcae52f649fdab8600103ac367aaabbdf551ce
[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
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
mlir/test/Dialect/Linalg/vectorization.mlir