【架构分析】TensorFlow CUDA Device Backend扩展分析

最新推荐文章于 2024-04-07 19:48:22 发布

HaoBBNuanMM

最新推荐文章于 2024-04-07 19:48:22 发布

阅读量372

点赞数 1

本文链接：https://blog.csdn.net/HaoBBNuanMM/article/details/114321922

版权

概述

Tensorflow Build系统

CUDA Device Backend扩展

概述

Tensorflow除了能运行在CPU上，更重要的是它支持运行在GPU TPU 乃至其它硬件加速器上，从而带来倍爽的Training和Inference运行速度，本文旨以Tensorflow r1.15版本（下载链接）为例，分析Nvidia GPU的CUDA软件栈是如何被扩展到Tensorflow上作为加速Tensorflow运行速度的设备后端（Device Backend）

Tensorflow Build系统

Google的工程师既是编程好手也是造轮子好手，Tensorflow不出意外的用了Google自建的Bazel 编译系统，整个编译分成了4个阶段，涉及的bazel核心概念及其关系如上图所示

关于Bazel的详细介绍可以参考链接，这里用上图示意Bazel中的核心概念在Tensorflow中的应用

这里特别要吐槽Bazel系统的依赖关系，如上图左侧部分所示，传统的Makefile或cmake系统对于一个编译对象T的所有依赖都是直接表达在描述T编译过程的脚本中，但是Google的工程师们抽象分层能力极强，硬是在Bazel中把编译对象T的依赖分层抽象为可以依赖另外一个编译对象t1，而t1又依赖t2 ...... 以此类推，导致T1 依赖的 D2 ... Dn 都变成了间接的依赖关系，所以当出现上图右侧的编译error提示T有头文件没找到，我们需要把整个T依赖的树型结构捋出来去判断是不是T间接依赖的某个目标（比如d3）没有添加error中提示的头文件。这种分层抽象把一个原本很简单的依赖错误排查搞的过于反人类了，可能也注定Bazel只有Google自己吃独食吧。

CUDA Device Backend扩展

Nvidia提供了完整的CUDA软件栈，包括CUDA的Toolchain（工具链比如：nvcc 等）和相关加速库libcuXXX.so，所以Tensorflow要支持运行在CUDA GPU上，第一步就是扩展Bazel的Toolchain Framework增加CUDA的Toolchain支持，从而在整个Build阶段能够使用nvcc去编译cu.cc的算子（kernel）代码，从而使得这些算子在runtime阶段部署到Nvidia GPU上加速运行