讀TensorFlow 源碼筆記(1): tensorflow的多線程機制eigen::threadpool

讀TensorFlow 源碼筆記(1): tensorflow的多線程機制eigen::threadpool

線程池的概念大家都知道，就是事先創建固定數量或者不困定數量的線程，計算並行任務是直接調用線程池裏的線程，從而減少了線程的頻繁創建與銷燬等動作對資源的消耗。
TensorFlow 是高效的深度學習計算框架，以計算子operator爲單位進行調度，底層依賴於第三方數學庫eigen，算子間與算子內都可以都可以進行並行計算。在ConfigProto 文件中有定義inter_/intra_op_parallelism_threads,其中inter_op_parallelism_threads表示算子間並行計算線程數；其中intra_op_parallelism_threads表示算子內並行計算線程數，如CONV算子等。
TensorFlow的線程池直接封裝Eigen 的threadpool,菜鳥今天從eigen到TensorFlow源碼一層層細細品味eigen::threadpool 機制源碼。

1.tensorflow 與eigen 的threadpool 界面

先看看在TensorFlow中的threadpool 類圖架構：
![在這裏插入圖片描述]直https://接上傳(blog.csdnimg.cn/20200306223qm6135.pn742Fg?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Z1emkyMDEy,size_16,color_FFFFFF,t_70)https:/在TensorFlow源碼
在TensorFlow源碼tensorflow/core/platform/threadpool.h文件中230-237行：

  // underlying_threadpool_ is the user_threadpool if user_threadpool is
  // provided in the constructor. Otherwise it is the eigen_threadpool_.
  Eigen::ThreadPoolInterface* underlying_threadpool_;
  // eigen_threadpool_ is instantiated and owned by thread::ThreadPool if
  // user_threadpool is not in the constructor.
  std::unique_ptr<Eigen::ThreadPoolTempl<EigenEnvironment>> eigen_threadpool_;
  std::unique_ptr<Eigen::ThreadPoolDevice> threadpool_device_;
  TF_DISALLOW_COPY_AND_ASSIGN(ThreadPool);

從以上類圖跟代碼中都能看到 TensorFlow裏面直接封裝了Eigen::ThreadPoolInterface、Eigen::ThreadPoolTempl、Eigen::ThreadPoolDevice。

• Eigen::ThreadPoolInterface：提供一個eigen線程池的接口，供用戶自定義線程池。
• Eigen::ThreadPoolTempl ： 提供一個eigen的線程池實現，TensorFlow中的線程池多用了這個實現。
• Eigen::ThreadPoolDevice ：這個結構體在eigen中./unsupported/Eigen/CXX11/src/Tensor/TensorReduction.h找到，表示能夠使用線程池的設備，這裏指的是CPU。
  在eigen 源碼中也沒有找到 Eigen::ThreadPoolTempl 這個類的定義，我們在eigen源碼目錄下執行：

 grep 'ThreadPoolTempl' -R ./

能找到一下兩行內容：

./unsupported/Eigen/CXX11/ThreadPool:template <typename Env> using ThreadPoolTempl = NonBlockingThreadPoolTempl<Env>;
./unsupported/Eigen/CXX11/ThreadPool:template <typename Env> using ThreadPoolTempl = SimpleThreadPoolTempl<Env>;

也就是ThreadPoolTempl是NonBlockingThreadPoolTempl<Env> 和 SimpleThreadPoolTempl<Env> 一個選擇，具體定義在接口文件./unsupported/Eigen/CXX11/ThreadPool中：

// Use the more efficient NonBlockingThreadPool by default.
namespace Eigen {
#ifndef EIGEN_USE_SIMPLE_THREAD_POOL
template <typename Env> using ThreadPoolTempl = NonBlockingThreadPoolTempl<Env>;
typedef NonBlockingThreadPool ThreadPool;
#else
template <typename Env> using ThreadPoolTempl = SimpleThreadPoolTempl<Env>;
typedef SimpleThreadPool ThreadPool;
#endif
}  // namespace Eigen

2 eigen 中的threadpool

從上面的描述中可以知道，在Eigen庫中，有兩種線程管理方式，即NonBlockThreadPool和SimpleThreadPool，在默認情況下TensorFlow使用NonBlockThreadPool。爲了簡單起見，這裏僅從比較簡潔明瞭的SimpleThreadPool簡單線程管理機制闡述相應的內容，它定義在SimpleThreadPool.h文件中。先看一下源碼的類圖：

我們首先分析它的構造函數。

 // Construct a pool that contains "num_threads" threads.
  explicit SimpleThreadPoolTempl(int num_threads, Environment env = Environment())
      : env_(env), threads_(num_threads), waiters_(num_threads) {
    for (int i = 0; i < num_threads; i++) {
      threads_.push_back(env.CreateThread([this, i]() { WorkerLoop(i); }));
    }
  }

