From f4a87bd49fd592d21a31500d49ffce702ac74206 Mon Sep 17 00:00:00 2001
From: Dmytro Ivchenko <divchenko@fireworks.ai>
Date: Thu, 26 Jan 2023 20:40:11 +0000
Subject: [PATCH] Add ability to override cuda archs for conda

---
 conda/pytorch-nightly/build.sh  | 12 ++++++++----
 conda/pytorch-nightly/meta.yaml |  1 +
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/conda/pytorch-nightly/build.sh b/conda/pytorch-nightly/build.sh
index 0de62c76e..ad1871ac4 100755
--- a/conda/pytorch-nightly/build.sh
+++ b/conda/pytorch-nightly/build.sh
@@ -52,31 +52,35 @@ if [[ -z "$USE_CUDA" || "$USE_CUDA" == 1 ]]; then
 fi
 if [[ -n "$build_with_cuda" ]]; then
     export TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
-    export TORCH_CUDA_ARCH_LIST="3.7+PTX;5.0"
+    TORCH_CUDA_ARCH_LIST="3.7+PTX;5.0"
     export USE_STATIC_CUDNN=1 # links cudnn statically (driven by tools/setup_helpers/cudnn.py)
 
     if [[ $CUDA_VERSION == 11.6* ]]; then
-        export TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;6.0;6.1;7.0;7.5;8.0;8.6"
+        TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;6.0;6.1;7.0;7.5;8.0;8.6"
         #for cuda 11.5 we use cudnn 8.3.2.44 https://docs.nvidia.com/deeplearning/cudnn/release-notes/rel_8.html
         #which does not have single static libcudnn_static.a deliverable to link with
         export USE_STATIC_CUDNN=0
         #for cuda 11.5 include all dynamic loading libraries
         DEPS_LIST=(/usr/local/cuda/lib64/libcudnn*.so.8 /usr/local/cuda-11.6/extras/CUPTI/lib64/libcupti.so.11.6)
     elif [[ $CUDA_VERSION == 11.7* ]]; then
-        export TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;6.0;6.1;7.0;7.5;8.0;8.6"
+        TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;6.0;6.1;7.0;7.5;8.0;8.6"
         #for cuda 11.7 we use cudnn 8.5
         #which does not have single static libcudnn_static.a deliverable to link with
         export USE_STATIC_CUDNN=0
         #for cuda 11.7 include all dynamic loading libraries
         DEPS_LIST=(/usr/local/cuda/lib64/libcudnn*.so.8 /usr/local/cuda-11.7/extras/CUPTI/lib64/libcupti.so.11.7)
     elif [[ $CUDA_VERSION == 11.8* ]]; then
-	export TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;6.0;6.1;7.0;7.5;8.0;8.6;9.0"
+	TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;6.0;6.1;7.0;7.5;8.0;8.6;9.0"
 	#for cuda 11.8 we use cudnn 8.7
 	#which does not have single static libcudnn_static.a deliverable to link with
 	export USE_STATIC_CUDNN=0
 	#for cuda 11.8 include all dynamic loading libraries
 	DEPS_LIST=(/usr/local/cuda/lib64/libcudnn*.so.8 /usr/local/cuda-11.8/extras/CUPTI/lib64/libcupti.so.11.8)
     fi
+    if [[ -n "$OVERRIDE_TORCH_CUDA_ARCH_LIST" ]]; then
+        TORCH_CUDA_ARCH_LIST="$OVERRIDE_TORCH_CUDA_ARCH_LIST"
+    fi
+    export TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST"
     export NCCL_ROOT_DIR=/usr/local/cuda
     export USE_STATIC_NCCL=1  # links nccl statically (driven by tools/setup_helpers/nccl.py, some of the NCCL cmake files such as FindNCCL.cmake and gloo/FindNCCL.cmake)
 
diff --git a/conda/pytorch-nightly/meta.yaml b/conda/pytorch-nightly/meta.yaml
index 6e0ba79ab..8acf9f0ad 100644
--- a/conda/pytorch-nightly/meta.yaml
+++ b/conda/pytorch-nightly/meta.yaml
@@ -84,6 +84,7 @@ build:
     - USE_PYTORCH_METAL_EXPORT # [osx]
     - USE_COREML_DELEGATE # [osx]
     - _GLIBCXX_USE_CXX11_ABI # [unix]
+    - OVERRIDE_TORCH_CUDA_ARCH_LIST
 
 test:
  imports: