llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp

   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