Split & remove extension_parallel

As per plan in #8932, we want to be able to include thread_parallel.h to build
libraries that are *capable* of parallelization, but don't *require* it. So, we
move the header to ExecuTorch core and add a fallback implementation (with
tests!) of `parallel_for` that just does a regular `for` loop. Then, targets
that link `extension_threadpool` will get parallelization automagically.

This PR doesn't add any optionally-parallelized code; that will be in the next
PR.

ghstack-source-id: 4f47374
ghstack-comment-id: 2702414287
Pull Request resolved: #8983

Author: swolchok

.lintrunner.toml

@@ -218,6 +218,8 @@ exclude_patterns = [
     'examples/**',
     'extension/**',
     'kernels/optimized/**',
+    # Justified <functional> include.
+    'runtime/kernel/thread_parallel_interface.h',
     'scripts/**',
     'third-party/**',
     'util/**',

CMakeLists.txt

@@ -751,7 +751,6 @@ if(EXECUTORCH_BUILD_PTHREADPOOL
    AND EXECUTORCH_BUILD_CPUINFO
 )
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/threadpool)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/parallel)
 endif()
 
 if(EXECUTORCH_BUILD_PYBIND)

Test.cmake

@@ -13,7 +13,6 @@ if(BUILD_TESTING)
   add_subdirectory(extension/evalue_util/test)
   add_subdirectory(extension/kernel_util/test)
   add_subdirectory(extension/memory_allocator/test)
-  add_subdirectory(extension/parallel/test)
   add_subdirectory(extension/pytree/test)
   add_subdirectory(kernels/portable/cpu/util/test)
   add_subdirectory(kernels/prim_ops/test)

build/cmake_deps.toml

