Support omp threadpool (#170)

This pull request introduces a new OpenMP thread pool, `OMPThreadPool`,
designed for users who wish to utilize OpenMP threading.

Author: Dian-Lun (Aaron) Lin

CMakeLists.txt

@@ -77,6 +77,10 @@ include("cmake/fmt.cmake")
 include("cmake/spdlog.cmake")
 include("cmake/toml.cmake")
 
+if(SVS_ENABLE_OMP)
+    include("cmake/openmp.cmake")
+endif()
+
 add_library(svs_x86_options_base INTERFACE)
 add_library(svs::x86_options_base ALIAS svs_x86_options_base)
 if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86)|(X86)|(amd64)|(AMD64)")

cmake/openmp.cmake

@@ -13,10 +13,8 @@
 # limitations under the License.
 
 find_package(OpenMP)
-if (OPENMP_FOUND)
-    set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
-    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
-    set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} ${OpenMP_EXE_LINKER_FLAGS}")
+if (OPENMP_CXX_FOUND)
+    target_link_libraries(${SVS_LIB}  INTERFACE OpenMP::OpenMP_CXX)
 else()
     message(FATAL_ERROR "no OpenMP support")
 endif()

cmake/options.cmake

@@ -73,6 +73,11 @@ option(SVS_INITIALIZE_LOGGER
     ON # enabled by default
 )
 
+option(SVS_ENABLE_OMP
+    "Support OpenMP threadpool."
+    OFF # disable by default
+)
+
 #####
 ##### Experimental
 #####
@@ -140,6 +145,12 @@ else()
     target_compile_options(${SVS_LIB} INTERFACE -DSVS_INITIALIZE_LOGGER=0)
 endif()
 
+if (SVS_ENABLE_OMP)
+    target_compile_options(${SVS_LIB} INTERFACE -DSVS_ENABLE_OMP=1)
+else()
+    target_compile_options(${SVS_LIB} INTERFACE -DSVS_ENABLE_OMP=0)
+endif()
+
 #####
 ##### Helper target to apply relevant compiler optimizations.
 #####

include/svs/lib/threads/threadpool.h

@@ -26,6 +26,11 @@
 #include <sstream>
 #include <vector>
 
+SVS_VALIDATE_BOOL_ENV(SVS_ENABLE_OMP);
+#if SVS_ENABLE_OMP
+#include <omp.h>
+#endif
+
 #include "svs/lib/numa.h"
 #include "svs/lib/threads/thread.h"
 #include "svs/lib/threads/thunks.h"
@@ -283,6 +288,28 @@ class SwitchNativeThreadPool {
     NativeThreadPool threadpool_;
 };
 
+SVS_VALIDATE_BOOL_ENV(SVS_ENABLE_OMP);
+#if SVS_ENABLE_OMP
+/////
+///// A thread pool that utilizes OpenMP for multithreading
+/////
+class OMPThreadPool {
+  public:
+    explicit OMPThreadPool(size_t num_threads) { omp_set_num_threads(num_threads); }
+
+    size_t size() const { return omp_get_num_threads(); }
+
+    void parallel_for(std::function<void(size_t)> f, size_t n) {
+#pragma omp parallel for
+        for (size_t i = 0; i < n; ++i) {
+            f(i);
+        }
+    }
+
+    void resize(size_t num_threads) { omp_set_num_threads(num_threads); }
+};
+#endif
+
 /////
 ///// A handy reference wrapper for situations where we only want to share a thread pool
 /////