(tvm) branch main updated: [Runtime] ParallelFor skipping thread backend for unit extent (#16508)

[hongyij]

This is an automated email from the ASF dual-hosted git repository.

hongyij pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new b5f44100e5 [Runtime] ParallelFor skipping thread backend for unit 
extent (#16508)
b5f44100e5 is described below

commit b5f44100e55ac79f1a5627243247097f8b26e020
Author: Ruihang Lai <ruiha...@cs.cmu.edu>
AuthorDate: Fri Feb 2 13:57:31 2024 -0500

    [Runtime] ParallelFor skipping thread backend for unit extent (#16508)
    
    This PR skips launching multiple parallel threads in
    `parallel_for_with_threading_backend` when the input
    extent is 1, in which case we can just use the current
    thread to run the given function.
---
 include/tvm/runtime/threading_backend.h | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/include/tvm/runtime/threading_backend.h 
b/include/tvm/runtime/threading_backend.h
index 3122b000e0..e56c130b2c 100644
--- a/include/tvm/runtime/threading_backend.h
+++ b/include/tvm/runtime/threading_backend.h
@@ -203,6 +203,11 @@ inline void parallel_launch_with_threading_backend(T 
flambda) {
 
 template <typename T>
 inline void parallel_for_with_threading_backend(T flambda, int64_t begin, 
int64_t end) {
+  if (end - begin == 1) {
+    flambda(begin);
+    return;
+  }
+
   auto flaunch = [begin, end, flambda](int task_id, int num_task) {
     // For each thread, do static division and call into flambda.
     int64_t total_len = end - begin;