JOVANA
Explore Library Glossary Getting Started Three Levels Fields How it works Mission
Join the mission
All guides

向量化與 BLAS 的威力

你的處理器能在名義上做一次乘法的時間裡做完八次乘法,而且早有一套調校好的函式庫,把這個事實榨到一滴不剩。我們會看看 SIMD 向量化與三個層級的 BLAS,如何把前兩篇那場「受記憶體限制」的惡夢,變成矩陣乘矩陣那種「受運算限制」的美夢。

一道指令,許多數字

這一級的第一篇給了你一句口號——浮點運算很便宜,記憶體很慢——第二篇則讓你看到記憶體階層與資料局部性如何決定你的資料是已經在快取裡待命,還是正從遙遠的 DRAM 一點一滴滲進來。本篇要談的是硬幣的另一面:就算資料「已經」準備好了,一個天真的迴圈仍會浪費掉晶片所能提供的大部分算力。現代處理器不是一次只加一對數字,而是一次就加一整條短向量。

這就是 SIMD——單指令多資料(Single Instruction, Multiple Data)。一個 SIMD 暫存器很寬:256 位元裝得下四個雙精度數字,512 位元裝得下八個。一道 SIMD 加法指令吃進兩個這樣的暫存器,一口氣產出四個(或八個)和,成本大約只跟一次純量加法相當。把你的運算安排成能餵養這些寬指令,這個技術就叫向量化,也就是你在這一級接下來會仰賴的SIMD 向量化。它的承諾很直接:一個完美向量化的迴圈,在同一顆核心上、不用額外執行緒或硬體,就能比純量版本快上四到八倍。

為什麼你手寫的迴圈把速度留在桌上

如果向量化是這麼白吃的好處,為什麼不乾脆信任編譯器?很多時候你確實可以——一個把兩個陣列相加的乾淨迴圈,會愉快地自動向量化。但編譯器保守,許多迴圈會抗拒它。如果第 n 次迭代依賴第 n-1 次(一個遞迴關係),晶片就無法一次做四個,因為它還不知道輸入。如果你的資料在記憶體裡是散落的而非連續的,那道用來填滿 SIMD 暫存器的寬載入,就得從四個遙遠的位址去蒐集,這很慢。如果兩個指標可能偷偷指向同一塊記憶體,編譯器就必須假設最壞情況而拒絕向量化。這每一項,都會悄悄把你打回純量的速度。

就算一個迴圈「確實」向量化了,純粹的 SIMD 也只攻擊帳本上浮點運算那一邊。回想第一篇的運算密度——所執行的浮點運算數,與從記憶體搬動的位元組數之比。向量化讓浮點運算發生得更快,卻完全沒讓位元組搬得更快。在屋頂線模型上,加快運算只是把你向右滑向那條斜的記憶體天花板;如果你的核心受記憶體限制,你就撞上那道天花板,而 SIMD 單元只能閒置,因缺資料而挨餓。所以光靠向量化並不是答案。真正的大獎,是把運算重新結構化以「提高」運算密度——讓每個位元組做許多浮點運算——而這正是 BLAS 變出來的把戲。

BLAS:三個運算密度的層級

BLAS——基礎線性代數副程式(Basic Linear Algebra Subprograms)——是幾乎所有數值軟體底層所說的標準化詞彙。它不是一套函式庫,而是一個介面,底下有許多手工調校的實作(OpenBLAS、Intel MKL、Apple Accelerate、BLIS)競相要在每顆晶片上跑得最快。關鍵的設計理念是:這些常式分成三個層級,而層級的編號告訴你它的運算密度。這些層級之所以存在,整個理由就是讓你能伸手去拿那個能解決你問題、又最省資料的運算。

level   example          data      flops      intensity (flops/byte)
-----   ---------------   -------   --------   ----------------------
BLAS-1  y = a*x + y       O(n)      O(n)       ~ 1/2     (memory-bound)
BLAS-2  y = A*x + y       O(n^2)    O(n^2)     ~ 1/4     (memory-bound)
BLAS-3  C = A*B + C       O(n^2)    O(n^3)     ~ n/8     (compute-bound!)

The win: BLAS-3 reuses each loaded number O(n) times.
A 1000x1000 matrix multiply does ~10^9 flops but reads
only ~10^6 numbers -- ~1000 flops per number moved.
三個 BLAS 層級與它們的運算密度。只有 BLAS-3(矩陣乘矩陣)能掙脫記憶體天花板,因為浮點運算量像 n^3 成長,而資料量只像 n^2 成長。

慢慢讀這張表,因為它是高效能數值計算的核心。第一級是向量對向量:縮放一個向量、把兩個向量相加、取一個內積。它碰觸 O(n) 的資料、做 O(n) 次浮點運算,所以每載入一個數字大約只做一次運算——徹底地受記憶體限制。第二級是矩陣對向量:它讀進一整個 n 乘 n 的矩陣,O(n^2) 個數字,去做 O(n^2) 次運算,所以每個矩陣元素只被用一次。仍然受記憶體限制。第三級是矩陣對矩陣:它讀 O(n^2) 個數字,卻執行 O(n^3) 次運算,所以每個抵達快取的數字,在被逐出之前大約會被重用 n 次。

