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

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

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

0 概述

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

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

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

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

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

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

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

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

1 研究背景

1.1 冷凍電鏡三維重建

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

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

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

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

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

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

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

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

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

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

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

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

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

2.1 手動(dòng)負(fù)載均衡

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

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

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

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

2.2 SIMD 矢量化

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

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

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

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

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

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

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

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

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

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

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

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

3 實(shí)驗(yàn)與結(jié)果分析

3.1 設(shè)備信息

本文使用的實(shí)驗(yàn)平臺(tái)是Intel 的Skylake 架構(gòu)CPU。對(duì)于SIMD 指令集，選擇CPU 支持的AVX-512 指令集。快速高效的OpenMP 并行語(yǔ)言用于多線程實(shí)現(xiàn)。硬件和軟件配置如表1 所示。

表1 實(shí)驗(yàn)平臺(tái)信息Table 1 Experimental platform information

3.2 性能比較與分析

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

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

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

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

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

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

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

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

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

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

4 結(jié)束語(yǔ)

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