東方算程譚

επιστημηがヨタをこく、弾幕とは無縁のCUDAなタワゴト

COO-format から CSR-format への変換

m行n列の行列、成分数は m*n個。m,nが大きくなるとメモリの使用量がシャレになりません。けどもほとんどの成分が0であるなら、0成分を省略することでメモリ消費をぐっと抑えることができます。

ほとんどの成分が0のスカスカな行列を疎行列(Sparse matrix)と称します。その疎行列を対象にした各種演算をCUDAでやらせるライブラリがcuSPARSEです。

疎行列を表現するには (行index, 列index, 成分) の組を非0成分数だけ並べればよいですな。SoA(Structure of Array)なレイアウトだと

  int   rowInd[nnz]; // 行index
  int   colInd[nnz]; // 列index
  float val[nnz];    // 成分

こんなカンジ、nnzは非0成分の数(number of nonzero)。この形式をCOO-format(Coodinate format)っていいます。

くSPARSEが提供する演算には疎行列を引数に与えるんですけど、上記COO-formatのままでは引数に与えることができず、CSR-format(Compressed Sparse Row format)に変換せにゃなりません。その方法。

まずCOO-formatの各組を行(row),列(col)の順に昇順でソートします。行(row)の小さい順、行が同じなら列(col)の小さい順。当然成分(val)も連動して入れ替えを行います。

こんな元データが

f:id:Episteme:20161116191559p:plain

このようにソートされます。

f:id:Episteme:20161116191605p:plain

これが row-major でソートされた COO-format。

行indexは同じ値が連続しますね、row-majorにソートしたんだから当然ですけど。 同じ値が連続するってことはすなわち冗長であり圧縮できるってことです。

rowInd[nnz] を rowPtr[m+1] に圧縮した結果がコチラ。

f:id:Episteme:20161116191611p:plain

これがCSR-format(Compressed Sparse matrix Row format)。

COO-format から CSR-format への変換関数群は cuSPARSE の中に用意されています。

#include <cuda_runtime.h>
#include <cusparse.h>
#include <algorithm>
#include <numeric>
#include <random>
#include <iostream>
#include <iomanip>

using namespace std;

int main() {
  const int m   = 10; // 行数
  const int n   = 20; // 列数
  const int nnz = 30; // 非0成分数(numner of non-zero)

  // 疎行列のCOO(Coodinate format)
  int*   rowInd;
  int*   colInd;
  float* cooVal;
  cudaMallocManaged(&rowInd, nnz*sizeof(int));
  cudaMallocManaged(&colInd, nnz*sizeof(int));
  cudaMallocManaged(&cooVal, nnz*sizeof(float));

  { // 乱数で疎行列を生成する
    int inx[m*n];
    iota(begin(inx), end(inx), 0);
    shuffle(begin(inx), end(inx), mt19937());
    for ( int i = 0; i < nnz; ++i ) {
      rowInd[i] = inx[i]/n;
      colInd[i] = inx[i]%n;
      cooVal[i] = rowInd[i] + colInd[i]*0.01f;
    }
  }

  cout << "COO-format\n"
          "# ( rowInd, colInd, val)\n";
  for ( int i = 0; i< nnz; ++i ) {
    cout << setw(3) << i << " ( "
         << setw(3) << rowInd[i] << ',' 
         << setw(3) << colInd[i] << ", " 
         << cooVal[i] 
         << " )" << endl;
  }
  cout << endl;

  // sort後の成分領域を確保
  float* csrVal;
  cudaMallocManaged(&csrVal, nnz*sizeof(float));

  // rowPtr(圧縮されたrowInd)領域を確保 大きさは m+1
  int* rowPtr;
  cudaMallocManaged(&rowPtr, (m+1)*sizeof(int));

  // cuSPARSE 初期化
  cusparseHandle_t handle;
  cusparseCreate(&handle);

  // coosortに必要なバッファの大きさを求め、
  size_t bufferSize;
  cusparseXcoosort_bufferSizeExt(handle, m, n, nnz, rowInd, colInd, &bufferSize);
  // バッファを確保する
  void* sortBuffer;
  cudaMalloc(&sortBuffer, bufferSize);

  // permutation領域を確保し、
  int* permutation;
  cudaMalloc(&permutation, nnz*sizeof(int));
  // 0, 1, 2,... で埋める
  cusparseCreateIdentityPermutation(handle, nnz, permutation);

  // 行、列の順で昇順にソートし、
  cusparseXcoosortByRow(handle, m, n, nnz, rowInd, colInd, permutation, sortBuffer);
  cudaFree(sortBuffer);

  // 成分がソート順に対応するようpermutationを基に配置する
  cusparseSgthr(handle, nnz, cooVal, csrVal, permutation, CUSPARSE_INDEX_BASE_ZERO);
  cudaFree(permutation);

  // rowIndを圧縮してrwoPtrに求める
  cusparseXcoo2csr(handle, rowInd, nnz, m, rowPtr, CUSPARSE_INDEX_BASE_ZERO);
  cudaDeviceSynchronize();

  cout << "COO-format row-major sorted\n"
          "# ( rowInd, colInd, val)\n";
  for ( int i = 0; i< nnz; ++i ) {
    cout << setw(3) << i << " ( "
         << setw(3) << rowInd[i] << ',' 
         << setw(3) << colInd[i] << ", " 
         << csrVal[i] 
         << " )" << endl;
  }
  cout << endl;

  cout << "CSR-format\n"
          "# ( rowPtr, colInd, val)\n";
  for ( int i = 0; i< nnz; ++i ) {
    cout << setw(3) << i; ;
    if ( i <= m )
      cout << " ( " << setw(3) << rowPtr[i] << ',';
    else
      cout << "     ( ";
    cout << setw(3) << colInd[i] << ", " 
         << csrVal[i] 
         << " )" << endl;
  }
  cout << endl;

  cudaFree(rowInd);
  cudaFree(rowPtr);
  cudaFree(colInd);
  cudaFree(cooVal);
  cudaFree(csrVal);
  cusparseDestroy(handle);
  cudaDeviceReset();

}