Merge pull request #645 from basnijholt/cuda-builds

CUDA build improvements
quantumlib · Jan 31, 2024 · 5416d95 · 5416d95
2 parents 83839cc + cad9662
commit 5416d95
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 13 deletions.
diff --git a/pybind_interface/cuda/CMakeLists.txt b/pybind_interface/cuda/CMakeLists.txt
@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 3.11)
+cmake_minimum_required(VERSION 3.18)
 project(qsim LANGUAGES CXX CUDA)
 
 if(WIN32)
@@ -18,12 +18,18 @@ INCLUDE(../GetPybind11.cmake)
 find_package(PythonLibs 3.7 REQUIRED)
 find_package(CUDA REQUIRED)
 
-include_directories(${PYTHON_INCLUDE_DIRS} ${pybind11_SOURCE_DIR}/include)
+include_directories(${PYTHON_INCLUDE_DIRS})
+if(pybind11_FOUND)
+    include_directories(${pybind11_INCLUDE_DIRS})
+else()  # means pybind11 has been fetched in GetPybind11.cmake
+    include_directories(${pybind11_SOURCE_DIR}/include)
+endif()
 
 cuda_add_library(qsim_cuda MODULE pybind_main_cuda.cpp)
 set_target_properties(qsim_cuda PROPERTIES
-       PREFIX "${PYTHON_MODULE_PREFIX}"
-       SUFFIX "${PYTHON_MODULE_EXTENSION}"
+    CUDA_ARCHITECTURES "$ENV{CUDAARCHS}"
+    PREFIX "${PYTHON_MODULE_PREFIX}"
+    SUFFIX "${PYTHON_MODULE_EXTENSION}"
 )
 set_source_files_properties(pybind_main_cuda.cpp PROPERTIES LANGUAGE CUDA)
 

diff --git a/pybind_interface/cuda/pybind_main_cuda.cpp b/pybind_interface/cuda/pybind_main_cuda.cpp
@@ -27,8 +27,10 @@ namespace qsim {
       unsigned num_sim_threads,
       unsigned num_state_threads,
       unsigned num_dblocks
-    ) : ss_params{num_state_threads, num_dblocks} {}
-
+    ) {
+      ss_params.num_threads = num_state_threads;
+      ss_params.num_dblocks = num_dblocks;
+    }
     StateSpace CreateStateSpace() const {
       return StateSpace(ss_params);
     }

diff --git a/pybind_interface/custatevec/CMakeLists.txt b/pybind_interface/custatevec/CMakeLists.txt
@@ -40,8 +40,8 @@ cuda_add_library(qsim_custatevec MODULE pybind_main_custatevec.cpp)
 target_link_libraries(qsim_custatevec -lcustatevec -lcublas)
 
 set_target_properties(qsim_custatevec PROPERTIES
-       PREFIX "${PYTHON_MODULE_PREFIX}"
-       SUFFIX "${PYTHON_MODULE_EXTENSION}"
+    PREFIX "${PYTHON_MODULE_PREFIX}"
+    SUFFIX "${PYTHON_MODULE_EXTENSION}"
 )
 set_source_files_properties(pybind_main_custatevec.cpp PROPERTIES LANGUAGE CUDA)
 

diff --git a/pybind_interface/decide/CMakeLists.txt b/pybind_interface/decide/CMakeLists.txt
@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 3.11)
+cmake_minimum_required(VERSION 3.18)
 
 if(WIN32)
     set(CMAKE_CXX_FLAGS "/O2 /openmp")
@@ -26,8 +26,9 @@ if(has_nvcc)
     find_package(Python3 3.7 REQUIRED COMPONENTS Interpreter Development)
     include_directories(${PYTHON_INCLUDE_DIRS} ${pybind11_SOURCE_DIR}/include)
     set_target_properties(qsim_decide PROPERTIES
-            PREFIX "${PYTHON_MODULE_PREFIX}"
-            SUFFIX "${PYTHON_MODULE_EXTENSION}"
+        CUDA_ARCHITECTURES "$ENV{CUDAARCHS}"
+        PREFIX "${PYTHON_MODULE_PREFIX}"
+        SUFFIX "${PYTHON_MODULE_EXTENSION}"
     )
     set_source_files_properties(decide.cpp PROPERTIES LANGUAGE CUDA)
     target_link_libraries(qsim_decide OpenMP::OpenMP_CXX)
@@ -39,8 +40,8 @@ elseif(has_hipcc)
     find_package(Python3 3.7 REQUIRED COMPONENTS Interpreter Development)
     include_directories(${PYTHON_INCLUDE_DIRS} ${pybind11_SOURCE_DIR}/include)
     set_target_properties(qsim_decide PROPERTIES
-            PREFIX "${PYTHON_MODULE_PREFIX}"
-            SUFFIX "${PYTHON_MODULE_EXTENSION}"
+        PREFIX "${PYTHON_MODULE_PREFIX}"
+        SUFFIX "${PYTHON_MODULE_EXTENSION}"
     )
     target_link_libraries(qsim_decide PUBLIC OpenMP::OpenMP_CXX)
 else()

diff --git a/setup.py b/setup.py
@@ -48,6 +48,10 @@ def build_extension(self, ext):
             "-DPYTHON_INCLUDE_DIR=" + python_include_dir,
         ]
 
+        additional_cmake_args = os.environ.get("CMAKE_ARGS", "")
+        if additional_cmake_args:
+            cmake_args += additional_cmake_args.split()
+
         cfg = "Debug" if self.debug else "Release"
         build_args = ["--config", cfg]