見出し画像

カーネル関数の作成|行列積高速化#9

この記事は、以下の記事を分割したものです。
[元の記事]行列積計算を高速化してみる
一括で読みたい場合は、元の記事をご覧ください。

ここで、最深ループ部分とコピー部分、及び行列Cのスケール計算を、カーネル関数として切り出します。この先、このカーネル関数を最適化して行くためです。ただし、現段階では、関数呼び出し処理が入るため、計算速度は落ちてしまうと予想されます。

9-1. コピー処理の関数化(行列A)

上記の行列AをバッファーA2へコピーする部分を、そのまま関数化すると次のようになります。ただし、タイルに関するサイズ情報は、block2d_info_t構造体を作成し、一つにまとめています。

void myblas_dgemm_copy_t(const double* A, size_t lda, double* A2, const block2d_info_t* info ){
                        
       size_t i2     = info->i2    ;
       size_t k2     = info->j2    ;
       size_t M      = info->M     ;
       size_t K      = info->N     ;
       size_t M2     = info->M2    ;
       size_t K2     = info->N2    ;
       size_t tile_M = info->tile_M;
       size_t tile_K = info->tile_N;
       
       // On L2-Cache Copy for A   
       for( size_t k1=k2; k1<k2+K2; k1+=MIN(K-k1,tile_K ) ){
         size_t K1 = MIN(tile_K ,K-k1);
         for( size_t i1=i2; i1<i2+M2; i1+=MIN(M-i1,tile_M )  ){
           size_t M1 = MIN(tile_M ,M-i1);
           for( size_t i =i1; i < i1+M1; i++ ){
             for( size_t k =k1; k < k1+K1; k++ ){
               (*A2) = (*A);
               A += lda ;
               A2++; 
             }
             A  = A  - lda *K1 + 1;
           } 
         }
         A = A  - M2 + lda *K1;
       }       
       A2 = A2  - M2*K2;
       A = A  - lda*K2;
                   
}          

ここで、k2,i2などのインデックス変数は計算自体には使われておらず、K2,M2などのサイズ情報だけで計算可能です。具体的には、フルサイズのタイル個数と端数の計算に場合分けします。プログラムは次のようになります。

void myblas_dgemm_copy_t(const double* A, size_t lda, double* A2, const block2d_info_t* info ){

       size_t M2     = info->M2    ;
       size_t K2     = info->N2    ;
       size_t tile_M = info->tile_M;
       size_t tile_K = info->tile_N;

       size_t MQ = M2/tile_M; // number of tiles in M axis
       size_t MR = M2%tile_M; // final tile size in M axis
       size_t KQ = K2/tile_K; // number of tiles in K axis
       size_t KR = K2%tile_K; // final tile size in K axis

       if( MR >  0 ){ MQ++; }
       if( MR == 0 ){ MR = tile_M; }
       if( KR >  0 ){ KQ++; }
       if( KR == 0 ){ KR = tile_K; }

       // L1-cache blocking
       size_t k1 = KQ;
       while( k1-- ){
         size_t K1 = tile_K; if( k1==0 ){ K1=KR; }
         size_t m1 = MQ;
         while( m1-- ){
           size_t M1 = tile_M; if( m1==0 ){ M1=MR; }
           size_t m = M1;
           while( m-- ){
             size_t k = K1;
             while( k-- ){
               (*A2) = (*A);
               A += lda ;
               A2++;
             }
             A  = A  - lda *K1 + 1;
           }
         }
         A = A  - M2 + lda *K1;
       }

}


9-2. コピー処理の関数化(行列B)

行列Bのコピー関数も同様の修正すると、次のようになります。

void myblas_dgemm_copy_n(const double* B, size_t ldb, double* B2, const block2d_info_t* info ){

       size_t K2     = info->M2    ;
       size_t N2     = info->N2    ;
       size_t tile_K = info->tile_M;
       size_t tile_N = info->tile_N;

       size_t NQ = N2/tile_N;
       size_t NR = N2%tile_N;
       size_t KQ = K2/tile_K;
       size_t KR = K2%tile_K;

       if( NR >  0 ){ NQ++; }
       if( NR == 0 ){ NR = tile_N; }
       if( KR >  0 ){ KQ++; }
       if( KR == 0 ){ KR = tile_K; }

       // L1-cache blocking
       size_t k1 = KQ;
       while( k1-- ){
         size_t K1 = tile_K; if( k1==0 ){ K1=KR; }
         size_t n1 = NQ;
         while( n1-- ){
           size_t N1 = tile_N; if( n1==0 ){ N1=NR; }
           size_t n = N1;
           while( n-- ){
             size_t k = K1;
             while( k-- ){
               *B2=*B;
               B++;
               B2++;
             }
             B  = B  - K1 + ldb ;
           }
         }
         B  = B  - ldb *N2 + K1;
       }

}


9-3. 行列積計算の関数化

さらに、最深ループ部分も同じく関数化します。そして、上記と同様の修正を三重ループに対して行います。すると、プログラムは、下記のようになります。

