torch_c_dlpack.yml

# 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/