TensorRT是如何做到架構(gòu)更快的？

作者：糯鹽 2023-11-30 10:13:17

人工智能智能汽車

文章總結(jié)了TensorRT工具鏈以及TensorRT后端優(yōu)化流程。

本文經(jīng)自動駕駛之心公眾號授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請聯(lián)系出處。

一. TensorRT是什么？

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

二. 為什么？

只有Nvidia最清楚自家GPU或DLA該如何優(yōu)化，所以TensorRT跑網(wǎng)絡(luò)的速度是最快的，比直接用Pytorch快N倍。

遙遙領(lǐng)先的TensorRT

三. 怎么做到的？

1. 搜索整個優(yōu)化空間

與Pytorch等其它訓(xùn)練框架最大區(qū)別是，TensorRT的網(wǎng)絡(luò)優(yōu)化算法是基于目標(biāo)GPU所做的推理性能優(yōu)化，而其它框架一方面需要綜合考慮訓(xùn)練和推理，更重要的是它們沒有在目標(biāo)GPU上做針對性的優(yōu)化。

TensorRT又是如何針對目標(biāo)GPU優(yōu)化的呢？

簡單講就是在可能的設(shè)計空間中搜索出全局最優(yōu)解。

這個搜索空間有哪些變量呢？

比如CUDA架構(gòu)中的編程模型所對應(yīng)的，將Tensor劃分為多少個block？以及這些block如何組織到Grid中。

任務(wù)被劃分為多個Block

Block以Grid的方式組織起來

不同的組織層次以對應(yīng)不同的存儲體系結(jié)構(gòu)中的不同存儲器

再舉例，使用什么樣的指令完成計算，可能是FFMA、FMMA，可能是TensorCore指令...

更難的部分可能是Tensor數(shù)據(jù)流的調(diào)度，把他們放在local、share還是global memory呢？如何擺放呢？

這些變量組合在一起是一個巨大的搜索空間，可能你的CPU計算幾天也得不出個結(jié)果來。

但是，我們知道神經(jīng)網(wǎng)絡(luò)的計算是由一個個粒度更大的算子組成的，算子上面還有粒度更大的層結(jié)構(gòu)。我們也清楚地知道層與層之間相對獨(dú)立，也就是說可以針對每層計算優(yōu)化，最后把優(yōu)化后的層串在一起大概率就是網(wǎng)絡(luò)的全局最優(yōu)解。

于是，TensorRT預(yù)先寫了很多算子和層（CUDA Kernel）。當(dāng)然這些算子的輸入和輸出tensor是可以配置的，以適應(yīng)網(wǎng)絡(luò)輸入和輸出的不同以及GPU資源的不同。

部分優(yōu)化好的算子

搜索空間變小了，從原來的指令級別的搜索，上升到了算子級別的搜索。因?yàn)檫@些實(shí)現(xiàn)都是用CUDA kernel所寫，更準(zhǔn)確的說是Kernel級別的搜索了。

但是Tensor數(shù)據(jù)流的調(diào)度問題并沒有解決，這也是最關(guān)鍵和復(fù)雜的地方。我們應(yīng)該將輸入Tensor劃分為多少個Block呢？這些Blocks應(yīng)該分配給多少個線程呢？Tensor存儲在哪呢？local/share/global memory的哪些地方呢？中間計算結(jié)果存儲在哪里呢？

對于計算部分是可以通過模擬的方式（類似指令集仿真器）計算得到性能的，但是Tensor數(shù)據(jù)流在share/L2/Global Memory的流動過程就很難通過仿真計算得到精確結(jié)果，因?yàn)橐荒M的數(shù)據(jù)量和線程數(shù)過大，何況要嘗試的可能性還很多，靠CPU仿真計算的思路就別想了。唯一辦法就是讓候選算子在目標(biāo)GPU上直接跑跑，統(tǒng)計出性能，最后通過比對選出最優(yōu)解。TensorRT把這個過程叫做Timing，TensorRT甚至可以將優(yōu)化的中間過程存儲下來供你分析，叫做timing caching（通過trtexec --timingCacheFile=<file>）。

Nvida GPU memory架構(gòu)

以上所描述的優(yōu)化過程可以叫做Hardware Aware Optimazation。

總結(jié)起來優(yōu)化器會重點(diǎn)分析：

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等框架不會去深入挖掘的，尤其是對存儲系統(tǒng)的優(yōu)化。

2. 強(qiáng)制選擇Kernel

