在 CFU Playground 上加速 MLPerf™ Tiny 影像分類模型 #6：軟硬體結合加速

接下來我們將會開始結合前面提到的乘法加速單元以及軟體加速技巧來最大化模型推理的效率。在筆者的加速器設計中包含五條與矩陣乘法單元有關的指令，以及三條跟 TFLite int8 量化處理有關的指令，我們將會看到筆者在設計這些指令時考慮到軟硬體複雜度而做的取捨。

這裡假設讀者已經讀過第四章提到的所有內容，如果還沒看過的話建議先回去大略翻過一遍，不然下面的內容將會很難看懂。

6.1 資料傳輸

與矩陣乘法相關的指令共有五條，三條指令跟資料傳輸有關，一條指令用來告知矩陣加速單元要執行的矩陣乘法大小，一條用來檢查矩陣乘法是否結束。

我們先從最直覺的方法開始來談談資料傳輸，考慮矩陣 A，由於要寫入暫存區一共需要兩個值，一是欲寫入的位址，二是欲寫入的值，回憶第一章提到指令格式為 cfu_op(func7, func3, a, b)，我們可以設計出一條指令 cfu_op(0, 0, addr, val) 用來傳輸資料，如果說我們的脈動陣列大小為 n×n，那麼暫存區的一個位址會儲存 n 個元素，總共需要 ceil(n/4) 次指令才能完整傳輸完資料，原因為 A 和 B 的元素為 int8，而 val 的大小為 32 位元，然而若 n 大於 4，我們需要想辦法將當前元素寫到 addr 內的特定位置，需要額外的硬體複雜度，並且這就表示 addr 需要加上關於定位的訊息，這就有些本末倒置了，所以說對於這樣的設計而言，令 n 為 4 最為適合，恰好能完美利用 32 位元的傳輸頻寬。

不過 addr 的大小也同樣是 32 位元，不利用起來的話有些可惜，但如果要將原先 addr 也用來傳資料，那麼位址就必然要由硬體計算，這裡我們可以讓軟體按位址自小到大將資料傳輸到硬體裡，那麼只要在硬體上維護一個每個週期會自增的計數器，其值就能拿來當作要寫入的位址，如此一來傳輸頻寬就由 32 位元翻倍為 64 位元，能夠一次傳送八個元素，此時令 n 為 8 最為適合。

舉例來說，若新指令為 cfu_op(0, 0, uval, lval)，第一次用該指令傳值時，硬體計數器的值應該為 0，以便將 uval 和 uval 串接起來存入暫存區的位址 0；而第二次用該指令傳值時，硬體計數器的值應該為 1，以便將 uval 和 lval 串接起來存入暫存區的位址 1，依此類推直到資料傳輸完畢。

這裡給出用來將兩個矩陣寫入暫存區的指令，這兩條指令的硬體實作邏輯相同，唯一差別為寫入的暫存區不同。

#define set_matrixA(uval, lval)      cfu_op0(0, uval, lval)
#define set_matrixB(uval, lval)      cfu_op0(1, uval, lval)

現在來看看具體該怎麼使用指令來傳送資料，以 im2col 矩陣舉例，卷積核矩陣道理相同，假設現在要做的是 mαnβkγ 的矩陣乘法，那麼顯然 im2col 矩陣的大小為 γ×α（別忘了 A 被轉置了），根據前面討論脈動陣列時提到的暫存區存放方法，可以得到傳輸資料的程式碼如下，此處有兩點需要注意，一是元素存放方式，二是越界補零的判斷（這裡的越界跟圖片讀取越界無關，而是跟暫存區存放方法有關）。如果像筆者用陣列方式處理資料的話，因為 RISC-V 使用小端序（little endian），要將位址最小的值放在陣列的最後一個，這樣在轉型為 int32 後最小的 int8 元素才會被放在最高位；如果不用陣列的話，也能透過位元運算來組合出正確的存放順序，兩種方法孰快孰慢留待讀者測試，筆者手上已無可用的開發板。

// Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
// Code Snippet for Transmitting Data

