BUILD(cfg): Add build toolchain for CUDA variant

MikeSWang · Sep 10, 2024 · 6f6f04e · 6f6f04e
1 parent c00b62f
commit 6f6f04e
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 29 deletions.
diff --git a/.github/workflows/cd_cuda.yml b/.github/workflows/cd_cuda.yml
@@ -76,12 +76,6 @@ jobs:
           package-dir: .
           output-dir: dist/
           config-file: pyproject.toml
-        env:
-          CIBW_ENVIRONMENT: >-
-            PY_CXX=nvcc
-            PY_OMP=1
-            PY_CUDA=1
-            PY_BUILD_PARALLEL='-j'
 
       - name: Verify built distribution
         run: python -m twine check --strict dist/*
@@ -124,13 +118,15 @@ jobs:
         with:
           fetch-depth: 0
 
-      - uses: Jimver/cuda-toolkit@master
+      - name: Set up CUDA Toolkit
+        uses: Jimver/cuda-toolkit@master
         id: cuda-toolkit
         with:
-          log-file-suffix: '${{matrix.os}}_cudatoolkit_log.txt'
+          # Match CUDA Toolkit version with conda-forge cudatookit package.
+          cuda: '11.8.0'
+          log-file-suffix: '${{matrix.os}}_log.txt'
           method: 'network'
           sub-packages: '["nvcc"]'
-          non-cuda-sub-packages: '["libcufft"]'
 
       - name: Set up (Mini)conda
         uses: conda-incubator/setup-miniconda@v3
@@ -149,24 +145,27 @@ jobs:
 
       - name: Override version
         if: github.event_name == 'workflow_dispatch'
-        uses: knicknic/os-specific-run@v1.0.4
         env:
           recipe_file: deploy/pkg/conda_recipe_cuda/meta.yaml
-        with:
-          linux: |
-            vers_tag=${{ github.event.inputs.version_tag }}
-            cuda_vers=${{ steps.cuda-toolkit.outputs.cuda }}
-            if [[ ! -z ${vers_tag} ]]; then
-              sed -i "s/# git_rev:.*/git_rev: ${vers_tag}/g" ${recipe_file}
-              sed -i "s/'CUDA_VERSION', ''/'CUDA_VERSION', '${cuda_vers}'/g" ${recipe_file}
-            fi
+        run: |
+          vers_tag=${{ github.event.inputs.version_tag }}
+          cuda_vers=${{ steps.cuda-toolkit.outputs.cuda }}
+          if [[ ! -z ${vers_tag} ]]; then
+            sed -i "s/# git_rev:.*/git_rev: ${vers_tag}/g" ${recipe_file}
+          fi
+          if [[ ! -z ${cuda_vers} ]]; then
+            sed -i "s/'CUDA_VERSION', '11.8.0'/'CUDA_VERSION', '${cuda_vers}'/g" ${recipe_file}
+          fi
 
       - name: Build Conda package
         env:
           recipe_dir: deploy/pkg/conda_recipe_cuda
           output_dir: dist/
           variants: "{'python': ['3.10', '3.11', '3.12']}"
         run: |
+          # Fetch environment variables.
+          export CUDA_PATH=${{ steps.cuda-toolkit.outputs.CUDA_PATH }}
+          export CUDA_VERSION=${{ steps.cuda-toolkit.outputs.cuda }}
           # Create output directory if non-existent.
           if [[ ! -d ${output_dir} ]]; then mkdir -p ${output_dir}; fi
           # Build.

diff --git a/.pyproject_cuda.toml b/.pyproject_cuda.toml
@@ -105,8 +105,19 @@ manylinux-aarch64-image = 'manylinux_2_28'
 
 [tool.cibuildwheel.linux]
 before-all = [
-    "yum install -y gsl-devel cuda-toolkit",
+    "yum install -y gsl-devel",
+# Install CUDA Toolkit inside Docker container, matching repository version
+# with image OS, and optionally cuda-toolkit version with conda-forge
+# cudatoolkit package version for Pip/Conda consistency.
+    "yum install -y yum-utils",
+    "yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo",
+    "yum install -y cuda-toolkit",
+    # "yum install -y cuda-toolkit-11-8",
 ]
