Revert "Split & remove extension_parallel (#8983)"

.ci/scripts/unittest-buck2.sh

@@ -17,10 +17,8 @@ buck2 query "//backends/apple/... + //backends/example/... + \
 //kernels/optimized/... + //kernels/portable/... + //kernels/quantized/... + \
 //kernels/test/... + //runtime/... + //schema/... + //test/... + //util/..."
 
-UNBUILDABLE_OPTIMIZED_OPS_REGEX="gelu|fft_r2c|log_softmax"
-BUILDABLE_OPTIMIZED_OPS=$(buck2 query //kernels/optimized/cpu/... | grep -E -v $UNBUILDABLE_OPTIMIZED_OPS_REGEX)
 # TODO: expand the covered scope of Buck targets.
 # //runtime/kernel/... is failing because //third-party:torchgen_files's shell script can't find python on PATH.
 # //runtime/test/... requires Python torch, which we don't have in our OSS buck setup.
-buck2 test $BUILDABLE_OPTIMIZED_OPS //kernels/portable/... //runtime/backend/... //runtime/core/... \
+buck2 test //kernels/portable/... //runtime/backend/... //runtime/core/... \
       //runtime/executor: //runtime/kernel/... //runtime/platform/...

.lintrunner.toml

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

CMakeLists.txt

@@ -751,6 +751,7 @@ 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,6 +13,7 @@ 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,6 +88,7 @@ excludes = [
 deps = [
   "executorch",
   "executorch_core",
+  "extension_parallel",
   "extension_threadpool",
   "portable_kernels",
 ]
@@ -130,7 +131,7 @@ excludes = [
 deps = [
   "executorch_core",
   "executorch",
-  "extension_threadpool",
+  "extension_parallel",
 ]
 
 [targets.optimized_native_cpu_ops]
@@ -145,6 +146,7 @@ excludes = [
 deps = [
   "executorch_core",
   "executorch",
+  "extension_parallel",
   "extension_threadpool",
   "portable_kernels",
 ]
@@ -225,6 +227,19 @@ 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",
@@ -364,7 +379,6 @@ excludes = [
 deps = [
   "executorch",
   "executorch_core",
-  "extension_threadpool",
   "xnnpack_backend",
   "portable_kernels",
 ]
@@ -379,7 +393,6 @@ filters = [
 deps = [
   "executorch",
   "executorch_core",
-  "extension_threadpool",
 ]
 
 [targets.xnnpack_schema]
@@ -414,6 +427,7 @@ deps = [
   "executorch",
   "executorch_core",
   "optimized_kernels",
+  "extension_parallel",
   "extension_threadpool",
   "reduce_util",
   "xnnpack_backend",
@@ -451,7 +465,7 @@ deps = [
   "executorch_core",
   "extension_data_loader",
   "extension_module",
-  "extension_threadpool",
+  "extension_parallel",
   "portable_kernels",
   "quantized_kernels",
   "xnnpack_backend",

build/executorch-config.cmake

@@ -75,6 +75,7 @@ set(lib_list
     custom_ops
     extension_module
     extension_module_static
+    extension_parallel
     extension_runner_util
     extension_tensor
     extension_threadpool
@@ -130,9 +131,14 @@ 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_threadpool
+    cpublas PROPERTIES INTERFACE_LINK_LIBRARIES extension_parallel
   )
 endif()
 if(TARGET extension_threadpool)

extension/llm/custom_ops/op_sdpa.cpp

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

extension/llm/custom_ops/targets.bzl

@@ -37,6 +37,7 @@ 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

@@ -0,0 +1,25 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Please keep this file formatted by running:
+# ~~~
+# cmake-format -i CMakeLists.txt
+# ~~~
+
+if(NOT (EXECUTORCH_BUILD_PTHREADPOOL AND EXECUTORCH_BUILD_CPUINFO))
+  message(FATAL_ERROR "extension/parallel requires extension/threadpool")
+endif()
+
+add_library(extension_parallel thread_parallel.cpp)
+
+target_link_libraries(extension_parallel PUBLIC executorch_core extension_threadpool)
+target_compile_options(extension_parallel PUBLIC ${_common_compile_options})
+
+install(
+  TARGETS extension_parallel
+  DESTINATION lib
+  INCLUDES
+  DESTINATION ${_common_include_directories})

extension/parallel/TARGETS

@@ -0,0 +1,8 @@
+# Any targets that should be shared between fbcode and xplat must be defined in
+# targets.bzl. This file can contain fbcode-only targets.
+
+load(":targets.bzl", "define_common_targets")
+
+oncall("executorch")
+
+define_common_targets()

extension/parallel/targets.bzl

@@ -0,0 +1,26 @@
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_aten_mode_options", "runtime")
+
+def define_common_targets():
+    """Defines targets that should be shared between fbcode and xplat.
+
+    The directory containing this targets.bzl file should also contain both
+    TARGETS and BUCK files that call this function.
+    """
+
+    runtime.cxx_library(
+        name = "thread_parallel",
+        srcs = [
+            "thread_parallel.cpp",
+        ],
+        exported_headers = [
+            "thread_parallel.h",
+        ],
+        visibility = [
+            "//executorch/...",
+            "@EXECUTORCH_CLIENTS",
+        ],
+        deps = [
+            "//executorch/extension/threadpool:threadpool",
+            "//executorch/runtime/core:core",
+        ],
+    )

extension/threadpool/test/CMakeLists.txt → extension/parallel/test/CMakeLists.txt

@@ -4,8 +4,6 @@
 # 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
@@ -14,14 +12,28 @@
 #
 
 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 threadpool_test.cpp)
+set(_test_srcs thread_parallel_test.cpp ../thread_parallel.cpp)
 
 et_cxx_test(
-  extension_threadpool_test SOURCES ${_test_srcs} EXTRA_LIBS
+  extension_parallel_test
+  SOURCES
+  ${_test_srcs}
+  EXTRA_LIBS
+  pthreadpool
+  cpuinfo
   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

@@ -0,0 +1,8 @@
+# Any targets that should be shared between fbcode and xplat must be defined in
+# targets.bzl. This file can contain fbcode-only targets.
+
+load(":targets.bzl", "define_common_targets")
+
+oncall("executorch")
+
+define_common_targets()

extension/parallel/test/targets.bzl

@@ -0,0 +1,19 @@
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+def define_common_targets():
+    """Defines targets that should be shared between fbcode and xplat.
+
+    The directory containing this targets.bzl file should also contain both
+    TARGETS and BUCK files that call this function.
+    """
+
+    runtime.cxx_test(
+        name = "thread_parallel_test",
+        srcs = [
+            "thread_parallel_test.cpp",
+        ],
+        deps = [
+            "//executorch/extension/parallel:thread_parallel",
+            "//executorch/runtime/platform:platform",
+        ],
+    )