gstore_fdw: GPUメモリをSQLで読み書き、そして…。

昨年、PGconf.ASIAで発表したPL/CUDAによる創薬ワークロードの高速化実験のテーマであるが、
kaigai.hatenablog.com


実測したベンチマークを見ると、奇妙な傾向が見てとれる。
このワークロードにおける計算量は「Qの行数×Dの行数」であるので、Dの行数が同じ1000万行であるならば、Qの行数が1000のケース(22.6s)に比べ、Qの行数が10のケース(13.4s)の実行時間はもっと顕著に短時間でなければならない。
計算量が1/100なのに、実行時間は半分弱にしかなっていない。

実はこれは、化合物同志の類似度を計算するための時間だけでなく、PL/CUDA関数に与える引数をセットアップするための時間に12秒程度を要しており、アムダールの法則を引用するまでもなく、類似度の計算を高速化するだけでは処理速度はこれ以上伸びないのである。

PL/CUDA関数の引数として行列(float4など固定長データでnullを含まない2次元配列型で代用)を使用するとして、これがGB単位の大きさになってくると、全くI/Oを伴わなくてもバッファからデータを読み出し、これをCPUで整形していくというのはそれなりに面倒な作業である。
また、PostgreSQL可変長データ形式の制約により、1GBを越える大きさのデータは複数に分割しなければ受け渡しができないという問題がある。

これらの問題に対応するため、PL/CUDA関数へのデータの受け渡しに新しいアプローチを考えてみた。FDWを用いる方法である。

FDW(Foreign Data Wrapper)を利用する事で、PostgreSQL管理外のデータソースをあたかもテーブルであるかのように扱う事ができる。最も一般的なケースでは、リモートのPostgreSQLOracleなどのテーブルを、外部表(Foreign Table)としてローカルのテーブルと同様に読み書きするといった使われ方をしている。また、PostgreSQLのcontribパッケージに含まれる file_fdw モジュールは、CSVファイルを外部表として扱う事が可能で、適切なFDWドライバさえ介在すれば、データソースはRDBMSのテーブルである必要はない。

PG-Stromの新機能 gstore_fdw は、GPUメモリ上に獲得した領域をSELECTやINSERTを用いて読み書きするためのFDWドライバである。

例えば、以下の外部表 ft はreal型の列を10個持っており、形式 pgstrom*1で、GPU番号0のデバイス上にメモリ領域を獲得する。

CREATE FOREIGN TABLE ft (
    id int,
    x0 real,
    x1 real,
    x2 real,
    x3 real,
    x4 real,
    x5 real,
    x6 real,
    x7 real,
    x8 real,
    x9 real
) SERVER gstore_fdw OPTIONS (pinning '0', format 'pgstrom');

適当にデータを流し込んでみる。どうやら420MB程度割り当てたようだ。

postgres=# INSERT INTO ft (SELECT x, 100*random(), 100*random(), 100*random(),
                                     100*random(), 100*random(), 100*random(),
                                     100*random(), 100*random(), 100*random(),
                                     100*random() FROM generate_series(1,10000000) x);
LOG:  alloc: preserved memory 440000320 bytes
INSERT 0 10000000

データロードの前。CUDA Contextの使用するリソースで171MBだけ消費されている。

$ nvidia-smi
Sun Nov 12 00:03:30 2017
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 384.81                 Driver Version: 384.81                    |
|-------------------------------+----------------------+----------------------+
| 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 |                    0 |
| N/A   36C    P0    52W / 250W |    171MiB / 22912MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

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

データロード後、171MB + 420MB = 591MB 消費されている。
この領域はPL/CUDA関数で参照する事ができるが、既にGPUメモリ上に留め置かれているので、呼び出し時に都度可変長データをセットアップする必要もなければ、データをGPUへ転送する必要もない。
これらは、冒頭の創薬ワークロードにおいてボトルネックとなっていた12秒に相当するものである。

$ nvidia-smi
Sun Nov 12 00:06:01 2017
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 384.81                 Driver Version: 384.81                    |
|-------------------------------+----------------------+----------------------+
| 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 |                    0 |
| N/A   36C    P0    51W / 250W |    591MiB / 22912MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

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

もちろん、外部テーブルであるのでSELECT文を用いて参照する事ができる。

postgres=# SELECT id,x0,x2,x3,x4 FROM ft LIMIT 10;
 id |   x0    |   x2    |   x3    |   x4
----+---------+---------+---------+---------
  1 | 27.1045 |  69.411 |  76.857 | 31.9964
  2 | 41.0264 |  74.365 | 15.4363 | 48.7949
  3 | 33.3703 | 62.7309 | 25.8918 | 8.37674
  4 | 72.0098 | 9.89505 | 41.4208 | 82.9163
  5 | 18.0448 | 49.7461 | 92.4098 | 16.0444
  6 | 25.1164 | 37.0391 | 87.8474 | 62.0111
  7 | 34.9195 | 62.4359 | 44.8145 | 8.56609
  8 | 86.4162 | 46.8959 | 36.1623 | 96.1458
  9 | 99.3865 | 73.5715 | 36.4256 | 36.2447
 10 | 70.0432 | 44.1223 | 32.4791 | 42.7989