for (int mm = 0; mm < α; mm += 8) {
  for (int kk = 0; kk < γ; kk++) {
    int8_t lower[] = {
      mm + 7 < α ? im2col[kk*α + (mm + 7)] : (int8_t)0,
      mm + 6 < α ? im2col[kk*α + (mm + 6)] : (int8_t)0,
      mm + 5 < α ? im2col[kk*α + (mm + 5)] : (int8_t)0,
      mm + 4 < α ? im2col[kk*α + (mm + 4)] : (int8_t)0
    };

    int8_t upper[] = {
      mm + 3 < α ? im2col[kk*α + (mm + 3)] : (int8_t)0,
      mm + 2 < α ? im2col[kk*α + (mm + 2)] : (int8_t)0,
      mm + 1 < α ? im2col[kk*α + (mm + 1)] : (int8_t)0,
      im2col[kk*α + (mm + 0)] : (int8_t)0
    };

    set_matrixA(*(int32_t*)&upper, *(int32_t*)&lower);
  }
}

將兩個矩陣傳送到硬體後，需要通過某種方法通知硬體開始運算，筆者利用一條指令將當前的矩陣運算大小以及偏移量傳入硬體，當接收到該指令時，硬體旋即開始計算，那麼要怎麼知道計算完成了沒呢？筆者採用非同步的方式來讓矩陣運算單元忙碌時能夠運行其他 CPU 運算，最後再用一條指令檢查運算單元是否完成運算，不過根據筆者最後並沒有利用到非同步的特性，因此理論上來說同步的方式一定不會比非同步還慢。

#define start_GEMM(M, N, K, off)     cfu_op0(2, M<<20 | N<<10 | K, off)
#define check_GEMM()                 cfu_op0(3, 0, 0)

運算完成後，我們需要將得到的矩陣讀回記憶體內儲存起來，然而這裡的問題是一條指令能傳輸 64 位元的資料給 CFU，而 CFU 只能傳輸 32 位元的資料回來，頻寬只有輸入的一半，更慘的是計算後得到的矩陣元素為 32 位元，也就是說一次只能回傳一個元素，這樣就有些麻煩了。

如果一次只能回傳一個元素，下個問題是該如何把正確地將資料傳回並儲存於記憶體。既然矩陣有兩個維度，而且一條指令剛好能傳入矩陣的行列之值，那麼將要取得的矩陣位置傳入，讓硬體將指定位置的值從暫存區中取出再回傳即可，換而言之，這樣的設計使得我們能夠隨機存取結果矩陣的元素。

然而暫存區的存取需要時間，按照前面章節提到的方法，暫存區的存取會在行為發生後的一個周期內完成，也就是說每次收到指令時，硬體需要計算出欲存取元素所在的暫存區位址，花一個週期取得該位址儲存的資料後再取值回傳，從速度和硬體複雜度來說不是太好的設計，況且我們必然需要存取所有元素才能將整個矩陣儲存回記憶體，因此隨機存取矩陣的必要性不大。

筆者採用硬體按順序回傳元素的方式來降低電路複雜程度，道理跟前面傳資料差不多，都會用到一個計數器，而在軟體層面需要保證程式碼將取出元素寫回記憶體的索引值正確。

#define get_matrixC()                cfu_op0(4, 0, 0)

設暫存區一個位址能儲存八個元素，那麼第一次 get_matrixC() 得到的元素會是 C₀₀，第二次會是 C₀₁，……，第八次會是 C₀₇，第九次會是 C₁₀，依此類推。

// Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
// Code Snippet for Reading Data

inline int get_buffer_rows(const int M, const int N) {
  return M * ((N + 7) >> 3);
}

for (int r = 0; r < get_buffer_rows(α, β); r++) {
  const int mm = r % α;
  const int nn = r / α * 8;

  for (int i = 0; i < 8; i++) {
    if (nn + i < β) {
      result[mm*β + (nn + i)] = get_matrixC();
    }
  }
}

以上我們看到如何將矩陣傳入和讀回記憶體中，接下來要將上面的邏輯融合進整個 m256n256k256 的矩陣乘法了，在筆者的設計中，脈動陣列的大小為 16×16、暫存區一個位址的寬度為 16 個元素且每個矩陣的暫存區分別最大能容納 256×256 = 65536 個元素（即一個暫存區有 4096 個位址），由於一次最多僅能寫入八個元素，硬體會將上半部分的八個元素暫存起來，等到得到下半部分的八個元素後，才憶起寫入暫存區對應的位址中。另外提一下其實不一定要把整個矩陣分塊成 m256n256k256 的矩陣乘法，僅須保證兩個輸入矩陣的元素個別不超過 65536 個元素即可，例如分塊成 m128n128k512 的矩陣乘法也行，不過程式碼要寫好就是了。

那為什麼要選 16×16 而不是更大的 32×32 或更小的 8×8 呢？一是 16×16 是筆者測試過能燒入開發版的最大尺寸，二來是模型中的所有矩陣乘法的 m 和 n 之值的最大公因數為 16，因此傳送資料時可以省去越界補零的判斷，能省下不少時間。

// Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
// Tiled Matrix Multiplication with Matrix Multiplication Unit

constexpr int T = 256;

inline int get_buffer_rows(int M, int N) {
  return M * ((N + 15) >> 4);
}

const int pad_M = M + (T - M&(T - 1));
const int pad_N = N + (T - N&(T - 1));
const int pad_K = K + (T - K&(T - 1));

// Construct im2col matrix (Chap. 4)

// Construct kernel matrix (Chap. 4)

for (int m = 0; m < pad_M; m += T) {
  const int MM = std::min(m + T, M) - m;

  for (int n = 0; n < pad_N; n += T) {
    const int NN = std::min(n + T, N) - n;

    for (int k = 0; k < pad_K; k += T) {
      const int KK = std::min(k + T, K) - k;

      for (int mm = 0; mm < MM; mm += 16) {
        for (int kk = 0; kk < KK; kk++) {
          int8_t lowerA[] = {
            im2col[(k + kk)*MM + (m + mm + 7)],
            im2col[(k + kk)*MM + (m + mm + 6)],
            im2col[(k + kk)*MM + (m + mm + 5)],
            im2col[(k + kk)*MM + (m + mm + 4)]
          };

          int8_t upperA[] = {
            im2col[(k + kk)*MM + (m + mm + 3)],
            im2col[(k + kk)*MM + (m + mm + 2)],
            im2col[(k + kk)*MM + (m + mm + 1)],
            im2col[(k + kk)*MM + (m + mm)]
          };

          int8_t lowerB[] = {
            im2col[(k + kk)*MM + (m + mm + 15)],
            im2col[(k + kk)*MM + (m + mm + 14)],
            im2col[(k + kk)*MM + (m + mm + 13)],
            im2col[(k + kk)*MM + (m + mm + 12)]
          };
 
          int8_t upperB[] = {
            im2col[(k + kk)*MM + (m + mm + 11)],
            im2col[(k + kk)*MM + (m + mm + 10)],
            im2col[(k + kk)*MM + (m + mm + 9)],
            im2col[(k + kk)*MM + (m + mm + 8)]
          };

          set_matrixA(*(int32_t*)&upperA, *(int32_t*)&lowerA);
          set_matrixA(*(int32_t*)&upperB, *(int32_t*)&lowerB);
        }
      }

      for (int nn = 0; nn < NN; nn += 16) {
        for (int kk = 0; kk < KK; kk++) {
          int8_t lowerA[] = {
            kernel[(k + kk)*NN + (n + nn + 7)],
            kernel[(k + kk)*NN + (n + nn + 6)],
            kernel[(k + kk)*NN + (n + nn + 5)],
            kernel[(k + kk)*NN + (n + nn + 4)]
          };

          int8_t upperA[] = {
            kernel[(k + kk)*NN + (n + nn + 3)],
            kernel[(k + kk)*NN + (n + nn + 2)],
            kernel[(k + kk)*NN + (n + nn + 1)],
            kernel[(k + kk)*NN + (n + nn)]
          };

          int8_t lowerB[] = {
            kernel[(k + kk)*NN + (n + nn + 15)],
            kernel[(k + kk)*NN + (n + nn + 14)],
            kernel[(k + kk)*NN + (n + nn + 13)],
            kernel[(k + kk)*NN + (n + nn + 12)]
          };

          int8_t upperB[] = {
            kernel[(k + kk)*NN + (n + nn + 11)],
            kernel[(k + kk)*NN + (n + nn + 10)],
            kernel[(k + kk)*NN + (n + nn + 9)],
            kernel[(k + kk)*NN + (n + nn + 8)]
          };

          set_matrixB(*(int32_t*)&upperA, *(int32_t*)&lowerA);
          set_matrixB(*(int32_t*)&upperB, *(int32_t*)&lowerB);
        }
      }

      start_GEMM(MM, NN, KK, 128);
      while (check_GEMM());

      for (int r = 0; r < get_buffer_rows(MM, NN); r++) {
        const int mm = r % MM;
        const int nn = r / MM * 16;

        for (int i = 0; i < 16; i++) {
          // result is zero-initialized
          result[mm*N + (nn + i)] += get_matrixC();
        }
      }
    }
  }
}

// Write back to output_data (Chap. 4)

使用矩陣加速單元：678492.32 微秒，約為 2.105 倍加速

6.2 暫存區累加設計

在上面的程式碼中可以看到結果的累加發生在硬體加速單元之外，但這件事應該能在加速單元內完成，因為值都儲存在暫存區內對吧？

