淺淡深度學習的發機機——張量計算

淺淡深度學習的發機機——張量計算

張量計算是個看似陌生，實際上很常用的事物，它包括圖形渲染的透明度混合、圖像處理的濾鏡、數學計算中的矩陣乘法、卷積等等，是圖形引擎、圖像算法、機器學習以及深度學習的基礎。如何進行高效的張量計算，是OpenCV之類的圖像庫、OpenBlas / Eigen之類的高性能計算庫以及MNN之類的深度學習推理引擎要解決的核心問題。
本文主要以端側深度學習推理引擎MNN爲示例，談一下張量計算及主要優化策略。

相關鏈接
https://github.com/alibaba/MNN

1. 概念說明

本節對張量計算做一些概念上的科普，熟悉Tensorflow 或者其他深度學習推理引擎的可以跳過本節。

1.1 張量

張量（英文Tensor）是標量、矢量、矩陣等概念的總稱與拓展，是機器學習領域的基礎數據結構。

我們這裏不討論數學意義上的張量，只考慮程序中的實現。程序中的張量是一個多維數組的數據結構，示例如下：

#define MAX_DIM 6
struct Tensor {
    // 維度信息
    size_t dim[MAX_DIM];
    uint8_t num_dim;

    // 數據信息
    float* data;
    size_t num_data;
};

0維張量，就是一個數。1維張量等同於一個向量。2維張量對應一個矩陣。3維張量則是一個立方體。

1.2 張量計算

張量集到張量集的映射稱爲張量計算。用編程語言來說，輸入是若干張量，輸出也是若干個張量，並且無副作用（參考函數式編程）的函數稱之爲張量計算。
由示例程序可以看出，張量有 “維度” 和 “數據” 兩個組成要素，張量計算，也就包含維度與數據這兩個組成要素的處理。

比如矩陣乘法C = MatMul(A, B)，首先是根據輸入的兩個張量A, B確定C的維度，然後根據A和B的數據再去計算C的數據。具體一些可參考下面的代碼：

Tensor* MatMul(Tensor* A, Tensor* B) {
    Tensor* C = new Tensor;
    // 計算維度
    C->num_dim = 2;
    C->dim[0] = A->dim[0];
    C->dim[1] = B->dim[1];

    // 分配內存
    C->data = malloc(C->dim[0]*C->dim[1]*sizeof(float));

    // 計算數據
    Matrix::multi(C, A, B);
    return C;
}

1.3 計算圖

在深度學習領域，一個模型是由一系列的張量計算與常量組合而得的計算圖，每一次張量計算稱爲一個算子（Op）。
比如手寫數字識別的Mnist模型：

class Mnist:
    def __init__(self):
        self.w0, self.w1, self.w2, self.w3 = LoadWeight()
    def forward(self, x):
        x = Conv(x, self.w0)
        x = Pool(x)
        x = Conv(x, self.w1)
        x = Pool(x)
        x = InnerProduct(x, self.w2)
        x = Relu(x)
        x = InnerProduct(x, self.w3)
        x = Softmax(x)
        return x

它由 8個算子組成，分別是 Conv -> Pool -> Conv -> Pool -> InnerProduct->Relu->InnerProduct->Softmax

計算圖涵蓋了一系列的算子，最壞情況下，它的計算時間是各算子時間之合。但我們可以對計算圖進行依賴分析，在資源管理、計算調度、冗餘清除等方面做文章，使計算圖的計算時間小於這個總合。

2. 張量計算的特點

2.1 變化的運算量

每種算子都是對一批數據進行計算，其運算量不僅取決於算子本身，也取決於數據本身的維度信息。
以下是一些常見算子的運算量的分析

算子	參數	輸入維度	輸出維度	運算量
Conv	$c_o, k_h, k_w$	$(n, c_i, h,w)$	$(n, c_o, h,w)$	$nc_oc_ihwk_hk_w$
ConvDw	$k_h, k_w$	$(n, c, h,w)$	$(n, c, h,w)$	$nchwk_hk_w$
MatMul		$(h_i, w), (w, w_o)$	$(h_i, w_o)$	$h_iww_o$
MatAdd		$(h, w), (h, w)$	$(h, w)$	$hw$
MatSub		$(h, w), (h, w)$	$(h, w)$	$hw$
Resize	$f=2$, Bilinear	$(n, c, h, w)$	$(n, c, 2h, 2w)$	$16nchw$
Resize	$f=0.5$, Bilinear	$(n, c, h, w)$	$(n, c, 0.5h, 0.5w)$	$nchw$
Resize	$f=2$, Nearest	$(n, c, h, w)$	$(n, c, 2h, 2w)$	$4nchw$

