PL/CUDAを使ってロジスティック回帰分析を実装してみた

GPU PostgreSQL

PostgreSQL Advent Calendar 2018の6日目です。

PG-Stromはアナリティクス向けにPL/CUDAというユーザ定義SQL関数を実装する機能を持っており、SQL処理の中で計算ヘビーな部分をCUDA Cで記述したGPUプログラムで実行させるという事ができる。
SQL関数としてPL/CUDA関数を呼び出すと、その引数は自動的にGPUメモリへロードされ、またGPUプログラムの処理結果は自動的にSQLの処理に渡される事になる。したがって。プログラマはデータの管理をDBに任せ、DBからデータを出し入れする事なく、GPUによる並列計算を自然な形でクエリに埋め込む事が可能となる。

なお、ここで言う”計算ヘビー”というのは、JOINやGROUP BY、SORTなどがCPUインテンシブな処理である、等々とは少し計算量の桁が違う話である事には留意していただきたい。
例えば、二年前のこのエントリは創薬領域での類似化合物探索のワークロードで、DBに登録された化合物同士の類似度をGPUで計算しているが、ｍその計算すべき類似度の組合せは100億通りもあった。
kaigai.hatenablog.com

で、時々お問合せを頂いていたのが『PL/CUDAでSQL関数を定義するにあたって、何かサンプルプログラムはないか？』という質問。
確かに、パブリックになっているPL/CUDA関数のサンプルはなかったので、折角なのでPGconf.ASIA 2018のLTあたりのネタにすべく、比較的簡単なアルゴリズムをPL/CUDAで実装してみた。

ソースコードはこちら。
https://github.com/heterodb/toybox/tree/master/logistic_regression

ロジスティック回帰分析

今回のテーマに選んでみたのは、ロジスティック回帰分析。

このアルゴリズムは、説明変数 x_j (j=1,...,m)とパラメータw_j (j=0,...,m)を、累積確率関数であるロジスティック関数\sigma(\alpha)=\frac{1}{1-e^{-\alpha}}に通して二値分類を行うためのもの。
学習フェーズでは、トレーニングセットのデータを用いてパラメータを推定する事となり、教師データである従属変数t_i (i=1,...,n)が実現する確率を最大化するパラメータを推定する必要がある。

なお、\alpha = w_0 + \sum_{j=1}^{m} w_j x_jとし、従属変数t_i = 1 である場合の確率をP(X) = \sigma(\alpha)としている。

あまりこの辺を踏み込んで書けるほど数学に強いわけではないので、詳しくはこちらの本をご参考に（逃げた）

ITエンジニアのための機械学習理論入門

ITエンジニアのための機械学習理論入門

最小二乗法の場合とは異なり、ロジスティック回帰分析のパラメータは解析的に数式を解いて求める事はできない。
そのため、実際に説明変数と従属変数から成るトレーニングセットをモデルに当てはめて計算し、その誤差を徐々に小さくしていくよう繰り返し計算を行っていく以外にパラメータを推定する方法がなく、それゆえに計算量が大きくなってしまう。

途中の説明は書籍に任せるとして、最終的には、以下の行列式を繰り返し解いてパラメータを更新していく事になる。

w_{new} = w_{old} - (\Phi^{T}R\Phi)^{-1}\Phi^{T}(z - t)

各変数の意味は以下の通り
トレーニングセット（従属変数）：
t=(t_1,...,t_n)
トレーニングセット（説明変数）：
\Phi = \left( \begin{array}{cccc} 1 & x_{11} & \ldots & x_{1m} \\ 1 & x_{21} & \ldots & x_{2m} \\ \vdots & \vdots & \ddots & \vdots \\ 1 & x_{n1} & \ldots & x_{nm} \end{array} \right)
累積確率関数：
z = (z_1,...,z_n)（但し、z_i = \sigma(w^T \Phi_i)）と、
それを対角成分に並べたもの：
R = diag[z_1(1-z_1), ..., z_n(1-z_n)]

