No Such Blog or Diary

更なる高速化のためにSSE2に手を出してみる．CIPスキームの最内ループの元がこんなもん：

  double const *v = VV + (La+Lb)*(k-1) + La;
  double *v2 = VV + (La+Lb)*(k-1);
  for(double const*const ve = VV + (La+Lb)*(k-1) + La + STEPLEN; ve != v; v2+=2, v+=2){
    v2[0] = CA2 * v[-2] + CA1 * v[-1] + C0 * *v + CB1 * v[1];
    v2[1] = DA2 * v[-2] + DA1 * v[-1] + D0 * *v + DB1 * v[1];
  }

配列をインターリーブして使ってるが，長さ 2 のベクトルを一つ左の値と現在地の値で更新している．ちょうど 128bit のデータになってるから SSE2 にぴったりなので，SSE2を使ったコードが下．当然のごとく配列の alignment を 16byte に合わせておく．

__m128d A2 = _mm_setr_pd(CA2, DA2);
__m128d A1 = _mm_setr_pd(CA1, DA1);
__m128d A0 = _mm_setr_pd(C0 , D0 );
__m128d B1 = _mm_setr_pd(CB1, DB1);
double const *v = VV + (La+Lb)*(k-1) + La;
double *v2 = VV + (La+Lb)*(k-1);
__m128d pv = _mm_load_pd(v-2);
for(double const*const ve = VV + (La+Lb)*(k-1) + La + STEPLEN; ve != v; v2+=2, v+=2){
  __m128d r = _mm_shuffle_pd(pv,pv,_MM_SHUFFLE2(0,0));
  __m128d op = _mm_shuffle_pd(pv,pv,_MM_SHUFFLE2(1,1));
  r = _mm_mul_pd(r, A2);
  op = _mm_mul_pd(op, A1);
  r = _mm_add_pd(r,op);
  pv = _mm_load_pd(v);
  op = _mm_shuffle_pd(pv,pv,_MM_SHUFFLE2(0,0));
  op = _mm_mul_pd(op, A0);
  r = _mm_add_pd(r,op);
  op = _mm_shuffle_pd(pv,pv,_MM_SHUFFLE2(1,1));
  op = _mm_mul_pd(op, B1);
  r = _mm_add_pd(r,op);			
  _mm_store_pd(v2, r);
}

xmm レジスタがもう一本あればもう少しスマートなのだが...　上のコードより場合によっては下のほうが速い．アセンブラ上は定数のうち一つをメモリから取るようになるが，並列に実行できる命令が増えるので速いのかな？

  __m128d op1 = _mm_shuffle_pd(pv,pv,_MM_SHUFFLE2(0,0));
  __m128d op2 = _mm_shuffle_pd(pv,pv,_MM_SHUFFLE2(1,1));
  op1 = _mm_mul_pd(op1, A2);
  op2 = _mm_mul_pd(op2, A1);
  __m128d r = _mm_add_pd(op1,op2);
  pv = _mm_load_pd(v);
  op1 = _mm_shuffle_pd(pv,pv,_MM_SHUFFLE2(0,0));
  op2 = _mm_shuffle_pd(pv,pv,_MM_SHUFFLE2(1,1));
  op1 = _mm_mul_pd(op1, A0);
  op2 = _mm_mul_pd(op2, B1);
  op1 = _mm_add_pd(op1,op2);
  r = _mm_add_pd(r,op1);

うーむ，ここら辺になると Pentium4 の仕様を良く調べないと全く分からない．

次なる最適化は SSE2 の SIMD 命令の使用だ，と．128bit レジスタが8本使えるからいろいろレジスタだけで出来そう．調べてみると xmmintrin.h, emmintrin.h, pmmintrin.h が SSE, SSE2, SSE3 を使うためのヘッダファイルらしい．インテルのコンパイラのリファレンスを見ないと細かい部分がどう使うのか良くわからないが使えるものは使ってみよう．ちなみに g++ に -msse2 のオプションをつけないと怒られた．

ある点の更新に左右の固定範囲の点の値を使う場合ようなスキーム全般に対して，先の例と同様のキャッシュを効くようにするループ変形を可能とする方法を思いついた．

ということで，とある人間が書いたCIPスキームとやらのプログラムを変形してみた．結果としてメモリを少々余計に喰うがやっぱ３倍速くなった．さて，忘れないうちに原理をまとめとくか．

キャッシュにヒットし続けるようにループを分割してみた（簡単のため配列の使い方も少し変えてる）．３倍から５倍程度速くなった．あとは sse2 の命令とか使ってみたいところだけどうまく使える形にプログラムを変形できない…　むりかなぁ？

