Add _unsafe_reset_threadpool to pybindings

Summary: Expose thread count setting in pybindings. Currently, pybindings can be very slow on some server machines (observed 300x slower - went from 12s for 100 iterations to under 40ms for 100 iterations). As a stopgap measure, this diff exposes _unsafe_reset_threadpool from python, which can be used to reduce the thread count. This should be done prior to loading a model.

Differential Revision: D71023514

Author: GregoryComer

extension/pybindings/pybindings.cpp

@@ -23,6 +23,7 @@
 #include <executorch/extension/data_loader/buffer_data_loader.h>
 #include <executorch/extension/data_loader/mmap_data_loader.h>
 #include <executorch/extension/memory_allocator/malloc_memory_allocator.h>
+#include <executorch/extension/threadpool/threadpool.h>
 #include <executorch/runtime/backend/interface.h>
 #include <executorch/runtime/core/data_loader.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
@@ -1064,6 +1065,12 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
       "_reset_profile_results",
       []() { EXECUTORCH_RESET_PROFILE_RESULTS(); },
       call_guard);
+  m.def("_unsafe_reset_threadpool", 
+      [](int num_threads) { 
+        executorch::extension::threadpool::get_threadpool()->_unsafe_reset_threadpool(num_threads);
+      },
+      py::arg("num_threads")
+   );
 
   py::class_<PyModule>(m, "ExecuTorchModule")
       .def("load_bundled_input", &PyModule::load_bundled_input, call_guard)

shim_et/xplat/executorch/extension/pybindings/pybindings.bzl

@@ -54,6 +54,7 @@ def executorch_pybindings(python_module_name, srcs = [], cppdeps = [], visibilit
         ],
         deps = [
             "//executorch/runtime/core:core",
+            "//executorch/extension/threadpool:threadpool",
         ] + cppdeps,
         external_deps = [
             "pybind11",