Enable CPU-only builds

.github/workflows/ci.yml

@@ -13,17 +13,24 @@ jobs:
       matrix:
         include:
           # Oldest supported versions
-          - cuda: 10.2.89
+          - enable_cuda: true
+            cuda: 10.2.89
             gcc: 8.5.0
             nvcc: 10.2
             python: 3.8
             pytorch: 1.8.0
           # Latest supported versions
-          - cuda: 11.2.2
+          - enable_cuda: true
+            cuda: 11.2.2
             gcc: 10.3.0
             nvcc: 11.2
             python: 3.9
             pytorch: 1.10.0
+          # Without CUDA
+          - enable_cuda: false
+            gcc: 10.3.0
+            python: 3.9
+            pytorch: 1.10.0
 
     steps:
     - name: Check out
@@ -33,6 +40,7 @@ jobs:
       uses: Jimver/[email protected]
       with:
         cuda: ${{ matrix.cuda }}
+      if: ${{ matrix.enable_cuda }}
 
     - name: Install Miniconda
       uses: conda-incubator/setup-miniconda@v2
@@ -41,23 +49,39 @@ jobs:
         auto-activate-base: true
         miniforge-variant: Miniforge3
 
-    - name: Install dependencies
+    - name: Prepare dependencies (with CUDA)
       shell: bash -l {0}
+      if: ${{ matrix.enable_cuda }}
       run: |
         sed -i -e "/cudatoolkit/c\  - cudatoolkit ${{ matrix.cuda }}" \
                -e "/gxx_linux-64/c\  - gxx_linux-64 ${{ matrix.gcc }}" \
                -e "/nvcc_linux-64/c\  - nvcc_linux-64 ${{ matrix.nvcc }}" \
                -e "/python/c\  - python ${{ matrix.python }}" \
                -e "/pytorch-gpu/c\  - pytorch-gpu ${{ matrix.pytorch }}" \
                environment.yml
-        conda env create -n nnpops -f environment.yml
+
+    - name: Prepare dependencies (without CUDA)
+      shell: bash -l {0}
+      if: ${{ !matrix.enable_cuda }}
+      run: |
+        sed -i -e "/cudatoolkit/c\  # - cudatoolkit" \
+               -e "/gxx_linux-64/c\  - gxx_linux-64 ${{ matrix.gcc }}" \
+               -e "/nvcc_linux-64/c\  # - nvcc_linux-64" \
+               -e "/python/c\  - python ${{ matrix.python }}" \
+               -e "/pytorch-gpu/c\  - pytorch-cpu ${{ matrix.pytorch }}" \
+               environment.yml
+
+    - name: Install dependencies
+      shell: bash -l {0}
+      run: conda env create -n nnpops -f environment.yml
 
     - name: Configure, compile, and install
       shell: bash -l {0}
       run: |
         conda activate nnpops
         mkdir build && cd build
         cmake .. \
+              -DENABLE_CUDA=${{ matrix.enable_cuda }} \
               -DTorch_DIR=$CONDA_PREFIX/lib/python${{ matrix.python }}/site-packages/torch/share/cmake/Torch \
               -DCMAKE_INSTALL_PREFIX=$CONDA_PREFIX
         make install

CMakeLists.txt

@@ -1,39 +1,57 @@
 cmake_minimum_required(VERSION 3.20 FATAL_ERROR)
 
+# Configure
 set(NAME NNPOps)
-set(LIBRARY ${NAME}PyTorch)
-project(${NAME} LANGUAGES CXX CUDA)
+project(${NAME} LANGUAGES CXX)
+option(ENABLE_CUDA "Enable biuilding CUDA components" ON)
+if(ENABLE_CUDA)
+    enable_language(CUDA)
+endif(ENABLE_CUDA)
 
+# Find dependencies
 find_package(Python REQUIRED)
 find_package(PythonLibs REQUIRED)
 find_package(Torch REQUIRED)
 enable_testing()
 
-add_library(${LIBRARY} SHARED src/ani/CpuANISymmetryFunctions.cpp
-                              src/ani/CudaANISymmetryFunctions.cu
-                              src/pytorch/BatchedNN.cpp
-                              src/pytorch/SymmetryFunctions.cpp
-                              src/schnet/CpuCFConv.cpp
-                              src/schnet/CudaCFConv.cu)
+# Source files of the library
+set(SRC_FILES src/ani/CpuANISymmetryFunctions.cpp
+              src/ani/CudaANISymmetryFunctions.cu
+              src/pytorch/BatchedNN.cpp
+              src/pytorch/SymmetryFunctions.cpp
+              src/schnet/CpuCFConv.cpp
+              src/schnet/CudaCFConv.cu)
+
+# Build the library
+set(LIBRARY ${NAME}PyTorch)
+add_library(${LIBRARY} SHARED ${SRC_FILES})
 target_include_directories(${LIBRARY} PRIVATE src/ani src/schnet)
 target_link_libraries(${LIBRARY} ${TORCH_LIBRARIES} ${PYTHON_LIBRARIES})
+if(ENABLE_CUDA)
+    target_compile_definitions(${LIBRARY} PRIVATE ENABLE_CUDA)
+endif(ENABLE_CUDA)
 
