[pytorch] use a mutex in initialize_torch_libraries (#151938)

Summary: The TORCH_LIBRARY_THREAD_UNSAFE_LAZY_INIT feature is thread unsafe for calling the initializers, but we want to allow the deferred initializer call to be safe from multiple threads. Add a mutex to ensure we have thread safe construction of the libraries post launch.

Differential Revision: D73457714

Pull Request resolved: https://github.com/pytorch/pytorch/pull/151938
Approved by: https://github.com/swolchok, https://github.com/zou3519

aten/src/ATen/core/library.cpp

@@ -61,9 +61,17 @@ void Library::reset() {
 
 #if defined(TORCH_LIBRARY_THREAD_UNSAFE_LAZY_INIT) && defined(C10_MOBILE)
 namespace detail {
+  // Insertion of library initializers into torch_library_initializers is not
+  // thread-safe as we expect this to be handled by the applications dynamic
+  // library loader, which would guarantee that only one thread is inserting
+  // libraries into the vector. We do require thread safety when calling
+  // initialize_torch_libraries however, as this can be called from any
+  // thread, and potentially race and corrupt the library initializer vector.
+  std::mutex torch_library_initializer_mutex;
   std::vector<TorchLibraryInit*> torch_library_initializers;
 } // namespace detail
 void initialize_torch_libraries() {
+  const std::lock_guard<std::mutex> lock(detail::torch_library_initializer_mutex);
   for (auto* initializer : detail::torch_library_initializers) {
     initializer->initialize();
   }