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

[Reubend]

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):

# 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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 1 Unclassified Failure

As of commit 1904c5c with merge base e285edf ():

NEW FAILURES - The following jobs have failed:

• pull / unittest / macos / macos-job (gh)
  export/tests/test_target_recipes.py::TestTargetRecipes::test_vit_model
• trunk / test-qnn-optimum-model (fp32, swin) / linux-job (gh)
  RuntimeError: Command docker exec -t 781707ff3b63ca2e2545875cbb4ac72c845a7edcab506b656bb15211e36ee3f0 /exec failed with exit code 1

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• trunk / test-models-macos-cpu (ic4, xnnpack-quantization-delegation) / macos-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: Reubend / name: Reuben Dunn (427011e)

[Reubend]

I took a look at the 3 test failures that the CI found, and I don't think that they're caused by any of my changes here. Please let me know if you find anything that indicates differently.

[SS-JIA]

ack. Taking a look soon.

[SS-JIA]

⏺ PR #20138 — changes by theme

  Theme 1: Build portability (make it compile off-Android)

  - backends/vulkan/CMakeLists.txt — per-compiler exception flag (/EHsc MSVC vs -fexceptions); submodule-existence check w/ actionable error.
  - backends/vulkan/cmake/ShaderLibrary.cmake — find glslc via VULKAN_SDK HINTS (Windows); graceful skip in wheel builds.
  - backends/vulkan/runtime/vk_api/memory/vma_api.h — warning suppression for GCC + MSVC (was clang-only).
  - backends/vulkan/runtime/gen_vulkan_spv.py — NameError guard when cache_dir is None.

  Theme 2: Shader compilation

  - backends/vulkan/runtime/graph/ops/glsl/coopmat_mm.yaml — target Vulkan 1.3 so GL_KHR_cooperative_matrix (needs SPIR-V 1.6) compiles.

  Theme 3: Runtime correctness on desktop GPUs

  - vk_api/Adapter.cpp — enable int16/int64/float64 device features shaders already use.
  - vk_api/Runtime.cpp — prefer discrete GPU over software/integrated; ETVK_DEVICE_INDEX override.
  - api/Context.cpp — guard blit against compute-only queues.
  - api/containers/Tensor.cpp — fix buffer-size check (bytes vs numel).

  Theme 4: AOT compile-option round-trip fix

  - vulkan_preprocess.py — deserialize small_texture_limits + skip_memory_planning (were silently dropped); fix tuple-mutation bug.
  - partitioner/vulkan_partitioner.py — wire small_texture_limits into partition texture limits.
  - utils.py — add SMALL_TEXTURE_LIMITS constant.
  - test/test_vulkan_compile_options.py — regression test for round-trip.

  Theme 5: Packaging (opt-in wheel)

  - setup.py — build vulkan_backend target when EXECUTORCH_BUILD_VULKAN.
  - tools/cmake/preset/pybind.cmake — enable Vulkan in Linux/Windows wheel if flag + glslc present.

  Theme 6: CI infra

  - .github/workflows/test-backend-vulkan.yml — new real-GPU (NVIDIA) job.
  - .github/workflows/vulkan-windows.yml — new MSVC build job.
  - .ci/scripts/test_backend.sh — pick real-gpu vs swiftshader ICD.
  - .ci/scripts/setup-vulkan-linux-deps.sh — arg-based SwiftShader/real-GPU setup.
  - .ci/scripts/setup-vulkan-windows-deps.ps1, setup-windows-msvc-vulkan.ps1 — Windows glslc + build.
  - .ci/scripts/wheel/pre_build_script.sh — provision Vulkan SDK/submodules when flag set.
  - .ci/scripts/wheel/test_linux.py, test_windows.py — check VulkanBackend registered.

All changes from Theme 1 through 5 look super solid! Thanks for putting together those fixes.

For Theme 6, just have a small request regarding the new CI tests added. Overall, the PR looks great!