由於這一性質，張量計算的耗時隨着輸入的變大而上升。我們在做性能分析時，不能一刀切地去說一個算子是不是耗時的，要具體地結合維度信息去判斷。比如Conv(卷積算子)一般情況下比較耗時，但如果它的輸入維度比較小，就很可能不如後面的一個Resize耗時多。

2.2 校驗困難與誤差容忍

張量計算的結果是一個龐大的數組，比如一個圖像算法，輸入輸出均是 400 x 300 像素的RGB圖像，那麼它的輸出數據量就是$400x300x3 = 360000$，這麼大的數據量我們很難人工去確認是否正確，有參考輸出時，可以用統計學的一些方法，編程序判斷。沒有參考輸出的情況，比如圖像濾鏡，只能人肉眼去看圖像。

大部分張量計算，即便是最簡單的矩陣相加，不同的計算順序，不同的硬件指令集也可能導致計算結果有非常小的偏差。在實際的應用中，我們往往沒有必要要求計算結果和標準輸出完全一致，如下面兩幅圖像，兩者像素值均存在一定偏差，但不仔細看基本分辨不出。

由於實際應用中對張量計算的結果有一定的誤差容忍，我們可以採用有一定誤差的優化方案，典型的就是定點計算，將原先浮點計算等效替換爲整型計算。

2.3 形實無關

對常用算子而言，維度的計算僅由輸入張量的維度與算子信息確定，因此可以先計算維度信息（“形”），再計算內容（“實”）。具體到一個張量計算圖來說，就是先計算圖中所有張量的維度，再去算各張量的內容。

一般而言，維度信息確定下來後，就能確定分配的內存大小和該算子所需要採用的最優計算流程，而很多情況下內容在不斷變化，而維度信息是相對穩定的，利用這一特性，我們可以做優化，減少運行時的耗時。

並非所有算子都有這個性質，比如unique算子需要從輸入張量中取互不相同的元素出來組成輸出張量，就無法單從輸入維度計算維度信息，對於這些情況需要特別處理。

2.4 "複雜"的優化

張量計算很多情況下需要適當添加計算步驟來加速，使得程序變得複雜而難懂。

2.4.1 普通程序的優化

對普通程序而言，性能優化是一個使代碼變"簡潔"的過程，刪除掉冗餘計算之後，性能肯定比之前好，代碼看上去也更舒服。
例如：

void lower1(char *s){
    long i;
    for(i = 0; i < strlen(s); i++){
        if(s[i] >= 'A' && s[i] <= 'Z'){
            s[i] -= ('A' - 'a');
        }
    }
}

strlen(s) 每次循環都計算，可以放到循環外面。優化後：

void lower1(char *s){
    long i;
    long sLength = strlen(s);
    for(i = 0; i < sLength; i++){
        if(s[i] >= 'A' && s[i] <= 'Z'){
            s[i] -= ('A' - 'a');
        }
    }
}

2.4.2 優化示例：兩種矩陣乘法

我們以兩種矩陣乘法爲例，來看張量計算的優化。

簡易版本V1

// 矩陣乘法 v1
void Matrix::multi(Tensor* C, const Tensor* A, const Tensor* B) {
    const auto a = A->host<float>();
    const auto b = B->host<float>();
    auto c       = C->host<float>();

    const int h = A->length(0);
    const int k = A->length(1);
    const int w = B->length(1);

    const int aw = A->stride(0);
    const int bw = B->stride(0);
    const int cw = C->stride(0);

    MNN_ASSERT(k == B->length(0));

    for (int y=0; y < h; ++y) {
        const auto aLine = a + y * aw;
        auto cLine       = c + y * cw;
        for (int x=0; x < w; ++x) {
            auto bColumn = b + x;
            float sum    = 0.0f;
            for (int i = 0; i < k; ++i) {
                sum += aLine[i] * bColumn[i * bw];
            }
            cLine[x] = sum;
        }
    }
}

複雜版本V2

具體代碼參見
https://github.com/alibaba/MNN/blob/master/source/backend/cpu/CPUMatMul.cpp

