目录
概述
本文基于Tensorflow r1.15源码 链接 ,重点分析抽象了计算设备算力并提供TF框架控制和使用计算设备的模块Stream Executor
架构
![](https://i-blog.csdnimg.cn/blog_migrate/66e9b1b8a79f873461cc8fbbaa7999e5.jpeg)
从上图可以看到,Stream Executor抽象了计算设备(各类xPU)的能力,提供TF框架特别是Op Kernel提供了统一使用计算设备的API,由Stream Executor将Op Kernel对计算设备Device/流Stream的创建、管理、各类加速库的使用统一封装对接到各个计算设备后端(比如 CUDA GPU)
![](https://i-blog.csdnimg.cn/blog_migrate/6fb31b75134b82386a2ff9da57474c50.jpeg)
Stream Executor的源码在tensorflow/stream_executor目录下面,Google写代码抽象层级很高(KPI导向逼的?),往往是封装再封装,直接阅读源码很容易迷失其中,不知道来龙去脉。所以我梳理总结了Stream Executor架构图,先按照架构图理解它的5层功能抽象,然后再读源码就会更加清楚当前调用处于哪一层功能模块
源码分析
数据结构
![](https://i-blog.csdnimg.cn/blog_migrate/8313b0c52779ecfcf7e37e2e400ee9a3.jpeg)
- Stream Executor中数据结构比较复杂,结合架构图的5个层次梳理其中的核心数据结构以及互相之间的关系如上图
CudaPlatform注册
![](https://i-blog.csdnimg.cn/blog_migrate/7a051bbea181346e67c8fedc94c4f6ab.jpeg)
- Tensorflow每支持一个计算设备,就要注册对应的设备Platform对象(比如上图的CudaPlatform),Tensorflow通过MultiPlatformManager统一管理所有的计算设备Platform,StreamExecutor的实例就是由对应的设备Platform创建
Stream Executor初始化
![](https://i-blog.csdnimg.cn/blog_migrate/ed9b52f1e6161417eb2d059bfa15912b.jpeg)
- 在TF应用程序启动创建GPUDevice的过程中,就会根据设备名字查找对应的Platform实例
- 根据找到的CudaPlatform创建StreamExecutor和多个Stream实例,提供TF框架中Op Kernel使用CUDA GPU的API 实现
Stream Executor API调用
![](https://i-blog.csdnimg.cn/blog_migrate/621c5275e23563149b547edb795a6f1d.jpeg)
- TF框架特别是Op Kernel在Compute实现中通过Stream API操控底层的CUDA GPU 设备
- 如Stream Executor架构图所示,该过程即按照架构图的几个层次最终调用到了CUDA SDK完成对GPU的算力使用,包括
- 通过Stream API 执行卷积算法、流同步、Device和Host的内存传输
- GPU Backend中的GPUExecutor GPUDriver执行抽象的GPU操作,转发cudnn、cuda driver 等API请求到CUDA Backend
- CUDA Backend转发API请求到CUDA API Wrapper
- CUDA API Wrapper load了CUDA SDK中driver、runtime、cuXXX等库后,调用了对应cuda API实现了对计算设备算力的使用,完成Op Kernel要求的计算