[ET-VK] Add 2D Reduction to Vulkan Backend

[alexdean08]

Summary: This change adds 2D reduction to the Vulkan delegate. Prior to this change, only 1D reduction was implemented. Models like MobileNetV3 and ResNet do 2D reduction, and their performance was being negatively impacted by the lack of a 2D reduction Vulkan implementation.

cc @SS-JIA @manuelcandales @cbilgin

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• ghstack-mergeability-check and Check labels failing with 'Resource not accessible by integration'

❌ 2 New Failures, 1 Unrelated Failure

As of commit 3ca809f with merge base 9d86cbe ():

NEW FAILURES - The following jobs have failed:

• Build documentation / build (buck2) / Build doc (gh)
  At least one of the pre-conditions you specified did not hold
• pull / test-moshi-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 7cceb3d712b5b88a8a9be73d32c01f8ff41403d776ce9b6d11b363ed6d215f5d /exec failed with exit code 1

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / test-openvino-linux / linux-job (gh) (trunk failure)
  AttributeError: '_OpNamespace' 'quantized_decomposed' object has no attribute 'convert_element_type'

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

[alexdean08]

@pytorchbot label "release notes: vulkan"

[SS-JIA]

LGTM with some small comments. Thanks for working on this!

[SS-JIA]

As I understand it, this implementation doesn't currently support reduction among the elements of a texel,.

So for example for a tensor of size (channels = 4, height = 16, width = 16) where reading one texel returns 4 channels, this function would only be able to handle reduction along the height and width dims, but not if channels is one of the dims being reduced - is that correct?

If so, we can leave the more generalized implementation for later, but we should add some checks to make sure we don't lower unsupported cases to Vulkan 😛

Would also recommend renaming this function to reduce_2d_non_texel_dim or reduce_2d_non_packed_dim.

[alexdean08]

Yup you're right, this doesn't support reduction among the elements of a texel. Just changed to reduce_2d_non_packed_dim. Let me know if you're aligned

[SS-JIA]

should also add a check that graph.packed_dim_of(in) and graph.packed_dim_of(out) is not one of the reduction dims.

[alexdean08]

Just added above this comment. Let me know if you're aligned