Add _unsafe_reset_threadpool to pybindings (#9196)

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. The long-term plan is tracked in #8763. The issue calls out mobile, but will also benefit server evals.

Differential Revision: D71023514

Author: GregoryComer

extension/pybindings/portable_lib.py

@@ -45,6 +45,7 @@
     _load_for_executorch_from_buffer,  # noqa: F401
     _load_for_executorch_from_bundled_program,  # noqa: F401
     _reset_profile_results,  # noqa: F401
+    _unsafe_reset_threadpool,  # noqa: F401
     BundledModule,  # noqa: F401
     ExecuTorchModule,  # noqa: F401
     MethodMeta,  # noqa: F401

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,14 @@ 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"),
+      call_guard);
 
   py::class_<PyModule>(m, "ExecuTorchModule")
       .def("load_bundled_input", &PyModule::load_bundled_input, call_guard)

extension/pybindings/pybindings.pyi

@@ -264,3 +264,12 @@ def _reset_profile_results() -> None:
         This API is experimental and subject to change without notice.
     """
     ...
+
+@experimental("This API is experimental and subject to change without notice.")
+def _unsafe_reset_threadpool(num_threads: int) -> None:
+    """
+    .. warning::
+
+        This API is experimental and subject to change without notice.
+    """
+    ...

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",