parallel_for should return true if precondition fails (#2240)

Summary:
Pull Request resolved: #2240

Don't crash, but log error and return false.

Reviewed By: SS-JIA

Differential Revision: D54505636

fbshipit-source-id: 1cdd2861fbae5bb355ef7bd61eb34c03f35418c5

Author: kirklandsign

extension/parallel/targets.bzl

@@ -24,5 +24,6 @@ def define_common_targets():
             deps = [
                 "//executorch/backends/xnnpack/threadpool:threadpool",
                 "//executorch/runtime/core:core",
+                "//executorch/runtime/core/exec_aten/util:tensor_util" + aten_suffix,
             ],
         )

extension/parallel/test/targets.bzl

@@ -14,5 +14,6 @@ def define_common_targets():
         ],
         deps = [
             "//executorch/extension/parallel:thread_parallel",
+            "//executorch/runtime/platform:platform",
         ],
     )

extension/parallel/test/thread_parallel_test.cpp

@@ -12,7 +12,7 @@
 #include <mutex>
 
 #include <executorch/extension/parallel/thread_parallel.h>
-#include <executorch/test/utils/DeathTest.h>
+#include <executorch/runtime/platform/platform.h>
 
 using namespace ::testing;
 
@@ -49,19 +49,19 @@ class ParallelTest : public ::testing::Test {
 };
 
 TEST_F(ParallelTest, TestAllInvoked) {
-  parallel_for(0, 10, 1, [this](int64_t begin, int64_t end) {
+  EXPECT_TRUE(parallel_for(0, 10, 1, [this](int64_t begin, int64_t end) {
     this->RunTask(begin, end);
-  });
+  }));
 
   for (int64_t i = 0; i < 10; ++i) {
     EXPECT_EQ(data_[i], i);
   }
 }
 
 TEST_F(ParallelTest, TestAllInvokedWithMutex) {
-  parallel_for(0, 10, 1, [this](int64_t begin, int64_t end) {
+  EXPECT_TRUE(parallel_for(0, 10, 1, [this](int64_t begin, int64_t end) {
     this->RunExclusiveTask(begin, end);
-  });
+  }));
 
   int expected_sum = 0;
   for (int64_t i = 0; i < 10; ++i) {
@@ -72,13 +72,10 @@ TEST_F(ParallelTest, TestAllInvokedWithMutex) {
 }
 
 TEST_F(ParallelTest, TestInvalidRange) {
-  ET_EXPECT_DEATH(
-      {
-        parallel_for(10, 0, 1, [this](int64_t begin, int64_t end) {
-          this->RunExclusiveTask(begin, end);
-        });
-      },
-      "");
+  et_pal_init();
+  EXPECT_FALSE(parallel_for(10, 0, 1, [this](int64_t begin, int64_t end) {
+    this->RunExclusiveTask(begin, end);
+  }));
 
   for (int64_t i = 0; i < 10; ++i) {
     EXPECT_EQ(data_[i], 0);
@@ -87,13 +84,10 @@ TEST_F(ParallelTest, TestInvalidRange) {
 }
 
 TEST_F(ParallelTest, TestInvalidRange2) {
-  ET_EXPECT_DEATH(
-      {
-        parallel_for(6, 5, 1, [this](int64_t begin, int64_t end) {
-          this->RunExclusiveTask(begin, end);
-        });
-      },
-      "");
+  et_pal_init();
+  EXPECT_FALSE(parallel_for(6, 5, 1, [this](int64_t begin, int64_t end) {
+    this->RunExclusiveTask(begin, end);
+  }));
 
   for (int64_t i = 0; i < 10; ++i) {
     EXPECT_EQ(data_[i], 0);
@@ -102,9 +96,9 @@ TEST_F(ParallelTest, TestInvalidRange2) {
 }
 
 TEST_F(ParallelTest, TestInvokePartialFromBeginning) {
-  parallel_for(0, 5, 1, [this](int64_t begin, int64_t end) {
+  EXPECT_TRUE(parallel_for(0, 5, 1, [this](int64_t begin, int64_t end) {
     this->RunTask(begin, end);
-  });
+  }));
 
   for (int64_t i = 0; i < 5; ++i) {
     EXPECT_EQ(data_[i], i);
@@ -115,9 +109,9 @@ TEST_F(ParallelTest, TestInvokePartialFromBeginning) {
 }
 
 TEST_F(ParallelTest, TestInvokePartialToEnd) {
-  parallel_for(5, 10, 1, [this](int64_t begin, int64_t end) {
+  EXPECT_TRUE(parallel_for(5, 10, 1, [this](int64_t begin, int64_t end) {
     this->RunTask(begin, end);
-  });
+  }));
 
   for (int64_t i = 0; i < 5; ++i) {
     EXPECT_EQ(data_[i], 0);
@@ -128,9 +122,9 @@ TEST_F(ParallelTest, TestInvokePartialToEnd) {
 }
 
 TEST_F(ParallelTest, TestInvokePartialMiddle) {
-  parallel_for(2, 8, 1, [this](int64_t begin, int64_t end) {
+  EXPECT_TRUE(parallel_for(2, 8, 1, [this](int64_t begin, int64_t end) {
     this->RunTask(begin, end);
-  });
+  }));
 
   for (int64_t i = 0; i < 2; ++i) {
     EXPECT_EQ(data_[i], 0);
@@ -144,19 +138,19 @@ TEST_F(ParallelTest, TestInvokePartialMiddle) {
 }
 
 TEST_F(ParallelTest, TestChunkSize2) {
-  parallel_for(0, 10, 2, [this](int64_t begin, int64_t end) {
+  EXPECT_TRUE(parallel_for(0, 10, 2, [this](int64_t begin, int64_t end) {
     this->RunTask(begin, end);
-  });
+  }));
 
   for (int64_t i = 0; i < 10; ++i) {
     EXPECT_EQ(data_[i], i);
   }
 }
 
 TEST_F(ParallelTest, TestChunkSize2Middle) {
-  parallel_for(3, 8, 2, [this](int64_t begin, int64_t end) {
+  EXPECT_TRUE(parallel_for(3, 8, 2, [this](int64_t begin, int64_t end) {
     this->RunTask(begin, end);
-  });
+  }));
 
   for (int64_t i = 0; i < 3; ++i) {
     EXPECT_EQ(data_[i], 0);
@@ -170,29 +164,29 @@ TEST_F(ParallelTest, TestChunkSize2Middle) {
 }
 
 TEST_F(ParallelTest, TestChunkSize3) {
-  parallel_for(0, 10, 3, [this](int64_t begin, int64_t end) {
+  EXPECT_TRUE(parallel_for(0, 10, 3, [this](int64_t begin, int64_t end) {
     this->RunTask(begin, end);
-  });
+  }));
 
   for (int64_t i = 0; i < 10; ++i) {
     EXPECT_EQ(data_[i], i);
   }
 }
 
 TEST_F(ParallelTest, TestChunkSize6) {
-  parallel_for(0, 10, 6, [this](int64_t begin, int64_t end) {
+  EXPECT_TRUE(parallel_for(0, 10, 6, [this](int64_t begin, int64_t end) {
     this->RunTask(begin, end);
-  });
+  }));
 
   for (int64_t i = 0; i < 10; ++i) {
     EXPECT_EQ(data_[i], i);
   }
 }
 
 TEST_F(ParallelTest, TestChunkSizeTooLarge) {
-  parallel_for(0, 10, 11, [this](int64_t begin, int64_t end) {
+  EXPECT_TRUE(parallel_for(0, 10, 11, [this](int64_t begin, int64_t end) {
     this->RunTask(begin, end);
-  });
+  }));
 
   for (int64_t i = 0; i < 10; ++i) {
     EXPECT_EQ(data_[i], i);

extension/parallel/thread_parallel.cpp

@@ -10,6 +10,7 @@
 
 #include <executorch/backends/xnnpack/threadpool/threadpool.h>
 #include <executorch/extension/parallel/thread_parallel.h>
+#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
 #include <executorch/runtime/platform/assert.h>
 
 namespace torch::executor {
@@ -34,14 +35,14 @@ calc_num_tasks_and_chunk_size(int64_t begin, int64_t end, int64_t grain_size) {
   return std::make_tuple(num_tasks, chunk_size);
 }
 
-void parallel_for(
+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) {
-  ET_CHECK_MSG(begin >= 0 && end >= 0, "Begin and end should be non-negative");
-  ET_CHECK_MSG(end >= begin, "end should be greater than or equal to begin");
-  ET_CHECK_MSG(grain_size > 0, "grain_size should be positive");
+  ET_LOG_AND_RETURN_IF_FALSE(begin >= 0 && end >= 0);
+  ET_LOG_AND_RETURN_IF_FALSE(end >= begin);
+  ET_LOG_AND_RETURN_IF_FALSE(grain_size > 0);
   int64_t num_tasks = 0, chunk_size = 0;
   std::tie(num_tasks, chunk_size) =
       calc_num_tasks_and_chunk_size(begin, end, grain_size);
@@ -57,6 +58,7 @@ void parallel_for(
   // Per protocol from threadpool (pthreadpool), when this returns, all tasks
   // are executed, so this is synchronous.
   get_threadpool()->run(task, num_tasks);
+  return true;
 }
 
 } // namespace torch::executor

extension/parallel/thread_parallel.h

@@ -23,12 +23,13 @@ namespace torch::executor {
  * 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.
  */
-void parallel_for(
+bool parallel_for(
     const int64_t begin,
     const int64_t end,
     const int64_t grain_size,