[fbsync] Use TORCH_CUDA_ARCH_LIST to specify which CUDA architectures to build for (#3399)

Summary:
* trying stuff

* Put flags back?

* remove second one just to see

* It worked but it's because it's not built by CI. Trying another one

* Try using TORCH_CUDA_ARCH_LIST instead

* oops

* Add new env variable to build/script_env

* set TORCH_CUDA_ARCH_LIST for the rest of the CUDA versions

* don't pass NVCC_FLAGS venv to conda, let's see if it works

Reviewed By: fmassa

Differential Revision: D27433927

fbshipit-source-id: e207f7fe6a1bab322e41d6dcabe7eb2b8e70b246

Co-authored-by: Francisco Massa <[email protected]>

Author: 2 people

packaging/pkg_helpers.bash

@@ -56,9 +56,7 @@ setup_cuda() {
         export CUDA_HOME=/usr/local/cuda-11.2/
       fi
       export FORCE_CUDA=1
-      # Hard-coding gencode flags is temporary situation until
-      # https://github.com/pytorch/pytorch/pull/23408 lands
-      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_50,code=compute_50"
+      export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
       ;;
     cu111)
       if [[ "$OSTYPE" == "msys" ]]; then
@@ -67,9 +65,7 @@ setup_cuda() {
         export CUDA_HOME=/usr/local/cuda-11.1/
       fi
       export FORCE_CUDA=1
-      # Hard-coding gencode flags is temporary situation until
-      # https://github.com/pytorch/pytorch/pull/23408 lands
-      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_50,code=compute_50"
+      export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
       ;;
     cu110)
       if [[ "$OSTYPE" == "msys" ]]; then
@@ -78,9 +74,7 @@ setup_cuda() {
         export CUDA_HOME=/usr/local/cuda-11.0/
       fi
       export FORCE_CUDA=1
-      # Hard-coding gencode flags is temporary situation until
-      # https://github.com/pytorch/pytorch/pull/23408 lands
-      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_50,code=compute_50"
+      export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0"
       ;;
     cu102)
       if [[ "$OSTYPE" == "msys" ]]; then
@@ -89,9 +83,7 @@ setup_cuda() {
         export CUDA_HOME=/usr/local/cuda-10.2/
       fi
       export FORCE_CUDA=1
-      # Hard-coding gencode flags is temporary situation until
-      # https://github.com/pytorch/pytorch/pull/23408 lands
-      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_50,code=compute_50"
+      export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5"
       ;;
     cu101)
       if [[ "$OSTYPE" == "msys" ]]; then
@@ -100,9 +92,7 @@ setup_cuda() {
         export CUDA_HOME=/usr/local/cuda-10.1/
       fi
       export FORCE_CUDA=1
-      # Hard-coding gencode flags is temporary situation until
-      # https://github.com/pytorch/pytorch/pull/23408 lands
-      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_50,code=compute_50"
+      export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5"
       ;;
     cu100)
       if [[ "$OSTYPE" == "msys" ]]; then
@@ -111,9 +101,7 @@ setup_cuda() {
         export CUDA_HOME=/usr/local/cuda-10.0/
       fi
       export FORCE_CUDA=1
-      # Hard-coding gencode flags is temporary situation until
-      # https://github.com/pytorch/pytorch/pull/23408 lands
-      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_50,code=compute_50"
+      export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5"
       ;;
     cu92)
       if [[ "$OSTYPE" == "msys" ]]; then
@@ -122,7 +110,7 @@ setup_cuda() {
         export CUDA_HOME=/usr/local/cuda-9.2/
       fi
       export FORCE_CUDA=1
-      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_50,code=compute_50"
+      export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0"
       ;;
     cpu)
       ;;

packaging/torchvision/meta.yaml

@@ -37,8 +37,8 @@ build:
   script_env:
     - CUDA_HOME
     - FORCE_CUDA
-    - NVCC_FLAGS
     - BUILD_VERSION
+    - TORCH_CUDA_ARCH_LIST
   features:
     {{ environ.get('CONDA_CPUONLY_FEATURE') }}