Add CUDA 12.1 workflows

.circleci/scripts/binary_linux_build.sh

@@ -11,7 +11,7 @@ NUM_CPUS=$(( $(nproc) - 2 ))
 # Defaults here for **binary** linux builds so they can be changed in one place
 export MAX_JOBS=${MAX_JOBS:-$(( ${NUM_CPUS} > ${MEMORY_LIMIT_MAX_JOBS} ? ${MEMORY_LIMIT_MAX_JOBS} : ${NUM_CPUS} ))}
 
-if [[ "${DESIRED_CUDA}" =~ cu11[0-9] ]]; then
+if [[ "${DESIRED_CUDA}" =~ cu1[1-2][0-9] ]]; then
   export BUILD_SPLIT_CUDA="ON"
 fi
 

.circleci/scripts/binary_windows_build.sh

@@ -10,7 +10,7 @@ export SCCACHE_BUCKET=ossci-compiler-cache
 export SCCACHE_IGNORE_SERVER_IO_ERROR=1
 export VC_YEAR=2022
 
-if [[ "${DESIRED_CUDA}" == *"cu11"* ]]; then
+if [[ "${DESIRED_CUDA}" =~ cu1[1-2][0-9] ]]; then
     export BUILD_SPLIT_CUDA=ON
 fi
 

.github/scripts/generate_binary_build_matrix.py

@@ -13,7 +13,7 @@
 from typing import Dict, List, Optional, Tuple
 
 
-CUDA_ARCHES = ["11.7", "11.8"]
+CUDA_ARCHES = ["11.7", "11.8", "12.1"]
 
 
 ROCM_ARCHES = ["5.3", "5.4.2"]
@@ -99,6 +99,9 @@ def generate_conda_matrix(os: str) -> List[Dict[str, str]]:
     python_versions = FULL_PYTHON_VERSIONS
     if os == "linux" or os == "windows":
         arches += CUDA_ARCHES
+        # skip CUDA 12.1 builds on Windows
+        if os == "windows" and "12.1" in arches:
+            arches.remove("12.1")
     for python_version in python_versions:
         # We don't currently build conda packages for rocm
         for arch_version in arches:
@@ -135,6 +138,9 @@ def generate_libtorch_matrix(
             arches += ROCM_ARCHES
         elif os == "windows":
             arches += CUDA_ARCHES
+            # skip CUDA 12.1 builds on Windows
+            if "12.1" in arches:
+                arches.remove("12.1")
 
     if libtorch_variants is None:
         libtorch_variants = [
@@ -199,6 +205,9 @@ def generate_wheels_matrix(
             arches += CPU_CXX11_ABI_ARCH + CUDA_ARCHES + ROCM_ARCHES
         elif os == "windows":
             arches += CUDA_ARCHES
+            # skip CUDA 12.1 builds on Windows
+            if "12.1" in arches:
+                arches.remove("12.1")
 
     ret: List[Dict[str, str]] = []
     for python_version in python_versions: