Use cpuinfo to determine c10::ThreadPool thread number (#107010)

cyyever · pytorchmergebot · commit ad0476540dca · 2023-08-15T17:26:24.000Z
This PR prefers "logical processor number" (the cpu cores shown in htop) returned by cpuinfo for determining c10 thread number. If that fails, it uses hardware_concurrency exactly. The motivation is that in a x86 host with 64 cores and Hyper-Threading disabled, the current behavior uses 32 threads, resulting half of cores being idle. Pull Request resolved: #107010 Approved by: https://github.com/ezyang
diff --git a/c10/CMakeLists.txt b/c10/CMakeLists.txt
@@ -89,6 +89,8 @@ if(${USE_GLOG})
 endif()
 target_link_libraries(c10 PRIVATE fmt::fmt-header-only)
 
+target_link_libraries(c10 PRIVATE cpuinfo)
+
 find_package(Backtrace)
 if(Backtrace_FOUND)
   target_include_directories(c10 PRIVATE ${Backtrace_INCLUDE_DIRS})
diff --git a/c10/core/build.bzl b/c10/core/build.bzl
@@ -82,6 +82,7 @@ def define_targets(rules):
         deps = [
             ":ScalarType",
             ":impl_cow_context",
+            "//third_party/cpuinfo",
             "//c10/macros",
             "//c10/util:TypeCast",
             "//c10/util:base",
diff --git a/c10/core/thread_pool.cpp b/c10/core/thread_pool.cpp
@@ -1,8 +1,22 @@
 #include <c10/core/thread_pool.h>
 #include <c10/util/Logging.h>
+#include <cpuinfo.h>
 
 namespace c10 {
 
+size_t TaskThreadPoolBase::defaultNumThreads() {
+  cpuinfo_initialize();
+  size_t num_threads = cpuinfo_get_processors_count();
+  if (num_threads > 0) {
+    return num_threads;
+  }
+  num_threads = std::thread::hardware_concurrency();
+  if (num_threads == 0) {
+    num_threads = 1;
+  }
+  return num_threads;
+}
+
 ThreadPool::ThreadPool(
     int pool_size,
     int numa_node_id,
diff --git a/c10/core/thread_pool.h b/c10/core/thread_pool.h
@@ -18,7 +18,6 @@ C10_CLANG_DIAGNOSTIC_IGNORE("-Wshorten-64-to-32")
 
 namespace c10 {
 
-// TODO: move this to C10 and make it C10_API
 class C10_API TaskThreadPoolBase {
  public:
   virtual void run(std::function<void()> func) = 0;
@@ -37,13 +36,7 @@ class C10_API TaskThreadPoolBase {
 
   virtual ~TaskThreadPoolBase() noexcept = default;
 
-  static size_t defaultNumThreads() {
-    auto num_threads = std::thread::hardware_concurrency();
-#if defined(_M_X64) || defined(__x86_64__)
-    num_threads /= 2;
-#endif
-    return num_threads;
-  }
+  static size_t defaultNumThreads();
 };
 
 class C10_API ThreadPool : public c10::TaskThreadPoolBase {
