xref: /llvm-project/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp (revision 2d10b7b750f97b42055d5b9b08a88c18ff811cd2)
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 char IdleTask::ID = 0;
19 
20 const char *GenericNamedTask::DefaultDescription = "Generic Task";
21 
22 void Task::anchor() {}
23 void IdleTask::anchor() {}
24 
25 TaskDispatcher::~TaskDispatcher() = default;
26 
27 void InPlaceTaskDispatcher::dispatch(std::unique_ptr<Task> T) { T->run(); }
28 
29 void InPlaceTaskDispatcher::shutdown() {}
30 
31 #if LLVM_ENABLE_THREADS
32 void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
33 
34   enum { Normal, Materialization, Idle } TaskKind;
35 
36   if (isa<MaterializationTask>(*T))
37     TaskKind = Materialization;
38   else if (isa<IdleTask>(*T))
39     TaskKind = Idle;
40   else
41     TaskKind = Normal;
42 
43   {
44     std::lock_guard<std::mutex> Lock(DispatchMutex);
45 
46     // Reject new tasks if they're dispatched after a call to shutdown.
47     if (Shutdown)
48       return;
49 
50     if (TaskKind == Materialization) {
51 
52       // If this is a materialization task and there are too many running
53       // already then queue this one up and return early.
54       if (!canRunMaterializationTaskNow())
55         return MaterializationTaskQueue.push_back(std::move(T));
56 
57       // Otherwise record that we have a materialization task running.
58       ++NumMaterializationThreads;
59     } else if (TaskKind == Idle) {
60       if (!canRunIdleTaskNow())
61         return IdleTaskQueue.push_back(std::move(T));
62     }
63 
64     ++Outstanding;
65   }
66 
67   std::thread([this, T = std::move(T), TaskKind]() mutable {
68     while (true) {
69 
70       // Run the task.
71       T->run();
72 
73       // Reset the task to free any resources. We need this to happen *before*
74       // we notify anyone (via Outstanding) that this thread is done to ensure
75       // that we don't proceed with JIT shutdown while still holding resources.
76       // (E.g. this was causing "Dangling SymbolStringPtr" assertions).
77       T.reset();
78 
79       // Check the work queue state and either proceed with the next task or
80       // end this thread.
81       std::lock_guard<std::mutex> Lock(DispatchMutex);
82 
83       if (TaskKind == Materialization)
84         --NumMaterializationThreads;
85       --Outstanding;
86 
87       if (!MaterializationTaskQueue.empty() && canRunMaterializationTaskNow()) {
88         // If there are any materialization tasks running then steal that work.
89         T = std::move(MaterializationTaskQueue.front());
90         MaterializationTaskQueue.pop_front();
91         TaskKind = Materialization;
92         ++NumMaterializationThreads;
93         ++Outstanding;
94       } else if (!IdleTaskQueue.empty() && canRunIdleTaskNow()) {
95         T = std::move(IdleTaskQueue.front());
96         IdleTaskQueue.pop_front();
97         TaskKind = Idle;
98         ++Outstanding;
99       } else {
100         if (Outstanding == 0)
101           OutstandingCV.notify_all();
102         return;
103       }
104     }
105   }).detach();
106 }
107 
108 void DynamicThreadPoolTaskDispatcher::shutdown() {
109   std::unique_lock<std::mutex> Lock(DispatchMutex);
110   Shutdown = true;
111   OutstandingCV.wait(Lock, [this]() { return Outstanding == 0; });
112 }
113 
114 bool DynamicThreadPoolTaskDispatcher::canRunMaterializationTaskNow() {
115   return !MaxMaterializationThreads ||
116          (NumMaterializationThreads < *MaxMaterializationThreads);
117 }
118 
119 bool DynamicThreadPoolTaskDispatcher::canRunIdleTaskNow() {
120   return !MaxMaterializationThreads ||
121          (Outstanding < *MaxMaterializationThreads);
122 }
123 
124 #endif
125 
126 } // namespace orc
127 } // namespace llvm
128