這裡筆者採用特殊設計過的暫存區來儲存矩陣 C，其行為為當要寫入資料時，寫入到當前位址的前一個位址，並且總是輸出當前位址儲存的資料，如此一來我們能夠在當前週期取得欲寫入位置之值，同時預取下次要寫入位址之值。

#define start_GEMM_acc(M, N, K, off)  \
  cfu_op0(2, 1<<30 | M<<20 | N<<10 | K, off)

接下來把前面的程式碼稍微修改一下，可以看到我們將 m×n×k 次的 get_matrixC() 變成只需 m×n 次，多少能省去一些時間。

// Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
// Tiled Matrix Multiplication with Matrix Multiplication Unit
// and Inplace Accumulation

// ...

for (int m = 0; m < pad_M; m += T) {
  const int MM = std::min(m + T, M) - m;

  for (int n = 0; n < pad_N; n += T) {
    const int NN = std::min(n + T, N) - n;

    for (int k = 0; k < pad_K; k += T) {
      const int KK = std::min(k + T, K) - k;

      // ...

      if (k == 0) cfu_start_GEMM(MM, NN, KK, 128);
      else cfu_start_GEMM_acc(MM, NN, KK, 128);
      while (cfu_check_GEMM());
    }
  }

  for (int r = 0; r < get_buffer_rows(MM, NN); r++) {
    const int mm = r % MM;
    const int nn = r / MM * 16;

    for (int i = 0; i < 16; i++) {
      result[mm*N + (nn + i)] = get_matrixC();
    }
  }
}

// ...

暫存區累加：669268.155 微秒，約為 1.013 倍加速

6.3 隱式 im2col

或許讀者看完上一小節發現 result[mm*N + (nn + i)] = get_matrixC(); 有些多餘，可以直接將 get_matrixC(); 讀出來後直接後處理儲存回 output_data這樣的觀察十分正確，並且接下來不僅要去除掉 result，連 im2col 和 kernel 矩陣我們都不打算用迴圈顯式地構造出來，而是利用位址翻譯的形式達到同樣效果，如此一來能夠節省記憶體和時間，這個概念在 NVIDIA 的 cuDNN 和 CUTLASS 中便有使用，有興趣的讀者可以用 implicit GEMM 當作關鍵詞搜尋。

還記得筆者在第二章的最後用粗體表示「如果要構造出 im2col 矩陣」嗎？現在要來解決這個懸念。提醒一下矩陣乘法是用軟體實作的話，因為 CFU Playground 用的是 VexRiscv，其快取設計相較現代處理器低效，所以隱式 im2col 會比較慢一些。

以較難的 im2col 矩陣為例，回憶第二章的 matrix_im2col[i][j] = data[y_in][x_in][ch_in]（請注意這裡並未轉置），當時我們令 i = y_out*W_out + x_out、j = ch_in*H_k*W_k + y_k*W_k + x_k、y_in = y_out*H_s + y_k*H_e - H_p 和 x_in = x_out*W_s + x_k*W_e - W_p，但在存取 im2col 矩陣時我們只有 i 和 j，沒有 y_in、x_in 和 ch_in，該怎麼辦呢？

其實也不難，因為剩下的值全部已知，按照式子從 i 和 j 分別反推出 y_out、x_out、y_k、x_k 和 ch_in 即可，如此一來便能得到 data[y_in][x_in][ch_in]。

// Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
// Implicit Im2col (Not Transposed)

auto im2col = [&](const int i, const int j) {
  const int y_out = i / W_out;
  const int x_out = i % W_out;
  const int ch_in = j / (H_k * W_k);
  const int kernel_no = j % (H_k * W_k);
  const int y_k = kernel_no / W_k;
  const int x_k = kernel_no % W_k;
  const int y_in = y_out*H_s + y_k*H_e - H_p;
  const int x_in = x_out*W_s + x_k*H_e - W_p;

  const bool is_inside = (uint32_t)y_in < (uint32_t)H_in
                       && (uint32_t)x_in < (uint32_t)W_in;

  return is_inside ? data[y_in][x_in][ch_in] : (uint8_t)-128;
};

從上面可以看到我們相當於是在 data 上多加了一層抽象，而不需要構造出 im2col 矩陣，不過要注意實際使用時 im2col 矩陣為轉置的狀態，因此要將上面程式碼中的 i 和 j 的角色互換，此外最後的 (uint8_t)-128 是為了抵銷模型的輸入偏移量，如此一來在加速單元運算時越界的值才會為 0。

