Skip to content

torch c dlpack

torch c dlpack #8

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
name: torch c dlpack
on:
workflow_dispatch:
inputs:
branch:
description: "Branch or tag to publish (manual run)"
required: true
jobs:
build_libs:
strategy:
fail-fast: false
matrix:
arch: ["x86_64", "aarch64"]
python-version: ["cp39", "cp310", "cp311", "cp312", "cp313"]
runs-on: ${{ matrix.arch == 'aarch64' && 'ubuntu-24.04-arm' || 'ubuntu-latest' }}
steps:
- uses: jlumbroso/free-disk-space@main
if: ${{ matrix.arch == 'x86_64'}}
- uses: actions/checkout@v5
with:
repository: apache/tvm-ffi
ref: main
fetch-depth: 0
submodules: recursive
path: tvm-ffi
- name: prepare docker environment
run: |
docker pull quay.io/pypa/manylinux_2_28_${{ matrix.arch }}:latest
docker run --name build_lib -d quay.io/pypa/manylinux_2_28_${{ matrix.arch }}:latest tail -f /dev/null
docker cp tvm-ffi build_lib:/tvm-ffi
wget https://developer.download.nvidia.com/compute/cuda/13.0.2/local_installers/cuda-repo-rhel8-13-0-local-13.0.2_580.95.05-1.${{ matrix.arch }}.rpm > /dev/null 2>&1
docker cp cuda-repo-rhel8-13-0-local-13.0.2_580.95.05-1.${{ matrix.arch }}.rpm build_lib:/
rm cuda-repo-rhel8-13-0-local-13.0.2_580.95.05-1.${{ matrix.arch }}.rpm
docker exec build_lib bash -c "rpm -i /cuda-repo-rhel8-13-0-local-13.0.2_580.95.05-1.${{ matrix.arch }}.rpm \
&& dnf clean all \
&& dnf -y install cuda-toolkit-13-0 \
&& rm /cuda-repo-rhel8-13-0-local-13.0.2_580.95.05-1.${{ matrix.arch }}.rpm \
&& dnf clean all \
&& mkdir /libs"
- name: build torch 2.4 lib
if: ${{ matrix.python-version != 'cp313' && matrix.python-version != 'cp314' && matrix.arch != 'aarch64' }}
run: |
docker exec build_lib bash -c "uv venv /torch2.4 --python ${{ matrix.python-version }} \
&& source /torch2.4/bin/activate \
&& uv pip install setuptools ninja \
&& uv pip install torch==2.4 --index-url https://download.pytorch.org/whl/cu124 \
&& cd /tvm-ffi \
&& uv pip install -v . \
&& python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir /libs \
&& python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir /libs --build-with-cuda \
&& rm -rf /torch2.4"
- name: build torch 2.5 lib
if: ${{ matrix.python-version != 'cp314' && matrix.arch != 'aarch64' }}
run: |
docker exec build_lib bash -c "uv venv /torch2.5 --python ${{ matrix.python-version }} \
&& source /torch2.5/bin/activate \
&& uv pip install setuptools ninja \
&& uv pip install torch==2.5 --index-url https://download.pytorch.org/whl/cu124 \
&& cd /tvm-ffi \
&& uv pip install -v . \
&& python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir /libs \
&& python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir /libs --build-with-cuda \
&& rm -rf /torch2.5"
- name: build torch 2.6 lib
if: ${{ matrix.python-version != 'cp314' && matrix.arch != 'aarch64' }}
run: |
docker exec build_lib bash -c "uv venv /torch2.6 --python ${{ matrix.python-version }} \
&& source /torch2.6/bin/activate \
&& uv pip install setuptools ninja \
&& uv pip install torch==2.6 --index-url https://download.pytorch.org/whl/cu126 \
&& cd /tvm-ffi \
&& uv pip install -v . \
&& python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir /libs \
&& python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir /libs --build-with-cuda \
&& rm -rf /torch2.6"
- name: build torch 2.7 lib
if: ${{ matrix.python-version != 'cp314' }}
run: |
docker exec build_lib bash -c "uv venv /torch2.7 --python ${{ matrix.python-version }} \
&& source /torch2.7/bin/activate \
&& uv pip install setuptools ninja \
&& uv pip install torch==2.7 --index-url https://download.pytorch.org/whl/cu128 \
&& cd /tvm-ffi \
&& uv pip install -v . \
&& python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir /libs \
&& python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir /libs --build-with-cuda \
&& rm -rf /torch2.7"
- name: build torch 2.8 lib
if: ${{ matrix.python-version != 'cp314' }}
run: |
docker exec build_lib bash -c "uv venv /torch2.8 --python ${{ matrix.python-version }} \
&& source /torch2.8/bin/activate \
&& uv pip install setuptools ninja \
&& uv pip install torch==2.8 --index-url https://download.pytorch.org/whl/cu129 \
&& cd /tvm-ffi \
&& uv pip install -v . \
&& python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir /libs \
&& python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir /libs --build-with-cuda \
&& rm -rf /torch2.8"
- name: build torch 2.9 lib
if: ${{ matrix.python-version != 'cp39' }}
run: |
docker exec build_lib bash -c "uv venv /torch2.9 --python ${{ matrix.python-version }} \
&& source /torch2.9/bin/activate \
&& uv pip install setuptools ninja \
&& uv pip install torch==2.9 --index-url https://download.pytorch.org/whl/cu129 \
&& cd /tvm-ffi \
&& uv pip install -v . \
&& python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir /libs \
&& python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir /libs --build-with-cuda \
&& rm -rf /torch2.9"
- name: collect built libs
run: |
mkdir libs
docker cp build_lib:/libs .
- uses: actions/upload-artifact@v4
with:
name: libs-${{ matrix.arch }}-${{ matrix.python-version }}
path: ./libs/*.so
build_wheels:
needs: build_libs
strategy:
matrix:
arch: ["x86_64", "aarch64"]
python-version: ["cp39", "cp310", "cp311", "cp312", "cp313"]
runs-on: ${{ matrix.arch == 'aarch64' && 'ubuntu-24.04-arm' || 'ubuntu-latest' }}
container: quay.io/pypa/manylinux_2_28_${{ matrix.arch }}
steps:
- uses: actions/checkout@v5
with:
repository: apache/tvm-ffi
ref: main
fetch-depth: 0
submodules: recursive
path: tvm-ffi
- uses: actions/download-artifact@v4
with:
pattern: libs-${{ matrix.arch }}-${{ matrix.python-version }}
path: ./tvm-ffi/addons/torch_c_dlpack_ext/torch_c_dlpack_ext
merge-multiple: true
- name: build wheels
run: |
/opt/python/${{ matrix.python-version }}-${{ matrix.python-version }}/bin/pip3 install build wheel auditwheel
cd ./tvm-ffi/addons/torch_c_dlpack_ext
/opt/python/${{ matrix.python-version }}-${{ matrix.python-version }}/bin/python -m build -w
/opt/python/${{ matrix.python-version }}-${{ matrix.python-version }}/bin/python -m wheel tags dist/*.whl --python-tag=${{ matrix.python-version }} --abi-tag=${{ matrix.python-version }} --remove
auditwheel repair --exclude libtorch.so --exclude libtorch_cpu.so --exclude libc10.so --exclude libtorch_python.so dist/*.whl -w wheelhouse
- uses: actions/upload-artifact@v4
with:
name: pypi-wheels-${{ matrix.arch }}-${{ matrix.python-version }}
path: ./tvm-ffi/addons/torch_c_dlpack_ext/wheelhouse/*.whl
build_source:
runs-on: ubuntu-latest
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
with:
python-version: "3.12"
activate-environment: true
- name: build source
run: |
uv pip install build
cd ./tvm-ffi/addons/torch_c_dlpack_ext
python -m build -s
- uses: actions/upload-artifact@v4
with:
name: pypi-source
path: ./tvm-ffi/addons/torch_c_dlpack_ext/dist/*tar.gz
upload_pypi:
needs: [build_wheels, build_source]
runs-on: ubuntu-latest
environment: pypi
permissions:
id-token: write
attestations: write
# if: github.event_name == 'workflow_dispatch' # <-- publish only on manual trigger
steps:
- uses: actions/download-artifact@v4
with:
pattern: pypi-*
path: dist
merge-multiple: true
- name: Generate artifact attestation for sdist and wheels
uses: actions/attest-build-provenance@v1
with:
subject-path: dist/*
- name: Publish package distributions to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
attestations: true
verbose: true
# testing publish url
# repository-url: https://test.pypi.org/legacy/