GFD セミナー 2011 分科会 2011 年 8 月 21 日

SJPACK の AVX 利用について
石岡圭一
E-mail: ishioka@gfd-dennou.org

まずは球面のスペクトル法についてのおさらい

偏微分方程式中の従属変数を球面調和関数で展開.
MXf (λ, µ, t) =nX

n=0 m=−nm(λ, µ)am(t)Ynn

ここに, λ: 経度, µ = sin θ, θ: 緯度, t: 時刻, M : 切断波数.

Ynm(λ, µ) = Pn|m|(µ)eimλ.
Pnm(µ) は Legendre 陪関数.

Pnm(µ) ≡vuuuut(n − m)! 1(2n + 1)
(n + m)! 2nn!
n+md×(1 − µ2)m/2 n+m (µ2 − 1)ndµ(0 ≤ m ≤ n).

ルジャンドル陪関数による展開

球面スペクトル法には球面調和関数変換が必要. そのうち計算量の大部分を占めるのはルジャンドル陪関数に関する以下の変換である:

逆変換

gjm =MXn=mm(µ )
sm
Pjn n

(m = 0, · · · , M ; j = 1, · · · , J)

正変換

smn =JX
j=1wj gjmPnm(µj )(m = 0, · · · , M ; n = m, · · · , M )ここに, M : 切断波数, J: ガウス緯度の個数, Pnm(µ): ルジャンドル陪関数, µj : ガウス緯度, wj : ガウス重み.

必要な計算量m ≥ 1 に対しては実部・虚部を扱わねばならないことと, ルジャンドル陪関数の対称性から添字 j に関する計算は半分の J/2 でよいことを考慮すると, 逆変換・正変換に必要な計算量は N はルジャンドル陪関数をすべて事前に計算しておくとしても (乗算と加算が必要なことを考慮して)

J
N = · (M + 1)2 · 2 = J(M + 1)2.
2
さらに, 正変換時の数値計算で誤差が出ないことを保証するためには3 M としておく必要があるので,
J>2

3 3
N ∼ M .
2
M が大のときは, この変換計算をいかに効率化するかが全体の計算スピードを決める.

ISPACK における球面調和関数変換のこれまでの実装
stpack(1995 年) 最も素朴な実装
smpack(1998 年) ベクトル化を追求
snpack(1999 年) ベクトル化を追求しつつ省メモリ化
sjpack(2009 年) スカラー計算機での高速化 (SSE2 に対応)
sjpack-cuda(2010 年) GPGPU への対応
sjpack(AVX 対応版)(2011 年) AVX 命令への対応 (ispack-0.94 として公開済み).

AVX 対応について
今年の 1 月に発売が開始された Intel の Sandy Bridge アーキテクチャの CPU では AVX (Advanced Vector eXtensions) 命令が新たに使えるようになった.

AVX 命令の主な特徴:• 256bit 長の SIMD 命令が使えるので, 倍精度変数 4 個に対する処理が一括でできる. それにより, 128bit 長の SSE2 命令に比べると理論的には倍速で処理ができる.

• 四則演算などに 3 オペランド命令が使えるので, C = A + B のような形の計算がレジスタの余分なコピーなしで可能.

コンパイラはそれなりに賢くなってきているが, AVX 命令を最大限活用しようとするならば (SSE2 対応の際と同様に), アセンブリ言語でのプログラミングは避けて通れない.ということで, Legendre 陪関数変換のコアの部分はアセンブリ言語で書いて最適化しておくべき.

実際コアとなる部分はどんなものか?
例: ljlswg-fort.f

