TensorFlow中的并行执行引擎——StreamExecutor框架

在前一篇文章中，我们梳理了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管理类的静态方法主动注册到系统中的，下面是这一层次的类图表示。