Merge pull request #60735 from Intel-tensorflow:mabuzain/tp-use-caller-thread-fix

PiperOrigin-RevId: 549652138

tensorflow/tsl/util/onednn_threadpool.h

@@ -114,32 +114,39 @@ class OneDnnThreadPool : public threadpool_iface {
         can_use_caller_thread_ && nthr == port::NumSchedulableCPUs();
     const int njobs_to_schedule = use_caller_thread ? njobs - 1 : njobs;
 
-    tsl::BlockingCounter counter(njobs_to_schedule);
-    std::function<void(int, int)> handle_range = [=, &handle_range, &counter](
-                                                     int first, int last) {
-      while (last - first > 1) {
-        const auto mid = first + (last - first) / 2;
-        // Find something near the midpoint which is a multiple of block size.
-        eigen_interface_->ScheduleWithHint([=]() { handle_range(mid, last); },
-                                           mid, mid + 1);
-        last = mid;
-      }
-      counter.DecrementCount();
-      run_jobs(balance, first, n, njobs, fn);
-    };
-
-    // Eigen avoids a thread hop by running the root of the tree on the main
-    // thread. We have disabled this because it actually slows things down
-    // relative to base because base cheats and uses n threads while letting
-    // main continue doing other work
-    eigen_interface_->ScheduleWithHint(
-        [=]() { handle_range(0, njobs_to_schedule); }, 0, 1);
-
     if (use_caller_thread) {
-      run_jobs(balance, njobs - 1, n, njobs, fn);
+      for (int i = 0; i < njobs_to_schedule; i++) {
+        eigen_interface_->ScheduleWithHint(
+            [balance, i, n, njobs, fn]() {
+              run_jobs(balance, i, n, njobs, fn);
+            },
+            i, i + 1);
+      }
+      run_jobs(balance, njobs_to_schedule, n, njobs, fn);
+    } else {
+      tsl::BlockingCounter counter(njobs);
+      std::function<void(int, int)> handle_range = [=, &handle_range, &counter](
+                                                       int first, int last) {
+        while (last - first > 1) {
+          const auto mid = first + (last - first) / 2;
+          // Find something near the midpoint which is a multiple of block size.
+          eigen_interface_->ScheduleWithHint([=]() { handle_range(mid, last); },
+                                             mid, mid + 1);
+          last = mid;
+        }
+        counter.DecrementCount();
+        run_jobs(balance, first, n, njobs, fn);
+      };
+
+      // Eigen avoids a thread hop by running the root of the tree on the main
+      // thread. We have disabled this because it actually slows things down
+      // relative to base because base cheats and uses n threads while letting
+      // main continue doing other work
+      eigen_interface_->ScheduleWithHint([=]() { handle_range(0, njobs); }, 0,
+                                         1);
+
+      counter.Wait();
     }
-
-    counter.Wait();
   }
 
   ~OneDnnThreadPool() {}