// 矩陣乘法 V2 （一部分）
    const Tensor* A = inputs[0];
    const Tensor* B = inputs[1];
    auto APtr = A->host<float>();
    auto BPtr = B->host<float>();
    Tensor* C       = outputs[0];
    auto CPtr = C->host<float>();
    auto w0         = inputs[0]->length(1);
    auto h0         = inputs[0]->length(0);
    mFunction.clear();
    auto e = C->length(0);
    auto h = C->length(1);
    auto l = w0;
    std::shared_ptr<Tensor> AT(Tensor::createDevice<float>({UP_DIV(l, 4), e, 4}));
    std::shared_ptr<Tensor> BT(Tensor::createDevice<float>({UP_DIV(h, 4), UP_DIV(l, 4), 16}));
    std::shared_ptr<Tensor> CT(Tensor::createDevice<float>({UP_DIV(h, 4), e, 4}));
    std::shared_ptr<Tensor> BTemp;
    if (l % 4 != 0) {
        BTemp.reset(Tensor::createDevice<float>({UP_DIV(h, 4), l, 4}));
        auto res = backend()->onAcquireBuffer(BTemp.get(), Backend::DYNAMIC);
        if (!res) {
            return OUT_OF_MEMORY;
        }
    }
    auto res = backend()->onAcquireBuffer(BT.get(), Backend::DYNAMIC);
    if (!res) {
        return OUT_OF_MEMORY;
    }
    auto BTPtr = BT->host<float>();
    float* BTempPtr = BTPtr;
    if(l % 4 != 0) {
        BTempPtr = BTemp->host<float>();
    }
    mFunction.emplace_back([BPtr, BTempPtr, l, h] {
        MNNTensorConvertNHWCToNC4HW4(BTempPtr, BPtr, l, h);
    });
    if (l % 4 != 0) {
        mFunction.emplace_back([BTPtr, BTempPtr, l, h] {
            auto hC4 = UP_DIV(h, 4);
            auto lC4 = UP_DIV(l, 4);
            for (int y=0; y<hC4; ++y) {
                auto dst = BTPtr + 16*lC4 * y;
                auto src = BTempPtr + 4 * l * y;
                ::memcpy(dst, src, 4*l*sizeof(float));
                ::memset(dst+4*l, 0, 4*(lC4*4-l) * sizeof(float));
            }
        });
        backend()->onReleaseBuffer(BTemp.get(), Backend::DYNAMIC);
    }
    res = backend()->onAcquireBuffer(AT.get(), Backend::DYNAMIC);
    res = res && backend()->onAcquireBuffer(CT.get(), Backend::DYNAMIC);
    if (!res) {
        return OUT_OF_MEMORY;
    }
    auto ATPtr = AT->host<float>();
    mFunction.emplace_back([ATPtr, APtr, e, l]() {
        MNNTensorConvertNHWCToNC4HW4(ATPtr, APtr, e, l);
    });
    std::shared_ptr<StrassenMatrixComputor> computor(new StrassenMatrixComputor(backend()));

    auto code = computor->onEncode({AT.get(), BT.get()}, {CT.get()});
    if (NO_ERROR != code) {
        return code;
    }
    auto CTPtr = CT->host<float>();
    mFunction.emplace_back([computor, CPtr, CTPtr, e, h]() {
        computor->onExecute();
        MNNTensorConvertNC4HW4ToNHWC(CPtr, CTPtr, e, h);
    });
    backend()->onReleaseBuffer(AT.get(), Backend::DYNAMIC);
    backend()->onReleaseBuffer(BT.get(), Backend::DYNAMIC);
    backend()->onReleaseBuffer(CT.get(), Backend::DYNAMIC);
    //限於篇幅，這裏只展示一部分代碼，實際還有另外一大半就不展示了。

初看上去，v2版本流程比v1複雜多了。但同樣單線程運行，v2版本大部分情況下速度是v1的10倍。

由於張量計算的優化是“複雜”的，針對單獨張量計算函數，我們不能以將像傳統程序優化一樣去找冗餘並清理，正確的方式請看下節。

3. 張量計算的優化

3.1 優化就是“修路”

以城市交通建設舉例說明一下張量計算優化的基本思路。
比如下圖，魏博下屬有魏州、相州、博州、貝州、澶州，范陽下屬有幽州、嬀州、易州、定州、恆州。現在要讓魏博的人民到范陽更方便，應該怎麼做呢？

我們可以每兩個州都修一條路，但這樣需要25條路，成本高昂。現實中的我們的解決方案一般是這樣的：
1、建一條高速路，起點站爲S，終點站爲D。
2、魏博各州與S連通。
3、范陽各州與D連通。

這樣做的好處在於：
1、保證各州居民均能較快地到達。
2、集資打造一條高速路就可以，成本較低。

缺點在於，有些相鄰比較近的州要繞遠路（如圖上的博州到易州）。

張量計算的優化就類似上面這個"修路"的過程，設計一系列高速計算模塊，然後原始的計算轉化爲高速計算模塊可解決的問題。

3.2 修高速——設計高性能計算模塊

高性能的來源第一部分是硬件層面，在開發者層面無法參與，但需要深入瞭解。