仮定としては、説明変数の数はそれほど多くない（数個～数百程度）に対し、トレーニングセットの数は数百万～になり得る。
そうすると、計算量として比較的大きくなるのはi = 1,...,n方向の行列積で、ここは数千コアを持つGPUの得意分野になってくる。

PL/CUDAでの実装を掘り下げる

前振りが長くなったが、ソースコードを見てみる事にする。
https://github.com/heterodb/toybox/tree/master/logistic_regression

エントリポイント

まず、PL/CUDA関数のエントリポイントは#plcuda_beginと#plcuda_endで挟まれた区間である。

CREATE OR REPLACE FUNCTION
logregr_train(reggstore,      -- source table (only Gstore foreign table)
              smallint,       -- column number of dependent variable
              smallint[],     -- columns number of independent variables
              int = 20,       -- max number of iteration
              real = 0.0001)  -- threashold to stop iteration
RETURNS real[]
AS $$
   :
#plcuda_begin
   /*
    * ここにエントリポイントの処理を記述する
    */
#plcuda_end
$$ LANGUAGE 'plcuda' STRICT;

ロジスティック回帰の学習関数logregr_trainは、第１引数がトレーニングセットを含むGstore_Fdw外部テーブル*1への参照で、第２引数は従属変数の列番号、第３引数は説明変数の列番号（配列）で、残りはパラメータ推定を繰り返す回数の最大値と、繰り返しを打ち切る閾値*2の設定。

エントリポイント内では、smallintの引数はshort型で、realの引数はfloat型でというように、それぞれSQLの型に対応するCUDA Cのデータ型でarg1、arg2．．．というように引数を参照する事ができる。

例えば、以下のコードはsmallint[]の配列であるarg3を参照し、指定された列がreal型であるかどうかをチェックしている部分である。
型が一致しない場合、EEXITマクロでスクリプトを終了させ、これはPostgreSQL側でエラーとしてユーザへ報告される。

    /* set of simple array of independent data */
    if (!VALIDATE_ARRAY_VECTOR_TYPE_STRICT(arg3, PG_INT2OID))
        EEXIT("independent variables must be 'smallint[]'");
    iattrs = (VectorTypeShort *)arg3;
    rc = cudaMallocManaged(&Xp, sizeof(float *) * iattrs->height);
    if (rc != cudaSuccess)
        CUEXIT(rc, "failed on cudaMallocManaged");
    for (i=0; i < iattrs->height; i++)
    {
        j = iattrs->values[i];
        if (j <= 0 || j > kds_h->ncols)
            EEXIT("independent variable is out of range");
        cmeta = &kds_h->colmeta[j-1];
        if (cmeta->atttypid != PG_FLOAT4OID)
            EEXIT("independent variables must be 'real'");
        Xp[i] = (float *)((char *)kds_d + __kds_unpack(cmeta->va_offset));
    }

既にGPUへロード済みのデータをreggstore型の引数で参照する場合だけは少々特殊で、以下のように定義された GstoreIpcMapping構造体へのポインタとなる。mapメンバはGPUデバイスメモリのポインタなので、ホストコードで触ってはいけない。*3

typedef struct
{
    GstoreIpcHandle h;      /* unique identifier of Gstore_Fdw */
    void           *map;    /* mapped device address
                             * in the CUDA program context */
} GstoreIpcMapping;
宣言部

CUDA CでGPUカーネルを呼び出す場合、例えばmy_kernel<<<10,1024>>>(x, y, z)というように、起動すべきブロック数／スレッド数を指定してデバイス関数をコールする。
つまり、エントリポイント内にベタっと処理を記述するだけで不十分で、少なくともGPUカーネル関数をどこかに定義しなければPL/CUDA関数を利用する意味はあまりない。
それを行うのが、#plcuda_decl以降のブロックで、ここに記述したホスト関数／デバイス関数をエントリポイントから呼び出すというのが一般的なPL/CUDA関数の実装という事になる。

