Enable Linux wheel/conda GPU package builds (#1730)

* Remove some obsolete conditions about CUDA-10.0 from `pkg_helpers.bash`.
* Use `USE_CUDA` instead of `FORCE_CUDA` in `pkg_helpers.bash`
* Do not define `NO_CUDA_PACKAGE` in build_wheel.sh and build_cuda.sh
* Add conda-forge for Win cuda-11.1 builds
* Pass USE_CUDA / TORCH_CUDA_ARCH_LIST to conda build
* Add selected CUDA to path
* Don't define USE_CUDA for ROCM

TODO: Fix Windows CUDA builds

.circleci/config.yml

@@ -60,10 +60,6 @@ commands:
           name: Install CUDA
           command: |
               packaging/windows/internal/cuda_install.bat
-              # cudatoolkit >= 11 isn't available for windows in the nvidia channel
-              if [[ "${CU_VERSION}" =~ cu11.* ]]; then
-                export CONDA_CHANNEL_FLAGS="-c conda-forge"
-              fi
 
 binary_common: &binary_common
   parameters:
@@ -267,6 +263,10 @@ jobs:
             eval "$('/C/tools/miniconda3/Scripts/conda.exe' 'shell.bash' 'hook')"
             conda activate base
             conda install -yq conda-build "conda-package-handling!=1.5.0"
+            # cudatoolkit >= 11 isn't available for windows in the nvidia channel
+            if [[ "${CU_VERSION}" =~ cu11.* ]]; then
+              export CONDA_CHANNEL_FLAGS="-c conda-forge"
+            fi
             bash packaging/build_conda.sh
       - store_artifacts:
           path: C:/tools/miniconda3/conda-bld/win-64

.circleci/config.yml.in

@@ -60,10 +60,6 @@ commands:
           name: Install CUDA
           command: |
               packaging/windows/internal/cuda_install.bat
-              # cudatoolkit >= 11 isn't available for windows in the nvidia channel
-              if [[ "${CU_VERSION}" =~ cu11.* ]]; then
-                export CONDA_CHANNEL_FLAGS="-c conda-forge"
-              fi
 
 binary_common: &binary_common
   parameters:
@@ -267,6 +263,10 @@ jobs:
             eval "$('/C/tools/miniconda3/Scripts/conda.exe' 'shell.bash' 'hook')"
             conda activate base
             conda install -yq conda-build "conda-package-handling!=1.5.0"
+            # cudatoolkit >= 11 isn't available for windows in the nvidia channel
+            if [[ "${CU_VERSION}" =~ cu11.* ]]; then
+              export CONDA_CHANNEL_FLAGS="-c conda-forge"
+            fi
             bash packaging/build_conda.sh
       - store_artifacts:
           path: C:/tools/miniconda3/conda-bld/win-64

packaging/build_conda.sh

@@ -5,9 +5,10 @@ script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 . "$script_dir/pkg_helpers.bash"
 
 export BUILD_TYPE="conda"
-export NO_CUDA_PACKAGE=1
 setup_env 0.10.0
 export SOURCE_ROOT_DIR="$PWD"
 setup_conda_pytorch_constraint
+setup_conda_cudatoolkit_constraint
 setup_visual_studio_constraint
-conda build $CONDA_CHANNEL_FLAGS --no-anaconda-upload --python "$PYTHON_VERSION" packaging/torchaudio
+# nvidia channel included for cudatoolkit >= 11
+conda build -c defaults -c nvidia $CONDA_CHANNEL_FLAGS --no-anaconda-upload --python "$PYTHON_VERSION" packaging/torchaudio

packaging/build_wheel.sh

@@ -5,7 +5,6 @@ script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 . "$script_dir/pkg_helpers.bash"
 
 export BUILD_TYPE="wheel"
-export NO_CUDA_PACKAGE=1
 setup_env 0.10.0
 setup_wheel_python
 pip_install numpy future cmake ninja

packaging/pkg_helpers.bash

@@ -14,8 +14,8 @@
 #   PYTORCH_VERSION_SUFFIX (e.g., +cpu)
 #   WHEEL_DIR (e.g., cu100/)
 #   CUDA_HOME (e.g., /usr/local/cuda-9.2, respected by torch.utils.cpp_extension)
-#   FORCE_CUDA (respected by torchvision setup.py)
-#   NVCC_FLAGS (respected by torchvision setup.py)
+#   USE_CUDA (respected by torchaudio setup.py)
+#   NVCC_FLAGS (respected by torchaudio setup.py)
 #
 # Precondition: CUDA versions are installed in their conventional locations in
 # /usr/local/cuda-*
@@ -35,18 +35,12 @@ setup_cuda() {
   export WHEEL_DIR="cpu/"
   # Wheel builds need suffixes (but not if they're on OS X, which never has suffix)
   if [[ "$BUILD_TYPE" == "wheel" ]] && [[ "$(uname)" != Darwin ]]; then
-    # The default CUDA has no suffix
-    if [[ "$CU_VERSION" != "cu100" ]]; then
-      export PYTORCH_VERSION_SUFFIX="+$CU_VERSION"
-    fi
+    export PYTORCH_VERSION_SUFFIX="+$CU_VERSION"
     # Match the suffix scheme of pytorch, unless this package does not have
     # CUDA builds (in which case, use default)
     if [[ -z "$NO_CUDA_PACKAGE" ]]; then
       export VERSION_SUFFIX="$PYTORCH_VERSION_SUFFIX"
-      # If the suffix is non-empty, we will use a wheel subdirectory
-      if [[ -n "$PYTORCH_VERSION_SUFFIX" ]]; then
-        export WHEEL_DIR="$PYTORCH_VERSION_SUFFIX/"
-      fi
+      export WHEEL_DIR="$CU_VERSION/"
     fi
   fi
 
@@ -58,7 +52,6 @@ setup_cuda() {
       else
         export CUDA_HOME=/usr/local/cuda-11.2/
       fi
-      export FORCE_CUDA=1
       export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
       ;;
     cu111)
@@ -67,7 +60,6 @@ setup_cuda() {
       else
         export CUDA_HOME=/usr/local/cuda-11.1/
       fi
-      export FORCE_CUDA=1
       export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
       ;;
     cu110)
@@ -76,7 +68,6 @@ setup_cuda() {
       else
         export CUDA_HOME=/usr/local/cuda-11.0/
       fi
-      export FORCE_CUDA=1
       export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0"
       ;;
     cu102)
@@ -85,7 +76,6 @@ setup_cuda() {
       else
         export CUDA_HOME=/usr/local/cuda-10.2/
       fi
-      export FORCE_CUDA=1
       export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5"
       ;;
     cu101)
@@ -94,21 +84,17 @@ setup_cuda() {
       else
         export CUDA_HOME=/usr/local/cuda-10.1/
       fi
-      export FORCE_CUDA=1
       export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5"
       ;;
     cu100)
       export CUDA_HOME=/usr/local/cuda-10.0/
-      export FORCE_CUDA=1
       export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5"
       ;;
     cu92)
       export CUDA_HOME=/usr/local/cuda-9.2/
-      export FORCE_CUDA=1
       export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0"
       ;;
     rocm*)
