[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
Loading
Please register or sign in to comment