Update on "[ET-VK] Split up prepack command buffer"

## Changes

* Introduce `run_prepack()` API which combines the functionality of `encode_prepack()` and `prepack()`, but submits prepacking shaders incrementally rather than all at once.
* Introduce graph config options to control command buffer submission behaviour during prepacking.

Note that the current default values for the prepack submission thresholds were determined through experimentation. I will leave determining optimal values for specific devices as a later exercise. The goal of this diff is simply to introduce this mechanism to fix the Llama model loading crash on Samsung S24 (described below).

## Context

Currently, ET-VK will encode all prepacking shaders, and then perform prepacking by submitting only one command buffer.

However, this approach has some drawbacks:

* CPU/GPU parallelism is decreased, since the command buffer is submitted only after all commands have been encoded.
* There can be performance issues at the Vulkan API level when processing a single "large" command buffer.

By splitting up prepacking to occur over multiple command buffers, performance can be improved by avoiding both the aforementioned issues.

## Llama 3.2 1B crash on Samsung S24

I have also noticed that running large models (i.e. Llama 3.2 1B) on the Samsung S24 with ET-VK, the device's display will crash (causing the screen to go black and become unresponsive), and sometimes the device will shut down entirely.

Fortunately, this change also fixes this behaviour, in addition to providing a significant performance boost to model load time for Llama models (from 9s to 3s).

## Performance Impact

* Improves model load time, especially on larger models.

## Future Work

* Deprecate the `encode_prepack()` + `prepack()` pattern in favor of the `run_prepack()` pattern

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

[ghstack-poisoned]

Author: SS-JIA

.ci/scripts/unittest-buck2.sh

@@ -10,7 +10,7 @@ set -eux
 # TODO: can't query cadence & vulkan backends
 # TODO: can't query //kernels/prim_ops because of non-buckified stuff in OSS.
 buck2 query "//backends/apple/... + //backends/example/... + \
-//backends/mediatek/... + //backends/test/... + //backends/transforms/... + \
+//backends/mediatek/... + //backends/transforms/... + \
 //backends/xnnpack/... + //configurations/... + //kernels/aten/... + \
 //kernels/optimized/... + //kernels/portable/... + //kernels/quantized/... + \
 //kernels/test/... + //runtime/... + //schema/... + //test/... + //util/..."

.github/workflows/android-release-artifacts.yml

@@ -47,15 +47,13 @@ jobs:
     name: build-aar
     needs: check-if-aar-exists
     if: ${{ !github.event.pull_request.head.repo.fork }}
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@release/2.7
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     secrets: inherit
     permissions:
       id-token: write
       contents: read
     with:
       secrets-env: EXECUTORCH_MAVEN_SIGNING_KEYID EXECUTORCH_MAVEN_SIGNING_PASSWORD EXECUTORCH_MAVEN_CENTRAL_PASSWORD EXECUTORCH_MAVEN_CENTRAL_USERNAME EXECUTORCH_MAVEN_SIGNING_GPG_KEY_CONTENTS
-      # As this job has access to Maven credential, run this on a fresh ephemeral runner
-      runner: ephemeral.linux.2xlarge
       docker-image: executorch-ubuntu-22.04-clang12-android
       submodules: 'recursive'
       ref: ${{ github.sha }}

backends/arm/scripts/mlsdk_utils.sh

