[Mlir-commits] [mlir] [mlir][vector] Add unroll patterns for vector.load and vector.store (PR #143420)

[Nishant Patel]

nbpatel wrote:

> Thanks for moving this from linearization to unrolling, i.e. from #136193 and #137558
> 
> As far as I can tell the existing 9 unroll patterns support unrolling to any tile size of the vector. See for example
> 
> elementwise:
> 
> https://github.com/llvm/llvm-project/blob/13fe07d670e8a115929c9e595c4490ef5c75f583/mlir/test/Dialect/Vector/vector-transforms.mlir#L187
> 
> transfer_read:
> 
> https://github.com/llvm/llvm-project/blob/13fe07d670e8a115929c9e595c4490ef5c75f583/mlir/test/Dialect/Vector/vector-transfer-unroll.mlir#L29
> 
> and the tests in the file you added your tests to in this PR. Can you motivate deviating from this design? For example, I would expect based on the other unroll patterns for provide a 'target shape' of 2x2 to convert a load of 6x6 to 9 loads of 2x2.
> 
> I think that transfer_read is a more complex op than vector load. For example, transfer_read is lowered to load [here](https://github.com/llvm/llvm-project/blob/f12b1ed11672bc40a53fb1180541b2fda6e7d9fc/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp#L406). So I think it should be possible to unroll load to any tile size, if that's what happens for transfer_reads.
> 
> In terms of achieving what is I think your goad of only having rank-1 loads, I propose the following
> 
> 1. use unrolling to go from loads of `6x6` to 6 loads of `1x6`
> 2. linearize the loads of `1x6` to loads of `6`
> 
> Inn my opinion, having a linearization pattern for 2 above is fine, because it converts **one** rank-2 op into **one** rank-1 op. That's what I was kind of driving at when I added this comment:
> 
> https://github.com/llvm/llvm-project/blob/a53003fe23cb6c871e72d70ff2d3a075a7490da2/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h#L411
> 
> I acknowledge that end result is exactly what you had in #136193 !! But I think it makes Vector more modular and intuitive for other users if unrolling and linearization have a clear boundary.

The motivation for this unrolling was to have 1D loads and stores and hence the target shape is set accordingly. But, I agree with your point of making it more modular to align with vector dialect's flexible unrolling design. Can I address this in a follow up PR?

https://github.com/llvm/llvm-project/pull/143420