基于SIMD 的并行傅里葉空間圖像相似度計算

郭渝洛，邊浩東，董潤婷，唐嘉豪，王曉英，黃建強

（青海大學 計算機技術與應用系，西寧 810016）

0 概述

二十世紀七十年代，TAYLOR 和GLAESER 研發了冷凍電鏡（cryo-Electron Microscopy，cryo-EM）。該項技術發展至今，已經成為研究生物大分子結構與功能強有力的手段，其簡化了生物細胞的成像過程，提高了成像質量，將生物化學帶入一個新時代［1-3］。

冷凍電鏡是在低溫下使用透射電子顯微鏡觀察樣品的顯微技術，即把樣品保存在液氮或液氦溫度下，通過對其某方向的投影顯微像在時空中經調整后進行疊加，從而提高信噪比，最終將不同投影方向的顯微像在三維空間進行重構，獲得大分子的三維結構。

實際上，基于2D 投影的3D 重構具有挑戰性。首先，樣品中生物分子在隨機分布后具有不同的角度，很難確定它們正確的角度，因此，需要多次迭代步驟才能使最終的模型收斂；其次，為了避免損壞樣品，通常使用低劑量的電子束進行輻照，而這樣獲得的電子顯微像具有高噪聲和低對比度，需要大量的電子顯微圖像來生成具有高分辨率的3D 模型。由此可見，冷凍電鏡模型重構依賴于數以萬計的投影圖片進行分析、組裝和優化，冷凍電鏡三維重構需要超過1 petaflops 的計算能力［4］。

目前許多研究學者對冷凍電鏡三維重建計算的優化提出了許多有效的解決方案，如文獻［5］基于分塊流模型的GPU 集群并行化冷凍電鏡三維重建，文獻［6］在CPU-GPU 異構系統上并行化冷凍電鏡三維重建。

在冷凍電鏡三維重構過程中，傅里葉空間圖像相似度計算是調用最為頻繁的計算，其計算成本較高，存在應用程序尚未充分利用并行化優勢、編譯器自動對代碼進行矢量化的優化效果不佳和數據存儲架構導致不必要內存消耗等問題。隨著每個芯片處理核數的增加，負載均衡的重要性也在不斷提高，處理核的利用率是多核系統性能的一個關鍵指標［7］。文獻［8-9］對經典的負載平衡方法進行了回顧和分類，基于此，本文對程序進行手動負載均衡優化操作，使各處理核對計算任務均衡處理，提高程序的并行效率。

在現代微處理器中，單指令多數據流（Single Instruction Multiple Data，SIMD）作為一種利用通用核中數據級并行性的方法［10］得到了廣泛的認可［11］，SIMD 指令集允許程序員與編譯器通過CPU 實現數據級并行性，提高運算效率。SIMD 矢量化在稀疏矩陣向量乘法上有著優異的性能提升效果［12］。因此，本文借鑒SpMV 中用SIMD 矢量化對傅里葉空間圖像相似度計算進行優化。

隨著CPU 計算速度越來越快，處理器速度與訪問主內存速度差距日益增大，以至于程序將大部分CPU時間花費在等待內存訪問上。緩存（cache）［13］技術能夠有效減少磁盤訪問次數，提高數據訪問效率。但硬件緩存大多依賴于空間局部性來實現高效的內存訪問，且多數內存布局優化算法都是基于深入理解程序的算法以及對原有數據結構進行調整。文獻［14］指出算法的內存訪問模式已成為決定算法運行時間的關鍵因素，文獻［15］則對大網格數據布局進行優化，提高了緩存的一致性，減少了緩存未命中率。

更改數據存儲結構能夠獲得更優的存儲分配效果［16］，但原程序中的數據結構在地址空間中分散，不利于內存訪問，因此，需要設計新的數據結構實現對空間局部性的有效利用。本文提出一種基于SIMD的并行傅里葉空間圖像相似度算法。手動對任務進行均勻分配操作，從而在多核處理器上實現負載均衡，同時使用AVX-512 指令集進行高效的矢量化操作。此外，通過優化原程序中的數據結構并使其字節對齊，降低緩存缺失率。

1 研究背景

1.1 冷凍電鏡三維重建

冷凍電鏡結構解析的理論基礎是電鏡三維重構原理［17］，該原理基于數學中的中央截面定理和傅立葉變換的性質。中央截面定理的含義是：一個函數沿某方向投影函數的傅立葉變換等于此函數的傅立葉變換通過原點且垂直于此投影方向的截面函數，即一個生物樣品的電鏡圖像的傅里葉變換等于該三維物體的傅里葉變換通過樣品中心（設定為坐標原點）并垂直于攝像方向的截面，如圖1 所示。

圖1 中央截面定理示意圖Fig.1 Schematic diagram of central section theorem

由于一個函數的傅立葉變換的逆傅立葉變換等價于原來的函數，因此利用傅里葉變換可以很容易地實現圖像的旋轉縮放和平移不變性。根據中央截面定理和傅里葉變換的性質，當利用透射電子顯微鏡獲得的生物樣品多角度放大電子顯微圖像足夠多時，就能得到生物樣品在傅里葉空間的三維信息，再經過傅里葉逆變換便能得到物體的三維結構，如圖2 所示［18］。

