TensorFlow中的並行執行引擎——StreamExecutor框架

背景

[做者：DeepLearningStack，阿里巴巴算法工程師，開源TensorFlow Contributor]

在前一篇文章中，咱們梳理了TensorFlow中各類異構Device的添加和註冊機制，經過使用預先定義好的宏，各類自定義好的Device可以將本身註冊到全局表中。TensorFlow指望經過這種模式，可以讓Device的添加和註冊於系統自己更好的解耦，從而體現了較好的模塊化特性。在這篇文章中，咱們選擇直接去窺探TensorFlow底層架構較爲複雜的一個部分——StreamExecutor框架。咱們已經知道TensorFlow是一個異構的並行執行框架，對於異構Device的管理是一件很是複雜的事，不只包括Device的添加、註冊、刪除、屬性的管理，還必需要對Device的並行執行過程作進一步抽象造成統一的框架，才能實現更好的解耦。經過閱讀這部分源碼不但能夠對執行引擎的管理有很深的理解，還能夠體驗學習到各類設計模式。若是想要對TensorFlow底層甚至是XLA作一些性能上的深度優化，那麼這一部分則是必需要了解的內容。

Stream

Stream存在於計算機相關的各類技術中，好比在操做系統、流式計算、計算機網絡傳輸或是CUDA編程中都有涉及。Stream從抽象角度來看其本質是定義了一個操做序列，處於同一個Stream的操做必須按順序執行，不一樣Stream之間的並沒有順序關係。在TensorFlow中存在一些高性能的並行編程設備，因此須要有一套抽象框架對這些設備的執行過程管理起來，這就是StreamExecutor的用武之地了。

StreamExecutor簡介

其實StreamExecutor自己就是一個在Google內部爲並行編程模型開發的單獨的庫，感興趣的能夠直接參考GitHub。在TensorFlow中的StreamExecutor是一個開源StreamExecutor的簡版，而且並非以第三方庫的形式出現，而是在源碼中單獨放了一個stream_executor的文件夾，裏面的代碼很是的精簡，目錄結構部分截圖以下圖所示。

StreamExecutor爲TensorFlow的執行層面提供了較爲統一的抽象，而在底層各類Device的執行管理細節卻徹底不一樣。咱們能夠看到stream_executor下面有cuda和host兩個子目錄，他們分別是GPU執行引擎和CPU執行引擎所使用的子模塊。下面咱們先從統一的抽象層面來梳理該框架的結構。

StreamExecutor對外提供的句柄——Stream對象

爲了隱藏StreamExecutor框架管理的複雜性，它對外暴露的handler必須足夠簡單。事實也確實如此，StreamExecutor經過暴露Stream對象做爲操做底層的handler。通常而言，在TensorFlow的框架中都是使用Stream對象來調用底層計算庫，進行設備間數據拷貝操做等過程。好比調用Stream對象的ThenMemcpy便可完成異步的數據傳輸拷貝過程，調用ThenConvolveXXX等函數便可完成DNN庫中的卷積調用。事實上，TensorFlow中不少Op的C++實現中，其Compute函數內就是經過使用Stream對象來完成某些實際計算或數據拷貝的過程，下圖展現了Stream對象、StreamExecutor框架以及其餘模塊的關係。

Stream對象是經過持有StreamInterface的具體實現對象來得到實際平臺的Stream，進而經過Stream這個統一的handler完成與底層的交互，下面試這一子模塊的類圖結構。

 

StreamExecutor框架內的層次結構

熟悉GPU編程的同窗都知道，CUDA程序的編寫是相對複雜的，不但要針對某種任務設計特定的並行編程思路，還要管理Event，Stream等較爲底層的對象。爲了可以減輕StreamExecutor用戶的使用負擔，也爲了可以給上層調用者即TensorFlow引擎提供更加統一的接口，一些抽象分層的工做是很是有必要的。整體上StreamExecutor框架由三個層次組成，從上到下依次爲Platform層（平臺描述）、StreamExecutor Core層（執行引擎）和LibrarySupport層（基礎庫）。若是須要爲TensorFlow添加新的計算設備種類，不但要向TensorFlow中註冊Device的定義，還須要在StreamExecutor框架中提供負責管理該Device計算的代碼。

Platform層

在StreamExecutor中Platform指的是計算所使用設備平臺的抽象，每種Device對應一種Platform。好比GPU對應的是CudaPlatform，而CPU對應的是HostPlatform等。一旦得到了某種Device的Platform，就能夠獲取和該Platform對應的StreamExecutor Core以及相應的LibrarySupport。在TensorFlow的代碼實現中，全部Platform類都是經過宏定義和MultiPlatformManager管理類的靜態方法主動註冊到系統中的，下面是這一層次的類圖表示。

