転置しながらコピーする

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);

}