圖2 冷凍電鏡三維重構過程Fig.2 Three dimensional reconstruction process of cryo-EM

在冷凍電鏡結構解析的實踐中，可以根據生物樣品的性質及其特點，選擇不同的顯微鏡成像及三維重構方法，目前主要使用單顆粒重構、電子斷層掃描重構等方法針對不同的生物大分子復合體及亞細胞結構進行解析［19］。

冷凍電鏡單粒子法是獲得三維重構圖像的重要辦法［20-21］。單粒子法就是對分離純化的顆粒分子分別成像，其適合的樣品分子量范圍為80～50 MD，最高分辨率約3 ?，其基本原理就是基于分子結構同一性的假設，利用不同投影角度的多個分子的顯微放大二維圖像進行同一分析，通過對正、加和平均等圖像操作提高信噪比，最后將不同投影方向的單顆粒顯微像在三維空間進行重構，獲得三維模型。

1.2 傅里葉空間圖像相似度算法

在冷凍電鏡數據分析處理中，傅里葉空間圖像相似度計算被頻繁調用，其主要使用計算樣品的二維真實圖像與空間中三維重構模型的投影圖像之間的相似度，計算公式如下：

傅里葉空間圖像相似度計算在程序中的實現代碼如算法1 所示，主要計算步驟如圖3 所示。

圖3 傅里葉空間圖像相似度計算的主要步驟Fig.3 Main steps of Fourier space image similarity calculation

首先將disturb0 與ctf 相乘，然后與復數pri 相乘，接著計算dat 與步驟2 結果相減后的范數，與上一步的result 相加后，i加1，重復執行步驟1，直至i=n-1。

算法1串行傅里葉空間圖像相似度算法

2 傅里葉空間圖像相似度計算優化方法

2.1 手動負載均衡

經過OpenMP 對程序進行并行化處理后，程序具有良好的并行性。但是通過vtune 工具進行分析后，有大量資源花費在了Spin time 上。Spin time 是CPU 繁忙的等待時間，是等待其他同步資源處理的自旋等待時間。當同步API 導致CPU 在軟件線程等待進行輪詢時，通常會發生這種情況。對于并行程序，并行循環占大部分執行時間。因此，減少并行循環的耗費也是提升整個程序性能的關鍵，需要對其進行手動負載均衡優化。

首先根據運行環境CPU 的線程數，盡可能將任務均分給每個線程上。如算法2 所示，本文用Begin與end 數組記錄下每個線程處理數據的起始點與終點。在函數logDataVSPrior 中將任務均勻地分給每個線程，盡可能使程序負載均衡。在并行循環計算完成后，再對每個線程計算的result 進行累加，如算法3 所示。

算法2根據線程數對任務進行分配

算法3手動負載均衡后的傅里葉空間圖像相似度算法

2.2 SIMD 矢量化

SIMD 矢量化操作在優化稀疏矩陣向量乘法上表現優異，其采用一個控制器來控制多個處理器，同時對一組數據中的每一個分別執行相同操作，從而實現空間并行性。利用perf 工具分析特定程序后，編譯器已經自動矢量化程序中主要的計算，但仍占據了大量的總體開銷。為了提高系統整體性能，需要在程序中手動添加SIMD 內聯函數對其進行矢量化優化。從算法4 中可以看出，logDataVSPrior 函數中有大量數據執行單個單元的運算，如乘法、加法和減法，適合使用SIMD 向量化對其進行優化。因此，本文將SIMD 矢量化操作添加到logDataVSPrior 函數中。下文選擇INTEL 的AVX-512 指令集，并給出詳細優化步驟。

算法4更新數據結構后傅里葉空間圖像相似度算法的實現

為了方便使用SIMD 指令集，將復數數組拆分為實部數組和虛部數組，分別對復數的實部和虛部進行計算，SIMD 優化步驟如圖4 所示。首先利用_mm512_mul_pd 指令實現實數ctf 與實數disturb0 相乘的操作，并使用_mm512_fnmadd_pd 指令分別計算實部與虛部的值；然后利用_mm512_mul_pd 指令和_mm512_fmadd_pd 指令計算算法 4 中的“real*real+comp*comp”，通 過_mm512_fmadd_pd 指令將步驟4 的計算結果與sigRcp 相乘再與mark 進行相加，并把結果存儲在mark 中；重復執行上述步驟，直到本線程上的任務處理完畢；最后利用_mm512_add_pd 指令將每個線程上得到的計算結果進行累加，得到二范數之和。

圖4 SIMD 優化步驟Fig.4 SIMD optimization step

2.3 內存訪問優化

目前CPU 計算速度遠超內存訪問速度，緩慢的內存訪問速度成為限制系統整體性能的主要瓶頸。利用Linux 系統的Perf 工具進行分析，根據分析結果可以推斷出，由于較高的LLC 缺失率，CPU 需耗費時間等待從本地或遠程DRAM 中獲取數據。由于原程序的數據存儲架構不利于內存訪問，導致不必要的內存消耗，因此提高內存訪問效率也是提升整個程序性能的關鍵。

