CUBLASを使ってみたぞ

「はじめてのCUDAプログラミング」*1を読んで、CUDA用の線形計算ライブラリCUBLASを使ってみた。サンプルコードそのままではちょっと分かりにくいのでメモ。
CUBLASの元になってるBLASライブラリはそもそもFortran用のライブラリで、C/C++で使うときには注意が必要であることが知られている。それは、行列演算を行う場合の数値のメモリ上の配置がFortranとC/C++で違うこと。行列Aのi行j列目を普通A_ijと書くんだけど、これをFortranでA(i,j)と書いた場合はメモリ上でこのお隣はA(i+1,j)になる。一方で、C/C++でA[i][j]と書いた場合お隣はA[i][j+1]((一次元配列でA_ijを書く場合でも、A[i*N+j]と書きたくなるのがC/C++プログラマなのだ))。前者を行指向の配置(column-major storage)、後者を列指向(row-major storage)と呼ぶらしい。
さて、BLASが元になってるCUBLASは当然行指向が想定されているけれど、「はじめてのCUDAプログラミング」はいまいちこの件に関する言及が微妙で、慣れていないと容易に間違える。間違えないようにするには、行指向で普段から考えておくか、

#define IDX2C(i,j,ld) (((j)*(ld))+(i))

みたいなマクロを定義しておくこと*2。ここでldはleading dimensionで、早い話が行の数だ。
このマクロを使って書き直したコードを以下に貼っておく。

#include<stdio.h>
#include<stdlib.h>
#include<cublas.h>
#define N 1000
#define M 1500
#define K 500
#define IDX2C(i,j,ld) (((j)*(ld)+(i)))
int main(int argc,char **argv){
  double alpha = 3.0, beta = 1.0;
  double *A,*B,*C;
  double *dA,*dB,*dC;
  int LDA = M, LDB = K, LDC = M;
  int i,j;
  cudaSetDevice(0);
  cublasInit();
  cudaMallocHost((void **)&A,sizeof(double) * M * K);
  cudaMallocHost((void **)&B,sizeof(double) * K * N);
  cudaMallocHost((void **)&C,sizeof(double) * M * N);
  for(i=0;i<M;++i) 
    for(j=0;j<K;++j) A[IDX2C(i,j,M)] = i*K+j + 1;
  for(i=0;i<K;++i)
    for(j=0;j<N;++j) B[IDX2C(i,j,K)] = i*N+j + 1;
  for(i=0;i<M;++i)
    for(j=0;j<N;++j) C[IDX2C(i,j,M)] = 0.0;
  cublasAlloc(M*K,sizeof(double),(void **)&dA);
  cublasAlloc(K*N,sizeof(double),(void **)&dB);
  cublasAlloc(M*N,sizeof(double),(void **)&dC);
  cublasSetMatrix(M,K,sizeof(double),A,LDA,dA,M);
  cublasSetMatrix(K,N,sizeof(double),B,LDB,dB,K);
  cublasSetMatrix(M,N,sizeof(double),C,LDC,dC,M);
  
  cublasDgemm('N','N',M,N,K,alpha,dA,LDA,dB,LDB,beta,dC,LDC);
  cublasGetMatrix(M,N,sizeof(double),dC,M,C,LDC);
  cublasFree(dA);
  cublasFree(dB);
  cublasFree(dC);
  cublasShutdown();
  return 0;
}

ちなみに、Intel Math Kernel LibraryについてるCBLASは、引数で行指向と列指向を切り替えられて便利なので、CUBLASも将来のバージョンでは見習って欲しいものである。

*1:

*2:CUDA Toolkit Archive | NVIDIA DeveloperにあるCUBLAS User Guideを参考にしました。