CudaPlatform和HostPlatform繼承自公共父類Platform，若是有新的Platform出現，依然能夠沿用這樣的設計直接繼承並給出實現。全部的Platform都經過MultiPlaftormManager調用RegsiterPlatform函數主動註冊到系統中並作初始化，下面代碼段是CudaPlaftorm的註冊過程，註冊使用了Initializer模塊及相應的宏定義，這些代碼比較簡單，這裏就再也不詳細展開了。

 1 static void InitializeCudaPlatform() {
 2   // Disabling leak checking, MultiPlatformManager does not destroy its
 3   // registered platforms.
 4 
 5   std::unique_ptr<cuda::CudaPlatform> platform(new cuda::CudaPlatform);
 6   SE_CHECK_OK(MultiPlatformManager::RegisterPlatform(std::move(platform)));
 7 }
 8 
 9 }  // namespace stream_executor
10 
11 REGISTER_MODULE_INITIALIZER(cuda_platform,
12                             stream_executor::InitializeCudaPlatform());
13 
14 // Note that module initialization sequencing is not supported in the
15 // open-source project, so this will be a no-op there.
16 REGISTER_MODULE_INITIALIZER_SEQUENCE(cuda_platform, multi_platform_manager);
17 REGISTER_MODULE_INITIALIZER_SEQUENCE(multi_platform_manager_listener,
18                                      cuda_platform);

MultiPlatformManager提供了兩種獲取具體Platform的方式，一種是經過name，另外一種是經過Id，以下代碼段所示。

 1   // Retrieves the platform registered with the given platform name (e.g.
 2   // "CUDA", "OpenCL", ...) or id (an opaque, comparable value provided by the
 3   // Platform's Id() method).
 4   //
 5   // If the platform has not already been initialized, it will be initialized
 6   // with a default set of parameters.
 7   //
 8   // If the requested platform is not registered, an error status is returned.
 9   // Ownership of the platform is NOT transferred to the caller --
10   // the MultiPlatformManager owns the platforms in a singleton-like fashion.
11   static port::StatusOr<Platform*> PlatformWithName(absl::string_view target);
12   static port::StatusOr<Platform*> PlatformWithId(const Platform::Id& id);

StreamExecutor Core層

從源代碼上看這一層很是複雜，由於它涉及到的類最多，可是當咱們把Platform層和Library層分開看待後，這一層就變得很是簡單了。對於外部使用者來講，獲取Platform就是爲了獲取對應的執行引擎。對於TensorFlow這種存在多種Platform和執行引擎的異構框架來講，必須爲每一種執行引擎提供完整的實現，這具備必定的複雜度。爲了讓代碼結構更有層次感，也爲了向Platform層隱藏底層的設計複雜度，該層選擇只向上層暴露StreamExecutor類，而涉及到具體實現的StreamExecutorInterface以及各類具體的實現將由StreamExecutor類統一控制，這種代理的方式讓這一層的架構更加乾淨，下面是涉及到這一層的類圖。

CudaExecutor和HostExecutor繼承自StreamExecutorInterface後，由StreamExecutor持有，並暴露給上一層Platform使用。同各類Platform相似，每一個具體的StreamExecutor也須要註冊到系統中，但他們卻沒有依賴於任何控制類，直接經過宏定義將本身註冊到全局工廠中，註冊過程也是藉助Initializer模塊實現的。下面的代碼段展現了CudaExecutor的註冊過程。

 1 void initialize_cuda_gpu_executor() {
 2   *internal::MakeCUDAExecutorImplementation() = [](const PluginConfig &config) {
 3     return new cuda::CUDAExecutor{config};
 4   };
 5 }
 6 
 7 }  // namespace stream_executor
 8 
 9 REGISTER_MODULE_INITIALIZER(cuda_gpu_executor, {
10   stream_executor::initialize_cuda_gpu_executor();
11 });

initialize_cuda_gpu_executor函數中定義了一個建立CUDAExecutor的匿名函數，而MakeCUDAExecutorImplementation函數實際上建立了一個全局的table，中間的等號賦值操做實際上就是把該匿名函數放到了全局instance中，這實際上就是一種簡單的工廠模式，在StreamExecutor中存在多種相似的工廠，下面代碼段展現了這些工廠的本質。

1 using StreamExecutorFactory =
2     std::function<StreamExecutorInterface *(const PluginConfig &)>;
3 using EventFactory = std::function<EventInterface *(StreamExecutor *)>;
4 using StreamFactory = std::function<StreamInterface *(StreamExecutor *)>;
5 using TimerFactory = std::function<TimerInterface *(StreamExecutor *)>;
6 using KernelFactory = std::function<KernelInterface*()>;
7 
8 StreamExecutorFactory* MakeCUDAExecutorImplementation();

StreamExecutor框架使用Cache機制避免爲同一種StreamExecutor Core被重複建立，這個Cache就是ExecutorCache，下面代碼展現了Platform從Cache獲取StreamExecutor Core的內容，當Cache中不存在所須要的StreamExecutor時，會建立新的對象並放入cache中，並以config做爲key。

1 port::StatusOr<StreamExecutor*> CudaPlatform::GetExecutor(
2     const StreamExecutorConfig& config) {
3   return executor_cache_.GetOrCreate(
4       config, [&]() { return GetUncachedExecutor(config); });
5 }

Library層

這一層提供的是各類底層加速庫的接入，當前該層主要負責接入Dnn，Blas，Rng和Fft模塊，每一個模塊和對應的類說明以下表所示 。

