迭代式 MRI 重建 練習題 (Practice - MRI Background and the Iterative Reconstruction Problem)

兩階段:acquisition (scan) 與 reconstruction。acquisition 階段掃描器沿一條預定的 trajectory 在 k-space(即空間頻率域 / Fourier transform 域)取樣資料 s(k);reconstruction 階段把這些樣本轉換成目標影像 m(r),即估計組織的形狀與紋理。trajectory 的選擇同時影響重建影像品質、演算法時間複雜度與掃描時間。

Cartesian 是均勻間隔的網格取樣:此時 weighting function W(k) 為常數可提出加總,且 Eq.(17.1) 的指數項在 k-space 均勻間隔 → 整個加總恰好是 inverse FFT(極有效率)。non-Cartesian 的指數項不再均勻間隔,加總不具 FFT 形式,不能直接套 inverse FFT;需先用 gridding(內插到均勻網格後 FFT)或改用 iterative linear solver 重建。

因為矩陣尺寸過大:即使中等解析度 128³ voxel,F 就有 128³ = 2,097,152 欄,每欄有 N 個元素(N = k-space 樣本數)。如此巨大的矩陣使 Gaussian elimination 等直接求逆法實務上不可行 (intractable)。改用 Conjugate Gradient (CG) 迭代法逼近,不顯式求逆;每次迭代的主成本是 matrix-vector 乘 (FᴴF+λWᴴW)·ρ。

FHD(= FᴴD,matrix-vector 乘)與 Q(加速 FᴴF 的 matrix-matrix 乘)核心計算結構完全相同,Q 只是計算量大得多,故討論其一即可。選 FHD 因為它 每次影像擷取都要跑(D 每次不同,~3h CPU);而 Q 每個 scanner+trajectory 只算一次可重用,CG solver(find ρ)因 Q 已預算只占 CPU 時間 <1%,故暫不平行化。

每個 FHD[n] 累加所有 k-space 樣本 m 的貢獻,但任一 voxel 不依賴其他 voxel → 外、內迴圈所有 iteration 皆可平行。唯一依賴:同一外迴圈 iteration 內,內迴圈用到該迭代先算好的 rMu[m]/iMu[m]。兩瓶頸:(1) 記憶體頻寬,compute-to-memory 僅 0.25~0.75 OP/B;(2) 三角函數,FP 算術:trig ≈ 13:2,sin/cos 長延遲低吞吐。

這是 scatter 結構:每個 thread(一個樣本 m)會寫進所有 rFhD/iFhD voxel,多個 thread 撞同一輸出,必須用 atomicAdd 防止互相覆蓋。不能 privatization 是因為 rFhD/iFhD 大小 = 影像 voxel 總數(百萬級),塞不進 shared memory。出路是改成 gather。

loop interchange 要求完美巢狀 (perfectly nested) 迴圈(外、內 for 之間無其他 statement),但 Fig. 17.4 外迴圈在內迴圈前有 rMu/iMu 計算 → 非完美巢狀。Loop fission 把外迴圈 body 拆成「只算 rMu/iMu 的 cmpMu 迴圈」與「完美巢狀的 FHD 迴圈」。fission 後變成「先做完全部 pre,再做全部 inner」;因為 pre 部分(rMu/iMu)不依賴任何前次迭代內迴圈的結果,故改變相對順序不影響結果。

看 index 與 threadIdx 的關係。x[n] y[n] z[n] 與 rFhD[n] iFhD[n] 的 index 是 n(thread 固定),整個 m-loop 不變(座標)或僅累加(輸出)→ 可在迴圈前載入 register、圈後寫回一次。kx[m] ky[m] kz[m] 的 index 是 m,每次迭代都變,無法整圈重用 register → register 救不了,改放 constant memory。效果:每圈 global accesses 14→7,OP/B 0.23→0.46。

適合的原因:(1) k-space 唯讀;(2) index = m 與 threadIdx 無關 → warp 內 32 個 thread 在同一迭代存取同一元素 → constant cache 可broadcast 給全 warp,每 32 次存取 ≥31 次由 cache 服務,消除 96%+ 的 global accesses。限制:constant memory 僅 64KB,但樣本可達數百萬 → 必須 chunking:host 端用迴圈把資料切成 ≤64KB 的塊,每次 cudaMemcpyToSymbol 搬一塊再啟動一次 kernel。加上 constant memory 後僅剩 rMu/iMu 兩次 global access,OP/B 升至 1.63。

__sin/__cos(前綴兩底線)是 compiler 認得的 intrinsic,翻成 SFU (Special Function Units) 硬體指令,throughput 遠高於軟體版(軟體版用 ~13/12 個 FP op 的 5 項 Taylor 級數),但精度略低。因 MRI 是醫學影像,須驗證:用 phantom (I0) 反向合成 k-space → 重建出 I → 代入 PSNR 公式 PSNR = 20·log10(max(I0)/√MSE) 比較;HW trig 使 PSNR 由 27.6 僅降到 27.5 dB(可接受),必要時再目視檢查 artifacts。

