[ET-VK][ez] Fix 8 bit linear compute shader dispatch

[SS-JIA]

Stack from ghstack (oldest at bottom):

• -> [ET-VK][ez] Fix 8 bit linear compute shader dispatch #9531

Context

Currently, for the q_8w_linear shader, both the texture and the buffer variants use the same global work group and local work group setting.

Specially, the global work group is set to {out.numel(), 1, 1} and the local work group is set to {64, 1, 1}.

However, I believe this results in a very poor memory re-use for the texture shader. In this configuration:

• Within a work group each invocation will be requesting a different row of A - 64 rows of A requested in total
• All work groups will be requesting the same row of B
• One work group will load 65 unique rows from A and B

Compare this to a local work group size of {8, 8, 1}

• Across the work group, 8 rows will be loaded from A and 8 rows will be loaded from B
• One work group will load 16 unique rows total from A and B

Evidently, there is better memory re-use in the latter work group as fewer unique rows are loaded.

Changes

Modify the q_8w_linear shader to use {8, 8, 1} local wg if possible. If M is small, then instead use {4, 16, 1} or {2, 32, 1} to reduce the number of inactive invocations.

Differential Revision: D71706489

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 1 Unrelated Failure

As of commit 44aace8 with merge base 76ae537 ():

NEW FAILURE - The following job has failed:

• Check Labels / Check labels (gh)
  RuntimeError: Error checking labels: PR does not have required labels

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

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

• pull / unittest-buck / macos / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 3

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

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D71706489

[github-actions]

This PR needs a release notes: label

If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

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