Avoid nested thread-pool locking in parallel_for

extension/threadpool/test/thread_parallel_test.cpp

@@ -1,3 +1,3 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
@@ -70,6 +70,30 @@
   }
 }
 
+TEST_P(ParallelTest, NestedCallsPreserveOuterThreadNumber) {
+  std::array<std::array<int, 3>, 10> visits{};
+  EXPECT_TRUE(parallel_for(0, 10, 1, [&](int64_t begin, int64_t end) {
+    const auto outer_thread = executorch::extension::get_thread_num();
+    for (int64_t row = begin; row < end; ++row) {
+      EXPECT_TRUE(parallel_for(2, 5, 1, [&](int64_t inner_begin, int64_t inner_end) {
+        EXPECT_EQ(executorch::extension::get_thread_num(), outer_thread);
+        for (int64_t column = inner_begin; column < inner_end; ++column) {
+          ++visits[row][column - 2];
+        }
+      }));
+      EXPECT_TRUE(parallel_for(2, 2, 1, [&](int64_t, int64_t) {
+        ADD_FAILURE() << "Empty nested range invoked its callback";
+      }));
+      EXPECT_EQ(executorch::extension::get_thread_num(), outer_thread);
+    }
+  }));
+  for (const auto& row : visits) {
+    for (const auto count : row) {
+      EXPECT_EQ(count, 1);
+    }
+  }
+}
+
 TEST_P(ParallelTest, TestAllInvokedWithMutex) {
   EXPECT_TRUE(parallel_for(0, 10, 1, [this](int64_t begin, int64_t end) {
     this->RunExclusiveTask(begin, end);

extension/threadpool/thread_parallel.cpp

@@ -11,6 +11,7 @@
 #include <tuple>
 
 #include <executorch/extension/threadpool/threadpool.h>
+#include <executorch/extension/threadpool/threadpool_guard.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/kernel/thread_parallel_interface.h>
 #include <executorch/runtime/platform/assert.h>
@@ -61,6 +62,13 @@ bool parallel_for(
       begin,
       end);
   ET_CHECK_OR_RETURN_FALSE(grain_size > 0, "grain_size = %" PRId64, grain_size);
+  if (NoThreadPoolGuard::is_enabled()) {
+    // Querying the pool's size would lock the mutex held by the outer call.
+    if (begin < end) {
+      f(begin, end);
+    }
+    return true;
+  }
   int64_t num_tasks = 0, chunk_size = 0;
   std::tie(num_tasks, chunk_size) =
       calc_num_tasks_and_chunk_size(begin, end, grain_size);