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.
- torchvulkan-0.1.0/.github/actions/setup-toolchain/action.yml +90 -0
- torchvulkan-0.1.0/.github/scripts/install-toolchain-linux.sh +27 -0
- torchvulkan-0.1.0/.github/workflows/release.yaml +98 -0
- torchvulkan-0.1.0/.github/workflows/test_torchvulkan.yml +60 -0
- torchvulkan-0.1.0/.gitignore +42 -0
- torchvulkan-0.1.0/CMakeLists.txt +28 -0
- torchvulkan-0.1.0/LICENSE +21 -0
- torchvulkan-0.1.0/PKG-INFO +200 -0
- torchvulkan-0.1.0/README.md +184 -0
- torchvulkan-0.1.0/benchmark/mnist_benchmark.py +276 -0
- torchvulkan-0.1.0/examples/main.py +10 -0
- torchvulkan-0.1.0/examples/mnist.py +62 -0
- torchvulkan-0.1.0/extern/vma/vk_mem_alloc.h +19845 -0
- torchvulkan-0.1.0/extern/volk/volk.c +4004 -0
- torchvulkan-0.1.0/extern/volk/volk.h +3267 -0
- torchvulkan-0.1.0/media/benchmark-cuda.png +0 -0
- torchvulkan-0.1.0/media/benchmark-mps.png +0 -0
- torchvulkan-0.1.0/pyproject.toml +54 -0
- torchvulkan-0.1.0/requirements.txt +1 -0
- torchvulkan-0.1.0/scripts/build_shaders.py +211 -0
- torchvulkan-0.1.0/scripts/config.py +106 -0
- torchvulkan-0.1.0/scripts/pack_shaders.py +214 -0
- torchvulkan-0.1.0/src/api/backend.cpp +34 -0
- torchvulkan-0.1.0/src/api/guard.cpp +38 -0
- torchvulkan-0.1.0/src/api/hook.cpp +35 -0
- torchvulkan-0.1.0/src/api/ops/binary.cpp +360 -0
- torchvulkan-0.1.0/src/api/ops/binary.h +76 -0
- torchvulkan-0.1.0/src/api/ops/factory.cpp +511 -0
- torchvulkan-0.1.0/src/api/ops/factory.h +84 -0
- torchvulkan-0.1.0/src/api/ops/helpers.h +196 -0
- torchvulkan-0.1.0/src/api/ops/matmul.cpp +500 -0
- torchvulkan-0.1.0/src/api/ops/matmul.h +70 -0
- torchvulkan-0.1.0/src/api/ops/nllloss.cpp +150 -0
- torchvulkan-0.1.0/src/api/ops/nllloss.h +28 -0
- torchvulkan-0.1.0/src/api/ops/reduce.cpp +175 -0
- torchvulkan-0.1.0/src/api/ops/reduce.h +61 -0
- torchvulkan-0.1.0/src/api/ops/softmax.cpp +26 -0
- torchvulkan-0.1.0/src/api/ops/softmax.h +14 -0
- torchvulkan-0.1.0/src/api/ops/unary.cpp +216 -0
- torchvulkan-0.1.0/src/api/ops/unary.h +39 -0
- torchvulkan-0.1.0/src/api/ops.cpp +122 -0
- torchvulkan-0.1.0/src/shaders/CMakeLists.txt +44 -0
- torchvulkan-0.1.0/src/shaders/binaryop.slang.j2 +183 -0
- torchvulkan-0.1.0/src/shaders/cast.slang.j2 +54 -0
- torchvulkan-0.1.0/src/shaders/copy.slang.j2 +53 -0
- torchvulkan-0.1.0/src/shaders/fill.slang.j2 +65 -0
- torchvulkan-0.1.0/src/shaders/helpers/int_divider.slang.j2 +10 -0
- torchvulkan-0.1.0/src/shaders/helpers/vulkan_arithmetic.slang.j2 +170 -0
- torchvulkan-0.1.0/src/shaders/matmul_coop.slang.j2 +780 -0
- torchvulkan-0.1.0/src/shaders/matmul_simd.slang.j2 +317 -0
- torchvulkan-0.1.0/src/shaders/nllloss.slang.j2 +80 -0
- torchvulkan-0.1.0/src/shaders/reduce.slang.j2 +80 -0
- torchvulkan-0.1.0/src/shaders/unaryop.slang.j2 +138 -0
- torchvulkan-0.1.0/src/vulkan/allocator.cpp +215 -0
- torchvulkan-0.1.0/src/vulkan/allocator.h +113 -0
- torchvulkan-0.1.0/src/vulkan/cache.cpp +487 -0
- torchvulkan-0.1.0/src/vulkan/cache.h +81 -0
- torchvulkan-0.1.0/src/vulkan/device_context.cpp +55 -0
- torchvulkan-0.1.0/src/vulkan/device_context.h +42 -0
- torchvulkan-0.1.0/src/vulkan/memory.cpp +48 -0
- torchvulkan-0.1.0/src/vulkan/memory.h +36 -0
- torchvulkan-0.1.0/src/vulkan/shader.cpp +74 -0
- torchvulkan-0.1.0/src/vulkan/shader.h +34 -0
- torchvulkan-0.1.0/src/vulkan/vulkan_context.cpp +384 -0
- torchvulkan-0.1.0/src/vulkan/vulkan_context.h +43 -0
- torchvulkan-0.1.0/tests/test_pytorch_ops.py +185 -0
- torchvulkan-0.1.0/torchvulkan/__init__.py +31 -0
- torchvulkan-0.1.0/torchvulkan/_version.py +24 -0
- 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.
|