SUBROUTINE LJLSWG(JH,S,R,Y,QA,QB,W)
IMPLICIT REAL*8(A-H,O-Z)
DIMENSION W(JH,2),Y(JH),S(2)DIMENSION QA(JH),QB(JH)
DO J=1,JH
W(J,1)=W(J,1)+S(1)*QA(J)W(J,2)=W(J,2)+S(2)*QA(J)QB(J)=QB(J)+R*Y(J)*QA(J)END DO
ENDと短いサブルーチンなので, これを (AVX 命令等を使って) ガシガシに最適化すればいい (というか, このコアの部分をサブルーチンとして独立させてAVX 命令を使って最適化した例
例: ljlswg-avx.s.globl ljlswg_ljlswg_:
movl (%rdi), %edi
vbroadcastsd (%rsi), %ymm0
vbroadcastsd 8(%rsi), %ymm5
vbroadcastsd (%rdx), %ymm1
movq 8(%rsp), %r10
shlq $3,%rdi
xorq %rdx,%rdx
subq %rdi,%rdx
movq %r10,%r11addq %rdi,%r11
subq %rdx,%r8
subq %rdx,%r9
subq %rdx,%r10
subq %rdx,%r11
subq %rdx,%rcx

.L0:
vmulpd (%rcx,%rdx),%ymm1,%ymm4
vmovapd (%r8,%rdx), %ymm2
vmulpd %ymm2,%ymm4,%ymm4
vmulpd %ymm0,%ymm2,%ymm3
vmulpd %ymm5,%ymm2,%ymm2
vaddpd (%r9,%rdx),%ymm4,%ymm4
vaddpd (%r10,%rdx),%ymm3,%ymm3
vaddpd (%r11,%rdx),%ymm2,%ymm2
vmovapd %ymm4,(%r9,%rdx)
vmovapd %ymm3,(%r10,%rdx)
vmovapd %ymm2,(%r11,%rdx)
addq $32,%rdx
jnz .L0
retベンチマーク結果

とりあえず T170 の普通の設定 (J = 256, I = 512) での snpack との速度比較 (なお, sjpack では FFT として FFTJ を使っている). な
お, Flops は実際の計算量 (Legendre 陪関数の毎回計算を含む) ではな
く, Legendre 陪関数をもし事前に計算してメモリに格納していたとする
場合の変換そのものに絶対必要な計算量で算出している).

Intel Core-i7 2620M(Sandy Bridge 2.7GHz, turbo burst 時 3.4GHz),
OS: Debian squeeze 64bit, Compiler: ifort 12.0.5, option xAVXにおいてSJPACK: 6.7 GFlops
SNPACK: 2.9 GFlops

sjpack を SSE2 対応のものにあえて変えた場合
SJPACK: 5.1 GFlops

評価

上記の GFlops は Legendre 陪関数の on the fly での計算に必要な分を除いたものであるので, それも込みの実際の計算量で GFlops を見積るなら, 7/4 倍すべき. すると, AVX 版の SJPACK での実際の GFlops
は, 6.7 × 7/4 = 11.7 GFlops.

Intel Core-i7 2620M の理論性能は, 3.4 × 4 × 2 = 27.2GFlos なので, AVX 版の SJPACK では理論性能の 11.7/27.2 × 100 = 43% はでていることになる.

なお, 次ページでコメントするが, FFT 等を除いた Legendre 陪関数変換の部分だけなら理論性能の約 50%程度はでている.

結論

・sjpack の AVX 対応はそれなりにできた. 球面調和関数変換全体としては SSE2 版の 1.3 倍速くらいにはなっている. 倍率が 1.3 倍止まりなのは, AVX 対応にしたのが Legendre 陪関数変換のところだけで, FFTJなどは AVX に対応していないことが一因. Legendre 陪関数変換のところだけなら 1.5 倍速にはなっている.

・AVX 化した Legendre 陪関数変換が 1.5 倍速どまりなのは, AVX 命令を使っても, load/store の部分が SSE2 命令に比べて速くなるわけではないため. SSE2 版では乗算のところが律速になっているが, AVX 版では乗算の部分が倍速になっている分, load/store が律速になり, その関係で全体で 1.5 倍速どまりになっている.

課題

・FFTJ の部分の AVX 化. ただ, ここはあまり頑張ってもそんなに意味が無いかも (上に書いたように AVX 化しても load/store が速くなるわけではないので, 多分 SSE2 版の倍速には絶対にならない).

・それよりは, もっと簡単にできる sjpack の MPI 化が TODO スタッ
クの一番上かな.

ネット上の情報源- Agner Fog 氏のページ (x86 アセンブリでの高速化についての情報源)(http://www.agner.org/optimize/)- Intel のデベロッパーズマニュアル(ftp://download.intel.co.jp/jp/developer/jpdoc/ )「ネットは広大だわ……」(by 草薙素子, 1989 年)