Vulkan: support Linux/Windows desktop GPUs and opt-in wheel builds

The Vulkan backend was developed with a focus on Android GPUs. This change
makes it build and run correctly on Linux and Windows desktop discrete GPUs
(NVIDIA/AMD/Intel) and adds an opt-in path to produce pre-built Vulkan binaries,
without changing behavior on Android.

The architecture was already largely portable (volk loader, headless compute,
optional extensions gated by availability, staging-buffer transfers), so the
work here is concentrated in build portability, a few discrete-GPU correctness
fixes, and CI/packaging plumbing. Every change is Android-safe by construction:
build divergence sits behind compile-time compiler/OS guards, AOT/policy changes
are opt-in and default to current behavior, and runtime device-behavior changes
key off queried capabilities so they resolve identically on a single-GPU
Adreno/Mali device.

Suggested review order:

1. Build portability -- backends/vulkan/CMakeLists.txt (per-compiler exception
   flag, submodule check), cmake/ShaderLibrary.cmake (glslc discovery, graceful
   skip in wheel builds), and runtime/vk_api/memory/vma_api.h (GCC/MSVC warning
   suppression alongside the existing clang block).

2. Shader compilation -- runtime/graph/ops/glsl/coopmat_mm.yaml targets Vulkan
   1.3 so GL_KHR_cooperative_matrix compiles, plus a NameError-safety fix in
   runtime/gen_vulkan_spv.py.

3. Runtime correctness on discrete GPUs -- vk_api/Adapter.cpp enables the
   shaderInt16/Int64/Float64 features the shaders already use; vk_api/Runtime.cpp
   prefers a real GPU over software/integrated devices (with an ETVK_DEVICE_INDEX
   override); api/Context.cpp guards blit against compute-only queues; and
   api/containers/Tensor.cpp compares the storage-buffer size in bytes.

4. Ahead-of-time -- vulkan_preprocess.py, partitioner/vulkan_partitioner.py, and
   utils.py wire the previously-dropped small_texture_limits option through the
   compile-spec round-trip; test/test_vulkan_compile_options.py covers it.

