Vulkan: support Linux/Windows desktop GPUs and opt-in wheel builds (#20138)

## Summary

We already provide good desktop support for NVIDIA GPUs through our CUDA
backend, which works well on both Linux and Windows. However, our Vulkan
backend only provides solid support for Android, leaving AMD GPUs
without sufficient support. That's a shame since Vulkan is well
supported on every operating system and with every major GPU
manufacturer.

This PR gets the Vulkan backend building and running correctly on Linux
and Windows desktop GPUs (NVIDIA, AMD, Intel), and adds an opt-in way to
build pre-built Vulkan binaries. It leaves everything Android related
the same so that we don't regress anything for that platform.

Most of the backend was already portable, So this is mostly build fixes,
a few small fixes for desktop GPUs, and packaging/CI plumbing.

fixes:#20140

## Changes

- Picks the right exception flag for MSVC vs GCC/Clang, finds glslc on
Windows, suppresses third-party (VMA) warnings on GCC/MSVC, and gives a
clear error if the Vulkan submodules aren't checked out.
- Compiles the cooperative-matrix shader. It needs a newer Vulkan target
than the default, so that one shader now builds against Vulkan 1.3.
- Fixes correctness on desktop GPUs. Turns on the device features the
shaders actually use (int16/int64/float64), which were never enabled
before; picks a discrete GPU instead of the first one found if one
exists (with an `ETVK_DEVICE_INDEX` override for multi-GPU machines);
avoids an invalid image-copy on compute-only queues; and fixes a
buffer-size check that compared the wrong units.
- Makes the `small_texture_limits` export option work as intended. It
was being silently dropped; now it round-trips so you can target GPUs
with smaller texture limits. Adds a small unit test for it.
- Adds opt-in packaging and CI. Behind the `EXECUTORCH_BUILD_VULKAN`
flag, the wheel build can include Vulkan; adds a real-GPU (NVIDIA) test
job and a Windows/MSVC build job. The default wheel and other backends
are untouched.

## Android Safety

Every change is gated so Android behaves exactly as before: build
differences are behind compiler/OS checks, the new export and runtime
options are opt-in and default to today's behavior, and the
device-selection / feature changes are based on what the GPU reports .
The existing SwiftShader CI job is unchanged.

## Testing

Built with the Vulkan SDK (glslc on PATH) and run on an NVIDIA A100
(driver 580.126.09, Vulkan 1.4.312):

```bash
# Build the backend + a runner (the Vulkan SDK's glslc must be on PATH)
cmake -B cmake-out -S . -GNinja -DCMAKE_BUILD_TYPE=Release \
    -DEXECUTORCH_BUILD_VULKAN=ON \
    -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \
    -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=ON

cmake --build cmake-out --target vulkan_backend executor_runner   # 402/402, 0 errors

# Export a model to Vulkan and run it on the GPU
python -m examples.vulkan.export -m mv2 -o .
./cmake-out/executor_runner --model_path mv2.pte
```

I ran a small fp32 model and an int8 model on the A100 and matched the
reference output (fp32 to 5 decimals, int8 to 4 decimals). The int8 run
exercises the integer-dot-product / int16 shaders that SwiftShader can't
run.
All the shaders compiled fine and the new unit tests added here passed.

I didn't test the windows build yet, so I'll be relying on CI for that.

## TODO

We also need to publish a Vulkan wheel to PyPI. The build supports it
(`EXECUTORCH_BUILD_VULKAN=1` + glslc), but we need a Vulkan entry added
to the shared `build-wheels-*.yml` workflows.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin

---------

Co-authored-by: Sicheng Stephen Jia <ssjia@meta.com>

Author: Reubend

.ci/scripts/setup-vulkan-linux-deps.sh

@@ -1,4 +1,3 @@
-
 #!/bin/bash
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
@@ -22,7 +21,7 @@ install_swiftshader() {
   tar -C "${_swiftshader_dir}" -xzf "${_tmp_archive}"
 
   export VK_ICD_FILENAMES="${_swiftshader_dir}/swiftshader/build/Linux/vk_swiftshader_icd.json"
-  export LD_LIBRARY_PATH="${_swiftshader_dir}/swiftshader/build/Linux/"
+  export LD_LIBRARY_PATH="${_swiftshader_dir}/swiftshader/build/Linux/:${LD_LIBRARY_PATH:-}"
   export ETVK_USING_SWIFTSHADER=1
 }
 
@@ -43,7 +42,125 @@ install_vulkan_sdk() {
   export PATH="${PATH}:${_vulkan_sdk_dir}/${VULKAN_SDK_VERSION}/x86_64/bin/"
 }
 
+_maybe_sudo() {
+  if [ "$(id -u)" -eq 0 ]; then
+    "$@"
+  else
+    sudo "$@"
+  fi
+}
+
+install_glslc() {
+  # The glslc shipped in the LunarG SDK is dynamically linked against a newer
+  # glibc/libstdc++ than the manylinux_2_28 / AlmaLinux 8 CUDA runner image
+  # provides (glibc 2.28), where it fails to load with "GLIBC_2.29 not found".
+  # conda-forge's shaderc is built against an old sysroot, runs there, and is
+  # recent enough for the GL_EXT_integer_dot_product / GL_KHR_cooperative_matrix
+  # extensions the Vulkan shaders use. Install it into an isolated prefix so the
+  # base conda env that builds ExecuTorch is left untouched, then put it on PATH.
+  _glslc_prefix=/tmp/shaderc
+  conda create -y -p "${_glslc_prefix}" -c conda-forge shaderc
+  export PATH="${_glslc_prefix}/bin:${PATH}"
+}
+
+install_vulkan_loader() {
+  # libvulkan.so.1 (the Khronos loader that volk dlopen()s at runtime) is not part
+  # of the NVIDIA driver and is absent from the CUDA builder image; vulkan-tools
+  # provides vulkaninfo for the device sanity check. Both ship as native el8 RPMs.
+  if command -v dnf >/dev/null 2>&1; then
+    _maybe_sudo dnf install -y vulkan-loader vulkan-tools
+  fi
+}
+
+_find_nvidia_vulkan_library() {
+  # NVIDIA implements its Vulkan ICD inside libGLX_nvidia.so.0. The NVIDIA
+  # container runtime mounts this library into the container (it is pulled from
+  # the driver's ldcache when NVIDIA_DRIVER_CAPABILITIES includes graphics/all),
+  # so prefer ldconfig and fall back to the usual mount locations.
+  local lib cand
+  lib="$(ldconfig -p 2>/dev/null | awk '/libGLX_nvidia\.so\.0/ {print $NF; exit}')"
+  if [ -z "${lib}" ]; then
+    for cand in /usr/lib64/libGLX_nvidia.so.0 \
+        /usr/lib/x86_64-linux-gnu/libGLX_nvidia.so.0 \
+        /usr/lib/libGLX_nvidia.so.0; do
+      if [ -e "${cand}" ]; then
+        lib="${cand}"
+        break
+      fi
+    done
+  fi
+  printf '%s' "${lib}"
+}
+
+_vulkan_has_real_device() {
+  # True if the loader enumerates a hardware GPU. vulkaninfo can exit non-zero
+  # for unrelated reasons (no display/WSI), so key off the reported deviceType.
+  command -v vulkaninfo >/dev/null 2>&1 || return 0
+  vulkaninfo --summary 2>/dev/null |
+    grep -qE 'PHYSICAL_DEVICE_TYPE_(DISCRETE|INTEGRATED|VIRTUAL)_GPU'
+}
+
+setup_real_gpu_icd() {
+  # Select a Vulkan ICD so the runtime exercises the real GPU when one is usable.
+  # Two quirks of the CUDA CI image make this non-trivial:
+  #   1. The NVIDIA container runtime mounts the driver's Vulkan library but does
+  #      not register its ICD manifest, so the loader never discovers the GPU on
+  #      its own. We synthesize the manifest and pin the loader to it.
+  #   2. Installing vulkan-loader/vulkan-tools pulls in mesa-vulkan-drivers,
+  #      which drop Intel/AMD/lavapipe manifests for absent hardware. lavapipe
+  #      fails vkCreateInstance on this image and, because the loader walks every
+  #      manifest in icd.d, that poisons device enumeration for the whole
+  #      process. Pinning VK_ICD_FILENAMES makes the loader ignore icd.d, so the
+  #      broken stubs cannot interfere.
+  local nvidia_lib
+  nvidia_lib="$(_find_nvidia_vulkan_library)"
+  if [ -n "${nvidia_lib}" ]; then
+    local icd=/tmp/nvidia_icd.json
+    cat >"${icd}" <<JSON
+{
+    "file_format_version": "1.0.0",
+    "ICD": {
+        "library_path": "${nvidia_lib}",
+        "api_version": "1.3.0"
+    }
+}
+JSON
+    export VK_ICD_FILENAMES="${icd}"
+    unset ETVK_USING_SWIFTSHADER || true
+    if _vulkan_has_real_device; then
+      echo "Real NVIDIA GPU selected; pinned Vulkan ICD to ${nvidia_lib}"
+      return
+    fi
+    echo "WARNING: ${nvidia_lib} present but no GPU enumerated; using SwiftShader."
+    # Surface why the NVIDIA driver did not enumerate (e.g. a missing dependency
+    # of libGLX_nvidia, or no render node) so the fallback is diagnosable in CI.
+    if command -v vulkaninfo >/dev/null 2>&1; then
+      echo "--- NVIDIA Vulkan ICD diagnostic ---"
+      VK_LOADER_DEBUG=warn vulkaninfo --summary 2>&1 | head -40 || true
+      echo "--- end diagnostic ---"
+    fi
+    unset VK_ICD_FILENAMES
+  else
+    echo "WARNING: no NVIDIA Vulkan driver library found; using SwiftShader."
+  fi
+  install_swiftshader
+}
+
 VULKAN_SDK_VERSION="1.4.321.1"
 
-install_swiftshader
-install_vulkan_sdk "${VULKAN_SDK_VERSION}"
+# The no-argument default installs SwiftShader so the existing CPU-runner CI is
+# unchanged. Pass "real-gpu" to prefer a real system ICD when one is present.
+case "${1:-swiftshader}" in
+  real-gpu)
+    # Do not download the LunarG SDK here: its prebuilt glslc cannot run on the
+    # old-glibc CUDA image. glslc comes from conda-forge and the loader from the
+    # system package manager instead.
+    install_vulkan_loader
+    install_glslc
+    setup_real_gpu_icd
+    ;;
+  swiftshader | *)
+    install_swiftshader
+    install_vulkan_sdk "${VULKAN_SDK_VERSION}"
+    ;;
+esac