例えば、以下のlogregt_update_Zというデバイス関数は、上記の累積確率関数：z = (z_1,...,z_n) の値を格納するVectorTypeFloat型((PostgreSQLのreal[]型とバイナリ互換な単純配列))のバッファを更新するためのもので、これを（ブロック数×ブロックサイズ）個のスレッドで並列に実行する。
例えばブロック数40、ブロックサイズ1024の場合、get_global_id()は0～40959の間の値を返し、全てのスレッドが処理を完了した時にはZ->values[]配列には\frac{1}{1-e^{-\alpha}}の計算結果が格納されている事になる。

KERNEL_FUNCTION_MAXTHREADS(void)
logregr_update_Z(VectorTypeFloat *Z,
                 cl_float  **Xp,
                 VectorTypeFloat *W)
{
    cl_uint        nitems = Z->height;
    cl_uint        width = W->height;
    cl_uint        i, j;

    for (i = get_global_id(); i < nitems; i += get_global_size())
    {
        cl_float    x, w, sum = 0.0;
        for (j=0; j < width; j++)
        {
            w = W->values[j];
            x = (j == 0 ? 1.0 : Xp[j-1][i]);
            sum += w * x;
        }
        Z->values[i] = 1.0 / (1.0 + __expf(-sum));
    }
}

このように定義されたデバイス関数を以下のように繰り返し呼び出す事で、徐々にロジスティック回帰のパラメータを更新し、目的とする値に近付くというステップを踏む。

    for (loop=0; loop < nloops; loop++)
    {
               :
        /* compute Z vector */
        logregr_update_Z<<<gridSz_Z, blockSz>>>(Z, Xp, W);
               :
        /* compute P matrix */
        logregr_update_P<<<gridSz_P, blockSz>>>(Preg, Xp, width, Z);
               :
        cudaStreamSynchronize(NULL);
               :
    }
CUDA用ライブラリの利用

ところで、この繰り返し計算の過程では、単純な行列・ベクトル積だけでなく、一ヵ所逆行列を計算する箇所があった事を思い出してほしい。
このように典型的な計算であれば、CUDA Toolkitでライブラリが提供されている場合があり、逆行列の計算であればBlasのCUDA版であるcuBlasを使う事ができる。

ライブラリ関数を使うには、まず普通に#include ...でヘッダファイルをインクルードすると共に、GPUバイナリをビルドする時にリンクするライブラリを指定する必要がある。

以下のように#plcuda_libraryでライブラリ名を指定すると、

#plcuda_library cublas
#plcuda_decl
#include <cublas_v2.h>
#include "cuda_matrix.h"
     :

これはGPUバイナリをビルドする時に

% nvcc .... -lcublas

と展開され、ライブラリ関数をリンクする。

若干、引数が多くてアレなのだが、今回の実装でも以下のcuBlas関数を呼び出して逆行列を計算している。

        /* compute P-inverse */
        status = cublasDgetrfBatched(handle,
                                     width,
                                     Preg,
                                     width,
                                     Pivot,
                                     Info,
                                     1);
        if (status != CUBLAS_STATUS_SUCCESS)
            EEXIT("failed on cublasSgetrfBatched: %d", (int)status);
        status = cublasDgetriBatched(handle,
                                     width,
                                     Preg,
                                     width,
                                     Pivot,
                                     Pinv,
                                     width,
                                     Info,
                                     1);
        if (status != CUBLAS_STATUS_SUCCESS)
            EEXIT("failed on cublasSgetriBatched: %d", (int)status);

CPU実装（MADLib）との比較

最後に、CPU実装と比べてどれくらい速くなっているのかを調べるため、MADLibのロジスティック回帰と性能値を比較してみる事にした。

とは言っても、これを使って何を分類するか、というのはあまり興味がないので、ひとまずランダムに4000万件程度のデータを生成し、人為的に1+2x_1-3x_2+x_3+0.5x_4 > 0である場合を真(true)、それ以外を偽(false)とするテストデータを作って、それを学習させてみる事にする。

postgres=# CREATE TABLE logreg (
                    t  bool,
                    x1 float,
                    x2 float,
                    x3 float,
                   x4 float );
CREATE TABLE

