カスタムロジックをWHERE句で使う - KaiGaiの俺メモ

カスタムロジックをWHERE句で使う

しばらくSSD-to-GPUダイレクトSQL実行の開発にどっぷり時間を突っ込んでいたので、久々にPL/CUDAネタ。

この辺のネタや、
kaigai.hatenablog.com

この辺のネタで
kaigai.hatenablog.com

紹介したように、PG-Stromはユーザ定義関数をCUDA Cで記述するための機能（PL/CUDA）を持っており、これを使えば、データベースから読み出したデータをGPUに流し込み、GPU上でカスタムのロジックを実行した後、結果をまたSQLの世界へ戻すという事ができる。

この仕組みはPostgreSQLの手続き言語ハンドラの機能を用いて実装されており、ユーザ定義のPL/CUDA関数が呼び出される毎に、手続き言語ハンドラが以下の処理を行う。

  1. ユーザ記述のCUDA Cコードブロックをテンプレートに埋め込んで、ビルド可能なソースコードを作成。
  2. NVRTC(NVIDIA Run-Time Compiler)を用いてソースコードをビルド。既にビルド済みのバイナリがキャッシュされていればそれを使用する。
  3. 関数の引数をGPUに転送する。（場合によっては数百MBの配列データである事も）
  4. ユーザの指定したスレッド数、ブロックサイズ、共有メモリ割当てなどのパラメータに従ってGPU Kernelを起動する。
  5. 予め設定しておいた結果バッファの内容をホスト側に書き戻す。
  6. 手続き言語ハンドラの戻り値として、GPUから戻された結果の値を返す。

これは、一個のGPU Kernelが巨大なマトリックスを引数に取り、数千～数万スレッドが協調動作をしながら処理を進めるようなタイプのワークロードに望ましい。典型的な利用シーンにおいては、SQL構文の中で数回呼び出され、一回あたりの応答時間が少なくとも数秒はあるといったタイプの呼び出し方である。

実際、以前に創薬の領域で類似化合物検索のロジックをPL/CUDA関数で作成した時も、PL/CUDA関数であるところのknn_gpu_similarity()が呼び出されるのは、以下のクエリ実行中に一回だけであった*1。

PREPARE knn_sim_rand_10m_gpu_v2(int)	-- arg1:@k-value
AS
SELECT float4_as_int4(key_id) key_id, similarity
  FROM matrix_unnest((SELECT rbind(knn_gpu_similarity($1,Q.matrix,D.matrix))
                        FROM (SELECT cbind(array_matrix(id),
                                           array_matrix(bitmap)) matrix
                                FROM finger_print_query
                               LIMIT 99999) Q,
                             (SELECT matrix
                              FROM finger_print_10m_matrix) D))
    AS sim(key_id real, similarity real)
ORDER BY similarity DESC
LIMIT 1000;

ただ、こういった前提を置いてしまうと、例えば数億件のレコードをスキャンし、GPUで条件句を処理するといったシチュエーションでカスタムのロジックを使うのが難しくなってしまう。

そこで、PG-Stromのコードジェネレータに手を加え、いくつかの制約を満たす場合に限り、カスタムのPL/CUDA関数をWHERE句などで利用できるようにしてみた。

  • PL/CUDA関数が可変長データ(varlena)を返さない
  • PL/CUDA関数は単一のコードブロックしか持たない（prep/postブロックを持たない）
  • ブロックサイズは1、スレッド数も1
  • GPU shared memoryを要求しない
  • 作業バッファ(workbuf)、結果バッファ(resultbuf)を要求しない

これだと全くマルチスレッドを活用できないように見えるが、PG-Stromの場合、テーブルから64MB毎にデータブロックを切り出し、各レコード毎にGPUスレッドを立ち上げて並列にWHERE句を処理するため、結果として数千～数万スレッドで並列処理する事になる。

簡単なPL/CUDA関数を定義してみる。int型の引数を4つ取り、その合計を返すだけというシロモノ。

create function test_1(int,int,int,int)
returns int
as
$$
#plcuda_begin
  if (!arg1.isnull && !arg2.isnull && !arg3.isnull && !arg4.isnull)
  {
    retval->isnull = false;
    retval->value = arg1.value + arg2.value + arg3.value + arg4.value;
  }
#plcuda_end
$$ language 'plcuda';

で、WHERE句の中でこの関数を使ってみる。