第二部分是軟件層面，主要分成三類策略：
第一類是相對固定的計算套路，比如彙編排列技巧、SIMD使用技巧等、GPU的調度技巧等等，這些套路需要舉一反三，應用到自己的代碼中。

第二類是內存訪問與併發設計，有較簡單的針對硬件特性的引用（比如GPU加速在高通GPU上用Image而不是用Buffer存儲），也需要反覆試驗調試，在並行度與緩存友好中取折中的Schedule過程。

第三類是張量計算優化算法，如 Winograd 卷積計算，Strassen 矩陣乘，這些算法限定於固定的某類張量計算，但可以結合第一、二類策略廣泛應用於不同的硬件上。

“修高速”是優化過程中最爲重要與困難的任務，極其考驗研發人員的智商與耐力。近些年來，多面體模型編譯技術得到了不少發展，並以 Halide / TVM 知名度最廣，這種技術可以自動地產生“高速路”，但性能與人們手工設計的仍然有差距，並且要達到較好的性能，人工且入的成本目前也不低。

3.3 連接高速——原始計算與高性能計算模塊之間的轉換

類似於高速路搭完之後，我們需要將居民點與高速連接，由於軟件層面前兩類策略的實施，高性能計算需要一些觸發條件，或者稱入口，隨着所使用的具體硬件而不同，比如：
（1）ARM的 SIMD 運算浮點矩陣乘法要求把矩陣重排爲一系列 (1, 4), (4, 4) 的小塊
（2）ARMv8.2 的 SDOT 指令計算要求把矩陣重排爲(4, 4), (4, 4)的小塊
（3）分塊後，需要申請若干緩存，以便逐塊複製數據與計算
（4）使用 GPU 加速，需要創建 Kernel / Buffer，並上傳數據到顯存中

在高性能計算完成後，其產生的數據往往也不便於用戶直接使用，需要作一些轉換，使之用戶可見；另外，在高性能計算過程產生的一些緩存也需要清理，這類似於一個下高速的過程，對應上文“上高速”的例子，分別的“下高速”操作爲：
（1）將重排成 (1, 4) 小塊的矩陣輔平
（2）將重排成 (4, 4) 小塊的矩陣輔平
（3）銷燬緩存
（4）將顯存中的數據複製出來，銷燬之前創建的Kernel/Buffer 對象。

3.4 優化方案實踐——矩陣乘法

這裏以上一節的矩陣乘法爲例說明一個完整的張量計算優化方案。回到上節矩陣乘法的V2版本，那段複雜的代碼是按如下流程編寫的：

首先把 A , B 分別轉到適宜計算的數據佈局AT和BT，然後執行該佈局下的矩陣計算，其中包含Strassen分解，分塊，彙編等具體優化手段，最後把計算完成的矩陣CT轉回原始的數據佈局C。
這樣我們付出 $O(el) + O(eh) + O(lh)$的佈局轉換代價，換來核心計算$O(elh)$的數倍性能提升。

可能有人會注意到，在一些情況下，比如$e, l, h$中有某個數爲1，那麼佈局轉換的代價就會超過運算的代價。因此，在必要的時候我們需要提供多種方案，在情況不同時採用。

4. 計算圖的優化

對於單個張量計算而言，“上高速”——“跑高速”——“下高速”是一個完整鏈路，“上高速”、“下高速”相對於“跑高速”而言，時間一般比較少，也沒有太多優化空間。而在深度學習模型的推理過程中，由於多個張量計算先後執行，就產生了與普通程序相似的冗餘計算，在“不關注中間結果”的前提下，可以進行優化以進一步提升。

A = Input();
B = Const();
C = MatMul(A, B);
D = Const();
E = MatMul(C, D);

比如上面代碼，矩陣乘之後接矩陣乘，且每個矩陣乘的輸入之一是常量。展開之後類似這樣：

A = Input();
B = Const();
AT = Convert(A);
BT = Convert(B)
CT = Strassen_MatMul(AT, BT);
C = Revert(CT);
CT = Convert(C);
D = Const();
DT = Convert(D);
ET = Strassen_MatMul(CT, DT);
E = Revert(ET);

很自然地我們會發現如下的冗餘計算：

1、B、D 是常量，BT、DT 可以預先計算
2、CT -> C 和 C->CT 兩步可以去掉

推理引擎除了提供一系列高效張量計算實現之外，冗餘計算的清除也是非常重要的，主要的策略是：
1、“車同軌”：各類張量計算儘量按照相似的內存佈局進行計算，減少切換成本。
2、“書同文”：將各類張量計算按一定的接口要求改造，便於調度模塊根據這些接口進行冗餘計算的清除。

