TargetListの計算をGPUで行う。 - KaiGaiの俺メモ

TargetListの計算をGPUで行う。

正月休みの宿題だった機能を実装できた。(注:ちゃんと動くとは言っていない)

PG-Stromを使って数式の評価をGPUにオフロードする場合、WHERE句やJOIN..ON句のオフロードには対応していたものの、TargetListに複雑な演算式を含む場合、これは完全にCPU側で処理せざるを得なかった。
これはPostgreSQLオプティマイザがこれらの数式をPlanノードにアタッチするタイミングの問題で、別に本質的な問題がある訳ではない。(実際、コミュニティの側でも改善に向けた動きはある)

例えば、以下のようなクエリは数式の評価をGPUで実行できた。

-- 任意の二点間の4次元空間上の距離が 10 未満であるペアを出力
SELECT a.id a_id, b.id b_id
  FROM test a, test b
 WHERE a.id > b.id AND sqrt((a.x1 - b.x1)^2 +
                            (a.x2 - b.x2)^2 +
                            (a.x3 - b.x3)^2 +
                            (a.x4 - b.x4)^2) < 10;

EXPLAIN文で確認すると、以下のように、GpuNestLoopのJoinQualに上記のWHERE句が収まっていることが分かる。

# EXPLAIN
  SELECT a.id a_id, b.id b_id
    FROM test a, test b
   WHERE a.id > b.id AND sqrt((a.x1 - b.x1)^2 +
                              (a.x2 - b.x2)^2 +
                              (a.x3 - b.x3)^2 +
                              (a.x4 - b.x4)^2) < 10;

                        QUERY PLAN

------------------------------------------------------------
 Custom Scan (GpuJoin)  (cost=14050655.37..25161766.48 rows=1111111111 width=8)
   Bulkload: On (density: 100.00%)
   Depth 1: GpuNestLoop, JoinQual: ((id > id) AND (sqrt(((((((x1 - x1))::double precision ^ '2'::double precision) + (((x2 - x2))::double precision ^ '2'::double precision)) + (((x3 - x3))::double precision ^ '2'::double precision)) + (((x4 - x4))::double precision ^ '2'::double precision))) < '10'::double precision))
            Nrows (in/out: 1111111.13%), KDS-Heap (size: 8.11MB, nbatches: 1)
   ->  Custom Scan (BulkScan) on test a  (cost=0.00..1637.00 rows=100000 width=20)
   ->  Seq Scan on test b  (cost=0.00..1637.00 rows=100000 width=20)
(6 rows)

なお、テーブル test は 4次元空間上の特定の点を表現するとして、以下のように定義。

postgres=# \d test
     Table "public.test"
 Column |  Type   | Modifiers
--------+---------+-----------
 id     | integer | not null
 x1     | real    |
 x2     | real    |
 x3     | real    |
 x4     | real    |
Indexes:
    "test_pkey" PRIMARY KEY, btree (id)

しかし、同じ式を含んでいても、このパターンだと無理だった。

-- 任意の二点間の4次元空間上の距離を算出してIDのペアと共に出力
SELECT a.id a_id, b.id b_id, sqrt((a.x1 - b.x1)^2 +
                                  (a.x2 - b.x2)^2 +
                                  (a.x3 - b.x3)^2 +
                                  (a.x4 - b.x4)^2) dist
  FROM test a, test b
 WHERE a.id > b.id;

同様にEXPLAIN文で確認すると、a.id > b.id の部分しかGPUで実行されていない。計算ロジックとして重いのは距離計算の部分であるにも関わらず、だ。

# EXPLAIN
  SELECT a.id a_id, b.id b_id, sqrt((a.x1 - b.x1)^2 +
                                    (a.x2 - b.x2)^2 +
                                    (a.x3 - b.x3)^2 +
                                    (a.x4 - b.x4)^2) dist
            FROM test a, test b
           WHERE a.id > b.id;
                                    QUERY PLAN
-----------------------------------------------------------------------------------
 Custom Scan (GpuJoin)  (cost=776405.37..167443072.02 rows=3333333333 width=40)
   Bulkload: On (density: 100.00%)
   Depth 1: GpuNestLoop, JoinQual: (id > id)
            Nrows (in/out: 3333333.20%), KDS-Heap (size: 8.11MB, nbatches: 1)
   ->  Custom Scan (BulkScan) on test a  (cost=0.00..1637.00 rows=100000 width=20)
   ->  Seq Scan on test b  (cost=0.00..1637.00 rows=100000 width=20)
(6 rows)

では、どのタイミングで計算が行われているのか?これはEXPLAIN VERBOSEで確認することができる。

# EXPLAIN VERBOSE
  SELECT a.id a_id, b.id b_id, sqrt((a.x1 - b.x1)^2 +
                                    (a.x2 - b.x2)^2 +
                                    (a.x3 - b.x3)^2 +
                                    (a.x4 - b.x4)^2) dist
    FROM test a, test b
   WHERE a.id > b.id;

                        QUERY PLAN

