Skip to content

[ET-VK] Reductions dispatch a fixed 16 threads regardless of the reduction length#22350

Description

@msluszniak

馃悰 Describe the bug

reduce_gwg() in backends/vulkan/runtime/graph/ops/impl/Reduce.cpp sizes every texture-storage reduction from the output extents and then hardcodes the thread count:

constexpr uint32_t max_nthreads = 16u;
constexpr uint32_t nworkers_per_group = 4u;
constexpr uint32_t ngroups = 4u;

For a global average pool the output is 1x1xC/4, so the dispatch is global {1,1,4} local {4,4,1}. Four threads reduce an entire HxW plane and the whole dispatch runs 16 threads, independent of how much work there is.

Measured with the shader query pool on MediaPipe selfie segmentation @256 (the MobileNetV3 squeeze-excitation pools land here), Adreno 840, 120 dispatches, 2.194 ms GPU total:

kernel count total share
mean2d_float 10 0.767 ms 35.0%
conv2d_pw_tiled_float 42 0.677 ms 30.8%

Worst single dispatch 447 us, the rest 113-121 us, every one at global {1,1,4}.

reduce_gwg() backs both add_reduce_node and add_reduce2d_node, and DEFINE_REDUCE_FN registers aten.sum.dim_IntList, aten.mean.dim, aten.amax.default and aten.amin.default, so every global pool and texture-storage reduction pays it.

Fix in #22348: scale the worker count with the reduction extent. Selfie segmentation goes 2.306 -> 1.625 ms on an Adreno 840 (-28.5%) and 30.3 -> 25.4 ms on a Mali-G76 (-17%), bit-exact against the reference on both.

Versions

main @ 60cb889

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions