転置しながらコピーする
cuBLAS が扱う行列要するに二次元配列のメモリレイアウトはFortranとC/C++では大きく異なります。 行列
│ 1 2 3 │ │ 4 5 6 │
がメモリ上にリニアに展開されたとき、
Fortranでは 1 4 2 5 3 6 の順、上から下/左から右に並ぶのに対し、
C/C++ では 1 2 3 4 5 6 の順、左から右/上から下に並びます。
上から下(行:row方向)が左から右(列:column方向)より優先するFortran流レイアウトをrow-major、反対のC/C++流レイアウトをcolumn-majorっていいます。
cuBLASの基となった本家BLASはもともとFortranで実装されたライブラリであるがため、行列表現はFortranに倣ってrow-majorになってます。したがって上記の行列をC/C++の二次元配列で表現すると、
// 列 行 float A[3][2] = { { 1.0f, 4.0f }, { 2.0f, 5.0f }, { 3.0f, 6.0f } };
となります。行列を転置(行/列を交換)せにゃならんのです。
これがどうにも居心地よろしくありません。Host上ではcolumn-major、cuBLASが扱うDevice上ではrow-majorで表現し、Host-Device間のコピーついでに転置する transpose_copy
をこしらえてみました。
#include <cuda_runtime.h> #include <cublas_v2.h> #include <iostream> template<typename T> cudaError_t transpose_copy( T* dst, int ldDst, // コピー元 const T* src, int ldSrc, // コピー元 int m, int n, // 行, 列 cudaStream_t stream = nullptr) { cudaError_t status = cudaSuccess; for ( int i = 0; i < n && status == cudaSuccess; ++i ) { status = cudaMemcpy2DAsync(dst, sizeof(T), src, sizeof(T)*ldSrc, sizeof(T), m, cudaMemcpyDefault, stream); ++src; dst += ldDst; } return status; } int main() { const int LD = 8; // 2行3列のA と float A[LD][LD] = { { 1.0f, 1.1f, 1.2f }, { 1.3f, 1.4f, 1.5f } }; // 3行2列のB の積を float B[LD][LD] = { { 2.0f, 2.1f }, { 2.2f, 2.3f }, { 2.4f, 2.5f } }; // 2行2列のC に求める float C[LD][LD]; float* dA; float* dB; float* dC; cudaMalloc(&dA, LD*LD*sizeof(float)); cudaMalloc(&dB, LD*LD*sizeof(float)); cudaMalloc(&dC, LD*LD*sizeof(float)); // Host->Device transpose_copy(dA, LD, (float*)A, LD, 2, 3); transpose_copy(dB, LD, (float*)B, LD, 3, 2); // Cは0で埋めておく cudaMemset(dC, 0, LD*LD*sizeof(float)); cublasHandle_t handle; cublasCreate(&handle); // GEMM: C = αAB + βC (α=1, β=0 なら C = AB) float alpha = 1.0f; float beta = 0.0f; cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, 2, 2, 3, &alpha, dA, LD, dB, LD, &beta, dC, LD ); // Device->Host transpose_copy((float*)C, LD, dC, LD, 2, 2); cudaDeviceSynchronize(); for ( int y = 0; y < 2; ++y ) { for ( int x = 0; x < 2; ++x ) { std::cout << C[y][x] << ' '; } std::cout << std::endl; } cublasDestroy(handle); cudaFree(dA); cudaFree(dB); cudaFree(dC); }