[GitHub] [tvm] YuchenJin commented on a change in pull request #9056: [Support] Add `parallel_for_dynamic` with dynamic schedules

[GitBox <gi...@apache.org>]

YuchenJin commented on a change in pull request #9056:
URL: https://github.com/apache/tvm/pull/9056#discussion_r713396111



##########
File path: src/support/parallel_for.cc
##########
@@ -93,5 +93,52 @@ void parallel_for(int begin, int end, const std::function<void(int)>& f, int ste
   }
 }
 
+void parallel_for_dynamic(int begin, int end, int num_threads,
+                          const std::function<void(int thread_id, int task_id)>& f) {
+  // Step 1. Sanity checks
+  if (begin == end) {
+    return;
+  }
+  CHECK_LT(begin, end) << "ValueError: The interval [begin, end) requires `begin <= end`";
+  CHECK_GT(num_threads, 0) << "ValueError: `num_threads` should be positive";
+  // Step 2. Launch threads
+  // Step 2.1. Launch worker 1 to worker `num_threads - 1`
+  std::atomic<int> counter{begin};
+  std::vector<std::future<void>> futures;
+  std::vector<std::thread> threads;
+  futures.reserve(num_threads - 1);
+  threads.reserve(num_threads - 1);
+  auto worker = [end, &counter, &f](int thread_id) -> void {
+    for (int task_id; (task_id = counter++) < end;) {
+      f(thread_id, task_id);
+    }
+  };
+  for (int thread_id = 1; thread_id < num_threads; ++thread_id) {
+    std::packaged_task<void(int)> task(worker);
+    futures.emplace_back(task.get_future());
+    threads.emplace_back(std::move(task), thread_id);
+  }
+  // Step 2.2. Launch worker 0 inplace
+  try {
+    worker(0);
+  } catch (const std::exception& e) {
+    for (auto&& thread : threads) {
+      thread.join();
+    }
+    LOG(FATAL) << "Parallel_for error with " << e.what();

Review comment:
       Use `Parallel_for_dynamic error` for clarity and to distinguish it from the static schedules?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org