[Mlir-commits] [mlir] [mlir] Add optional layout attribute to VectorType (PR #71916)

[llvmlistbot at llvm.org]

harsh-nod wrote:

> Hey Hash, thanks for the proposal. This is an interesting change! Some initial questions that come to mind:
> 
> * How would we lower something like this to LLVM?
> * Should this be a property of the type or of the operations supporting these swizzles? In the doc you mention that this is mostly for SIMT matmuls, which make me wonder if we should only make this available for those ops.
> * Assuming this should be a property of the type, should we consider a "LayoutVectorType", where the attribute is not optional but mandatory?
> * It would be very helpful to add a bunch for IR examples. I would help me get a clearer idea of what we are trying to achieve here.
> 
> This looks like a big change so I would give it enough time to settle. Perhaps we should discuss some of these questions around an RFC?

So this PR is un-opinionated in terms of how the vector layout can be represented. The idea behind that is that many people will have a preference on what layout representation best suits them and we would like upstream to be able to support everyone's ideas. So regarding lowering to LLVM, this PR does not prescribe a way to utilize the layout information to lower to LLVM. That will be user dependent. That being said, I am a proponent of a specific form of the layout representation that can be used to lower to LLVM. For more details on how that works as well as IR examples, please see the following slides and videos -
https://github.com/nod-ai/techtalks/blob/main/high_dimensional_layout_flash_attn_harsh_23.pdf
https://www.youtube.com/watch?v=fqeyTCmqO4g
https://www.youtube.com/watch?v=5i7xrBUCD38

Regarding whether this should be a property of the type or the operations that support these swizzles -
I believe this should be a property of the type but it stems from constraints/requirements on the operations. While my work is currently mostly driven by GPU/SIMT requirements, I believe this layout is more general than that and could be used for other execution models as well.

Regarding a new vector with layout type, that is an interesting idea. However, my concerns would be the scope of the changes required to have all the vector, arithmetic etc. ops support this new type. Also, this new type would be a superset of the existing vector type so would seem like a lot of duplication.

Hope the links above provide you with enough IR examples. If not, happy to get on a GVC to give more info. 

Happy to continue the conversation anywhere. I am hoping that this is not a big change with the goal that everyone should hopefully be able to continue using vector ops and types as usual with extra work required from those who want to use this layout.




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