diff --git a/include/onnxruntime/core/platform/EigenNonBlockingThreadPool.h b/include/onnxruntime/core/platform/EigenNonBlockingThreadPool.h index b191a98e8c..04a47ca5d2 100644 --- a/include/onnxruntime/core/platform/EigenNonBlockingThreadPool.h +++ b/include/onnxruntime/core/platform/EigenNonBlockingThreadPool.h @@ -455,6 +455,11 @@ class ThreadPoolTempl : public onnxruntime::concurrency::ExtendedThreadPoolInter for (size_t i = 0; i < worker_data_.size(); ++i) worker_data_[i].thread.reset(); } + // Run fn(). Ordinarily, the function will be added to the thread pool and executed + // by a worker thread. If the thread pool rejects the work then fn() will instead + // execute synchronously during Schedule(fn). Currently the thread pool will only + // reject work if the queue of pending work is full. + void Schedule(std::function fn) override { Task t = env_.CreateTask(std::move(fn)); PerThread* pt = GetPerThread(); @@ -469,14 +474,16 @@ class ThreadPoolTempl : public onnxruntime::concurrency::ExtendedThreadPoolInter WorkerData &td = worker_data_[q_idx]; Queue& q = td.queue; t = q.PushBack(std::move(t)); - if (t.f) { - // The queue rejected the work; run it directly - env_.ExecuteTask(t); - } else { + if (!t.f) { // The queue accepted the work; ensure that the thread will pick it up td.EnsureAwake(); } } + + // Run the work directly if the queue rejected the work + if (t.f) { + env_.ExecuteTask(t); + } } // The thread pool maintains a set of hints for which threads will be good to distribute diff --git a/onnxruntime/test/platform/threadpool_test.cc b/onnxruntime/test/platform/threadpool_test.cc index d7ab072be6..fda2cbbc49 100644 --- a/onnxruntime/test/platform/threadpool_test.cc +++ b/onnxruntime/test/platform/threadpool_test.cc @@ -104,6 +104,39 @@ void TestMultipleParallelFor(const std::string& name, int num_threads, int num_c } } +void TestBurstScheduling(const std::string& name, int num_tasks) { + // Test submitting a burst of functions for executing. The aim is to provoke cases such + // as the thread pool's work queues being full. + for (int rep = 0; rep < 5; rep++) { + std::atomic ctr{0}; + // Schedule a burst of num_tasks back-to-back, and then cleanly shut down the thread + // pool. The synchronization barrier during shut down should ensure that all of the + // tasks are complete. Note that if the thread pool's work queues are full, then a + // call to tp->Schedule() may run its argument synchronously. In any case, we expect + // ctr==num_tasks. + CreateThreadPoolAndTest(name, 2, [&](ThreadPool* tp) { + // First variant : schedule from outside the pool + for (int tasks = 0; tasks < num_tasks; tasks++) { + tp->Schedule([&]() { + ctr++; + }); + } + }); + ASSERT_TRUE(ctr == num_tasks); + CreateThreadPoolAndTest(name, 2, [&](ThreadPool* tp) { + // Second variant : schedule from inside the pool + tp->Schedule([&, tp]() { + for (int tasks = 0; tasks < num_tasks; tasks++) { + tp->Schedule([&]() { + ctr++; + }); + } + }); + }); + ASSERT_TRUE(ctr == num_tasks*2); + } +} + } // namespace namespace onnxruntime { @@ -202,6 +235,24 @@ TEST(ThreadPoolTest, TestMultipleParallelFor_4Thread_4Conc_8Tasks) { TEST(ThreadPoolTest, TestMultipleParallelFor_4Thread_4Conc_1MTasks) { TestMultipleParallelFor("TestMultipleParallelFor_4Thread_4Conc_1MTasks", 4, 4, 1000000); } + +TEST(ThreadPoolTest, TestBurstScheduling_0Tasks) { + TestBurstScheduling("TestBurstScheduling_0Tasks", 0); +} + +TEST(ThreadPoolTest, TestBurstScheduling_1Task) { + TestBurstScheduling("TestBurstScheduling_1Task", 1); +} + +TEST(ThreadPoolTest, TestBurstScheduling_16Tasks) { + TestBurstScheduling("TestBurstScheduling_16Tasks", 16); +} + +TEST(ThreadPoolTest, TestBurstScheduling_65536Task) { + // Attempt to exhaust the size of the queues used in the thread pool to + // buffer tasks. + TestBurstScheduling("TestBurstScheduling_65536Tasks", 65536); +} #ifdef _WIN32 TEST(ThreadPoolTest, TestStackSize) { ThreadOptions to;