子模塊名稱	功能說明
DNNSupport	DNN計算模塊，主要包含DNN計算的基本操做。在GPU實現中，它將做爲CuDNN的封裝
RngSupport	隨機數生成模塊
BlasSupport	基礎線性代數庫模塊，主要包含矩陣系列的計算，在CPU實現中它能夠是Eigen，mkl等；在GPU實現中，它將做爲CuBLAS的封裝
FFTSupport	FFT系列運算模塊

由於這些基礎庫同StreamExecutor相似，都具備平臺屬性，例如在CUDAHostPlatform中使用的Blas庫應爲CuBLAS，而HostPlatform中對應的多是OpenBlas，MKL等。雖然StreamExecutorInterface建立出來的各類Library指針均由StreamExecutor持有，可是他們卻由StreamExecutorInterface的實現類負責建立，因此從邏輯上看他們處於StreamExecutor Core的下一層，下圖展現了Library層的類圖。

Library層將這些基礎庫統一做爲插件（Plugin）來管理，用以應對將來出現的各類各樣的基礎庫。他們經過PluginRegister模塊註冊。和StreamExecutor Core中的管理方式相同，依然要先建立插件的Factory，Factory的建立也經過宏實現。以CudnnSupport爲例，經過向通用初始化模塊Intializer傳入initialize_cudnn函數並調用，將建立CudnnSupport的函數做爲DnnFactory放到PluginRegister模塊中，至此完成了DnnFactory的建立。使用時，只須要拿到PluginRegister的key（即要求拿到何種插件）便可取出對應的LibrarySupport。下面展現了CudnnSupport的工廠註冊代碼。

 1 void initialize_cudnn() {
 2   port::Status status =
 3       PluginRegistry::Instance()->RegisterFactory<PluginRegistry::DnnFactory>(
 4           cuda::kCudaPlatformId, cuda::kCuDnnPlugin, "cuDNN",
 5           [](internal::StreamExecutorInterface* parent) -> dnn::DnnSupport* {
 6             cuda::CUDAExecutor* cuda_executor =
 7                 dynamic_cast<cuda::CUDAExecutor*>(parent);
 8             if (cuda_executor == nullptr) {
 9               LOG(ERROR) << "Attempting to initialize an instance of the cuDNN "
10                          << "support library with a non-CUDA StreamExecutor";
11               return nullptr;
12             }
13 
14             cuda::CudnnSupport* dnn = new cuda::CudnnSupport(cuda_executor);
15             if (!dnn->Init().ok()) {
16               // Note: Init() will log a more specific error.
17               delete dnn;
18               return nullptr;
19             }
20             return dnn;
21           });
22 
23   if (!status.ok()) {
24     LOG(ERROR) << "Unable to register cuDNN factory: "
25                << status.error_message();
26   }
27 
28   PluginRegistry::Instance()->SetDefaultFactory(
29       cuda::kCudaPlatformId, PluginKind::kDnn, cuda::kCuDnnPlugin);
30 }
31 
32 }  // namespace stream_executor
33 
34 REGISTER_MODULE_INITIALIZER(register_cudnn,
35                             { stream_executor::initialize_cudnn(); });

再看整體類圖

在StreamExecutor框架中還存在其餘模塊，好比XLA的支持，好比Event的管理，在逐個梳理StreamExecutor框架的三個層次後再看其他部分就很是清晰明瞭了，下面的兩張圖展現了總體類圖和一些繼承結構。

 

StreamExecutor的調用棧

在完整的理解了StreamExecutor框架的內部結構和外部句柄後，咱們就能夠很是清晰地trace其調用棧了。最後，咱們以調用Cudnn中的FusedConvolveWIthAlgorithm爲例，畫出完整的調用時序圖。FusedConvolveWIthAlgorithm是將Convolution計算，Bias計算以及Activation計算fuse在一塊兒的優化版本CUDA kernel，它的效率相對於分開調用相比更高。

總結

StreamExecutor是一個相對獨立的項目，在TensorFlow中所使用的StreamExecutor是精簡以後的版本。正是由於異構框架管理每種Device的並行執行過程很是繁雜，因此須要StreamExecutor向上層調用者隱藏底層的複雜性。在架構設計上，StreamExecutor選擇向上層暴露簡單的Stream對象handler實現了這一封裝。事實上，TensorFlow中全部須要調用與Device相關的第三方高性能計算庫的Op都使用Stream這一handler輕鬆完成Op的編寫。從StreamExecutor框架內部看，能夠分爲Platform層、StreamExecutor Core層和LibrarySupport層，每層的核心組件都經過Initializer模塊和宏定義主動註冊到系統Factory中，從上層Op對Stream的調用棧中也能夠清晰地感覺到這井井有條的架構設計。掌握並理解StreamExecutor的調用棧是很是重要的，由於不管是爲TensorFlow底層作XLA優化仍是爲某些Op提供Int8計算支持，都須要改寫這一部分。未來咱們在梳理XLA總體框架時還會回過頭來窺探StreamExecutor框架中的其餘部分。