Skip to content

Commit d1595c4

Browse files
cyx-6cyx-6
authored
[ADDON] Add torch c dlpack ext wheels for windows (#252)
This PR adds the torch c dlpack ext wheels for windows.
1 parent 752ac8e commit d1595c4

4 files changed

Lines changed: 202 additions & 4 deletions

File tree

‎.github/workflows/torch_c_dlpack.yml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ on:
2525
required: true
2626

2727
jobs:
28-
build_wheels:
28+
build_wheels_linux:
2929
strategy:
3030
fail-fast: false
3131
matrix:
@@ -64,7 +64,7 @@ jobs:
6464
docker cp build_lib:/tvm-ffi/addons/torch_c_dlpack_ext/wheelhouse/ .
6565
- uses: actions/upload-artifact@v4
6666
with:
67-
name: pypi-wheels-${{ matrix.arch }}-${{ matrix.python-version }}
67+
name: pypi-wheels-linux-${{ matrix.arch }}-${{ matrix.python-version }}
6868
path: ./wheelhouse/*.whl
6969
build_source:
7070
runs-on: ubuntu-latest
@@ -90,7 +90,7 @@ jobs:
9090
name: pypi-source
9191
path: ./tvm-ffi/addons/torch_c_dlpack_ext/dist/*tar.gz
9292
upload_pypi:
93-
needs: [build_wheels, build_source]
93+
needs: [build_wheels_linux, build_source]
9494
runs-on: ubuntu-latest
9595
environment: pypi
9696
permissions:

‎.github/workflows/torch_c_dlpack_windows.yml‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
name: torch c dlpack
19+
20+
on:
21+
workflow_dispatch:
22+
inputs:
23+
branch:
24+
description: "Branch or tag to publish (manual run)"
25+
required: true
26+
27+
jobs:
28+
build_wheels_windows:
29+
strategy:
30+
fail-fast: false
31+
matrix:
32+
arch: ["x86_64"]
33+
python-version: ["cp39", "cp310", "cp311", "cp312", "cp313", "cp314"]
34+
runs-on: windows-latest
35+
steps:
36+
- uses: actions/checkout@v5
37+
with:
38+
repository: apache/tvm-ffi
39+
ref: main
40+
fetch-depth: 0
41+
submodules: recursive
42+
path: tvm-ffi
43+
- uses: Jimver/cuda-toolkit@6008063726ffe3309d1b22e413d9e88fed91a2f2
44+
id: cuda-toolkit
45+
- uses: astral-sh/setup-uv@v7
46+
- name: build wheels
47+
env:
48+
CUDA_HOME: ${{steps.cuda-toolkit.outputs.CUDA_PATH}}
49+
shell: cmd
50+
working-directory: ./tvm-ffi
51+
run: ./addons/torch_c_dlpack_ext/build_aot_wheels.bat ${{ matrix.arch }} ${{ matrix.python-version }}
52+
- uses: actions/upload-artifact@v4
53+
with:
54+
name: pypi-wheels-windows-${{ matrix.arch }}-${{ matrix.python-version }}
55+
path: ./tvm-ffi/addons/torch_c_dlpack_ext/wheelhouse/*.whl
56+
upload_pypi:
57+
needs: [build_wheels_windows]
58+
runs-on: ubuntu-latest
59+
environment: pypi
60+
permissions:
61+
id-token: write
62+
attestations: write
63+
if: github.event_name == 'workflow_dispatch' # <-- publish only on manual trigger
64+
steps:
65+
- uses: actions/download-artifact@v4
66+
with:
67+
pattern: pypi-*
68+
path: dist
69+
merge-multiple: true
70+
- name: Generate artifact attestation for sdist and wheels
71+
uses: actions/attest-build-provenance@v1
72+
with:
73+
subject-path: dist/*
74+
- name: Publish package distributions to PyPI
75+
uses: pypa/gh-action-pypi-publish@release/v1
76+
with:
77+
attestations: true
78+
verbose: true
79+
# testing publish url
80+
# repository-url: https://test.pypi.org/legacy/

‎addons/torch_c_dlpack_ext/build_aot_wheels.bat‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
@REM Licensed to the Apache Software Foundation (ASF) under one
2+
@REM or more contributor license agreements. See the NOTICE file
3+
@REM distributed with this work for additional information
4+
@REM regarding copyright ownership. The ASF licenses this file
5+
@REM to you under the Apache License, Version 2.0 (the
6+
@REM "License"); you may not use this file except in compliance
7+
@REM with the License. You may obtain a copy of the License at
8+
@REM
9+
@REM http://www.apache.org/licenses/LICENSE-2.0
10+
@REM
11+
@REM Unless required by applicable law or agreed to in writing,
12+
@REM software distributed under the License is distributed on an
13+
@REM "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
@REM KIND, either express or implied. See the License for the
15+
@REM specific language governing permissions and limitations
16+
@REM under the License.
17+
@echo off
18+
setlocal enabledelayedexpansion
19+
20+
set arch=%~1
21+
set python_version=%~2
22+
23+
set tvm_ffi=%cd%
24+
set torch_c_dlpack_ext=%tvm_ffi%\addons\torch_c_dlpack_ext
25+
26+
if not exist "%tvm_ffi%\.venv" mkdir "%tvm_ffi%\.venv"
27+
if not exist "%tvm_ffi%\lib" mkdir "%tvm_ffi%\lib"
28+
for %%P in (2.4 2.5 2.6 2.7 2.8 2.9) do (
29+
call :build_libs %%P
30+
)
31+
32+
copy %tvm_ffi%\lib\*.dll %torch_c_dlpack_ext%\torch_c_dlpack_ext
33+
uv venv %tvm_ffi%\.venv\build --python %python_version%
34+
call %tvm_ffi%\.venv\build\Scripts\activate
35+
uv pip install build wheel
36+
cd %torch_c_dlpack_ext%
37+
python -m build -w
38+
dir dist
39+
for %%f in (dist\*.whl) do python -m wheel tags "%%f" --python-tag=%python_version% --abi-tag=%python_version% --platform-tag=win_amd64
40+
dir dist
41+
mkdir wheelhouse
42+
copy dist\*-win_amd64.whl wheelhouse
43+
dir wheelhouse
44+
endlocal
45+
exit /b
46+
47+
:build_libs
48+
set torch_version=%1
49+
call :check_availability
50+
if %errorlevel%==0 (
51+
call :get_torch_url
52+
uv venv %tvm_ffi%\.venv\torch%torch_version% --python %python_version%
53+
call %tvm_ffi%\.venv\torch%torch_version%\Scripts\activate
54+
uv pip install setuptools ninja
55+
uv pip install torch==%torch_version% --index-url !torch_url!
56+
uv pip install -v .
57+
python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir %tvm_ffi%\lib
58+
python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir %tvm_ffi%\lib --build-with-cuda
59+
call deactivate
60+
rmdir -s -q %tvm_ffi%\.venv\torch%torch_version%
61+
) else (
62+
echo Skipping build for torch %torch_version% on %arch% with python %python_version% as it is not available.
63+
)
64+
exit /b 0
65+
66+
67+
:check_availability
68+
if %torch_version%==2.4 (
69+
if %python_version%==cp313 exit /b 1
70+
if %python_version%==cp314 exit /b 1
71+
exit /b 0
72+
)
73+
if %torch_version%==2.5 (
74+
if %python_version%==cp314 exit /b 1
75+
exit /b 0
76+
)
77+
if %torch_version%==2.6 (
78+
if %python_version%==cp314 exit /b 1
79+
exit /b 0
80+
)
81+
if %torch_version%==2.7 (
82+
if %python_version%==cp314 exit /b 1
83+
exit /b 0
84+
)
85+
if %torch_version%==2.8 (
86+
if %python_version%==cp314 exit /b 1
87+
exit /b 0
88+
)
89+
if %torch_version%==2.9 (
90+
if %python_version%==cp39 exit /b 1
91+
exit /b 0
92+
)
93+
echo Unknown or unsupported torch version: %torch_version% >&2
94+
exit /b 1
95+
96+
:get_torch_url
97+
set cuda_version=
98+
if %torch_version%==2.4 set cuda_version=cu124
99+
if %torch_version%==2.5 set cuda_version=cu124
100+
if %torch_version%==2.6 set cuda_version=cu126
101+
if %torch_version%==2.7 set cuda_version=cu128
102+
if %torch_version%==2.8 set cuda_version=cu129
103+
if %torch_version%==2.9 set cuda_version=cu129
104+
if defined cuda_version (
105+
set torch_url=https://download.pytorch.org/whl/%cuda_version%
106+
exit /b 0
107+
)
108+
echo Unknown or unsupported torch version: %torch_version% >&2
109+
exit /b 1

‎python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -627,7 +627,12 @@ def _generate_ninja_build(
627627

628628
# append include paths
629629
for path in include_paths:
630-
cflags.append("-I{}".format(path.replace(":", "$:")))
630+
path_str = str(path)
631+
if " " in path_str:
632+
path_str = f'"{path_str}"'
633+
if IS_WINDOWS:
634+
path_str = path_str.replace(":", "$:")
635+
cflags.append(f"-I{path_str}")
631636

632637
# flags
633638
ninja = []
@@ -795,9 +800,13 @@ def main() -> None: # noqa: PLR0912, PLR0915
795800
if IS_WINDOWS:
796801
# On Windows, use .lib format for linking
797802
ldflags.extend(["c10.lib", "torch.lib", "torch_cpu.lib", "torch_python.lib"])
803+
if args.build_with_cuda:
804+
ldflags.extend(["torch_cuda.lib", "c10_cuda.lib"])
798805
else:
799806
# On Unix/macOS, use -l format for linking
800807
ldflags.extend(["-lc10", "-ltorch", "-ltorch_cpu", "-ltorch_python"])
808+
if args.build_with_cuda:
809+
ldflags.extend(["-ltorch_cuda", "-lc10_cuda"])
801810

802811
# Add Python library linking
803812
if IS_WINDOWS:

0 commit comments

Comments
 (0)