Skip to content

Commit f65bfcf

Browse files
zcbenzzcbenz
committed
Add CI for Windows CUDA build
1 parent 96296e9 commit f65bfcf

10 files changed

Lines changed: 313 additions & 143 deletions

File tree

‎.github/actions/build-backend-wheel/action.yml‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
name: 'Build backend wheel'
2+
description: 'Build the backend-only Python wheel for release on Linux and Windows'
3+
4+
inputs:
5+
cmake-args:
6+
description: 'The args for generating CMake project'
7+
required: true
8+
arch:
9+
description: 'Platform architecture tag'
10+
required: false
11+
default: |-
12+
${{ case(runner.arch == 'x64', 'x86_64',
13+
runner.arch == 'x86', 'i686',
14+
runner.arch == 'arm', 'armv7l',
15+
runner.arch == 'arm64', 'aarch64',
16+
'unknown')
17+
}}
18+
19+
runs:
20+
using: 'composite'
21+
steps:
22+
- name: Build package
23+
shell: bash
24+
env:
25+
CMAKE_ARGS: ${{ inputs.cmake-args }}
26+
run: |
27+
uv pip install auditwheel "build<=1.4.2" patchelf setuptools
28+
python setup.py clean --all
29+
MLX_BUILD_STAGE=2 python -m build -w
30+
31+
- name: Post process
32+
if: runner.os == 'Linux'
33+
shell: bash
34+
run: |
35+
auditwheel repair dist/mlx_cuda*.whl \
36+
--plat manylinux_2_35_${{ inputs.arch }} \
37+
--exclude libcublas* \
38+
--exclude libcuda* \
39+
--exclude libcudnn* \
40+
--exclude libnccl* \
41+
--exclude libnvrtc*

‎.github/actions/build-cuda-release/action.yml‎

Lines changed: 0 additions & 31 deletions
This file was deleted.

‎.github/actions/build-linux-release/action.yml‎

Lines changed: 0 additions & 42 deletions
This file was deleted.

‎.github/actions/build-linux/action.yml‎

Lines changed: 11 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,38 +1,26 @@
1-
name: 'Build and Test on Linux'
1+
name: 'Build on Linux'
2+
description: 'Build C++ and Python binaries for testing on Linux'
23

34
inputs:
4-
toolkit:
5-
description: 'The toolkit to build with'
6-
required: false
7-
default: 'cpu'
5+
cmake-args:
6+
description: 'The args for generating CMake project'
7+
required: true
88

99
runs:
10-
using: "composite"
10+
using: 'composite'
1111
steps:
12-
1312
- name: Install Python package
14-
id: python_build
15-
shell: sh
13+
shell: bash
1614
env:
1715
DEBUG: 1
18-
CMAKE_ARGS: >-
19-
-DCMAKE_COMPILE_WARNING_AS_ERROR=ON
20-
-DMLX_BUILD_CUDA=${{ startsWith(inputs.toolkit, 'cuda') && 'ON' || 'OFF' }}
16+
CMAKE_ARGS: ${{ inputs.cmake-args }}
2117
run: |
22-
if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then
23-
# There is no GPU in arm64 runner, use a common arch.
24-
CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=80"
25-
# Can not build tests and stubs when the built executables can not run.
26-
CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF -DMLX_BUILD_PYTHON_STUBS=OFF"
27-
fi
2818
# Install cpu-only torch to save space
29-
pip install torch --index-url https://download.pytorch.org/whl/cpu
30-
pip install --no-build-isolation -e ".[dev]" -v
31-
# Pass the CMAKE_ARGS to following steps.
32-
echo CMAKE_ARGS="$CMAKE_ARGS" >> $GITHUB_OUTPUT
19+
uv pip install torch --index-url https://download.pytorch.org/whl/cpu
20+
uv pip install --no-build-isolation -e ".[dev]" -v
3321
3422
- name: Build CPP only
3523
shell: bash
3624
run: |
37-
cmake . -B build -DCMAKE_BUILD_TYPE=Debug ${{ steps.python_build.outputs.CMAKE_ARGS }}
25+
cmake . -B build -DCMAKE_BUILD_TYPE=Debug ${{ inputs.cmake-args }}
3826
cmake --build build -j $(nproc)

‎.github/actions/build-wheel/action.yml‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
name: 'Build wheel'
2+
description: 'Build the Python wheel for release on Linux and Windows'
3+
4+
inputs:
5+
cmake-args:
6+
description: 'The args for generating CMake project'
7+
required: true
8+
arch:
9+
description: 'Platform architecture tag'
10+
required: false
11+
default: |-
12+
${{ case(runner.arch == 'x64', 'x86_64',
13+
runner.arch == 'x86', 'i686',
14+
runner.arch == 'arm', 'armv7l',
15+
runner.arch == 'arm64', 'aarch64',
16+
'unknown')
17+
}}
18+
build-backend:
19+
description: 'Build the backend mlx-cpu package'
20+
type: boolean
21+
required: false
22+
default: true
23+
24+
runs:
25+
using: 'composite'
26+
steps:
27+
- name: Build frontend package
28+
shell: bash
29+
run: |
30+
uv pip install auditwheel "build<=1.4.2" patchelf setuptools
31+
python setup.py clean --all
32+
MLX_BUILD_STAGE=1 python -m build -w
33+
34+
- name: Post process
35+
if: runner.os == 'Linux'
36+
shell: bash
37+
run: |
38+
auditwheel repair dist/mlx-*.whl \
39+
--plat manylinux_2_35_${{ inputs.arch }} \
40+
--exclude libmlx.so* \
41+
--only-plat
42+
43+
- name: Build backend package
44+
if: inputs.build-backend
45+
shell: bash
46+
run: |
47+
python setup.py clean --all
48+
MLX_BUILD_STAGE=2 python -m build -w
49+
50+
- name: Post process
51+
if: inputs.build-backend && runner.os == 'Linux'
52+
shell: bash
53+
run: |
54+
auditwheel repair dist/mlx_cpu*.whl \
55+
--plat manylinux_2_35_${{ inputs.arch }}

