Vulkan: classify conv2d method by weight shape in conv2d_local_wg_size

[msluszniak]

Summary

Fixes #21942.

conv2d_local_wg_size() picked the convolution method from the shader name alone, and the condition it used matched every conv2d shader:

if (kernel_name.find("conv2d_pw") != npos ||
    (kernel_name.find("conv2d") != npos &&
     kernel_name.find("conv_transpose2d") == npos)) {
  method = Conv2dMethod::Pointwise;
} else {
  method = Conv2dMethod::SlidingWindow;
}

The sliding window shader is itself named conv2d, so it matched and was labelled Pointwise, making the SlidingWindow branch unreachable for every conv2d variant. Only conv_transpose2d reached the else.

The sibling conv2d_global_wg_size() directly above uses the identical outer name test but then disambiguates by inspecting the weight's spatial extent, so the two could disagree about the same dispatch: the global size computed as sliding window while the local size was computed as pointwise. This factors that classification into one function used by both, so they cannot drift apart again.

Why this is a restoration, not a new heuristic

The weight-shape check arrived with #13173, which introduced the tuned {64 / y, y, 1} local size for pointwise convolutions. Before that commit the dispatch used create_local_wg_size(global_size) for every method:

-  const utils::uvec3 local_size = graph.create_local_wg_size(global_size);

The name test swept sliding window convolutions into the new pointwise size along with the pointwise ones, so this restores the local size they had before #13173.

Scope of the behavior change

Only sliding window conv2d changes.

• Pointwise was classified correctly before and after, so it keeps the tuned size.
• Depthwise is routed to conv2d_dw_impl() before this dispatch and never reaches either function.
• Transposed resolves to the same non-pointwise branch as before (conv_transpose2d does not contain the substring conv2d, so it already fell through to create_local_wg_size).

create_local_wg_size() is pure arithmetic on the global workgroup size, so the exact before/after can be computed on the host. Across 7452 realistic non-pointwise conv2d output shapes (spatial extents 1 to 512, 3 to 2048 output channels):

• every case is 64 threads per group, before and after, so occupancy per dispatch is unchanged and nothing exceeds 64
• only the group shape differs, most commonly {8, 8, 1} to one of {8, 4, 2}, {4, 8, 2}, {2, 8, 4}, {4, 2, 8}

For example, a 128x128x64 output has global size {128, 128, 16} and goes from {8, 8, 1} to {8, 4, 2}.

Since sliding window convolutions were last benchmarked under create_local_wg_size before #13173, and the pointwise tuning in that commit was not aimed at them, I would expect this to be neutral to positive. I do not have perf numbers across GPUs, so if you would like this validated on a specific device before landing, say which and I will run it.

Test plan

vulkan_backend builds clean. The change is a refactor of method classification plus the restored branch; existing conv2d correctness coverage applies unchanged, since the local workgroup size affects scheduling and not results.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22051

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[SS-JIA]

also, just to make sure this function covers all possible Conv2dMethod entries, would recommend detecting Conv2dMethod::Depthwise as well.

[msluszniak]

I guess that now every enum entry is reachable.

[SS-JIA]

Overall, LGTM! Just have a few small comments. Thanks for the fix!