Fix omp assert issue (apache#17039)

CMakeLists.txt

@@ -240,7 +240,7 @@ if(USE_TENSORRT)
 endif()
 
 # please note that when you enable this, you might run into an linker not being able to work properly due to large code injection.
-# you can find more information here https://github.com/apache/incubator-mxnet/issues/15971 
+# you can find more information here https://github.com/apache/incubator-mxnet/issues/15971
 if(ENABLE_TESTCOVERAGE)
   message(STATUS "Compiling with test coverage support enabled. This will result in additional files being written to your source directory!")
   find_program( GCOV_PATH gcov )
@@ -436,18 +436,23 @@ endif()
 
 # ---[ OpenMP
 if(USE_OPENMP)
+
+  function(load_omp)
+    # Intel/llvm OpenMP: https://github.com/llvm-mirror/openmp
+    set(OPENMP_STANDALONE_BUILD TRUE)
+    set(LIBOMP_ENABLE_SHARED TRUE)
+    set(CMAKE_BUILD_TYPE Release)
+    add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/openmp)
+  endfunction()
+
   find_package(OpenMP REQUIRED)
   # This should build on Windows, but there's some problem and I don't have a Windows box, so
   # could a Windows user please fix?
   if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/openmp/CMakeLists.txt
      AND SYSTEM_ARCHITECTURE STREQUAL "x86_64"
      AND NOT MSVC
      AND NOT CMAKE_CROSSCOMPILING)
-
-    # Intel/llvm OpenMP: https://github.com/llvm-mirror/openmp
-    set(OPENMP_STANDALONE_BUILD TRUE)
-    set(LIBOMP_ENABLE_SHARED TRUE)
-    add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/openmp)
+    load_omp()
     list(REMOVE_ITEM mxnet_LINKER_LIBS iomp5)
     list(APPEND mxnet_LINKER_LIBS omp)
     if(UNIX)

src/engine/openmp.cc

@@ -41,6 +41,7 @@ OpenMP *OpenMP::Get() {
 OpenMP::OpenMP()
   : omp_num_threads_set_in_environment_(is_env_set("OMP_NUM_THREADS")) {
 #ifdef _OPENMP
+  initialize_process();
   const int max = dmlc::GetEnv("MXNET_OMP_MAX_THREADS", INT_MIN);
   if (max != INT_MIN) {
     omp_thread_max_ = max;
@@ -61,6 +62,12 @@ OpenMP::OpenMP()
 #endif
 }
 
+void OpenMP:: initialize_process() {
+#ifdef _OPENMP
+  omp_get_num_procs();  // will force OpenMP to be initialized
+#endif
+}
+
 void OpenMP::on_start_worker_thread(bool use_omp) {
 #ifdef _OPENMP
   if (!omp_num_threads_set_in_environment_) {

src/engine/openmp.h

@@ -74,6 +74,13 @@ class OpenMP {
    */
   void on_start_worker_thread(bool use_omp);
 
+  /*!
+   * \brief Initialize a new process to use omp (after a fork,
+   *        in case you're starting threads in the atfork() that may interfere
+   *        with the initialization. Can serialize the init with this first.
+   */
+  void initialize_process();
+
   /*!
    * \brief Get the OpenMP object's singleton pointer
    * \return Singleton OpenMP object pointer

src/initialize.cc

@@ -209,6 +209,7 @@ void LibraryInitializer::atfork_child() {
 #if MXNET_USE_OPENCV && !__APPLE__
   cv::setNumThreads(mp_cv_num_threads_);
 #endif  // MXNET_USE_OPENCV
+  engine::OpenMP::Get()->initialize_process();
   engine::OpenMP::Get()->set_thread_max(1);
   engine::OpenMP::Get()->set_enabled(false);
   Engine::Get()->Start();
@@ -218,6 +219,7 @@ void LibraryInitializer::atfork_child() {
 
 void LibraryInitializer::install_pthread_atfork_handlers() {
 #ifndef _WIN32
+  engine::OpenMP::Get()->initialize_process();  // force omp to set its atfork handler first
   pthread_atfork(pthread_atfork_prepare, pthread_atfork_parent, pthread_atfork_child);
 #endif
 }