webgpu: use naive reduction - #28174
Conversation
Use naive reduction when the output size of ReduceMean is far greater than reduce size. Shared reduction method may need to transpose input, which costs much time.
|
When I run florence-2-base-vision-encoder-fp16-0.5-fp32 with input [batch_size:1, height:768, width:768], 10 ReduceMean nodes using shared reduction method cost much time, It is heavy to transpose the input. If we use the naive method, The model's total time is from ~340ms to ~270ms. |
There was a problem hiding this comment.
Pull request overview
Updates the WebGPU reduction kernel dispatch heuristic to prefer the naive reduction implementation for certain ReduceMean shapes, aiming to avoid transpose overhead that can dominate runtime when the reduced dimension is small but the output is very large.
Changes:
- Adjusts the
use_naive_reductionselection logic inReduceKernel::ComputeInternal. - Adds a
ReduceMean-specific threshold to route some cases toReduceNaiveProgram.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
Jie Chen (@jchen10) PTAL |
|
Jiajia Qin (@qjia7) PTAL, thanks |
|
Guenther Schmuelling (@guschmue) It seems that the CI failure is not related with this PR, please re-trigger the bots again, thanks. |
|
Rebase the code. |
|
Guenther Schmuelling (@guschmue) Please take a look it, thanks. |
|
Guenther Schmuelling (@guschmue) Just a gentle ping. |
a534347
into
microsoft:main
Use naive reduction when the output size of ReduceMean is far greater than reduce size. Shared reduction method may need to transpose input, which costs much time.
Description
Motivation and Context