Init threadpool with user defined num_threads before default (#136793)

janeyx99 · pytorchmergebot · commit adbcaee950af · 2024-09-27T22:22:37.000Z
Fixes #134714 (or attempts to, idk how to test yet) For posterity, how one can test: 1. make sure you have USE_PTHREADPOOL=1 or pull a packaged binary 2. run gdb --args python, with `r` to enter, `Ctrl-C` to pause, and `c` to get back into Python 3. import torch 4. torch.set_num_threads(1), make sure this does not trigger any additional threads getting created. Pull Request resolved: #136793 Approved by: https://github.com/albanD
diff --git a/aten/src/ATen/ParallelNative.cpp b/aten/src/ATen/ParallelNative.cpp
@@ -209,8 +209,8 @@ void init_num_threads() {
 }
 
 void set_num_threads(int nthreads) {
-#ifndef C10_MOBILE
   TORCH_CHECK(nthreads > 0, "Expected positive number of threads");
+#ifndef C10_MOBILE
   int no_value = NOT_SET;
   if (!num_intraop_threads.compare_exchange_strong(no_value, nthreads)) {
     // num_intraop_threads either stores a positive integer or CONSUMED,
@@ -229,9 +229,8 @@ void set_num_threads(int nthreads) {
     }
   }
 #else
-  caffe2::PThreadPool* const pool = caffe2::pthreadpool();
+  caffe2::PThreadPool* const pool = caffe2::pthreadpool(nthreads);
   TORCH_INTERNAL_ASSERT(pool, "Invalid thread pool!");
-  pool->set_thread_count(nthreads);
 #endif // C10_MOBILE
 }
 
diff --git a/aten/src/ATen/ParallelOpenMP.cpp b/aten/src/ATen/ParallelOpenMP.cpp
@@ -61,9 +61,8 @@ void set_num_threads(int nthreads) {
 #endif
 #ifdef USE_PTHREADPOOL
   // because PyTorch uses caffe2::pthreadpool() in QNNPACK
-  caffe2::PThreadPool* const pool = caffe2::pthreadpool();
+  caffe2::PThreadPool* const pool = caffe2::pthreadpool(nthreads);
   TORCH_INTERNAL_ASSERT(pool, "Invalid thread pool!");
-  pool->set_thread_count(nthreads);
 #endif
 #if AT_MKLDNN_ENABLED()
   at::native::mkldnn::clear_computation_cache();
diff --git a/caffe2/utils/threadpool/pthreadpool-cpp.cc b/caffe2/utils/threadpool/pthreadpool-cpp.cc
@@ -82,12 +82,9 @@ void PThreadPool::run(
       0u);
 }
 
-// Forward declaration
-size_t getDefaultNumThreads();
-
-PThreadPool* pthreadpool() {
+PThreadPool* pthreadpool(size_t thread_count) {
   static auto threadpool =
-    std::make_unique<PThreadPool>(getDefaultNumThreads());
+    std::make_unique<PThreadPool>(thread_count);
 #if !(defined(WIN32))
   static std::once_flag flag;
   std::call_once(flag, []() {
@@ -105,6 +102,13 @@ PThreadPool* pthreadpool() {
   return threadpool.get();
 }
 
+// Forward declaration
+size_t getDefaultNumThreads();
+
+PThreadPool* pthreadpool() {
+  return pthreadpool(getDefaultNumThreads());
+}
+
 pthreadpool_t pthreadpool_() {
   if (caffe2::_NoPThreadPoolGuard::is_enabled()) {
     return nullptr;
diff --git a/caffe2/utils/threadpool/pthreadpool-cpp.h b/caffe2/utils/threadpool/pthreadpool-cpp.h
@@ -42,6 +42,7 @@ class PThreadPool final {
 
 // Return a singleton instance of PThreadPool for ATen/TH multithreading.
 PThreadPool* pthreadpool();
+PThreadPool* pthreadpool(size_t thread_count);
 
 // Exposes the underlying implementation of PThreadPool.
 // Only for use in external libraries so as to unify threading across

Original file line number	Diff line number	Diff line change
`@@ -209,8 +209,8 @@ void init_num_threads() {`
`209`	`209`	`}`
`210`	`210`
`211`	`211`	`void set_num_threads(int nthreads) {`
`212`		`-#ifndef C10_MOBILE`
`213`	`212`	`TORCH_CHECK(nthreads > 0, "Expected positive number of threads");`
	`213`	`+#ifndef C10_MOBILE`
`214`	`214`	`int no_value = NOT_SET;`
`215`	`215`	`if (!num_intraop_threads.compare_exchange_strong(no_value, nthreads)) {`
`216`	`216`	`// num_intraop_threads either stores a positive integer or CONSUMED,`
`@@ -229,9 +229,8 @@ void set_num_threads(int nthreads) {`
`229`	`229`	`}`
`230`	`230`	`}`
`231`	`231`	`#else`
`232`		`- caffe2::PThreadPool* const pool = caffe2::pthreadpool();`
	`232`	`+ caffe2::PThreadPool* const pool = caffe2::pthreadpool(nthreads);`
`233`	`233`	`TORCH_INTERNAL_ASSERT(pool, "Invalid thread pool!");`
`234`		`- pool->set_thread_count(nthreads);`
`235`	`234`	`#endif // C10_MOBILE`
`236`	`235`	`}`
`237`	`236`