Update on "[ET-VK] Better work group sizes for matmul"

## Context

Currently `default_pick_local_wg_size()` (which internally calls `ComputeGraph::create_local_wg_size`) is used to select the local work group size for matrix multiplication ops. However, these functions currently bias the size of the local work group towards the largest dim of the global work group producing local wg sizes like

```
shader                                                                          globalwg size            localwg size
===========                                                                     =====================    ====================                 =============
linear_qga4w_tiled_texture3d_texture3d_texture2d_float                          {256, 29, 1}             {32, 2, 1}                                    1487
matmul_naive_texture3d_float                                                    {29, 115, 32}            {4, 2, 8}                                      712
```

for matrix multiplication shaders. This behaviour was introduced in D64418632 / #6409.

However, through experimental testing a "square" work group size of `{8, 8, 1}` works a lot better for matrix multiplication shaders. The theoretical analysis for this behaviour is that the local work group size determines the memory locations that need to be loaded to compute the overall work group. For a work group with size `{W, H, 1}` the data required to compute the output would be `W * OUTPUT_TILE_W` columns of the weight tensor and `H * OUTPUT_TILE_H` rows of the input tensor. Note that all work group items in the same W index will be requesting the same columns from the weight tensor, and all work group items in the same H index will be requesting the same rows from the input tensor.

If `H==W`, then that "balances" the amount of data needed to loaded from each input tensor and may result in better data sharing behaviour among all work group items. Assuming `OUTPUT_TILE_W == OUTPUT_TILE_H == 1`, a local work group of size `{64, 1, 1}` would require 1 unique row from the input tensor an 64 unique columns to be loaded from the weight tensor, resulting in `(1 + 64) * K = 65K` elements to be loaded in total, where K is the size of the shared reduction dim. Conversely, a local work group of size `{8, 8, 1}` would require 8 unique rows / 8 unique columns resulting in only `(8 + 8) * K = 16K` unique elements to be loaded.

This highlights the need to use dedicated logic to compute work group sizes for matrix multiplication shaders.

## Changes

* Introduce `pick_hw_square_wg_size`
* Use the new local work group size determination function for Quantized Linear, Matmul, and Linear

