From 416ca1c2bb430ff8064010f464bdd1c008f395cc Mon Sep 17 00:00:00 2001
From: Mike S Wang <32841762+MikeSWang@users.noreply.github.com>
Date: Tue, 10 Sep 2024 17:27:53 +0100
Subject: [PATCH] BUILD(cfg): Add build toolchain for CUDA variant

---
 .github/workflows/cd_cuda.yml                 | 33 +++++++++++--------
 .pyproject_cuda.toml                          | 10 +++++-
 .../conda_recipe_cuda/conda_build_config.yaml |  2 +-
 deploy/pkg/conda_recipe_cuda/meta.yaml        | 16 +++++----
 4 files changed, 39 insertions(+), 22 deletions(-)

diff --git a/.github/workflows/cd_cuda.yml b/.github/workflows/cd_cuda.yml
index 11fafa58..e37cf156 100644
--- a/.github/workflows/cd_cuda.yml
+++ b/.github/workflows/cd_cuda.yml
@@ -78,7 +78,9 @@ jobs:
           config-file: pyproject.toml
         env:
           CIBW_ENVIRONMENT: >-
-            PY_CXX=nvcc
+            PY_CXX=/usr/local/cuda/bin/nvcc
+            PY_CXXFLAGS=-I/usr/local/cuda/include
+            PY_LDFLAGS=-L/usr/local/cuda/lib64
             PY_OMP=1
             PY_CUDA=1
             PY_BUILD_PARALLEL='-j'
@@ -124,13 +126,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,17 +153,17 @@ 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', ''/'CUDA_VERSION', '${cuda_vers}'/g" ${recipe_file}
+          fi
 
       - name: Build Conda package
         env:
@@ -167,6 +171,9 @@ jobs:
           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
index 752e7bb9..b5e6ff12 100644
--- a/.pyproject_cuda.toml
+++ b/.pyproject_cuda.toml
@@ -105,7 +105,15 @@ 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",
+    "echo CUDA_PATH=$CUDA_PATH",
 ]
 
 [tool.autopep8]
diff --git a/deploy/pkg/conda_recipe_cuda/conda_build_config.yaml b/deploy/pkg/conda_recipe_cuda/conda_build_config.yaml
index 311f7da7..3fa18852 100644
--- 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
index 4cd00606..d2c7494d 100644
--- a/deploy/pkg/conda_recipe_cuda/meta.yaml
+++ b/deploy/pkg/conda_recipe_cuda/meta.yaml
@@ -15,20 +15,21 @@ build:
   number: 0
   string: cuda{{ cuda_vers }}_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 +48,7 @@ requirements:
     - extension-helpers >=1.1
     - numpy >=2.0
   run:
+    - cudatoolkit >={{ cuda_vers }}
     - python >=3.10
     - numpy >=1.23
     - scipy >=1.13