Add build option to gpu-system.sh

cuda.sh

@@ -0,0 +1,82 @@
+package: CUDA
+version: "13.3.1_cudnn_9.25.1.1"
+license: CUDA
+build_requires:
+  - alibuild-recipe-tools
+---
+#!/bin/bash -e
+DRIVER_VERSION=610.43.02
+CUDA_VERSION="${FOO%%_cudnn_*}"
+CUDNN_VERSION="${FOO#*_cudnn_}"
+FILENAME=cuda_${CUDA_VERSION}_${DRIVER_VERSION}_linux.run
+if [[ -n $ALIBUILD_O2_FORCE_GPU_BUILDSOURCES ]]; then
+  cp "${ALIBUILD_O2_FORCE_GPU_BUILDSOURCES}/${FILENAME}" ./
+else
+  curl -o "${FILENAME}" "https://developer.download.nvidia.com/compute/cuda/${CUDA_VERSION}/local_installers/${FILENAME}"
+fi
+
+exclude_list=(
+  "cuda-installer"
+  "*-uninstaller"
+  "NVIDIA-Linux-${narch}-${DRIVER_PV}.run"
+  "builds/NVIDIA-Linux*"
+  "builds/*.txt"
+  "builds/*.json"
+  "builds/cuda_documentation"
+  "builds/cuda_gdb"
+  "builds/cuda_sanitizer_api"
+  "builds/integration"
+  "builds/cuda_nsight"
+  "builds/cuda_nvvp"
+  "builds/nsight_compute"
+  "builds/nsight_systems"
+  "builds/nvidia_fs"
+)
+
+bash "${FILENAME}" --tar xf -X <(printf "%s\n" "${exclude_list[@]}")
+
+merge_move()
+{
+  local src=$1
+  local dst=$2
+
+  mkdir -p "$dst"
+  shopt -s dotglob nullglob
+
+  local x
+  for x in "$src"/*; do
+    local name=${x##*/}
+    local target="$dst/$name"
+
+    if [[ ! -L "$x" && -d "$x" && -d "$target" ]]; then
+      merge_move "$x" "$target"
+      rmdir "$x"
+    elif [[ -L "$x" && -L "$target" ]]; then
+      rm -- "$x"
+    else
+      mv -- "$x" "$target"
+    fi
+  done
+}
+
+for i in builds/*/; do
+  merge_move "$i" "$INSTALLROOT"
+done
+
+rm -Rf builds ${FILENAME}
+
+FILENAME=cudnn-linux-x86_64-${CUDNN_VERSION}_cuda13-archive.tar.xz
+if [[ -n $ALIBUILD_O2_FORCE_GPU_BUILDSOURCES ]]; then
+  cp "${ALIBUILD_O2_FORCE_GPU_BUILDSOURCES}/${FILENAME}" ./
+else
+  curl -o "${FILENAME}" "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/${FILENAME}"
+fi
+
+tar -C $INSTALLROOT ${FILENAME}
+
+# Modulefile
+mkdir -p "$INSTALLROOT/etc/modulefiles"
+alibuild-generate-module --bin --lib > "$INSTALLROOT/etc/modulefiles/$PKGNAME"
+cat >> "$INSTALLROOT/etc/modulefiles/$PKGNAME" <<EoF
+setenv CUDA_PATH \$PKG_ROOT
+EoF

gpu-system.sh

@@ -35,8 +35,14 @@
     # - ci: for now defaults to 1
     # - disable: disable all backends
     # - manual: disable auto-detection, set features manually
+    # - build: build GPU tools from alidist recipe
     # 0 / unset: defaults to fullauto
 
+    if [[ ${ALIBUILD_O2_FORCE_GPU} == "build" ]]; then
+      GPU_FEATURES=build
+      break
+    fi
+
     if [[ -z ${ALIBUILD_O2_FORCE_GPU} || ${ALIBUILD_O2_FORCE_GPU} == "0" ]]; then
       ALIBUILD_O2_FORCE_GPU=fullauto
     fi
@@ -264,13 +270,33 @@
       #%Module1.0
       echo "ERROR: gpu-system.sh GPU detection failed: ${ALIBUILD_PREFER_SYSTEM_KEY}" | sed "s/error-//" 1>&2
       exit 1
+  "build":
+    version: "build-from-alidist"
+    requires:
+      - CUDA
+      - ROCm
+    recipe: |
+      #!/bin/bash -e
+      mkdir -p "$INSTALLROOT/etc/modulefiles"
+      alibuild-generate-module > "$INSTALLROOT/etc/modulefiles/$PKGNAME"
+      {
+        echo 'export O2_GPU_CUDA_AVAILABLE=1'
+        echo 'export O2_GPU_ROCM_AVAILABLE=1'
+        echo 'export O2_GPU_CUDA_HOME="${CUDA_PATH}"'
+        echo 'export O2_GPU_ROCM_HOME="${ROCM_PATH}"'
+        echo 'export O2_GPU_MIOPEN_AVAILABLE=1'
+        echo 'export O2_GPU_CUDNN_AVAILABLE=1'
+        echo 'export O2_GPU_MIGRAPHX_AVAILABLE=1'
+        echo 'export O2_GPU_TENSORRT_AVAILABLE=0'
+        echo 'O2_GPU_CUDA_AVAILABLE_ARCH="'${ALIBUILD_O2_FORCE_GPU_CUDA_ARCH:-default}'"'
+        echo 'O2_GPU_ROCM_AVAILABLE_ARCH="'${ALIBUILD_O2_FORCE_GPU_ROCM_ARCH:-default}'"'
+      } > "$INSTALLROOT"/etc/gpu-features-available.sh
   ".*":
     version: "%(key)s"
     recipe: |
       #!/bin/bash -e
-      #%Module1.0
-      mkdir -p "$INSTALLROOT"/etc
-      rm -f "$INSTALLROOT"/etc/gpu-features-available.sh
+      mkdir -p "$INSTALLROOT/etc/modulefiles"
+      alibuild-generate-module > "$INSTALLROOT/etc/modulefiles/$PKGNAME"
       {
         echo "export O2_GPU_CUDA_AVAILABLE=\"$( ( [[ "$PKG_VERSION" =~ (^|-)cuda(-|_|$) ]] && echo 1 ) || ( [[ "$PKG_VERSION" =~ (^|-)auto(-|_|$) ]] && echo auto || echo 0 ) )\""
         echo "export O2_GPU_ROCM_AVAILABLE=\"$( ( [[ "$PKG_VERSION" =~ (^|-)rocm(-|_|$) ]] && echo 1 ) || ( [[ "$PKG_VERSION" =~ (^|-)auto(-|_|$) ]] && echo auto || echo 0 ) )\""
@@ -286,4 +312,16 @@
         [[ "$PKG_VERSION" =~ (^|-)rocm([^-]*)_arch@([^@]*)@(-|_|$) ]] && echo 'export O2_GPU_ROCM_AVAILABLE_ARCH="'$(sed -e 's/#/;/g' -e 's/_/-/g' <<< "${BASH_REMATCH[3]}" 2> /dev/null)'"'
         true
       } > "$INSTALLROOT"/etc/gpu-features-available.sh
+  "build":
+    requires:
+      - CUDA
+      - ROCm
+    recipe: |
+      #!/bin/bash -e
+      #%Module1.0
+      mkdir -p "$INSTALLROOT"/etc
+      rm -f "$INSTALLROOT"/etc/gpu-features-available.sh
+      {
+        true
+      } > "$INSTALLROOT"/etc/gpu-features-available.sh
 ---

python-modules.sh

@@ -5,6 +5,7 @@ requires:
   - "FreeType:(?!osx)"
   - libpng
   - hdf5
+  - gpu-system
 license: Python-2.0
 build_requires:
   - Python-modules-list
@@ -38,11 +39,16 @@ prefer_system_replacement_specs:
       # import name BuildDependencyInstallError". --force-reinstall uninstalls
       # first, so nothing survives, while still reinstalling a version the
       # environment already satisfies (which is why -I was here).
+      $GPU_SYSTEM_VERSION
       echo "$PIP_BASE_REQUIREMENTS" > base-requirements.txt
       python3 -m pip install --force-reinstall -U -r base-requirements.txt
       # The above updates pip and setuptools, so install the rest of the packages separately.
       echo "$PIP_REQUIREMENTS" > requirements.txt
       python3 -m pip install -IU -r requirements.txt
+      if [[ -n $PIP_ROCM_REQUIREMENTS ]]; then
+        echo "$PIP_ROCM_REQUIREMENTS" > requirements.txt
+        python3 -m pip install -IU -r requirements.txt
+      fi
       # We do not need anything else, because python is going to be in path
       # if we are inside a virtualenv so no need to pretend we know where
       # the correct python is.

rocm.sh

@@ -0,0 +1,65 @@
+package: ROCm
+version: "7.14.1"
+license: MIT
+build_requires:
+  - alibuild-recipe-tools
+env:
+  PIP_ROCM_REQUIREMENTS: "migraphx==2.17.0+rocm7.14.1"
+---
+#!/bin/bash -e
+ROCM_SUBVERSION=1
+FULLVERSION=${PKGVERSION}
+if [[ ! $FULLVERSION =~ \..*\. ]]; then FULLVERSION+=.0; fi
+
+FILENAME=rocm-installer-${FULLVERSION}-${ROCM_SUBVERSION}.run
+if [[ -n $ALIBUILD_O2_FORCE_GPU_BUILDSOURCES ]]; then
+  cp "${ALIBUILD_O2_FORCE_GPU_BUILDSOURCES}/${FILENAME}" ./
+else
+  curl -o "${FILENAME}" "https://repo.radeon.com/rocm/installer/rocm-runfile-installer/rocm-rel-${PKGVERSION}/${FILENAME}"
+fi
+
+merge_move()
+{
+  local src=$1
+  local dst=$2
+
+  mkdir -p "$dst"
+  shopt -s dotglob nullglob
+
+  local x
+  for x in "$src"/*; do
+    local name=${x##*/}
+    local target="$dst/$name"
+
+    if [[ ! -L "$x" && -d "$x" && -d "$target" ]]; then
+      merge_move "$x" "$target"
+      rmdir "$x"
+    elif [[ -L "$x" && -L "$target" ]]; then
+      rm -- "$x"
+    else
+      mv -- "$x" "$target"
+    fi
+  done
+}
+
+bash rocm-installer-7.14.1-1.run rocm gfx=all compo=core-sdk nopostrocm verbose noexec noexec-cleanup || [[ -d rocm-installer ]]
+
+for i in $(ls rocm-installer/component-rocm/ | grep '.xz$'); do
+  mkdir unpack
+  tar -C unpack -Jxf rocm-installer/component-rocm/$i
+  for k in unpack/*; do
+    pushd $k
+    for j in *; do
+      merge_move $j/*/*/ $INSTALLROOT
+    done
+    popd
+  done
+  rm -Rf unpack
+done
+
+# Modulefile
+mkdir -p "$INSTALLROOT/etc/modulefiles"
+alibuild-generate-module --bin --lib > "$INSTALLROOT/etc/modulefiles/$PKGNAME"
+cat >> "$INSTALLROOT/etc/modulefiles/$PKGNAME" <<EoF
+setenv ROCM_PATH \$PKG_ROOT
+EoF