@@ -0,0 +1,121 @@
+#!/usr/bin/env bash
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+set -euo pipefail
+
+# TODO
+mlsdk_manifest_url=""
+
+script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+
+source ${script_dir}/utils.sh
+
+usage() { echo "Usage: $0 [-u <mlsdk-manifest-url>]" 1>&2; exit 1; }
+
+while getopts ":u:" opt; do
+    case "${opt}" in
+        u)
+            mlsdk_manifest_url=${OPTARG}
+            ;;
+        *)
+            usage
+            ;;
+    esac
+done
+
+function download_ai_mlsdk_manifest() {
+    local _dada_dir="$1"
+
+    if [[ -z "${_dada_dir}" ]]; then
+        echo "Error: _dada_dir parameter missing?"
+        return 1
+    fi
+
+    if [[ -z "${mlsdk_manifest_url}" ]]; then
+        echo "Error: mlsdk_manifest_url parameter missing?"
+        return 1
+    fi
+
+    if [[ ! -d "${_dada_dir}" ]]; then
+        mkdir -p "$_dada_dir"
+        pushd "$_dada_dir" || exit 1
+
+        curl https://storage.googleapis.com/git-repo-downloads/repo > repo
+        chmod u+x repo
+        ./repo init  --no-repo-verify --depth=1  --manifest-url  ${mlsdk_manifest_url} -g model-converter,emulation-layer,vgf-library
+        ./repo sync
+
+        popd
+    fi
+}
+
+function setup_model_converter() {
+    local work_dir="$1"
+    local manifest_dir="$2"
+    local enable_vgf_lib="$3"
+    local enable_emulation_layer="$4"
+
+    if [[ -z "$work_dir" ]]; then
+        echo "Error: work_dir parameter is required."
+        return 1
+    fi
+
+    if [[ -z "$manifest_dir" ]]; then
+        echo "Error: manifest_dir parameter is required."
+        return 1
+    fi
+
+    mkdir -p "$work_dir"
+    pushd "$work_dir" || exit 1
+
+    download_ai_mlsdk_manifest ${manifest_dir}
+
+    pushd "$manifest_dir"
+
+    # model-converter
+    # TODO: Remove macOS patch after mlsdk fully supports macOS
+    if [[ "$(uname)" == "Darwin" ]]; then
+    sed -i '' '/^ *print(f"Unsupported host platform/ i\
+            if system == "Darwin":\
+                # Use default Apple toolchain (Clang) on macOS\
+                return True\
+\
+' sw/model-converter/scripts/build.py
+    fi
+    python sw/model-converter/scripts/build.py -j$(nproc)
+
+    # libvgf
+    if [[ "${enable_vgf_lib}" -eq 1 ]]; then
+    # TODO: Remove macOS patch after mlsdk fully supports macOS
+        if [[ "$(uname)" == "Darwin" ]]; then
+    sed -i '' '/^ *print(f"ERROR: Unsupported host platform/ i\
+            if system == "Darwin":\
+                # Use default Apple toolchain (Clang) on macOS\
+                return True\
+\
+' sw/vgf-lib/scripts/build.py
+        fi
+        python sw/vgf-lib/scripts/build.py -j$(nproc)
+    fi
+
+    # emu layer
+    if [[ "${enable_emulation_layer}" -eq 1 ]]; then
+        pushd sw/emulation-layer
+        cmake -B build                                               \
+            -DGLSLANG_PATH=../../dependencies/glslang                \
+            -DSPIRV_CROSS_PATH=../../dependencies/SPIRV-Cross        \
+            -DSPIRV_HEADERS_PATH=../../dependencies/SPIRV-Headers    \
+            -DSPIRV_TOOLS_PATH=../../dependencies/SPIRV-Tools        \
+            -DVULKAN_HEADERS_PATH=../../dependencies/Vulkan-Headers
+        cmake --build build
+        popd
+    fi
+
+    popd
+}
+
+#setup_model_converter() $1
+# `"$manifest_dir"'

backends/arm/tosa_quant_utils.py

@@ -146,7 +146,15 @@ def quantize_value(self, x: torch.Tensor | float) -> Tensor:
         x = x.to(torch.float32)
         if self.per_channel:
             q_op = exir_ops.edge.quantized_decomposed.quantize_per_channel.default
-            args = (x, self.scale, self.zp, self.axis, self.qmin, self.qmax, self.dtype)
+            args = (
+                x,
+                torch.tensor(self.scale),
+                torch.tensor(self.zp),
+                self.axis,
+                self.qmin,
+                self.qmax,
+                self.dtype,
+            )
         else:
             q_op = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default
             args = (x, self.scale, self.zp, self.qmin, self.qmax, self.dtype)  # type: ignore[assignment]
@@ -162,8 +170,8 @@ def dequantize_value(self, qx: torch.Tensor) -> torch.Tensor:
             dq_op = exir_ops.edge.quantized_decomposed.dequantize_per_channel.default
             args = (
                 qx,
-                self.scale,
-                self.zp,
+                torch.tensor(self.scale),
+                torch.tensor(self.zp),
                 self.axis,
                 self.qmin,
                 self.qmax,

backends/cadence/aot/TARGETS

@@ -184,6 +184,21 @@ python_library(
     ],
 )
 
+python_library(
+    name = "program_builder",
+    srcs = [
+        "program_builder.py",
+    ],
+    typing = True,
+    deps = [
+        ":graph_builder",
+        "fbcode//caffe2:torch",
+        "fbcode//executorch/exir:lib",
+        "fbcode//executorch/exir:pass_base",
+        "fbcode//executorch/exir/verification:verifier",
+    ],
+)
+
 python_library(
     name = "fuse_ops",
     srcs = [
@@ -508,6 +523,7 @@ python_unittest(
         ":typing_stubs",
         ":ops_registrations",
         ":pass_utils",
+        ":program_builder",
         "//caffe2:torch",
         "//executorch/exir:memory",
         "//executorch/exir/dialects:lib",