本文經自動駕駛之心公眾號授權轉載，轉載請聯系出處。

一. 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編譯器還不了解，可以看下面這篇入門文章

https://zhuanlan.zhihu.com/p/632648673

最后送上兩幅圖，作為總結

TensorRT工具鏈

TensorRT后端優化流程