[pytorch] add experimental TORCH_LIBRARY_THREAD_UNSAFE_LAZY_INIT (pytorch#150537)

Summary: Add an experimental feature to defer pytorch library initialization cost to post startup. As noted this feature is not thread safe, it requires the client to maintain thread safety at library load time.

Reviewed By: zou3519

Differential Revision: D71917841

Pull Request resolved: pytorch#150537
Approved by: https://github.com/zou3519

Author: rmaz

aten/src/ATen/core/library.cpp

@@ -58,6 +58,18 @@ void Library::reset() {
 
 #define ERROR_CONTEXT "(Error occurred while processing ", toString(kind_), " block at ", file_, ":", line_, ")"
 
+#ifdef TORCH_LIBRARY_THREAD_UNSAFE_LAZY_INIT
+namespace detail {
+  std::vector<TorchLibraryInit*> torch_library_initializers;
+} // namespace detail
+void initialize_torch_libraries() {
+  for (auto* initializer : detail::torch_library_initializers) {
+    initializer->initialize();
+  }
+  detail::torch_library_initializers.clear();
+}
+#endif
+
 Library::Library(Kind kind, std::string ns, std::optional<c10::DispatchKey> k, const char* file, uint32_t line)
   : kind_(kind)
   , ns_(ns == "_" ? std::nullopt : std::make_optional(std::move(ns)))

torch/csrc/jit/mobile/import.cpp

@@ -22,6 +22,7 @@
 #include <torch/csrc/jit/serialization/import_export_functions.h>
 #include <torch/csrc/jit/serialization/import_read.h>
 #include <torch/custom_class.h>
+#include <torch/library.h>
 #include <optional>
 #include <string>
 #include <vector>
@@ -646,6 +647,9 @@ mobile::Module _load_for_mobile(
     std::optional<at::Device> device,
     ExtraFilesMap& extra_files,
     uint64_t module_load_options) {
+#ifdef TORCH_LIBRARY_THREAD_UNSAFE_LAZY_INIT
+  torch::initialize_torch_libraries();
+#endif
   auto observer = torch::observerConfig().getModuleObserver();
   if (observer) {
     extra_files.insert(std::make_pair("model_path", filename));

torch/library.h

@@ -884,8 +884,42 @@ class TORCH_API Library final {
   at::OperatorName _parseNameForLib(const char* name_str) const;
 };
 
+#ifdef TORCH_LIBRARY_THREAD_UNSAFE_LAZY_INIT
+void initialize_torch_libraries();
+#endif
+
 namespace detail {
 
+#ifdef TORCH_LIBRARY_THREAD_UNSAFE_LAZY_INIT
+extern std::vector<TorchLibraryInit*> torch_library_initializers;
+class TorchLibraryInit final {
+    private:
+      using InitFn = void(Library&);
+      Library::Kind kind;
+      InitFn* init_function;
+      const char* ns;
+      std::optional<c10::DispatchKey> key;
+      const char* file;
+      uint32_t line;
+      std::unique_ptr<Library> lib = nullptr;
+
+    public:
+      TorchLibraryInit(
+            Library::Kind kind,
+            InitFn* fn,
+            const char* ns,
+            std::optional<c10::DispatchKey> k,
+            const char* file,
+            uint32_t line) : kind(kind), init_function(fn), ns(ns), key(k), file(file), line(line) {
+              torch_library_initializers.push_back(this);
+            }
+
+      void initialize() {
+        lib = std::unique_ptr<Library>(new Library(kind, ns, key, file, line));
+        init_function(*lib);
+      }
+};
+#else
 class TorchLibraryInit final {
  private:
   using InitFn = void(Library&);
@@ -903,6 +937,7 @@ class TorchLibraryInit final {
     fn(lib_);
   }
 };
+#endif
 
 } // namespace detail