Add CUDA 12.1 builds (#3284)

Summary:
CC atalman malfet

Pull Request resolved: https://github.com/pytorch/audio/pull/3284

Reviewed By: mthrok

Differential Revision: D45444670

Pulled By: atalman

fbshipit-source-id: d0cf8696a99000c2b9a7e41ceeb781f5a54daeda

.circleci/config.yml

@@ -934,6 +934,19 @@ workflows:
           python_version: '3.8'
           requires:
           - binary_windows_wheel_py3.8_cu118
+      - binary_windows_wheel:
+          cuda_version: cu121
+          name: binary_windows_wheel_py3.8_cu121
+          python_version: '3.8'
+          requires:
+          - build_ffmpeg_windows
+          wheel_docker_image: pytorch/manylinux-cuda121
+      - smoke_test_windows_pip:
+          cuda_version: cu121
+          name: binary_windows_wheel_py3.8_cu121_smoke_test_pip
+          python_version: '3.8'
+          requires:
+          - binary_windows_wheel_py3.8_cu121
       - binary_windows_wheel:
           cuda_version: cpu
           name: binary_windows_wheel_py3.9_cpu
@@ -972,6 +985,19 @@ workflows:
           python_version: '3.9'
           requires:
           - binary_windows_wheel_py3.9_cu118
+      - binary_windows_wheel:
+          cuda_version: cu121
+          name: binary_windows_wheel_py3.9_cu121
+          python_version: '3.9'
+          requires:
+          - build_ffmpeg_windows
+          wheel_docker_image: pytorch/manylinux-cuda121
+      - smoke_test_windows_pip:
+          cuda_version: cu121
+          name: binary_windows_wheel_py3.9_cu121_smoke_test_pip
+          python_version: '3.9'
+          requires:
+          - binary_windows_wheel_py3.9_cu121
       - binary_windows_wheel:
           cuda_version: cpu
           name: binary_windows_wheel_py3.10_cpu
@@ -1010,6 +1036,19 @@ workflows:
           python_version: '3.10'
           requires:
           - binary_windows_wheel_py3.10_cu118
+      - binary_windows_wheel:
+          cuda_version: cu121
+          name: binary_windows_wheel_py3.10_cu121
+          python_version: '3.10'
+          requires:
+          - build_ffmpeg_windows
+          wheel_docker_image: pytorch/manylinux-cuda121
+      - smoke_test_windows_pip:
+          cuda_version: cu121
+          name: binary_windows_wheel_py3.10_cu121_smoke_test_pip
+          python_version: '3.10'
+          requires:
+          - binary_windows_wheel_py3.10_cu121
       - binary_windows_wheel:
           cuda_version: cpu
           name: binary_windows_wheel_py3.11_cpu
@@ -1048,6 +1087,19 @@ workflows:
           python_version: '3.11'
           requires:
           - binary_windows_wheel_py3.11_cu118
+      - binary_windows_wheel:
+          cuda_version: cu121
+          name: binary_windows_wheel_py3.11_cu121
+          python_version: '3.11'
+          requires:
+          - build_ffmpeg_windows
+          wheel_docker_image: pytorch/manylinux-cuda121
+      - smoke_test_windows_pip:
+          cuda_version: cu121
+          name: binary_windows_wheel_py3.11_cu121_smoke_test_pip
+          python_version: '3.11'
+          requires:
+          - binary_windows_wheel_py3.11_cu121
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda117
           cuda_version: cu117
@@ -1100,6 +1152,19 @@ workflows:
           python_version: '3.8'
           requires:
           - binary_windows_conda_py3.8_cu118
+      - binary_windows_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          cuda_version: cu121
+          name: binary_windows_conda_py3.8_cu121
+          python_version: '3.8'
+          requires:
+          - build_ffmpeg_windows
+      - smoke_test_windows_conda_gpu:
+          cuda_version: cu121
+          name: binary_windows_conda_py3.8_cu121_smoke_test_conda
+          python_version: '3.8'
+          requires:
+          - binary_windows_conda_py3.8_cu121
       - binary_windows_conda:
           conda_docker_image: pytorch/conda-builder:cpu
           cuda_version: cpu
@@ -1139,6 +1204,19 @@ workflows:
           python_version: '3.9'
           requires:
           - binary_windows_conda_py3.9_cu118
+      - binary_windows_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          cuda_version: cu121
+          name: binary_windows_conda_py3.9_cu121
+          python_version: '3.9'
+          requires:
+          - build_ffmpeg_windows
+      - smoke_test_windows_conda_gpu:
+          cuda_version: cu121
+          name: binary_windows_conda_py3.9_cu121_smoke_test_conda
+          python_version: '3.9'
+          requires:
+          - binary_windows_conda_py3.9_cu121
       - binary_windows_conda:
           conda_docker_image: pytorch/conda-builder:cpu
           cuda_version: cpu
@@ -1178,6 +1256,19 @@ workflows:
           python_version: '3.10'
           requires:
           - binary_windows_conda_py3.10_cu118
+      - binary_windows_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          cuda_version: cu121
+          name: binary_windows_conda_py3.10_cu121
+          python_version: '3.10'
+          requires:
+          - build_ffmpeg_windows
+      - smoke_test_windows_conda_gpu:
+          cuda_version: cu121
+          name: binary_windows_conda_py3.10_cu121_smoke_test_conda
+          python_version: '3.10'
+          requires:
+          - binary_windows_conda_py3.10_cu121
       - binary_windows_conda:
           conda_docker_image: pytorch/conda-builder:cpu
           cuda_version: cpu
@@ -1217,6 +1308,19 @@ workflows:
           python_version: '3.11'
           requires:
           - binary_windows_conda_py3.11_cu118
+      - binary_windows_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          cuda_version: cu121
+          name: binary_windows_conda_py3.11_cu121
+          python_version: '3.11'
+          requires:
+          - build_ffmpeg_windows
+      - smoke_test_windows_conda_gpu:
+          cuda_version: cu121
+          name: binary_windows_conda_py3.11_cu121_smoke_test_conda
+          python_version: '3.11'
+          requires:
+          - binary_windows_conda_py3.11_cu121
       - build_docs:
           cuda_version: cu117
           filters:
@@ -1448,6 +1552,43 @@ workflows:
           python_version: '3.8'
           requires:
           - nightly_binary_windows_wheel_py3.8_cu118
+      - binary_windows_wheel:
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_wheel_py3.8_cu121
+          python_version: '3.8'
+          requires:
+          - build_ffmpeg_windows
+          wheel_docker_image: pytorch/manylinux-cuda121
+      - binary_wheel_upload:
+          context: org-member
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_wheel_py3.8_cu121_upload
+          requires:
+          - nightly_binary_windows_wheel_py3.8_cu121
+          subfolder: cu121/
+      - smoke_test_windows_pip:
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_wheel_py3.8_cu121_smoke_test_pip
+          python_version: '3.8'
+          requires:
+          - nightly_binary_windows_wheel_py3.8_cu121
       - binary_windows_wheel:
           cuda_version: cpu
           filters:
@@ -1558,6 +1699,43 @@ workflows:
           python_version: '3.9'
           requires:
           - nightly_binary_windows_wheel_py3.9_cu118
+      - binary_windows_wheel:
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_wheel_py3.9_cu121
+          python_version: '3.9'
+          requires:
+          - build_ffmpeg_windows
+          wheel_docker_image: pytorch/manylinux-cuda121
+      - binary_wheel_upload:
+          context: org-member
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_wheel_py3.9_cu121_upload
+          requires:
+          - nightly_binary_windows_wheel_py3.9_cu121
+          subfolder: cu121/
+      - smoke_test_windows_pip:
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_wheel_py3.9_cu121_smoke_test_pip
+          python_version: '3.9'
+          requires:
+          - nightly_binary_windows_wheel_py3.9_cu121
       - binary_windows_wheel:
           cuda_version: cpu
           filters:
@@ -1668,6 +1846,43 @@ workflows:
           python_version: '3.10'
           requires:
           - nightly_binary_windows_wheel_py3.10_cu118
+      - binary_windows_wheel:
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_wheel_py3.10_cu121
+          python_version: '3.10'
+          requires:
+          - build_ffmpeg_windows
+          wheel_docker_image: pytorch/manylinux-cuda121
+      - binary_wheel_upload:
+          context: org-member
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_wheel_py3.10_cu121_upload
+          requires:
+          - nightly_binary_windows_wheel_py3.10_cu121
+          subfolder: cu121/
+      - smoke_test_windows_pip:
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_wheel_py3.10_cu121_smoke_test_pip
+          python_version: '3.10'
+          requires:
+          - nightly_binary_windows_wheel_py3.10_cu121
       - binary_windows_wheel:
           cuda_version: cpu
           filters:
@@ -1778,6 +1993,43 @@ workflows:
           python_version: '3.11'
           requires:
           - nightly_binary_windows_wheel_py3.11_cu118
+      - binary_windows_wheel:
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_wheel_py3.11_cu121
+          python_version: '3.11'
+          requires:
+          - build_ffmpeg_windows
+          wheel_docker_image: pytorch/manylinux-cuda121
+      - binary_wheel_upload:
+          context: org-member
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_wheel_py3.11_cu121_upload
+          requires:
+          - nightly_binary_windows_wheel_py3.11_cu121
+          subfolder: cu121/
+      - smoke_test_windows_pip:
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_wheel_py3.11_cu121_smoke_test_pip
+          python_version: '3.11'
+          requires:
+          - nightly_binary_windows_wheel_py3.11_cu121
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda117
           cuda_version: cu117
@@ -1911,6 +2163,42 @@ workflows:
           python_version: '3.8'
           requires:
           - nightly_binary_windows_conda_py3.8_cu118
+      - binary_windows_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_conda_py3.8_cu121
+          python_version: '3.8'
+          requires:
+          - build_ffmpeg_windows
+      - binary_conda_upload:
+          context: org-member
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_conda_py3.8_cu121_upload
+          requires:
+          - nightly_binary_windows_conda_py3.8_cu121
+      - smoke_test_windows_conda_gpu:
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_conda_py3.8_cu121_smoke_test_conda
+          python_version: '3.8'
+          requires:
+          - nightly_binary_windows_conda_py3.8_cu121
       - binary_windows_conda:
           conda_docker_image: pytorch/conda-builder:cpu
           cuda_version: cpu
@@ -2019,6 +2307,42 @@ workflows:
           python_version: '3.9'
           requires:
           - nightly_binary_windows_conda_py3.9_cu118
+      - binary_windows_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_conda_py3.9_cu121
+          python_version: '3.9'
+          requires:
+          - build_ffmpeg_windows
+      - binary_conda_upload:
+          context: org-member
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_conda_py3.9_cu121_upload
+          requires:
+          - nightly_binary_windows_conda_py3.9_cu121
+      - smoke_test_windows_conda_gpu:
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_conda_py3.9_cu121_smoke_test_conda
+          python_version: '3.9'
+          requires:
+          - nightly_binary_windows_conda_py3.9_cu121
       - binary_windows_conda:
           conda_docker_image: pytorch/conda-builder:cpu
           cuda_version: cpu
@@ -2127,6 +2451,42 @@ workflows:
           python_version: '3.10'
           requires:
           - nightly_binary_windows_conda_py3.10_cu118
+      - binary_windows_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_conda_py3.10_cu121
+          python_version: '3.10'
+          requires:
+          - build_ffmpeg_windows
+      - binary_conda_upload:
+          context: org-member
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_conda_py3.10_cu121_upload
+          requires:
+          - nightly_binary_windows_conda_py3.10_cu121
+      - smoke_test_windows_conda_gpu:
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_conda_py3.10_cu121_smoke_test_conda
+          python_version: '3.10'
+          requires:
+          - nightly_binary_windows_conda_py3.10_cu121
       - binary_windows_conda:
           conda_docker_image: pytorch/conda-builder:cpu
           cuda_version: cpu
@@ -2235,3 +2595,39 @@ workflows:
           python_version: '3.11'
           requires:
           - nightly_binary_windows_conda_py3.11_cu118
+      - binary_windows_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_conda_py3.11_cu121
+          python_version: '3.11'
+          requires:
+          - build_ffmpeg_windows
+      - binary_conda_upload:
+          context: org-member
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_conda_py3.11_cu121_upload
+          requires:
+          - nightly_binary_windows_conda_py3.11_cu121
+      - smoke_test_windows_conda_gpu:
+          cuda_version: cu121
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_windows_conda_py3.11_cu121_smoke_test_conda
+          python_version: '3.11'
+          requires:
+          - nightly_binary_windows_conda_py3.11_cu121

.circleci/regenerate.py

@@ -23,8 +23,8 @@ from jinja2 import select_autoescape
 
 PYTHON_VERSIONS = ["3.8", "3.9", "3.10", "3.11"]
 CU_VERSIONS_DICT = {
-    "linux": ["cpu", "cu117", "cu118", "rocm5.2", "rocm5.3"],
-    "windows": ["cpu", "cu117", "cu118"],
+    "linux": ["cpu", "cu117", "cu118", "cu121", "rocm5.2", "rocm5.3"],
+    "windows": ["cpu", "cu117", "cu118", "cu121"],
     "macos": ["cpu"],
 }
 

packaging/pkg_helpers.bash

@@ -46,13 +46,21 @@ setup_cuda() {
 
   # Now work out the CUDA settings
   case "$CU_VERSION" in
+    cu121)
+      if [[ "$OSTYPE" == "msys" ]]; then
+        export CUDA_HOME="C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v12.1"
+      else
+        export CUDA_HOME=/usr/local/cuda-12.1/
+      fi
+      export TORCH_CUDA_ARCH_LIST="5.0+PTX;6.0;7.0;7.5;8.0;8.6;9.0"
+      ;;
     cu118)
       if [[ "$OSTYPE" == "msys" ]]; then
         export CUDA_HOME="C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.8"
       else
         export CUDA_HOME=/usr/local/cuda-11.8/
       fi