+before-build = [
+    "export CUDA_HOME=/usr/local/cuda",
+]
+environment = { PY_CXX="${CUDA_HOME}/bin/nvcc", PY_CXXFLAGS="-I${CUDA_HOME}/include", PY_LDFLAGS="-L${CUDA_HOME}/lib64", PY_OMP='1', PY_CUDA='1', PY_BUILD_PARALLEL='-j' }
 
 [tool.autopep8]
 in-place = true

diff --git a/deploy/pkg/conda_recipe_cuda/conda_build_config.yaml b/deploy/pkg/conda_recipe_cuda/conda_build_config.yaml
@@ -1,3 +1,3 @@
 target_platform:
   - linux-64       # [linux]
-  - linux-aarch64  # [linux]
+  # - linux-aarch64  # [linux]
diff --git a/deploy/pkg/conda_recipe_cuda/meta.yaml b/deploy/pkg/conda_recipe_cuda/meta.yaml
@@ -1,6 +1,9 @@
 {% set name = 'Triumvirate-CUDA' %}
 {% set version = environ.get('GIT_DESCRIBE_TAG', 'v0.5.0') %}
-{% set cuda_vers = environ.get('CUDA_VERSION', '') %}
+{% set cuda_vers = environ.get('CUDA_VERSION', '11.8.0') %}
+{% set cuda_vers_parts = cuda_vers.split('.') %}
+{% set cuda_vers_major_minor = cuda_vers_parts[0] ~ cuda_vers_parts[1] %}
+{% set cuda_vers_int = cuda_vers_major_minor|int %}
 
 package:
   name: "{{ name|lower }}"
@@ -13,22 +16,23 @@ source:
 
 build:
   number: 0
-  string: cuda{{ cuda_vers }}_py{{ CONDA_PY }}h{{ PKG_HASH }}_{{ PKG_BUILDNUM }}
+  string: cuda{{ cuda_vers_int }}py{{ CONDA_PY }}h{{ PKG_HASH }}_{{ PKG_BUILDNUM }}
   script:
-    # Use Conda-provided compiler.
-    - export PY_CXX=$CXX
-    - export PY_CXXFLAGS="${PY_CXXFLAGS} -D_LIBCPP_DISABLE_AVAILABILITY"
-    # Enforce OpenMP support.
-    - export PY_OMP=1
-    - export PY_CUDA=1
+    # Use host compiler.
+    - export PY_CXX=${CUDA_PATH}/bin/nvcc
+    - export PY_CXXFLAGS="${PY_CXXFLAGS} -I${CUDA_PATH}/include"
+    - export PY_LDFLAGS="${PY_LDFLAGS} -L${CUDA_PATH}/lib64"
     - cp .pyproject_cuda.toml pyproject.toml
     - {{ PYTHON }} -m pip install . -vvv
   script_env:
+    # Enforce OpenMP and CUDA support.
+    - PY_OMP=1
+    - PY_CUDA=1
     - PY_BUILD_PARALLEL="-j"
 
 requirements:
   build:
-    - {{ compiler('cuda') }}
+    - cudatoolkit >={{ cuda_vers }}
     - python                              # [build_platform != target_platform]
     - cross-python_{{ target_platform }}  # [build_platform != target_platform]
     - numpy >=2.0                         # [build_platform != target_platform]
@@ -47,6 +51,7 @@ requirements:
     - extension-helpers >=1.1
     - numpy >=2.0
   run:
+    - cudatoolkit >={{ cuda_vers }}
     - python >=3.10
     - numpy >=1.23
     - scipy >=1.13