gstore_fdwの圧縮オプション

昨年11月に、GPUメモリをSQLで読み書きするための新機能『gstore_fdw』というものを実装した。*1

これは、PostgreSQLのFDW(Foreign Data Wrapper)の機能を利用して、SELECTが実行された時にはGPUメモリから読み出した内容をPostgreSQLの内部データ表現に変換して出力、逆にINSERTやDELETEの実行時にはGPUメモリを更新するというものであった。

ただ、PG-Strom v2.0のリリースに向けて、少しクセの強い制限事項があったため、もう一度全体的なデザインを見直す事としてみた。
制限事項というのは

  • UPDATEには非対応
  • DELETEは条件なしのみ対応

というもので、要は『GPUメモリへの書込みは可能だが、完全に空の gstore_fdw にINSERTでデータを流し込むだけ』というシロモノである。
GPU側のデータストアは、PL/CUDA関数が単純配列であるかのようにアクセスしたいので行単位のMVCC可視性チェックは不可。
なので、行単位で異なるxmin/xmaxが混在するような状況はよろしくない。。。。という設計意図であった。

新しい実装では、UPDATEや条件付きDELETEを可能にするために、トランザクションのコミットまではホストメモリ上のバッファでデータを保持しておき、コミット時のコールバック処理で、DELETEした行を消去/INSERTした行を追加したGPUメモリのイメージを作成し、それをGPU側に転送するという構造に変更した。*2
そのため、GPU側には一時的に旧バージョン/新バージョンの2つのバッファが混在してしまう事になるが、あるトランザクションIDを持つ人からはどちらか片方しか見えないので、整合性は保たれている。

ちなみに、未コミット状態で gstore_fdw をPL/CUDA関数の引数に積もうとすると、CPU側で最新のイメージを作ってこれをGPUに転送してからPL/CUDA関数のGPUカーネルを起動する事になるので、あまりお勧めはしない。データのロードが終わったら、一度 commit してからPL/CUDA関数を実行する事をお勧めする。

この辺の改造に合わせて、もう一つ、内部的なデータ構造の持ち方を変えてみた。
圧縮フォーマットのサポートである。

圧縮オプション

外部テーブル(Foreign Table)にはFDWに固有のオプションを指定する事が可能で、gstore_fdwの場合は以下のオプションをサポートしている。最後の 'compression' オプションが今回追加したもの。

オプション名 対象 説明
pinning テーブル データを配置するGPUを番号で指定する。必須。
format テーブル GPU上のデータ形式を指定する。現在は 'pgstrom' (デフォルト)のみ対応。
compression カラム 可変長データを保存する時に圧縮を行う。現在は圧縮方式として 'pglz' のみ対応。

FDWの仕組み上、PostgreSQLとデータのやり取りをするときだけはデータ表現をPostgreSQLの内部形式に直してやる必要があるが、内部的なデータ表現はドライバが任意に決めてよい。

現在のデフォルトである 'pgstrom' 形式の場合、データ領域の先頭から単純配列のデータが順番に並んでいるような構造になっている。固定長データの場合はこれで十分であるし、列指向のデータ配置によってGPUのメモリバスの利用効率を最大限に引き出す事ができる。
可変長データ(textや配列など)の場合、少し工夫が必要で、データ領域の先頭から順に 32bit のインデックスがN個並んでいる。このインデックスは可変長データのバッファ(extra buffer)を指しており、可変長データの実体はこちらのバッファに保持されている事になる。

勘のいい人ならお気付きの通り、この構造は元々辞書圧縮を包含している。そのため、内容の重複する要素が多数存在する場合は一個の可変長データを複数の行から共有できるため、例えば複数のテーブルをJOINして正規化の崩れた(= 冗長度の高い)データであっても効率的に格納する事ができる。

ただ、扱う問題の種類によっては、それ以外にも効率的なデータの持ち方があり得る。
例えば、以前にPL/CUDAの検証で利用した化合物の特徴データや、あるいは機械学習の領域で使われるような特徴ベクトルは、基本的には疎行列である事が多い。つまり、ほとんどが 0 でちょびっと非零のデータが存在する。

こういったデータであれば、PostgreSQLが内部的に使用している LZ 圧縮形式でも十分に高い圧縮効率を実現できるはずである。

早速、GPU上のデータストアのサイズが圧縮の有無でどの程度変わるのかを調べてみる事にした。

-- GPUデータストア(圧縮オプションなし)
CREATE FOREIGN TABLE ft_flat (
    id int,
    signature smallint[]
)
SERVER gstore_fdw OPTIONS(pinning '0');

-- GPUデータストア(圧縮オプションあり)
CREATE FOREIGN TABLE ft_comp (
    id int,
    signature smallint[] OPTIONS (compression 'pglz'))
)
SERVER gstore_fdw OPTIONS(pinning '0');

