torchvulkan 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (69) hide show
  1. torchvulkan-0.1.0/.github/actions/setup-toolchain/action.yml +90 -0
  2. torchvulkan-0.1.0/.github/scripts/install-toolchain-linux.sh +27 -0
  3. torchvulkan-0.1.0/.github/workflows/release.yaml +98 -0
  4. torchvulkan-0.1.0/.github/workflows/test_torchvulkan.yml +60 -0
  5. torchvulkan-0.1.0/.gitignore +42 -0
  6. torchvulkan-0.1.0/CMakeLists.txt +28 -0
  7. torchvulkan-0.1.0/LICENSE +21 -0
  8. torchvulkan-0.1.0/PKG-INFO +200 -0
  9. torchvulkan-0.1.0/README.md +184 -0
  10. torchvulkan-0.1.0/benchmark/mnist_benchmark.py +276 -0
  11. torchvulkan-0.1.0/examples/main.py +10 -0
  12. torchvulkan-0.1.0/examples/mnist.py +62 -0
  13. torchvulkan-0.1.0/extern/vma/vk_mem_alloc.h +19845 -0
  14. torchvulkan-0.1.0/extern/volk/volk.c +4004 -0
  15. torchvulkan-0.1.0/extern/volk/volk.h +3267 -0
  16. torchvulkan-0.1.0/media/benchmark-cuda.png +0 -0
  17. torchvulkan-0.1.0/media/benchmark-mps.png +0 -0
  18. torchvulkan-0.1.0/pyproject.toml +54 -0
  19. torchvulkan-0.1.0/requirements.txt +1 -0
  20. torchvulkan-0.1.0/scripts/build_shaders.py +211 -0
  21. torchvulkan-0.1.0/scripts/config.py +106 -0
  22. torchvulkan-0.1.0/scripts/pack_shaders.py +214 -0
  23. torchvulkan-0.1.0/src/api/backend.cpp +34 -0
  24. torchvulkan-0.1.0/src/api/guard.cpp +38 -0
  25. torchvulkan-0.1.0/src/api/hook.cpp +35 -0
  26. torchvulkan-0.1.0/src/api/ops/binary.cpp +360 -0
  27. torchvulkan-0.1.0/src/api/ops/binary.h +76 -0
  28. torchvulkan-0.1.0/src/api/ops/factory.cpp +511 -0
  29. torchvulkan-0.1.0/src/api/ops/factory.h +84 -0
  30. torchvulkan-0.1.0/src/api/ops/helpers.h +196 -0
  31. torchvulkan-0.1.0/src/api/ops/matmul.cpp +500 -0
  32. torchvulkan-0.1.0/src/api/ops/matmul.h +70 -0
  33. torchvulkan-0.1.0/src/api/ops/nllloss.cpp +150 -0
  34. torchvulkan-0.1.0/src/api/ops/nllloss.h +28 -0
  35. torchvulkan-0.1.0/src/api/ops/reduce.cpp +175 -0
  36. torchvulkan-0.1.0/src/api/ops/reduce.h +61 -0
  37. torchvulkan-0.1.0/src/api/ops/softmax.cpp +26 -0
  38. torchvulkan-0.1.0/src/api/ops/softmax.h +14 -0
  39. torchvulkan-0.1.0/src/api/ops/unary.cpp +216 -0
  40. torchvulkan-0.1.0/src/api/ops/unary.h +39 -0
  41. torchvulkan-0.1.0/src/api/ops.cpp +122 -0
  42. torchvulkan-0.1.0/src/shaders/CMakeLists.txt +44 -0
  43. torchvulkan-0.1.0/src/shaders/binaryop.slang.j2 +183 -0
  44. torchvulkan-0.1.0/src/shaders/cast.slang.j2 +54 -0
  45. torchvulkan-0.1.0/src/shaders/copy.slang.j2 +53 -0
  46. torchvulkan-0.1.0/src/shaders/fill.slang.j2 +65 -0
  47. torchvulkan-0.1.0/src/shaders/helpers/int_divider.slang.j2 +10 -0
  48. torchvulkan-0.1.0/src/shaders/helpers/vulkan_arithmetic.slang.j2 +170 -0
  49. torchvulkan-0.1.0/src/shaders/matmul_coop.slang.j2 +780 -0
  50. torchvulkan-0.1.0/src/shaders/matmul_simd.slang.j2 +317 -0
  51. torchvulkan-0.1.0/src/shaders/nllloss.slang.j2 +80 -0
  52. torchvulkan-0.1.0/src/shaders/reduce.slang.j2 +80 -0
  53. torchvulkan-0.1.0/src/shaders/unaryop.slang.j2 +138 -0
  54. torchvulkan-0.1.0/src/vulkan/allocator.cpp +215 -0
  55. torchvulkan-0.1.0/src/vulkan/allocator.h +113 -0
  56. torchvulkan-0.1.0/src/vulkan/cache.cpp +487 -0
  57. torchvulkan-0.1.0/src/vulkan/cache.h +81 -0
  58. torchvulkan-0.1.0/src/vulkan/device_context.cpp +55 -0
  59. torchvulkan-0.1.0/src/vulkan/device_context.h +42 -0
  60. torchvulkan-0.1.0/src/vulkan/memory.cpp +48 -0
  61. torchvulkan-0.1.0/src/vulkan/memory.h +36 -0
  62. torchvulkan-0.1.0/src/vulkan/shader.cpp +74 -0
  63. torchvulkan-0.1.0/src/vulkan/shader.h +34 -0
  64. torchvulkan-0.1.0/src/vulkan/vulkan_context.cpp +384 -0
  65. torchvulkan-0.1.0/src/vulkan/vulkan_context.h +43 -0
  66. torchvulkan-0.1.0/tests/test_pytorch_ops.py +185 -0
  67. torchvulkan-0.1.0/torchvulkan/__init__.py +31 -0
  68. torchvulkan-0.1.0/torchvulkan/_version.py +24 -0
  69. torchvulkan-0.1.0/tox.ini +16 -0
