ExecutionEngine/Orc/TaskDispatch.h

0001 //===--------- TaskDispatch.h - ORC task dispatch utils ---------*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008 //
0009 // Task and TaskDispatch classes.
0010 //
0011 //===----------------------------------------------------------------------===//
0012
0013 #ifndef LLVM_EXECUTIONENGINE_ORC_TASKDISPATCH_H
0014 #define LLVM_EXECUTIONENGINE_ORC_TASKDISPATCH_H
0015
0016 #include "llvm/Config/llvm-config.h"
0017 #include "llvm/Support/Debug.h"
0018 #include "llvm/Support/ExtensibleRTTI.h"
0019 #include "llvm/Support/raw_ostream.h"
0020
0021 #include <cassert>
0022 #include <string>
0023
0024 #if LLVM_ENABLE_THREADS
0025 #include <condition_variable>
0026 #include <deque>
0027 #include <mutex>
0028 #include <thread>
0029 #endif
0030
0031 namespace llvm {
0032 namespace orc {
0033
0034 /// Represents an abstract task for ORC to run.
0035 class Task : public RTTIExtends<Task, RTTIRoot> {
0036 public:
0037   static char ID;
0038
0039   virtual ~Task() = default;
0040
0041   /// Description of the task to be performed. Used for logging.
0042   virtual void printDescription(raw_ostream &OS) = 0;
0043
0044   /// Run the task.
0045   virtual void run() = 0;
0046
0047 private:
0048   void anchor() override;
0049 };
0050
0051 /// Base class for generic tasks.
0052 class GenericNamedTask : public RTTIExtends<GenericNamedTask, Task> {
0053 public:
0054   static char ID;
0055   static const char *DefaultDescription;
0056 };
0057
0058 /// Generic task implementation.
0059 template <typename FnT> class GenericNamedTaskImpl : public GenericNamedTask {
0060 public:
0061   GenericNamedTaskImpl(FnT &&Fn, std::string DescBuffer)
0062       : Fn(std::forward<FnT>(Fn)), Desc(DescBuffer.c_str()),
0063         DescBuffer(std::move(DescBuffer)) {}
0064   GenericNamedTaskImpl(FnT &&Fn, const char *Desc)
0065       : Fn(std::forward<FnT>(Fn)), Desc(Desc) {
0066     assert(Desc && "Description cannot be null");
0067   }
0068   void printDescription(raw_ostream &OS) override { OS << Desc; }
0069   void run() override { Fn(); }
0070
0071 private:
0072   FnT Fn;
0073   const char *Desc;
0074   std::string DescBuffer;
0075 };
0076
0077 /// Create a generic named task from a std::string description.
0078 template <typename FnT>
0079 std::unique_ptr<GenericNamedTask> makeGenericNamedTask(FnT &&Fn,
0080                                                        std::string Desc) {
0081   return std::make_unique<GenericNamedTaskImpl<FnT>>(std::forward<FnT>(Fn),
0082                                                      std::move(Desc));
0083 }
0084
0085 /// Create a generic named task from a const char * description.
0086 template <typename FnT>
0087 std::unique_ptr<GenericNamedTask>
0088 makeGenericNamedTask(FnT &&Fn, const char *Desc = nullptr) {
0089   if (!Desc)
0090     Desc = GenericNamedTask::DefaultDescription;
0091   return std::make_unique<GenericNamedTaskImpl<FnT>>(std::forward<FnT>(Fn),
0092                                                      Desc);
0093 }
0094
0095 /// IdleTask can be used as the basis for low-priority tasks, e.g. speculative
0096 /// lookup.
0097 class IdleTask : public RTTIExtends<IdleTask, Task> {
0098 public:
0099   static char ID;
0100
0101 private:
0102   void anchor() override;
0103 };
0104
0105 /// Abstract base for classes that dispatch ORC Tasks.
0106 class TaskDispatcher {
0107 public:
0108   virtual ~TaskDispatcher();
0109
0110   /// Run the given task.
0111   virtual void dispatch(std::unique_ptr<Task> T) = 0;
0112
0113   /// Called by ExecutionSession. Waits until all tasks have completed.
0114   virtual void shutdown() = 0;
0115 };
0116
0117 /// Runs all tasks on the current thread.
0118 class InPlaceTaskDispatcher : public TaskDispatcher {
0119 public:
0120   void dispatch(std::unique_ptr<Task> T) override;
0121   void shutdown() override;
0122 };
0123
0124 #if LLVM_ENABLE_THREADS
0125
0126 class DynamicThreadPoolTaskDispatcher : public TaskDispatcher {
0127 public:
0128   DynamicThreadPoolTaskDispatcher(
0129       std::optional<size_t> MaxMaterializationThreads)
0130       : MaxMaterializationThreads(MaxMaterializationThreads) {}
0131
0132   void dispatch(std::unique_ptr<Task> T) override;
0133   void shutdown() override;
0134 private:
0135   bool canRunMaterializationTaskNow();
0136   bool canRunIdleTaskNow();
0137
0138   std::mutex DispatchMutex;
0139   bool Shutdown = false;
0140   size_t Outstanding = 0;
0141   std::condition_variable OutstandingCV;
0142
0143   std::optional<size_t> MaxMaterializationThreads;
0144   size_t NumMaterializationThreads = 0;
0145   std::deque<std::unique_ptr<Task>> MaterializationTaskQueue;
0146   std::deque<std::unique_ptr<Task>> IdleTaskQueue;
0147 };
0148
0149 #endif // LLVM_ENABLE_THREADS
0150
0151 } // End namespace orc
0152 } // End namespace llvm
0153
0154 #endif // LLVM_EXECUTIONENGINE_ORC_TASKDISPATCH_H