‎.github/actions/build-windows/action.yml‎

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,22 @@
11
name: 'Build on Windows'
2+
description: 'Build C++ and Python binaries for testing on Windows'
3+
4+
inputs:
5+
cmake-args:
6+
description: 'The args for generating CMake project'
7+
required: true
28

39
runs:
410
using: 'composite'
511
steps:
612
- name: Install Python package
7-
id: python-build
813
shell: cmd
914
env:
10-
# For MSVC, Ninja/Release is the only config supported by ccache.
11-
CMAKE_ARGS: >-
12-
-G Ninja
13-
-DCMAKE_BUILD_TYPE=Release
14-
-DCMAKE_C_COMPILER=cl
15-
-DCMAKE_CXX_COMPILER=cl
16-
-DCMAKE_RC_COMPILER=rc
17-
run: |
18-
uv pip install ".[dev]" -v
19-
:: Pass the CMAKE_ARGS to following steps.
20-
>>%GITHUB_OUTPUT% ECHO CMAKE_ARGS=%CMAKE_ARGS%
15+
CMAKE_ARGS: ${{ inputs.cmake-args }}
16+
run: uv pip install ".[dev]" -v
2117

2218
- name: Build CPP only
2319
shell: cmd
2420
run: |
25-
cmake . -B build ${{ steps.python-build.outputs.CMAKE_ARGS }}
21+
cmake . -B build ${{ inputs.cmake-args }}
2622
cmake --build build -j %NUMBER_OF_PROCESSORS%

‎.github/actions/setup-linux/action.yml‎

Lines changed: 35 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
name: 'Setup Linux Environment'
1+
name: 'Setup Linux environment'
22
description: 'Install dependencies for Linux builds'
33

44
inputs:
@@ -18,6 +18,11 @@ inputs:
1818
required: false
1919
default: 'ccache'
2020

21+
outputs:
22+
cmake-args:
23+
description: 'The args for generating CMake project'
24+
value: ${{ steps.cmake-args.outputs.cmakeArgs }}
25+
2126
runs:
2227
using: "composite"
2328
steps:
@@ -27,7 +32,7 @@ runs:
2732
echo "::group::Install common dependencies"
2833
sudo apt-get update
2934
sudo apt-get install -y --no-install-recommends \
30-
gdb zip \
35+
cmake gdb ninja-build zip \
3136
libblas-dev liblapack-dev liblapacke-dev \
3237
openmpi-bin openmpi-common libopenmpi-dev
3338
echo "::endgroup::"
@@ -50,17 +55,14 @@ runs:
5055
ptx-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}-
5156
${{ hashFiles('mlx/backend/cuda/**') }}
5257
53-
- uses: actions/setup-python@v6
54-
with:
55-
python-version: ${{ inputs.python-version }}
58+
- uses: astral-sh/setup-uv@v7
5659

5760
- name: Setup Python venv
5861
shell: bash
5962
run: |
6063
echo "::group::Setup Python venv"
61-
python -m venv .venv
64+
uv venv --python ${{ inputs.python-version }} --managed-python
6265
source .venv/bin/activate
63-
pip install setuptools cmake typing_extensions
6466
echo PATH=$PATH >> $GITHUB_ENV
6567
# Search python packages in .venv
6668
echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
@@ -109,3 +111,29 @@ runs:
109111
echo "::group::NVIDIA-SMI Status"
110112
nvidia-smi || true
111113
echo "::endgroup::"
114+
115+
- name: Generate CMake args
116+
id: cmake-args
117+
shell: bash
118+
run: |
119+
echo "::group::Generate CMake args"
120+
cmakeArgs=(
121+
"-G Ninja"
122+
"-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
123+
)
124+
if ${{ startsWith(inputs.toolkit, 'cuda') }} ; then
125+
cmakeArgs+=("-DMLX_BUILD_CUDA=ON")
126+
if ${{ runner.arch == 'arm64' }} ; then
127+
# There is no GPU in arm64 runner, use a common arch.
128+
cmakeArgs+=("-DMLX_CUDA_ARCHITECTURES=80")
129+
# Can not build tests and stubs when the built executables can not run.
130+
cmakeArgs+=("-DMLX_BUILD_TESTS=OFF")
131+
cmakeArgs+=("-DMLX_BUILD_PYTHON_STUBS=OFF")
132+
fi
133+
else
134+
cmakeArgs+=("-DMLX_BUILD_CUDA=OFF")
135+
fi
136+
# Pass to following steps.
137+
IFS=" "
138+
echo "cmakeArgs=${cmakeArgs[*]}" >> $GITHUB_OUTPUT
139+
echo "::endgroup::"

0 commit comments

Comments
 (0)