5. Distribution and CI, all gated behind EXECUTORCH_BUILD_VULKAN so default
   wheels and other backends are untouched -- tools/cmake/preset/pybind.cmake,
   setup.py, .ci/scripts/wheel/*, .ci/scripts/setup-vulkan-*.{sh,ps1},
   .github/workflows/test-backend-vulkan.yml (adds a real-GPU job), and
   .github/workflows/vulkan-windows.yml (MSVC build validation).

Tested by building the backend with the Vulkan SDK and running fp32 and int8
models on an NVIDIA A100: outputs match the reference and all shaders compile
(cooperative-matrix as SPIR-V 1.6). The existing SwiftShader CI path is
unchanged.

This change was authored with Claude.

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.
@@ -43,7 +42,35 @@ install_vulkan_sdk() {
   export PATH="${PATH}:${_vulkan_sdk_dir}/${VULKAN_SDK_VERSION}/x86_64/bin/"
 }
 
+setup_real_gpu_icd() {
+  # On a real-GPU runner the system Vulkan ICD is installed by the GPU driver.
+  # The loader searches both /etc/vulkan/icd.d and /usr/share/vulkan/icd.d, so
+  # check both. If a system ICD is present, do NOT use SwiftShader so the real
+  # device (and its fp16/int16/dot-product shader variants) is exercised. Fall
+  # back to SwiftShader if no system ICD is found so the job stays green either
+  # way.
+  if ls /etc/vulkan/icd.d/*.json /usr/share/vulkan/icd.d/*.json \
+      >/dev/null 2>&1; then
+    echo "System Vulkan ICD(s) detected:"
+    ls /etc/vulkan/icd.d/*.json /usr/share/vulkan/icd.d/*.json 2>/dev/null
+    unset ETVK_USING_SWIFTSHADER || true
+  else
+    echo "WARNING: no system Vulkan ICD found; using SwiftShader."
+    install_swiftshader
+  fi
+}
+
 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)
+    install_vulkan_sdk "${VULKAN_SDK_VERSION}"
+    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

@@ -69,3 +69,37 @@ 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[@]}"
+    echo "Installing Vulkan SDK for Linux wheel build..."
+    VULKAN_SDK_VERSION="1.4.341.1"
+    _vulkan_sdk_url="https://sdk.lunarg.com/sdk/download/${VULKAN_SDK_VERSION}/linux/vulkansdk-linux-x86_64-${VULKAN_SDK_VERSION}.tar.xz"
+    _vulkan_sdk_dir="${HOME}/.vulkan-sdk/${VULKAN_SDK_VERSION}"
+    mkdir -p "${_vulkan_sdk_dir}"
+    curl --silent --show-error --location --fail --retry 3 --retry-all-errors \
+      --output /tmp/vulkansdk.tar.xz "${_vulkan_sdk_url}"
+    tar -C "${_vulkan_sdk_dir}" -xJf /tmp/vulkansdk.tar.xz
+    VULKAN_SDK="${_vulkan_sdk_dir}/${VULKAN_SDK_VERSION}/x86_64"
+    export VULKAN_SDK
+    export PATH="${VULKAN_SDK}/bin:${PATH}"
+    echo "VULKAN_SDK=${VULKAN_SDK}" >> "${GITHUB_ENV}"
+    echo "${VULKAN_SDK}/bin" >> "${GITHUB_PATH}"
+    echo "Vulkan SDK installed to ${VULKAN_SDK}"
+  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,47 @@ jobs:
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       timeout: 120
       run-linux: true
+
+  # Real-GPU coverage on an NVIDIA runner (exercises the fp16/int16/dot-product
+  # shader variants that SwiftShader cannot). Standalone job (does not go through
+  # _test_backend.yml) so the GPU is attached via gpu-arch-type, mirroring
+  # cuda.yml. Not run on pull_request to avoid GPU cost and to keep it from
+  # blocking PR merges; runs on nightly/push/dispatch.
+  test-vulkan-real-gpu:
+    if: ${{ github.event_name != 'pull_request' }}
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    permissions:
+      id-token: write
+      contents: read
+    with:
+      timeout: 90
+      runner: linux.g5.4xlarge.nvidia.gpu
+      gpu-arch-type: cuda
+      gpu-arch-version: "12.6"
+      use-custom-docker-registry: false
+      submodules: recursive
+      ref: ${{ github.sha }}
+      script: |
+        set -eux
+
+        # Install the Vulkan SDK (glslc) and select a real system ICD. The NVIDIA
+        # driver on this runner provides the ICD; install the loader as well.
+        # NOTE: first-run check - inspect the vulkaninfo output below to confirm a
+        # real NVIDIA device is selected (not llvmpipe/SwiftShader). If no system
+        # ICD is present, setup-vulkan-linux-deps.sh falls back to SwiftShader.
+        sudo apt-get update && sudo apt-get install -y libvulkan1 vulkan-tools || true
+        source .ci/scripts/setup-vulkan-linux-deps.sh real-gpu
+        vulkaninfo --summary || true
+
+        PYTHON_EXECUTABLE=python ./install_executorch.sh
+
+        cmake -DCMAKE_BUILD_TYPE=Release \
+            -DEXECUTORCH_BUILD_VULKAN=ON \
+            -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \
+            -DPYTHON_EXECUTABLE=python \
+            -Bcmake-out .
+        cmake --build cmake-out -j4 --target executor_runner
+
+        # Export a model to the Vulkan backend and run it on the GPU.
+        python -m examples.vulkan.export -m mv2 -o .
+        ./cmake-out/executor_runner --model_path mv2.pte

.github/workflows/vulkan-windows.yml

@@ -0,0 +1,48 @@
+name: Test Vulkan Backend Windows Build
+
+# Build-validation for the Vulkan backend under MSVC on Windows. This is a
+# bring-up job (no GPU): it confirms the backend configures and compiles with
+# MSVC. Real-GPU Windows E2E is a follow-up once a Windows Vulkan GPU runner is
+# available. Path-filtered and not part of the required PR checks so it can be
+# iterated on without blocking unrelated work.
+
+on:
+  push:
+    branches:
+      - main
+      - release/*
+    tags:
+      - ciflow/nightly/*
+  pull_request:
+    paths:
+      - backends/vulkan/**
+      - .ci/scripts/setup-vulkan-windows-deps.ps1
+      - .ci/scripts/setup-windows-msvc-vulkan.ps1
+      - .github/workflows/vulkan-windows.yml
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
+  cancel-in-progress: true
+
+permissions:
+  contents: read
+
+jobs:
+  build-vulkan-windows-msvc:
+    name: build-vulkan-windows-msvc
+    uses: pytorch/test-infra/.github/workflows/windows_job.yml@main
+    with:
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      timeout: 90
+      script: |
+        git config --global http.sslBackend openssl
+        git submodule update --init backends/vulkan/third-party/Vulkan-Headers backends/vulkan/third-party/volk backends/vulkan/third-party/VulkanMemoryAllocator
+        git submodule update --init
+        conda init powershell
+        powershell -Command "& {
+          Set-PSDebug -Trace 1
+          \$ErrorActionPreference = 'Stop'
+          \$PSNativeCommandUseErrorActionPreference = \$true
+          .ci/scripts/setup-windows-msvc-vulkan.ps1
+        }"

backends/vulkan/CMakeLists.txt

@@ -41,6 +41,24 @@ set(VULKAN_HEADERS_PATH ${VULKAN_THIRD_PARTY_PATH}/Vulkan-Headers)
 set(VOLK_PATH ${VULKAN_THIRD_PARTY_PATH}/volk)
 set(VMA_PATH ${VULKAN_THIRD_PARTY_PATH}/VulkanMemoryAllocator)
 
+# These third-party dependencies are git submodules. They are not part of the
+# default submodule set checked out by install_executorch.py, so fail early with
+# an actionable message rather than a confusing missing-header error.
+if(NOT EXISTS "${VOLK_PATH}/volk.c"
+   OR NOT EXISTS "${VULKAN_HEADERS_PATH}/include/vulkan/vulkan.h"
+   OR NOT EXISTS "${VMA_PATH}/include/vk_mem_alloc.h"
+)
+  message(
+    FATAL_ERROR
+      "The Vulkan backend third-party submodules are missing. "
+      "Run the following from the repository root:\n"
+      "  git submodule update --init "
+      "backends/vulkan/third-party/Vulkan-Headers "
+      "backends/vulkan/third-party/volk "
+      "backends/vulkan/third-party/VulkanMemoryAllocator"
+  )
+endif()
+
 set(COMMON_INCLUDES
     $<BUILD_INTERFACE:${EXECUTORCH_ROOT}/..>
     $<BUILD_INTERFACE:${VULKAN_HEADERS_PATH}/include>
@@ -49,7 +67,11 @@ set(COMMON_INCLUDES
 
 # Compile settings
 
-set(VULKAN_CXX_FLAGS "-fexceptions")
+# Exceptions are required: the vk_api layer throws on Vulkan errors (see
+# vk_api/Exception.h). MSVC does not understand -fexceptions and enables C++
+# exceptions via /EHsc instead, so select the flag per compiler.
+set(VULKAN_CXX_FLAGS "$<$<NOT:$<CXX_COMPILER_ID:MSVC>>:-fexceptions>")
+list(APPEND VULKAN_CXX_FLAGS "$<$<CXX_COMPILER_ID:MSVC>:/EHsc>")
 list(APPEND VULKAN_CXX_FLAGS "-DUSE_VULKAN_WRAPPER")
 list(APPEND VULKAN_CXX_FLAGS "-DUSE_VULKAN_VOLK")