This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new d6bfb45 [ADDON] Add torch c dlpack for macos (#256)
d6bfb45 is described below
commit d6bfb45e0c9c3c6150da7885804a1d3918afbeff
Author: Yaxing Cai <[email protected]>
AuthorDate: Wed Nov 12 04:08:26 2025 -0800
[ADDON] Add torch c dlpack for macos (#256)
This PR adds the torch c dlpack for macos.
---
.github/workflows/torch_c_dlpack.yml | 26 +++++++++++++++-
addons/torch_c_dlpack_ext/build_aot_wheels.sh | 44 +++++++++++++++++++++------
addons/torch_c_dlpack_ext/build_backend.py | 2 --
3 files changed, 60 insertions(+), 12 deletions(-)
diff --git a/.github/workflows/torch_c_dlpack.yml
b/.github/workflows/torch_c_dlpack.yml
index 8d7296e..2054352 100644
--- a/.github/workflows/torch_c_dlpack.yml
+++ b/.github/workflows/torch_c_dlpack.yml
@@ -66,6 +66,30 @@ jobs:
with:
name: pypi-wheels-linux-${{ matrix.arch }}-${{ matrix.python-version
}}
path: ./wheelhouse/*.whl
+ build_wheels_macos:
+ strategy:
+ fail-fast: false
+ matrix:
+ arch: ["arm64"]
+ python-version: ["cp39", "cp310", "cp311", "cp312", "cp313", "cp314"]
+ runs-on: macos-14
+ steps:
+ - uses: actions/checkout@v5
+ with:
+ repository: apache/tvm-ffi
+ ref: main
+ fetch-depth: 0
+ submodules: recursive
+ path: tvm-ffi
+ - uses: astral-sh/setup-uv@v7
+ - name: build torch libs and wheels
+ working-directory: ./tvm-ffi
+ run: |
+ bash ./addons/torch_c_dlpack_ext/build_aot_wheels.sh ${{ matrix.arch
}} ${{ matrix.python-version }}
+ - uses: actions/upload-artifact@v4
+ with:
+ name: pypi-wheels-macos-${{ matrix.arch }}-${{ matrix.python-version
}}
+ path: ./tvm-ffi/addons/torch_c_dlpack_ext/wheelhouse/*.whl
build_source:
runs-on: ubuntu-latest
steps:
@@ -90,7 +114,7 @@ jobs:
name: pypi-source
path: ./tvm-ffi/addons/torch_c_dlpack_ext/dist/*tar.gz
upload_pypi:
- needs: [build_wheels_linux, build_source]
+ needs: [build_wheels_linux, build_wheels_macos, build_source]
runs-on: ubuntu-latest
environment: pypi
permissions:
diff --git a/addons/torch_c_dlpack_ext/build_aot_wheels.sh
b/addons/torch_c_dlpack_ext/build_aot_wheels.sh
index 0361b5f..f8a4885 100755
--- a/addons/torch_c_dlpack_ext/build_aot_wheels.sh
+++ b/addons/torch_c_dlpack_ext/build_aot_wheels.sh
@@ -21,6 +21,17 @@ set -eux
arch=$1
python_version=$2
+os=$(uname -s)
+
+case "$os" in
+ "Linux" | "Darwin")
+ ;;
+ *)
+ echo "Unknown OS: $os"
+ return 1
+ ;;
+esac
+
tvm_ffi="$PWD"
torch_c_dlpack_ext="$tvm_ffi"/addons/torch_c_dlpack_ext
@@ -54,7 +65,10 @@ function check_availability() {
"2.4")
! [[ "$arch" == "aarch64" || "$python_version" == "cp313" ||
"$python_version" == "cp314" ]]
;;
- "2.5" | "2.6")
+ "2.5")
+ ! [[ ("$os" == "Linux" && ("$arch" == "aarch64" ||
"$python_version" == "cp314")) || ("$os" == "Darwin" && ("$python_version" ==
"cp313" || "$python_version" == "cp314")) ]]
+ ;;
+ "2.6")
! [[ "$arch" == "aarch64" || "$python_version" == "cp314" ]]
;;
"2.7" | "2.8")
@@ -74,15 +88,19 @@ function check_availability() {
function build_libs() {
local torch_version=$1
if check_availability "$torch_version"; then
- mkdir "$tvm_ffi"/.venv -p
uv venv "$tvm_ffi"/.venv/torch"$torch_version" --python
"$python_version"
source "$tvm_ffi"/.venv/torch"$torch_version"/bin/activate
uv pip install setuptools ninja
- uv pip install torch=="$torch_version" --index-url "$(get_torch_url
"$torch_version")"
+ if [[ "$os" == "Linux" ]]; then
+ uv pip install torch=="$torch_version" --index-url
"$(get_torch_url "$torch_version")"
+ else
+ uv pip install torch=="$torch_version"
+ fi
uv pip install -v .
- mkdir "$tvm_ffi"/lib -p
python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir
"$tvm_ffi"/lib
- python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir
"$tvm_ffi"/lib --build-with-cuda
+ if [[ "$os" == "Linux" ]]; then
+ python -m tvm_ffi.utils._build_optional_torch_c_dlpack
--output-dir "$tvm_ffi"/lib --build-with-cuda
+ fi
ls "$tvm_ffi"/lib
deactivate
rm -rf "$tvm_ffi"/.venv/torch"$torch_version"
@@ -91,6 +109,8 @@ function build_libs() {
fi
}
+mkdir -p "$tvm_ffi"/.venv
+mkdir -p "$tvm_ffi"/lib
torch_versions=("2.4" "2.5" "2.6" "2.7" "2.8" "2.9")
for version in "${torch_versions[@]}"; do
build_libs "$version"
@@ -99,11 +119,17 @@ done
cp "$tvm_ffi"/lib/*.so "$torch_c_dlpack_ext"/torch_c_dlpack_ext
uv venv "$tvm_ffi"/.venv/build --python "$python_version"
source "$tvm_ffi"/.venv/build/bin/activate
-uv pip install build wheel auditwheel
+uv pip install build wheel
cd "$torch_c_dlpack_ext"
python -m build -w
ls dist
-python -m wheel tags dist/*.whl --python-tag="$python_version"
--abi-tag="$python_version" --remove
-ls dist
-auditwheel repair --exclude libtorch.so --exclude libtorch_cpu.so --exclude
libc10.so --exclude libtorch_python.so dist/*.whl -w wheelhouse
+if [[ "$os" == "Linux" ]]; then
+ python -m wheel tags dist/*.whl --python-tag="$python_version"
--abi-tag="$python_version" --remove
+ uv pip install auditwheel
+ auditwheel repair --exclude libtorch.so --exclude libtorch_cpu.so
--exclude libc10.so --exclude libtorch_python.so --exclude libtorch_cuda.so
--exclude libc10_cuda.so dist/*.whl -w wheelhouse
+else
+ python -m wheel tags dist/*.whl --python-tag="$python_version"
--abi-tag="$python_version" --platform-tag=macosx_11_0_arm64 --remove
+ uv pip install delocate
+ delocate-wheel -v --ignore-missing-dependencies --exclude
libtorch.dylib,libtorch_cpu.dylib,libc10.dylib,libtorch_python.dylib dist/*.whl
-w wheelhouse
+fi
ls wheelhouse
diff --git a/addons/torch_c_dlpack_ext/build_backend.py
b/addons/torch_c_dlpack_ext/build_backend.py
index 17bc2c4..e4504fa 100644
--- a/addons/torch_c_dlpack_ext/build_backend.py
+++ b/addons/torch_c_dlpack_ext/build_backend.py
@@ -40,8 +40,6 @@ build_editable = orig.build_editable
def _is_lib_prebuilt() -> bool:
if sys.platform.startswith("win32"):
extension = "dll"
- elif sys.platform.startswith("darwin"):
- extension = "dylib"
else:
extension = "so"
return next(_package_path.rglob(f"*.{extension}"), None) is not None