由于Block之間線程的運(yùn)行順序是隨機(jī)的，CPU可能在向GDDR/HBM讀寫數(shù)據(jù)，甚至GPU的時鐘頻率也在隨負(fù)載的變化而變化，這導(dǎo)致了不同系統(tǒng)運(yùn)行環(huán)境下GPU的性能表現(xiàn)會有差異。這種差異也可能導(dǎo)致TensorRT Timing的最優(yōu)解不是實(shí)際推理時的最優(yōu)解，可能選擇了次優(yōu)的Kernel。

TensorRT提供了一個補(bǔ)救方法，就是強(qiáng)制指定選擇某個Kernel實(shí)現(xiàn)，如果你很確信它是最優(yōu)解的話。

TensorRT提供的API叫做AlgorithmSelector。

3. Plugin

當(dāng)然，你對自己設(shè)計的算子更有把握，可以自己寫Kernel，然后指定使用它。

不過，更多情況下，是因?yàn)榘l(fā)現(xiàn)TensorRT不支持某個算子，你才被迫去寫Kernel，畢竟CUDA編程不簡單，何況性能還需要足夠好。

4. cuBLAS和cuDNN

TensorRT安裝指導(dǎo)要求你先安裝CUDA SDK和cuDNN。

CUDA SDK需要安裝是顯而易見的，因?yàn)門ensorRT所調(diào)用的Kernel需要NVCC編譯器來編譯成Nvidia GPU的匯編指令序列啊！

但是CUDA SDK中還有一個cuBLAS庫也是被TensorRT所依賴的，我們知道C++庫BLAS（Basic Linear Algebra Subprograms），它是針對CPU進(jìn)行的線性代數(shù)計算優(yōu)化，那么cuBLAS就是針對CUDA GPU開發(fā)的線性代數(shù)計算庫，它的底層當(dāng)然也就是用CUDA Kernel寫成的。典型的矩陣乘法算子就可以直接調(diào)用cuBLAS了。

cuBLAS開發(fā)的很早，應(yīng)該是CUDA生態(tài)最早的一批庫了吧，但是隨著深度學(xué)習(xí)的普及，Nvidia又在生態(tài)中加入了cuDNN庫，它的層次更高，封裝了到了網(wǎng)絡(luò)層，所以其實(shí)TensorRT也可以直接調(diào)用優(yōu)化好的cuDNN庫中的Kernel？是也不是。

TensorRT可以選擇所謂Tactic（策略）來決定是使用TensorRT寫的Kernel還是cuBLAS和cuDNN的。

5. Tactic

TensorRT的Tactic能決定很多優(yōu)化選項。

例如，每次timing某個算子時需要平均的運(yùn)行次數(shù)。缺省TensorRT會運(yùn)行四次，以降低不確定性帶來的誤差，但這個次數(shù)是可以修改的。

還可以決定上面提到的Kernel庫的選擇，Plugin的選擇，GPU時鐘頻率鎖定等。

6. 量化

TensorRT當(dāng)然具備網(wǎng)絡(luò)量化能力，提供了將全網(wǎng)都量化到int8的隱性量化方式，也提供了插入Q/DQ Layer的顯性量化方式。

混合量化是Nvidia做的很優(yōu)秀的地方，這對于高效利用計算資源起到了重要作用，不過，這個另外的話題，以后有機(jī)會再談。

7. 多應(yīng)用推理和多卡推理

其實(shí)這才是Nvidia強(qiáng)悍的地方，在友商都在談單卡性能時，其實(shí)多卡或多節(jié)點(diǎn)才是Nvidia的殺手锏

另外，對于單卡性能富余的情況下，可能希望有多個流并行推理，這個對于TensorRT來說也是必須支持的

四. TensorRT的內(nèi)核到底是什么？

答：根據(jù)網(wǎng)絡(luò)、輸入、輸出tensor、目標(biāo)GPU的資源，通過實(shí)際運(yùn)行，在候選Kernel庫中擇優(yōu)的一個Hardware Aware優(yōu)化器。

五. 編譯器

最后，如果非要套用編譯器前后端理論的話，上述談到的部分應(yīng)該屬于編譯器后端部分了，因?yàn)樗呀?jīng)和底層硬件息息相關(guān)了。只不過它邏輯上處于于NVCC這個實(shí)體編譯器的上層。而編譯器前端，也就是與硬件不相關(guān)的圖融合部分是也是在TensorRT的Builder內(nèi)完成的。

好了，如果你對AI編譯器還不了解，可以看下面這篇入門文章

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

最后送上兩幅圖，作為總結(jié)