void myblas_dgemm_kernel(double alpha, const double *A2, const double *B2,
                        double *C, size_t ldc, const block3d_info_t* info ){

       size_t M2     = info->M2    ;
       size_t N2     = info->N2    ;
       size_t K2     = info->K2    ;
       size_t tile_M = info->tile_M;
       size_t tile_N = info->tile_N;
       size_t tile_K = info->tile_K;

       size_t MQ = M2/tile_M;
       size_t MR = M2%tile_M;
       size_t NQ = N2/tile_N;
       size_t NR = N2%tile_N;
       size_t KQ = K2/tile_K;
       size_t KR = K2%tile_K;

       double       AB;

       if( MR >  0 ){ MQ++; }
       if( MR == 0 ){ MR = tile_M; }
       if( NR >  0 ){ NQ++; }
       if( NR == 0 ){ NR = tile_N; }
       if( KR >  0 ){ KQ++; }
       if( KR == 0 ){ KR = tile_K; }

       // L1-cache blocking
       size_t k1 = KQ;
       while( k1-- ){
         size_t K1 = tile_K; if( k1==0 ){ K1=KR; }

         size_t n1 = NQ;
         while( n1-- ){
           size_t N1 = tile_N; if( n1==0 ){ N1=NR; }

           size_t m1 = MQ;
           while( m1-- ){
             size_t M1 = tile_M; if( m1==0 ){ M1=MR; }

             // Kernel ----
             size_t n = N1;
             while( n-- ){
               size_t m = M1;
               while( m-- ){
                 AB=0e0;
                 size_t k = K1;
                 while( k-- ){
                   AB = AB + (*A2)*(*B2);
                   A2++;
                   B2++;
                 }
                 *C = (*C) + alpha*AB;
                 B2 = B2 - K1;
                 C++;
               }
               A2 = A2 - M1*K1;
               B2 = B2 + K1;
               C  = C - M1 + ldc;
             }
             A2 = A2 + M1*K1;
             B2 = B2 - K1*N1;
             C  = C - ldc*N1 + M1;
             // ---- Kernel

           } // end of m2-loop
           A2 = A2 - M2*K1;
           B2 = B2 + K1*N1;
           C  = C - M2 + ldc*N1;

         } // end of n2-loop
         A2 = A2 + M2*K1;
         C  = C - ldc*N2;

       } // end of k2-loop
       A2 = A2 - M2*K2;
       B2 = B2 - K2*N2;

}


9-4. スケール計算の関数化(行列C)

最後に、行列Cにbetaを乗算するスケール関数も作成しておきます。これは次のようになります。

void myblas_dgemm_scale2d(double beta, double *C, size_t ldc, const block2d_info_t* info ){

       size_t M      = info->M2    ;
       size_t N      = info->N2    ;
       size_t tile_M = info->tile_M;
       size_t tile_N = info->tile_N;
       
        // scaling beta*C
        size_t n = N;
        while( n-- ){
            size_t m = M;
            while( m-- ){
               *C=beta*(*C);
                C++;
            }
            C = C - M + ldc;
        }

}


9-5. 作成した構造体

なお、サイズ情報を上記のカーネル関数に渡しているblock2d_info_t構造体とblock3d_info_t構造体は、次のように宣言しています。

typedef struct _block3d_info_t {
       size_t M2;
       size_t N2;
       size_t K2;
       size_t tile_M;
       size_t tile_N;
       size_t tile_K;
} block3d_info_t;

typedef struct _block2d_info_t {
       size_t M2;
       size_t N2;
       size_t tile_M;
       size_t tile_N;
} block2d_info_t;


9-6. 計算速度チェック

では、最後に計算速度を確認してみます。

<カーネル化前>

Max  Peak MFlops per Core: 52800 MFlops 
Base Peak MFlops per Core: 46400 MFlops 
size  , elapsed time[s],          MFlops,   base ratio[%],    max ratio[%] 
   16,      2.5034E-05,         357.914,        0.771366,        0.677867 
   32,     0.000141144,         486.086,          1.0476,        0.920617 
   64,      0.00026679,         2011.23,         4.33454,         3.80914 
  128,      0.00189614,         2237.94,         4.82315,         4.23852 
  256,       0.0145819,         2314.58,         4.98832,         4.38368 
  512,        0.119027,         2261.86,         4.87469,         4.28382 
 1024,        0.903981,         2379.06,         5.12729,          4.5058 
 2048,         7.25661,         2369.21,         5.10606,         4.48715 

<カーネル化後>

Max  Peak MFlops per Core: 52800 MFlops 
Base Peak MFlops per Core: 46400 MFlops 
size  , elapsed time[s],          MFlops,   base ratio[%],    max ratio[%] 
   16,     2.00272E-05,         447.392,        0.964208,        0.847334 
   32,     0.000152111,         451.039,        0.972067,         0.85424 
   64,     0.000251055,         2137.29,         4.60622,         4.04789 
  128,      0.00179315,         2366.49,         5.10019,         4.48198 
  256,       0.0152731,         2209.84,         4.76258,          4.1853 
  512,        0.130602,         2061.39,         4.44265,         3.90415 
 1024,        0.916094,         2347.61,          5.0595,         4.44623 
 2048,         7.26869,         2365.28,         5.09758,         4.47969 

結果としては、行列サイズ2048×2048の場合、カーネル化前後で比較すると、2369 MFLOPS → 2365 MFLOPSと、わずかに速度が落ちてしまいました。原因としては、関数呼び出しコストが増加したためと考えられます。


次の記事

元の記事はこちらです。

ソースコードはGitHubで公開しています。


この記事が気に入ったらサポートをしてみませんか?