(10 rows)

更新系の対応は以下の通りである。

  • 外部表が空の場合に限り、INSERTを実行可能。
  • UPDATEには非対応。
  • DELETEは条件句を伴わない場合に限り実行可能。
  • PostgreSQLが再起動した場合、内容は全て消える。

つまり、マスターとなるデータはPostgreSQL管理下のテーブルに保持しつつ、統計解析や機械学習のワークロードを実行する際に、SQLを用いて母集団を選択したり前処理を行った上で、GPUメモリにロードする事に特化したFDWモジュールである。

間違っても、飛んだら困るデータを格納してはならない。

gstore_fdwにロードしたデータは、内部的にはカラム毎にまとめて保持される。これは、GPUの特性上、同じタイミングで参照するデータ(つまり同じ列のデータ)が隣接領域に存在する方がメモリバスの性能を引き出しやすいため、列指向データ形式を選択したという理由がある。
もちろん、gstore_fdwの外部テーブルオプションには 'format' 指定が存在しているため、将来的には他のデータ形式(例えば、NumPyの行列表現など)を選択できるようにもしようと考えている。

可変長データの場合、先頭から(sizeof(uint) * N)バイトは、後段のExtra Buffer領域を参照するためのオフセット値になっている。
同じ内容のテキストが重複する場合など、複数のレコードが同一のExtra Buffer領域を参照する事となるので、結果として辞書圧縮が効いているのと同じ状態となる。
この設計は、複数のテーブルをJOINして正規化が崩れた状態のデータが gstore_fdw に投入される事を意図している。可変長データに関してはこのような内部表現を持つため、基本的には、冗長なデータを持つ事を過分に恐れる必要はない。

ちなみに、地味な特徴としてトランザクションにちゃんと(??)対応していたりもする。

postgres=# SELECT id,x0,x2,x3,x4 FROM ft LIMIT 10;
 id |   x0    |   x2    |   x3    |   x4
----+---------+---------+---------+---------
  1 | 27.1045 |  69.411 |  76.857 | 31.9964
  2 | 41.0264 |  74.365 | 15.4363 | 48.7949
  3 | 33.3703 | 62.7309 | 25.8918 | 8.37674
  :      :         :         :         :
  9 | 99.3865 | 73.5715 | 36.4256 | 36.2447
 10 | 70.0432 | 44.1223 | 32.4791 | 42.7989
(10 rows)

postgres=# BEGIN;
BEGIN
postgres=# DELETE FROM ft;
DELETE 10000000
postgres=# SAVEPOINT sv1;
SAVEPOINT
postgres=# INSERT INTO ft (SELECT x + 200, 100*random(), 100*random(), 100*random(),
                                           100*random(), 100*random(), 100*random(),
                                           100*random(), 100*random(), 100*random(),
                                           100*random() FROM generate_series(1,20) x);
LOG:  alloc: preserved memory 1200 bytes
INSERT 0 20

postgres=# SELECT id,x0,x2,x3,x4 FROM ft;
 id  |   x0    |   x2    |   x3    |   x4
-----+---------+---------+---------+---------
 201 | 48.2404 | 31.7422 | 38.0771 | 53.8396
 202 | 2.20366 | 56.2501 | 67.7503 |  43.756
 203 | 92.6225 | 61.5537 | 6.36337 | 80.5079
  :       :         :         :         :
 218 | 59.8691 | 88.3846 | 73.4542 | 99.4559
 219 | 2.13907 | 58.8645 | 23.5547 | 21.6422
 220 |  25.592 | 29.1767 | 24.9466 | 40.3255
(20 rows)

postgres=# ROLLBACK TO SAVEPOINT sv1;
ROLLBACK
postgres=# SELECT id,x0,x2,x3,x4 FROM ft;
 id | x0 | x2 | x3 | x4
----+----+----+----+----
(0 rows)

postgres=# ABORT;
LOG:  free: preserved memory at 0x10228800000
ROLLBACK
postgres=# SELECT id,x0,x2,x3,x4 FROM ft LIMIT 10;
 id |   x0    |   x2    |   x3    |   x4
----+---------+---------+---------+---------
  1 | 27.1045 |  69.411 |  76.857 | 31.9964
  2 | 41.0264 |  74.365 | 15.4363 | 48.7949
  3 | 33.3703 | 62.7309 | 25.8918 | 8.37674
  :      :         :         :         :
  9 | 99.3865 | 73.5715 | 36.4256 | 36.2447
 10 | 70.0432 | 44.1223 | 32.4791 | 42.7989
(10 rows)

では実際に、単純なPL/CUDA関数を作成し、実行してみる事にする。