總而言之，應用隱式 im2col 後程式碼看起來如下，筆者利用前一章參數帶入的技巧特化了三種不同的 im2ocl 版本，可以再減少一些運算。

// Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
// Tiled Matrix Multiplication with Matrix Multiplication Unit,
// Inplace Accumulation, and Implicit Im2col

// ...

for (int m = 0; m < pad_M; m += T) {
  const int MM = std::min(m + T, M) - m;

  for (int n = 0; n < pad_N; n += T) {
    const int NN = std::min(n + T, N) - n;

    for (int k = 0; k < pad_K; k += T) {
      const int KK = std::min(k + T, K) - k;

      for (int mm = 0; mm < MM; mm += 16) {
        for (int kk = 0; kk < KK; kk++) {
          int8_t lowerA[] = {
            im2col(k + kk, m + mm + 7),  // assume transposed
            im2col(k + kk, m + mm + 6),
            im2col(k + kk, m + mm + 5),
            im2col(k + kk, m + mm + 4)
          };

          // ...
        }
      }

      for (int nn = 0; nn < NN; nn += 16) {
        for (int kk = 0; kk < KK; kk++) {
          int8_t lowerA[] = {
            kernel(k + kk, m + mm + 7),
            kernel(k + kk, m + mm + 6),
            kernel(k + kk, m + mm + 5),
            kernel(k + kk, m + mm + 4)
          };

          // ...
        }
      }

      if (k == 0) cfu_start_GEMM(MM, NN, KK, 128);
      else cfu_start_GEMM_acc(MM, NN, KK, 128);
      while (cfu_check_GEMM());
    }
  }

  for (int r = 0; r < get_buffer_rows(MM, NN); r++) {
    const int mm = r % MM;
    const int nn = r / MM * 16;

    for (int i = 0; i < 16; i++) {
      int32_t acc = get_matrixC();

      // post-processing logic
    }
  }
}

// ...

隱式 im2col：440558.235 微秒，約為 1.519 倍加速

6.4 資料後處理電路

這裡的話就只是單純把上一章 gemmlowp 的 RoundingDivideByPOT 和 SaturatingRoundingDoublingHighMul從軟體運算轉換成硬體運算而已，不過因為 SaturatingRoundingDoublingHighMul 沒辦法在一個周期內完成，因此筆者將其拆成兩條同步指令運算，當然讀者也能利用非同步的方式用一條指令完成，另外筆者還將 acc += output_offset; acc = max(acc, -128); acc = min(acc, 127); 改成一條指令完成，總之運用硬體計算部分 int8 量化後處理邏輯後，可以得到加速。

#define post_set_SRDHM(a, b)         cfu_op0(5, a, b)
#define post_get_SRDHM()             cfu_op0(6, 0, 0)
#define post_RDBPOT(x, exp)          cfu_op0(7, x, exp)
#define post_off_maxmin(val, off)    cfu_op0(8, val, off)

後處理電路：407715.35 微秒，約為 1.080 倍加速

6.5 向量乘積累加單元

筆者原先實作了第一章提到的 CFU Playground 教學中的向量乘積累加單元以及其改進，但是將其運用在 FullyConnected 後經過實測不會比較快，反倒是把裡面的迴圈展開八次會得到額外加速，理論上我們應該能用矩陣乘法加速單元來計算 FullyConnected，但為什麼筆者沒有這麼做呢？答案就留待讀者思考了。

只迴圈展開八次：406357.93 微秒，約為 1.003 倍加速

6.6 結語

從第一章開始到第六章結束，我們看到了卷積層的計算和如何構造出等價的矩陣乘法，並在此基礎上討論了矩陣乘法的軟體加速策略，進而引伸出脈動陣列以及其針對矩陣乘法運算的設計，最終利用對模型計算的理解，從軟硬體層面上協同加速得到最終約十倍的加速比。

若讀者從頭開始閱讀本文，應不難發現其篇幅之長需要大量時間支撐才能完成，實際上本文包含中英文至少有三萬字，是筆者前後花了兩個月約共三百小時的成果，希望能幫助到有需要的讀者，誠然不論對讀者是否有幫助，筆者可以肯定在未來相當長的一段時間內自己忘不掉這一系列的相關知識了，現在大概能夠理解為何會有小說寫到最後斷更或爛尾。

由於筆者帶著知識的詛咒撰寫文章，或許有些部分筆者認為已經充分論述，但讀者認為仍不夠詳盡，若是如此，歡迎在留言區詢問，雖然到時候筆者還記不記得就不好說了。