OP/B = 浮點運算數 ÷ (global accesses × 4 bytes/float)。
(a) 13/(14×4) = 0.23 OP/B(bandwidth bound)。
(b) 13/(7×4) = 0.46 OP/B(register promotion,代價每 thread +5 register)。
(c) 13/(2×4) = 1.63 OP/B(constant memory + 編譯器把 rMu/iMu 的 4 次存取合併成 2 次 load)。三步把 0.23 → 1.63,使 bandwidth 不再是唯一瓶頸。

(a) N = 128³ = 2,097,152 voxel → grid = 2,097,152/1024 = 2048 blocks(每 thread 算一個 voxel n)。
(b) grid = 1,000,000/1024 = 977 blocks(若非整除,需再補一輪)。
(c) 變成 12·CHUNK_SIZE:AoS 把 x/y/z 三個 single-precision(4 byte)分量併進一個 struct 元素,一次 memcpy 同時搬三分量 → 3×4 = 12 bytes/元素,且只需一個陣列、一次 memcpy。

(a) warps/SM = (1024/32) × 2 = 64 warps。
(b) worst-case working set = 64 × 3 = 192 cache lines(即使假設 3 個 warp 共享同迭代,仍約 64 條)。
(c) 32 條 < working set → 不同 warp 互相搶 line,下一圈要用的元素早被擠出 → cache thrashing,無法消除 global accesses。解法:把 kx/ky/kz 改成 array-of-structures,使一迭代的 x/y/z 三分量落在同一條 cache line,每 warp 只需 1 條 → working set ÷3。

① thread = 內迴圈 → N×M(百萬 × 十萬)個 thread,遠超裝置所需 → 不可行。② thread = 外迴圈(m)→ M 個 thread,平行度尚可,但與 Fig. 17.5 相同是 scatter:每 thread 寫進所有 voxel → 大量衝突 → 需 atomicAdd → 慢。③ thread = 輸出 voxel(n)→ 需先 loop interchange,每 thread 有唯一的 n,只累加自己的 rFhD[n]/iFhD[n] → 無衝突、無 atomics,且 N(=128³≈2.1M)平行度充足。故 ③ 是 gather,為三者最佳。核心心法:讓 thread 擁有唯一輸出,就能消滅 atomics。

關鍵在 warp 內 thread 的存取型態。global memory coalescing 的情境是「相鄰 thread 存取連續元素」(如 data[threadIdx.x])→ 此時 SoA 讓同一欄位連續排列,使一次 warp 存取落在連續位址 → coalesced。本章 k-space 的情境是「warp 內所有 thread 存取同一個元素 k[m]」(index m 與 threadIdx 無關)→ 走的是 constant cache 廣播,目標是讓一迭代用到的 x/y/z 三分量擠進同一條 cache line,故用 AoS。存取型態不同(連續 vs 同一元素),最佳佈局就相反。

(a) FHD 被大幅加速後,原本占 CPU <1% 的 CG solver(find ρ) 反而變成約 50% 的執行時間,FHD 只剩 ~50%。這是 Amdahl's Law 的常見現象:某階段被成功加速後,執行時間轉由原本不顯著的其他階段(此處 CG solver)主宰;要再快就得去加速 CG。
(b) OP/B = FP-op ÷ (memory bytes)。把 trig 從「~13 個 Taylor FP op」變成「單一 SFU 指令」會讓分子的 FP-op 計數下降 → OP/B 趨近 worst-case 0.25。但實際更快,因為 SFU 是獨立硬體單元,把 trig 從 FP pipeline 卸載且本身高吞吐、可與 FP 運算重疊。OP/B 數字不能脫離「是哪個單元在算」單獨解讀。

核心觀念鏈:non-Cartesian 軌跡 → 不能 inverse FFT → iterative CG → FHD 是瓶頸(~3h、每次都要)→ scatter→gather(loop fission + interchange)避開 atomics → register promotion → constant memory(broadcast + 64KB chunking)→ AoS 解 cache thrashing → HW trig 解 trig 吞吐 → 全章 ~10×,但 Amdahl 反噬使 CG solver 成新瓶頸。

關鍵公式:ρ=(FᴴF+λWᴴW)⁻¹FᴴD、OP/B = FP_ops/(accesses×4)、gridDim = N/THREADS_PER_BLOCK、working_set = (threads/32)×blocks/SM × cache_lines_per_warp、PSNR = 20·log10(max(I0)/√MSE)。

數字記憶:OP/B 0.23→0.46→1.63;FP:trig=13:2;constant memory 64KB;working set 64×3=192 vs ~32 lines;PSNR double=single=27.6、HW trig=27.5、gridding/iFFT=16.8 dB;全章 ~10×;FHD 占比 100%→50%。

一句話總結:迭代式 MRI 重建是「output-centric gather」案例——先用迴圈轉換把 scatter 變 gather 消滅 atomics,再層層抬高 arithmetic intensity 與 trig 吞吐,最後被 Amdahl 定律提醒新瓶頸何在。