具體到 MNN 中，表現爲如下幾個設計：

4.1 形實分離

MNN 要求所有算子實現的形狀計算與內容計算分離，形狀計算爲各種硬件實現共用，內容計算由各類硬件抽象實現。
如 Pooling 算子的計算，被拆解爲如下文件：
形狀計算：

source/shape/ShapePool.cpp

各種硬件下的內容計算

source/backend/cpu/CPUPool.hpp  // ARM / x86 等實現
source/backend/opencl/execution/PoolExecution.hpp // 基於 OpenCL 標準的實現 
source/backend/vulkan/execution/VulkanPool.hpp // 基於 Vulkan 標準實現
source/backend/opengl/GLPool.cpp // 基於 OpenGL ES 3.1 的實現

這樣做的好處是：
（1）減少異構計算支持的成本
（2）便於統一內存分配管理
（3）便於統一計算調度

4.2 NC4HW4佈局

基於ARM / GPU 上的4單元SIMD，及大部分圖像相關的算子可天然在通道並行的特性，MNN 對大部分CV相關算子採用NC4HW4 佈局計算，這個佈局設計可以在多數情況下減少內存佈局轉化的開銷。

4.3 Resize機制 / 預推理機制

爲了避免在推理中頻繁申請和釋放內存，減少異構計算中的冗餘，MNN設計時引入一個預推理過程（接口上爲 resize）。

這樣做，可以在支持動態形狀（即輸入的形狀可變的情況，允許運行時改變輸入形狀大小）的前提下達到如下目的：
（1）計算策略調度：根據輸入形狀決定一些算子的最優計算策略。
（2）進行內存管理：申請每個算子的輸入輸出Tensor內存與運算時所需的緩存，並按依賴關係複用中間算子的內存，這樣既在運算過程中無內存申請/釋放的損耗，也不會過多佔用系統內存。
（3）冗餘計算清除：在形狀確定的情況下，部分算子的輸出是固定的，如Priorbox，這些可以預先計算。另外對異構設備來說，如 Vulkan ，可以製作相關算子的命令緩衝（Command Buffer），填充參數等等，在執行過程中僅需提交 Command Buffer，將CPU-GPU的交互降到最低。

4.4 Fuse （算子融合）

在模型轉換階段，也即離線將一些算子合併或消除，如 scale 與 convolution 合併，relu 與 convolution 合併，這部分依賴於對具體算子的分析制定專家規則。

對於常用算子而言，MNN 在冗餘計算的清除上已經做得比較好，但對於更復雜的算子處理上，仍有較多優化空間。業界也有基於編譯技術進行冗餘計算清除的，比如XLA(Tensorflow)，TVM，但目前主要侷限於Fuse，且還是靠專業經驗去堆，不見得有更簡單的做法。此外除了冗餘計算清除外，業界也有不少研究異構調度，將計算拆解後分配到不同硬件上，由於未有成熟方案，不多敘述。

5. 展望

關於張量計算領域的未來發展，目前主要關注點還在於深度學習推理引擎的優化，像軟硬協同、編譯優化、模型壓縮被經常提及。這裏我想談一些不同的視角：
（1）近些年來高性能計算的硬件層出不窮，華爲、高通、蘋果、谷歌、MTK都有自研的NPU，大家都支持一些核心的Op，但支持粒度參差不齊，也沒幾家願意開放指令集。這種碎片化的現狀目前來看會長期存在，承認碎片化的現狀，梳理各類張量計算的邏輯，使它們能基於幾個相對固定的被硬件所支持的高性能計算模塊去實現，對接各個廠商去實現這少量模塊，相對於統一標準、基於編譯技術實現跨硬件，目前來看是更現實更可行的做法。
（2）由於歷史原因，在圖像與機器學習領域，張量被定義爲多維數組，而這種定義在處理形變時顯得繁瑣，我們需要大量的算子去支持縮放、平移、裁剪功能，是否有更好的定義與計算方式，值得進一步探討。
（3）除了深度學習以外，還有更多應用張量計算的場景，如圖形圖像、科學計算等，但這些領域往往是使用單獨的張量計算庫，少見像深度學習框架那樣的清除冗餘的機制，也沒有求導的能力。未來深度學習框架與這些領域打通，可以賦予它們學習能力，也可以進一步優化性能，統一調度硬件。Tensorflow 已經推出了 Tensorflow Graphics 以做嘗試，但後效如何還需觀察。

參考鏈接：
http://www.jos.org.cn/html/2018/8/5563.htm
https://halide-lang.org/docs/index.html
https://github.com/alibaba/MNN