CI cuda test speedup (#573)

* Update install_cuda.sh

* Update install_cuda.sh

* Update install_cuda.sh

* update

* update

* update

* update

* update

* update

* update

* update

.github/workflows/cuda.yml

@@ -32,10 +32,12 @@ jobs:
     - name: Fetch submodules
       run: git submodule update --init
     - name: Install cuda
-      run: ci/install_cuda.sh
+      run: |
+        source ci/install_cuda.sh
+        echo "CUDA_HOME=${CUDA_HOME}" >> $GITHUB_ENV
     - name: Install dependencies
       run: ci/install_dependencies.sh
     - name: Install TorchANI
-      run: python setup.py install --cuaev
+      run: echo $CUDA_HOME; python setup.py install --cuaev --only-sm80
     - name: Run cuaev test
       run: pytest tests/test_cuaev.py -v

ci/install_cuda.sh

@@ -7,4 +7,9 @@ sudo mv cuda-ubuntu1804.pin /etc/apt/preferences.d/cuda-repository-pin-600
 sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
 sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/ /"
 sudo apt-get update
-sudo apt-get -y install cuda
+
+# from https://github.com/ptheywood/cuda-cmake-github-actions/blob/master/scripts/actions/install_cuda_ubuntu.sh
+sudo apt-get -y install cuda-command-line-tools-11-0 cuda-libraries-dev-11-0
+export CUDA_HOME=/usr/local/cuda-11.0
+export PATH="$CUDA_HOME/bin:$PATH"
+nvcc -V

setup.py

@@ -13,6 +13,11 @@ FAST_BUILD_CUAEV = '--cuaev' in sys.argv
 if FAST_BUILD_CUAEV:
     sys.argv.remove('--cuaev')
 
+# Use along with --cuaev for CI test to reduce compilation time on Non-GPUs system
+ONLY_BUILD_SM80 = '--only-sm80' in sys.argv
+if ONLY_BUILD_SM80:
+    sys.argv.remove('--only-sm80')
+
 if not BUILD_CUAEV_ALL_SM and not FAST_BUILD_CUAEV:
     log.warn("Will not install cuaev")  # type: ignore
 
@@ -70,7 +75,10 @@ def cuda_extension(build_all=False):
     if SMs:
         for sm in SMs:
             nvcc_args.append(f"-gencode=arch=compute_{sm},code=sm_{sm}")
-    else:
+    elif len(SMs) == 0 and ONLY_BUILD_SM80:  # --cuaev --only-sm80
+        nvcc_args.append("-gencode=arch=compute_80,code=sm_80")
+    else:  # no gpu detected
+        print('NO gpu detected, will build for all SMs')
         nvcc_args.append("-gencode=arch=compute_60,code=sm_60")
         nvcc_args.append("-gencode=arch=compute_61,code=sm_61")
         nvcc_args.append("-gencode=arch=compute_70,code=sm_70")