【架构分析】Tensorflow Internals 源码分析3 - Stream Executor 连接计算设备

目录

概述

架构

源码分析

数据结构

CudaPlatform注册

Stream Executor初始化

Stream Executor API调用

---

概述

本文基于Tensorflow r1.15源码 链接 ，重点分析抽象了计算设备算力并提供TF框架控制和使用计算设备的模块Stream Executor

架构

Stream Executor与Tensorflow 整体架构

从上图可以看到，Stream Executor抽象了计算设备（各类xPU）的能力，提供TF框架特别是Op Kernel提供了统一使用计算设备的API，由Stream Executor将Op Kernel对计算设备Device/流Stream的创建、管理、各类加速库的使用统一封装对接到各个计算设备后端（比如 CUDA GPU）

 

Stream Executor 架构图

 

Stream Executor的源码在tensorflow/stream_executor目录下面，Google写代码抽象层级很高（KPI导向逼的？），往往是封装再封装，直接阅读源码很容易迷失其中，不知道来龙去脉。所以我梳理总结了Stream Executor架构图，先按照架构图理解它的5层功能抽象，然后再读源码就会更加清楚当前调用处于哪一层功能模块

源码分析

数据结构

Stream Executor核心数据结构

• Stream Executor中数据结构比较复杂，结合架构图的5个层次梳理其中的核心数据结构以及互相之间的关系如上图

CudaPlatform注册

Cuda Platform注册时序

•  Tensorflow每支持一个计算设备，就要注册对应的设备Platform对象（比如上图的CudaPlatform），Tensorflow通过MultiPlatformManager统一管理所有的计算设备Platform，StreamExecutor的实例就是由对应的设备Platform创建

Stream Executor初始化

StreamExecutor和Stream初始化时序

•  在TF应用程序启动创建GPUDevice的过程中，就会根据设备名字查找对应的Platform实例
• 根据找到的CudaPlatform创建StreamExecutor和多个Stream实例，提供TF框架中Op Kernel使用CUDA GPU的API 实现

Stream Executor API调用

TF框架使用stream_executor-API时序

 

• 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要求的计算