CUDAで行列演算：乗算(シェアードメモリ使用版)

CUDAで行列演算 その3 乗算(シェアードメモリ使用版)†

CUDAで行列演算：乗算で行列の積をやりましたが，メモリアクセスについて問題点があるのと， CUDAで使えるメモリについて勉強するために，シェアードメモリを使って改良してみます．

CUDAで行列演算：乗算の問題点†

CUDAで行列演算：乗算でもかなり高速に実行できていますが，単純にそのまま実装しただけでしたので，さらに高速化することができます． 一番手っ取り早い方法はシェアードメモリを使う方法です． カーネル関数において，行列Aの列数をMとすると

for(int k = 0; k < A.width; ++k){
			x += A.elements[row*A.width+k]*B.elements[k*B.width+col];
		}

の部分だけで，2M回のグローバルメモリ(GPUから見るとグローバルメモリ，ホストから見るとデバイスメモリ)へのアクセスが発生しています． 全体では，スレッド数×2Mで，正方行列を仮定すれば，O(M^3)でアクセス数が増えます． ブロックごとに部分行列を処理していると考え，ブロックごとに必要なAとBの部分行列をシェアードメモリに格納しておけば， グローバルメモリへのアクセスをO(M^2)に抑えることができそうです．

シェアードメモリ†

シェアードメモリ(shared memory)はブロック内のスレッドが共用して使えるオンチップメモリで， グローバルメモリに比べて非常に高速です．シェアードメモリはバンクと呼ばれるメモリモジュールに分割されており， スレッド間のバンク・コンフリクトがなければレジスタアクセスと同じ速さで処理できます． コンフリクトなしのメモリアクセスについてはCUDA Programming Guideの5.1.2.5節を参照してください．

下図は行列の積のイメージです．青の部分が1スレッドがアクセスし，計算する行列の要素であり， この結果C_ijが計算され結果がグローバルメモリに書き込まれます． ブロック単位で考えると，黄色で示した部分が1ブロックが処理する部分であり， それぞれ元の行列の部分行列と考えることができます． 行列AとBの部分行列(黄色部分)をシェアードメモリに格納すれば，グローバルメモリのアクセスを1/BLOCK_SIZEにすることができます． ただし，1ブロック内のスレッドが1度に処理できるのはBLOCK_SIZE×BLOCK_SIZEなので， 行列AとBの部分行列をBLOCK_SIZE×BLOCK_SIZEに分割して処理する必要はあります．

同期処理†

シェアードメモリを用いるとき，スレッド間の同期が必要になります． 例えば，行列の内容をシェアードメモリに格納して，行列の要素の積を計算して足し込んでいくとき， 行列の要素の積を計算する前に，シェアードメモリ内には部分行列のすべてが書き込まれていなければなりません． 各スレッドごとに部分行列の各要素をシェアードメモリへ格納する場合， ブロック内の全スレッドで格納が完全に終了したことを確かにするために同期が必要になります． CUDAではブロック内スレッド同期点の設定のために以下の関数が用意されています．

__syncthreads()

スレッドは，CUDAで行列演算：加減算#l7a8f65aで述べたワープごとにグループ化されて実行されるため， ワープ内のスレッドは暗黙的に同期されます．ワープを意識した実装をすれば，__syncthreads()を省略して同期することも可能です．

実装†

メモリの確保やデータ転送は前回と同じなので省略します．

行列を格納するためのMatrix構造体ですが，部分行列へのアクセスに対応するために， 行列の大きさを示すwidthとheightに加えて，行列要素へのアクセス時に使用するstride変数を追加します．

struct Matrix
{
	int width;
	int height;
	int stride;
	float *elements;
};

部分行列でない場合は，stride = widthとします． 部分行列の場合は，strideは元の行列のwidthとし，elementsに部分行列の 先頭アドレスを格納します(図1の行列Cでいうと黄色部分の左上の要素のアドレス)．

元の行列から部分行列を作成する関数は以下です．

__device__ 
Matrix GetSubMatrix(Matrix A, int row, int col)
{
	Matrix Asub;
	Asub.width  = BLOCK_SIZE;
	Asub.height = BLOCK_SIZE;
	Asub.stride = A.stride;
	Asub.elements = &A.elements[A.stride*BLOCK_SIZE*row+BLOCK_SIZE*col];

	return Asub;
}

GetSubMatrix関数では第1引数に元の行列，第2,3引数に元の行列をBLOCK_SIZEで分割したときの 部分行列の位置(0スタート)を渡します．ここで，部分行列は基本的にBLOCK_SIZE×BLOCK_SIZEであるとしています．

シェアードメモリを用いたカーネルの変更版は以下です．

__global__ 
void matrixMulShared(Matrix A, Matrix B, Matrix C)
{
	int row = blockIdx.y*blockDim.y+threadIdx.y;
	int col = blockIdx.x*blockDim.x+threadIdx.x;

	if(row < C.height && col < C.width){
		int brow = blockIdx.y;
		int bcol = blockIdx.x;

		Matrix Csub = GetSubMatrix(C, brow, bcol);

		int trow = threadIdx.y;
		int tcol = threadIdx.x;

		float x = 0.0f;
		for(int l = 0; l < (A.width+BLOCK_SIZE-1)/BLOCK_SIZE; ++l){
			Matrix Asub = GetSubMatrix(A, brow, l);
			Matrix Bsub = GetSubMatrix(B, l, bcol);

			__shared__ float As[BLOCK_SIZE][BLOCK_SIZE];
			__shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];

			// サブ行列の内容をシェアードメモリへ
			As[trow][tcol] = Asub.elements[trow*Asub.stride+tcol];
			Bs[trow][tcol] = Bsub.elements[trow*Bsub.stride+tcol];

			// 他のスレッドがシェアードへの書き込みを終了するのを待つ
			__syncthreads();

			for(int k = 0; k < BLOCK_SIZE; ++k){
				x += As[trow][k]*Bs[k][tcol];
			}

			// 次の反復でサブ行列が書き換えられるため，ここで処理待ち
			__syncthreads();
		}

		Csub.elements[trow*Csub.stride+tcol] = x;
	}
}

同期点を2点設定しています(ループ内なので実際に実行される数は2×(A.width/BLOCK_SIZE)です)． シェアードメモリを使用するためには，修飾子 __shared__ を用います．

計算時間†

CUDAで行列演算：乗算と同じく，ブロックサイズを16×16として，行列の行と列のサイズをその1倍,10倍,100倍として計算してみました． 計算時間は以下の表です．

時間の単位は ms(ミリ秒)です．一回計測しただけです． シェアードメモリを使うことで高速化ができていることが分かります．