東方算程譚

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

cuSOLVER: 鶴亀算を解く

「つる と かめ があわせて3匹います。
 足の数はあわせて10本でした。
 問1: つる と かめ はそれぞれ何匹ですか?
 問2: 足の数があわせて8本なら、それぞれ何匹ですか?」

小学校の算数でやりましたよね、鶴亀算
コ難しくいえば 連立一次方程式の解を求める問題です。

行列で表現するとこんな。

f:id:Episteme:20161025212200p:plain

与えられた A と B から X を求めよ、ってわけ。

CUDA 7.0からこんな問題をCUDAで解いてくれるライブラリ: cuSOLVER がついてます。早速使ってみましたよ。

#include <cuda_runtime.h>
#include <cusolverDn.h> // dense LAPACK

#include <cassert>
#include <iostream>
using namespace std;

template<typename T>
inline size_t bytesof(unsigned int s) { return s * sizeof(T); }

template<typename T>
T* allocate(unsigned int size) {
  T* result = nullptr;
  cudaError_t status = cudaMalloc(&result, bytesof<T>(size));
  assert( status == cudaSuccess );
  return result;
}

int main() {
  cusolverStatus_t status;

  // dense LAPACK
  cusolverDnHandle_t handle;
  status = cusolverDnCreate(&handle);
  assert( status == CUSOLVER_STATUS_SUCCESS );

  int n = 2; // 2x2 正方行列

  float A[] = {  1.0f,  1.0f, 
                 2.0f,  4.0f };
  float* dA = allocate<float>(n*n);
  cudaMemcpy(dA, A, bytesof<float>(n*n), cudaMemcpyHostToDevice);
  int lda = 2;

  // 必要なバッファ量を求め、確保する
  int worksize;
  status = cusolverDnSgetrf_bufferSize(
             handle,
             n,   // 行
             n,   // 列
             dA,  // A
             lda, // Aのヨコハバ
             &worksize);
  assert( status == CUSOLVER_STATUS_SUCCESS );
#ifdef _DEBUG
  cout << "worksize = " << worksize << endl;
#endif

  float* workspace = allocate<float>(worksize);

  // 計算結果に関する情報
  int* devInfo = allocate<int>(1);
  // ピボット
  int* pivot = allocate<int>(n);

  // LU分解 : dAに結果が求まる(それとpivot)
  status = cusolverDnSgetrf(
             handle,
             n,   // 行
             n,   // 列
             dA,  // A
             lda, // Aのヨコハバ
             workspace,
             pivot,
             devInfo);
#ifdef _DEBUG
  int info;
  cudaMemcpy(&info, devInfo, sizeof(int), cudaMemcpyDeviceToHost);
  cout << "info = " << info << endl;
#endif
  assert( status == CUSOLVER_STATUS_SUCCESS );

  //         鶴と亀の総数, 足の総数
  float B[] = {  3.0f,       10.0f , 
                 3.0f,        8.0f };
  int nrhs = 2; // 問題数
  float* dB = allocate<float>(n*nrhs);
  cudaMemcpy(dB, B, bytesof<float>(n*nrhs), cudaMemcpyHostToDevice);
  int ldb = 2;

  // AX = B を解く (解XはBをoverrideする)
  status = cusolverDnSgetrs(
             handle,
             CUBLAS_OP_T,
             n,     // 行(=列)
             nrhs,  // 問題数
             dA,    // A
             lda,   // Aのヨコハバ
             pivot, // LU分解で得られたピボット
             dB,    // B
             ldb,   // Bのヨコハバ
             devInfo);
#ifdef _DEBUG
  cudaMemcpy(&info, devInfo, sizeof(int), cudaMemcpyDeviceToHost);
  cout << "info = " << info << endl;
#endif
  assert( status == CUSOLVER_STATUS_SUCCESS );

  // 結果を取得し、出力する
  float X[16];
  cudaMemcpy(X, dB, bytesof<float>(n*nrhs), cudaMemcpyDeviceToHost);
  for ( int i = 0; i < nrhs; ++i ) {
    float* q = B + i*2;
    float* a = X + i*2;
    cout << "総数 = " << q[0] << ", 足の数 = " << q[1]
         << "\t 解 : "
         <<"鶴 = " << a[0] << " , 亀 = " << a[1] <<endl;
  }

  cudaFree(workspace);
  cudaFree(dA);
  cudaFree(dB);
  cudaFree(devInfo);
  cudaFree(pivot);

  cusolverDnDestroy(handle);
}

実行結果:

f:id:Episteme:20161025212525p:plain

(original: 2016-06-03 #40)