転置行列コピー関数の改良｜行列積高速化#27

行列コピー関数は、行列積の計算量O(N^3)よりも計算コストが圧倒的に小さいことを利用して、以下の目的で導入しました。

（１）あらかじめ、行列データをキャッシュメモリに書き込んでおく
（２）行列データを直列化し、行列積計算カーネルのキャッシュ利用効率を最大化
（３）DGEMMの転置オプションを行列コピー関数の差し替えで実現する

これらに関しては、以下の記事をご参照ください。ただ、これらの記事は、実装や高速化にとって重要なポイントなので、有料記事にしています。

転置行列コピーの問題点

行列コピー関数は、転置せずにコピーする関数と転置しながらコピーする関数の２種類を用意しています。転置を行わない行列コピー関数は特に問題はないのですが、転置操作を伴いながら行う行列コピー関数には、高速化を阻む問題点が１つあります。

それは、転置行列コピー関数が、ストライド・アクセスになっていることです。

アンロールブロッキングを行う前のソースコードを見てみましょう。

size_t m = M1;
while( m-- ){
  size_t k = K1;
  while( k-- ){
    (*A2) = (*A);
    A += lda ;
    A2++;
  }
  A  = A  - lda *K1 + 1;
}

ここで、M1,K1はL1キャッシュブロッキングサイズ、ldaは行列Aの行サイズ（リーディングディメンジョン）です。

このコードでは、ループ最深部でデータ読み込みに使用するポインタAは、ldaずつポインタシフトしています。これは、lda=1であればシーケンシャル・アクセスしていることになりますが、lda>1の場合は全てストライド・アクセスになります。

もし、ldaがキャッシュラインサイズ64Bを超えると、次のデータがハードウェア・プリフェッチされておらず、キャッシュヒットミスをし続けることになります。

そこで、下記の記事では、これを防ぐためにソフトウェア・プリフェッチを使って、キャッシュヒットミスを削減していました。

しかしながら、これをシーケンシャル・アクセスに変更した方が、より性能が高いであろうことは予想できるでしょう。

シーケンシャル・アクセス化

上記のコードでポインタAをシーケンシャル・アクセスに変更するには、mループとkループの入れ替えること（＝ループ交換）で実現できます。

実際に入れ換えると次のようなコードになります。

size_t k = K1;
while( k-- ){
  size_t m = M1;
  while( m-- ){
    (*A2) = (*A);
    A++;
    A2+=K;
  }
  A2 = A2 - K1*M1 + 1;
  A  = A  - M1  + lda;
}

ポインタAがシーケンシャル・アクセスになった代わりに、書き込み先のポインタA2がストライド・アクセスになります。

この段階だと、とてもシンプルな変更ですみます。

しかしながら、このコードを最適化するために、アンロールブロッキングを行うと別の問題が現れます。

その問題とは、ブロッキングに伴う端数の取り扱いです。

端数処理の方法

なぜ端数の取り扱いが問題になるかというと、ポインタA2がストライド・アクセスになっているためです。

行列コピー関数は、コピーと同時にデータの直列化（シリアライズ）を行っていますが、これがストライド・アクセスと非常に相性が悪いのです。

具体的には、ポインタA2のポインタシフトが端数を含めると、シフトサイズをうまく計算することができません。散々悩みましたが、シフトサイズの計算を諦めました。

ポインタシフトを諦めると、解決策は意外と簡単でした。

具体的には、端数専用のポインタを用意すると解決できます。

下記のコードをご覧ください。

double* A2_2 = A2   + K*( M & ~3 );
double* A2_1 = A2_2 + K*( M &  2 );;


if( K >> 3 ){       
 size_t k8 = ( K >> 3 );
 while( k8-- ){  
   if( M >> 2 ){
     size_t m4 = ( M >> 2 );
     while( m4-- ){ 

       for( size_t l=0; l<8; l++ ){
         for( size_t i=0; i<4; i++ ){
           (*A2) = *(A+i+l*lda);
           A2++; 
         }
       }       
       A  = A  + 4;
       A2 = A2 - 4*8 + 4*K;

     }
     A2 = A2 - (M&~3)*K + 4*8; // move to next row block
   }
   if( M & 2 ){

       for( size_t l=0; l<8; l++ ){
         for( size_t i=0; i<2; i++ ){
           (*A2_2) = *(A+i+l*lda);
           A2_2++;
         }         
       }           
       A  = A  + 2;

   }
   if( M & 1 ){

       for( size_t l=0; l<8; l++ ){
         for( size_t i=0; i<1; i++ ){
           (*A2_1) = *(A+i+l*lda);
           A2_1++;
         }
       }
       A  = A  + 1;

   }
   A  = A - M + 8*lda;
 }
}
...以下省略...

このコードは、K（=K1）を8でブロッキングした部分を示しています。Kの端数であるK=4,2,1の場合は省略しています。M（=M1）は、４でブロッキングし、端数であるM=2,1の場合も示してあります。

ポイントは、端数専用のポインタA2_2、A2_1を用意していることです。

ポインタA2, A2_2, A2_1が、それぞれ担当している範囲は次のようになります。

上図は、Kを４でブロッキングした例ですが、Kを８でブロッキングした場合でも同じです。行列内部の矢印は、メモリの連続方向を示しています。

ポインタA2の直列化の方向は、上図では横方向になります。行列積カーネル関数を高速化するためには、まずA2の範囲内で直列化し、次にA2_2の範囲で直列化、最後にA2_1の範囲で直列化する、という処理が必要になります。つまり、A2_2とA2_1は、A2とは独立に直列化していく必要があります。

コードの中の下記２行は、ポインタA2_2とA2_1の先頭アドレスを計算しています。

double* A2_2 = A2   + K*( M & ~3 );
double* A2_1 = A2_2 + K*( M &  2 );;

A2_2とA2_1は、A2とは独立に直列化するため、一度先頭アドレスが計算できれば、あとは単純にインクリメント（A2_2++,A2_1++など）するだけで良くなります。上記のコードでもそうなっていることが確認できると思います。

性能比較

あとは、上記のブロッキングしたコードを最適化の手順に従って、アンロールし、アセンブラでベクトル化し、プリフェッチを挿入したのち、ループウィンドウをずらしてアセンブラ命令順序を最適化しました。

ストライド・アクセスとシーケンシャル・アクセスの性能比較すると、下記のようになりました。

測定の結果、行列サイズが小さいとほとんど差がありませんが、行列サイズが1600x1600を超えたあたりからシーケンシャル・アクセス（オレンジ線）の方が明らかに高速になっていることが分かります。

行列積計算全体に対しては、1~2%の差でしかありませんが、シーケンシャル・アクセスの方が高性能になるようです。

まとめ

転置を伴う行列コピー関数では、データの読み込みがストライド・アクセスになっている点が問題でした。

データの読み込みをシーケンシャル・アクセスにすると、データの書き込みがストライド・アクセスになり、これが直列化との相性が悪く、ブロッキングを行った際の端数の取り扱いを難しくしていました。

端数部分の直列化には、端数専用のポインタを別途用意することで、実装が容易になりました。

ストライド・アクセスの場合とシーケンシャル・アクセスの場合で性能比較を行ったところ、行列サイズが大きい場合にはシーケンシャル・アクセスの方が早いことが分かりました。

次の記事

ソースコード

jyam45/dgemm