postgres=# INSERT INTO logreg (SELECT (1.0+2.0*x1-3.0*x2+x3+0.5*x4) > 0 t, x1, x2, x3, x4
                                 FROM (SELECT random() x1,
                                              random() x2,
                                              random() x3,
                                              random() x4
                                         FROM generate_series(1,40000000)) x);
INSERT 0 40000000

次に、このテストデータをGPUデバイスメモリにロードする。
PL/CUDA関数はreal型を前提にしているので、説明変数のデータ型はrealへキャストする。

postgres=# CREATE FOREIGN TABLE ft (
               t   bool,
               x1  real,
               x2  real,
               x3  real,
               x4  real
           ) SERVER gstore_fdw
             OPTIONS (pinning '0');
CREATE FOREIGN TABLE

postgres=# INSERT INTO ft (SELECT * FROM logreg);
INSERT 0 40000000

bool(1byte) + float(4bytes)*4 = 17bytes
17bytes * 40M行 = 約680MB
デバイス管理用に必要な約120MBと併せて、計800MBのデータがGPUにロードされているのが分かる。

[kaigai@saba src]$ nvidia-smi
Thu Dec  6 12:10:56 2018
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 410.72       Driver Version: 410.72       CUDA Version: 10.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla P40           Off  | 00000000:02:00.0 Off |                  N/A |
| N/A   42C    P0    52W / 250W |    817MiB / 22919MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0     27650      C   ...bgworker: PG-Strom GPU memory keeper      807MiB |
+-----------------------------------------------------------------------------+

PL/CUDA関数を呼び出す。attnum_of()は指定された列番号を返すユーティリティ関数。

postgres=# SELECT logregr_train('ft',
                                attnum_of('ft','t'),
                                attnums_of('ft','{x1,x2,x3,x4}'));
              logregr_train
------------------------------------------
 {3376.4,6752.71,-10129.1,3376.3,1688.27}
(1 row)

Time: 3647.059 ms (00:03.647)

約3.6secで完了。パラメータである w ベクトルを推定しているので、5個の要素が含まれている。

次に、MADLibのロジスティック回帰のトレーニング。（関数名や引数に与えるパラメータはこれを参考にした）

postgres=# SELECT madlib.logregr_train('logreg','hoge', 't','ARRAY[1,x1,x2,x3,x4]',NULL, 20);
 logregr_train
---------------

(1 row)

Time: 1301307.361 ms (21:41.307)

21分41秒。元々の実装がそこまで最適化されているか微妙なところではあるものの、広く使われている実装と比べて350倍速く実行できたのはGPUの面目躍如と言ったところ。

postgres=# SELECT coef FROM hoge;
                                          coef
------------------------------------------------------
 {3041.82722783601,6083.57794939209,-9125.44857123801,3041.73992459095,1520.98287953044}
(1 row)

一応、学習結果であるパラメータを比べてみると

俺々実装 MADLib
w0 3376.4 3041.83
w1 6752.71 6083.58
w2 -10129.1 -9125.45
w3 3376.3 3041.74
w4 1688.27 1520.98

それほど遜色ない値となっている。

これを、説明変数と学習パラメータを引数に取るlogregr_predict関数に与えて（人為的に作ったデータなので）きちんと分類できている事を示そうと思ったが、残念ながらgit addを忘れてgit cleanしてしまい、推論系の実装は闇に葬られてしまった。

こちらはPGconf.ASIAのLTでご報告しようと思うので、少々お待ちを。

明日は @hmatsu47 の『Aurora PostgreSQL 互換版（9.6.9）でフェイルオーバー時にバッファ／キャッシュの生きのこり方を雑に確かめてみた』です。

*1:Gstore_Fdw: GPUデバイスメモリの読書きを可能にするFDW

*2:パラメータの変動は\frac{\|w_{new} - w_{old}\|^2}{\|w_{old}\|^2}で計算します

*3:Managed MemoryというPage Faultハンドラを使ってCPU/GPU間でページを入れ替える機能があるが、CUDAがプロセス間通信に対応してくれていない。早く何とかならないものか。