xref: /llvm-project/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp (revision 2b3aff8f92f8713c27791997df3ad5f2a8b51301)
1 //===------------ TaskDispatch.cpp - ORC task dispatch utils --------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "llvm/ExecutionEngine/Orc/TaskDispatch.h"
10 #include "llvm/Config/llvm-config.h" // for LLVM_ENABLE_THREADS
11 #include "llvm/ExecutionEngine/Orc/Core.h"
12 
13 namespace llvm {
14 namespace orc {
15 
16 char Task::ID = 0;
17 char GenericNamedTask::ID = 0;
18 const char *GenericNamedTask::DefaultDescription = "Generic Task";
19 
20 void Task::anchor() {}
21 TaskDispatcher::~TaskDispatcher() = default;
22 
23 void InPlaceTaskDispatcher::dispatch(std::unique_ptr<Task> T) { T->run(); }
24 
25 void InPlaceTaskDispatcher::shutdown() {}
26 
27 #if LLVM_ENABLE_THREADS
28 void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
29   bool IsMaterializationTask = isa<MaterializationTask>(*T);
30 
31   {
32     std::lock_guard<std::mutex> Lock(DispatchMutex);
33 
34     if (IsMaterializationTask) {
35 
36       // If this is a materialization task and there are too many running
37       // already then queue this one up and return early.
38       if (MaxMaterializationThreads &&
39           NumMaterializationThreads == *MaxMaterializationThreads) {
40         MaterializationTaskQueue.push_back(std::move(T));
41         return;
42       }
43 
44       // Otherwise record that we have a materialization task running.
45       ++NumMaterializationThreads;
46     }
47 
48     ++Outstanding;
49   }
50 
51   std::thread([this, T = std::move(T), IsMaterializationTask]() mutable {
52     while (true) {
53 
54       // Run the task.
55       T->run();
56 
57       std::lock_guard<std::mutex> Lock(DispatchMutex);
58       if (!MaterializationTaskQueue.empty()) {
59         // If there are any materialization tasks running then steal that work.
60         T = std::move(MaterializationTaskQueue.front());
61         MaterializationTaskQueue.pop_front();
62         if (!IsMaterializationTask) {
63           ++NumMaterializationThreads;
64           IsMaterializationTask = true;
65         }
66       } else {
67         // Otherwise decrement work counters.
68         if (IsMaterializationTask)
69           --NumMaterializationThreads;
70         --Outstanding;
71         if (Outstanding == 0)
72           OutstandingCV.notify_all();
73         return;
74       }
75     }
76   }).detach();
77 }
78 
79 void DynamicThreadPoolTaskDispatcher::shutdown() {
80   std::unique_lock<std::mutex> Lock(DispatchMutex);
81   Running = false;
82   OutstandingCV.wait(Lock, [this]() { return Outstanding == 0; });
83 }
84 #endif
85 
86 } // namespace orc
87 } // namespace llvm
88