Windows cuda export ci and set up script

.ci/docker/build.sh

@@ -67,6 +67,13 @@ case "${IMAGE_NAME}" in
     # From https://developer.android.com/ndk/downloads
     ANDROID_NDK_VERSION=r28c
     ;;
+  executorch-ubuntu-22.04-cuda-windows)
+    LINTRUNNER=""
+    GCC_VERSION=11
+    CUDA_WINDOWS_CROSS_COMPILE=yes
+    CUDA_VERSION=12.8
+    SKIP_PYTORCH=yes
+    ;;
   *)
     echo "Invalid image name ${IMAGE_NAME}"
     exit 1
@@ -101,6 +108,8 @@ docker build \
   --build-arg "MEDIATEK_SDK=${MEDIATEK_SDK:-}" \
   --build-arg "ANDROID_NDK_VERSION=${ANDROID_NDK_VERSION:-}" \
   --build-arg "SKIP_PYTORCH=${SKIP_PYTORCH:-}" \
+  --build-arg "CUDA_WINDOWS_CROSS_COMPILE=${CUDA_WINDOWS_CROSS_COMPILE:-}" \
+  --build-arg "CUDA_VERSION=${CUDA_VERSION:-}" \
   -f "${OS}"/Dockerfile \
   "$@" \
   .

.ci/docker/common/install_conda.sh

@@ -54,25 +54,6 @@ install_pip_dependencies() {
   popd
 }
 
-fix_conda_ubuntu_libstdcxx() {
-  cat /etc/issue
-  # WARNING: This is a HACK from PyTorch core to be able to build PyTorch on 22.04.
-  # Specifically, ubuntu-20+ all comes lib libstdc++ newer than 3.30+, but anaconda
-  # is stuck with 3.29. So, remove libstdc++6.so.3.29 as installed by
-  # https://anaconda.org/anaconda/libstdcxx-ng/files?version=11.2.0
-  #
-  # PyTorch sev: https://github.com/pytorch/pytorch/issues/105248
-  # Ref: https://github.com/pytorch/pytorch/blob/main/.ci/docker/common/install_conda.sh
-  if grep -e "2[02].04." /etc/issue >/dev/null; then
-    rm /opt/conda/envs/py_${PYTHON_VERSION}/lib/libstdc++.so*
-  fi
-}
-
 install_miniconda
 install_python
 install_pip_dependencies
-# Hack breaks the job on aarch64 but is still necessary everywhere
-# else.
-if [ "$(uname -m)" != "aarch64" ]; then
-    fix_conda_ubuntu_libstdcxx
-fi

