カーネル関数の作成｜行列積高速化#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で公開しています。