Update

[ghstack-poisoned]

Author: swolchok

extension/threadpool/thread_parallel.cpp

@@ -54,7 +54,7 @@ bool parallel_for(
     const int64_t begin,
     const int64_t end,
     const int64_t grain_size,
-    const std::function<void(int64_t, int64_t)>& f) {
+    runtime::FunctionRef<void(int64_t, int64_t)> f) {
   ET_CHECK_OR_RETURN_FALSE(
       begin >= 0 && end >= 0 && end >= begin,
       "begin = %" PRId64 ", end = %" PRId64,

runtime/kernel/thread_parallel_interface.h

@@ -13,6 +13,7 @@
 
 #include <c10/util/irange.h>
 #include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/function_ref.h>
 #include <executorch/runtime/platform/assert.h>
 
 namespace executorch {
@@ -69,7 +70,7 @@ bool parallel_for(
     const int64_t begin,
     const int64_t end,
     const int64_t grain_size,
-    const std::function<void(int64_t, int64_t)>& f);
+    runtime::FunctionRef<void(int64_t, int64_t)> f);
 
 int64_t get_thread_num();