Use pthread_atfork to avoid master thread affinity be derived by child.

apache · Nov 19, 2019 · 4539e89 · 4539e89
1 parent ace4e7a
commit 4539e89
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 9 deletions.
diff --git a/src/runtime/thread_pool.cc b/src/runtime/thread_pool.cc
@@ -284,8 +284,8 @@ class ThreadPool {
       queues_.emplace_back(std::unique_ptr<SpscTaskQueue>(new SpscTaskQueue()));
     }
     const char* exclude_worker0 = getenv("TVM_EXCLUDE_WORKER0");
-    if (exclude_worker0 && atoi(exclude_worker0) == 1) {
-      exclude_worker0_ = true;
+    if (exclude_worker0 && atoi(exclude_worker0) == 0) {
+      exclude_worker0_ = false;
     }
     threads_ = std::unique_ptr<tvm::runtime::threading::ThreadGroup>(
         new tvm::runtime::threading::ThreadGroup(
@@ -374,7 +374,11 @@ class ThreadPool {
   // number of workers used (can be restricted with affinity pref)
   int num_workers_used_;
   // if excluding worker 0 and using master to run task 0
+#ifndef _LIBCPP_SGX_CONFIG
+  bool exclude_worker0_{true};
+#else
   bool exclude_worker0_{false};
+#endif
   std::vector<std::unique_ptr<SpscTaskQueue> > queues_;
   std::unique_ptr<tvm::runtime::threading::ThreadGroup> threads_;
 };

diff --git a/src/runtime/threading_backend.cc b/src/runtime/threading_backend.cc
@@ -134,21 +134,28 @@ class ThreadGroup::Impl {
 #endif
     }
     if (exclude_worker0) {  // bind the master thread to core 0
-      cpu_set_t cpuset;
-      CPU_ZERO(&cpuset);
       if (reverse) {
-        CPU_SET(sorted_order_[sorted_order_.size() - 1], &cpuset);
+        core0_id_ = sorted_order_[sorted_order_.size() - 1];
       } else {
-        CPU_SET(sorted_order_[0], &cpuset);
+        core0_id_ = sorted_order_[0];
       }
+     pthread_atfork(nullptr, nullptr, ThreadGroup::Impl::BindMasterThreadToCore0);
+    }
+#endif
+  }
+
+  static void BindMasterThreadToCore0() {
+#if defined(__linux__) || defined(__ANDROID__)
+    cpu_set_t cpuset;
+    CPU_ZERO(&cpuset);
+    CPU_SET(ThreadGroup::Impl::core0_id_, &cpuset);
 #if defined(__ANDROID__)
-      sched_setaffinity(pthread_self(),
+    sched_setaffinity(pthread_self(),
         sizeof(cpu_set_t), &cpuset);
 #else
-      pthread_setaffinity_np(pthread_self(),
+    pthread_setaffinity_np(pthread_self(),
         sizeof(cpu_set_t), &cpuset);
 #endif
-    }
 #endif
   }
 
@@ -198,8 +205,11 @@ class ThreadGroup::Impl {
   std::vector<unsigned int> sorted_order_;
   int big_count_ = 0;
   int little_count_ = 0;
+  static int core0_id_;
 };
 
+int ThreadGroup::Impl::core0_id_ = 0;
+
 ThreadGroup::ThreadGroup(int num_workers,
                          std::function<void(int)> worker_callback,
                          bool exclude_worker0)