diff --git a/.github/actions/build-cuda/action.yml b/.github/actions/build-cuda/action.yml new file mode 100644 index 000000000000..156632d81da2 --- /dev/null +++ b/.github/actions/build-cuda/action.yml @@ -0,0 +1,117 @@ +# 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: Build CUDA Runtime +description: Build libtvm_runtime_cuda for the TVM wheel packaging flow. + +inputs: + arch: + description: "Target Linux architecture for CUDA builds (x86_64 or aarch64)" + required: true + linux_image: + description: "Manylinux image tag to use on Linux runners" + required: false + default: "" + linux_image_tag: + description: "Pinned manylinux container tag shared with cibuildwheel" + required: false + default: "" + cuda_architectures: + description: "CMake CUDA architectures for libtvm_runtime_cuda.so" + required: false + default: "75" + include_cuda_runtime: + description: "Set to true to build the CUDA runtime library" + required: false + default: "false" + +outputs: + cuda_runtime_path: + description: "Absolute path to the built libtvm_runtime_cuda.so, or empty for CPU-only wheels" + value: ${{ steps.cuda_runtime.outputs.path }} + +runs: + using: "composite" + steps: + - uses: ./.github/actions/detect-env-vars + id: env_vars + + - name: Detect CUDA inputs + id: cuda_inputs + shell: bash -l {0} + env: + INPUT_INCLUDE_CUDA_RUNTIME: ${{ inputs.include_cuda_runtime }} + run: | + set -eux + include_cuda_runtime="$(printf '%s' "${INPUT_INCLUDE_CUDA_RUNTIME}" | tr '[:upper:]' '[:lower:]')" + case "${include_cuda_runtime}" in + 1|true|yes|on) include_cuda_runtime=1 ;; + 0|false|no|off) include_cuda_runtime=0 ;; + *) + echo "include_cuda_runtime must be a boolean value" >&2 + exit 1 + ;; + esac + echo "include_cuda_runtime=${include_cuda_runtime}" >> "${GITHUB_OUTPUT}" + + - name: Build CUDA runtime in manylinux + if: runner.os == 'Linux' && steps.cuda_inputs.outputs.include_cuda_runtime == '1' + shell: bash -l {0} + env: + TVM_MANYLINUX_IMAGE: ${{ inputs.linux_image }} + TVM_MANYLINUX_IMAGE_TAG: ${{ inputs.linux_image_tag }} + TVM_ARCH: ${{ inputs.arch }} + TVM_CUDA_ARCHITECTURES: ${{ inputs.cuda_architectures }} + TVM_CUDA_BUILD_DIR: ${{ runner.temp }}/tvm-wheel-cuda + TVM_INCLUDE_CUDA_RUNTIME: "1" + TVM_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} + CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} + run: ci/scripts/package/tvm_wheel_helper.sh manylinux-cuda + + - name: Reject non-Linux CUDA runtime builds + if: runner.os != 'Linux' && steps.cuda_inputs.outputs.include_cuda_runtime == '1' + shell: bash -l {0} + run: | + echo "CUDA runtime wheels are only enabled on Linux in this workflow" >&2 + exit 1 + + - name: Report CUDA runtime output + id: cuda_runtime + shell: bash -l {0} + env: + INCLUDE_CUDA_RUNTIME: ${{ steps.cuda_inputs.outputs.include_cuda_runtime }} + TVM_CUDA_BUILD_DIR: ${{ runner.temp }}/tvm-wheel-cuda + TVM_CUDA_RUNTIME_PATH: "" + run: | + set -eux + if [[ "${INCLUDE_CUDA_RUNTIME}" != "1" ]]; then + echo "path=" >> "${GITHUB_OUTPUT}" + exit 0 + fi + cuda_runtime="$(ci/scripts/package/tvm_wheel_helper.sh cuda-path)" + if [[ -z "${cuda_runtime}" ]]; then + echo "CUDA runtime build did not produce libtvm_runtime_cuda.so" >&2 + exit 1 + fi + case "${cuda_runtime}" in + "${TVM_CUDA_BUILD_DIR}"/*) ;; + *) + echo "CUDA runtime path is outside the expected build directory: ${cuda_runtime}" >&2 + exit 1 + ;; + esac + echo "path=${cuda_runtime}" >> "${GITHUB_OUTPUT}" diff --git a/.github/actions/build-wheel-for-publish/action.yml b/.github/actions/build-wheel-for-publish/action.yml new file mode 100644 index 000000000000..d6daf503816e --- /dev/null +++ b/.github/actions/build-wheel-for-publish/action.yml @@ -0,0 +1,311 @@ +# 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: Build TVM Wheel +description: > + Build and test the LLVM-enabled TVM wheel for a given OS/architecture + combination using cibuildwheel. + +inputs: + arch: + description: "Target architecture for cibuildwheel (e.g., x86_64, aarch64, arm64, AMD64)" + required: true + build: + description: "cibuildwheel build selector (e.g., cp310-manylinux_x86_64)" + required: true + linux_image: + description: "Manylinux image family to use on Linux runners; required on Linux" + required: false + default: "" + linux_image_tag: + description: "Pinned manylinux container tag shared with the CUDA runtime build; required on Linux" + required: false + default: "" + distribution_name: + description: "Optional wheel distribution name override, useful for TestPyPI" + required: false + default: "" + distribution_version: + description: "Optional wheel distribution version override, useful for TestPyPI" + required: false + default: "" + cuda_architectures: + description: "CMake CUDA architectures for libtvm_runtime_cuda.so" + required: false + default: "75" + include_cuda_runtime: + description: "Set to true to inject the CUDA runtime library" + required: false + default: "false" + cuda_runtime_path: + description: "Absolute path to libtvm_runtime_cuda.so produced by build-cuda" + required: false + default: "" + +runs: + using: "composite" + steps: + - uses: ./.github/actions/detect-env-vars + id: env_vars + + - name: Detect wheel inputs + id: wheel_inputs + shell: bash -l {0} + env: + INPUT_ARCH: ${{ inputs.arch }} + INPUT_INCLUDE_CUDA_RUNTIME: ${{ inputs.include_cuda_runtime }} + INPUT_LINUX_IMAGE: ${{ inputs.linux_image }} + INPUT_LINUX_IMAGE_TAG: ${{ inputs.linux_image_tag }} + INPUT_CUDA_RUNTIME_PATH: ${{ inputs.cuda_runtime_path }} + run: | + set -eux + wheel_platform_tag="" + manylinux_container_image="" + cibw_manylinux_x86_64_image="${INPUT_LINUX_IMAGE:-manylinux_2_28}" + cibw_manylinux_aarch64_image="${INPUT_LINUX_IMAGE:-manylinux_2_28}" + cibw_container_engine="docker" + include_cuda_runtime="$(printf '%s' "${INPUT_INCLUDE_CUDA_RUNTIME}" | tr '[:upper:]' '[:lower:]')" + case "${include_cuda_runtime}" in + 1|true|yes|on) include_cuda_runtime=1 ;; + 0|false|no|off) include_cuda_runtime=0 ;; + *) + echo "include_cuda_runtime must be a boolean value" >&2 + exit 1 + ;; + esac + if [[ "${RUNNER_OS}" == "Linux" ]]; then + if [[ -z "${INPUT_LINUX_IMAGE}" ]]; then + echo "linux_image is required on Linux runners" >&2 + exit 1 + fi + if [[ -z "${INPUT_LINUX_IMAGE_TAG}" ]]; then + echo "linux_image_tag is required on Linux runners" >&2 + exit 1 + fi + if [[ "${include_cuda_runtime}" == "1" ]]; then + if [[ -z "${INPUT_CUDA_RUNTIME_PATH}" ]]; then + echo "cuda_runtime_path is required when include_cuda_runtime=true" >&2 + exit 1 + fi + if [[ ! -f "${INPUT_CUDA_RUNTIME_PATH}" ]]; then + echo "cuda_runtime_path does not exist: ${INPUT_CUDA_RUNTIME_PATH}" >&2 + exit 1 + fi + fi + wheel_platform_tag="${INPUT_LINUX_IMAGE}_${INPUT_ARCH}" + manylinux_container_image="quay.io/pypa/${INPUT_LINUX_IMAGE}_${INPUT_ARCH}:${INPUT_LINUX_IMAGE_TAG}" + if [[ "${INPUT_ARCH}" == "x86_64" ]]; then + cibw_manylinux_x86_64_image="${manylinux_container_image}" + elif [[ "${INPUT_ARCH}" == "aarch64" ]]; then + cibw_manylinux_aarch64_image="${manylinux_container_image}" + fi + cibw_container_engine="docker; create_args: --volume /opt/llvm:/opt/llvm:ro" + if [[ "${include_cuda_runtime}" == "1" && -n "${INPUT_CUDA_RUNTIME_PATH}" ]]; then + cuda_runtime_dir="$(dirname "${INPUT_CUDA_RUNTIME_PATH}")" + cibw_container_engine+=" --volume ${cuda_runtime_dir}:${cuda_runtime_dir}:ro" + fi + elif [[ "${include_cuda_runtime}" == "1" ]]; then + echo "CUDA runtime injection is only enabled on Linux in this workflow" >&2 + exit 1 + fi + echo "wheel_platform_tag=${wheel_platform_tag}" >> "${GITHUB_OUTPUT}" + echo "manylinux_container_image=${manylinux_container_image}" >> "${GITHUB_OUTPUT}" + echo "cibw_manylinux_x86_64_image=${cibw_manylinux_x86_64_image}" >> "${GITHUB_OUTPUT}" + echo "cibw_manylinux_aarch64_image=${cibw_manylinux_aarch64_image}" >> "${GITHUB_OUTPUT}" + echo "cibw_container_engine=${cibw_container_engine}" >> "${GITHUB_OUTPUT}" + echo "include_cuda_runtime=${include_cuda_runtime}" >> "${GITHUB_OUTPUT}" + + # Single source of truth for the LLVM toolchain version, shared by the cache + # key and the conda install steps below. + - name: Set LLVM version + shell: bash + run: echo "LLVM_VERSION=22.1.0" >> "$GITHUB_ENV" + + - name: Prepare LLVM cache path (Unix) + if: runner.os != 'Windows' + shell: bash + run: | + set -eux + sudo mkdir -p /opt/llvm + sudo chown -R "$(whoami)" /opt/llvm + + # ---- Cache LLVM prefix ---- + - name: Cache LLVM + uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + id: llvm-cache + with: + path: ${{ runner.os == 'Windows' && 'C:/opt/llvm' || '/opt/llvm' }} + key: tvm-wheel-llvm-${{ env.LLVM_VERSION }}-${{ runner.os }}-${{ inputs.arch }}-v4 + + # ---- Install LLVM via conda (cache miss only) ---- + - name: Setup conda + if: steps.llvm-cache.outputs.cache-hit != 'true' + uses: conda-incubator/setup-miniconda@8ee1f361103df19b6f8c8655fd3967a8ecb162d5 # v4.0.1 + continue-on-error: true + id: conda1 + with: + miniforge-version: latest + conda-remove-defaults: true + + - name: Setup conda (retry with tar.bz2) + if: steps.llvm-cache.outputs.cache-hit != 'true' && steps.conda1.outcome == 'failure' + uses: conda-incubator/setup-miniconda@8ee1f361103df19b6f8c8655fd3967a8ecb162d5 # v4.0.1 + with: + miniforge-version: latest + use-only-tar-bz2: true + conda-remove-defaults: true + + - name: Install LLVM (Unix) + if: steps.llvm-cache.outputs.cache-hit != 'true' && runner.os != 'Windows' + shell: bash -l {0} + run: | + set -eux + if [[ "${RUNNER_OS}" == "Linux" ]]; then + sudo mkdir -p /opt/llvm + sudo chown -R "$(whoami)" /opt/llvm + fi + conda create -q -p /opt/llvm -c conda-forge \ + "llvmdev=${LLVM_VERSION}" "clangdev=${LLVM_VERSION}" "compiler-rt=${LLVM_VERSION}" zlib zstd-static libxml2-devel \ + -y + + - name: Install LLVM (Windows) + if: steps.llvm-cache.outputs.cache-hit != 'true' && runner.os == 'Windows' + shell: cmd /C call {0} + run: | + conda create -q -p C:\opt\llvm -c conda-forge llvmdev=%LLVM_VERSION% zlib zstd-static libxml2-devel -y + + - name: Create static llvm-config wrapper (Unix) + if: runner.os != 'Windows' + shell: bash + run: | + set -eux + printf '%s\n' \ + '#!/usr/bin/env bash' \ + 'exec "$(dirname "$0")/llvm-config" --link-static "$@"' \ + | sudo tee /opt/llvm/bin/llvm-config-static >/dev/null + sudo chmod +x /opt/llvm/bin/llvm-config-static + + - name: Create static llvm-config wrapper (Windows) + if: runner.os == 'Windows' + shell: pwsh + run: | + @' + @echo off + "C:\opt\llvm\Library\bin\llvm-config.exe" --link-static %* + '@ | Set-Content -Path 'C:\opt\llvm\Library\bin\llvm-config-static.bat' + + - name: Print build inputs + shell: bash -l {0} + run: | + set -eux + git log -1 --oneline + if [[ "${RUNNER_OS}" == "Windows" ]]; then + "C:/opt/llvm/Library/bin/llvm-config-static.bat" --version + "C:/opt/llvm/Library/bin/llvm-config-static.bat" --system-libs + else + /opt/llvm/bin/llvm-config-static --version + /opt/llvm/bin/llvm-config-static --system-libs + fi + + # ---- Build and test wheels ---- + - name: Build and test wheels + uses: pypa/cibuildwheel@298ed2fb2c105540f5ed055e8a6ad78d82dd3a7e # v3.3.1 + with: + package-dir: . + output-dir: wheelhouse + env: + CIBW_BUILD: ${{ inputs.build }} + CIBW_ARCHS_LINUX: ${{ inputs.arch }} + CIBW_ARCHS_MACOS: ${{ inputs.arch }} + CIBW_ARCHS_WINDOWS: ${{ inputs.arch }} + CIBW_MANYLINUX_X86_64_IMAGE: ${{ steps.wheel_inputs.outputs.cibw_manylinux_x86_64_image }} + CIBW_MANYLINUX_AARCH64_IMAGE: ${{ steps.wheel_inputs.outputs.cibw_manylinux_aarch64_image }} + CIBW_BUILD_VERBOSITY: 1 + CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} + CIBW_CONTAINER_ENGINE: ${{ steps.wheel_inputs.outputs.cibw_container_engine }} + CIBW_BEFORE_BUILD_LINUX: >- + if command -v dnf >/dev/null 2>&1; then dnf -y install libxml2; fi && + python -m pip install -U pip cmake ninja scikit-build-core wheel auditwheel && + python -m pip install -v "{project}/3rdparty/tvm-ffi" + CIBW_BEFORE_BUILD_MACOS: >- + python -m pip install -U pip cmake ninja scikit-build-core wheel delocate && + python -m pip install -v "{project}/3rdparty/tvm-ffi" + CIBW_BEFORE_BUILD_WINDOWS: >- + python -m pip install -U pip cmake ninja scikit-build-core wheel && + python -m pip install -v "{project}/3rdparty/tvm-ffi" + CIBW_ENVIRONMENT: >- + TVM_USE_LLVM="/opt/llvm/bin/llvm-config-static" + CMAKE_PREFIX_PATH="/opt/llvm" + CMAKE_ARGS="-DUSE_LLVM=/opt/llvm/bin/llvm-config-static -DUSE_CUDA=OFF -DBUILD_TESTING=OFF -DTVM_BUILD_PYTHON_MODULE=ON -DCMAKE_PREFIX_PATH=/opt/llvm" + TVM_CUDA_ARCHITECTURES="${{ inputs.cuda_architectures }}" + TVM_WHEEL_DIST_NAME="${{ inputs.distribution_name }}" + TVM_WHEEL_DIST_VERSION="${{ inputs.distribution_version }}" + TVM_INCLUDE_CUDA_RUNTIME="${{ steps.wheel_inputs.outputs.include_cuda_runtime }}" + TVM_CUDA_RUNTIME_PATH="${{ steps.wheel_inputs.outputs.include_cuda_runtime == '1' && inputs.cuda_runtime_path || '' }}" + TVM_AUDITWHEEL_PLAT="${{ steps.wheel_inputs.outputs.wheel_platform_tag }}" + TVM_BUILD_PARALLEL_LEVEL="${{ steps.env_vars.outputs.cpu_count }}" + TVM_EXPECT_LLVM_ENABLED=1 + TVM_EXPECT_STATIC_LLVM=1 + TVM_EXPECT_CUDA_RUNTIME="${{ steps.wheel_inputs.outputs.include_cuda_runtime }}" + TVM_EXPECT_CUDA_ENABLED="${{ steps.wheel_inputs.outputs.include_cuda_runtime != '1' && '0' || '' }}" + CIBW_ENVIRONMENT_WINDOWS: >- + TVM_USE_LLVM="C:/opt/llvm/Library/bin/llvm-config-static.bat" + CMAKE_PREFIX_PATH="C:/opt/llvm/Library" + CMAKE_ARGS="-DUSE_LLVM=C:/opt/llvm/Library/bin/llvm-config-static.bat -DUSE_CUDA=OFF -DBUILD_TESTING=OFF -DTVM_BUILD_PYTHON_MODULE=ON -DCMAKE_PREFIX_PATH=C:/opt/llvm/Library" + TVM_CUDA_ARCHITECTURES="${{ inputs.cuda_architectures }}" + TVM_WHEEL_DIST_NAME="${{ inputs.distribution_name }}" + TVM_WHEEL_DIST_VERSION="${{ inputs.distribution_version }}" + TVM_INCLUDE_CUDA_RUNTIME="${{ steps.wheel_inputs.outputs.include_cuda_runtime }}" + TVM_CUDA_RUNTIME_PATH="${{ steps.wheel_inputs.outputs.include_cuda_runtime == '1' && inputs.cuda_runtime_path || '' }}" + TVM_BUILD_PARALLEL_LEVEL="${{ steps.env_vars.outputs.cpu_count }}" + TVM_EXPECT_LLVM_ENABLED=1 + TVM_EXPECT_STATIC_LLVM=1 + TVM_EXPECT_CUDA_RUNTIME="${{ steps.wheel_inputs.outputs.include_cuda_runtime }}" + TVM_EXPECT_CUDA_ENABLED="${{ steps.wheel_inputs.outputs.include_cuda_runtime != '1' && '0' || '' }}" + CIBW_REPAIR_WHEEL_COMMAND_LINUX: >- + bash "{project}/ci/scripts/package/tvm_wheel_helper.sh" cibw-repair "{wheel}" "{dest_dir}" + CIBW_REPAIR_WHEEL_COMMAND_MACOS: >- + TVM_DELOCATE_ARCHS="{delocate_archs}" + bash "{project}/ci/scripts/package/tvm_wheel_helper.sh" cibw-repair "{wheel}" "{dest_dir}" + CIBW_REPAIR_WHEEL_COMMAND_WINDOWS: >- + python "{project}/ci/scripts/package/rewrite_wheel.py" "{wheel}" + --output-dir "{dest_dir}" + --distribution-name "${{ inputs.distribution_name }}" + --distribution-version "${{ inputs.distribution_version }}" + --extra-library-dir "C:/opt/llvm/Library/bin" + --extra-library-pattern "libxml2*.dll" + --extra-library-pattern "zstd*.dll" + --extra-library-pattern "zlib*.dll" + --extra-library-pattern "*lzma*.dll" + --extra-library-pattern "*iconv*.dll" + --extra-library-pattern "*charset*.dll" + CIBW_TEST_COMMAND: >- + python -u -X faulthandler "{project}/ci/scripts/package/verify_tvm_install.py" + CIBW_TEST_COMMAND_LINUX: >- + bash "{project}/ci/scripts/package/verify_tvm_linux.sh" + "{project}/ci/scripts/package/verify_tvm_install.py" + + - name: Verify final wheel + shell: bash -l {0} + env: + TVM_PYTHON: ${{ runner.os == 'Windows' && 'python' || 'python3' }} + TVM_EXPECT_WHEEL_PLATFORM_TAG: ${{ steps.wheel_inputs.outputs.wheel_platform_tag }} + TVM_EXPECT_LLVM_ENABLED: "1" + TVM_EXPECT_STATIC_LLVM: "1" + TVM_EXPECT_CUDA_RUNTIME: ${{ steps.wheel_inputs.outputs.include_cuda_runtime }} + TVM_EXPECT_CUDA_ENABLED: ${{ steps.wheel_inputs.outputs.include_cuda_runtime != '1' && '0' || '' }} + run: ci/scripts/package/tvm_wheel_helper.sh verify diff --git a/.github/actions/detect-env-vars/action.yml b/.github/actions/detect-env-vars/action.yml new file mode 100644 index 000000000000..e20b15746dac --- /dev/null +++ b/.github/actions/detect-env-vars/action.yml @@ -0,0 +1,38 @@ +# 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: Detect Environment Variables +description: Detects environment variables such as CPU count and sets them as outputs. +runs: + using: "composite" + steps: + - name: Run Python to detect environment variables + shell: python + id: detect + run: | + import multiprocessing + import os + + with open(os.environ["GITHUB_OUTPUT"], "a", encoding="utf-8") as output_file: + value = multiprocessing.cpu_count() + output_file.write(f"cpu_count={value}\n") + print(f"Detected environment variable: cpu_count={value}") + +outputs: + cpu_count: + description: "The number of CPU cores" + value: "${{ steps.detect.outputs.cpu_count }}" diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml index e78ce2f66d7a..c0fb205c111e 100644 --- a/.github/actions/setup/action.yml +++ b/.github/actions/setup/action.yml @@ -1,13 +1,13 @@ runs: using: "composite" steps: - - uses: actions/cache@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3 + - uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 env: CACHE_NUMBER: 2 with: path: ~/conda_pkgs_dir key: ${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-${{ hashFiles('tests/conda/build-environment.yaml') }} - - uses: conda-incubator/setup-miniconda@fc2d68f6413eb2d87b895e92f8584b5b94a10167 # v3.3.0 + - uses: conda-incubator/setup-miniconda@8ee1f361103df19b6f8c8655fd3967a8ecb162d5 # v4.0.1 continue-on-error: true id: conda1 with: @@ -18,7 +18,8 @@ runs: miniforge-version: latest python-version: "3.10" condarc-file: tests/conda/condarc - - uses: conda-incubator/setup-miniconda@fc2d68f6413eb2d87b895e92f8584b5b94a10167 # v3.3.0 + conda-remove-defaults: true + - uses: conda-incubator/setup-miniconda@8ee1f361103df19b6f8c8655fd3967a8ecb162d5 # v4.0.1 if: steps.conda1.outcome == 'failure' with: activate-environment: tvm-build @@ -29,6 +30,7 @@ runs: use-only-tar-bz2: true python-version: "3.10" condarc-file: tests/conda/condarc + conda-remove-defaults: true - name: Conda info shell: pwsh run: | diff --git a/.github/workflows/publish_wheel.yml b/.github/workflows/publish_wheel.yml new file mode 100644 index 000000000000..4ac52c25d91e --- /dev/null +++ b/.github/workflows/publish_wheel.yml @@ -0,0 +1,286 @@ +# 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: Publish TVM wheels + +on: + workflow_dispatch: + inputs: + tag: + description: "Tag, branch, or SHA to build; PyPI publishes require refs/tags/" + required: true + type: string + publish_repository: + description: "Where to publish after the wheel build succeeds" + required: true + default: "none" + type: choice + options: + - none + - testpypi + - pypi + distribution_name: + description: "Optional package name override, e.g. tvm-yourname-test for TestPyPI" + required: false + default: "" + type: string + distribution_version: + description: "Optional package version override for TestPyPI validation builds" + required: false + default: "" + type: string + cuda_architectures: + description: "CMake CUDA architectures for libtvm_runtime_cuda.so" + required: false + default: "75" + type: string + verify_from_repository: + description: "Install the uploaded package from the selected repository and import-test it" + required: true + default: true + type: boolean + +permissions: + contents: read + +jobs: + build_wheels: + name: ${{ matrix.name }} + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + include: + - name: "Linux x86_64 wheel with CUDA runtime (manylinux_2_28)" + os: ubuntu-latest + arch: x86_64 + build: cp310-manylinux_x86_64 + linux_image: manylinux_2_28 + linux_image_tag: "2026.01.04-1" + include_cuda_runtime: "true" + artifact_suffix: linux-x86_64-manylinux_2_28 + - name: "Linux aarch64 wheel with CUDA runtime (manylinux_2_28)" + os: ubuntu-24.04-arm + arch: aarch64 + build: cp310-manylinux_aarch64 + linux_image: manylinux_2_28 + linux_image_tag: "2026.01.04-1" + include_cuda_runtime: "true" + artifact_suffix: linux-aarch64-manylinux_2_28 + - name: "macOS arm64 CPU wheel" + os: macos-14 + arch: arm64 + build: cp310-macosx_arm64 + linux_image: "" + linux_image_tag: "" + include_cuda_runtime: "false" + artifact_suffix: macos-arm64 + - name: "Windows AMD64 CPU wheel" + os: windows-latest + arch: AMD64 + build: cp310-win_amd64 + linux_image: "" + linux_image_tag: "" + include_cuda_runtime: "false" + artifact_suffix: windows-amd64 + steps: + - name: Validate publish inputs + shell: bash + env: + TVM_PUBLISH_REPOSITORY: ${{ inputs.publish_repository }} + TVM_PUBLISH_REF: ${{ inputs.tag }} + TVM_VERIFY_FROM_REPOSITORY: ${{ inputs.verify_from_repository }} + TVM_WHEEL_DIST_NAME: ${{ inputs.distribution_name }} + TVM_WHEEL_DIST_VERSION: ${{ inputs.distribution_version }} + run: | + set -eux + if [[ -n "${TVM_WHEEL_DIST_NAME}" && ! "${TVM_WHEEL_DIST_NAME}" =~ ^[A-Za-z0-9]([A-Za-z0-9._-]*[A-Za-z0-9])?$ ]]; then + echo "distribution_name must be a valid Python package name override" >&2 + exit 1 + fi + if [[ -n "${TVM_WHEEL_DIST_VERSION}" && ! "${TVM_WHEEL_DIST_VERSION}" =~ ^[A-Za-z0-9][A-Za-z0-9._!+-]*$ ]]; then + echo "distribution_version must be a valid Python package version override" >&2 + exit 1 + fi + if [[ "${TVM_PUBLISH_REPOSITORY}" == "pypi" && -n "${TVM_WHEEL_DIST_NAME}" ]]; then + echo "distribution_name must be empty when publishing to PyPI" >&2 + exit 1 + fi + if [[ "${TVM_PUBLISH_REPOSITORY}" == "pypi" && -n "${TVM_WHEEL_DIST_VERSION}" ]]; then + echo "distribution_version must be empty when publishing to PyPI" >&2 + exit 1 + fi + if [[ "${TVM_PUBLISH_REPOSITORY}" == "pypi" && "${TVM_PUBLISH_REF}" != refs/tags/* ]]; then + echo "PyPI publishes must use an immutable refs/tags/ ref" >&2 + exit 1 + fi + if [[ "${TVM_PUBLISH_REPOSITORY}" == "pypi" && "${TVM_VERIFY_FROM_REPOSITORY}" != "true" ]]; then + echo "verify_from_repository must be enabled when publishing to PyPI" >&2 + exit 1 + fi + + - name: Checkout source + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + ref: ${{ inputs.tag }} + submodules: recursive + fetch-depth: 1 + fetch-tags: true + + - name: Build CUDA runtime + id: build_cuda + uses: ./.github/actions/build-cuda + with: + arch: ${{ matrix.arch }} + linux_image: ${{ matrix.linux_image }} + linux_image_tag: ${{ matrix.linux_image_tag }} + cuda_architectures: ${{ inputs.cuda_architectures }} + include_cuda_runtime: ${{ matrix.include_cuda_runtime }} + + - name: Build TVM wheel + uses: ./.github/actions/build-wheel-for-publish + with: + arch: ${{ matrix.arch }} + build: ${{ matrix.build }} + linux_image: ${{ matrix.linux_image }} + linux_image_tag: ${{ matrix.linux_image_tag }} + distribution_name: ${{ inputs.distribution_name }} + distribution_version: ${{ inputs.distribution_version }} + cuda_architectures: ${{ inputs.cuda_architectures }} + include_cuda_runtime: ${{ matrix.include_cuda_runtime }} + cuda_runtime_path: ${{ steps.build_cuda.outputs.cuda_runtime_path }} + + - name: Upload wheel artifact + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: tvm-wheel-${{ matrix.artifact_suffix }} + path: wheelhouse/*.whl + if-no-files-found: error + + upload_pypi: + name: Upload package distributions + needs: [build_wheels] + if: ${{ inputs.publish_repository != 'none' }} + runs-on: ubuntu-latest + environment: ${{ inputs.publish_repository }} + permissions: + actions: read + contents: read + id-token: write + attestations: write + steps: + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + pattern: tvm-wheel-* + path: dist + merge-multiple: true + + - name: Check wheel sizes + shell: bash + run: | + set -euo pipefail + limit_bytes=100000000 + shopt -s nullglob + wheels=(dist/*.whl) + if [[ "${#wheels[@]}" -eq 0 ]]; then + echo "No wheel artifacts found under dist/" >&2 + exit 1 + fi + if [[ "${#wheels[@]}" -ne 4 ]]; then + echo "Expected 4 wheel artifacts, found ${#wheels[@]}" >&2 + printf '%s\n' "${wheels[@]}" >&2 + exit 1 + fi + failed=0 + for wheel in "${wheels[@]}"; do + size="$(stat -c '%s' "$wheel")" + printf '%s %s bytes\n' "$wheel" "$size" + if (( size > limit_bytes )); then + echo "Wheel exceeds 100 MB PyPI/TestPyPI upload limit: ${wheel}" >&2 + failed=1 + fi + done + exit "$failed" + + - name: Generate artifact attestation for wheels + uses: actions/attest-build-provenance@v1 + with: + subject-path: dist/* + + - name: Publish package distributions to TestPyPI + if: ${{ inputs.publish_repository == 'testpypi' }} + uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0 + with: + attestations: true + verbose: true + repository-url: https://test.pypi.org/legacy/ + + - name: Publish package distributions to PyPI + if: ${{ inputs.publish_repository == 'pypi' }} + uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0 + with: + attestations: true + verbose: true + + verify_pypi: + name: Verify uploaded package + needs: [upload_pypi] + if: ${{ inputs.publish_repository != 'none' && inputs.verify_from_repository }} + runs-on: ubuntu-latest + permissions: + actions: read + contents: read + steps: + - name: Check out source + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + ref: ${{ inputs.tag }} + submodules: recursive + fetch-depth: 0 + fetch-tags: true + + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: tvm-wheel-linux-x86_64-manylinux_2_28 + path: wheelhouse + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Verify package from TestPyPI + if: ${{ inputs.publish_repository == 'testpypi' }} + env: + TVM_PYTHON: python + TVM_TEST_INDEX_URL: https://test.pypi.org/simple/ + TVM_EXTRA_INDEX_URL: https://pypi.org/simple + TVM_EXPECT_LLVM_ENABLED: "1" + TVM_EXPECT_STATIC_LLVM: "1" + TVM_EXPECT_CUDA_RUNTIME: "1" + run: ci/scripts/package/tvm_wheel_helper.sh verify-pypi + + - name: Verify package from PyPI + if: ${{ inputs.publish_repository == 'pypi' }} + env: + TVM_PYTHON: python + TVM_TEST_INDEX_URL: https://pypi.org/simple/ + TVM_EXTRA_INDEX_URL: https://pypi.org/simple + TVM_EXPECT_LLVM_ENABLED: "1" + TVM_EXPECT_STATIC_LLVM: "1" + TVM_EXPECT_CUDA_RUNTIME: "1" + run: ci/scripts/package/tvm_wheel_helper.sh verify-pypi diff --git a/.gitignore b/.gitignore index 9e734b0be06d..0ee1eb241807 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,8 @@ __pycache__/ env/ build/ build-*/ +!.github/actions/build-*/ +!.github/actions/build-*/action.yml develop-eggs/ dist/ downloads/ diff --git a/CMakeLists.txt b/CMakeLists.txt index a11a8729700f..422d7112a6ea 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -104,6 +104,24 @@ tvm_option(USE_NVSHMEM "Build with NVSHMEM support" OFF) # Python package options tvm_option(TVM_BUILD_PYTHON_MODULE "Build Python module with scikit-build-core" OFF) +function(tvm_set_python_module_relative_rpath target_name) + if(NOT TVM_BUILD_PYTHON_MODULE OR NOT TARGET ${target_name}) + return() + endif() + + if(APPLE) + set_target_properties(${target_name} PROPERTIES + BUILD_RPATH "@loader_path" + INSTALL_RPATH "@loader_path" + ) + elseif(UNIX) + set_target_properties(${target_name} PROPERTIES + BUILD_RPATH "\$ORIGIN" + INSTALL_RPATH "\$ORIGIN" + ) + endif() +endfunction() + # include directories include_directories(${CMAKE_INCLUDE_PATH}) include_directories("include") @@ -531,6 +549,7 @@ set_target_properties(tvm_runtime_extra PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ) +tvm_set_python_module_relative_rpath(tvm_runtime_extra) install(TARGETS tvm_runtime_extra DESTINATION lib${LIB_SUFFIX}) if(TVM_BUILD_PYTHON_MODULE) @@ -592,6 +611,20 @@ target_include_directories(tvm_compiler PUBLIC "$/lib so that # tvm_ffi.libinfo.load_lib_ctypes(package="tvm", target_name=...) can find @@ -865,9 +893,11 @@ if(TVM_BUILD_PYTHON_MODULE) # Install third-party compiled dependencies into the same lib/ dir. if(TARGET fpA_intB_gemm) + tvm_set_python_module_relative_rpath(fpA_intB_gemm) install(TARGETS fpA_intB_gemm DESTINATION "lib") endif() if(TARGET flash_attn) + tvm_set_python_module_relative_rpath(flash_attn) install(TARGETS flash_attn DESTINATION "lib") endif() diff --git a/ci/scripts/package/README.md b/ci/scripts/package/README.md new file mode 100644 index 000000000000..3602f15f9737 --- /dev/null +++ b/ci/scripts/package/README.md @@ -0,0 +1,162 @@ + + + + + + + + + + + + + + + + + +# TVM wheel packaging helper + +This directory contains the helper scripts used to build, repair, verify, and +publish TVM Python wheels. The GitHub Actions workflow keeps orchestration in +YAML and puts platform-specific packaging behavior in focused composite actions +and shell/Python helpers. + +The wheel build flow is: + +1. Optionally build `libtvm_runtime_cuda.so` in a CUDA-enabled Linux CMake build. +2. Build the main Python wheel with `cibuildwheel`, LLVM enabled, and CUDA + disabled. +3. When requested, inject the CUDA runtime DSO into `tvm/lib/` during the + `cibuildwheel` repair hook. +4. Repair the wheel, excluding CUDA toolkit/driver DSOs and `libtvm_ffi`. + `libtvm_runtime_cuda.so`, when requested, is the TVM CUDA runtime that is + intentionally injected into the wheel. + On Windows, copy the small runtime DLLs required by LLVM support libraries + into `tvm/lib/` because there is no auditwheel-style repair tool. +5. Validate ELF links so intra-wheel TVM DSOs resolve through relative rpaths. + LLVM is expected to be linked statically; the final wheel must not bundle + or dynamically depend on `libLLVM`. +6. Verify the wheel in a fresh virtualenv. +7. Optionally upload and verify the uploaded package. + +GitHub Actions flow: + +1. The `Publish TVM wheels` workflow builds a platform wheel matrix: + - Linux x86_64 in a pinned `manylinux_2_28` container, with the CUDA runtime. + - Linux aarch64 in a pinned `manylinux_2_28` container, with the CUDA runtime. + - macOS arm64 CPU-only. + - Windows AMD64 CPU-only. +2. The Linux CUDA runtime action exposes the built DSO path as an action output. + The wheel action receives that path explicitly and mounts it into the + `cibuildwheel` container for the repair hook. +3. The optional publishing jobs upload the artifacts and can verify the package + from the selected package index. PyPI publishing requires a `refs/tags/` + input and keeps post-upload verification enabled. + +Linux wheels are built inside manylinux images. This avoids accidentally +publishing a wheel tagged for the GitHub runner's host glibc, such as +`manylinux_2_39`, which would not install on older supported Linux systems. + +Workflow structure: + +- `.github/workflows/publish_wheel.yml`: defines the platform matrix, + artifact upload, optional publishing, and post-upload verification. +- `.github/actions/detect-env-vars`: shared environment detection. +- `.github/actions/build-cuda`: builds only the optional CUDA runtime library. + On Linux this action owns the pinned manylinux Docker/CUDA setup and exposes + the runtime DSO path as an action output. +- `.github/actions/build-wheel-for-publish`: installs the cached LLVM prefix + and runs `pypa/cibuildwheel` for the LLVM-enabled runtime wheel. Its custom + repair hook injects the CUDA runtime before `auditwheel`/`delocate`/Windows + dependency-copy repair. +- `ci/scripts/package/tvm_wheel_helper.sh`: implements reusable local and CI + entrypoints around the `cibuildwheel` build, such as `cuda`, + `manylinux-cuda`, `cibw-repair`, `verify`, `upload`, and `verify-pypi`. +- `ci/scripts/package/rewrite_wheel.py`: rewrites wheel metadata and injects + extra runtime files, including the CUDA runtime library when CUDA is enabled. +- `ci/scripts/package/verify_tvm_install.py`: imports the installed wheel and + checks that the runtime library was loaded from the wheel, expected runtime + DSOs are present, dynamic LLVM libraries are not bundled when static LLVM is + required, and minimal TIRX/Relax programs compile and run through LLVM. + +To test the workflow from a fork without publishing: + +```bash +git push origin HEAD: +git tag -a tvm-wheel-test0 -m "Test TVM wheel workflow" +git push origin tvm-wheel-test0 + +gh workflow run publish_wheel.yml \ + --repo / \ + --ref \ + -f tag=tvm-wheel-test0 \ + -f publish_repository=none \ + -f distribution_name= \ + -f cuda_architectures=75 \ + -f verify_from_repository=false +``` + +If the workflow is not visible in the GitHub UI yet, push or merge these files +to the fork's default branch first. GitHub only lists manually dispatched +workflows once the workflow file exists in the repository. + +Local debugging: + +The main wheel build is owned by `cibuildwheel`. The shell helper is used for +the build pieces around `cibuildwheel`: CUDA runtime construction, the +`CIBW_REPAIR_WHEEL_COMMAND` hook, final wheel verification, and optional +publish verification. + +For the exact `cibuildwheel` environment, use +`.github/actions/build-wheel-for-publish/action.yml` as the source of truth. +For local checks after a wheel exists under `wheelhouse/`, run: + +```bash +TVM_PYTHON=/tmp/tvm-wheel-tools/bin/python \ +ci/scripts/package/tvm_wheel_helper.sh verify + +TVM_PYTHON=/tmp/tvm-wheel-tools/bin/python \ +ci/scripts/package/tvm_wheel_helper.sh verify-pypi +``` + +For a manual or local upload with the helper, leave `TVM_WHEEL_DIST_NAME` +unset and set the normal Twine credentials: + +```bash +TWINE_USERNAME=__token__ \ +TWINE_PASSWORD="$PYPI_TOKEN" \ +TVM_PYTHON=/tmp/tvm-wheel-tools/bin/python \ +ci/scripts/package/tvm_wheel_helper.sh upload +``` + +Useful knobs: + +- `TVM_USE_LLVM`: LLVM config for the CIBW build and repair helpers, default + `llvm-config --link-static`. +- `TVM_USE_CUDA`: CUDA root or `ON` for the CUDA build, default `ON`. +- `TVM_CUDA_RUNTIME_PATH`: explicit path to `libtvm_runtime_cuda.so` for repair. +- `TVM_CUDA_ARCHITECTURES`: CMake CUDA architectures, default `75`. +- `TVM_WHEEL_DIST_NAME`: optional distribution rename for TestPyPI. +- `TVM_WHEEL_DIST_VERSION`: optional distribution version rewrite. +- `TVM_INCLUDE_CUDA_RUNTIME=1`: build or repair a wheel with the CUDA runtime. + Do not set this to a value that conflicts with `TVM_SKIP_CUDA`. +- `TVM_SKIP_REPAIR=1`: leave the injected wheel unrepaired. +- `TVM_SKIP_CUDA=1`: build or repair a wheel without the CUDA runtime. +- `TVM_KEEP_BUILD_DIRS=1`: reuse the CMake build directories. +- `TVM_MANYLINUX_IMAGE`: manylinux image family for `manylinux-cuda`, such as + `manylinux_2_28`. +- `TVM_MANYLINUX_IMAGE_TAG`: pinned manylinux image tag for `manylinux-cuda`. +- `TVM_ARCH`: target architecture for `manylinux-cuda`, such as `x86_64` or + `aarch64`. +- `TVM_AUDITWHEEL_PLAT`: optional `auditwheel repair --plat` override. +- `TVM_AUDITWHEEL_LIBRARY_PATH`: optional, explicit library search path for + `auditwheel repair`. +- `TVM_EXPECT_WHEEL_PLATFORM_TAG`: require the final wheel filename to include + a specific platform tag, such as `manylinux_2_28_x86_64`. +- `TVM_EXPECT_CUDA_RUNTIME`: verify whether the installed wheel ships a CUDA + runtime library. +- `TVM_EXPECT_STATIC_LLVM`: verify that the installed wheel does not ship a + dynamic LLVM library. +- `TVM_TEST_INDEX_URL`: package index for `verify-pypi`, default TestPyPI. +- `TVM_EXTRA_INDEX_URL`: extra package index for dependencies, default PyPI. diff --git a/ci/scripts/package/rewrite_wheel.py b/ci/scripts/package/rewrite_wheel.py new file mode 100755 index 000000000000..7584edfcdf8c --- /dev/null +++ b/ci/scripts/package/rewrite_wheel.py @@ -0,0 +1,399 @@ +#!/usr/bin/env python3 +# 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. + +"""Rewrite TVM wheel metadata and inject extra runtime files.""" + +from __future__ import annotations + +import argparse +import base64 +import csv +import hashlib +import io +import re +import shutil +import subprocess +import sys +import tempfile +import zipfile +from email.parser import Parser +from pathlib import Path + + +def _wheel_escape_distribution(value: str) -> str: + """Escape a distribution component for wheel filenames and dist-info dirs.""" + + return re.sub(r"[^\w\d.]+", "_", value).lower() + + +def _wheel_escape_version(value: str) -> str: + """Escape a version component while preserving PEP 440 local version markers.""" + + return re.sub(r"[^\w\d.!+]+", "_", value).lower() + + +def _hash_record(data: bytes) -> tuple[str, str]: + digest = hashlib.sha256(data).digest() + encoded = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii") + return f"sha256={encoded}", str(len(data)) + + +def _copy_info(info: zipfile.ZipInfo, filename: str) -> zipfile.ZipInfo: + copied = zipfile.ZipInfo(filename=filename, date_time=info.date_time) + copied.compress_type = info.compress_type + copied.comment = info.comment + copied.extra = info.extra + copied.internal_attr = info.internal_attr + copied.external_attr = info.external_attr + return copied + + +def _replace_header(metadata: bytes, key: str, value: str) -> bytes: + text = metadata.decode("utf-8") + lines = text.splitlines(keepends=True) + prefix = f"{key.lower()}:" + for index, line in enumerate(lines): + if line.lower().startswith(prefix): + newline = "\r\n" if line.endswith("\r\n") else "\n" + lines[index] = f"{key}: {value}{newline}" + return "".join(lines).encode("utf-8") + raise ValueError(f"METADATA does not contain {key!r}") + + +def _find_dist_info(names: list[str]) -> str: + dist_infos = sorted({name.split("/", 1)[0] for name in names if ".dist-info/" in name}) + dist_infos = [name for name in dist_infos if name.endswith(".dist-info")] + if len(dist_infos) != 1: + raise ValueError(f"Expected one .dist-info directory, found {dist_infos}") + return dist_infos[0] + + +def _metadata_headers(metadata: bytes) -> tuple[str, str]: + headers = Parser().parsestr(metadata.decode("utf-8")) + name = headers.get("Name") + version = headers.get("Version") + if not name or not version: + raise ValueError("METADATA must contain Name and Version") + return name, version + + +def _is_elf_shared_lib(name: str, data: bytes) -> bool: + return ( + name.startswith("tvm/lib/") + and re.search(r"\.so(?:\.|$)", Path(name).name) is not None + and data.startswith(b"\x7fELF") + ) + + +def _set_rpath(data: bytes, rpath: str, name: str) -> bytes: + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / Path(name).name + path.write_bytes(data) + path.chmod(0o755) + try: + subprocess.run(["patchelf", "--set-rpath", rpath, str(path)], check=True) + except subprocess.CalledProcessError as err: + raise ValueError(f"patchelf failed while setting rpath on {name}") from err + return path.read_bytes() + + +def _retag_wheel_filename( + wheel: Path, + dist_name: str, + version: str, +) -> str: + parts = wheel.name.removesuffix(".whl").split("-") + if len(parts) not in (5, 6): + raise ValueError(f"Unsupported wheel filename: {wheel.name}") + tags = parts[2:] + return ( + f"{_wheel_escape_distribution(dist_name)}-" + f"{_wheel_escape_version(version)}-" + f"{'-'.join(tags)}.whl" + ) + + +def _normalize_wheel_path(value: str, label: str) -> str: + raw = value.replace("\\", "/") + if raw.startswith("/") or re.match(r"^[A-Za-z]:", raw) is not None: + raise argparse.ArgumentTypeError( + f"{label} must be a relative wheel path without drive, empty, '.' or '..' segments" + ) + normalized = raw + parts = normalized.split("/") + if ( + not normalized + or any(part in {"", ".", ".."} for part in parts) + ): + raise argparse.ArgumentTypeError( + f"{label} must be a relative wheel path without drive, empty, '.' or '..' segments" + ) + return normalized + + +def _validate_wheel_member_path(value: str) -> str: + if "\\" in value: + raise ValueError(f"Wheel member path must use forward slashes: {value}") + try: + normalized = _normalize_wheel_path(value, "wheel member path") + except argparse.ArgumentTypeError as err: + raise ValueError(str(err)) from err + if normalized != value: + raise ValueError(f"Wheel member path is not normalized: {value}") + return normalized + + +def _parse_extra_file(value: str) -> tuple[Path, str]: + if "=" not in value: + raise argparse.ArgumentTypeError("extra files must use SOURCE=TARGET format") + source, target = value.split("=", 1) + if not source or not target: + raise argparse.ArgumentTypeError("extra files must use SOURCE=TARGET format") + return Path(source), _normalize_wheel_path(target, "extra file target") + + +def _extra_library_files( + library_dirs: list[Path], + patterns: list[str], + target_dir: str, +) -> list[tuple[Path, str]]: + target_dir = _normalize_wheel_path(target_dir, "extra library target dir") + extra_files: list[tuple[Path, str]] = [] + missing_dirs = [str(library_dir) for library_dir in library_dirs if not library_dir.is_dir()] + if missing_dirs: + raise ValueError(f"extra library dirs do not exist: {', '.join(missing_dirs)}") + for library_dir in library_dirs: + for pattern in patterns: + for source in sorted(library_dir.glob(pattern)): + if source.is_file(): + extra_files.append((source, f"{target_dir}/{source.name}")) + if library_dirs and patterns and not extra_files: + raise ValueError( + "extra library patterns did not match any files: " + ", ".join(patterns) + ) + return sorted(extra_files, key=lambda item: (item[1], str(item[0]))) + + +def _check_duplicate_targets(targets: list[str]) -> None: + seen: set[str] = set() + duplicates: set[str] = set() + for target in targets: + if target in seen: + duplicates.add(target) + seen.add(target) + if duplicates: + joined = ", ".join(sorted(duplicates)) + raise ValueError(f"Duplicate wheel target paths are not allowed: {joined}") + + +def rewrite_wheel( + wheel: Path, + output_dir: Path, + cuda_runtime: Path | None, + target_path: str, + distribution_name: str | None, + distribution_version: str | None, + set_rpath: str | None, + extra_files: list[tuple[Path, str]], +) -> Path: + output_dir.mkdir(parents=True, exist_ok=True) + extra_targets = {target for _, target in extra_files} + with zipfile.ZipFile(wheel, "r") as zin: + original_infos = [info for info in zin.infolist() if not info.is_dir()] + original_names = [_validate_wheel_member_path(info.filename) for info in original_infos] + _check_duplicate_targets(original_names) + original_dist_info = _find_dist_info(original_names) + metadata_path = f"{original_dist_info}/METADATA" + if metadata_path not in original_names: + raise ValueError(f"Wheel metadata is missing: {metadata_path}") + original_name, original_version = _metadata_headers(zin.read(metadata_path)) + + final_name = distribution_name or original_name + final_version = distribution_version or original_version + final_dist_info = ( + f"{_wheel_escape_distribution(final_name)}-" + f"{_wheel_escape_version(final_version)}.dist-info" + ) + record_path = f"{final_dist_info}/RECORD" + target_paths = [target for _, target in extra_files] + if cuda_runtime is not None: + target_paths.append(target_path) + target_paths.append(record_path) + _check_duplicate_targets(target_paths) + output_path = output_dir / _retag_wheel_filename(wheel, final_name, final_version) + + entries: list[tuple[zipfile.ZipInfo, bytes]] = [] + entry_names: list[str] = [] + for info in original_infos: + mapped_name = info.filename + if mapped_name == f"{original_dist_info}/RECORD": + continue + if mapped_name.startswith(f"{original_dist_info}/"): + mapped_name = f"{final_dist_info}/{mapped_name.split('/', 1)[1]}" + mapped_name = _validate_wheel_member_path(mapped_name) + if ( + cuda_runtime is not None and mapped_name == target_path + ) or mapped_name in extra_targets: + continue + + data = zin.read(info) + if mapped_name == f"{final_dist_info}/METADATA": + if distribution_name is not None: + data = _replace_header(data, "Name", final_name) + if distribution_version is not None: + data = _replace_header(data, "Version", final_version) + if set_rpath is not None and _is_elf_shared_lib(mapped_name, data): + data = _set_rpath(data, set_rpath, mapped_name) + entries.append((_copy_info(info, mapped_name), data)) + entry_names.append(mapped_name) + + if cuda_runtime is not None: + data = cuda_runtime.read_bytes() + if set_rpath is not None and _is_elf_shared_lib(target_path, data): + data = _set_rpath(data, set_rpath, target_path) + info = zipfile.ZipInfo(target_path) + info.compress_type = zipfile.ZIP_DEFLATED + info.external_attr = 0o644 << 16 + entries.append((info, data)) + entry_names.append(target_path) + + for source, target in extra_files: + data = source.read_bytes() + if set_rpath is not None and _is_elf_shared_lib(target, data): + data = _set_rpath(data, set_rpath, target) + info = zipfile.ZipInfo(target) + info.compress_type = zipfile.ZIP_DEFLATED + info.external_attr = 0o644 << 16 + entries.append((info, data)) + entry_names.append(target) + + _check_duplicate_targets([*entry_names, record_path]) + + record_buffer = io.StringIO() + writer = csv.writer(record_buffer, lineterminator="\n") + for info, data in entries: + digest, size = _hash_record(data) + writer.writerow([info.filename, digest, size]) + writer.writerow([record_path, "", ""]) + + record_info = zipfile.ZipInfo(record_path) + record_info.compress_type = zipfile.ZIP_DEFLATED + record_info.external_attr = 0o644 << 16 + entries.append((record_info, record_buffer.getvalue().encode("utf-8"))) + + with zipfile.ZipFile(output_path, "w", compression=zipfile.ZIP_DEFLATED) as zout: + for info, data in entries: + zout.writestr(info, data) + return output_path + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("wheel", type=Path) + parser.add_argument("--cuda-runtime", type=Path) + parser.add_argument("--target-path", default=None) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--distribution-name") + parser.add_argument("--distribution-version") + parser.add_argument("--set-rpath") + parser.add_argument( + "--extra-file", + action="append", + default=[], + type=_parse_extra_file, + help="Additional file to place in the wheel, using SOURCE=TARGET format.", + ) + parser.add_argument( + "--extra-library-dir", + action="append", + default=[], + type=Path, + help="Directory to scan for extra runtime libraries.", + ) + parser.add_argument( + "--extra-library-pattern", + action="append", + default=[], + help="Glob pattern for files under --extra-library-dir.", + ) + parser.add_argument( + "--extra-library-target-dir", + default="tvm/lib", + help="Wheel directory for files matched by --extra-library-pattern.", + ) + args = parser.parse_args() + + cuda_runtime = args.cuda_runtime + if cuda_runtime is not None and not cuda_runtime.is_file(): + parser.error(f"CUDA runtime DSO does not exist: {cuda_runtime}") + if args.set_rpath and shutil.which("patchelf") is None: + parser.error("--set-rpath requires patchelf on PATH") + + target_path = args.target_path + if target_path is None: + if cuda_runtime is None: + target_path = "tvm/lib/libtvm_runtime_cuda.so" + else: + target_path = f"tvm/lib/{cuda_runtime.name}" + else: + try: + target_path = _normalize_wheel_path(target_path, "target path") + except argparse.ArgumentTypeError as err: + parser.error(str(err)) + + extra_files = list(args.extra_file) + try: + extra_files.extend( + _extra_library_files( + library_dirs=args.extra_library_dir, + patterns=args.extra_library_pattern, + target_dir=args.extra_library_target_dir, + ) + ) + except (argparse.ArgumentTypeError, ValueError) as err: + parser.error(str(err)) + missing_extra_files = [str(source) for source, _ in extra_files if not source.is_file()] + if missing_extra_files: + parser.error(f"extra files do not exist: {', '.join(missing_extra_files)}") + target_paths = [target for _, target in extra_files] + if cuda_runtime is not None: + target_paths.append(target_path) + try: + _check_duplicate_targets(target_paths) + except ValueError as err: + parser.error(str(err)) + + try: + output_path = rewrite_wheel( + wheel=args.wheel, + output_dir=args.output_dir, + cuda_runtime=cuda_runtime, + target_path=target_path, + distribution_name=args.distribution_name or None, + distribution_version=args.distribution_version or None, + set_rpath=args.set_rpath, + extra_files=extra_files, + ) + except (ValueError, zipfile.BadZipFile, KeyError) as err: + parser.error(str(err)) + print(output_path) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/ci/scripts/package/tvm_wheel_helper.sh b/ci/scripts/package/tvm_wheel_helper.sh new file mode 100755 index 000000000000..7752d39b08c5 --- /dev/null +++ b/ci/scripts/package/tvm_wheel_helper.sh @@ -0,0 +1,627 @@ +#!/usr/bin/env bash +# 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. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" + +TVM_PYTHON="${TVM_PYTHON:-python}" +TVM_WHEELHOUSE="${TVM_WHEELHOUSE:-${REPO_ROOT}/wheelhouse}" +TVM_CUDA_BUILD_DIR="${TVM_CUDA_BUILD_DIR:-${REPO_ROOT}/build-wheel-cuda}" +TVM_CUDA_RUNTIME_PATH="${TVM_CUDA_RUNTIME_PATH:-}" +TVM_USE_LLVM="${TVM_USE_LLVM:-llvm-config --link-static}" +TVM_USE_CUDA="${TVM_USE_CUDA:-ON}" +TVM_CUDA_ARCHITECTURES="${TVM_CUDA_ARCHITECTURES:-75}" +TVM_BUILD_PARALLEL_LEVEL="${TVM_BUILD_PARALLEL_LEVEL:-${CMAKE_BUILD_PARALLEL_LEVEL:-$(getconf _NPROCESSORS_ONLN 2>/dev/null || echo 4)}}" +TVM_WHEEL_DIST_NAME="${TVM_WHEEL_DIST_NAME:-}" +TVM_WHEEL_DIST_VERSION="${TVM_WHEEL_DIST_VERSION:-}" +TVM_INCLUDE_CUDA_RUNTIME="${TVM_INCLUDE_CUDA_RUNTIME:-}" + +normalize_bool() { + local name="$1" + local value="$2" + local normalized + normalized="$(printf '%s' "$value" | tr '[:upper:]' '[:lower:]')" + case "$normalized" in + 1|true|yes|on) echo 1 ;; + 0|false|no|off) echo 0 ;; + *) + echo "error: ${name} must be a boolean value" >&2 + return 1 + ;; + esac +} + +if [[ -n "$TVM_INCLUDE_CUDA_RUNTIME" ]]; then + tvm_include_cuda_runtime_normalized="$(normalize_bool TVM_INCLUDE_CUDA_RUNTIME "$TVM_INCLUDE_CUDA_RUNTIME")" + if [[ -n "${TVM_SKIP_CUDA+x}" ]]; then + tvm_skip_cuda_normalized="$(normalize_bool TVM_SKIP_CUDA "$TVM_SKIP_CUDA")" + if [[ "$tvm_include_cuda_runtime_normalized" == "$tvm_skip_cuda_normalized" ]]; then + echo "error: TVM_INCLUDE_CUDA_RUNTIME conflicts with TVM_SKIP_CUDA" >&2 + exit 1 + fi + fi + if [[ "$tvm_include_cuda_runtime_normalized" == "1" ]]; then + TVM_SKIP_CUDA=0 + else + TVM_SKIP_CUDA=1 + fi +else + TVM_SKIP_CUDA="${TVM_SKIP_CUDA:-0}" +fi +TVM_SKIP_CUDA="$(normalize_bool TVM_SKIP_CUDA "$TVM_SKIP_CUDA")" +TVM_SKIP_REPAIR="${TVM_SKIP_REPAIR:-0}" +TVM_KEEP_BUILD_DIRS="${TVM_KEEP_BUILD_DIRS:-0}" +TVM_SKIP_REPAIR="$(normalize_bool TVM_SKIP_REPAIR "$TVM_SKIP_REPAIR")" +TVM_KEEP_BUILD_DIRS="$(normalize_bool TVM_KEEP_BUILD_DIRS "$TVM_KEEP_BUILD_DIRS")" + +usage() { + cat <<'EOF' +Usage: ci/scripts/package/tvm_wheel_helper.sh [cuda|cuda-path|manylinux-cuda|cibw-repair|validate|verify|verify-installed|upload|verify-pypi] + +Environment knobs: + TVM_USE_LLVM LLVM config used by repair helpers, default "llvm-config --link-static" + TVM_USE_CUDA CUDA root or ON for the CUDA build, default ON + TVM_CUDA_RUNTIME_PATH Explicit libtvm_runtime_cuda.so path for repair + TVM_CUDA_ARCHITECTURES CMake CUDA arch list, default 75 + TVM_WHEEL_DIST_NAME Optional distribution rename for TestPyPI + TVM_WHEEL_DIST_VERSION Optional distribution version rewrite + TVM_UPLOAD_REPOSITORY_URL Twine repository URL, e.g. TestPyPI legacy URL + TVM_INCLUDE_CUDA_RUNTIME=1 Build/inject libtvm_runtime_cuda.so + TVM_SKIP_CUDA=1 Do not build/inject libtvm_runtime_cuda.so + TVM_SKIP_REPAIR=1 Keep injected wheel as final wheel + TVM_KEEP_BUILD_DIRS=1 Reuse CMake build dirs instead of cleaning them + TVM_MANYLINUX_IMAGE manylinux image tag for manylinux-cuda + TVM_MANYLINUX_IMAGE_TAG pinned image tag for manylinux-cuda + TVM_ARCH Target architecture for manylinux-cuda + TVM_AUDITWHEEL_PLAT Optional auditwheel --plat value + TVM_AUDITWHEEL_LIBRARY_PATH Optional library search path for auditwheel repair + TVM_EXPECT_WHEEL_PLATFORM_TAG + Require the final wheel filename to include this tag + TVM_EXPECT_CUDA_RUNTIME Verify whether the installed wheel ships a CUDA runtime DSO + TVM_EXPECT_STATIC_LLVM Verify that the installed wheel does not ship libLLVM + TVM_DELOCATE_ARCHS Optional delocate --require-archs value for macOS repair + TVM_TEST_INDEX_URL Package index for verify-pypi, default TestPyPI + TVM_EXTRA_INDEX_URL Extra package index for dependencies, default PyPI +EOF +} + +require_cmd() { + if ! command -v "$1" >/dev/null 2>&1; then + echo "error: required command not found: $1" >&2 + return 1 + fi +} + +single_wheel() { + local dir="$1" + local wheels=() + local wheel + while IFS= read -r wheel; do + wheels+=("$wheel") + done < <(find "$dir" -maxdepth 1 -type f -name '*.whl' | sort) + if [[ "${#wheels[@]}" -ne 1 ]]; then + echo "error: expected exactly one wheel under ${dir}, found ${#wheels[@]}" >&2 + printf '%s\n' "${wheels[@]}" >&2 + return 1 + fi + echo "${wheels[0]}" +} + +wheel_metadata_field() { + local wheel="$1" + local field="$2" + "$TVM_PYTHON" - "$wheel" "$field" <<'PY' +from email.parser import Parser +from pathlib import Path +import sys +import zipfile + +wheel = Path(sys.argv[1]) +field = sys.argv[2] +with zipfile.ZipFile(wheel) as zf: + metadata_name = next(name for name in zf.namelist() if name.endswith(".dist-info/METADATA")) + metadata = Parser().parsestr(zf.read(metadata_name).decode("utf-8")) +print(metadata[field]) +PY +} + +cuda_runtime_path() { + if [[ -n "$TVM_CUDA_RUNTIME_PATH" ]]; then + if [[ -f "$TVM_CUDA_RUNTIME_PATH" ]]; then + echo "$TVM_CUDA_RUNTIME_PATH" + else + echo "error: TVM_CUDA_RUNTIME_PATH does not exist: ${TVM_CUDA_RUNTIME_PATH}" >&2 + return 1 + fi + return 0 + fi + if [[ ! -d "$TVM_CUDA_BUILD_DIR" ]]; then + return 0 + fi + find "$TVM_CUDA_BUILD_DIR" -type f -name 'libtvm_runtime_cuda.so' | sort | tail -n 1 +} + +manylinux_image_name() { + local base="$1" + local arch="$2" + local tag="${3:-}" + if [[ "$base" == *"/"* || "$base" == *":"* ]]; then + if [[ "$base" != *@sha256:* && "${base##*/}" != *":"* ]]; then + echo "error: fully qualified TVM_MANYLINUX_IMAGE must include a tag or digest" >&2 + return 1 + fi + echo "$base" + elif [[ -n "$tag" ]]; then + echo "quay.io/pypa/${base}_${arch}:${tag}" + else + echo "error: TVM_MANYLINUX_IMAGE_TAG is required when TVM_MANYLINUX_IMAGE is not fully qualified" >&2 + return 1 + fi +} + +validate_manylinux_cuda_image() { + local image="$1" + local image_name="${image##*/}" + if [[ "$image_name" != manylinux_2_28_*:* && "$image_name" != manylinux_2_28_*@sha256:* ]]; then + echo "error: manylinux-cuda currently supports only pinned manylinux_2_28 images" >&2 + return 1 + fi +} + +run_manylinux_cuda_container() { + if [[ "$TVM_SKIP_CUDA" == "1" ]]; then + echo "Skipping manylinux CUDA build because TVM_SKIP_CUDA=1" + return 0 + fi + + require_cmd docker + require_cmd curl + if [[ -z "${TVM_MANYLINUX_IMAGE:-}" ]]; then + echo "error: TVM_MANYLINUX_IMAGE is required for manylinux-cuda" >&2 + return 1 + fi + if [[ -z "${TVM_ARCH:-}" ]]; then + echo "error: TVM_ARCH is required for manylinux-cuda" >&2 + return 1 + fi + + local image + image="$(manylinux_image_name "$TVM_MANYLINUX_IMAGE" "$TVM_ARCH" "${TVM_MANYLINUX_IMAGE_TAG:-}")" + validate_manylinux_cuda_image "$image" + local container="tvm_wheel_cuda_${GITHUB_RUN_ID:-local}_${GITHUB_RUN_ATTEMPT:-1}_${TVM_ARCH}" + local host_cuda_build_dir="$TVM_CUDA_BUILD_DIR" + local container_cuda_root="/workspace-cuda-build" + local container_cuda_build_dir="${container_cuda_root}/build" + mkdir -p "$host_cuda_build_dir" + local cuda_rpm="/tmp/cuda-repo-rhel8-13-0-local-13.0.2_580.95.05-1.${TVM_ARCH}.rpm" + trap "rm -f '${cuda_rpm}'; docker exec '${container}' bash -lc 'chown -R $(id -u):$(id -g) ${container_cuda_root} || true' >/dev/null 2>&1 || true; docker rm -f '${container}' >/dev/null 2>&1 || true" EXIT + docker pull "$image" + docker rm -f "$container" >/dev/null 2>&1 || true + docker run --name "$container" -d \ + --workdir /workspace \ + --volume "${REPO_ROOT}:/workspace" \ + --volume "${host_cuda_build_dir}:${container_cuda_root}" \ + "$image" tail -f /dev/null + + local cuda_rpm_name + cuda_rpm_name="$(basename "$cuda_rpm")" + curl -fsSLo "$cuda_rpm" "https://developer.download.nvidia.com/compute/cuda/13.0.2/local_installers/${cuda_rpm_name}" + docker cp "$cuda_rpm" "${container}:/${cuda_rpm_name}" + rm "$cuda_rpm" + docker exec "$container" bash -lc " + rpm -i /${cuda_rpm_name} && \ + dnf clean all && \ + dnf -y --disablerepo=epel install cuda-toolkit-13-0 && \ + rm /${cuda_rpm_name} && \ + dnf clean all" + + docker exec \ + -e TVM_PYTHON=/opt/python/cp310-cp310/bin/python \ + -e TVM_USE_CUDA=/usr/local/cuda \ + -e TVM_CUDA_ARCHITECTURES="$TVM_CUDA_ARCHITECTURES" \ + -e TVM_CUDA_BUILD_DIR="$container_cuda_build_dir" \ + -e TVM_INCLUDE_CUDA_RUNTIME=1 \ + -e CMAKE_BUILD_PARALLEL_LEVEL="$TVM_BUILD_PARALLEL_LEVEL" \ + -e TVM_BUILD_PARALLEL_LEVEL="$TVM_BUILD_PARALLEL_LEVEL" \ + "$container" bash -lc ' + set -eux + export PATH=/opt/python/cp310-cp310/bin:/usr/local/cuda/bin:$PATH + python -m pip install -U pip cmake ninja + python --version + cmake --version + nvcc --version + ci/scripts/package/tvm_wheel_helper.sh cuda' + + docker exec "$container" bash -lc \ + "chown -R $(id -u):$(id -g) ${container_cuda_root} || true" +} + +build_cuda_runtime() { + if [[ "$TVM_SKIP_CUDA" == "1" ]]; then + echo "Skipping CUDA build because TVM_SKIP_CUDA=1" + return 0 + fi + + require_cmd cmake + echo "Building libtvm_runtime_cuda.so in ${TVM_CUDA_BUILD_DIR}" + if [[ "$TVM_KEEP_BUILD_DIRS" != "1" ]]; then + rm -rf "$TVM_CUDA_BUILD_DIR" + fi + local cuda_compiler_args=() + if [[ -x "${TVM_USE_CUDA}/bin/nvcc" ]]; then + cuda_compiler_args+=(-DCMAKE_CUDA_COMPILER="${TVM_USE_CUDA}/bin/nvcc") + fi + cmake -S "$REPO_ROOT" -B "$TVM_CUDA_BUILD_DIR" \ + -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_TESTING=OFF \ + -DTVM_BUILD_PYTHON_MODULE=ON \ + -DUSE_CUDA="$TVM_USE_CUDA" \ + -DUSE_LLVM=OFF \ + -DUSE_CUBLAS=OFF \ + -DUSE_CUDNN=OFF \ + -DUSE_CUTLASS=OFF \ + -DUSE_NCCL=OFF \ + -DUSE_NVTX=OFF \ + -DCMAKE_CUDA_ARCHITECTURES="$TVM_CUDA_ARCHITECTURES" \ + "${cuda_compiler_args[@]}" + + cmake --build "$TVM_CUDA_BUILD_DIR" --target tvm_runtime tvm_runtime_cuda --parallel "$TVM_BUILD_PARALLEL_LEVEL" + + local cuda_lib + cuda_lib="$(cuda_runtime_path)" + if [[ -z "$cuda_lib" ]]; then + echo "error: libtvm_runtime_cuda.so was not produced" >&2 + return 1 + fi + if [[ "$(uname -s)" == "Linux" ]]; then + require_cmd patchelf + patchelf --set-rpath '$ORIGIN' "$cuda_lib" + fi + echo "CUDA runtime: ${cuda_lib}" +} + +inject_wheel_file() { + local raw_wheel="$1" + local output_dir="$2" + rm -rf "$output_dir" + mkdir -p "$output_dir" + + local inject_args=(--output-dir "$output_dir") + if [[ "$TVM_SKIP_CUDA" != "1" ]]; then + local cuda_lib + cuda_lib="$(cuda_runtime_path)" + if [[ -z "$cuda_lib" ]]; then + echo "error: CUDA runtime missing; run the 'cuda' step first" >&2 + return 1 + fi + inject_args+=(--cuda-runtime "$cuda_lib") + fi + if [[ -n "$TVM_WHEEL_DIST_NAME" ]]; then + inject_args+=(--distribution-name "$TVM_WHEEL_DIST_NAME") + fi + if [[ -n "$TVM_WHEEL_DIST_VERSION" ]]; then + inject_args+=(--distribution-version "$TVM_WHEEL_DIST_VERSION") + fi + if [[ "$(uname -s)" == "Linux" ]]; then + inject_args+=(--set-rpath '$ORIGIN') + fi + + echo "Rewriting wheel metadata/runtime files for ${raw_wheel}" + "$TVM_PYTHON" "$SCRIPT_DIR/rewrite_wheel.py" "$raw_wheel" "${inject_args[@]}" +} + +auditwheel_excludes() { + local cuda_lib="$1" + local seen + seen=" libtvm_ffi.so " + seen+="libcuda.so.1 libcuda.so libcudart.so.11.0 libcudart.so.12 libcudart.so.12.0 " + + printf '%s\n' "--exclude" "libtvm_ffi.so" + printf '%s\n' "--exclude" "libcuda.so.1" + printf '%s\n' "--exclude" "libcuda.so" + printf '%s\n' "--exclude" "libcudart.so.11.0" + printf '%s\n' "--exclude" "libcudart.so.12" + printf '%s\n' "--exclude" "libcudart.so.12.0" + + if [[ -n "$cuda_lib" ]] && command -v readelf >/dev/null 2>&1; then + while IFS= read -r needed; do + case "$needed" in + libcuda.so*|libcudart.so*|libnv*.so*) + if [[ "$seen" != *" ${needed} "* ]]; then + seen+="${needed} " + printf '%s\n' "--exclude" "$needed" + fi + ;; + esac + done < <(readelf -d "$cuda_lib" | sed -n 's/.*Shared library: \[\(.*\)\].*/\1/p') + fi +} + +llvm_libdir() { + if [[ "$TVM_USE_LLVM" == "OFF" || "$TVM_USE_LLVM" == "0" ]]; then + return 0 + fi + local -a llvm_config + read -r -a llvm_config <<<"$TVM_USE_LLVM" + if [[ "${#llvm_config[@]}" -eq 0 ]]; then + return 0 + fi + if command -v "${llvm_config[0]}" >/dev/null 2>&1; then + "${llvm_config[@]}" --libdir + elif [[ -x "${llvm_config[0]}" ]]; then + "${llvm_config[@]}" --libdir + fi +} + +llvm_prefix() { + if [[ "$TVM_USE_LLVM" == "OFF" || "$TVM_USE_LLVM" == "0" ]]; then + return 0 + fi + local -a llvm_config + read -r -a llvm_config <<<"$TVM_USE_LLVM" + if [[ "${#llvm_config[@]}" -eq 0 ]]; then + return 0 + fi + if command -v "${llvm_config[0]}" >/dev/null 2>&1 || [[ -x "${llvm_config[0]}" ]]; then + "${llvm_config[@]}" --prefix + fi +} + +prepare_repair_libdir() { + local source_dir="$1" + shift + local repair_libdir + repair_libdir="$(mktemp -d)" + shopt -s nullglob + local pattern lib + for pattern in "$@"; do + for lib in "$source_dir"/$pattern; do + ln -sf "$lib" "$repair_libdir/$(basename "$lib")" + done + done + shopt -u nullglob + if find "$repair_libdir" -type l -print -quit | grep -q .; then + echo "$repair_libdir" + else + rm -rf "$repair_libdir" + fi +} + +diagnose_wheel_elf() ( + local wheel="$1" + if ! command -v readelf >/dev/null 2>&1; then + return 0 + fi + local tmpdir + tmpdir="$(mktemp -d)" + trap 'rm -rf "$tmpdir"' EXIT + "$TVM_PYTHON" - "$wheel" "$tmpdir" <<'PY' +from pathlib import Path +import sys +import zipfile + +wheel = Path(sys.argv[1]) +target = Path(sys.argv[2]) +with zipfile.ZipFile(wheel) as zf: + for name in zf.namelist(): + if name.endswith(".so") or ".so." in name: + zf.extract(name, target) +PY + local lib rel + while IFS= read -r lib; do + rel="${lib#"$tmpdir"/}" + echo "::group::ELF diagnostics: ${rel}" + readelf -d "$lib" | sed -n 's/.*Shared library: \[\(.*\)\].*/NEEDED \1/p; s/.*Library .*path: \[\(.*\)\].*/RPATH \1/p' || true + readelf --version-info "$lib" \ + | sed -n 's/.*Name: \(GLIBC[^ ]*\|GLIBCXX[^ ]*\|CXXABI[^ ]*\).*/VERSION \1/p' \ + | sort -Vu || true + echo "::endgroup::" + done < <(find "$tmpdir" -type f \( -name '*.so' -o -name '*.so.*' \) | sort) +) + +repair_wheel_to_dir() { + local injected_wheel="$1" + local output_dir="$2" + mkdir -p "$output_dir" + local existing_wheel + while IFS= read -r existing_wheel; do + rm -f "$existing_wheel" + done < <(find "$output_dir" -maxdepth 1 -type f -name '*.whl') + if [[ "$TVM_SKIP_REPAIR" == "1" ]]; then + cp "$injected_wheel" "$output_dir/" + echo "Repair skipped; final wheel copied to ${output_dir}" + return 0 + fi + + case "$(uname -s)" in + Linux) + require_cmd auditwheel + local cuda_lib="" + if [[ "$TVM_SKIP_CUDA" != "1" ]]; then + cuda_lib="$(cuda_runtime_path)" + fi + local exclude_args=() + local exclude_arg + while IFS= read -r exclude_arg; do + exclude_args+=("$exclude_arg") + done < <(auditwheel_excludes "$cuda_lib") + echo "Repairing Linux wheel with auditwheel" + ( + auditwheel_plat_args=() + if [[ -n "${TVM_AUDITWHEEL_PLAT:-}" ]]; then + auditwheel_plat_args+=(--plat "$TVM_AUDITWHEEL_PLAT") + fi + if [[ -n "${TVM_AUDITWHEEL_LIBRARY_PATH:-}" ]]; then + echo "Adding explicit library path to LD_LIBRARY_PATH for auditwheel: ${TVM_AUDITWHEEL_LIBRARY_PATH}" + export LD_LIBRARY_PATH="${TVM_AUDITWHEEL_LIBRARY_PATH}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + fi + if ! auditwheel -v repair "${auditwheel_plat_args[@]}" "${exclude_args[@]}" \ + -w "$output_dir" "$injected_wheel"; then + echo "auditwheel repair failed; printing diagnostics for ${injected_wheel}" >&2 + auditwheel -v show "$injected_wheel" >&2 || true + diagnose_wheel_elf "$injected_wheel" >&2 || true + return 1 + fi + ) + ;; + Darwin) + require_cmd delocate-wheel + echo "Repairing macOS wheel with delocate" + ( + repair_libdir="" + trap '[[ -z "${repair_libdir:-}" ]] || rm -rf "$repair_libdir"' EXIT + llvm_dir="$(llvm_libdir || true)" + if [[ -n "${llvm_dir:-}" && -d "$llvm_dir" ]]; then + repair_libdir="$(prepare_repair_libdir "$llvm_dir" \ + 'libxml2*.dylib' 'libz*.dylib' 'libzstd*.dylib' 'liblzma*.dylib' 'libiconv*.dylib' || true)" + if [[ -n "${repair_libdir:-}" ]]; then + echo "Adding filtered LLVM libdir to DYLD_LIBRARY_PATH for delocate: ${repair_libdir}" + export DYLD_LIBRARY_PATH="${repair_libdir}${DYLD_LIBRARY_PATH:+:${DYLD_LIBRARY_PATH}}" + fi + fi + delocate_arch_args=() + if [[ -n "${TVM_DELOCATE_ARCHS:-}" ]]; then + delocate_arch_args+=(--require-archs "$TVM_DELOCATE_ARCHS") + fi + delocate-wheel \ + "${delocate_arch_args[@]}" \ + --ignore-missing-dependencies \ + --exclude libtvm_ffi.dylib \ + -w "$output_dir" \ + -v "$injected_wheel" + ) + ;; + *) + cp "$injected_wheel" "$output_dir/" + echo "No repair step for this platform; final wheel copied to ${output_dir}" + ;; + esac + + single_wheel "$output_dir" >/dev/null +} + +cibw_repair_wheel() ( + local raw_wheel="$1" + local dest_dir="$2" + local injected_dir + injected_dir="$(mktemp -d)" + trap 'rm -rf "$injected_dir"' EXIT + inject_wheel_file "$raw_wheel" "$injected_dir" + local injected_wheel + injected_wheel="$(single_wheel "$injected_dir")" + repair_wheel_to_dir "$injected_wheel" "$dest_dir" +) + +validate_wheel_elf() { + local final_wheel + final_wheel="$(single_wheel "$TVM_WHEELHOUSE")" + if [[ "$(uname -s)" == "Linux" ]]; then + "$TVM_PYTHON" "$SCRIPT_DIR/validate_wheel_elf.py" "$final_wheel" + fi +} + +verify_wheel() { + local final_wheel + final_wheel="$(single_wheel "$TVM_WHEELHOUSE")" + if [[ -n "${TVM_EXPECT_WHEEL_PLATFORM_TAG:-}" ]]; then + if [[ "$(basename "$final_wheel")" != *"${TVM_EXPECT_WHEEL_PLATFORM_TAG}"* ]]; then + echo "error: expected final wheel tag ${TVM_EXPECT_WHEEL_PLATFORM_TAG}, got ${final_wheel}" >&2 + return 1 + fi + fi + validate_wheel_elf + + local venv="${TVM_VERIFY_VENV:-${REPO_ROOT}/build-wheel-verify-venv}" + rm -rf "$venv" + "$TVM_PYTHON" -m venv "$venv" + + local venv_python="${venv}/bin/python" + if [[ "$(uname -s)" == MINGW* || "$(uname -s)" == CYGWIN* ]]; then + venv_python="${venv}/Scripts/python.exe" + fi + + "$venv_python" -m pip install --upgrade pip + "$venv_python" -m pip install --extra-index-url "${TVM_EXTRA_INDEX_URL:-https://pypi.org/simple}" "$final_wheel" + "$venv_python" "$SCRIPT_DIR/verify_tvm_install.py" +} + +upload_wheel() { + require_cmd twine + local repo_args=() + if [[ -n "${TVM_UPLOAD_REPOSITORY_URL:-}" ]]; then + repo_args+=(--repository-url "$TVM_UPLOAD_REPOSITORY_URL") + fi + twine upload "${repo_args[@]}" "$TVM_WHEELHOUSE"/* +} + +verify_pypi_wheel() { + local final_wheel + final_wheel="$(single_wheel "$TVM_WHEELHOUSE")" + + local package_name package_version + package_name="$(wheel_metadata_field "$final_wheel" Name)" + package_version="$(wheel_metadata_field "$final_wheel" Version)" + + local index_url="${TVM_TEST_INDEX_URL:-https://test.pypi.org/simple/}" + local extra_index_url="${TVM_EXTRA_INDEX_URL:-https://pypi.org/simple}" + local venv="${TVM_VERIFY_PYPI_VENV:-${REPO_ROOT}/build-wheel-verify-pypi-venv}" + rm -rf "$venv" + "$TVM_PYTHON" -m venv "$venv" + + local venv_python="${venv}/bin/python" + if [[ "$(uname -s)" == MINGW* || "$(uname -s)" == CYGWIN* ]]; then + venv_python="${venv}/Scripts/python.exe" + fi + + "$venv_python" -m pip install --upgrade pip + "$venv_python" -m pip install \ + --index-url "$index_url" \ + --extra-index-url "$extra_index_url" \ + "${package_name}==${package_version}" + "$venv_python" "$SCRIPT_DIR/verify_tvm_install.py" +} + +main() { + local step="${1:-help}" + case "$step" in + cuda) build_cuda_runtime ;; + cuda-path) cuda_runtime_path ;; + manylinux-cuda) run_manylinux_cuda_container ;; + cibw-repair) + if [[ "$#" -ne 3 ]]; then + echo "error: cibw-repair requires " >&2 + return 1 + fi + cibw_repair_wheel "$2" "$3" + ;; + validate) validate_wheel_elf ;; + verify) verify_wheel ;; + verify-installed) "$TVM_PYTHON" "$SCRIPT_DIR/verify_tvm_install.py" ;; + upload) upload_wheel ;; + verify-pypi) verify_pypi_wheel ;; + -h|--help|help) usage ;; + *) + usage >&2 + return 1 + ;; + esac +} + +main "$@" diff --git a/ci/scripts/package/validate_wheel_elf.py b/ci/scripts/package/validate_wheel_elf.py new file mode 100644 index 000000000000..7ff81234d73a --- /dev/null +++ b/ci/scripts/package/validate_wheel_elf.py @@ -0,0 +1,169 @@ +# 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. +"""Validate ELF linkage inside a repaired TVM wheel.""" + +from __future__ import annotations + +import argparse +import os +import re +import shutil +import subprocess +import sys +import tempfile +import zipfile +from pathlib import Path + +TVM_EXTERNAL_LIBS = { + "libtvm_ffi.so", +} + + +def _run(command: list[str], env: dict[str, str] | None = None) -> str: + try: + proc = subprocess.run( + command, + check=True, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + except subprocess.CalledProcessError as err: + output = err.stdout or "" + raise RuntimeError(f"{' '.join(command)} failed:\n{output}") from err + return proc.stdout + + +def _dynamic_entries(path: Path) -> tuple[list[str], list[str]]: + output = _run(["readelf", "-d", str(path)]) + needed = re.findall(r"Shared library: \[(.*?)\]", output) + rpaths = re.findall(r"Library (?:rpath|runpath): \[(.*?)\]", output) + return needed, [entry for rpath in rpaths for entry in rpath.split(":") if entry] + + +def _ldd(path: Path) -> dict[str, str]: + output = _run(["ldd", str(path)], env={**os.environ, "LD_LIBRARY_PATH": ""}) + resolved: dict[str, str] = {} + for line in output.splitlines(): + line = line.strip() + if "=>" not in line: + continue + name, target = line.split("=>", 1) + target = target.strip() + resolved[name.strip()] = target.split(" ", 1)[0] + return resolved + + +def validate(wheel: Path) -> None: + if sys.platform != "linux": + print("ELF wheel validation skipped on non-Linux platform") + return + for command in ("readelf", "ldd"): + if shutil.which(command) is None: + raise RuntimeError(f"{command} is required for ELF wheel validation") + + with tempfile.TemporaryDirectory() as tmpdir: + root = Path(tmpdir) + with zipfile.ZipFile(wheel) as zf: + zf.extractall(root) + + libdir = root / "tvm" / "lib" + if not libdir.is_dir(): + raise RuntimeError(f"wheel does not contain {libdir.relative_to(root)}") + + bundled_tvm_ffi = sorted( + str(path.relative_to(root)) for path in root.rglob("libtvm_ffi*.so*") if path.is_file() + ) + if bundled_tvm_ffi: + raise RuntimeError( + "TVM wheel must depend on tvm_ffi instead of bundling libtvm_ffi: " + + ", ".join(bundled_tvm_ffi) + ) + + libs = {path.name: path for path in sorted(libdir.glob("*.so*")) if path.is_file()} + if "libtvm_runtime.so" not in libs: + raise RuntimeError("wheel does not contain tvm/lib/libtvm_runtime.so") + bundled_llvm = sorted( + str(path.relative_to(root)) for path in root.rglob("libLLVM*.so*") if path.is_file() + ) + if bundled_llvm: + raise RuntimeError( + "TVM wheel must link LLVM statically instead of bundling libLLVM: " + + ", ".join(bundled_llvm) + ) + + errors: list[str] = [] + for lib in libs.values(): + needed, rpaths = _dynamic_entries(lib) + llvm_needed = sorted(name for name in needed if name.startswith("libLLVM")) + if llvm_needed: + errors.append( + f"{lib.relative_to(root)} links dynamic LLVM libraries: {llvm_needed}" + ) + + internal_needed = sorted(name for name in needed if name in libs) + if internal_needed and "$ORIGIN" not in rpaths: + errors.append( + f"{lib.relative_to(root)} needs {internal_needed} but RUNPATH/RPATH is {rpaths}" + ) + + resolved = _ldd(lib) + for name in internal_needed: + target = resolved.get(name) + if target is None: + errors.append(f"{lib.relative_to(root)}: ldd did not report {name}") + continue + if target == "not": + errors.append(f"{lib.relative_to(root)}: {name} is not found") + continue + if Path(target).resolve() != libs[name].resolve(): + errors.append( + f"{lib.relative_to(root)}: {name} resolved to {target}, " + f"expected {libs[name].relative_to(root)}" + ) + + unexpected_tvm_deps = sorted( + name + for name in needed + if name.startswith("libtvm_") and name not in libs and name not in TVM_EXTERNAL_LIBS + ) + if unexpected_tvm_deps: + errors.append( + f"{lib.relative_to(root)} has unresolved TVM deps: {unexpected_tvm_deps}" + ) + + if errors: + raise RuntimeError("ELF wheel validation failed:\n" + "\n".join(errors)) + + print(f"ELF wheel validation passed for {wheel.name}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("wheel", type=Path) + args = parser.parse_args() + try: + validate(args.wheel) + except RuntimeError as err: + print(err, file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ci/scripts/package/verify_tvm_install.py b/ci/scripts/package/verify_tvm_install.py new file mode 100644 index 000000000000..bdb27678be73 --- /dev/null +++ b/ci/scripts/package/verify_tvm_install.py @@ -0,0 +1,275 @@ +# 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. +"""Verify an installed TVM wheel imports and ships the expected runtime DSO.""" + +from __future__ import annotations + +import faulthandler +import os +from pathlib import Path +import signal +import sys + +faulthandler.enable(all_threads=True) +if hasattr(sys.stdout, "reconfigure"): + sys.stdout.reconfigure(line_buffering=True) + +import numpy as np + + +def log(*args: object) -> None: + print(*args, flush=True) + + +def _enable_python_fault_handler() -> None: + """Install Python's signal handler after native libraries may install theirs.""" + faulthandler.enable(all_threads=True) + try: + faulthandler.register(signal.SIGUSR1, all_threads=True) + except (AttributeError, RuntimeError, ValueError): + pass + + +def expect_bool(name: str) -> bool | None: + value = os.environ.get(name) + if value is None or value == "": + return None + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off"}: + return False + raise RuntimeError(f"{name} must be a boolean value, got {value!r}") + + +def _clear_external_library_overrides() -> None: + for name in ("TVM_LIBRARY_PATH", "LD_LIBRARY_PATH", "DYLD_LIBRARY_PATH"): + if name in os.environ: + log(f"clearing {name} before importing tvm") + os.environ.pop(name, None) + + +def _first_existing(candidates: list[Path]) -> Path: + for candidate in candidates: + if candidate.exists(): + return candidate + return candidates[0] + + +def _assert_loaded_runtime_from_wheel(libdir: Path, runtime_candidates: list[Path]) -> None: + import tvm.base as tvm_base # pylint: disable=import-outside-toplevel + + loaded_runtime = Path(tvm_base._LIB_RUNTIME._name).resolve() # pylint: disable=protected-access + expected_runtime_paths = {candidate.resolve() for candidate in runtime_candidates} + log("loaded runtime library:", loaded_runtime) + if loaded_runtime not in expected_runtime_paths: + expected = ", ".join(str(path) for path in sorted(expected_runtime_paths)) + raise RuntimeError( + f"loaded runtime library is not from the installed wheel: " + f"got {loaded_runtime}, expected one of {expected}" + ) + + +def _dynamic_llvm_libs(libdir: Path) -> list[Path]: + if sys.platform == "darwin": + patterns = ["libLLVM*.dylib"] + elif sys.platform == "win32": + patterns = ["LLVM*.dll", "libLLVM*.dll"] + else: + patterns = ["libLLVM*.so", "libLLVM*.so.*"] + found: set[Path] = set() + for pattern in patterns: + found.update(libdir.glob(pattern)) + return sorted(found) + + +def _log_tvm_ffi_details() -> None: + import tvm_ffi # pylint: disable=import-outside-toplevel + + ffi_lib = getattr(tvm_ffi, "LIB", None) + ffi_lib_path = getattr(ffi_lib, "_name", None) + log("tvm_ffi version:", getattr(tvm_ffi, "__version__", "")) + log("tvm_ffi package:", Path(tvm_ffi.__file__).resolve().parent) + if ffi_lib_path: + log("tvm_ffi library:", Path(ffi_lib_path).resolve()) + + +def _log_loaded_native_libraries() -> None: + if sys.platform != "linux": + return + maps_path = Path("/proc/self/maps") + if not maps_path.exists(): + return + interesting_names = ( + "libtvm", + "libLLVM", + "libstdc++", + "libgcc_s", + "libxml2", + "libzstd", + "liblzma", + ) + loaded: set[str] = set() + for line in maps_path.read_text(encoding="utf-8", errors="ignore").splitlines(): + path = line.rsplit(maxsplit=1)[-1] + if "/" not in path: + continue + if any(name in path for name in interesting_names): + loaded.add(path) + log("loaded native libraries:") + for path in sorted(loaded): + log(" ", path) + + +def _verify_llvm_tirx_compile() -> None: + import tvm # pylint: disable=import-outside-toplevel + from tvm import te # pylint: disable=import-outside-toplevel + + log("llvm tirx compile smoke: starting") + extent = 8 + log("llvm tirx compile smoke: create numpy inputs") + lhs_np = np.arange(extent, dtype="float32") + rhs_np = np.arange(extent, dtype="float32") * np.float32(2) + out_np = np.zeros(extent, dtype="float32") + + log("llvm tirx compile smoke: create placeholders") + lhs = te.placeholder((extent,), name="lhs", dtype="float32") + rhs = te.placeholder((extent,), name="rhs", dtype="float32") + log("llvm tirx compile smoke: create compute") + out = te.compute((extent,), lambda i: lhs[i] + rhs[i], name="out") + log("llvm tirx compile smoke: create prim func") + prim_func = te.create_prim_func([lhs, rhs, out]) + log("llvm tirx compile smoke: compile") + executable = tvm.compile(prim_func, target="llvm") + + log("llvm tirx compile smoke: create tensors") + dev = tvm.cpu() + lhs_t = tvm.runtime.tensor(lhs_np, dev) + rhs_t = tvm.runtime.tensor(rhs_np, dev) + out_t = tvm.runtime.tensor(out_np, dev) + log("llvm tirx compile smoke: execute") + executable(lhs_t, rhs_t, out_t) + log("llvm tirx compile smoke: check output") + np.testing.assert_allclose(out_t.numpy(), lhs_np + rhs_np, rtol=1e-6) + log("llvm tirx compile smoke: passed") + + +def _verify_relax_compile() -> None: + import tvm # pylint: disable=import-outside-toplevel + from tvm import relax # pylint: disable=import-outside-toplevel + + log("llvm relax compile smoke: starting") + log("llvm relax compile smoke: create numpy inputs") + lhs_np = np.arange(8, dtype="float32") + rhs_np = np.arange(8, dtype="float32") * np.float32(3) + dev = tvm.cpu() + + log("llvm relax compile smoke: create vars") + lhs = relax.Var("lhs", relax.TensorStructInfo((8,), "float32")) + rhs = relax.Var("rhs", relax.TensorStructInfo((8,), "float32")) + log("llvm relax compile smoke: create module") + builder = relax.BlockBuilder() + with builder.function("main", [lhs, rhs]): + out = builder.emit(relax.op.add(lhs, rhs)) + builder.emit_func_output(out) + + log("llvm relax compile smoke: compile") + executable = tvm.compile(builder.get(), target="llvm") + log("llvm relax compile smoke: create vm") + vm = relax.VirtualMachine(executable, dev) + log("llvm relax compile smoke: execute") + out = vm["main"](tvm.runtime.tensor(lhs_np, dev), tvm.runtime.tensor(rhs_np, dev)) + log("llvm relax compile smoke: check output") + np.testing.assert_allclose(out.numpy(), lhs_np + rhs_np, rtol=1e-6) + log("llvm relax compile smoke: passed") + + +def main() -> int: + _clear_external_library_overrides() + + log("import tvm: starting") + import tvm # pylint: disable=import-outside-toplevel + log("import tvm: passed") + _enable_python_fault_handler() + _log_tvm_ffi_details() + + root = Path(tvm.__file__).resolve().parent + libdir = root / "lib" + if sys.platform == "darwin": + runtime_candidates = [libdir / "libtvm_runtime.dylib"] + cuda_runtime_candidates = [libdir / "libtvm_runtime_cuda.dylib"] + elif sys.platform == "win32": + runtime_candidates = [libdir / "tvm_runtime.dll", libdir / "libtvm_runtime.dll"] + cuda_runtime_candidates = [ + libdir / "tvm_runtime_cuda.dll", + libdir / "libtvm_runtime_cuda.dll", + ] + else: + runtime_candidates = [libdir / "libtvm_runtime.so"] + cuda_runtime_candidates = [libdir / "libtvm_runtime_cuda.so"] + + log("tvm version:", tvm.__version__) + log("tvm package:", root) + llvm_enabled = bool(tvm.runtime.enabled("llvm")) + cuda_enabled = bool(tvm.runtime.enabled("cuda")) + runtime_lib = _first_existing(runtime_candidates) + cuda_runtime = _first_existing(cuda_runtime_candidates) + runtime_present = any(candidate.exists() for candidate in runtime_candidates) + cuda_runtime_present = any(candidate.exists() for candidate in cuda_runtime_candidates) + dynamic_llvm_libs = _dynamic_llvm_libs(libdir) + + log("llvm enabled:", llvm_enabled) + log("cuda runtime enabled:", cuda_enabled) + log("runtime library:", runtime_lib) + if not runtime_present: + raise RuntimeError( + "runtime library is missing; checked " + + ", ".join(str(candidate) for candidate in runtime_candidates) + ) + _assert_loaded_runtime_from_wheel(libdir, runtime_candidates) + log("cuda runtime present:", cuda_runtime_present) + if cuda_runtime_present: + log("cuda runtime library:", cuda_runtime) + log("dynamic LLVM libraries:", [str(path) for path in dynamic_llvm_libs]) + _log_loaded_native_libraries() + + expected_llvm = expect_bool("TVM_EXPECT_LLVM_ENABLED") + if expected_llvm is not None and llvm_enabled != expected_llvm: + raise RuntimeError(f"llvm enabled: expected {expected_llvm}, got {llvm_enabled}") + if llvm_enabled: + _verify_llvm_tirx_compile() + _verify_relax_compile() + expected_static_llvm = expect_bool("TVM_EXPECT_STATIC_LLVM") + if expected_static_llvm and dynamic_llvm_libs: + raise RuntimeError( + "expected LLVM to be linked statically, but dynamic LLVM libraries are present: " + + ", ".join(str(path) for path in dynamic_llvm_libs) + ) + expected_cuda_runtime = expect_bool("TVM_EXPECT_CUDA_RUNTIME") + if expected_cuda_runtime is not None and cuda_runtime_present != expected_cuda_runtime: + raise RuntimeError( + f"cuda runtime present: expected {expected_cuda_runtime}, got {cuda_runtime_present}" + ) + expected_cuda = expect_bool("TVM_EXPECT_CUDA_ENABLED") + if expected_cuda is not None and cuda_enabled != expected_cuda: + raise RuntimeError(f"cuda runtime enabled: expected {expected_cuda}, got {cuda_enabled}") + log("verify tvm install: passed") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ci/scripts/package/verify_tvm_linux.sh b/ci/scripts/package/verify_tvm_linux.sh new file mode 100644 index 000000000000..54b5515170a2 --- /dev/null +++ b/ci/scripts/package/verify_tvm_linux.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# 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. + +# Linux wheel verification entrypoint. Clears any development library-path +# overrides so the verifier loads the libraries bundled in the installed wheel, +# then runs the cross-platform verifier. + +set -euo pipefail + +VERIFY_SCRIPT="${1:?usage: verify_tvm_linux.sh /path/to/verify_tvm_install.py}" + +for name in TVM_LIBRARY_PATH LD_LIBRARY_PATH DYLD_LIBRARY_PATH; do + if [[ -n "${!name:-}" ]]; then + echo "clearing ${name} before importing tvm" + unset "${name}" + fi +done + +exec python -u -X faulthandler "${VERIFY_SCRIPT}" diff --git a/cmake/modules/CUDA.cmake b/cmake/modules/CUDA.cmake index ec6160e7afaf..42bd72d3b773 100644 --- a/cmake/modules/CUDA.cmake +++ b/cmake/modules/CUDA.cmake @@ -79,6 +79,7 @@ if(USE_CUDA) RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ) + tvm_set_python_module_relative_rpath(tvm_runtime_cuda) install(TARGETS tvm_runtime_cuda DESTINATION lib${LIB_SUFFIX}) if(TVM_BUILD_PYTHON_MODULE) install(TARGETS tvm_runtime_cuda DESTINATION "lib") @@ -102,7 +103,7 @@ if(USE_CUDA AND USE_CUDNN) add_library(tvm_cudnn_objs OBJECT ${CONTRIB_CUDNN_SRCS}) target_link_libraries(tvm_cudnn_objs PRIVATE tvm_runtime_extra_defs) target_link_libraries(tvm_runtime_extra PRIVATE tvm_cudnn_objs ${CUDA_CUDNN_LIBRARY}) -endif(USE_CUDNN) +endif(USE_CUDA AND USE_CUDNN) if(USE_CUDA AND USE_CUDNN_FRONTEND) message(STATUS "Build with cuDNN Frontend support") diff --git a/cmake/modules/Hexagon.cmake b/cmake/modules/Hexagon.cmake index 431b15b13ac6..254cda10971f 100644 --- a/cmake/modules/Hexagon.cmake +++ b/cmake/modules/Hexagon.cmake @@ -351,6 +351,7 @@ elseif(USE_HEXAGON) RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ) + tvm_set_python_module_relative_rpath(tvm_runtime_hexagon) install(TARGETS tvm_runtime_hexagon DESTINATION lib${LIB_SUFFIX}) if(TVM_BUILD_PYTHON_MODULE) install(TARGETS tvm_runtime_hexagon DESTINATION "lib") diff --git a/cmake/modules/LLVM.cmake b/cmake/modules/LLVM.cmake index f944b4130415..fcb31481afc8 100644 --- a/cmake/modules/LLVM.cmake +++ b/cmake/modules/LLVM.cmake @@ -50,7 +50,7 @@ if(NOT ${USE_LLVM} MATCHES ${IS_FALSE_PATTERN}) ) list(APPEND TVM_LINKER_LIBS ${LLVM_LIBS}) list(APPEND COMPILER_SRCS ${COMPILER_LLVM_SRCS}) - if(NOT MSVC) + if(NOT MSVC AND NOT TVM_LLVM_HAS_RTTI) set_source_files_properties(${COMPILER_LLVM_SRCS} PROPERTIES COMPILE_FLAGS "-fno-rtti") endif() diff --git a/cmake/modules/Metal.cmake b/cmake/modules/Metal.cmake index 72e7585534bb..7271ccfcae2e 100644 --- a/cmake/modules/Metal.cmake +++ b/cmake/modules/Metal.cmake @@ -35,6 +35,7 @@ if(USE_METAL) RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ) + tvm_set_python_module_relative_rpath(tvm_runtime_metal) install(TARGETS tvm_runtime_metal DESTINATION lib${LIB_SUFFIX}) if(TVM_BUILD_PYTHON_MODULE) install(TARGETS tvm_runtime_metal DESTINATION "lib") diff --git a/cmake/modules/OpenCL.cmake b/cmake/modules/OpenCL.cmake index 9a1c20a5a5ab..76eb6381e632 100644 --- a/cmake/modules/OpenCL.cmake +++ b/cmake/modules/OpenCL.cmake @@ -51,6 +51,7 @@ if(USE_OPENCL) RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ) + tvm_set_python_module_relative_rpath(tvm_runtime_opencl) install(TARGETS tvm_runtime_opencl DESTINATION lib${LIB_SUFFIX}) if(TVM_BUILD_PYTHON_MODULE) install(TARGETS tvm_runtime_opencl DESTINATION "lib") diff --git a/cmake/modules/ROCM.cmake b/cmake/modules/ROCM.cmake index b974aa412959..1e04026c00d4 100644 --- a/cmake/modules/ROCM.cmake +++ b/cmake/modules/ROCM.cmake @@ -53,6 +53,7 @@ if(USE_ROCM) RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ) + tvm_set_python_module_relative_rpath(tvm_runtime_rocm) install(TARGETS tvm_runtime_rocm DESTINATION lib${LIB_SUFFIX}) if(TVM_BUILD_PYTHON_MODULE) install(TARGETS tvm_runtime_rocm DESTINATION "lib") diff --git a/cmake/modules/Vulkan.cmake b/cmake/modules/Vulkan.cmake index 6821b4419b1a..d08b3353b57c 100644 --- a/cmake/modules/Vulkan.cmake +++ b/cmake/modules/Vulkan.cmake @@ -60,6 +60,7 @@ if(USE_VULKAN) RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ) + tvm_set_python_module_relative_rpath(tvm_runtime_vulkan) install(TARGETS tvm_runtime_vulkan DESTINATION lib${LIB_SUFFIX}) if(TVM_BUILD_PYTHON_MODULE) install(TARGETS tvm_runtime_vulkan DESTINATION "lib") diff --git a/cmake/utils/FindLLVM.cmake b/cmake/utils/FindLLVM.cmake index 8aa9c8b1b959..a947f05da0bd 100644 --- a/cmake/utils/FindLLVM.cmake +++ b/cmake/utils/FindLLVM.cmake @@ -59,6 +59,10 @@ macro(find_llvm use_llvm) message(STATUS "Fall back to using llvm-config") set(LLVM_CONFIG "${LLVM_TOOLS_BINARY_DIR}/llvm-config") endif() + set(TVM_LLVM_HAS_RTTI 0) + if(LLVM_ENABLE_RTTI) + set(TVM_LLVM_HAS_RTTI 1) + endif() endif() if(LLVM_LIBS) # check if defined, not if it is true @@ -122,6 +126,12 @@ macro(find_llvm use_llvm) message(FATAL_ERROR "Fatal error executing: ${LLVM_CONFIG} --libdir") endif() message(STATUS "LLVM libdir: ${__llvm_libdir}") + set(__llvm_lib_hints + "${__llvm_libdir}" + "${__llvm_prefix}/lib" + "${__llvm_prefix}/lib64" + "${__llvm_prefix}/Library/lib" + ) execute_process(COMMAND ${LLVM_CONFIG} --cmakedir RESULT_VARIABLE __llvm_exit_code OUTPUT_VARIABLE __llvm_cmakedir @@ -136,6 +146,13 @@ macro(find_llvm use_llvm) if(NOT "${__llvm_exit_code}" STREQUAL "0") message(FATAL_ERROR "Fatal error executing: ${LLVM_CONFIG} --targets-built") endif() + execute_process(COMMAND ${LLVM_CONFIG} --has-rtti + RESULT_VARIABLE __llvm_exit_code + OUTPUT_VARIABLE __llvm_has_rtti + OUTPUT_STRIP_TRAILING_WHITESPACE) + if(NOT "${__llvm_exit_code}" STREQUAL "0") + message(FATAL_ERROR "Fatal error executing: ${LLVM_CONFIG} --has-rtti") + endif() cmake_path(SET "__llvm_cmakedir" "${__llvm_cmakedir}") message(STATUS "LLVM cmakedir: ${__llvm_cmakedir}") # map prefix => $ @@ -175,6 +192,13 @@ macro(find_llvm use_llvm) if("AArch64" IN_LIST BUILT_TARGET_LIST) set(TVM_LLVM_HAS_AARCH64_TARGET 1) endif() + string(TOUPPER "${__llvm_has_rtti}" __llvm_has_rtti_upper) + set(TVM_LLVM_HAS_RTTI 0) + if("${__llvm_has_rtti_upper}" STREQUAL "YES" + OR "${__llvm_has_rtti_upper}" STREQUAL "ON" + OR "${__llvm_has_rtti_upper}" STREQUAL "1") + set(TVM_LLVM_HAS_RTTI 1) + endif() if (${USE_MLIR}) if (EXISTS "${__llvm_libdir}/libMLIRPresburger.a") if (EXISTS "${__llvm_libdir}/libMLIRSupport.a") @@ -193,22 +217,61 @@ macro(find_llvm use_llvm) message(STATUS "LLVM links against math") list(APPEND LLVM_LIBS "m") elseif(("${__flag}" STREQUAL "-lz") OR ("${__flag}" STREQUAL "z.lib")) - message(STATUS "LLVM links against zlib") - find_package(ZLIB REQUIRED) - list(APPEND LLVM_LIBS "ZLIB::ZLIB") + find_library(ZLIB_STATIC + NAMES libz.a zlibstatic z + HINTS ${__llvm_lib_hints} + NO_DEFAULT_PATH) + if (ZLIB_STATIC) + message(STATUS "LLVM links against static zlib: ${ZLIB_STATIC}") + list(APPEND LLVM_LIBS "${ZLIB_STATIC}") + else() + message(STATUS "LLVM links against zlib") + find_package(ZLIB REQUIRED) + list(APPEND LLVM_LIBS "ZLIB::ZLIB") + endif() elseif("${__flag}" STREQUAL "-lzstd") - list(APPEND CMAKE_MODULE_PATH "${__llvm_cmakedir}") - find_package(zstd REQUIRED) - if (TARGET "zstd::libzstd_static") - message(STATUS "LLVM links against static zstd") - list(APPEND LLVM_LIBS "zstd::libzstd_static") + find_library(ZSTD_STATIC + NAMES libzstd.a zstd_static zstd + HINTS ${__llvm_lib_hints} + NO_DEFAULT_PATH) + if (ZSTD_STATIC) + message(STATUS "LLVM links against static zstd: ${ZSTD_STATIC}") + list(APPEND LLVM_LIBS "${ZSTD_STATIC}") else() - message(STATUS "LLVM links against shared zstd") - list(APPEND LLVM_LIBS "zstd::libzstd_shared") + list(APPEND CMAKE_MODULE_PATH "${__llvm_cmakedir}") + find_package(zstd REQUIRED) + if (TARGET "zstd::libzstd_static") + message(STATUS "LLVM links against static zstd") + list(APPEND LLVM_LIBS "zstd::libzstd_static") + else() + message(STATUS "LLVM links against shared zstd") + list(APPEND LLVM_LIBS "zstd::libzstd_shared") + endif() endif() elseif("${__flag}" STREQUAL "-lxml2") - message(STATUS "LLVM links against xml2") - list(APPEND LLVM_LIBS "-lxml2") + if (UNIX AND NOT APPLE) + find_library(LIBXML2_SYSTEM_LIBRARY + NAMES libxml2.so.2 xml2 libxml2 + PATHS /usr/lib64 /usr/lib /lib64 /lib + NO_DEFAULT_PATH) + endif() + if (LIBXML2_SYSTEM_LIBRARY) + message(STATUS "LLVM links against system xml2: ${LIBXML2_SYSTEM_LIBRARY}") + list(APPEND LLVM_LIBS "${LIBXML2_SYSTEM_LIBRARY}") + else() + find_library(LIBXML2_LIBRARY + NAMES libxml2.a xml2 libxml2 + HINTS ${__llvm_lib_hints} + NO_DEFAULT_PATH) + if (LIBXML2_LIBRARY) + message(STATUS "LLVM links against xml2: ${LIBXML2_LIBRARY}") + list(APPEND LLVM_LIBS "${LIBXML2_LIBRARY}") + else() + message(STATUS "LLVM links against xml2") + find_package(LibXml2 REQUIRED) + list(APPEND LLVM_LIBS "LibXml2::LibXml2") + endif() + endif() elseif("${__flag}" STREQUAL "zstd.dll.lib") message(STATUS "LLVM linker flag under LLVM libdir: ${__llvm_libdir}/zstd.lib") list(APPEND LLVM_LIBS "${__llvm_libdir}/zstd.lib") @@ -234,6 +297,7 @@ macro(find_llvm use_llvm) message(FATAL_ERROR "TVM requires LLVM 15.0 or higher.") endif() message(STATUS "Found TVM_LLVM_HAS_AARCH64_TARGET=" ${TVM_LLVM_HAS_AARCH64_TARGET}) + message(STATUS "Found TVM_LLVM_HAS_RTTI=" ${TVM_LLVM_HAS_RTTI}) # Detect whether DIBuilder insertion APIs (insertDeclare, # insertDbgValueIntrinsic) accept BasicBlock::iterator as the insertion point diff --git a/pyproject.toml b/pyproject.toml index 2e400f0609e6..e1edbf5c6db0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,7 @@ classifiers = [ ] # Core dependencies - these are the minimum required for basic TVM functionality dependencies = [ - "apache-tvm-ffi", + "apache-tvm-ffi>=0.1.11", "cloudpickle", "ml_dtypes", "numpy", diff --git a/python/tvm/libinfo.py b/python/tvm/libinfo.py index 136d85a49f30..a8dd1ea34144 100644 --- a/python/tvm/libinfo.py +++ b/python/tvm/libinfo.py @@ -265,7 +265,11 @@ def find_include_path(name=None, search_path=None, optional=False): source_dir = os.environ["TVM_HOME"] else: ffi_dir = os.path.dirname(os.path.abspath(os.path.expanduser(__file__))) - for source_dir in ["..", "../..", "../../.."]: + # "." covers the installed-wheel layout, where the bundled headers live + # in ``/include``; the remaining entries cover a source tree, + # where the package is at ``/python/tvm`` and headers at + # ``/include``. + for source_dir in [".", "..", "../..", "../../.."]: source_dir = os.path.join(ffi_dir, source_dir) if os.path.isdir(os.path.join(source_dir, "include")): break @@ -310,6 +314,28 @@ def find_include_path(name=None, search_path=None, optional=False): include_found += [p for p in tvm_ffi_include_path if os.path.exists(p) and os.path.isdir(p)] include_found += [p for p in dlpack_include_path if os.path.exists(p) and os.path.isdir(p)] + # In a wheel install the tvm-ffi (and bundled dlpack) headers ship with + # the separate apache-tvm-ffi package rather than under 3rdparty/, so + # ask it where its headers live. This lets C/C++ source compilation + # (e.g. target="c" / Module.export_library) work from installed wheels. + try: + import tvm_ffi.libinfo as _ffi_libinfo # pylint: disable=import-outside-toplevel + + for _finder in ( + getattr(_ffi_libinfo, "find_include_path", None), + getattr(_ffi_libinfo, "find_dlpack_include_path", None), + ): + if _finder is None: + continue + try: + _p = _finder() + except Exception: # pylint: disable=broad-except + continue + if _p and os.path.isdir(_p) and _p not in include_found: + include_found.append(_p) + except Exception: # pylint: disable=broad-except + pass + if not include_found: message = ( "Cannot find the files.\n" diff --git a/python/tvm/rpc/server.py b/python/tvm/rpc/server.py index 099cb8f1f1e7..5ffcb878f185 100644 --- a/python/tvm/rpc/server.py +++ b/python/tvm/rpc/server.py @@ -48,9 +48,16 @@ from tvm.support.popen_pool import PopenWorker # pylint: disable=unused-import -from . import _ffi_api, base, testing +from . import _ffi_api, base from .base import TrackerCode +try: + # Register RPC unit-test helpers when the testing dependencies are present. + from . import testing +except ModuleNotFoundError as err: + if err.name != "pytest": + raise + logger = logging.getLogger("RPCServer") console_handler = logging.StreamHandler() console_handler.setFormatter( diff --git a/tests/conda/build-environment.yaml b/tests/conda/build-environment.yaml index ebd45ff4c422..3b2c4dd16751 100644 --- a/tests/conda/build-environment.yaml +++ b/tests/conda/build-environment.yaml @@ -33,6 +33,8 @@ dependencies: - pip - git - bzip2 + - zlib + - zstd-static - pytest - numpy - scipy diff --git a/tests/python/tirx/transform/test_transform_lower_tirx.py b/tests/python/tirx/transform/test_transform_lower_tirx.py index 3e20d61f8059..80e68243d0b3 100644 --- a/tests/python/tirx/transform/test_transform_lower_tirx.py +++ b/tests/python/tirx/transform/test_transform_lower_tirx.py @@ -24,7 +24,7 @@ from tvm.tirx.layout import laneid, warpid, wg_local_layout from tvm.tirx.stmt import ExecScopeStmt from tvm.tirx.stmt_functor import post_order_visit -from tvm.tirx.transform import LowerTIRx, Simplify +from tvm.tirx.transform import LowerTIRx, StmtSimplify def _contains_exec_scope(mod): @@ -1000,7 +1000,7 @@ def before(A_ptr: Tx.handle): with tvm.target.Target("cuda"): lowered = LowerTIRx()(tvm.IRModule({"main": before})) - simplified = Simplify()(lowered) + simplified = StmtSimplify()(lowered) script = simplified.script(extra_config={"tirx.prefix": "Tx"}) assert "if warp_id_in_cta // 4 == 0:" in script