BigQuery (Standard SQL) でレイトレーシングをしてみました。
レイトレーシングとは
レイトレーシングとは、光の輸送(屈折や反射)を物理シミュレーションして現実的なCG画像を作りだす技術です。
最近では RTX や PS5 など、リアルタイムレイトレーシングが台頭してきています。
レイトレーシングではピクセルごとにレイを飛ばして計算するため計算量が膨大になりがちですが、
ピクセルごとに独立に計算することができるので、処理の高速化が期待できます。
それなら BigQuery が得意分野じゃないか？と思い今回の挑戦をしてみました。
BigQuery とは
超高速でSQLを分散実行し数秒でペタバイト級データに対しても結果が返ってくるデータ分析向けサーバーレス・データウェアハウスです。詳細は以下をごらんください。
BigQuery による画像生成
テキスト形式で画像を記述する Portable Any Map というフォーマットがあります。
BigQuery は CSV や JSON などのテキスト形式での出力をサポートしていますので、
今回は PNM として CSV を出力して画像を生成してみます。
画像出力の例としてマンデルブロ集合を描画してみます。
CREATE TEMPORARY FUNCTION F(s STRUCT<a FLOAT64, b FLOAT64, x FLOAT64, y FLOAT64, c INT64>)
AS (
IF(s.a*s.a + s.b*s.b <= 4,
STRUCT(s.a*s.a - s.b*s.b - s.x AS a,
2*s.a*s.b - s.y AS b,
s.x AS x,
s.y AS y,
s.c+1 AS c),
s)
);
-- Fを4回呼び出す
CREATE TEMPORARY FUNCTION F4(s STRUCT<a FLOAT64, b FLOAT64, x FLOAT64, y FLOAT64, c INT64>)
AS (F(F(F(F(s)))));
WITH data AS (
SELECT
x, y,
-- Fを16回呼び出す
F4(F4(F4(F4(STRUCT(
-1 + (j/400)*3 AS a,
-1.5 + (i/400)*3 AS b,
-1 + (j/400)*3 AS x,
-1.5 + (i/400)*3 AS y, 0 AS c))))) AS c
FROM
-- (1,1) から (400,400) の座標を生成する
UNNEST(GENERATE_ARRAY(1, 400)) AS i,
UNNEST(GENERATE_ARRAY(1, 400)) AS j
)
-- PGMとして出力
SELECT "P2 400 400 16 #", -1 AS x, -1 AS y UNION ALL
SELECT FORMAT("%d #", c.c), i, j
FROM
data
ORDER BY
i, j
各ピクセル i,j に対し、色(モノクロなので輝度) c を計算するコードになっています。
BigQuery SQL では UDF が使えますが、再帰はできないので不恰好な呼び出し方でループを実現しています。今回のレイトレーシングでもUDFを活用します。
この結果としてCSVとして結果をヘッダなしで出力すると以下のようなファイルになります。
P2 400 400 16 #,-1,-1
1 #,1,1
1 #,1,2
1 #,1,3
1 #,1,4
1 #,1,5
1 #,1,6
1 #,1,7
1 #,1,8
...
# 以降はコメントなのでこれは valid な pnm フォーマットです。
拡張子 pgm で保存すれば、Windows の場合は IfranView、macOS の場合は Preview.app で表示できます。
これで BigQuery で画像を出力できることが確認できました。
BigQuery によるレイトレーシング
というわけで、BigQueryでレイトレーシングをやってみましょう。
実際のSQLコードは以下のようになります。
-- Vec3のドット積
CREATE TEMPORARY FUNCTION DOT
(a STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>,
b STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>)
AS (
a.x*b.x + a.y*b.y + a.z*b.z
)
;
-- 線形結合 aP + bQ
CREATE TEMPORARY FUNCTION COMB
(a FLOAT64, p STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>,
b FLOAT64, q STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>)
AS (
STRUCT( a*p.x+b*q.x AS x, a*p.y+b*q.y AS y, a*p.z+b*q.z AS z)
)
;
-- 単位ベクトルを得る
CREATE TEMPORARY FUNCTION UNIT
(r STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>)
AS (
STRUCT( r.x/SQRT(DOT(r,r)) AS x, r.y/SQRT(DOT(r,r)) AS y, r.z/SQRT(DOT(r,r)) AS z)
)
;
-- [0,255]の範囲におさめる
CREATE TEMPORARY FUNCTION CLAMP255
(f FLOAT64)
AS (
CASE WHEN f < 0 THEN 0
WHEN f > 1 THEN 255
ELSE CAST(ROUND(IFNULL(f, 0) * 255.) AS INT64)
END
)
;
-- レイの反射を計算する
CREATE TEMPORARY FUNCTION STEP
(arg STRUCT<
-- レイ情報
rt ARRAY<STRUCT<
i INT64, -- ピクセル座標
j INT64,
r STRUCT<
o STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>, -- 原点
d STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>, -- 方向
af STRUCT<x FLOAT64, y FLOAT64, z FLOAT64> -- 色の累積値
>,
final STRUCT<x FLOAT64, y FLOAT64, z FLOAT64> -- レイが無限遠方に飛んだときの色
>>,
world ARRAY<STRUCT< -- オブジェクト設定
sc STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>, -- 球の中心
sr FLOAT64, -- 球の半径
att STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>, -- 色
mat STRING -- 材質 (diffuse, metal)
>>,
sample_num INT64 -- サンプル番号
>)
AS ((
-- (い) レイの衝突を計算
WITH hitinfo AS (
SELECT
i, j, r, s,
IF(t IS NOT NULL,
-- 衝突したので交点情報を返す
STRUCT(
COMB(1.0, r.o, t, r.d) AS p,
COMB(1.0/s.sr, COMB(1.0, r.o, t, r.d), -1.0/s.sr, s.sc) AS n
),
NULL
) AS hit
FROM (
SELECT
i, j, r,
(
SELECT AS STRUCT
-- 線分と球面の交点を計算する
s, t
FROM (SELECT
s, (-b - SQRT(b*b - a*c))/a AS t
FROM (SELECT
s,
DOT(r.d, r.d) AS a,
DOT(oc, r.d) AS b,
DOT(oc, oc) - s.sr*s.sr AS c
FROM
(SELECT
s, COMB(1.0, r.o, -1.0, s.sc) AS oc
FROM UNNEST(arg.world) AS s)
)
WHERE b*b - a*c > 0
)
WHERE t > 1e-7
ORDER BY t
LIMIT 1
).*
FROM
UNNEST(arg.rt)
WHERE
final IS NULL
)
)
SELECT
STRUCT(ARRAY(
SELECT AS STRUCT
i, j,
-- 衝突した場合材質に応じてレイを反射させる
IF(hit IS NULL, NULL,
CASE s.mat
WHEN 'diffuse' THEN
-- Lambert則(厳密ではない)
STRUCT(
hit.p AS o,
-- arg.sample_num/arg.sample_num はバグっぽい最適化がかかって RAND() が同じ値を返すのに対する workaround
COMB(1.0, hit.n, 1.0,
UNIT(STRUCT(2*RAND() - arg.sample_num/arg.sample_num AS x,
2*RAND() - arg.sample_num/arg.sample_num AS y,
2*RAND() - arg.sample_num/arg.sample_num AS z))) AS d,
STRUCT(r.af.x * s.att.x AS x, r.af.y * s.att.y AS y, r.af.z * s.att.z AS z) AS af
)
WHEN 'metal' THEN
-- 鏡面反射
STRUCT(
hit.p AS o,
COMB(1,UNIT(r.d),-2*DOT(UNIT(r.d), hit.n), hit.n) AS d,
STRUCT(r.af.x * s.att.x AS x, r.af.y * s.att.y AS y, r.af.z * s.att.z AS z) AS af
)
-- (屈折は時間がなかったため省略)
END) AS r, -- (う)
-- 何にもヒットしなかったので背景を表示
IF(hit IS NULL,
(SELECT STRUCT(
(1-t/2)*r.af.x AS x,
(1-t*0.3)*r.af.y AS y,
r.af.z AS z
)
FROM (SELECT (1+r.d.y/SQRT(DOT(r.d, r.d)))/2 AS t))
, NULL) AS final
FROM hitinfo
UNION ALL
-- 無限遠方にあるレイについては衝突判定をしない
SELECT AS STRUCT i, j, r, final
FROM UNNEST(arg.rt)
WHERE final IS NOT NULL
) AS rt,
arg.world AS world,
arg.sample_num AS sample_num
)
));
WITH pixels AS (
-- ピクセルを生成
SELECT
i, j
FROM
UNNEST(GENERATE_ARRAY(0, 399)) AS j,
UNNEST(GENERATE_ARRAY(0, 399)) AS i
),
world AS (
-- オブジェクトの定義
-- c: 中心, r: 半径, att: 色, mat: 材質
SELECT
STRUCT(0.0 AS x, 0.0 AS y, -1.0 AS z) AS sc, 0.5 AS sr,
STRUCT(0.8 AS x, 0.5 AS y, 0.2 AS z) AS att, "diffuse" AS mat
UNION ALL SELECT
STRUCT(1.0 AS x, 0.0 AS y, -1.0 AS z) AS sc, 0.5 AS sr,
STRUCT(0.8 AS x, 0.7 AS y, 0.9 AS z) AS att, "metal" AS mat
UNION ALL SELECT
STRUCT(-1.0 AS x, 0.0 AS y, -1.0 AS z) AS sc, 0.5 AS sr,
STRUCT(0.6 AS x, 0.5 AS y, 0.9 AS z) AS att, "metal" AS mat
UNION ALL SELECT
STRUCT(0.4 AS x, -0.4 AS y, -0.5 AS z) AS sc, 0.1 AS sr,
STRUCT(0 AS x, 1 AS y, 0.5 AS z) AS att, "diffuse" AS mat
UNION ALL SELECT
STRUCT(-0.4 AS x, -0.4 AS y, -0.5 AS z) AS sc, 0.1 AS sr,
STRUCT(1 AS x, 1 AS y, 1 AS z) AS att, "metal" AS mat
UNION ALL SELECT
STRUCT(0.0 AS x, -100.5 AS y, -1.0 AS z) AS sc, 100.0 AS sr,
STRUCT(0.5 AS x, 0.5 AS y, 0.2 AS z) AS att, "diffuse" AS mat
),
result AS (
WITH samples AS (
SELECT
-- 8回反射させる (力技)
-- BigQuery Scripting を使えばもっとスマートにできるが、スキャン料金がかかる
STEP(STEP(STEP(STEP(STEP(STEP(STEP(STEP(
STRUCT(
ARRAY(SELECT AS STRUCT
i, j,
-- (あ) カメラからのレイを生成
STRUCT(
STRUCT(0.0 AS x, 0.0 AS y, 0.0 AS z) AS o,
STRUCT(2.0*(i+2*RAND()-s/s)/400-1 AS x, 1.0-2.0*(j+2*RAND()-s/s)/400 AS y, -1.0 AS z) AS d,
STRUCT(1.0 AS x, 1.0 AS y, 1.0 AS z) AS af
) AS r,
CAST(NULL AS STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>) AS final
FROM pixels) AS rt,
ARRAY(SELECT AS STRUCT * FROM world) AS world,
s AS sample_num
)
)))))))).rt AS rt
FROM
-- 50サンプル取る
UNNEST(GENERATE_ARRAY(1,50)) AS s
)
SELECT i, j, final
FROM samples, UNNEST(rt)
)
-- PMX形式で出力 (ヘッダなしでCSV出力すればそのまま読み込めます)
SELECT
"P3 400 400 255 #", -1 AS i, -1 AS j,
UNION ALL
(
SELECT
FORMAT("%d %d %d #",
-- 各サンプルを加算してガンマ補正
CLAMP255(SQRT(SUM(final.x)/COUNT(final))),
CLAMP255(SQRT(SUM(final.y)/COUNT(final))),
CLAMP255(SQRT(SUM(final.z)/COUNT(final)))), i, j
FROM result
GROUP BY i, j -- (え)
)
ORDER BY j, i
詳細については割愛しますが、流れとしては普通のレイトレと同様な流れになっています
- I. 各ピクセルごとに以下を S 回 (サンプル数) 実行する
- 1. カメラから出るレイを生成 (アンチエイリアスのためにランダムに少しずらす) (あ)
- 2. 各物体との当たり判定を計算する
- 物体テーブル world と CROSS JOIN して t が最小となる物体を SQL で抽出しています
- 3. 当たり判定がある場合は、物体の材質に応じて反射の方向を計算する (い)
- 4. 当たり判定がない場合は、背景の色を参照する
- final というカラムに格納し、これ以上の当たり判定をしないようにします。。
- 5. 反射が起きた場合、2 に戻る
- 繰り返しは STEP という関数を何回も呼びだすことで実現しています
- コメントにもあるように BigQuery では SCRIPTING というプロシージャ機能にループがあるので、それを使うのもありです。
- II. ピクセルごとにサンプルを集計し、出力
- GROUP BY で集計しています (え)
出力
上のコードを実行して ppm で保存したものが冒頭の画像になります。
なおこのスクリプトの実行には 10 分 24 秒かかってしまいました1。手元の Python でのリファレンス実装では 10秒かからないのでちょっと遅いですね。
改善案としては、そもそも BigQuery では JavaScript で UDF が書けるため、それをレイの計算に利用するという方法があります。おそらくは普通の JavaScript で書くと同じかそれ以上のスピードは出ると予想されるため、数秒で終了するはずです。ただ今回はピュアな SQL で書いてみたいということで JavaScirpt は使いませんでした。2
結論
RTX を使おう。
参考文献
参考文献としては以下を参照しました。
-
https://raytracing.github.io/books/RayTracingInOneWeekend.html
- レンダリングのアルゴリズムはほとんどこれを利用しています
-
http://kagamin.net/hole/simple/index.htm
- 非再帰のコードの参考にしました