iree-org / iree

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

[codegen] Add max_workgroup_counts to TargetWgpAttr #17771

Closed krzysz00 closed 1 month ago

krzysz00 commented 3 months ago

This commit adds a max_workgroup_counts to the workgroup processor information attribute and sets values for the known targets. Some of these values may be underestimates as I was not able to locate information on their values.

This field is added so that we can annotate calls to workgroup.id and workgroup.count with upper bounds, neabling range inference and strength reduction.

Note that in some cases (for instance, AMD) we give a max_workgroup_counts value lower than what is actually supported because a grid dimension greater than int32_max would be sign-extended to a negative number to meet the 64-bit nature of index.

(This PR is split out of #17707)

antiagainst commented 1 month ago

BTW please feel free to directly ping me via chats for re-reviews. I'm a bit bad at tracking GitHub notifications..