Use subgroup operations when possible

[beaufortfrancois]

Subgroups can substantially enhance performance and adaptability for machine learning tasks on GPUs. Since they're now available on origin trial, it means https://webllm.mlc.ai/ could take advantage of them.

I'm not sure what is needed yet to make it work... I assume some work in Apache TVM as well.

I highly recommend you check out the quick-start guide at https://developer.chrome.com/blog/new-in-webgpu-128#experimenting_with_subgroups. For info, only subgroupBallot and subgroupBroadcast are there for now but more built-in functions such as subgroupAdd, subgroupAll, subgroupElect, subgroupShuffle will be added in a near future.

[beaufortfrancois]

@CharlieFRuan @tqchen What are your thoughts on this?

[tqchen]

This is great, subgroup shuffle can be useful for reduction operations. We did have warp shuffle support for metal backend, so maybe we can try add codegen backend for webgpu

[beaufortfrancois]

The following subgroup shuffle functions are actually in Chrome 129 (currently beta):

• subgroupShuffle(value, id): Returns value from the active invocation whose subgroup_invocation_id matches id.
• subgroupShuffleXor(value, mask): Returns value from the active invocation whose subgroup_invocation_id matches subgroup_invocation_id ^ mask. mask must be dynamically uniform.
• subgroupShuffleUp(value, delta): Returns value from the active invocation whose subgroup_invocation_id matches subgroup_invocation_id - delta.
• subgroupShuffleDown(value, delta): Returns value from the active invocation whose subgroup_invocation_id matches subgroup_invocation_id + delta.

[beaufortfrancois]

@tqchen @CharlieFRuan Is this being implemented in Apache TVM?

[CharlieFRuan]

Hi @beaufortfrancois Really appreciate the info and suggestions! We think it is a good idea to have it implemented in the TVM flow. Unfortunately, we are a bit out of bandwidth as of now. We'll revisit in the future!