[GraphBolt][CUDA] Better `TORCH_CUDA_ARCH_LIST` initialization. (#7293)

graphbolt/build.sh

@@ -16,7 +16,19 @@ fi
 # TORCH_CUDA_ARCH_LIST and we need to at least compile for Volta. Until
 # https://github.com/NVIDIA/cccl/issues/1083 is resolved, we need to compile the
 # cuda/extension folder with Volta+ CUDA architectures.
-CMAKE_FLAGS="-DCUDA_TOOLKIT_ROOT_DIR=$CUDA_TOOLKIT_ROOT_DIR -DUSE_CUDA=$USE_CUDA -DGPU_CACHE_BUILD_DIR=$BINDIR -DTORCH_CUDA_ARCH_LIST=Volta"
+TORCH_CUDA_ARCH_LIST="Volta"
+if ! [[ -z "${CUDAARCHS}" ]]; then
+  # The architecture list is passed as an environment variable, we set
+  # TORCH_CUDA_ARCH_LIST to the latest architecture.
+  CUDAARCHSARR=(${CUDAARCHS//;/ })
+  LAST_ARCHITECTURE=${CUDAARCHSARR[-1]}
+  # TORCH_CUDA_ARCH_LIST has to be at least 70 to override Volta default.
+  if (( $LAST_ARCHITECTURE >= 70 )); then
+    # Convert "75" to "7.5".
+    TORCH_CUDA_ARCH_LIST=${LAST_ARCHITECTURE:0:-1}'.'${LAST_ARCHITECTURE: -1}
+  fi
+fi
+CMAKE_FLAGS="-DCUDA_TOOLKIT_ROOT_DIR=$CUDA_TOOLKIT_ROOT_DIR -DUSE_CUDA=$USE_CUDA -DGPU_CACHE_BUILD_DIR=$BINDIR -DTORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST"
 echo $CMAKE_FLAGS
 
 if [ $# -eq 0 ]; then