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框架中的其他部分。