.ci/docker/common/install_cuda.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Install Linux CUDA toolkit
+# This installs nvcc and other CUDA development tools needed for compiling CUDA code
+
+set -ex
+
+# CUDA version must be specified (e.g., 12.8)
+CUDA_VERSION="${CUDA_VERSION:?CUDA_VERSION must be set}"
+
+# Convert version format (e.g., 12.8 -> 12-8 for package names)
+CUDA_VERSION_DASH=$(echo "${CUDA_VERSION}" | tr '.' '-')
+
+# Add NVIDIA package repository
+apt-get update
+apt-get install -y --no-install-recommends \
+    gnupg2 \
+    ca-certificates \
+    wget
+
+# Download and install the CUDA keyring
+wget -q "https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb" -O /tmp/cuda-keyring.deb
+dpkg -i /tmp/cuda-keyring.deb
+rm /tmp/cuda-keyring.deb
+
+apt-get update
+
+# Install CUDA toolkit (nvcc and development libraries)
+# We install a minimal set of packages needed for compilation:
+# - cuda-nvcc: The CUDA compiler
+# - cuda-cudart-dev: CUDA runtime development files
+# - cuda-nvrtc-dev: CUDA runtime compilation library
+# - libcublas-dev: cuBLAS development files
+# - libcusparse-dev: cuSPARSE development files
+# - libcufft-dev: cuFFT development files
+apt-get install -y --no-install-recommends \
+    "cuda-nvcc-${CUDA_VERSION_DASH}" \
+    "cuda-cudart-dev-${CUDA_VERSION_DASH}" \
+    "cuda-nvrtc-dev-${CUDA_VERSION_DASH}" \
+    "libcublas-dev-${CUDA_VERSION_DASH}" \
+    "libcusparse-dev-${CUDA_VERSION_DASH}" \
+    "libcufft-dev-${CUDA_VERSION_DASH}"
+
+# Clean up
+apt-get clean
+rm -rf /var/lib/apt/lists/*
+
+# Verify installation
+/usr/local/cuda-${CUDA_VERSION}/bin/nvcc --version
+
+echo "CUDA ${CUDA_VERSION} toolkit installation complete"
+echo "CUDA_HOME=/usr/local/cuda-${CUDA_VERSION}"

.ci/docker/common/install_cuda_windows_cross_compile.sh

@@ -0,0 +1,146 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Install mingw-w64 cross-compiler and Windows CUDA toolkit for cross-compilation
+
+set -ex
+
+INSTALL_DIR="${WINDOWS_CUDA_INSTALL_DIR:-/opt/cuda-windows}"
+
+# Mapping of CUDA versions to their corresponding driver versions for Windows installers
+# Source: https://developer.nvidia.com/cuda-toolkit-archive
+declare -A CUDA_DRIVER_MAP=(
+    ["12.6"]="12.6.3:561.17"
+    ["12.8"]="12.8.1:572.61"
+    ["12.9"]="12.9.1:576.57"
+)
+
+install_mingw() {
+    echo "Installing mingw-w64 cross-compiler..."
+
+    apt-get update
+    apt-get install -y --no-install-recommends \
+        g++-mingw-w64-x86-64 \
+        mingw-w64-tools \
+        p7zip-full \
+        wget
+
+    # Verify installation
+    x86_64-w64-mingw32-g++ --version
+
+    # Cleanup
+    apt-get clean
+    rm -rf /var/lib/apt/lists/*
+
+    echo "mingw-w64 installation complete"
+}
+
+get_torch_cuda_version() {
+    # Query PyTorch for its CUDA version using conda environment
+    conda run -n "py_${PYTHON_VERSION}" python3 -c "import torch; print(torch.version.cuda)" 2>/dev/null || echo ""
+}
+
+install_windows_cuda() {
+    # Get CUDA version from torch
+    TORCH_CUDA_VERSION=$(get_torch_cuda_version)
+
+    if [ -z "${TORCH_CUDA_VERSION}" ] || [ "${TORCH_CUDA_VERSION}" = "None" ]; then
+        echo "ERROR: Could not detect CUDA version from PyTorch."
+        echo "Make sure PyTorch with CUDA support is installed before running this script."
+        exit 1
+    fi
+
+    echo "Detected PyTorch CUDA version: ${TORCH_CUDA_VERSION}"
+
+    # Extract major.minor version (e.g., "12.8" from "12.8.1" or "12.8")
+    CUDA_MAJOR_MINOR=$(echo "${TORCH_CUDA_VERSION}" | cut -d. -f1,2)
+
+    # Look up the full version and driver version
+    if [ -z "${CUDA_DRIVER_MAP[${CUDA_MAJOR_MINOR}]}" ]; then
+        echo "ERROR: CUDA version ${CUDA_MAJOR_MINOR} is not in the known version map."
+        echo "Known versions: ${!CUDA_DRIVER_MAP[*]}"
+        exit 1
+    fi
+
+    CUDA_INFO="${CUDA_DRIVER_MAP[${CUDA_MAJOR_MINOR}]}"
+    CUDA_VERSION=$(echo "${CUDA_INFO}" | cut -d: -f1)
+    CUDA_DRIVER_VERSION=$(echo "${CUDA_INFO}" | cut -d: -f2)
+
+    echo "Using CUDA ${CUDA_VERSION} with driver ${CUDA_DRIVER_VERSION}"
+
+    echo "Installing Windows CUDA toolkit ${CUDA_VERSION}..."
+
+    mkdir -p "${INSTALL_DIR}"
+    cd "${INSTALL_DIR}"
+
+    CUDA_INSTALLER="cuda_${CUDA_VERSION}_${CUDA_DRIVER_VERSION}_windows.exe"
+    CUDA_URL="https://developer.download.nvidia.com/compute/cuda/${CUDA_VERSION}/local_installers/${CUDA_INSTALLER}"
+
+    # Check if already downloaded and extracted
+    if [ -d "${INSTALL_DIR}/extracted/cuda_cudart" ]; then
+        echo "Windows CUDA toolkit already installed, skipping download..."
+        return 0
+    fi
+
+    echo "Downloading CUDA installer from ${CUDA_URL}..."
+    wget -q "${CUDA_URL}" -O "${CUDA_INSTALLER}"
+
+    echo "Extracting CUDA toolkit..."
+    7z x "${CUDA_INSTALLER}" -o"extracted" -y
+
+    # Fix permissions so ci-user can access the files
+    chmod -R a+rX "${INSTALL_DIR}"
+
+    # Clean up installer to save space
+    rm -f "${CUDA_INSTALLER}"
+
+    echo "Windows CUDA toolkit installation complete"
+    echo "WINDOWS_CUDA_HOME=${INSTALL_DIR}/extracted/cuda_cudart/cudart"
+}
+
+# Parse command line arguments
+INSTALL_MINGW=false
+INSTALL_CUDA=false
+
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --mingw)
+            INSTALL_MINGW=true
+            shift
+            ;;
+        --cuda)
+            INSTALL_CUDA=true
+            shift
+            ;;
+        --all)
+            INSTALL_MINGW=true
+            INSTALL_CUDA=true
+            shift
+            ;;
+        *)
+            echo "Unknown option: $1"
+            echo "Usage: $0 [--mingw] [--cuda] [--all]"
+            exit 1
+            ;;
+    esac
+done
+
+# Default to installing everything if no options specified
+if [ "${INSTALL_MINGW}" = false ] && [ "${INSTALL_CUDA}" = false ]; then
+    INSTALL_MINGW=true
+    INSTALL_CUDA=true
+fi
+
+if [ "${INSTALL_MINGW}" = true ]; then
+    install_mingw
+fi
+
+if [ "${INSTALL_CUDA}" = true ]; then
+    install_windows_cuda
+fi
+
+echo "Installation complete"

.ci/docker/common/install_pytorch_cuda.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Install PyTorch with CUDA support from prebuilt wheels
+# This is used for the cuda-windows Docker image to get a specific CUDA version
+
+set -ex
+
+# shellcheck source=/dev/null
+source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
+
+# Default CUDA version if not specified
+CUDA_VERSION="${CUDA_VERSION:-12.8}"
+
+# Ensure PYTHON_VERSION is set (should be set by Dockerfile ENV)
+if [ -z "${PYTHON_VERSION}" ]; then
+    echo "ERROR: PYTHON_VERSION environment variable is not set"
+    exit 1
+fi
+
+echo "Using Python version: ${PYTHON_VERSION}"
+
+# Convert CUDA version to PyTorch wheel suffix (e.g., 12.8 -> cu128)
+CUDA_SUFFIX="cu$(echo ${CUDA_VERSION} | tr -d '.')"
+
+echo "Installing PyTorch with CUDA ${CUDA_VERSION} (${CUDA_SUFFIX})..."
+
+# Install PyTorch from nightly with specific CUDA version into the conda environment
+pip_install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/nightly/${CUDA_SUFFIX}"
+
+# Verify installation
+conda_run python -c "import torch; print(f'PyTorch {torch.__version__} installed with CUDA {torch.version.cuda}')"
+
+echo "PyTorch CUDA installation complete"

.ci/docker/ubuntu/Dockerfile

@@ -98,5 +98,25 @@ ARG QNN_SDK
 
 ARG MEDIATEK_SDK
 
+ARG CUDA_WINDOWS_CROSS_COMPILE
+ARG CUDA_VERSION
+COPY ./common/install_cuda.sh install_cuda.sh
+COPY ./common/install_pytorch_cuda.sh install_pytorch_cuda.sh
+COPY ./common/install_cuda_windows_cross_compile.sh install_cuda_windows_cross_compile.sh
+COPY ./common/utils.sh utils.sh
+RUN if [ -n "${CUDA_WINDOWS_CROSS_COMPILE}" ]; then \
+    CUDA_VERSION=${CUDA_VERSION} bash ./install_cuda.sh && \
+    CUDA_VERSION=${CUDA_VERSION} bash ./install_pytorch_cuda.sh && \
+    bash ./install_cuda_windows_cross_compile.sh; \
+    fi
+RUN rm -f install_cuda.sh install_pytorch_cuda.sh install_cuda_windows_cross_compile.sh utils.sh
+# Set up CUDA environment for Linux compilation (nvcc, etc.)
+ENV CUDA_HOME=/usr/local/cuda
+ENV PATH=${CUDA_HOME}/bin:${PATH}
+# Ensure system libstdc++ is found before conda's (GLIBCXX_3.4.30 compatibility)
+ENV LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
+# Windows CUDA for cross-compilation
+ENV WINDOWS_CUDA_HOME=/opt/cuda-windows/extracted/cuda_cudart/cudart
+
 USER ci-user
 CMD ["bash"]

.ci/scripts/export_model_artifact.sh

@@ -58,11 +58,13 @@ OUTPUT_DIR="${4:-.}"
 case "$DEVICE" in
   cuda)
     ;;
+  cuda-windows)
+    ;;
   metal)
     ;;
   *)
     echo "Error: Unsupported device '$DEVICE'"
-    echo "Supported devices: cuda, metal"
+    echo "Supported devices: cuda, cuda-windows, metal"
     exit 1
     ;;
 esac
@@ -147,7 +149,7 @@ if [ -n "$MAX_SEQ_LEN" ]; then
 fi
 
 DEVICE_ARG=""
-if [ "$DEVICE" = "cuda" ]; then
+if [ "$DEVICE" = "cuda" ] || [ "$DEVICE" = "cuda-windows" ]; then
   DEVICE_ARG="--device cuda"
 fi
 
@@ -169,8 +171,15 @@ if [ -n "$PREPROCESSOR_OUTPUT" ]; then
       --output_file $PREPROCESSOR_OUTPUT
 fi
 
+# Determine blob file name - cuda and cuda-windows both use aoti_cuda_blob.ptd
+if [ "$DEVICE" = "cuda" ] || [ "$DEVICE" = "cuda-windows" ]; then
+  BLOB_FILE="aoti_cuda_blob.ptd"
+else
+  BLOB_FILE="aoti_${DEVICE}_blob.ptd"
+fi
+
 test -f model.pte
-test -f aoti_${DEVICE}_blob.ptd
+test -f $BLOB_FILE
 if [ -n "$PREPROCESSOR_OUTPUT" ]; then
   test -f $PREPROCESSOR_OUTPUT
 fi
@@ -179,7 +188,7 @@ echo "::endgroup::"
 echo "::group::Store $MODEL_NAME Artifacts"
 mkdir -p "${OUTPUT_DIR}"
 mv model.pte "${OUTPUT_DIR}/"
-mv aoti_${DEVICE}_blob.ptd "${OUTPUT_DIR}/"
+mv $BLOB_FILE "${OUTPUT_DIR}/"
 if [ -n "$PREPROCESSOR_OUTPUT" ]; then
   mv $PREPROCESSOR_OUTPUT "${OUTPUT_DIR}/"
 fi