TensorRT~内核

这里讲述TRT是如何做到比其他框架更快的,总结了TensorRT工具链以及TensorRT后端优化流程。

2016年Nvidia为自家GPU加速推理而提供的SDK，人们有时也把它叫做推理框架。

#为什么？

只有Nvidia最清楚自家GPU或DLA该如何优化，所以TensorRT跑网络的速度是最快的，比直接用Pytorch快N倍。

遥遥领先的TensorRT

#怎么做到的？

1. 搜索整个优化空间

与Pytorch等其它训练框架最大区别是，TensorRT的网络优化算法是基于目标GPU所做的推理性能优化，而其它框架一方面需要综合考虑训练和推理，更重要的是它们没有在目标GPU上做针对性的优化。

TensorRT又是如何针对目标GPU优化的呢？

简单讲就是在可能的设计空间中搜索出全局最优解。

这个搜索空间有哪些变量呢？

比如CUDA架构中的编程模型所对应的，将tensor划分为多少个block？以及这些block如何组织到Grid中

任务被划分为多个Block

Block以Grid的方式组织起来

不同的组织层次以对应不同的存储体系结构中的不同存储器

再举例，使用什么样的指令完成计算，可能是FFMA、FMMA，可能是TensorCore指令...

更难的部分可能是tensor数据流的调度，把他们放在local、share还是global memory呢？如何摆放呢？

这些变量组合在一起是一个巨大的搜索空间，可能你的CPU计算几天也得不出个结果来。

但是我们知道神经网络的计算是由一个个粒度更大的算子组成的，算子上面还有粒度更大的层结构。我们也清楚地知道层与层之间相对独立，也就是说可以针对每层计算优化，最后把优化后的层串在一起大概率就是网络的全局最优解。 

于是，TensorRT预先写了很多算子和层（CUDA Kernel）。当然这些算子的输入和输出tensor是可以配置的，以适应网络输入和输出的不同以及GPU资源的不同。

部分优化好的算子

搜索空间变小了，从原来的指令级别的搜索，上升到了算子级别的搜索。因为这些实现都是用CUDA kernel所写，更准确的说是Kernel级别的搜索了。

但是tensor数据流的调度问题并没有解决，这也是最关键和复杂的地方。我们应该将输入tensor划分为多少个Block呢？这些Blocks应该分配给多少个线程呢？tensor存储在哪呢？local/share/global memory的哪些地方呢？中间计算结果存储在哪里呢？

对于计算部分是可以通过模拟的方式（类似指令集仿真器）计算得到性能的，但是tensor数据流在share/L2/Global Memory的流动过程就很难通过仿真计算得到精确结果，因为要被模拟的数据量和线程数过大，何况要尝试的可能性还很多，靠CPU仿真计算的思路就别想了。唯一办法就是让候选算子在目标GPU上直接跑跑，统计出性能，最后通过比对选出最优解。TensorRT把这个过程叫做Timing，TensorRT甚至可以将优化的中间过程存储下来供你分析，叫做timing caching（通过trtexec --timingCacheFile=<file>）。

Nvida GPU memory架构

以上所描述的优化过程可以叫做Hardware Aware Optimazation

总结起来优化器会重点分析：

• Type of hardware（Hardware capability...）
• Memory footprint（Share, Cache, Global...）
• Input and output shape
• Weight shapes
• Weight sparsity
• Level of quantization （so, reconsider memory)

而这些是Pytorch等框架不会去深入挖掘的，尤其是对存储系统的优化。

2. 强制选择Kernel

由于Block之间线程的运行顺序是随机的，CPU可能在向GDDR/HBM读写数据，甚至GPU的时钟频率也在随负载的变化而变化，这导致了不同系统运行环境下GPU的性能表现会有差异。这种差异也可能导致TensorRT Timing的最优解不是实际推理时的最优解，可能选择了次优的Kernel。

TensorRT提供了一个补救方法，就是强制指定选择某个Kernel实现，如果你很确信它是最优解的话。

