iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.85k stars 621 forks source link

[LLVMGPUVectorDistribute] Fix vector step distribute #19227

Closed manupak closed 5 days ago

manupak commented 5 days ago

Currently, the 'thread_stride' of NestedLayoutAttr is misinterpreted as the access stride of multi-dimensional vector.

However, it turns out it correspond to tid -> vtid mapping and the undistributed vector is packed as : subgroup x batch x outer x thread x element where vtid is used to index 'thread' dimension.

Therefore, this commit removes the usage of 'thread_stride's and 'subgroups_stride' when calculating the base constant offset and rather obtain them from packed undistributed vector shape.