構造函數有兩個參數：

• num_threads用來指定線程池中線程的個數；
• Environment類在此處是一個模板類，是一個運行環境的接口，裏面要實現EnvThread* CreateThread(std::function<void()> f) 方法和Task CreateTask(std::function<void()> f) 方法。目前Eigen有一相應的實現類StlThreadEnvironment類封裝了一些與線程相關的操作，而這些操作都是調用了C++的標準庫，但是這樣實現，其實也是爲了以後擴展和移植方便。事實上，在TensorFlow中也有Environment類，其中定義了一些與操作系統、文件系統、線程等相關的方法。

構造函數主要在完成一些初始化操作，對應上邊的類圖看：
包括初始化Environment對象env_，初始化線程數num_threads並且創建線程並保存在線程數組threads_中，以及初始化waiters_數組。需要說明的是，waiters_數組是用於存放線程池中那些空閒線程，這就是它長度與線程個數相同的原因。結構體：

struct Waiter {
    std::condition_variable cv;  //信號條件變量，作爲掛起線程與激活激活線程的信號量
    Task task;                   // 需要被線程調度的任務(封裝可執行的函數)
    bool ready;                  // 線程是否繁忙的標誌
  };

在類圖中可以看到，除了上面提及的三個成員變量外，SimpleThreadPool還有一個重要的成員叫做pending_，定義如下。

std::deque<Task> pending_;  

pending_是一個雙向隊列，它用來保存那些等待被線程領取的任務。pending_隊列元素類型是Task類，它是一個函數對象的封裝，見ThreadEnvironment.h。

struct Task {
    std::function<void()> f;
  };

SimpleThreadPool的工作流程大致如下。當一個任務（Task）過來時，會首先從waiters_數組中查看有沒有空閒的線程。如果有，那麼從waiters_數組中拿一個線程出來，並把該任務分配給線程處理，並激活該線程幹活。如果waiters_數組爲空，說明所有線程都忙碌，此時會將Task放到pending_隊列中。上面的邏輯定義在SimpleThreadPool.h文件的Schedule函數中。

void Schedule(std::function<void()> fn) final {
    Task t = env_.CreateTask(std::move(fn)); // 創建一個Task類的對象
    std::unique_lock<std::mutex> l(mu_);
    if (waiters_.empty()) {                  //判斷是否有空閒線程
      pending_.push_back(std::move(t));      // 若沒有空閒線程，放進隊列裏等待
    } else {                                 //若有空閒進程
      Waiter* w = waiters_.back();          //在隊列中取出該waiter對象
      waiters_.pop_back();  
      w->ready = true;                     // 設置可執行標誌 
      w->task = std::move(t);              //分配任務  
      w->cv.notify_one();                  //激活線程去幹活
    }
  }

從上面的代碼我們還可以猜到，waiters_數組實現了一個棧，採用的是後進先出策略，在線程池中充當一個監工的角色。那些人沒活兒幹，我給你記着；等活兒來了，我給你分配，並叫你去幹活。

線程的關鍵邏輯：
在構建函數裏就看到，在初始化創建線程時就開始執行一個WorkerLoop(i)函數，

threads_.push_back(env.CreateThread([this, i]() { WorkerLoop(i); }));

比作一個員工，每個線程的工作邏輯是WorkerLoop函數定義的，在SimpleThreadPool.h文件中。

void WorkerLoop(int thread_id) {
    std::unique_lock<std::mutex> l(mu_);
    PerThread* pt = GetPerThread();
    pt->pool = this;
    pt->thread_id = thread_id;    // 定義自己的ID ，也就是該員工的工號
    Waiter w;   
    Task t;
    while (!exiting_) {          // 沒有被叫停
      if (pending_.empty()) {     //   如果隊列爲空 
        // Wait for work to be assigned to me
        w.ready = false;         //
        waiters_.push_back(&w);  // 把該空閒線程隊列中等待被喚醒
        while (!w.ready) {
          w.cv.wait(l);          //
        }
        t = w.task;
        w.task.f = nullptr;
      } else {                 // 等待隊列中有任務
        // Pick up pending work
        t = std::move(pending_.front());   //取出任務
        pending_.pop_front();    
        if (pending_.empty()) {
          empty_.notify_all();    //通知等待隊列有空位，可以插值
        }
      }
      if (t.f) {
        mu_.unlock();
        env_.ExecuteTask(t);    //執行任務
        t.f = nullptr;
        mu_.lock();
      }
    }
  }

線程池的大體流程如下：