@@ -88,7 +88,6 @@ excludes = [
 deps = [
   "executorch",
   "executorch_core",
-  "extension_parallel",
   "extension_threadpool",
   "portable_kernels",
 ]
@@ -131,7 +130,6 @@ excludes = [
 deps = [
   "executorch_core",
   "executorch",
-  "extension_parallel",
 ]
 
 [targets.optimized_native_cpu_ops]
@@ -146,7 +144,6 @@ excludes = [
 deps = [
   "executorch_core",
   "executorch",
-  "extension_parallel",
   "extension_threadpool",
   "portable_kernels",
 ]
@@ -227,19 +224,6 @@ deps = [
   "extension_runner_util",
 ]
 
-[targets.extension_parallel]
-buck_targets = [
-  "//extension/parallel:thread_parallel",
-]
-filters = [
-  ".cpp$",
-]
-deps = [
-  "executorch",
-  "executorch_core",
-  "extension_threadpool",
-]
-
 [targets.extension_tensor]
 buck_targets = [
   "//extension/tensor:tensor",
@@ -393,6 +377,7 @@ filters = [
 deps = [
   "executorch",
   "executorch_core",
+  "extension_threadpool",
 ]
 
 [targets.xnnpack_schema]
@@ -427,7 +412,6 @@ deps = [
   "executorch",
   "executorch_core",
   "optimized_kernels",
-  "extension_parallel",
   "extension_threadpool",
   "reduce_util",
   "xnnpack_backend",
@@ -465,7 +449,6 @@ deps = [
   "executorch_core",
   "extension_data_loader",
   "extension_module",
-  "extension_parallel",
   "portable_kernels",
   "quantized_kernels",
   "xnnpack_backend",

build/executorch-config.cmake

@@ -68,7 +68,6 @@ set(lib_list
     custom_ops
     extension_module
     extension_module_static
-    extension_parallel
     extension_runner_util
     extension_tensor
     extension_threadpool
@@ -119,14 +118,9 @@ endforeach()
 
 # TODO: investigate use of install(EXPORT) to cleanly handle
 # target_compile_options/target_compile_definitions for everything.
-if(TARGET extension_parallel)
-  set_target_properties(
-    extension_parallel PROPERTIES INTERFACE_LINK_LIBRARIES extension_threadpool
-  )
-endif()
 if(TARGET cpublas)
   set_target_properties(
-    cpublas PROPERTIES INTERFACE_LINK_LIBRARIES extension_parallel
+    cpublas PROPERTIES INTERFACE_LINK_LIBRARIES extension_threadpool
   )
 endif()
 target_compile_definitions(extension_threadpool INTERFACE ET_USE_THREADPOOL)

extension/llm/custom_ops/op_sdpa.cpp

@@ -19,7 +19,7 @@
 #include <vector>
 
 #ifdef ET_USE_THREADPOOL
-#include <executorch/extension/parallel/thread_parallel.h>
+#include <executorch/runtime/kernel/thread_parallel_interface.h>
 #include <executorch/extension/threadpool/threadpool.h>
 #endif
 #include <executorch/extension/kernel_util/make_boxed_from_unboxed_functor.h>

extension/llm/custom_ops/targets.bzl

@@ -37,7 +37,6 @@ def define_common_targets():
                 "//executorch/kernels/optimized:libblas{}".format(mkl_dep),
                 "//executorch/kernels/optimized:libvec",
                 "//executorch/extension/kernel_util:kernel_util",
-                "//executorch/extension/parallel:thread_parallel",
                 "//executorch/extension/threadpool:threadpool",
             ],
             deps = [

extension/parallel/CMakeLists.txt

@@ -8,46 +8,9 @@
 
 #pragma once
 
-#include <cstdint>
-#include <functional>
-
-namespace executorch {
-namespace extension {
-
-/**
- * A helper to run function in parallel.
- *
- * begin, end: describe the extent of the workitems via first and last workitem
- * to be processed
- * grain_size: number of workitems processed by user callback which is
- * described below
- * f: user function applied in parallel to the chunks, signature:
- *   void f(int64_t begin, int64_t end)
- * Returns true if all work items are processed successfully, false otherwise
- *
- * Warning: parallel_for does NOT copy thread local states from the current
- * thread to the worker threads. Users need to protect the access to captured
- * data if they mutate them in f.
- */
-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);
-
-int64_t get_thread_num();
-
-void set_thread_num(int64_t thread_num);
-
-} // namespace extension
-} // namespace executorch
-
-namespace torch {
-namespace executor {
-// TODO(T197294990): Remove these deprecated aliases once all users have moved
-// to the new `::executorch` namespaces.
-using ::executorch::extension::get_thread_num;
-using ::executorch::extension::parallel_for;
-using ::executorch::extension::set_thread_num;
-} // namespace executor
-} // namespace torch
+// This header is a stub left behind after the move to
+// executorch/runtime/kernel. Depend on this target and include this
+// header if you have a hard requirement for threading; if you want to
+// cleanly use parallelization if available, then depend on and use
+// the below header instead.
+#include <executorch/runtime/kernel/thread_parallel_interface.h>

extension/parallel/TARGETS

@@ -21,7 +21,8 @@ if(NOT CMAKE_CXX_STANDARD)
 endif()
 
 add_library(
-  extension_threadpool threadpool.cpp threadpool_guard.cpp cpuinfo_utils.cpp
+  extension_threadpool threadpool.cpp threadpool_guard.cpp thread_parallel.cpp
+                       cpuinfo_utils.cpp
 )
 target_link_libraries(
   extension_threadpool PUBLIC executorch_core cpuinfo pthreadpool
@@ -42,3 +43,7 @@ install(
   INCLUDES
   DESTINATION ${_common_include_directories}
 )
+
+if(BUILD_TESTING)
+  add_subdirectory(test)
+endif()

extension/parallel/targets.bzl

@@ -9,6 +9,7 @@ def define_common_targets():
     """
 
     _THREADPOOL_SRCS = [
+        "thread_parallel.cpp",
         "threadpool.cpp",
         "threadpool_guard.cpp",
     ] + (["fb/threadpool_use_n_threads.cpp"] if not runtime.is_oss else [])
@@ -29,6 +30,8 @@ def define_common_targets():
         exported_deps = [
             third_party_dep("pthreadpool"),
             third_party_dep("cpuinfo"),
+            # Allow users to use the header without an extra deps entry.
+            "//executorch/runtime/kernel:thread_parallel_interface",
         ],
         exported_preprocessor_flags = [
             "-DET_USE_THREADPOOL",

extension/parallel/test/TARGETS

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# @generated by test/utils/generate_gtest_cmakelists.py
+#
 # This file should be formatted with
 # ~~~
 # cmake-format -i CMakeLists.txt
@@ -12,28 +14,14 @@
 #
 
 cmake_minimum_required(VERSION 3.19)
-project(extension_parallel_test)
-
-# Use C++17 for test.
-set(CMAKE_CXX_STANDARD 17)
 
 set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..)
 
 include(${EXECUTORCH_ROOT}/build/Test.cmake)
 
-set(_test_srcs thread_parallel_test.cpp ../thread_parallel.cpp)
+set(_test_srcs thread_parallel_test.cpp threadpool_test.cpp)
 
 et_cxx_test(
-  extension_parallel_test
-  SOURCES
-  ${_test_srcs}
-  EXTRA_LIBS
-  pthreadpool
-  cpuinfo
+  extension_threadpool_test SOURCES ${_test_srcs} EXTRA_LIBS
   extension_threadpool
 )
-target_include_directories(
-  extension_parallel_test
-  PRIVATE ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/cpuinfo/include
-          ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/pthreadpool/include
-)

extension/parallel/test/targets.bzl

@@ -18,3 +18,15 @@ def define_common_targets():
             "//executorch/extension/threadpool:threadpool",
         ],
     )
+
+    runtime.cxx_test(
+        name = "thread_parallel_test",
+        srcs = [
+            "thread_parallel_test.cpp",
+        ],
+        deps = [
+            "//executorch/extension/threadpool:threadpool",
+            "//executorch/runtime/kernel:thread_parallel_interface",
+            "//executorch/runtime/platform:platform",
+        ],
+    )