+# Test of the library
 set(TEST_PATHS src/ani/TestCpuANISymmetryFunctions.cpp
-               src/ani/TestCudaANISymmetryFunctions.cpp
-               src/schnet/TestCpuCFConv.cpp
-               src/schnet/TestCudaCFConv.cu)
+               src/schnet/TestCpuCFConv.cpp)
+if(ENABLE_CUDA)
+    list(APPEND TEST_PATHS src/ani/TestCudaANISymmetryFunctions.cu
+                           src/schnet/TestCudaCFConv.cu)
+endif(ENABLE_CUDA)
 foreach(TEST_PATH ${TEST_PATHS})
     cmake_path(GET TEST_PATH STEM TEST_NAME)
     add_executable(${TEST_NAME} ${TEST_PATH})
     target_link_libraries(${TEST_NAME} ${LIBRARY})
     add_test(${TEST_NAME} ${TEST_NAME})
 endforeach()
 
+# Tests of PyTorch wrappers
 add_test(TestBatchedNN         pytest -v ${CMAKE_SOURCE_DIR}/src/pytorch/TestBatchedNN.py)
 add_test(TestSpeciesConverter  pytest -v ${CMAKE_SOURCE_DIR}/src/pytorch/TestSpeciesConverter.py)
 add_test(TestEnergyShifter     pytest -v ${CMAKE_SOURCE_DIR}/src/pytorch/TestEnergyShifter.py)
 add_test(TestSymmetryFunctions pytest -v ${CMAKE_SOURCE_DIR}/src/pytorch/TestSymmetryFunctions.py)
 
+# Installation
 install(TARGETS ${LIBRARY} DESTINATION ${Python_SITEARCH}/${NAME})
 install(FILES src/pytorch/__init__.py
               src/pytorch/BatchedNN.py

src/pytorch/SymmetryFunctions.cpp

@@ -23,14 +23,16 @@
 
 #include <stdexcept>
 #include <torch/script.h>
-#include <c10/cuda/CUDAStream.h>
 #include "CpuANISymmetryFunctions.h"
+#ifdef ENABLE_CUDA
+#include <c10/cuda/CUDAStream.h>
 #include "CudaANISymmetryFunctions.h"
 
 #define CHECK_CUDA_RESULT(result) \
     if (result != cudaSuccess) { \
         throw std::runtime_error(std::string("Encountered error ")+cudaGetErrorName(result)+" at "+__FILE__+":"+std::to_string(__LINE__));\
     }
+#endif
 
 namespace NNPOps {
 namespace ANISymmetryFunctions {
@@ -87,17 +89,23 @@ class Holder : public torch::CustomClassHolder {
         const torch::Device& device = tensorOptions.device();
         if (device.is_cpu())
             symFunc = std::make_shared<CpuANISymmetryFunctions>(numAtoms, numSpecies, Rcr, Rca, false, atomSpecies, radialFunctions, angularFunctions, true);
-        if (device.is_cuda()) {
+#ifdef ENABLE_CUDA
+        else if (device.is_cuda()) {
             // PyTorch allow to chose GPU with "torch.device", but it doesn't set as the default one.
             CHECK_CUDA_RESULT(cudaSetDevice(device.index()));
             symFunc = std::make_shared<CudaANISymmetryFunctions>(numAtoms, numSpecies, Rcr, Rca, false, atomSpecies, radialFunctions, angularFunctions, true);
         }
+#endif
+        else
+            throw std::runtime_error("Unsupported device: " + device.str());
 
         radial  = torch::empty({numAtoms, numSpecies * (int)radialFunctions.size()}, tensorOptions);
         angular = torch::empty({numAtoms, numSpecies * (numSpecies + 1) / 2 * (int)angularFunctions.size()}, tensorOptions);
         positionsGrad = torch::empty({numAtoms, 3}, tensorOptions);
 
+#ifdef ENABLE_CUDA
         cudaSymFunc = dynamic_cast<CudaANISymmetryFunctions*>(symFunc.get());
+#endif
     };
 
     tensor_list forward(const Tensor& positions_, const optional<Tensor>& periodicBoxVectors_) {
@@ -111,10 +119,12 @@ class Holder : public torch::CustomClassHolder {
             float* periodicBoxVectorsPtr = periodicBoxVectors.data_ptr<float>();
         }
 
+#ifdef ENABLE_CUDA
         if (cudaSymFunc) {
             const torch::cuda::CUDAStream stream = torch::cuda::getCurrentCUDAStream(tensorOptions.device().index());
             cudaSymFunc->setStream(stream.stream());
         }
+#endif
 
         symFunc->computeSymmetryFunctions(positions.data_ptr<float>(), periodicBoxVectorsPtr, radial.data_ptr<float>(), angular.data_ptr<float>());
 
@@ -126,10 +136,12 @@ class Holder : public torch::CustomClassHolder {
         const Tensor radialGrad = grads[0].clone();
         const Tensor angularGrad = grads[1].clone();
 
+#ifdef ENABLE_CUDA
         if (cudaSymFunc) {
             const torch::cuda::CUDAStream stream = torch::cuda::getCurrentCUDAStream(tensorOptions.device().index());
             cudaSymFunc->setStream(stream.stream());
         }
+#endif
 
         symFunc->backprop(radialGrad.data_ptr<float>(), angularGrad.data_ptr<float>(), positionsGrad.data_ptr<float>());
 
@@ -146,7 +158,9 @@ class Holder : public torch::CustomClassHolder {
     Tensor radial;
     Tensor angular;
     Tensor positionsGrad;
+#ifdef ENABLE_CUDA
     CudaANISymmetryFunctions* cudaSymFunc;
+#endif
 };
 
 class AutogradFunctions : public torch::autograd::Function<AutogradFunctions> {