データの詳細はお知らせできないが、これは実際にR&Dの現場で使われているもので、2048次元の特徴ベクトルの中に平均で10~20個程度の非零の要素が存在する。*3

まず、圧縮オプションなしの場合。

test=# INSERT INTO ft_flat(SELECT id,signature FROM source_db);
LOG:  alloc: preserved memory 3110899992 bytes
INSERT 0 1000000

3110899992 bytes = 約3.0GB のデバイスメモリを消費している。
確認してみると、クエリ実行前はこうだったのが

$ nvidia-smi
Thu Jan 11 11:01:08 2018
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 387.26                 Driver Version: 387.26                    |
|-------------------------------+----------------------+----------------------+
| 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:04:00.0 Off |                    0 |
| N/A   41C    P0    52W / 250W |    179MiB / 22912MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0     24010      C   ...bgworker: PG-Strom GPU memory keeper      161MiB |   <--- 初期状態
+-----------------------------------------------------------------------------+

こうなっている

$ nvidia-smi
Thu Jan 11 11:01:18 2018
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 387.26                 Driver Version: 387.26                    |
|-------------------------------+----------------------+----------------------+
| 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:04:00.0 Off |                    0 |
| N/A   41C    P0    52W / 250W |   3147MiB / 22912MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0     24010      C   ...bgworker: PG-Strom GPU memory keeper     3129MiB |   <--- デバイスメモリ確保済
+-----------------------------------------------------------------------------+

次に、圧縮オプションありの場合。

test=# INSERT INTO ft_comp(SELECT id,signature FROM source_db);
LOG:  alloc: preserved memory 158944912 bytes
INSERT 0 1000000

特徴ベクトルである signature が非常に sparse な構造を持っているため、158944912 bytes = 151MB にまでデータを圧縮できている事が分かる。

$ nvidia-smi
Thu Jan 11 11:18:04 2018
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 387.26                 Driver Version: 387.26                    |
|-------------------------------+----------------------+----------------------+
| 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:04:00.0 Off |                    0 |
| N/A   41C    P0    52W / 250W |    331MiB / 22912MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0     24199      C   ...bgworker: PG-Strom GPU memory keeper      313MiB |   <--- 少ないメモリ消費
+-----------------------------------------------------------------------------+

現状、最もデバイスメモリ搭載量の多い NVIDIA Tesla P40 でも 24GB(ECCに少し取られるので実際に使えるのは22GBちょい)しかデバイスメモリを持っていないので、この比率だと、非圧縮では100万件 → 1000万件にデータが増えた時点で破綻するが、圧縮の場合はまだまだ余裕がある事になる。
もちろんデータの特性によるので、これが一般化してどの程度データ圧縮が効くのかを断言はできないが、多くの場合に sparse な行列/ベクトルを扱うという事らしいので、一つの効果的な機能ではある事だろう。

PostgreSQLの可変長データ形式

PostgreSQLの可変長データ形式は元々圧縮をサポートしており、今回、gstore_fdw の圧縮オプションもこれを踏襲している。

可変長データを参照する時、まず先頭バイトの最下位ビットを参照する。
ここが1なら "short format" と呼ばれる形式で、データ長が1~126バイトの比較的短い可変長データの表現か、あるいは逆に、外部のtoastテーブルを参照するためのオブジェクトIDが格納される事になる。
最下位ビットが0の場合、ヘッダは32bitの "long format" と呼ばれる形式で、次のビットが0か1かで圧縮の有無を判定する。残りの30bitで可変長データの長さを表現するため、PostgreSQLにおける可変長データの最大長は1GBとなる。
圧縮ありの場合、ここの30bitに格納されているのは『圧縮後』のイメージの長さで、次の32bitで圧縮前のデータ長を、残りの領域が圧縮済みのイメージである。
このデータ形式だけを見ると、データ圧縮さえ行えば4GBまで表現できそうなものだが、PostgreSQLが内部的に使っているメモリアロケーション関数が 1GB でサイズチェックを行っている箇所が多々あり、実際にはうまく動かないだろう。

圧縮済みのデータも、実際にテーブルから読み出し、PostgreSQLの外部へ出力する際には透過的に展開され、ユーザからは圧縮/非圧縮の差があるようには見えない。
世の中には sparse な行列/ベクトルの表現に特化した専用のデータ構造もあるが、データ形式PostgreSQL と揃えておくことでの利便性を優先した。

*1:http://kaigai.hatenablog.com/entry/2017/11/12/092355

*2:同時に2つ以上のトランザクションが更新処理を走らせる場合、このデザインはうまく機能しないが、gstore_fdwの更新系処理は完全に排他的である。同時実行性能は低下するが、そもそも gstore_fdw の想定するワークロードではほとんど必要ではないので。

*3:作者のご好意で評価に使わせていただきました。ありがとうございます。