CREATE OR REPLACE FUNCTION gstore_test(reggstore)
RETURNS float
AS
$$
#plcuda_begin
#plcuda_num_threads     gstore_fdw_height
  kern_data_store *kds = arg1.kds;
  int      i, ncols = kds->ncols;
  __shared__ double temp[1024];
  double   psum = 0.0;
  double   total;

  if (get_global_id() < kds->nitems)
  {
    for (i=0; i < ncols; i++)
    {
      kern_colmeta cmeta = kds->colmeta[i];
      Datum       datum;
      cl_bool     isnull;

      if (cmeta.atttypid == PG_FLOAT4OID)
      {
        datum = KDS_COLUMN_GET_VALUE(kds, i, get_global_id(),
                                     &isnull);
        if (!isnull)
          psum += __int_as_float(datum & 0xffffffff);
      }
    }
  }
  temp[get_local_id()] = psum;
  total = pgstromTotalSum(temp, get_local_size());
  if (get_global_id() == 0)
    retval->isnull = false;
  if (get_local_id() == 0)
    atomicAdd(&retval->value, total);
#plcuda_end
$$ LANGUAGE 'plcuda';

この人は、引数として与えられた gstore_fdw 表の中からreal型のカラムの内容をとにかく全部足して、その総和を返すというだけのPL/CUDA関数。
実行してみると、以下のような結果を返す。
平均値50のデータが10列×1000万件あるので、全部足したら50億ちょいというのは妥当なところだろう。

postgres=# select gstore_test('ft');
   gstore_test
------------------
 5000140834.18597
(1 row)

Time: 548.382 ms

統計解析・機械学習向けデータ管理の基盤として

さて、である。
ここまでの機能とCUDA APIこの辺の機能を一緒に使えば、統計解析・機械学習をやっている方にとって大変便利な機能として利用する事はできないだろうか?という所に考えが至った。

CUDAには以下のようなAPIがあり、あるプロセスが獲得したGPUバイスメモリを、別のプロセスと共有して利用する事ができる。

-- GPUデバイスメモリ上の領域の識別子を取得する
CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle,
                           CUdeviceptr dptr);

-- 取得した識別子を用いて、他プロセスに当該領域をマップする
CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr,
                            CUipcMemHandle handle,
                            unsigned int flags);

-- マップした領域を解放する
CUresult cuIpcCloseMemHandle(CUdeviceptr dptr);

この CUipcMemHandle というのは64byteのキーで、簡単にエクスポートする事ができる。

postgres=# select gstore_export_ipchandle('ft');
                        gstore_export_ipchandle
------------------------------------------------------------------------
 \xd00e220100000000963000000000000040df391a000000000000401a000000000000.
.000000000000000200000000000042000000000000002500d0c1ac00005c
(1 row)

gstore_fdwにロードしたデータ、つまり、既にGPUメモリに載った状態のデータをPythonやRなどスクリプトから直接使用する事ができれば、統計解析・機械学習に用いるデータの選択や前処理・後処理にSQLを使用する事ができる一方で、中核ロジックに関しては、データサイエンティストが使い慣れたPythonやRスクリプトを利用することできる。ゼロ・コピーなので、データの移動に関わるコストはほとんど省略する事ができる。

さらにSQLである事から、厳密な型チェックや整合性の検査を含める事ができ、これらは前処理における頭痛の種を相当程度に軽減してくれる事だろう。

加えて、PostgreSQLのData Federation機能と連携する事で、更に面白い構想に繋がってくる。
FDWの本来の使い方、つまり、リモートのPostgreSQLOracleデータベースのテーブルを読み書きする機能を用いて、データサイエンティストが作業で利用するワークステーションと、大量のデータを集積する役割のデータベースとを接続する。
(集積するデータの規模によっては、データベースを多段構成にするかもしれないが)

そうすると、一元管理によってデータの散逸を防ぎつつ、データサイエンティストが利用するデータの種類や母集団を選択したり、サマリ作成やデータ正規化といった前処理はワークステーションSQLを使用して実行*2し、機械学習のコア部分は使い慣れたRやPythonを利用して実行する事ができる。

これまで、In-database Analytics対応機能として、ユーザ定義関数(PL/CUDA)を使用して全てをSQLの世界で完結する事を考えていたが、適切なデータ形式に成形した上で、これを外部のスクリプトから共有・参照できるようにすれば、

  • 統計解析・機械学習ワークロードを実行するために、わざわざデータベースからデータをエクスポートさせない。
  • 中核アルゴリズムの前後で、母集団の選択やデータの前処理・後処理のためにSQLを使用した柔軟な記述を可能にする。

といった、従来からのコンセプトだけでなく

  • PythonやRといった、ユーザが使い慣れたツールを用いて In-database Analytics を実現できる。

という、新たな軸を打ち出せるような気がした。
この辺、自分はデータ解析を生業にしている訳ではないので、どの程度の価値がある機能なのかいま一つ想定し難いが、統計解析・機械学習といった領域で仕事をされている方の意見を伺ってみたいところではある。

今のところ、GPU上のデータ形式はPG-Stromの独自形式のみだが、Python向けにはNumPyとバイナリ互換な形式、R向けには行列とバイナリ互換な形式といった対応を拡充する事は検討している。

*1:現時点で唯一実装しているデータ形式

*2:PG-Stromが入っている前提なので、場合によってはSQL自体も高速化されることも