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.
|
@jchen10 PTAL |
|
@qjia7 PTAL, thanks |
|
@guschmue PTAL |
|
@guschmue It seems that the CI failure is not related with this PR, please re-trigger the bots again, thanks. |
|
Rebase the code. |
|
@guschmue Please take a look it, thanks. |
|
@guschmue Just a gentle ping. |
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