tensorflow StreamExecutor

 

StreamExecutor框架

 

Plaftorm

Platform是StreamExecutor一个抽象的设备平台层，它是一个抽象类，具体实现的子类有CUDA ， HOST， OpenCL。

通过GetExecutor() 就可以获取到相应的StreamExecutor指针。而对于platform的管理是借用MultiPlatformManager来进行管理的。

MultiPlatformManager的类主要函数如下：

                                                          

主要通过RegisterPlatform中，其实这个类的内部会维护两张表，一张是（name, platorm),另外一张是（id, platorm)。 外部就可以

通过PlatformWithName 和 PlatformWithId 这两个函数来获取到相应的platform。而RegisterPlatform是在每个Platform模块在初始化的时候就注册。 以HostPlatform为例：

                                        static void InitializeHostPlatform() {

                                                   std::unique_ptr<Platform> platform(new host::HostPlatform);

                                                   SE_CHECK_OK(MultiPlatformManager::RegisterPlatform(std::move(platform)));

                                        }

                                    REGISTER_MODULE_INITIALIZER(host_platform,
                                    stream_executor::host::InitializeHostPlatform());

这种插件式的管理方式可以使得模块间很好的解耦， 当需要加入其它的模块的时候只需要创建一个Platform的子类模块，然后采用上面的方式注册进系统中就可以。

 

StreamExecutor 核心

StreamExecutorInterface具体的事例模块是采用注册回调函数的方式进行管理的，注册方式如下：

 

而MakeCUDAExecutorImplementaion的定义如下：

using StreamExecutorFactory =
    std::function<StreamExecutorInterface *(const PluginConfig &)>;

那么在StreamExecutor 就可以通过platform_kind来进行获取对应的实例化的StreamExecutorInterface

通过这种方式可以很好的扩充新的设备相应的StreamExecutorInterface。

StreamExecutor 主要是获取Dnn， Blas,   Fft,   Rng(随机数)的操作对象， 而获取这些是通过一个StreamExecutorInterface类型的

implementation_来创建这些对象。而对于其它一些操作也可以通过自定义kernel,然后通过lauch去执行。StreamExecutorInterface

也是一个父类，具体使用是通过其子类来完成，这样根据不同的设备，不同的库可以进行很好的扩展。StreamExecutorInterface

会针对具体的设备进行事件，内存的管理（如cudaExecutor将帮我们完成复杂的cuda的事件管理， 显存管理）。

 

对于Dnn, Blas, Fft, Rng 的管理是采用PluginRegistry来管理的，这个类采用了单例模式来存储所有的插件模块。

                            

               

查看以上的代码就可以看出是通过Platform::Id  PluginId 构成一个二维的表来保存插件。

模块的注册如下所示：

 

以上的代码是cuda:cudnn 模块的插件初始化。

最终就可以在StreamExecutorInterface中通过PluginRegistry来使用具体的插件

 

最终的UML流程图