TensorRT提供的API叫做AlgorithmSelector

3. Plugin

当然，你对自己设计的算子更有把握，可以自己写Kernel，然后指定使用它

不过更多情况下，是因为发现TensorRT不支持某个算子，你才被迫去写Kernel，毕竟CUDA编程不简单，何况性能还需要足够好。

4. cuBLAS和cuDNN

TensorRT安装指导要求你先安装CUDA SDK和cuDNN

CUDA SDK需要安装是显而易见的，因为TensorRT所调用的Kernel需要NVCC编译器来编译成Nvidia GPU的汇编指令序列啊

但是CUDA SDK中还有一个cuBLAS库也是被TensorRT所依赖的，我们知道C++库BLAS（Basic Linear Algebra Subprograms），它是针对CPU进行的线性代数计算优化，那么cuBLAS就是针对CUDA GPU开发的线性代数计算库，它的底层当然也就是用CUDA Kernel写成的。典型的矩阵乘法算子就可以直接调用cuBLAS了。

cuBLAS开发的很早，应该是CUDA生态最早的一批库了吧，但是随着深度学习的普及，Nvidia又在生态中加入了cuDNN库，它的层次更高，封装了到了网络层，所以其实TensorRT也可以直接调用优化好的cuDNN库中的Kernel？是也不是

TensorRT可以选择所谓Tactic（策略）来决定是使用TensorRT写的Kernel还是cuBLAS和cuDNN的

5. Tactic

TensorRT的Tactic能决定很多优化选项

例如，每次timing某个算子时需要平均的运行次数。缺省TensorRT会运行四次，以降低不确定性带来的误差，但这个次数是可以修改的。

还可以决定上面提到的Kernel库的选择，Plugin的选择，GPU时钟频率锁定等。

6. 量化

TensorRT当然具备网络量化能力，提供了将全网都量化到int8的隐性量化方式，也提供了插入Q/DQ Layer的显性量化方式。

混合量化是Nvidia做的很优秀的地方，这对于高效利用计算资源起到了重要作用，不过，这个另外的话题，以后有机会再谈。

7. 多应用推理和多卡推理

其实这才是Nvidia强悍的地方，在友商都在谈单卡性能时，其实多卡或多节点才是Nvidia的杀手锏

另外，对于单卡性能富余的情况下，可能希望有多个流并行推理，这个对于TensorRT来说也是必须支持的

#TensorRT的内核到底是什么？

答：根据网络、输入、输出tensor、目标GPU的资源，通过实际运行，在候选Kernel库中择优的一个Hardware Aware优化器。

#编译器

最后，如果非要套用编译器前后端理论的话，上述谈到的部分应该属于编译器后端部分了，因为它已经和底层硬件息息相关了。只不过它逻辑上处于于NVCC这个实体编译器的上层。而编译器前端，也就是与硬件不相关的图融合部分是也是在TensorRT的Builder内完成的。

最后送上两幅图，作为总结

最后在说一下这个

AI编译器入门

坏消息是

框架太多了，Pytorch、TensorFlow、Caffe、Mxnet、Keras、PaddlePaddle...

语言太多了，C/C++、Python、Java、Fortran...

XPU太多了，CPU、GPU、APU、DPU、TPU、NPU、GPGPU、FPGA、ASIC...

好消息是

AI开发框架大部分是基于Python语言的，当然很多也支持C、Java等，但至少我们的组合少了很多。

但是无论如何，AI开发者在采用不同框架、不同语言在不同目标device上部署的组合情况还是非常多，如果每个组合都开发一套编译器框架实在是太费力。那么该如何简化整个行业编译器的开发负担呢？ 答案就是IR（Intermediate representation）。

宏观上讲，不同AI框架虽然有很多不同，但是他们最终都会生成神经网络结构，这些结构适合用图表示；不同编程语言语法有很多不同，但是他们都是在表达程序运行的数据移动、控制、计算、输出等；不同处理器虽然也有很大不同，但是他们大部分都是基于冯洛伊曼架构，即便FPGA、ASIC与众不同，也逃不过指令、数据流、控制、计算等基本要素。

