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

郭渝洛，邊浩東，董潤婷，唐嘉豪，王曉英，黃建強(qiáng)

（青海大學(xué) 計算機(jī)技術(shù)與應(yīng)用系，西寧 810016）

0 概述

二十世紀(jì)七十年代，TAYLOR 和GLAESER 研發(fā)了冷凍電鏡（cryo-Electron Microscopy，cryo-EM）。該項技術(shù)發(fā)展至今，已經(jīng)成為研究生物大分子結(jié)構(gòu)與功能強(qiáng)有力的手段，其簡化了生物細(xì)胞的成像過程，提高了成像質(zhì)量，將生物化學(xué)帶入一個新時代［1-3］。

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

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

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

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

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

隨著CPU 計算速度越來越快，處理器速度與訪問主內(nèi)存速度差距日益增大，以至于程序?qū)⒋蟛糠諧PU時間花費(fèi)在等待內(nèi)存訪問上。緩存（cache）［13］技術(shù)能夠有效減少磁盤訪問次數(shù)，提高數(shù)據(jù)訪問效率。但硬件緩存大多依賴于空間局部性來實現(xiàn)高效的內(nèi)存訪問，且多數(shù)內(nèi)存布局優(yōu)化算法都是基于深入理解程序的算法以及對原有數(shù)據(jù)結(jié)構(gòu)進(jìn)行調(diào)整。文獻(xiàn)［14］指出算法的內(nèi)存訪問模式已成為決定算法運(yùn)行時間的關(guān)鍵因素，文獻(xiàn)［15］則對大網(wǎng)格數(shù)據(jù)布局進(jìn)行優(yōu)化，提高了緩存的一致性，減少了緩存未命中率。

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

1 研究背景

1.1 冷凍電鏡三維重建

冷凍電鏡結(jié)構(gòu)解析的理論基礎(chǔ)是電鏡三維重構(gòu)原理［17］，該原理基于數(shù)學(xué)中的中央截面定理和傅立葉變換的性質(zhì)。中央截面定理的含義是：一個函數(shù)沿某方向投影函數(shù)的傅立葉變換等于此函數(shù)的傅立葉變換通過原點(diǎn)且垂直于此投影方向的截面函數(shù)，即一個生物樣品的電鏡圖像的傅里葉變換等于該三維物體的傅里葉變換通過樣品中心（設(shè)定為坐標(biāo)原點(diǎn)）并垂直于攝像方向的截面，如圖1 所示。

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

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

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

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

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

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

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

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

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

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

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

2 傅里葉空間圖像相似度計算優(yōu)化方法

2.1 手動負(fù)載均衡

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

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

算法2根據(jù)線程數(shù)對任務(wù)進(jìn)行分配

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

2.2 SIMD 矢量化

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

算法4更新數(shù)據(jù)結(jié)構(gòu)后傅里葉空間圖像相似度算法的實現(xiàn)

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

圖4 SIMD 優(yōu)化步驟Fig.4 SIMD optimization step

2.3 內(nèi)存訪問優(yōu)化

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

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

圖5 SIMD 優(yōu)化后的數(shù)據(jù)結(jié)構(gòu)Fig.5 Data structure after SIMD optimization

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

圖6 本文提出的數(shù)據(jù)存儲結(jié)構(gòu)Fig.6 Data storage structure proposed in this paper

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

圖7 數(shù)據(jù)字節(jié)對齊操作Fig.7 Data byte alignment operation

3 實驗與結(jié)果分析

3.1 設(shè)備信息

本文使用的實驗平臺是Intel 的Skylake 架構(gòu)CPU。對于SIMD 指令集，選擇CPU 支持的AVX-512 指令集?？焖俑咝У腛penMP 并行語言用于多線程實現(xiàn)。硬件和軟件配置如表1 所示。

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

3.2 性能比較與分析

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

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

在圖8 中，原代碼表示使用OenMP 進(jìn)行并行化處理的傅里葉空間圖像相似度計算的時間成本，優(yōu)化1 表示僅實現(xiàn)手動負(fù)載均衡優(yōu)化的時間成本，優(yōu)化2 表示在優(yōu)化1 的基礎(chǔ)上同時實現(xiàn)手動SIMD 矢量化優(yōu)化的時間成本，優(yōu)化3 表示實現(xiàn)第2 節(jié)提到的所有優(yōu)化方法。

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

1）在經(jīng)過手動負(fù)載均衡優(yōu)化后，各處理核對任務(wù)均衡處理，進(jìn)行負(fù)載均衡處理的時間成本可以忽略不計，并且能夠有效減少負(fù)載不均衡導(dǎo)致的回滾次數(shù)，提高程序并行效率。

2）實現(xiàn)SIMD 矢量化的優(yōu)化2 比優(yōu)化1 具有更快的計算速度，因為自動矢量化并不能充分利用機(jī)器中SIMD 的性能，經(jīng)過手動矢量化后達(dá)到了更好的并行效果。

3）優(yōu)化3 中使用了創(chuàng)新的數(shù)據(jù)結(jié)構(gòu)以及字節(jié)對齊操作，使即將要用到的數(shù)據(jù)連續(xù)存儲，實現(xiàn)了優(yōu)秀的空間局部性，字節(jié)對齊操作提高了讀取速度。因此，優(yōu)化3 相較于之前的優(yōu)化有較明顯的性能提升。

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

為更直觀地反映本文方法的優(yōu)化性能，對比不同數(shù)據(jù)集上的程序加速比性能。如表2 所示?？梢钥闯觯c原程序相比，優(yōu)化后的程序運(yùn)行時間平均提高了5.132 倍加速比。隨著數(shù)據(jù)集增大，程序性能依舊得到穩(wěn)定的提高，這表明本文優(yōu)化方法不會因為數(shù)據(jù)集增大而導(dǎo)致性能下降。

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

4 結(jié)束語

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