Differential Revision: [D79813236](https://our.internmc.facebook.com/intern/diff/D79813236/)

[ghstack-poisoned]

Author: ssjia

.ci/scripts/test_model.sh

@@ -201,7 +201,6 @@ test_model_with_qnn() {
     EXPORT_SCRIPT=bert
   elif [[ "${MODEL_NAME}" == "conv_former" ]]; then
     EXPORT_SCRIPT=conv_former
-    EXTRA_FLAGS="--dataset imagenet-mini/val"
   elif [[ "${MODEL_NAME}" == "cvt" ]]; then
     EXPORT_SCRIPT=cvt
   elif [[ "${MODEL_NAME}" == "distilbert" ]]; then

.ci/scripts/test_qnn_static_llama.sh

@@ -33,12 +33,12 @@ echo "Creating tokenizer.bin"
 $PYTHON_EXECUTABLE -m pytorch_tokenizers.tools.llama2c.convert -t tokenizer.model -o tokenizer.bin
 
 set +e
-# Compile only as weight sharing is not applicable on x86
-$PYTHON_EXECUTABLE backends/qualcomm/tests/test_qnn_delegate.py -k TestExampleLLMScript.test_llama_stories_110m --model SM8650 --build_folder build-android/ --executorch_root . --artifact_dir . --llama_artifacts . --compile_only
+# Compile only as weight sharing is not applicable on x86.
+$PYTHON_EXECUTABLE backends/qualcomm/tests/test_qnn_delegate.py -k TestExampleLLMScript.test_llama_stories_110m --model SM8650 --build_folder build-android/ --executorch_root . --artifact_dir ./stories_110m_pte_size --llama_artifacts . --compile_only
 exit_code1=$?
 
 # Checks accuracy with weight sharing disabled since x86 does not support weight sharing.
-$PYTHON_EXECUTABLE backends/qualcomm/tests/test_qnn_delegate.py -k TestExampleLLMScript.test_llama_stories_110m --model SM8650 --build_folder build-x86/ --executorch_root . --artifact_dir . --llama_artifacts . --enable_x86_64
+$PYTHON_EXECUTABLE backends/qualcomm/tests/test_qnn_delegate.py -k TestExampleLLMScript.test_llama_stories_110m --model SM8650 --build_folder build-x86/ --executorch_root . --artifact_dir ./stories_110m_accuracy --llama_artifacts . --enable_x86_64
 exit_code2=$?
 
 # Check BC

.github/workflows/add-unanswered-to-project.yml

@@ -0,0 +1,93 @@
+name: Add Open External Contributor PRs and Issues to PyTorch Org Project 136
+
+on:
+  schedule:
+    - cron: '0 * * * *'
+  workflow_dispatch:
+
+jobs:
+  add_to_project:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Add open issues and open, non-draft PRs to org project (excluding certain authors)
+        uses: actions/github-script@v7
+        with:
+          github-token: ${{ secrets.PYTORCH_PROJECT_PAT }}
+          script: |
+            const projectId = "PVT_kwDOAUB9vs4A_PUL"; // PyTorch org project 136
+            const owner = 'pytorch';
+            const repo = 'executorch';
+
+            // List of authors to exclude
+            const excludedAuthors = new Set([
+              "nil-is-all", "cbilgin", "KimishPatel", "psiddh", "digantdesai", "SS-JIA", "ahmtox", "mcr229", "shoumikhin",
+              "manuelcandales", "metascroy", "cccclai", "rohansjoshi", "kirklandsign", "abhinaykukkadapu", "JacobSzwejbka",
+              "Conarnar", "lucylq", "larryliu0820", "BujSet", "Gasoonjia", "Juntian777", "guangy10", "jackzhxng",
+              "GregoryComer", "leafs1", "swolchok", "mergennachin", "tarun292", "byjlw", "jathu", "Jack-Khuu", "georgehong",
+              "zhenyan-zhang-meta", "silverguo", "dbort", "jorgep31415", "huydhn", "mcremon-meta", "trivedivivek", "angelayi",
+              "helunwencser", "hsharma35", "zhxchen17", "iseeyuan", "svekars", "nathanaelsee", "dulinriley", "jerryzh168",
+              "cmodi-meta", "bigfootjon", "sxu", "ydwu4", "Riandy", "tugsbayasgalan", "bsoyluoglu", "yangw-dev", "YIWENX14",
+              "namanahuja", "yushangdi", "limintang", "pianpwk", "viveknayakatmeta", "andreanicastro", "JakeStevens",
+              "gmagogsfm", "zonglinpeng", "eigen-k", "derekxu", "salilsdesai", "skrtskrtfb", "pssrawat", "r-barnes", "pytorchbot",
+              "pytorchmergebot", "pytorchupdatebot", "facebook-github-bot", "Erik-Lundell", "zingo", "AdrianLundell",
+              "oscarandersson8218", "per", "Sebastian-Larsson", "SaoirseARM", "robell", "mansnils", "martinlsm", "freddan80",
+              "YufengShi-dudu", "tom-arm", "perheld", "Jerry-Ge", "gggekov", "fumchin", "wwwind", "haowhsu-quic", "shewu-quic",
+              "winskuo-quic", "chunit-quic", "DannyYuyang-quic", "chuntl", "cymbalrush", "DenisVieriu97", "billmguo",
+              "StrycekSimon", "jirioc", "robert-kalmar", "skywall", "neuropilot-captain"
+            ]);
+
+            async function addItem(contentId, type, number) {
+              try {
+                await github.graphql(`
+                  mutation {
+                    addProjectV2ItemById(input: {projectId: "${projectId}", contentId: "${contentId}"}) {
+                      item { id }
+                    }
+                  }
+                `);
+                console.log(`Added ${type} #${number} to project`);
+              } catch (error) {
+                if (error.message && error.message.includes("A project item already exists for this content")) {
+                  // Ignore if already exists
+                  console.log(`${type} #${number} already in project`);
+                } else {
+                  console.log(`Error adding ${type} #${number}: ${error.message}`);
+                }
+              }
+            }
+
+            try {
+              // Add open issues (not PRs) and exclude by author
+              const issues = await github.paginate(
+                github.rest.issues.listForRepo,
+                {
+                  owner,
+                  repo,
+                  state: 'open',
+                  filter: 'all'
+                }
+              );
+              for (const issue of issues) {
+                if (!issue.pull_request && !excludedAuthors.has(issue.user.login)) {
+                  await addItem(issue.node_id, 'issue', issue.number);
+                }
+              }
+
+              // Add open, non-draft PRs (regardless of review state), exclude by author
+              const prs = await github.paginate(
+                github.rest.pulls.list,
+                {
+                  owner,
+                  repo,
+                  state: 'open',
+                  draft: false,
+                }
+              );
+              for (const pr of prs) {
+                if (!excludedAuthors.has(pr.user.login)) {
+                  await addItem(pr.node_id, 'pr', pr.number);
+                }
+              }
+            } catch (error) {
+              core.setFailed(`Workflow failed: ${error.message}`);
+            }

.github/workflows/pull.yml

@@ -315,7 +315,7 @@ jobs:
         bash examples/models/moshi/mimi/install_requirements.sh
 
         # reinstall executorch
-        bash ./install_executorch.sh
+        bash ./install_executorch.sh --minimal
 
         # run python unittest
         python -m unittest examples.models.moshi.mimi.test_mimi

.github/workflows/trunk.yml

@@ -60,7 +60,7 @@ jobs:
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     strategy:
       matrix:
-        model: [add]
+        model: [add, softmax, mv2]
       fail-fast: false
     with:
       runner: linux.2xlarge
@@ -72,6 +72,16 @@ jobs:
         MODEL_NAME=${{ matrix.model }}
         CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
         conda activate "${CONDA_ENV}"
+        if [[ ${{ matrix.model}} == "add" ]]; then
+          SIM_LIMIT_SEC=60
+        elif [[ ${{ matrix.model}} == "softmax" ]]; then
+          SIM_LIMIT_SEC=60
+        elif [[ ${{ matrix.model}} == "mv2" ]]; then
+          SIM_LIMIT_SEC=5000
+        else
+          echo "Failed unsupported model selection ${{ matrix.model }}"
+          exit 1
+        fi
 
         source .ci/scripts/utils.sh
         source .ci/scripts/zephyr-utils.sh
@@ -80,20 +90,23 @@ jobs:
         export ZEPHYR_PROJ_ROOT=$(realpath $(pwd))
         export ARM_FVP_TUTORIALS_ROOT=$ZEPHYR_PROJ_ROOT/zephyr/samples/modules/executorch/arm-fvp-tutorials
 
+        # TODO @Bujji: Should see if this can be moved into the docker image itself
         download_arm_zephyr_sdk
         ./zephyr-sdk-0.16.0/setup.sh -c -t arm-zephyr-eabi
-
         cd $ZEPHYR_PROJ_ROOT
         setup_zephyr_et_module
 
+        # Run setup scripts for Arm FVP and Arm AOT Compilation
         cd $ZEPHYR_PROJ_ROOT/modules/lib/executorch
         install_executorch "--use-pt-pinned-commit"
         .ci/scripts/setup-arm-baremetal-tools.sh --target-toolchain zephyr
         source examples/arm/ethos-u-scratch/setup_path.sh
         source $ZEPHYR_PROJ_ROOT/zephyr/zephyr-env.sh
 
         # Get the model as PTE
-        python -m examples.arm.aot_arm_compiler --model_name="${MODEL_NAME}" --output="${MODEL_NAME}.pte"
+        python -m examples.arm.aot_arm_compiler \
+            --model_name="${MODEL_NAME}" \
+            --output="${MODEL_NAME}.pte"
 
         # Generate the C-style header
         cd $ARM_FVP_TUTORIALS_ROOT
@@ -105,7 +118,8 @@ jobs:
         cd $ARM_FVP_TUTORIALS_ROOT/models/${MODEL_NAME}/
 
         # Build the zephyr elf
-        west build -p always -b mps3/corstone300/fvp
+        west build -p always -b mps3/corstone300/fvp -- \
+            -DET_PTE_FILE_PATH_FOR_SELECTIVE_BUILD=$ZEPHYR_PROJ_ROOT/modules/lib/executorch/${MODEL_NAME}.pte
 
         # Run the simulation
         FVP_Corstone_SSE-300_Ethos-U55 -a build/zephyr/zephyr.elf \
@@ -114,23 +128,29 @@ jobs:
             -C mps3_board.uart0.out_file='sim.out'  \
             -C cpu0.CFGITCMSZ=15 \
             -C cpu0.CFGDTCMSZ=15 \
-            --simlimit 120
+            --simlimit ${SIM_LIMIT_SEC}
 
+        # Disable exit on error
+        set +e
         # Report failure if any of the ouptut verification checks fail
         grep -qF "ERROR" sim.out
         exit_status=$? #store 0 if found (failure), 1 if not (success)
         if [[ "$exit_status" -eq "0" ]]; then
-          cat sim.out
-          exit 1
+            cat sim.out
+            set -e
+            exit 1
         fi
 
         # Report fail if simulation does not complete successfully
         grep -qF "SUCCESS: Program complete, exiting." sim.out
         exit_status=$? #store 0 if found (success), 1 if not (failure)
         if [[ "$exit_status" -eq "1" ]]; then
-          cat sim.out
-          exit 1
+            cat sim.out
+            set -e
+            exit 1
         fi
+        # Re-enable exit on error
+        set -e
 
   test-models-linux-aarch64:
     name: test-models-linux-aarch64

backends/apple/coreml/TARGETS

@@ -60,6 +60,26 @@ runtime.python_library(
     ],
 )
 
+runtime.python_library(
+    name = "recipes",
+    srcs = glob([
+        "recipes/*.py",
+    ]),
+    visibility = [
+        "@EXECUTORCH_CLIENTS",
+    ],
+    deps = [
+        "fbsource//third-party/pypi/coremltools:coremltools",
+        ":backend",
+        "//caffe2:torch",
+        "//executorch/exir:lib",
+        "//executorch/exir/backend:compile_spec_schema",
+        "//executorch/exir/backend:partitioner",
+        "//executorch/exir/backend:utils",
+        "//executorch/export:lib",
+    ],
+)
+
 runtime.cxx_python_extension(
     name = "executorchcoreml",
     srcs = [
@@ -103,6 +123,7 @@ runtime.python_test(
         "fbsource//third-party/pypi/pytest:pytest",
         ":partitioner",
         ":quantizer",
+        ":recipes",
         "//caffe2:torch",
         "//pytorch/vision:torchvision",
     ],

backends/apple/coreml/recipes/__init__.py

@@ -0,0 +1,17 @@
+# Copyright © 2025 Apple Inc. All rights reserved.
+#
+# Please refer to the license found in the LICENSE file in the root directory of the source tree.
+
+
+from executorch.export import recipe_registry
+
+from .coreml_recipe_provider import CoreMLRecipeProvider
+from .coreml_recipe_types import CoreMLRecipeType
+
+# Auto-register CoreML backend recipe provider
+recipe_registry.register_backend_recipe_provider(CoreMLRecipeProvider())
+
+__all__ = [
+    "CoreMLRecipeProvider",
+    "CoreMLRecipeType",
+]