ONNXRuntime的線程池接口在Eigen線程池接口基礎之上擴展而來(題外話:TensorFlow中的線程池同樣是建立在Eigen線程池基礎上),以下是線程池的繼承關系,其中 ThreadPoolTempl 是對接口的實現:
在 \(Environment::Initialize()\) 函數中,通過調用 \(onnxruntime::concurrency::CreateThreadPool\) 分別構建算子內和算子間的線程池($ intra_op_thread_pool_ $ & \(inter_op_thread_pool_\))。\(CreateThreadPool\)內部通過調用 \(onnxruntime::concurrency::CreateThreadPoolHelper(onnxruntime::Env *env, OrtThreadPoolParams options)\) ,進而初始化 \(ThreadPool\) 實例,並使用 \(unique\_ptr\) 裝飾該實例,作為返回結果。以上只是對類的繼承關系進行了簡單介紹,下面根據代碼中的注釋文檔解釋線程池的內部構造。
ORT線程池在構造上可分為底層和上層,其中底層可以窺探線程內部工作機制(EigenNonBlockingThreadPool.h),上層提供方便易用的接口(threadpool.h)。我們的目的是了解底層原理,因此接下來對這一層進行特別介紹。
ORT線程池派生於Eigen的非阻塞線程池,但是很多細節已經隨着迭代而更新很多。ORT主要內容包含以下幾點:
- 線程池維護一組系統線程(OS threads),用於執行ThreadPoolTempl::WorkerLoop。每個線程都擁有自己的運行隊列(RunQueue),運行隊列中都是被push進來的待執行的任務(Task)。主要的工作任務是從隊列中彈出一個任務並執行至結束。如果線程的運行隊列為空,則線程陷入自旋(spin)等待任務到達,並且嘗試從其它線程的運行隊列中“偷取”任務來執行,如果沒有偷來任務,則阻塞在系統中。在創建線程池時會通過配置標志(flag)和常量 spin_count 來實現這種“spin-then-block”操作;
- 雖然所有的任務(Task)都是簡單的 void()->void 函數,但是從概念上來看共有三種不同的類型:
- 通過Schedule()方法從外部提交的一次性任務(one-shot task),被用於支持異步工作。這些任務在並行處理器中使用,但是在測試工具(參考threadpool_test.cc中的一些樣例)之外沒有被廣泛使用;
- 運行 parallel loop 的任務。這些任務在 threadpool.cc 中被定義,並通過 RunParallel->SummonWorkers()提交到運行隊列中。每個任務將在內部循環,通過 atoic-fetch-and-add 從用戶代碼中提取迭代,直到循環完成。這種兩層方法讓我們將超輕量級的per-iteration-batch工作與管理任務對象的成本更高的per-loop工作分開。
- 運行 parallel section 的任務,這是對上面 parallel loop 方法的擴展,這些任務定義在 RunInParallelSection->SummonWorkers。parallel sections的附加層是進一步的攤銷成本的方法:創建任務完成的work可以執行一次,然后在一系列循環中利用。
此外,修改后的 Eigen 線程池有幾個方面需要強調:
- 運行隊列遵從常規設計,在隊頭/隊尾提供push/pop操作,並且針對擁有運行隊列的線程優化PopFront的情況。
(未完待續)