I'll also import this PR as a diff just so that these changes can run against our Meta internal CI flows. Note that this diff won't be landed, I will abandon and unlink it from the PR once we get the CI signals.

[SS-JIA]

Potentially valid point raised by Claude:

  Context.cpp blit guard — ⚠️  one real question. Guard allows GRAPHICS | TRANSFER. But per Vulkan spec vkCmdBlitImage requires GRAPHICS specifically — transfer queues do not support blit.
  Compute+transfer-without-graphics queue passes check but blit still invalid usage. In practice selected compute queue on desktop usually has graphics, so works. Still: check should likely be
  VK_QUEUE_GRAPHICS_BIT only. Strictly better than no guard (today), but the TRANSFER_BIT could let an invalid case slip silently.

If true, seems that only GRAPHICS_BIT should be checked.

[SS-JIA]

optional, but with the updated behaviour it seems First is no longer accurate. Would recommend renaming to Auto, perhaps.

[SS-JIA]

I would recommend creating a new vulkan.yml workflow file and just putting all vulkan CI jobs that require special runners there (i.e. build-vulkan-windows-msvc and test-vulkan-real-gpu, which perhaps should be renamed to test-vulkan-nvidia), for example:

# .github/workflows/vulkan.yml
name: Test Vulkan Backend (specialized runners)

on:
  push:
    branches: [main, release/*]
    tags: [ciflow/nightly/*]
  pull_request:
    paths:
      - .github/workflows/vulkan.yml
      - backends/vulkan/**
      - .ci/scripts/setup-vulkan-linux-deps.sh
      - .ci/scripts/setup-vulkan-windows-deps.ps1
      - .ci/scripts/setup-windows-msvc-vulkan.ps1
  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:
  changed-files:
    name: Get changed files
    uses: ./.github/workflows/_get-changed-files.yml
    with:
      include-push-diff: true # so push commits can also be path-filtered

  run-decision:
    name: CI run decision
    uses: ./.github/workflows/_ci-run-decision.yml

  test-vulkan-nvidia:
    needs: [changed-files, run-decision]
    # Path-filtered: skip commits that don't touch Vulkan-relevant paths,
    # except on sampled full runs (see _ci-run-decision.yml).
    if: |
      github.event_name != 'pull_request' && (
        contains(needs.changed-files.outputs.changed-files, 'backends/vulkan/') ||
        contains(needs.changed-files.outputs.changed-files, '.ci/scripts/setup-vulkan-linux-deps.sh') ||
        contains(needs.changed-files.outputs.changed-files, '.github/workflows/vulkan.yml') ||
        needs.run-decision.outputs.is-full-run == 'true'
      )
    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
    # ... (existing real-gpu config)

  build-vulkan-windows-msvc:
    needs: [changed-files, run-decision]
    if: |
      contains(needs.changed-files.outputs.changed-files, 'backends/vulkan/') ||
      contains(needs.changed-files.outputs.changed-files, '.ci/scripts/setup-vulkan-windows-deps.ps1') ||
      contains(needs.changed-files.outputs.changed-files, '.ci/scripts/setup-windows-msvc-vulkan.ps1') ||
      contains(needs.changed-files.outputs.changed-files, '.github/workflows/vulkan.yml') ||
      needs.run-decision.outputs.is-full-run == 'true'
    uses: pytorch/test-infra/.github/workflows/windows_job.yml@main
    # ... (existing windows config)

Also on the topic of test-vulkan-real-gpu / test-vulkan-nvidia, I think we can also consider:

1. Have it replicate the workflow of both test-vulkan-models-linux and test-vulkan-operators-linux for full coverage
2. Remove the github.event_name != 'pull_request' check, since the changed-files mechanism + on:pull_request:paths mechanism should drastically reduce the frequency at which this job is triggered.

[meta-codesync]

@SS-JIA has imported this pull request. If you are a Meta employee, you can view this in D108371716.