-      export FORCE_CUDA=1
       export USE_ROCM=1
       ;;
     cpu)
@@ -118,6 +104,15 @@ setup_cuda() {
       exit 1
       ;;
   esac
+  if [[ -n "$CUDA_HOME" ]]; then
+    # Adds nvcc binary to the search path so that CMake's `find_package(CUDA)` will pick the right one
+    export PATH="$CUDA_HOME/bin:$PATH"
+    # TODO: Fix Windows CUDA builds
+    if [[ "$OSTYPE" != "msys" ]]; then
+      # Force GPU builds on CPU runner, when `torch.cuda.is_available()` returns false
+      export USE_CUDA=1
+    fi
+  fi
 }
 
 # Populate build version if necessary, and add version suffix
@@ -252,6 +247,24 @@ setup_conda_cudatoolkit_constraint() {
     export CONDA_CUDATOOLKIT_CONSTRAINT=""
   else
     case "$CU_VERSION" in
+      cu113)
+        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.3,<11.4 # [not osx]"
+        ;;
+      cu112)
+        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.2,<11.3 # [not osx]"
+        ;;
+      cu111)
+        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.1,<11.2 # [not osx]"
+        ;;
+      cu110)
+        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.0,<11.1 # [not osx]"
+        ;;
+      cu102)
+        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=10.2,<10.3 # [not osx]"
+        ;;
+      cu101)
+        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=10.1,<10.2 # [not osx]"
+        ;;
       cu100)
         export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=10.0,<10.1 # [not osx]"
         ;;

packaging/torchaudio/meta.yaml

 package:
   name: torchaudio
-  version: "{{ environ.get('BUILD_VERSION') }}"
+  version: "{{ environ.get('BUILD_VERSION', '0.0.0') }}"
 
 source:
-  path: "{{ environ.get('SOURCE_ROOT_DIR') }}"
+  path: "{{ environ.get('SOURCE_ROOT_DIR', '../..') }}"
 
 requirements:
   build:
@@ -13,22 +13,28 @@ requirements:
   host:
     - python
     - setuptools
-    - cpuonly
     - cmake
     - ninja
     - defaults::numpy >=1.11
     {{ environ.get('CONDA_PYTORCH_BUILD_CONSTRAINT', 'pytorch') }}
     {{ environ.get('CONDA_EXTRA_BUILD_CONSTRAINT', '') }}
+    {{ environ.get('CONDA_CPUONLY_FEATURE', '') }}
+    {{ environ.get('CONDA_CUDATOOLKIT_CONSTRAINT', '') }}
 
   run:
     - python
     - defaults::numpy >=1.11
     {{ environ.get('CONDA_PYTORCH_CONSTRAINT', 'pytorch') }}
+    {{ environ.get('CONDA_CUDATOOLKIT_CONSTRAINT', '') }}
 
 build:
-  string: py{{py}}
+  string: py{{py}}_{{ environ.get('CU_VERSION', 'cpu') }}
   script_env:
     - BUILD_VERSION
+    - USE_CUDA  # [not win]
+    - TORCH_CUDA_ARCH_LIST  # [not win]
+  features:
+    {{ environ.get('CONDA_CPUONLY_FEATURE', '') }}
 
 test:
   imports:
@@ -46,7 +52,7 @@ test:
     # Ideally we would test this, but conda doesn't provide librosa
     # - librosa >=0.4.3
     - scipy
-    - cpuonly
+    {{ environ.get('CONDA_CPUONLY_FEATURE', '') }}
 
 about:
   home: https://github.com/pytorch/audio