這個重用就是一切。矩陣乘矩陣的乘積 C = A*B + C——傳統上稱為 GEMM 的那道常式——是唯一一個浮點運算遠多於位元組的核心,所以一個聰明的實作能讓 SIMD 單元持續有料可吃,真正逼近晶片的尖峰速度。它靠的正是你在第二篇見過的分塊與鋪磚:把矩陣切成能塞進快取的小塊,逐塊相乘,並趁每塊還熱的時候多次重用它。GEMM 正是向量化與資料局部性終於合作、而非互鬥的地方。

把你的問題改寫成 BLAS-3

這裡有一個會重塑「好的數值程式怎麼寫」的實用課題:只要能做到,就把你的演算法表達成讓大部分工作落進 BLAS-3 的矩陣乘矩陣呼叫。這就是為什麼現代稠密線性代數是「分塊的」。你最初學的那種樸素高斯消去法,是一片純量更新的汪洋——本質上是 BLAS-1 與 BLAS-2——以受記憶體限制的速度運行。同一個分解,若重新組織成一次更新整個子區塊,就能把絕大部分運算變成 GEMM 呼叫,跑得快上許多倍。拿你用來解 A x = b 的 LU 分解 A = L U 來說:教科書版本很慢,但分塊改寫後幾乎全是 GEMM。

  1. 從同一個 LU 分解 A = L U 出發。教科書版本一次只更新一列:運算密度低,受記憶體限制。
  2. 把 A 切成一格格的區塊。先用那段慢的、偏純量的程式碼分解一條細長的「直欄面板」,但只在一「小」片上做。
  3. 用那塊已分解的面板,以一次大型的 C = C - A*B 矩陣乘矩陣(一個 BLAS-3 的 GEMM),更新矩陣剩下的整個「尾隨區塊」。
  4. 沿著對角線重複。慢的面板工作總共是 O(n^2);快的 GEMM 工作是 O(n^3)——所以最昂貴的那部分以受運算限制的速度運行,整個求解便飛了起來。

這個分塊設計,正是 LAPACK 所做的事。LAPACK 是建構在 BLAS 「之上」的較高層函式庫——求解器、各種分解、特徵值與最小平方常式——而它刻意被寫成讓繁重的工作向下流進 BLAS-3 呼叫。當你在 NumPy、MATLAB、Julia 或 R 裡呼叫一個求解器,你幾乎總是在呼叫 LAPACK,而它底下又在呼叫調校過的 BLAS。整個 BLAS 與 LAPACK 堆疊,正是為什麼在高階語言裡一行求解,能跑出你硬體尖峰速度的一大部分。

別自己造輪子——以及誠實的小字條款

這篇指南能留給你的最有價值的習慣:別自己造輪子去寫矩陣核心。一個調校過的 GEMM 代表了好幾個人年的工作——配合 L1、L2、L3 確切大小的快取分塊、暫存器鋪磚、SIMD 內建函式、預取提示,有時甚至是手寫組合語言——全都是為了追求一個你那誠實的三層巢狀迴圈永遠看不到的速度。你照定義寫出來的三層迴圈,通常只跑出尖峰的百分之幾;函式庫版本則跑到近九成。這個別自己造輪子原則不是懶惰——它是正確的工程判斷:這個問題已經被解決了,而且解得遠比你一個下午能做到的好。

現在說誠實的告誡,因為 BLAS 很強大,但不是魔法。第一,改寫成 BLAS-3 並不改變你「算什麼」,只改變算得多快——而那個「什麼」仍是浮點。不同的 BLAS 實作可能以不同的順序去加總一個內積,而既然浮點加法不滿足結合律,兩個都正確的函式庫,給出的答案可能在最後幾個位元上有差。這是正常的,不是 bug;就是你最初學到「0.1 沒有精確的二進位形式」時遇過的那同一種非結合性。第二,這一切榮耀都是給「稠密」矩陣的。一個大多為零的稀疏矩陣,無論你怎麼分塊,運算密度都很低——根本沒有足夠的浮點運算去攤平資料的搬動——所以稀疏問題住得離記憶體天花板更近,需要其他級數裡的無矩陣法與迭代法。

帶上階梯的東西

把線索收攏。SIMD 向量化靠著在一顆核心裡「一道指令處理許多數字」贏得四到八倍,但它只加速運算——光靠自己打不破記憶體天花板。出路是提高運算密度,而 BLAS 遞給你一座三層的梯子來做這件事,其中 BLAS-3 的矩陣乘矩陣是唯一能跑近尖峰的受運算限制運算。把你的演算法寫成讓繁重工作流進 BLAS-3 與 LAPACK,並仰賴別人花了多年調校的函式庫。

攀登時請守住一個誠實的視角。這裡的每一次加速,只改變你「多快」得到答案,從不改變它對不對——結果仍是浮點、仍是近似,而對一個病態求解跑一次快的 BLAS,不過是更快地抵達一個糟糕的答案。速度與準確度是兩條獨立的軸。在接下來的指南裡,這同一套「向量化並重用」的思維會向外擴展:用執行緒與訊息傳遞跨越許多核心與機器,以及進到 GPU——在那裡,寬向量化就是整套硬體哲學,而阿姆達爾定律則為這一切平行所能換得的東西設下天花板。