-------------------------------------------------------------------------
 Custom Scan (GpuJoin)  (cost=776405.37..167443072.02 rows=3333333333 width=40)
   Output: a.id, b.id, sqrt(((((((a.x1 - b.x1))::double precision ^ '2'::double precision) + (((a.x2 - b.x2))::double precision ^ '2'::double precision)) + (((a.x3 - b.x3))::double precision ^ '2'::double precision)) + (((a.x4 - b.x4))::double precision ^ '2'::double precision)))
   Pseudo Scan: a.id::integer, a.x1::real, a.x2::real, a.x3::real, a.x4::real, b.id::integer, b.x1::real, b.x2:
:real, b.x3::real, b.x4::real
   Bulkload: On (density: 100.00%)
   Depth 1: GpuNestLoop, JoinQual: (a.id > b.id)
            Nrows (in/out: 3333333.20%), KDS-Heap (size: 8.11MB, nbatches: 1)
   Features: format: tuple-slot, bulkload: supported
   Kernel Source: /opt/pgsql/base/pgsql_tmp/pgsql_tmp_strom_31576.1.gpu
   ->  Custom Scan (BulkScan) on public.test a  (cost=0.00..1637.00 rows=100000 width=20)
         Output: a.id, a.x1, a.x2, a.x3, a.x4
         Features: format: tuple-slot, bulkload: supported
   ->  Seq Scan on public.test b  (cost=0.00..1637.00 rows=100000 width=20)
         Output: b.id, b.x1, b.x2, b.x3, b.x4
(13 rows)

やや煩雑な出力結果となってしまっているが、見るべきポイントは GpuJoin の Output: と Pseudo Scan: の差分。
GpuJoinはGPUでのJOIN処理結果として、Pseudo Scan: の定義に従ってレコードを生成し、これをCPU側へ返却する。
一方、それを受け取ったGpuJoinノードは、より上位へ結果を返すために Output: の定義に従ってレコードの内容を書き換える。これをProjection処理と呼ぶ。
Pseudo Scan == Output であれば実際にはProjectionは必要ではないが、上記のように複雑な計算式を伴うケースであれば、ScanやJoin処理よりもProjectionが処理時間の支配項となってしまう。


では、どうすればよいか?
CPUでのProjectionが支配項となってしまうなら、Projectionが発生しないよう、予めGPU側で計算した結果を書き戻してやればよい。これはPostgreSQLオプティマイザがパスを検討する時点ではどういった数式を処理すべきか判断できないため、planner_hook を使って無理やりプラン木を書き換える事になる。(v9.7辺りでもう少しマシな手段が使えるようになればいいけどー)

で、正月休みにシコシコと作業をした結果が以下の通り。

# EXPLAIN VERBOSE
  SELECT a.id a_id, b.id b_id, sqrt((a.x1 - b.x1)^2 +
                                    (a.x2 - b.x2)^2 +
                                    (a.x3 - b.x3)^2 +
                                    (a.x4 - b.x4)^2) dist
    FROM test a, test b
   WHERE a.id > b.id;

                        QUERY PLAN

-------------------------------------------------------------------------
 Custom Scan (GpuJoin) on public.test a  (cost=777399.00..167444065.65 rows=3333333333 width=40)
   Output: a.id, b.id, (sqrt(((((((a.x1 - b.x1))::double precision ^ '2'::double precision) + (((a.x2 - b.x2))::double precision ^ '2'::double precision)) + (((a.x3 - b.x3))::double precision ^ '2'::double precision)) + (((a.x4 - b.x4))::double precision ^ '2'::double precision))))
   GPU Projection: a.id::integer, b.id::integer, sqrt(((((((a.x1 - b.x1))::double precision ^ '2'::double precision) + (((a.x2 - b.x2))::double precision ^ '2'::double precision)) + (((a.x3 - b.x3))::double precision ^ '2'::double precision)) + (((a.x4 - b.x4))::double precision ^ '2'::double precision)))::double precision
   Depth 1: GpuNestLoop, JoinQual: (a.id > b.id)
            Nrows (in/out: 3333333.20%), KDS-Heap (size: 8.11MB, nbatches: 1)
   Extra: bulk-exec-support
   Kernel Source: /opt/pgsql/base/pgsql_tmp/pgsql_tmp_strom_913.0.gpu
   ->  Seq Scan on public.test b  (cost=0.00..1637.00 rows=100000 width=20)
         Output: b.id, b.x1, b.x2, b.x3, b.x4
(9 rows)

少しEXPLAIN出力結果が変わっているが、GpuJoinの GPU Projection: が従前の Pseudo Scan: に相当する。
つまり、この実行計画は、GPU側で計算処理を行った上でそれをCPU側に返却。CPU側では計算結果を参照するだけ(= 計算は行わない)なので、Projection処理を省略することができるようになる。

実際のところ、もう少し適切な問題サイズの推定と分割無しには、上記のクエリのように GpuNestLoop が入力行数に対して3.3万倍の出力行を生成するようなワークロードが適切に働くとは思えないが、この辺は、後日GpuJoinの問題領域分割にDynamic Parallelismを適用する辺りでなんとかするとして、ひとまずインフラとしてはこれで良いように思える。