|
| 1 | +#!/bin/bash |
| 2 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 3 | +# All rights reserved. |
| 4 | +# |
| 5 | +# This source code is licensed under the BSD-style license found in the |
| 6 | +# LICENSE file in the root directory of this source tree. |
| 7 | + |
| 8 | +# Install mingw-w64 cross-compiler and Windows CUDA toolkit for cross-compilation |
| 9 | + |
| 10 | +set -ex |
| 11 | + |
| 12 | +INSTALL_DIR="${WINDOWS_CUDA_INSTALL_DIR:-/opt/cuda-windows}" |
| 13 | + |
| 14 | +# Mapping of CUDA versions to their corresponding driver versions for Windows installers |
| 15 | +# Source: https://developer.nvidia.com/cuda-toolkit-archive |
| 16 | +declare -A CUDA_DRIVER_MAP=( |
| 17 | + ["12.6"]="12.6.3:561.17" |
| 18 | + ["12.8"]="12.8.1:572.61" |
| 19 | + ["12.9"]="12.9.1:576.57" |
| 20 | +) |
| 21 | + |
| 22 | +install_mingw() { |
| 23 | + echo "Installing mingw-w64 cross-compiler..." |
| 24 | + |
| 25 | + apt-get update |
| 26 | + # Install the POSIX threads version of mingw-w64 which supports C++11 threading |
| 27 | + # primitives (std::mutex, std::condition_variable, std::shared_mutex). |
| 28 | + # The default win32 threads version does not support these. |
| 29 | + apt-get install -y --no-install-recommends \ |
| 30 | + g++-mingw-w64-x86-64-posix \ |
| 31 | + mingw-w64-tools \ |
| 32 | + p7zip-full \ |
| 33 | + wget |
| 34 | + |
| 35 | + # Verify installation shows POSIX threads |
| 36 | + x86_64-w64-mingw32-g++ --version |
| 37 | + |
| 38 | + # Cleanup |
| 39 | + apt-get clean |
| 40 | + rm -rf /var/lib/apt/lists/* |
| 41 | + |
| 42 | + echo "mingw-w64 installation complete (POSIX threads version)" |
| 43 | +} |
| 44 | + |
| 45 | +get_torch_cuda_version() { |
| 46 | + # Query PyTorch for its CUDA version using conda environment |
| 47 | + conda run -n "py_${PYTHON_VERSION}" python3 -c "import torch; print(torch.version.cuda)" 2>/dev/null || echo "" |
| 48 | +} |
| 49 | + |
| 50 | +install_windows_cuda() { |
| 51 | + # Use CUDA_VERSION env var if set (from Docker build arg), otherwise query PyTorch |
| 52 | + if [ -n "${CUDA_VERSION:-}" ]; then |
| 53 | + echo "Using CUDA version from environment: ${CUDA_VERSION}" |
| 54 | + CUDA_MAJOR_MINOR=$(echo "${CUDA_VERSION}" | cut -d. -f1,2) |
| 55 | + else |
| 56 | + TORCH_CUDA_VERSION=$(get_torch_cuda_version) |
| 57 | + |
| 58 | + if [ -z "${TORCH_CUDA_VERSION}" ] || [ "${TORCH_CUDA_VERSION}" = "None" ]; then |
| 59 | + echo "ERROR: Could not detect CUDA version from PyTorch." |
| 60 | + echo "Make sure PyTorch with CUDA support is installed or set CUDA_VERSION." |
| 61 | + exit 1 |
| 62 | + fi |
| 63 | + |
| 64 | + echo "Detected PyTorch CUDA version: ${TORCH_CUDA_VERSION}" |
| 65 | + CUDA_MAJOR_MINOR=$(echo "${TORCH_CUDA_VERSION}" | cut -d. -f1,2) |
| 66 | + fi |
| 67 | + |
| 68 | + # Look up the full version and driver version |
| 69 | + if [ -z "${CUDA_DRIVER_MAP[${CUDA_MAJOR_MINOR}]}" ]; then |
| 70 | + echo "ERROR: CUDA version ${CUDA_MAJOR_MINOR} is not in the known version map." |
| 71 | + echo "Known versions: ${!CUDA_DRIVER_MAP[*]}" |
| 72 | + exit 1 |
| 73 | + fi |
| 74 | + |
| 75 | + CUDA_INFO="${CUDA_DRIVER_MAP[${CUDA_MAJOR_MINOR}]}" |
| 76 | + CUDA_VERSION=$(echo "${CUDA_INFO}" | cut -d: -f1) |
| 77 | + CUDA_DRIVER_VERSION=$(echo "${CUDA_INFO}" | cut -d: -f2) |
| 78 | + |
| 79 | + echo "Using CUDA ${CUDA_VERSION} with driver ${CUDA_DRIVER_VERSION}" |
| 80 | + |
| 81 | + echo "Installing Windows CUDA toolkit ${CUDA_VERSION}..." |
| 82 | + |
| 83 | + mkdir -p "${INSTALL_DIR}" |
| 84 | + cd "${INSTALL_DIR}" |
| 85 | + |
| 86 | + CUDA_INSTALLER="cuda_${CUDA_VERSION}_${CUDA_DRIVER_VERSION}_windows.exe" |
| 87 | + CUDA_URL="https://developer.download.nvidia.com/compute/cuda/${CUDA_VERSION}/local_installers/${CUDA_INSTALLER}" |
| 88 | + |
| 89 | + # Check if already downloaded and extracted |
| 90 | + if [ -d "${INSTALL_DIR}/extracted/cuda_cudart" ]; then |
| 91 | + echo "Windows CUDA toolkit already installed, skipping download..." |
| 92 | + return 0 |
| 93 | + fi |
| 94 | + |
| 95 | + echo "Downloading CUDA installer from ${CUDA_URL}..." |
| 96 | + wget -q "${CUDA_URL}" -O "${CUDA_INSTALLER}" |
| 97 | + |
| 98 | + echo "Extracting CUDA toolkit..." |
| 99 | + 7z x "${CUDA_INSTALLER}" -o"extracted" -y |
| 100 | + |
| 101 | + # Fix permissions so ci-user can access the files |
| 102 | + chmod -R a+rX "${INSTALL_DIR}" |
| 103 | + |
| 104 | + # Clean up installer to save space |
| 105 | + rm -f "${CUDA_INSTALLER}" |
| 106 | + |
| 107 | + echo "Windows CUDA toolkit installation complete" |
| 108 | + echo "WINDOWS_CUDA_HOME=${INSTALL_DIR}/extracted/cuda_cudart/cudart" |
| 109 | +} |
| 110 | + |
| 111 | +# Parse command line arguments |
| 112 | +INSTALL_MINGW=false |
| 113 | +INSTALL_CUDA=false |
| 114 | + |
| 115 | +while [[ $# -gt 0 ]]; do |
| 116 | + case $1 in |
| 117 | + --mingw) |
| 118 | + INSTALL_MINGW=true |
| 119 | + shift |
| 120 | + ;; |
| 121 | + --cuda) |
| 122 | + INSTALL_CUDA=true |
| 123 | + shift |
| 124 | + ;; |
| 125 | + --all) |
| 126 | + INSTALL_MINGW=true |
| 127 | + INSTALL_CUDA=true |
| 128 | + shift |
| 129 | + ;; |
| 130 | + *) |
| 131 | + echo "Unknown option: $1" |
| 132 | + echo "Usage: $0 [--mingw] [--cuda] [--all]" |
| 133 | + exit 1 |
| 134 | + ;; |
| 135 | + esac |
| 136 | +done |
| 137 | + |
| 138 | +# Default to installing everything if no options specified |
| 139 | +if [ "${INSTALL_MINGW}" = false ] && [ "${INSTALL_CUDA}" = false ]; then |
| 140 | + INSTALL_MINGW=true |
| 141 | + INSTALL_CUDA=true |
| 142 | +fi |
| 143 | + |
| 144 | +if [ "${INSTALL_MINGW}" = true ]; then |
| 145 | + install_mingw |
| 146 | +fi |
| 147 | + |
| 148 | +if [ "${INSTALL_CUDA}" = true ]; then |
| 149 | + install_windows_cuda |
| 150 | +fi |
| 151 | + |
| 152 | +echo "Installation complete" |
0 commit comments