初期の単純なネストループ：

        for(int p = R; p >0; p--) {
            int q = (ISROOT()) ? p-1: 0;
            double *u = U+q;
            double *v = u++;
            register double cv = *v;
            for(double const *e = U+Len+p; u != e; u++,v++){
                register const double cu = *u;
                *v = (1 - MU) * cu + MU * cv;
                cv = cu;
            }
        }

内側のループを分割してキャッシュ上のデータにアクセスしまくる：

        double *uu = U;
        for(int p = R; p >0; p--) {
            int q = (ISROOT()) ? p-1: 0;
            double *u = uu+q;
            double *v = u++;
            register double cv = *v;
            for(double const *e = uu+p; u != e; u++,v++){
                register const double cu = *u;
                *v = (1 - MU) * cu + MU * cv;
                cv = cu;
            }
        }
        double const *uue = uu + Len - STEPLEN;
        // iteration - STEPLEN * floor(Len/STEPLEN) steps
        for(; uu < uue; uu+=STEPLEN) {
            for(int p = R; p >0; p--) {
                double *u = uu+p-1;
                double *v = u++;
                register double cv = *v;
                for(double const *e = uu+STEPLEN+p; u != e; u++,v++){
                    register const double cu = *u;
                    *v = (1 - MU) * cu + MU * cv;
                    cv = cu;
                }
            }
        }
        // last - rest steps
        int rlen = Len - STEPLEN * (Len/STEPLEN);
        for(int p = R; p >0; p--) {
            double *u = uu+p-1;
            double *v = u++;
            register double cv = *v;
            for(double const *e = uu+rlen+p; u != e; u++,v++){
                register const double cu = *u;
                *v = (1 - MU) * cu + MU * cv;
                cv = cu;
            }
        }

多少の余計な計算を必要とするが，数回分の通信をまとめてやってしまう．今回のプログラムは値の更新に左の値しか使わないので楽．ＣＰＵの力がある場合は２倍から３倍くらい速くなった（１６回位まとめて）．次はキャッシュでも考えて…　

ループ始めに毎回通信：

    MPI_Request reqs1, reqr1;
    for(int t = 0; t < T; t++){
        if(!ISLAST())
            MPI_Isend(U + Len, 1, MPI_DOUBLE, Rank + 1, TAG1, MPI_COMM_WORLD, &reqs1);
        if(!ISROOT())
            MPI_Irecv(U      , 1, MPI_DOUBLE, Rank - 1, TAG1, MPI_COMM_WORLD, &reqr1);
        else
            U[0] = 1;
    
        if(!ISLAST())
            MPI_Wait(&reqs1, MPI_STATUS_IGNORE);
        if(!ISROOT())
            MPI_Wait(&reqr1, MPI_STATUS_IGNORE);
        register double *u = U;
        register double pu = *u++;
        for(double const *e = u+Len; u != e; u++){
            register double cu = *u;
            *u = (1 - MU) * cu + MU * pu;
            pu = cu;
        }            
    }

ループ始めに R 回分通信しとく：

    MPI_Request reqs1, reqr1;
    for(int t = 0; t < T; t+=R){
        if(!ISLAST())
            MPI_Isend(U + Len, R, MPI_DOUBLE, Rank + 1, TAG1, MPI_COMM_WORLD, &reqs1);
        if(!ISROOT())
            MPI_Irecv(U      , R, MPI_DOUBLE, Rank - 1, TAG1, MPI_COMM_WORLD, &reqr1);
        else
            U[R-1] = 1;
       
        if(!ISLAST())
            MPI_Wait(&reqs1, MPI_STATUS_IGNORE);
        if(!ISROOT())
            MPI_Wait(&reqr1, MPI_STATUS_IGNORE);
        for(int p = 0; p < R; p++) {
            int q = (ISROOT()) ? R-1 : p;
            double *u = U+q;
            register double pu = *u++;
            for(double const *e = u+Len+R-q-1; u != e; u++){
                register const double cu = *u;
                *u = (1 - MU) * cu + MU * pu;
                pu = cu;
            }
        }
    }

遅そうなプログラムの最適化でもしようかねぇ．とりあえずシーケンシャル部分から書き換え．無駄に配列使わず上書きで行く．次は通信をまとめてみようか？

配列２本を使う：

        for(int i = 1; i <= Len; i++)
            V[i] = (1 - MU) * U[i] + MU * U[i - 1];      
        double *tmp = U; U = V; V = tmp;

こんなものは一本で十分：

        register double *u = U;
        register double pu = *u++;
        for(double const *e = u+Len; u != e; u++){
            register double cu = *u;
            *u = (1 - MU) * cu + MU * pu;
            pu = cu;
        }            
 

これで一回のスッテプでのメモリアクセスが一箇所になった．良し．