@@ -0,0 +1,90 @@
1
+ name: Setup torchvulkan toolchain
2
+ description: >
3
+ Install the build/runtime toolchain for the current runner OS: the Slang
4
+ compiler, a Vulkan driver + SDK headers, and spirv-cross. Exports the needed
5
+ PATH / env entries so later steps in the job can build and run torchvulkan.
6
+
7
+ inputs:
8
+ slang-version:
9
+ description: Slang release version to install.
10
+ default: "2026.8.1"
11
+ vulkan-sdk-version:
12
+ description: Vulkan SDK version (Windows only).
13
+ default: "1.3.296.0"
14
+
15
+ runs:
16
+ using: composite
17
+ steps:
18
+ - name: Resolve Slang arch
19
+ shell: bash
20
+ run: |
21
+ case "${{ runner.os }}-${{ runner.arch }}" in
22
+ Linux-X64) echo "SLANG_ARCH=linux-x86_64" >> "$GITHUB_ENV" ;;
23
+ macOS-ARM64) echo "SLANG_ARCH=macos-aarch64" >> "$GITHUB_ENV" ;;
24
+ macOS-X64) echo "SLANG_ARCH=macos-x86_64" >> "$GITHUB_ENV" ;;
25
+ Windows-X64) echo "SLANG_ARCH=windows-x86_64" >> "$GITHUB_ENV" ;;
26
+ *) echo "::error::Unsupported runner ${{ runner.os }}-${{ runner.arch }}"; exit 1 ;;
27
+ esac
28
+
29
+ - name: Install Slang (macOS & Linux)
30
+ if: runner.os != 'Windows'
31
+ shell: bash
32
+ run: |
33
+ curl -L -o slang.zip "https://github.com/shader-slang/slang/releases/download/v${{ inputs.slang-version }}/slang-${{ inputs.slang-version }}-${SLANG_ARCH}.zip"
34
+ unzip slang.zip -d "$RUNNER_TEMP/slang"
35
+ echo "$RUNNER_TEMP/slang/bin" >> "$GITHUB_PATH"
36
+ echo "SLANG_DIR=$RUNNER_TEMP/slang/bin" >> "$GITHUB_ENV"
37
+
38
+ - name: Install Slang (Windows)
39
+ if: runner.os == 'Windows'
40
+ shell: pwsh
41
+ run: |
42
+ $zip = "slang-${{ inputs.slang-version }}-$env:SLANG_ARCH.zip"
43
+ Invoke-WebRequest -Uri "https://github.com/shader-slang/slang/releases/download/v${{ inputs.slang-version }}/$zip" -OutFile $zip
44
+ Expand-Archive -Path $zip -DestinationPath C:\slang
45
+ "C:\slang\bin" | Out-File -FilePath $env:GITHUB_PATH -Append
46
+ "SLANG_DIR=C:\slang\bin" | Out-File -FilePath $env:GITHUB_ENV -Append
47
+
48
+ - name: Install Vulkan (Linux)
49
+ if: runner.os == 'Linux'
50
+ shell: bash
51
+ run: |
52
+ wget -qO- https://packages.lunarg.com/lunarg-signing-key-pub.asc | sudo tee /etc/apt/trusted.gpg.d/lunarg.asc
53
+ sudo wget -qO /etc/apt/sources.list.d/lunarg-vulkan-$(lsb_release -cs).list http://packages.lunarg.com/vulkan/lunarg-vulkan-$(lsb_release -cs).list
54
+ sudo apt-get update
55
+ sudo apt-get install -y vulkan-sdk mesa-vulkan-drivers
56
+ vulkaninfo --summary
57
+
58
+ - name: Install Vulkan (macOS)
59
+ if: runner.os == 'macOS'
60
+ shell: bash
61
+ run: |
62
+ brew install molten-vk vulkan-tools ninja libxshmfence spirv-cross
63
+ MVK_ICD=$(find $(brew --prefix)/share $(brew --prefix)/etc $(brew --prefix molten-vk) -name "MoltenVK_icd.json" 2>/dev/null | head -n 1)
64
+ echo "VK_ICD_FILENAMES=$MVK_ICD" >> "$GITHUB_ENV"
65
+ echo "MVK_CONFIG_USE_METAL_ARGUMENT_BUFFERS=0" >> "$GITHUB_ENV"
66
+ echo "DYLD_LIBRARY_PATH=$(brew --prefix)/lib:$DYLD_LIBRARY_PATH" >> "$GITHUB_ENV"
67
+ export VK_ICD_FILENAMES="$MVK_ICD"
68
+ export MVK_CONFIG_USE_METAL_ARGUMENT_BUFFERS=0
69
+ export DYLD_LIBRARY_PATH="$(brew --prefix)/lib:$DYLD_LIBRARY_PATH"
70
+ vulkaninfo --summary
71
+
72
+ - name: Install Vulkan (Windows)
73
+ if: runner.os == 'Windows'
74
+ shell: pwsh
75
+ run: |
76
+ Invoke-WebRequest -Uri "https://github.com/pal1000/mesa-dist-win/releases/download/23.3.2/mesa3d-23.3.2-release-msvc.7z" -OutFile "mesa.7z"
77
+ 7z x mesa.7z -omesa -y | Out-Null
78
+
79
+ $lvpJson = (Resolve-Path "mesa\x64\lvp_icd.x86_64.json").Path
80
+ $driverKey = "HKLM:\SOFTWARE\Khronos\Vulkan\Drivers"
81
+ if (!(Test-Path $driverKey)) { New-Item -Path $driverKey -Force | Out-Null }
82
+ New-ItemProperty -Path $driverKey -Name $lvpJson -PropertyType DWord -Value 0 -Force | Out-Null
83
+
84
+ $sdkVersion = "${{ inputs.vulkan-sdk-version }}"
85
+ $installer = "VulkanSDK-$sdkVersion-Installer.exe"
86
+ Invoke-WebRequest -Uri "https://sdk.lunarg.com/sdk/download/$sdkVersion/windows/$installer" -OutFile $installer
87
+ Start-Process -FilePath ".\$installer" -ArgumentList "--accept-licenses --default-answer --confirm-command install" -Wait -NoNewWindow
88
+
89
+ "VULKAN_SDK=C:\VulkanSDK\$sdkVersion" | Out-File -FilePath $env:GITHUB_ENV -Append
90
+ "C:\VulkanSDK\$sdkVersion\Bin" | Out-File -FilePath $env:GITHUB_PATH -Append
@@ -0,0 +1,27 @@
1
+ #!/usr/bin/env bash
2
+ # Installs the torchvulkan build toolchain inside the cibuildwheel manylinux
3
+ # container (AlmaLinux 8 / manylinux_2_28): the Vulkan SDK (headers + loader +
4
+ # the spirv-cross executable) and the Slang compiler (slangc).
5
+ #
6
+ # The install paths here MUST match the VULKAN_SDK / SLANG_DIR / PATH entries in
7
+ # [tool.cibuildwheel.linux].environment in pyproject.toml.
8
+ set -euo pipefail
9
+
10
+ SLANG_VERSION="2026.8.1"
11
+ VULKAN_SDK_VERSION="1.3.296.0"
12
+
13
+ dnf install -y wget tar xz unzip
14
+
15
+ wget -q -O /tmp/vulkansdk.tar.xz \
16
+ "https://sdk.lunarg.com/sdk/download/${VULKAN_SDK_VERSION}/linux/vulkansdk-linux-x86_64-${VULKAN_SDK_VERSION}.tar.xz"
17
+ mkdir -p /opt/vulkan
18
+ tar -xf /tmp/vulkansdk.tar.xz -C /opt/vulkan --strip-components=1
19
+
20
+ wget -q -O /tmp/slang.zip \
21
+ "https://github.com/shader-slang/slang/releases/download/v${SLANG_VERSION}/slang-${SLANG_VERSION}-linux-x86_64.zip"
22
+ mkdir -p /opt/slang
23
+ unzip -q /tmp/slang.zip -d /opt/slang
24
+
25
+ test -f /opt/vulkan/x86_64/include/vulkan/vulkan.h
26
+ test -x /opt/vulkan/x86_64/bin/spirv-cross
27
+ test -x /opt/slang/bin/slangc
@@ -0,0 +1,98 @@
1
+ name: Release
2
+
3
+ on:
4
+ push:
5
+ tags:
6
+ - "v*"
7
+ workflow_dispatch: {}
8
+
9
+ permissions:
10
+ contents: read
11
+
12
+ jobs:
13
+ build_sdist:
14
+ name: Build sdist
15
+ runs-on: ubuntu-latest
16
+ steps:
17
+ - uses: actions/checkout@v4
18
+ with:
19
+ submodules: recursive
20
+ fetch-depth: 0
21
+
22
+ - uses: actions/setup-python@v5
23
+ with:
24
+ python-version: "3.10"
25
+
26
+ - name: Build sdist
27
+ run: |
28
+ python -m pip install --upgrade pip build
29
+ python -m build --sdist
30
+
31
+ - uses: actions/upload-artifact@v4
32
+ with:
33
+ name: dist-sdist
34
+ path: dist/*.tar.gz
35
+
36
+ build_wheels:
37
+ name: Wheels (${{ matrix.name }})
38
+ runs-on: ${{ matrix.os }}
39
+ strategy:
40
+ fail-fast: false
41
+ matrix:
42
+ include:
43
+ - name: Linux x86_64
44
+ os: ubuntu-latest
45
+ - name: macOS Apple Silicon
46
+ os: macos-15
47
+ - name: Windows x86_64
48
+ os: windows-latest
49
+ steps:
50
+ - uses: actions/checkout@v4
51
+ with:
52
+ submodules: recursive
53
+ fetch-depth: 0
54
+
55
+ - uses: actions/setup-python@v5
56
+ with:
57
+ python-version: "3.12"
58
+
59
+ - name: Compute version for setuptools_scm
60
+ shell: bash
61
+ run: |
62
+ python -m pip install --upgrade setuptools_scm
63
+ echo "SETUPTOOLS_SCM_PRETEND_VERSION_FOR_TORCHVULKAN=$(python -m setuptools_scm)" >> "$GITHUB_ENV"
64
+
65
+ - name: Install toolchain (Slang, Vulkan, spirv-cross)
66
+ if: runner.os != 'Linux'
67
+ uses: ./.github/actions/setup-toolchain
68
+
69
+ - name: Build wheels
70
+ uses: pypa/cibuildwheel@v3.2.1
71
+ env:
72
+ CIBW_BUILD_VERBOSITY: "1"
73
+
74
+ - uses: actions/upload-artifact@v4
75
+ with:
76
+ name: dist-${{ matrix.os }}
77
+ path: wheelhouse/*.whl
78
+
79
+ pypi-publish:
80
+ name: Upload release to PyPI
81
+ needs: [build_sdist, build_wheels]
82
+ if: startsWith(github.ref, 'refs/tags/')
83
+ runs-on: ubuntu-latest
84
+ environment:
85
+ name: pypi
86
+ url: https://pypi.org/p/torchvulkan
87
+ permissions:
88
+ id-token: write
89
+ steps:
90
+ - uses: actions/download-artifact@v4
91
+ with:
92
+ path: dist
93
+ pattern: dist-*
94
+ merge-multiple: true
95
+
96
+ - uses: pypa/gh-action-pypi-publish@release/v1
97
+ with:
98
+ packages-dir: dist
@@ -0,0 +1,60 @@
1
+ name: torchvulkan CI
2
+
3
+ on:
4
+ push:
5
+ branches: [ main ]
6
+ pull_request:
7
+ branches: [ main ]
8
+
9
+ jobs:
10
+ build-and-test:
11
+ name: ${{ matrix.os }} / Py${{ matrix.python-version }} / ${{ matrix.torch-factor }}
12
+ runs-on: ${{ matrix.os }}
13
+ env:
14
+ SKBUILD_CMAKE_GENERATOR: Ninja
15
+ strategy:
16
+ fail-fast: false
17
+ matrix:
18
+ os: [ubuntu-latest, macos-15, windows-latest]
19
+ python-version: ["3.10", "3.12", "3.14"]
20
+ torch-factor: ["torch210", "torch213"]
21
+
22
+ steps:
23
+ - name: Checkout Code
24
+ uses: actions/checkout@v4
25
+ with:
26
+ submodules: recursive
27
+
28
+ - name: Set up Python ${{ matrix.python-version }}
29
+ uses: actions/setup-python@v5
30
+ with:
31
+ python-version: ${{ matrix.python-version }}
32
+
33
+ - name: Install toolchain (Slang, Vulkan, spirv-cross)
34
+ uses: ./.github/actions/setup-toolchain
35
+
36
+ - name: Install Tox
37
+ run: |
38
+ python -m pip install --upgrade pip
39
+ pip install tox
40
+
41
+ - name: Run Tests via Tox
42
+ env:
43
+ PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu
44
+ SKBUILD_CMAKE_GENERATOR: Ninja
45
+ MVK_CONFIG_LOG_LEVEL: 3
46
+ VK_LOADER_DEBUG: all
47
+ run: |
48
+ tox -e py-${{ matrix.torch-factor }}
49
+
50
+ testing-success:
51
+ name: testing success
52
+ needs: build-and-test
53
+ if: always()
54
+ runs-on: ubuntu-latest
55
+ steps:
56
+ - name: Verify all test jobs passed
57
+ if: needs.build-and-test.result != 'success'
58
+ run: |
59
+ echo "build-and-test result: ${{ needs.build-and-test.result }}"
60
+ exit 1
@@ -0,0 +1,42 @@
1
+ # torchvulkan
2
+ torchvulkan/_version.py
3
+ src/shaders/shader_registry.h
4
+
5
+ # testing data
6
+ data/
7
+
8
+ # python
9
+ .venv/
10
+ venv/
11
+ __pycache__/
12
+ *.py[cod]
13
+ *.pyo
14
+ *.pyd
15
+ .pytest_cache/
16
+ .mypy_cache/
17
+
18
+ # build files
19
+ build/
20
+ dist/
21
+ *.egg-info/
22
+ *.so
23
+ *.dll
24
+ *.dylib
25
+ *.o
26
+ *.a
27
+ *.out
28
+
29
+ # cmake
30
+ CMakeCache.txt
31
+ CMakeFiles/
32
+ cmake_install.cmake
33
+ Makefile
34
+
35
+ # IDE files
36
+ .vscode/
37
+ .idea/
38
+ *.swp
39
+
40
+ # operating system
41
+ .DS_Store
42
+ Thumbs.db
@@ -0,0 +1,28 @@
1
+ cmake_minimum_required(VERSION 3.18)
2
+ project(torchvulkan LANGUAGES CXX)
3
+
4
+ find_package(Torch REQUIRED)
5
+ find_package(Vulkan REQUIRED)
6
+ find_package(Python3 COMPONENTS Interpreter Development.Module REQUIRED)
7
+ find_program(SLANGC_EXE slangc HINTS "$ENV{SLANG_DIR}" REQUIRED)
8
+ find_program(SPIRV_CROSS_EXE spirv-cross REQUIRED)
9
+
10
+ set(CMAKE_CXX_STANDARD 20)
11
+ set(CMAKE_CXX_STANDARD_REQUIRED ON)
12
+
13
+ add_subdirectory(src/shaders)
14
+
15
+ file(GLOB_RECURSE SOURCES CONFIGURE_DEPENDS "src/*.cpp")
16
+
17
+ python3_add_library(torchvulkan_lib MODULE ${SOURCES} WITH_SOABI)
18
+ add_dependencies(torchvulkan_lib compile_shaders)
19
+
20
+ set_target_properties(torchvulkan_lib PROPERTIES OUTPUT_NAME "_C" LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/torchvulkan")
21
+ target_link_libraries(torchvulkan_lib PRIVATE ${TORCH_LIBRARIES} Python3::Module ${CMAKE_DL_LIBS})
22
+ target_include_directories(torchvulkan_lib PRIVATE src/ extern/vma extern/volk ${Vulkan_INCLUDE_DIRS})
23
+
24
+ target_compile_definitions(torchvulkan_lib PRIVATE VK_NO_PROTOTYPES TORCH_EXTENSION_NAME=_C)
25
+
26
+ target_compile_features(torchvulkan_lib PRIVATE cxx_std_20)
27
+
28
+ install(TARGETS torchvulkan_lib DESTINATION torchvulkan)
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025 torchvulkan
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,200 @@
1
+ Metadata-Version: 2.4
2
+ Name: torchvulkan
3
+ Version: 0.1.0
4
+ Summary: Cross-Platform Vulkan Backend for PyTorch
5
+ Author: Zain Sharief
6
+ License-Expression: MIT
7
+ License-File: LICENSE
8
+ Project-URL: Homepage, https://github.com/zainsharief/torchvulkan
9
+ Project-URL: Repository, https://github.com/zainsharief/torchvulkan
10
+ Project-URL: Issues, https://github.com/zainsharief/torchvulkan/issues
11
+ Requires-Dist: torch>=2.10.0
12
+ Provides-Extra: test
13
+ Requires-Dist: pytest; extra == "test"
14
+ Requires-Dist: tox; extra == "test"
15
+ Description-Content-Type: text/markdown
16
+
17
+ <h1 align="center">torchvulkan</h1>
18
+ <p align="center">
19
+ <i>A Cross-Platform Vulkan Backend for PyTorch</i>
20
+ </p>
21
+
22
+ <p align="center">
23
+ <a href="https://github.com/zainsharief/torchvulkan/actions/workflows/test_torchvulkan.yml"><img alt="CI" src="https://github.com/zainsharief/torchvulkan/actions/workflows/test_torchvulkan.yml/badge.svg"></a>
24
+ <a href="LICENSE"><img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-blue.svg"></a>
25
+ <img alt="Python 3.10+" src="https://img.shields.io/badge/python-3.10%2B-blue">
26
+ <img alt="Status: experimental" src="https://img.shields.io/badge/status-experimental-orange">
27
+ </p>
28
+
29
+ `torchvulkan` registers a `vulkan` device with PyTorch and runs tensor operations as [Vulkan](https://www.vulkan.org/) compute shaders. Vulkan runs on GPUs from AMD, Intel, NVIDIA, and Apple (via MoltenVK), plus software rasterizers, so you get one vendor-neutral GPU backend across Linux, macOS, and Windows.
30
+
31
+ You use it like any other PyTorch device:
32
+
33
+ ```python
34
+ import torch
35
+ import torchvulkan # registers the 'vulkan' device on import
36
+
37
+ x = torch.randn(1024, 1024).to("vulkan")
38
+ y = torch.randn(1024, 1024).to("vulkan")
39
+
40
+ z = (x @ y).relu() # runs on the GPU via Vulkan compute shaders
41
+ print(z.sum().item()) # move a scalar back to the CPU
42
+ ```
43
+
44
+ > [!WARNING]
45
+ > **Experimental.** `torchvulkan` is an early-stage project. A small set of operators is implemented natively; everything else falls back to the CPU. Expect breaking changes before `1.0`.
46
+
47
+ ## Table of contents
48
+
49
+ - [Installation](#installation)
50
+ - [Quickstart](#quickstart)
51
+ - [Example: training MNIST](#example-training-mnist)
52
+ - [Benchmarks](#benchmarks)
53
+ - [Python API](#python-api)
54
+ - [CPU fallback](#cpu-fallback)
55
+ - [Limitations](#limitations)
56
+ - [How it works](#how-it-works)
57
+
58
+ ## Installation
59
+
60
+ ```bash
61
+ pip install torchvulkan
62
+ ```
63
+
64
+ At runtime you need a **Vulkan driver**: your GPU's driver on Linux and Windows (or a software rasterizer such as Mesa [lavapipe](https://docs.mesa3d.org/drivers/llvmpipe.html)), or [MoltenVK](https://github.com/KhronosGroup/MoltenVK) on macOS. Verify it worked:
65
+
66
+ ```python
67
+ import torchvulkan
68
+ print(torchvulkan.is_available(), torchvulkan.device_count())
69
+ ```
70
+
71
+ <details>
72
+ <summary><b>Building from source</b></summary>
73
+
74
+ Prebuilt wheels are published for Linux (x86_64), macOS (Apple Silicon), and Windows (x86_64) on Python 3.10 to 3.14. On other platforms `pip` builds from the source distribution, which needs the shader toolchain below. The same steps build a local checkout for development.
75
+
76
+ | Requirement | Notes |
77
+ | --- | --- |
78
+ | A Vulkan driver / loader | Your GPU's Vulkan driver, or a software rasterizer such as Mesa [lavapipe](https://docs.mesa3d.org/drivers/llvmpipe.html). On macOS this is [MoltenVK](https://github.com/KhronosGroup/MoltenVK). |
79
+ | [Vulkan SDK](https://vulkan.lunarg.com/) | Provides Vulkan headers and `spirv-cross`. |
80
+ | [`slangc`](https://github.com/shader-slang/slang/releases) | The Slang shader compiler, on your `PATH`. |
81
+ | `spirv-cross` | Ships with the Vulkan SDK (or your package manager). |
82
+ | CMake ≥ 3.18 + a C++20 compiler | Ninja is recommended (`SKBUILD_CMAKE_GENERATOR=Ninja`). |
83
+ | `torch==2.13.0` | Install before building; the extension is compiled against this exact version. |
84
+
85
+ > The exact per-OS commands CI uses to install this toolchain live in [`.github/actions/setup-toolchain`](.github/actions/setup-toolchain/action.yml).
86
+
87
+ ```bash
88
+ pip install torch==2.13.0
89
+ pip install . # editable install for development: pip install -e .
90
+ ```
91
+
92
+ </details>
93
+
94
+ ## Quickstart
95
+
96
+ ```python
97
+ import torch
98
+ import torchvulkan
99
+
100
+ # Move tensors to the Vulkan device
101
+ a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device="vulkan")
102
+ b = torch.ones(2, 2, device="vulkan")
103
+
104
+ c = a @ b + 1.0 # matmul + broadcasted add
105
+ d = c.exp().sum(dim=0) # elementwise + reduction
106
+
107
+ print(d.cpu()) # bring the result back to the CPU
108
+ ```
109
+
110
+ `nn.Module`s move to the device the usual way:
111
+
112
+ ```python
113
+ model = torch.nn.Linear(128, 10).to("vulkan")
114
+ out = model(torch.randn(32, 128).to("vulkan"))
115
+ ```
116
+
117
+ ## Example: training MNIST
118
+
119
+ [`examples/mnist.py`](examples/mnist.py) trains a small fully-connected network on the Vulkan device end to end. Forward, loss, backward, and an Adam step all run through the backend:
120
+
121
+ ```bash
122
+ pip install torchvision
123
+ python examples/mnist.py
124
+ ```
125
+
126
+ ```python
127
+ device = torch.device("vulkan" if torchvulkan.is_available() else "cpu")
128
+ model = SimpleNN().to(device)
129
+ criterion = torch.nn.CrossEntropyLoss()
130
+ optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
131
+
132
+ for data, targets in train_loader:
133
+ data, targets = data.to(device), targets.to(device)
134
+ loss = criterion(model(data), targets)
135
+ optimizer.zero_grad()
136
+ loss.backward()
137
+ optimizer.step()
138
+ ```
139
+
140
+ ## Benchmarks
141
+
142
+ [`benchmark/mnist_benchmark.py`](benchmark/mnist_benchmark.py) times a combined forward + backward pass of the MNIST network (in `float16`), sweeping the hidden width so the parameter count grows from ~50K up to ~312M.
143
+
144
+ Compared against the platform-native GPU backend, the `vulkan` device tracks closely and stays orders of magnitude ahead of the CPU as the model grows:
145
+
146
+ <p align="center">
147
+ <img alt="Forward + backward pass time vs. model size on an NVIDIA GPU: vulkan tracks CUDA and beats CPU" src="media/benchmark-cuda.png" width="49%">
148
+ <img alt="Forward + backward pass time vs. model size on Apple Silicon: vulkan tracks MPS and beats CPU" src="media/benchmark-mps.png" width="49%">
149
+ </p>
150
+
151
+ <p align="center">
152
+ <i>Left: NVIDIA GPU (vulkan vs. CUDA). Right: Apple Silicon (vulkan vs. MPS).</i>
153
+ </p>
154
+
155
+ ```bash
156
+ pip install matplotlib
157
+ python benchmark/mnist_benchmark.py
158
+ ```
159
+
160
+ ## Python API
161
+
162
+ ```python
163
+ import torchvulkan
164
+
165
+ torchvulkan.is_available() # -> bool: a usable Vulkan device was found
166
+ torchvulkan.device_count() # -> int: number of Vulkan devices
167
+ torchvulkan.synchronize() # block until queued GPU work has finished
168
+ torchvulkan.empty_cache() # release cached device memory back to the driver
169
+ torchvulkan.__version__ # package version
170
+ ```
171
+
172
+ Importing `torchvulkan` registers the backend, after which the device is addressed as `"vulkan"` / `torch.device("vulkan")` throughout PyTorch.
173
+
174
+ ## CPU fallback
175
+
176
+ By default, any operation, dtype, or tensor shape the backend doesn't implement natively runs on the **CPU** (with a one-time warning), so models keep working despite partial coverage.
177
+
178
+ Set `TORCHVULKAN_STRICT=1` to turn those silent fallbacks into errors. This is useful for finding which operators a workload actually needs:
179
+
180
+ ```bash
181
+ TORCHVULKAN_STRICT=1 python your_script.py
182
+ ```
183
+
184
+ ## Limitations
185
+
186
+ - `bfloat16` and complex dtypes not natively supported; falls back to the CPU.
187
+ - Tensor rank ≤ 4. Tensors with more than 4 dimensions fall back to the CPU.
188
+ - `torch>=2.10.0` only. The extension is compiled against 2.13.0 but has testing support for >=2.10.0.
189
+ - APIs and coverage may change before `1.0.0`.
190
+
191
+ ## How it works
192
+
193
+ 1. Importing the package registers a PyTorch `PrivateUse1` backend and renames it to `vulkan`.
194
+ 2. Operations dispatched to a `vulkan` tensor are routed to native implementations that record and submit **Vulkan compute shaders**. The shaders are written in [Slang](https://github.com/shader-slang/slang), compiled to SPIR-V at build time, and embedded into the extension.
195
+ 3. Device memory is managed with [VMA](https://github.com/GPUOpen-LibrariesAndSDKs/VulkanMemoryAllocator); Vulkan entry points are loaded with [volk](https://github.com/zeux/volk).
196
+ 4. Anything without a native implementation is handled by [CPU fallback](#cpu-fallback).
197
+
198
+
199
+ > [!NOTE]
200
+ > **Disclaimer:** `torchvulkan` is an independent open-source project and is not affiliated with or endorsed by PyTorch or the Linux Foundation.