postgres=# explain verbose select * from s0 where test_1(aid,bid,cid,did) < 50000;
                                      QUERY PLAN
--------------------------------------------------------------------------------------
 Custom Scan (GpuScan) on public.s0  (cost=24031.02..223694.11 rows=2666680 width=40)
   Output: id, cat, aid, bid, cid, did, eid, fid, gid, hid
   GPU Filter: (test_1(s0.aid, s0.bid, s0.cid, s0.did) < 50000)
   Kernel Source: /opt/pgdata/base/pgsql_tmp/pgsql_tmp_strom_2118.2.gpu
(4 rows)

自動生成されたGPUコード*2を眺めてみると、WHERE句のCUDA Cに落とし込んだgpuscan_quals_evalの中で、PL/CUDA関数のロジックを実装した pgfn_plcuda_86452() が呼び出されている。
この関数の実行結果は、そのまま大小比較の演算 pgfn_int4lt に渡され、KPARAM_0 (定数 50000) との比較を行っている。

STATIC_FUNCTION(pg_int4_t)
pgfn_plcuda_86452(kern_context *kcxt,
                  pg_int4_t arg1, pg_int4_t arg2,
                  pg_int4_t arg3, pg_int4_t arg4)
{
  pg_int4_t __retval = { true, 0 };
  pg_int4_t *retval = &__retval;
  /* #plcuda_begin */
  if (!arg1.isnull && !arg2.isnull && !arg3.isnull && !arg4.isnull)
  {
    retval->isnull = false;
    retval->value = arg1.value + arg2.value + arg3.value + arg4.value;
  }

  /* #plcuda_end */
  return __retval;
}

STATIC_FUNCTION(cl_bool)
gpuscan_quals_eval(kern_context *kcxt,
                   kern_data_store *kds,
                   ItemPointerData *t_self,
                   HeapTupleHeaderData *htup)
{
  pg_int4_t KPARAM_0 = pg_int4_param(kcxt,0);
  pg_int4_t KVAR_3;
  pg_int4_t KVAR_4;
  pg_int4_t KVAR_5;
  pg_int4_t KVAR_6;
  char *addr;

  assert(htup != NULL);
  EXTRACT_HEAP_TUPLE_BEGIN(addr, kds, htup);
  EXTRACT_HEAP_TUPLE_NEXT(addr);
  EXTRACT_HEAP_TUPLE_NEXT(addr);
  KVAR_3 = pg_int4_datum_ref(kcxt,addr);
  EXTRACT_HEAP_TUPLE_NEXT(addr);
  KVAR_4 = pg_int4_datum_ref(kcxt,addr);
  EXTRACT_HEAP_TUPLE_NEXT(addr);
  KVAR_5 = pg_int4_datum_ref(kcxt,addr);
  EXTRACT_HEAP_TUPLE_NEXT(addr);
  KVAR_6 = pg_int4_datum_ref(kcxt,addr);
  EXTRACT_HEAP_TUPLE_END();

  return EVAL(pgfn_int4lt(kcxt, pgfn_plcuda_86452(kcxt, KVAR_3, KVAR_4,
                                                        KVAR_5, KVAR_6), KPARAM_0));
}

今回の例は非常に単純な計算をPL/CUDAで実装したもので、この程度の条件句の評価であればあまり恩恵はないだろう。
ただ、同じ仕組みを使う事で、機械学習による"推論"をWHERE句の条件評価に混ぜて使う事ができるようにもなるだろう。

そうすると、種々多様なソースから上がってきたデータをPostgreSQLに蓄積し、これをSSD-to-GPUダイレクトSQL実行を使ってサマリ作成・前処理を行った上で、従来型のPL/CUDA関数で実装した機械学習ロジックによりパターンを学習する。さらに、今回の新しいPL/CUDA関数の使い方で、前ステップで生成した学習モデルを引数とする"推論系"の関数をWHERE句に加えてテーブルをスキャンする事ができるようになる。

応用としては、例えば平常パターンで学習させたモデルを元に、異常値くさいレコードを検出するアノマリー検知など。
(NoSQL系DBやらあるとはいえ、)金融系なら間違いなく決済データをRDBMSで保持しているだろう。そうすると、クレジットカードの不正利用や、オレオレ詐欺など「普段と違う」振る舞いの検出をSQLを用いてin-databaseで実行できるようになる、という事を意味する。

*1:その一回の呼び出しが処理時間の大半を占めるというワークロードであるワケだが

*2:見やすくするため、一部改行追加