-      export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
+      export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6;9.0"
       ;;
     cu117)
       if [[ "$OSTYPE" == "msys" ]]; then
@@ -246,6 +254,9 @@ setup_conda_cudatoolkit_constraint() {
     export CONDA_BUILD_VARIANT="cpu"
   else
     case "$CU_VERSION" in
+      cu121)
+        export CONDA_CUDATOOLKIT_CONSTRAINT="- pytorch-cuda=12.1 # [not osx]"
+        ;;
       cu118)
         export CONDA_CUDATOOLKIT_CONSTRAINT="- pytorch-cuda=11.8 # [not osx]"
         ;;

packaging/windows/internal/cuda_install.bat

@@ -26,6 +26,7 @@ set CUDNN_LIB_FOLDER="lib\x64"
 if %CUDA_VER% EQU 116 goto cuda116
 if %CUDA_VER% EQU 117 goto cuda117
 if %CUDA_VER% EQU 118 goto cuda118
+if %CUDA_VER% EQU 121 goto cuda121
 
 echo CUDA %CUDA_VERSION_STR% is not supported
 exit /b 1
@@ -108,6 +109,32 @@ if not exist "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" (
 
 goto cuda_common
 
+:cuda121
+
+set CUDA_INSTALL_EXE=cuda_12.1.0_531.14_windows.exe
+if not exist "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%" (
+    curl -k -L "https://ossci-windows.s3.amazonaws.com/%CUDA_INSTALL_EXE%" --output "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%"
+    if errorlevel 1 exit /b 1
+    set "CUDA_SETUP_FILE=%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%"
+    set "ARGS=cuda_profiler_api_12.1 thrust_12.1 nvcc_12.1 cuobjdump_12.1 nvprune_12.1 nvprof_12.1 cupti_12.1 cublas_12.1 cublas_dev_12.1 cudart_12.1 cufft_12.1 cufft_dev_12.1 curand_12.1 curand_dev_12.1 cusolver_12.1 cusolver_dev_12.1 cusparse_12.1 cusparse_dev_12.1 npp_12.1 npp_dev_12.1 nvrtc_12.1 nvrtc_dev_12.1 nvml_dev_12.1"
+)
+
+set CUDNN_FOLDER=cudnn-windows-x86_64-8.8.1.3_cuda12-archive
+set CUDNN_LIB_FOLDER="lib"
+set "CUDNN_INSTALL_ZIP=%CUDNN_FOLDER%.zip"
+if not exist "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" (
+    curl -k -L "http://s3.amazonaws.com/ossci-windows/%CUDNN_INSTALL_ZIP%" --output "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%"
+    if errorlevel 1 exit /b 1
+    set "CUDNN_SETUP_FILE=%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%
+
+    rem Make sure windows path contains zlib dll
+    curl -k -L "http://s3.amazonaws.com/ossci-windows/zlib123dllx64.zip" --output "%SRC_DIR%\temp_build\zlib123dllx64.zip"
+    7z x "%SRC_DIR%\temp_build\zlib123dllx64.zip" -o"%SRC_DIR%\temp_build\zlib"
+    xcopy /Y "%SRC_DIR%\temp_build\zlib\dll_x64\*.dll" "C:\Windows\System32"
+)
+
+goto cuda_common
+
 :cuda_common
 
 if not exist "%SRC_DIR%\temp_build\NvToolsExt.7z" (