.ci/scripts/setup-vulkan-windows-deps.ps1

@@ -0,0 +1,37 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Install glslc (the Vulkan shader compiler) on Windows via conda-forge's
+# shaderc package, and make sure it is on PATH. glslc is the only build-time
+# Vulkan dependency -- the Vulkan headers and the volk loader come from the
+# in-tree submodules -- so this avoids depending on the heavyweight LunarG SDK
+# installer. Requires conda to be available (the callers create/activate an env).
+
+$ErrorActionPreference = "Stop"
+
+Write-Host "Installing shaderc (provides glslc) from conda-forge..."
+conda install -y -c conda-forge shaderc
+if ($LASTEXITCODE -ne 0) {
+    Write-Error "Failed to install shaderc from conda-forge (exit ${LASTEXITCODE})"
+    exit 1
+}
+
+$glslc = Get-Command glslc -ErrorAction SilentlyContinue
+if (-not $glslc) {
+    Write-Error "glslc not found on PATH after installing shaderc"
+    exit 1
+}
+
+# Expose glslc to the current process and, when running as a GitHub Actions step,
+# to subsequent steps.
+$glslcDir = Split-Path -Parent $glslc.Source
+$env:PATH = "${glslcDir};${env:PATH}"
+if ($env:GITHUB_PATH) {
+    Add-Content -Path $env:GITHUB_PATH -Value $glslcDir
+}
+
+Write-Host "glslc available at $($glslc.Source)"
+& glslc --version

.ci/scripts/setup-windows-msvc-vulkan.ps1

@@ -0,0 +1,51 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Build-validation for the Vulkan backend under MSVC on Windows. Mirrors
+# setup-windows-msvc.ps1 but installs glslc (the Vulkan shader compiler) and
+# configures/builds the vulkan_backend target. This is a bring-up job: it exists
+# to surface MSVC portability issues in the Vulkan/volk/VMA code, so it may need
+# iteration.
+
+conda create --yes --quiet -n et python=3.12
+conda activate et
+
+# Install cmake
+conda install -y cmake
+
+# Activate the VS environment - this is required for MSVC to work.
+& "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\Common7\Tools\Launch-VsDevShell.ps1" -Arch amd64
+
+# Install glslc (via conda-forge shaderc) and put it on PATH in this process.
+.ci/scripts/setup-vulkan-windows-deps.ps1
+
+# Install CI requirements
+pip install -r .ci/docker/requirements-ci.txt
+
+$buildDir = "cmake-out-vulkan"
+if (Test-Path -Path $buildDir) {
+    Remove-Item -Path $buildDir -Recurse -Force
+}
+New-Item -Path $buildDir -ItemType Directory
+
+cmake -S . -B $buildDir `
+    -DCMAKE_BUILD_TYPE=Release `
+    -DEXECUTORCH_BUILD_VULKAN=ON `
+    -DPYTHON_EXECUTABLE=python
+
+if ($LASTEXITCODE -ne 0) {
+    Write-Host "CMake configuration failed. Exit code: $LASTEXITCODE."
+    exit $LASTEXITCODE
+}
+
+cmake --build $buildDir --config Release --target vulkan_backend -j16
+
+if ($LASTEXITCODE -ne 0) {
+    Write-Host "Vulkan backend MSVC build failed. Exit code: $LASTEXITCODE."
+    exit $LASTEXITCODE
+}
+
+Write-Host "Vulkan backend MSVC build completed successfully!"

.ci/scripts/test_backend.sh

@@ -51,8 +51,15 @@ if [[ "$FLOW" == *qnn* ]]; then
 fi
 
 if [[ "$FLOW" == *vulkan* ]]; then
-    # Setup swiftshader and Vulkan SDK which are required to build the Vulkan delegate.
-    source .ci/scripts/setup-vulkan-linux-deps.sh
+    # Setup the Vulkan SDK and select an ICD: use the real system GPU ICD when one
+    # is present (real-GPU runner), otherwise fall back to SwiftShader (CPU
+    # runner). The Vulkan loader searches both standard ICD directories.
+    if ls /etc/vulkan/icd.d/*.json /usr/share/vulkan/icd.d/*.json \
+        >/dev/null 2>&1; then
+        source .ci/scripts/setup-vulkan-linux-deps.sh "real-gpu"
+    else
+        source .ci/scripts/setup-vulkan-linux-deps.sh "swiftshader"
+    fi
 
     EXTRA_BUILD_ARGS+=" -DEXECUTORCH_BUILD_VULKAN=ON"
 fi

.ci/scripts/wheel/pre_build_script.sh

@@ -83,3 +83,34 @@ if [[ "$(uname -s)" == "Linux" && "$(uname -m)" == "x86_64" ]]; then
   echo "QNN_SDK_ROOT=${QNN_SDK_ROOT}" >> "${GITHUB_ENV}"
   echo "QNN SDK downloaded to ${QNN_SDK_ROOT}"
 fi
+
+# Provision the Vulkan SDK (glslc) and submodules ONLY when explicitly requested
+# via EXECUTORCH_BUILD_VULKAN. The default wheel build leaves this unset, so it
+# does no extra work (no submodule fetch, no SDK download) and is unaffected.
+if [[ "${EXECUTORCH_BUILD_VULKAN:-0}" != "0" \
+      && "${EXECUTORCH_BUILD_VULKAN:-OFF}" != "OFF" ]]; then
+  echo "Initializing Vulkan backend third-party submodules..."
+  VULKAN_SUBMODULES=(
+    backends/vulkan/third-party/Vulkan-Headers
+    backends/vulkan/third-party/volk
+    backends/vulkan/third-party/VulkanMemoryAllocator
+  )
+  if [[ $UNAME_S == *"MINGW"* || $UNAME_S == *"MSYS"* ]]; then
+    git -c http.sslBackend=openssl submodule update --init "${VULKAN_SUBMODULES[@]}"
+    echo "Installing Vulkan SDK for Windows wheel build..."
+    powershell -ExecutionPolicy Bypass -File .ci/scripts/setup-vulkan-windows-deps.ps1
+  else
+    git submodule update --init "${VULKAN_SUBMODULES[@]}"
+    # Install glslc from conda-forge rather than the LunarG SDK: the manylinux
+    # wheel image uses an old glibc where the SDK's prebuilt glslc cannot run
+    # ("GLIBC_2.29 not found"). conda-forge's shaderc is built against an old
+    # sysroot and runs there. Vulkan headers come from the submodules above and
+    # volk dlopen()s the loader at runtime, so only glslc is needed to build.
+    echo "Installing glslc (conda-forge shaderc) for Linux wheel build..."
+    _glslc_prefix="${HOME}/.shaderc"
+    conda create -y -p "${_glslc_prefix}" -c conda-forge shaderc
+    export PATH="${_glslc_prefix}/bin:${PATH}"
+    echo "${_glslc_prefix}/bin" >> "${GITHUB_PATH}"
+    echo "glslc installed: $(command -v glslc)"
+  fi
+fi

.ci/scripts/wheel/test_linux.py

@@ -31,6 +31,13 @@
         ), f"OpenvinoBackend not found in registered backends: {registered}"
         print("✓ OpenvinoBackend is registered")
 
+        # Vulkan backend is optional: only present when the wheel was built with
+        # EXECUTORCH_BUILD_VULKAN=1 and the Vulkan SDK (glslc) was available.
+        if "VulkanBackend" in registered:
+            print("✓ VulkanBackend is registered")
+        else:
+            print("⚠ VulkanBackend not registered (expected for the default wheel)")
+
     test_base.run_tests(
         model_tests=[
             test_base.ModelTest(

.ci/scripts/wheel/test_windows.py

@@ -5,6 +5,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import platform
 from typing import List
 
 import torch
@@ -15,6 +16,7 @@
 from executorch.examples.xnnpack.quantization.utils import quantize as quantize_xnn
 from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
 from executorch.extension.pybindings.portable_lib import (
+    _get_registered_backend_names,
     _load_for_executorch_from_buffer,
 )
 from test_base import ModelTest
@@ -63,6 +65,15 @@ def run_tests(model_tests: List[ModelTest]) -> None:
 
 
 if __name__ == "__main__":
+    if platform.system() == "Windows":
+        registered = _get_registered_backend_names()
+        # Vulkan backend is optional: only present when the wheel was built with
+        # EXECUTORCH_BUILD_VULKAN=1 and the Vulkan SDK (glslc) was available.
+        if "VulkanBackend" in registered:
+            print("✓ VulkanBackend is registered")
+        else:
+            print("⚠ VulkanBackend not registered (expected for the default wheel)")
+
     run_tests(
         model_tests=[
             ModelTest(

.github/workflows/test-backend-vulkan.yml

@@ -17,6 +17,8 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  # Default coverage: builds + runs on SwiftShader (software Vulkan) on CPU
+  # runners. Runs on every PR and nightly.
   test-vulkan:
     uses: ./.github/workflows/_test_backend.yml
     with:
@@ -28,3 +30,6 @@ jobs:
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       timeout: 120
       run-linux: true
+
+  # Real-GPU (NVIDIA) and Windows MSVC coverage live in vulkan.yml, which gates
+  # those scarce/expensive runners behind path filtering.