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模塊及相應的宏定義，這些代碼比較簡單，這里就不再詳細展開了。

static void InitializeCudaPlatform() {
  // Disabling leak checking, MultiPlatformManager does not destroy its
  // registered platforms.

  std::unique_ptr<cuda::CudaPlatform> platform(new cuda::CudaPlatform);
  SE_CHECK_OK(MultiPlatformManager::RegisterPlatform(std::move(platform)));
}

}  // namespace stream_executor

REGISTER_MODULE_INITIALIZER(cuda_platform,
                            stream_executor::InitializeCudaPlatform());

// Note that module initialization sequencing is not supported in the
// open-source project, so this will be a no-op there.
REGISTER_MODULE_INITIALIZER_SEQUENCE(cuda_platform, multi_platform_manager);
REGISTER_MODULE_INITIALIZER_SEQUENCE(multi_platform_manager_listener,
                                     cuda_platform);

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

  // Retrieves the platform registered with the given platform name (e.g.
  // "CUDA", "OpenCL", ...) or id (an opaque, comparable value provided by the
  // Platform's Id() method).
  //
  // If the platform has not already been initialized, it will be initialized
  // with a default set of parameters.
  //
  // If the requested platform is not registered, an error status is returned.
  // Ownership of the platform is NOT transferred to the caller --
  // the MultiPlatformManager owns the platforms in a singleton-like fashion.
  static port::StatusOr<Platform*> PlatformWithName(absl::string_view target);
  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的注冊過程。

void initialize_cuda_gpu_executor() {
  *internal::MakeCUDAExecutorImplementation() = [](const PluginConfig &config) {
    return new cuda::CUDAExecutor{config};
  };
}

}  // namespace stream_executor

REGISTER_MODULE_INITIALIZER(cuda_gpu_executor, {
  stream_executor::initialize_cuda_gpu_executor();
});

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

using StreamExecutorFactory =
    std::function<StreamExecutorInterface *(const PluginConfig &)>;
using EventFactory = std::function<EventInterface *(StreamExecutor *)>;
using StreamFactory = std::function<StreamInterface *(StreamExecutor *)>;
using TimerFactory = std::function<TimerInterface *(StreamExecutor *)>;
using KernelFactory = std::function<KernelInterface*()>;

StreamExecutorFactory* MakeCUDAExecutorImplementation();

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

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

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的工廠注冊代碼。

void initialize_cudnn() {
  port::Status status =
      PluginRegistry::Instance()->RegisterFactory<PluginRegistry::DnnFactory>(
          cuda::kCudaPlatformId, cuda::kCuDnnPlugin, "cuDNN",
          [](internal::StreamExecutorInterface* parent) -> dnn::DnnSupport* {
            cuda::CUDAExecutor* cuda_executor =
                dynamic_cast<cuda::CUDAExecutor*>(parent);
            if (cuda_executor == nullptr) {
              LOG(ERROR) << "Attempting to initialize an instance of the cuDNN "
                         << "support library with a non-CUDA StreamExecutor";
              return nullptr;
            }

            cuda::CudnnSupport* dnn = new cuda::CudnnSupport(cuda_executor);
            if (!dnn->Init().ok()) {
              // Note: Init() will log a more specific error.
              delete dnn;
              return nullptr;
            }
            return dnn;
          });

  if (!status.ok()) {
    LOG(ERROR) << "Unable to register cuDNN factory: "
               << status.error_message();
  }

  PluginRegistry::Instance()->SetDefaultFactory(
      cuda::kCudaPlatformId, PluginKind::kDnn, cuda::kCuDnnPlugin);
}

}  // namespace stream_executor

REGISTER_MODULE_INITIALIZER(register_cudnn,
                            { 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框架中的其他部分。