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

目录

概述

架构

源码分析

数据结构

CudaPlatform注册

Stream Executor初始化

Stream Executor API调用


概述

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

架构

【架构分析】Tensorflow Internals 源码分析3 - Stream Executor 连接计算设备_第1张图片 Stream Executor与Tensorflow 整体架构

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

 

【架构分析】Tensorflow Internals 源码分析3 - Stream Executor 连接计算设备_第2张图片 Stream Executor 架构图

 

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

源码分析

数据结构

【架构分析】Tensorflow Internals 源码分析3 - Stream Executor 连接计算设备_第3张图片 Stream Executor核心数据结构
  • Stream Executor中数据结构比较复杂,结合架构图的5个层次梳理其中的核心数据结构以及互相之间的关系如上图

CudaPlatform注册

【架构分析】Tensorflow Internals 源码分析3 - Stream Executor 连接计算设备_第4张图片 Cuda Platform注册时序
  •  Tensorflow每支持一个计算设备,就要注册对应的设备Platform对象(比如上图的CudaPlatform),Tensorflow通过MultiPlatformManager统一管理所有的计算设备Platform,StreamExecutor的实例就是由对应的设备Platform创建

Stream Executor初始化

【架构分析】Tensorflow Internals 源码分析3 - Stream Executor 连接计算设备_第5张图片 StreamExecutor和Stream初始化时序
  •  在TF应用程序启动创建GPUDevice的过程中,就会根据设备名字查找对应的Platform实例
  • 根据找到的CudaPlatform创建StreamExecutor和多个Stream实例,提供TF框架中Op Kernel使用CUDA GPU的API 实现

Stream Executor API调用

【架构分析】Tensorflow Internals 源码分析3 - Stream Executor 连接计算设备_第6张图片 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要求的计算

你可能感兴趣的:(tensorflow,深度学习,人工智能)