SIMD 矢量化中提出的數據結構，將復數分為實部向量與虛部向量，在logDataVSPrior 函數計算中存在大量的相對較長數據訪問操作，這對內存訪問非常不利。如圖5 所示，對于需要計算的6 個數組，需要獲取6 個位置的數據值，然后進行計算。假設每個數據之間的步長比當前CPU 的緩存行大得多，那么需要執行6 次高速緩存行替換操作才能獲得6 個不同位置的值。然而，多次替換高速緩存行的操作會影響緩存性能。

圖5 SIMD 優化后的數據結構Fig.5 Data structure after SIMD optimization

為解決這個問題，本文提出一種有效的數據結構。如圖6 所示，將6 個不同的數組鄰接存放，這樣可以將一次循環內需要使用的數據合并到新結構向量中，僅需訪問一次索引位置，即可獲得需要計算的連續數據。與原來的數據結構相比，合并后的數據結構能夠更好地實現緩存數據的重用，提高緩存命中率。這是因為合并后的數據結構使將來用到的數據與當前正在使用的數據在空間地址上相鄰，從而減少了替換高速緩存行的操作。

圖6 本文提出的數據存儲結構Fig.6 Data storage structure proposed in this paper

目前，計算機的內存空間都是按照字節進行劃分的，多數CPU 只從對齊的地址開始加載數據，如果數據不按照一定的規則存儲，則會降低讀取速度，從而影響計算效率。因此，字節對齊在CPU 架構中起著重要的作用。當高速緩存行獲取未字節對齊的數據時，會發生圖7（a）所示的情況。從圖中可以看出，int 類型占用4 Byte 的內存空間，當一個變量int 的起始地址是1 時，那么CPU 需要進行2 次數據讀取操作，這是因為僅一次讀取操作無法完全獲得所需的數據，所以需要多次讀取來獲取剩余的數據。圖7（b）顯示了數據按照內存對齊的方式存放后，只需要1 次讀取操作就能完全獲得所需數據，提高了讀取速度。

圖7 數據字節對齊操作Fig.7 Data byte alignment operation

3 實驗與結果分析

3.1 設備信息

本文使用的實驗平臺是Intel 的Skylake 架構CPU。對于SIMD 指令集，選擇CPU 支持的AVX-512 指令集。快速高效的OpenMP 并行語言用于多線程實現。硬件和軟件配置如表1 所示。

表1 實驗平臺信息Table 1 Experimental platform information

3.2 性能比較與分析

如圖8 所示，本文演示了10 000、20 000、40 000、80 000、和100 000 張真實圖像與投影圖像的所有像素在傅里葉空間中的二范數之和的性能，并比較了4 種不同狀態下的加速比性能，即原代碼運行時間/優化后代碼運行時間。

圖8 不同大小數據集上的加速比Fig.8 Speedup on different size datasets

在圖8 中，原代碼表示使用OenMP 進行并行化處理的傅里葉空間圖像相似度計算的時間成本，優化1 表示僅實現手動負載均衡優化的時間成本，優化2 表示在優化1 的基礎上同時實現手動SIMD 矢量化優化的時間成本，優化3 表示實現第2 節提到的所有優化方法。

在同一實驗平臺上，可以發現優化1 的時間成本相比原代碼有明顯的性能提升，主要有以下3 方面原因：

1）在經過手動負載均衡優化后，各處理核對任務均衡處理，進行負載均衡處理的時間成本可以忽略不計，并且能夠有效減少負載不均衡導致的回滾次數，提高程序并行效率。

2）實現SIMD 矢量化的優化2 比優化1 具有更快的計算速度，因為自動矢量化并不能充分利用機器中SIMD 的性能，經過手動矢量化后達到了更好的并行效果。

3）優化3 中使用了創新的數據結構以及字節對齊操作，使即將要用到的數據連續存儲，實現了優秀的空間局部性，字節對齊操作提高了讀取速度。因此，優化3 相較于之前的優化有較明顯的性能提升。

對于不同大小的數據集，性能影響會有所不同。從圖8 中可以看出，在不同大小的數據集上都有相同的性能優化趨勢，在數據集規模為10 000 和20 000 時，程序具有更出色的性能提升效果，這是因為其空間開銷接近高速緩存空間大小，較小的內存消耗更適合利用高速緩存的空間局部性。

為更直觀地反映本文方法的優化性能，對比不同數據集上的程序加速比性能。如表2 所示。可以看出，與原程序相比，優化后的程序運行時間平均提高了5.132 倍加速比。隨著數據集增大，程序性能依舊得到穩定的提高，這表明本文優化方法不會因為數據集增大而導致性能下降。

表2 在不同數據集上的程序加速比Table 2 Program speedup on different datasets

4 結束語

針對傅里葉空間圖像相似度算法，本文在單個節點上通過手動負載均衡、手動SIMD 矢量化和內存訪問優化這3 種優化方法來獲得更高效的性能。測試結果表明，優化后的程序具有更穩定的性能。本文僅對單節點的傅里葉空間圖像相似度算法進行優化，下一步將針對多節點上的程序做優化處理和測試。