【架构分析】Tensorflow Internals 源码分析4 - TF Core之MatMul Kernel 生命周期

目录

概述

TF Core

MatMul Kernel的生命周期

---

概述

本文基于Tensorflow r1.15源码 链接 ，以MatMul Kernel的整个生命周期为例来分析Tensorflow 最核心的Core是如何实现Kernel的注册、查找、运行

TF Core

tensorflow/core 目录结构如上图，其中包含的代码相当多，如果直接看进去很快就迷失其中，以MatMul Kernel的整个生命周期为例可以将其中核心的数据结构及其关系梳理为下图

TF Core 核心数据结构

• GPU设备的控制与使用模块：把GPU设备抽象为StreamExecutor和Stream API（关于这部分内容的详细介绍，可以参考我的上一篇文章链接）和第三方的Eigen API，用于创建/销毁GPU流、启动GPU kernel、分配GPU内存、调动GPU加速库的等等，真正实现对GPU算力的使用与控制
• 计算设备抽象模块：Google工程师写代码喜欢封装再封装，所以考虑到计算设备的种类很多，所以再框架中封装一层设备抽象模块，用工厂模块创建实例
• 内存管理模块：将分配和管理计算设备内存的操作做了抽象与封装，用工厂模块创建实例
• 执行计算的上下文和Tensor：执行Tensorflow的OpKernel依赖上下文，包括Kernel的上下文和Device设备的上下文，这个上下文中包含了诸如设备对象、Kernel对象、Input Tensor、设备流 等实例，在执行OpKernel计算的时候提供对GPU设备的抽象控制

MatMul Kernel的生命周期

MatMul Kernel 生命周期核心时序

 

上图从MatMul Kernel的注册、计算图执行的时候查找MatMul Kernel实例、MatMul运行并调用到CUDA libcublas库三个层次，完整展示了MatMul Kernel的生命周期（其它Kernel也类似）

代码调用关系比较复杂，要结合TF Core核心的数据结构来理解它们的关系，这样就容易理解整个代码逻辑了。