既然有这么多共性，那么就可以在每一层次上做一个兼容的表达，所有框架都可以输出成一个图表达，所有编程语言都可以输出成某个数据结构表达。

IR不是个新东西，我们可以来看看历史。

以LLVM（low level virtual machine）为例来，它的输入是高级编程语言，输出是机器码。还含一系列模块和工具链。初略来说，LLVM由3个部分组成，前端、中端和后端。这个中端就是我们说的IR。

如果有一个新的编程语言需要在各种处理器上运行程序，那么它可以借用LLVM，并且只开发前端模块，前端模块将高级语言转化为LLVM IR，然后优化和各种处理器机器码生产的事就可以省了。

另一方面，如果有个新的处理器需要适配各种编程语言，那么它可以借用LLVM，并且只开发后端模块，并且和LLVM IR连接起来，然后前端和优化的事就可以省了。

这种方法，极大简化了编译器开发工作量，尤其在扩展性上非常优秀。

AI编译器分层思路

借鉴了LLVM传统编译器后，我们可以把AI编译器也按类似的思路分解为“语言/框架”“IR”“硬件”的层次结构，只不过，为了让IR更细分（LLVM也如此），还可以把IR分为与硬件无关和有关两个部分，这样的好处是能让框架输出的高层IR与硬件无关，高层IR的优化结果就可以被所有硬件所共享。说白了还是避免重复劳动嘛，人懒也是一种美德：）

高层IR的抽象层次是以神经网络的“层”为出发点的，因为网络图优化的一个重要发力点就是“层”之间的融合，层之间的运行顺序，并行，以层为单位的数据载入、输出、存储分配等。

TensorRT层融合的例子

题外话，有人会问，层融合有什么好处？好处是，每层都是由输入和输出的，这个输入输出在处理器上会表现为数据的远距离搬运，可能是在处理器计算单元和Cache之间，也可能是计算单元和DDR之间，这取决于数据量的大小和Cache的大小，但是无论如何，数据搬运都是需要消耗时间和能量的。如果能合并某些层，数据就可能在计算单元内流动，而不会到Cache甚至DDR。

再往AI编译器层次结构下层走，是“算子”层也可以叫kernal层（大概率来自CUDA的语法）。网络的“层”是由“算子”组合而成的，所以这样的分法挺符合逻辑。但是我们理解其本质的区别还应该从是否硬件相关来看，因为这是IR分层的目的。算子层已经开始和底层硬件打交道，这部分的IR对接和优化需要处理器开发者提供，例如Intel的MKL, Nvidia的CuDNN和TensorRT，这些“libraries”是处理器开发者经过精细优化的高效算法，它们深度结合了处理器的特点。例如：CuDNN就会采用适合的精度来调用Tensor Core来进行矩阵运算，这比直接用CUDA core内其它计算单元快N倍。

最底层就不用说了，各种处理器硬件，包括FPGA、ASIC这些定制化加速器。

自从AI处理器兴起，AI编译器也开始迅速发展，其中最主流的是TVM（Tensor virtual machine）和MLIR（Multi-level IR，曾经以为是machine learning IR）

至于这两个IR哪个更好或更合适，优缺点是什么，没有发言权，上面这个链接算是我觉得靠谱的干货，供大家一起学习！

比较同意上文作者观点在于，如果想针对NPU开发编译器，后端IR适配硬件的活基本还是需要深度做的，需要你了解硬件底层细节，或许NPU开发者更合适，当然也必须具备基本的软件修养。朋友告诉我，到了中后期，大量的工作就是写算子，优化算子，这些活做着做着就从高大上变成了苦力活，需要一个一个算子，一个一个网络啃下来。所以对于想入坑的朋友做好心里准备罗，优势是市面上这类夸专用软硬件的人才极少，稀缺，坏处是做NPU的也就那几家，并且很辛苦：）

原创作者: whaosoft 转载于: https://blog.51cto.com/whaosoft/11637901