🔖

3D Gaussian Splatting ビューワを作る

に公開

この記事では、Rust と wgpu で実装した 3D Gaussian Splatting ビューワについて、描画パイプラインの仕組みを説明します。

GitHubリポジトリはこちらになります。
記事の内容は v-0.1.0 ブランチに対応しています。
https://github.com/abist-co-ltd/wgpu-gs-viewer/tree/v-0.1.0

今回作ったビューワは、3DGS PLY ファイルを読み込み、GPU 上の Tile-Based Rendering パイプラインによって高速に描画します。
また、ブラウザ/デスクトップ上で 3DGS シーンを表示できます。

デモ

ここでは、実装の細部を追うのではなく、3DGS ラスタライザを構成するアルゴリズムやデータの流れを理解することを目的とします。
そのため、wgpu 自体の詳しい説明は扱いません。
wgpu に慣れていない場合は、先に Learn wgpu などを確認しておくと理解しやすいと思います。

3D Gaussian Splatting (3DGS) とは

3DGSは、3Dシーン再構成に使われる技術の一つです。
複数の写真とカメラ位置をもとに、3D空間に半透明の3D Gaussianを大量に配置・最適化し、それらをスクリーンに投影して合成することで、任意の視点からリアルタイムに3Dシーンを描画する技術です。

https://www.youtube.com/watch?v=mD0oBE9LJTQ

3D Gaussian とは

3DGS における Gaussian は、3D空間上に置かれた半透明な楕円体状のプリミティブです。
単なる点とは異なり、位置・形状・向き・透明度・色を持ちます。

下の画像は、シーン内に配置された 各 Gaussian の形状を不透明な楕円体メッシュとして可視化したものです。

通常のGS描画

Gaussian形状の可視化

3D Gaussianは以下の特徴があります。

1. 面の表現しやすさ

Gaussianは広がりを持ち、スクリーン上では楕円として描画されます。
各Gaussianは周囲のピクセルに滑らかに寄与するため、点群よりも連続した面として描画することが可能です。

2. 透明度をもつ

Gaussianは半透明です。

3DGSを描画する際、各Gaussianをスクリーンに投影し、ピクセルごとに奥行き順にアルファ合成します。
これにより複数のGaussianの寄与を積み重ねて色を決定できます。

3. 視点依存の色を表現できる

Gaussianの色は、単純な固定RGBではなく、視線方向に応じて変化する色として持ちます。
これにより、単一の固定色ではなく、視点方向によって変化する色を表現できます。


1つのガウシアンでの視点依存の色変化

3DGS ラスタライザの概要

3DGS の描画では、3D空間上の Gaussian を現在のカメラ視点からスクリーンへ投影します。


Gaussianの2D投影のイメージ (ChatGPTで生成)

各ピクセルでは、そのピクセルを覆っている Gaussian の色と透明度を計算し、奥行き順にアルファ合成します。


Gaussianのalpha合成のイメージ (ChatGPTで生成)

これを全ピクセルについて処理すると、計算量はおおよそ 画面のピクセル数 \times Gaussian数 となります。

3DGS では、GPU上で並列に計算することに加えて、Tile-Based Rendringによる計算範囲の絞り込みによって、リアルタイム描画を実現します。

3DGS ラスタライザのステップは大きく分けると、次の4ステップです。

  1. PLYを読み込む

    学習済みの 3DGS データを PLY から読み込みます。
    各 Gaussian の位置、スケール、回転、透明度、SH 係数などを取得します。

  2. Preprocess pass

    各 Gaussian を並列に処理し、描画に必要な中間データを計算します。

    ここでは、Gaussian が現在のカメラから見てどこにあり、画面上でどのような楕円として見えるかを求めます。

    主に次の処理を行います。

    • camera space への変換
    • culling
    • 3D covariance 作成
    • 2D covariance への投影
    • screen center / radius / conic の計算
    • SH による色計算
  3. Tile-Based Rendering

    すべての pixel が画面全体の Gaussian を処理すると非常に重いため、
    画面を小さなタイルに分割し、各タイルで処理すべき Gaussian を絞り込みます。

    実装では、Preprocess pass で求めた画面上の中心位置や半径を使い、各 Gaussian がどのタイルに影響するかを計算します。


タイルとGaussianイメージ

  1. Gaussianを合成して描画する

    最後に、各 pixel の色を計算します。

    各 pixel では、属するタイルの Gaussian だけを調べます。
    そして、カメラに近いものから順番に色と透明度を合成します。

1. 3DGS PLY の読み込み

3DGS の学習済みデータは、PLY ファイルとして保存されていることが多いです。

PLY には、通常、各 Gaussian の位置、形状、向き、透明度、色に関するパラメータが格納されています。

PLYのヘッダーを見てみると、以下のようになっており、各 property は Gaussian を構成する情報です。

ply
format binary_little_endian 1.0
element vertex 139410
property float x
property float y
property float z
property float f_dc_0
property float f_dc_1
property float f_dc_2
property float f_rest_0
property float f_rest_1
property float f_rest_2
property float f_rest_3
...
property float f_rest_44
property float opacity
property float scale_0
property float scale_1
property float scale_2
property float rot_0
property float rot_1
property float rot_2
property float rot_3
end_header

property は以下のように分類できます。

property 用途
x, y, z 位置
scale スケール
rot 向き
opacity 透明度
f_dc, f_rest 色計算に使用される係数

PLY から読み込んだ各 property は、以下のようなGPU処理で扱いやすい構造体へ変換します。

#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
pub struct Gaussian3d {
    pub position: [f32; 3],
    pub opacity: f32,

    pub scale: [f32; 3],
    pub _pad0: u32,

    pub rotation: [f32; 4],
    pub sh: [f32; 48],
}

#[repr(C)]bytemuck::Podbytemuck::Zeroable を使うことで、Rust 側の構造体配列を GPU buffer に転送できるバイト列として扱えるようにします。[1]

実装

PLY 読み込み処理の実装は以下を参照してください。

src/ply_loader.rs

2. Preprocess Pass

Preprocess Pass では、PLY から読み込んだ Gaussian 配列を GPU の compute shader で並列に処理します。

この pass では、1 thread が 1 つの Gaussian を担当します。

各 thread は、自分が担当する Gaussian のGaussian の位置、スケール、回転、透明度、SH 係数を読み取り、現在のカメラ視点で描画に必要な情報へ変換します。

主に求める情報は次の通りです。

  • 画面上の中心座標
  • depth
  • 2D 共分散行列
  • 画面上での半径
  • 視線方向に応じた色
  • 透明度
  • Gaussian が重なるタイル範囲

つまり Preprocess Pass は、3D Gaussian を、現在のカメラから見た 描画用の 2D Gaussian 情報へ変換する pass です。

2.1 Gaussian形状の表現

各 Gaussian は中心位置 \mu を持つだけでなく、3D 空間上で広がりを持つ楕円体として扱われます。

この楕円体形状は、3次元共分散行列 \Sigma によって表現されます。

G_{\text{3D}}(x) = \exp\left( -\frac{1}{2} (x - \mu)^T \Sigma^{-1} (x - \mu) \right)
\mu:Gaussian の中心位置
\Sigma:Gaussian の広がりを表す3次元共分散行列

2.1.1 scale と rotation から 3D 共分散行列を作る

3DGS では、各 Gaussian の共分散行列 \Sigma を scale と rotation から構成します。

scale S は Gaussian の各ローカル軸方向の広がりを表し、rotation R はその楕円体の向きを表します。

まず、ローカル空間に標準的な Gaussian があると考えます。

この Gaussian の共分散行列を単位行列 I とします。

\Sigma_{\text{local}} = I

この標準 Gaussian に対して、scale と rotation による変換を適用します。

変換行列を

A = RS

とすると、A は線形変換のため、変換後の共分散行列は

\Sigma = A I A^T

になります。

ここに A = RS を代入すると、

\Sigma = A I A^T = (RS)(RS)^T = R S S^T R^T

が得られます。

つまり、RS は Gaussian の形状を変形する行列であり、RSS^TR^T はその変形後の広がりを表す共分散行列です。

2.1.2 3D 共分散行列を 2D に投影する

3D Gaussian はワールド空間上の楕円体ですが、描画時には画面上の 2D Gaussian として扱う必要があります。

そのため、3D 共分散行列 \Sigma を、画面上の 2D 共分散行列へ投影します。

3DGS では、投影後の共分散行列を次のように計算します。

\Sigma' = J W \Sigma W^T J^T

ここで、W は View 変換、J は透視投影変換のヤコビアンです。

まず、W によってワールド空間の共分散行列をカメラ空間へ変換します。

\Sigma_{\text{camera}} = W \Sigma W^T

次に、カメラ空間の Gaussian を画面上へ投影します。

ここで、透視投影は線形変換ではないことに注意が必要です。

カメラ空間の点を (x, y, z) とすると、画面上の座標はおおよそ次のように表されます。

u = f_x \frac{x}{z}
v = f_y \frac{y}{z}

このように、透視投影では z で割るため、投影は非線形になります。

そこで、Gaussian の中心付近で透視投影を局所的に線形近似します。

この局所的な線形近似がヤコビアン J です。

J の各成分は偏微分によって求められます。

J_{2 \times 3} = \frac{\partial (u, v)}{\partial (x, y, z)} = \begin{bmatrix} \dfrac{\partial u}{\partial x} & \dfrac{\partial u}{\partial y} & \dfrac{\partial u}{\partial z} \\ \dfrac{\partial v}{\partial x} & \dfrac{\partial v}{\partial y} & \dfrac{\partial v}{\partial z} \end{bmatrix} = \begin{bmatrix} \frac{f_x}{z} & 0 & -\frac{f_x x}{z^2} \\ 0 & \frac{f_y}{z} & -\frac{f_y y}{z^2} \end{bmatrix}

ここで、f_x, f_y は pixel 単位の焦点距離です。

実装では、後続の行列計算に合わせて次のような 3x3 行列として扱うことがあります。

J = \begin{bmatrix} \frac{f_x}{z} & 0 & -\frac{f_x x}{z^2} \\ 0 & \frac{f_y}{z} & -\frac{f_y y}{z^2} \\ 0 & 0 & 0 \end{bmatrix}

よって、

\Sigma_{\text{screen}} = J \Sigma_{\text{camera}} J^T = J W \Sigma W^T J^T

となります。

最後に、\Sigma_{\text{screen}} のうち、画面の x, y 方向に対応する 2x2 部分を取り出します。

\Sigma_{\text{2D}} = \begin{bmatrix} \Sigma_{\text{screen},00} & \Sigma_{\text{screen},01} \\ \Sigma_{\text{screen},10} & \Sigma_{\text{screen},11} \end{bmatrix}

この 2D 共分散行列が、画面上での Gaussian の楕円形状を表します。

投影後、画面上の点 p に対する Gaussian の値を評価するときは、投影後の中心 \mu_{\text{2D}} からのずれ

\delta_{\text{2D}} = p - \mu_{\text{2D}}

を使います。

投影された 2D Gaussian は次のように表せます。

G_{\text{2D}}(p) = \exp\left( -\frac{1}{2} (p - \mu_{\text{2D}})^T \Sigma_{\text{2D}}^{-1} (p - \mu_{\text{2D}}) \right)

つまり、画面上の Gaussian は、2D 共分散行列 \Sigma_{\text{2D}} によって定まる楕円形状として評価されます。

2.1.3 実装

実装は以下を参照してください。

2.2 色表現

各 Gaussian は見る方向によって異なる色を返すような仕組みを持っています。

これを実現する SH (Spherical Harmonics) 係数について説明します。

2.2.1 SHとは

例えば、球の表面について

  • 上方向は明るい
  • 横方向は暗い
  • 前方向は赤っぽい
  • 後ろ方向は青っぽい

のように、方向ごとの変化を表現したいときがあります。

SH は、球面上で方向によって変わる値を、複数の基本パターンの組み合わせで表す方法です。

2.2.2 SH の考え方

SH では、方向を表すベクトル

d = (x,y,z)

を入力すると値を返す関数 Y_{i}(d) を使います。

そして、方向によって変わる値 f(d) を、複数の Y_{i}(d) の足し合わせで近似します。

f(d) \approx c_0Y_0(d) + c_1Y_1(d) + c_2Y_2(d) + c_3Y_3(d) + c_4Y_4(d) + ...
c_{i}: 重み係数(SH 係数)

ここで

Y{_0}:どの方向でも同じ値になる
Y{_1}:上下に沿って変化
Y{_2}:前後に沿って変化
Y{_3}:左右に沿って変化
Y{_4}以降:斜め方向や軸方向の差を表す、より複雑なパターン

つまり SH では、複雑な方向変化を少数の基本パターンと係数だけで近似できます。

2.2.3 3DGS での SH

先ほど確認した PLY の property のうち、f_dc, f_rest が SH 係数 になります。

f_dc_0, f_dc_1, f_dc_2 は Y_{0} で使われるベースカラーであり、

f_rest_0 から f_rest_44 以降は、Y_{1}, Y_{2}, ..., Y_{15} で使われる SH 係数です。

R, G, B は以下のように計算されます。

\begin{aligned} R(d) &= \mathrm{f\_dc\_0} Y_0(d) + \mathrm{f\_rest\_0} Y_1(d) + \mathrm{f\_rest\_3} Y_2(d) + \mathrm{f\_rest\_6} Y_3(d) + \cdots + \mathrm{f\_rest\_42} Y_{15}(d) \\ G(d) &= \mathrm{f\_dc\_1} Y_0(d) + \mathrm{f\_rest\_1} Y_1(d) + \mathrm{f\_rest\_4} Y_2(d) + \mathrm{f\_rest\_7} Y_3(d) + \cdots + \mathrm{f\_rest\_43} Y_{15}(d) \\ B(d) &= \mathrm{f\_dc\_2} Y_0(d) + \mathrm{f\_rest\_2} Y_1(d) + \mathrm{f\_rest\_5} Y_2(d) + \mathrm{f\_rest\_8} Y_3(d) + \cdots + \mathrm{f\_rest\_44} Y_{15}(d) \end{aligned}

ここで用いる Y_i は、以下のようになります。[2]

\begin{aligned} Y_{0} &= \frac{1}{2\sqrt{\pi}} \\ Y_{1} &= -\sqrt{\frac{3}{4\pi}}\, y \\ Y_{2} &= \sqrt{\frac{3}{4\pi}}\, z \\ Y_{3} &= -\sqrt{\frac{3}{4\pi}}\, x \\ Y_{4} &= \frac{1}{2}\sqrt{\frac{15}{\pi}}\, xy \\ Y_{5} &= -\frac{1}{2}\sqrt{\frac{15}{\pi}}\, yz \\ Y_{6} &= \frac{1}{4}\sqrt{\frac{5}{\pi}}\, (2z^2 - x^2 - y^2) \\ Y_{7} &= -\frac{1}{2}\sqrt{\frac{15}{\pi}}\, zx \\ Y_{8} &= \frac{1}{4}\sqrt{\frac{15}{\pi}}\, (x^2 - y^2) \\ Y_{9} &= -\frac{1}{4}\sqrt{\frac{35}{2\pi}}\, (3x^2 - y^2)y \\ Y_{10} &= \frac{1}{2}\sqrt{\frac{105}{\pi}}\, xyz \\ Y_{11} &= -\frac{1}{4}\sqrt{\frac{21}{2\pi}}\, (4z^2 - x^2 - y^2)y \\ Y_{12} &= \frac{1}{4}\sqrt{\frac{7}{\pi}}\, z(2z^2 - 3x^2 - 3y^2) \\ Y_{13} &= -\frac{1}{4}\sqrt{\frac{21}{2\pi}}\, x(4z^2 - x^2 - y^2) \\ Y_{14} &= \frac{1}{4}\sqrt{\frac{105}{\pi}}\, (x^2 - y^2)z \\ Y_{15} &= -\frac{1}{4}\sqrt{\frac{35}{2\pi}}\, x(x^2 - 3y^2) \end{aligned}

2.2.4 実装

実装は以下を参照してください。

2.3 Preprocess Pass の出力

ここまでで、各 Gaussian について

  • 画面上での楕円形状
  • 視点方向に応じた色

を計算できました。

これらの値を後続の Tile-Based Rendering で使いやすい形にまとめて buffer へ書き込みます。

出力は以下の構造体配列です。

struct PreprocessOutput {
    conic_opacity: vec4<f32>, // 2D共分散行列の逆行列と透明度
    color_radius: vec4<f32>,  // 色, 画面上の Gaussian の半径
    tile_rect: vec4<u32>,     // Gaussian が重なる tile 範囲
    uv_depth: vec4<f32>,      // 画面上の中心座標と depth
};

@group(0) @binding(2)
var<storage, read_write> outputs: array<PreprocessOutput>;

※ 単純のため、1つのフィールドに複数の情報をまとめて格納しています

2.3.1 2D共分散行列の逆行列

後続の処理では、2D共分散行列の逆行列を使ってGaussian の寄与を計算します。

ここでは対称行列であることを活かし、vec3<f32> で出力します。

let det_cov2d = cov2d[0][0] * cov2d[1][1] - cov2d[0][1] * cov2d[0][1];
let inv_det = 1.0 / det_cov2d;
let conic = vec3<f32>(
    cov2d[1][1] * inv_det,
    -cov2d[0][1] * inv_det,
    cov2d[0][0] * inv_det,
);

2.3.2 透明度

PLY から読み込んだ opacity はそのまま使わず、シグモイド関数で 0 から 1 の範囲に変換します。

let opacity = 1.0 / (1.0 + exp(-gaussians[idx].opacity));

2.3.3 画面上の Gaussian の半径

2D共分散行列の最大固有値を求めることで、画面上の楕円が最も大きく広がる方向の分散を計算できます。

let tr = 0.5 * (cov2d[0][0] + cov2d[1][1]);
let disc = max(0.1, tr * tr - det_cov2d);
let s = sqrt(disc);
let lmax = max(tr + s, tr - s);
let radius = min(ceil(SIGMA_SCALE * sqrt(lmax)), MAX_RADIUS);

これは、Gaussian がどのタイルに影響するかを求めるために使います。

2.3.4 Gaussian が重なるタイル範囲

ここでは、画面上の中心座標 uv と Gaussian 半径を使って、まず Gaussian が重なる可能性のあるタイル範囲を矩形で求めます。

let min_tx = i32(floor((uv.x - radius) / f32(TILE_W)));
let min_ty = i32(floor((uv.y - radius) / f32(TILE_H)));
let max_tx = i32(ceil((uv.x + radius) / f32(TILE_W)));
let max_ty = i32(ceil((uv.y + radius) / f32(TILE_H)));

このとき、以下の矩形範囲が重なる可能性のある範囲となります。

x: [uv.x - radius, ..., uv.x + radius]
y: [uv.y - radius, ..., uv.y + radius]

この矩形範囲は、後続の Duplicate Pass で どのタイルが Gaussian と対応するか を決めるために使います。

ただし、Gaussian は本来は矩形ではなく楕円です。
そのため、矩形範囲をそのまま使うと、実際には楕円と重ならないタイルまで含まれてしまいます。

そこで、矩形範囲の中の各タイルについて、Gaussian の楕円と重なる可能性があるかを追加で判定 [3] します。

var tiles_touched_count = 0u;
for (var ty = tile_rect.y; ty < tile_rect.w; ty = ty + 1u) {
    for (var tx = tile_rect.x; tx < tile_rect.z; tx = tx + 1u) {
        // 楕円と重なる可能性があるタイルのみ残す
        if tile_may_intersect_conic(uv, conic, tx, ty) {
            tiles_touched_count = tiles_touched_count + 1u;
        }
    }
}

これにより、矩形範囲内の全タイルを登録する場合に比べて、不要なタイルへの登録を減らし、後続の pass を軽くできます。

2.3.5 出力バッファへの書き込み

最後に、ここまで計算した値を PreprocessOutput に書き込みます。

ただし、すべての Gaussian が出力対象になるわけではありません。
画面上に映らない Gaussian は、ここで除外されます。

そのため、出力バッファに書き込む index は、入力 Gaussian index ではなく、visible_count を atomic に加算して取得します。

let visible_idx = atomicAdd(&visible_count, 1u);

そして、後続の pass で使う値を outputs に書き込みます。

outputs[visible_idx].conic_opacity = vec4<f32>(
    conic.x,
    conic.y,
    conic.z,
    opacity,
);

outputs[visible_idx].color_radius = vec4<f32>(
    rgb.x,
    rgb.y,
    rgb.z,
    radius,
);

outputs[visible_idx].tile_rect = tile_rect;

outputs[visible_idx].uv_depth = vec4<f32>(
    uv.x,
    uv.y,
    view_pos.z,
    0.0
);

tiles_touched[visible_idx] = tiles_touched_count;


画面に映るGaussianだけ出力

これで、画面に映る Gaussian だけを格納した outputs が得られました。

3. Dispatch Indirect

Preprocess Pass では、各 Gaussian を処理し、画面上に映る Gaussian だけを outputs に書き込みました。

次の pass では、この Gaussian だけを処理したいです。

通常、compute shader を実行するときは、CPU 側 (Rust) から workgroup 数を指定します。

compute_pass.dispatch_workgroups(x, y, z);

今回の visible_count は GPU 上で計算されます。
そのため、CPU 側では処理すべき Gaussian の数を知りません。

GPU から visible_count を CPU に readback してから workgroup 数を決めることもできます。
しかし、毎フレーム GPU の結果を CPU に戻すと同期負荷が発生してしまいます。

そこで使うのが Dispatch Indirect です。
Dispatch Indirect は、compute shader の workgroup 数を CPU が直接指定するのではなく、GPU buffer に書かれた値を使って dispatch する仕組みです。

シェーダーは以下のように記述できます。


// compute_pass.dispatch_workgroups(x, y, z)の x, y, z に対応する
struct DispatchIndirectArgs {
    x: u32,
    y: u32,
    z: u32,
    _pad: u32,
};

@group(0) @binding(1)
var<storage, read_write> dispatch_args: DispatchIndirectArgs;

fn div_ceil(x: u32, y: u32) -> u32 {
    return (x + y - 1u) / y;
}

// 後続のpassのworkgroupサイズ
const WORKGROUP_SIZE: u32 = 256u;

@compute @workgroup_size(1, 1, 1)
fn main() {
    let count = ...;
    let num_workgroups = div_ceil(count, WORKGROUP_SIZE);

    dispatch_args.x = num_workgroups;
    dispatch_args.y = 1u;
    dispatch_args.z = 1u;
    dispatch_args._pad = 0u;
}

この pass 自体は、workgroup数が書き込まれたバッファを1つ作るだけなので、CPU 側からは 1 workgroup で実行します。

build_dispatch_args_pass.dispatch_workgroups(1, 1, 1);

その後、作成した dispatch_args_buffer を使って、後続の compute pass を実行します。

compute_pass.dispatch_workgroups_indirect(&dispatch_args_buffer, 0);

4. Tile-Based Rendering

ある pixel を描画するとき、処理すべきは、画面に映るすべての Gaussian ではありません。
その pixel が属するタイルに重なっている Gaussian だけを処理すれば十分です。

これは、以下のようなステップになります。

  1. 属するタイルを求める
  2. そのタイルに属する Gaussian 一覧を取得する
  3. それらの Gaussian だけを評価する
  4. 奥行き順に alpha 合成する

これを可能にするのが、タイルごとに整理された Gaussian 配列です。

Preprocess Pass の段階では、Gaussian 配列の並びはタイルとは無関係です。

Gaussian 0
Gaussian 1
Gaussian 2
Gaussian 3
...

この状態では、あるタイルに属する Gaussian だけをすぐに取り出すことができません。
タイルに重なっている Gaussian を調べるには、すべての Gaussian を確認する必要があります。
これでは、Tile Based Rendering の利点がありません。

そこで、後続の pass では Gaussian 単位の情報を、タイル単位で扱いやすい形に変換します。

目指す形は、タイルごとに Gaussian をすぐ取り出せるデータです。
そのために、各タイルと Gaussian の対応を 1次元配列として並べます。

tile_gaussian_list = 
  [
    tile 0 -> Gaussian A,
    tile 0 -> Gaussian B,
    tile 0 -> Gaussian C,
    tile 1 -> Gaussian D,
    tile 1 -> Gaussian E,
    tile 2 -> Gaussian F,
    ...
  ]

また、Gaussian Splatting では alpha 合成を行うため、同じタイル内の Gaussian は奥行き順 (カメラから近い順) に処理する必要があります。

そのため、対応リストは単にタイルごとにまとまっているだけでなく、同じタイル内では奥行き順に並べておく必要があります。

tile_gaussian_list = 
  [
    tile 0 -> Gaussian A  // 手前
    tile 0 -> Gaussian B
    tile 0 -> Gaussian C  // 奥
    tile 1 -> Gaussian D  // 手前
    tile 1 -> Gaussian E  // 奥
    tile 2 -> Gaussian F
    ...
  ]

これにより、タイルに属する Gaussian を先頭から順に処理するだけで、奥行き順の alpha 合成を行えます。

これに加えて、各タイルがリストのどの範囲に対応するかを表す tile_ranges を作成します。

例えば、

tile_ranges[5] = { start: 120, end: 145 }

であれば、タイル 5 に属する Gaussian は、

tile_gaussian_list[120..145]

に並んでいる、という意味になります。

これにより、pixel が属するタイルを求め、tile_ranges から処理すべき Gaussian を取り出すことができます。
また、取り出した順番は depth順 に一致するため、奥行き順の alpha 合成を簡単に実現できます。

range = tile_ranges[tile_id]

for i in range.start .. range.end {
    gaussian = tile_gaussian_list[i]
    
    // Gaussian を pixel 上で評価する
    
    // alpha 合成する
}

このように、Tile-Based Rendering では、

  • タイルと Gaussian の対応配列
  • 各タイルが対応配列のどの範囲を読むかを表す tile_ranges

を作ることで、各 pixel が自分のタイルに属する Gaussian だけを効率よく処理できるようにします。

ただし、各 Gaussian が重なるタイル数は異なるため、GPU 上でそのまま可変長のリストを作るのは難しいです。
そこで Prefix Scan、Duplicate、Radix Sort、Tile Range、Tile Render という複数の pass に分け、GPU が並列に処理しやすい形でデータを構築します。

それぞれの pass の役割は次の通りです。

Prefix Scan Pass:
  各 Gaussian が、タイルとの対応リストのどの場所に書き込めばよいかを求める

Duplicate Pass:
  Gaussian とタイルの対応リストを作る

Radix Sort Pass:
  対応リストをタイルごとにまとめ、
  同じタイル内ではカメラに近い順に並べる

Tile Range Pass:
  tile_ranges を作成する

Tile Render:
  各ピクセルのレンダリング

4.1 Prefix Scan

タイルと Gaussian の対応配列を作成する際、
問題は、1つの Gaussian が重なるタイル数は一定ではないことです。

例えば、ある Gaussian は 3 個のタイルに重なり、別の Gaussian は 10 個のタイルに重なるかもしれません。

tile 0 -> Gaussian 0
tile 1 -> Gaussian 0
tile 2 -> Gaussian 0
tile 5 -> Gaussian 1
tile 8 -> Gaussian 2
...

後続の Duplicate Pass では、各 Gaussian を重なるタイルの数だけ、タイルと Gaussian の対応配列へ書き込みます。

このとき問題になるのは、各 Gaussian が対応配列のどこに書き込めばよいかです。

そこで必要になるのが Prefix Scan Pass です。

Prefix Scan Pass では、各 Gaussian が重なるタイル数をもとに、対応配列内でその Gaussian に割り当てる範囲を求めます。

4.1.1 Prefix Scan とは

Prefix Scan とは、配列の先頭から順に累積和を作る処理です。

Prefix Scan には、大きく分けて inclusive scan と exclusive scan があります。

例えば、入力が次の配列だとします。

input = [3, 1, 5, 2]

inclusive scan は、自分自身の値を含めた累積和を出力します。

inclusive = [3, 4, 9, 11]

一方、exclusive scan は、自分自身の値を含めず、自分より前の要素の合計を出力します。

exclusive = [0, 3, 4, 9]

Tile-Based Renderingでは、 exclusive scan を応用します。

先ほどの入力配列[3, 1, 5, 2]は、それぞれの Gaussian が重なるタイル数だと考えることができます。

Gaussian 0 -> 3 個のタイルに重なる
Gaussian 1 -> 1 個のタイルに重なる
Gaussian 2 -> 5 個のタイルに重なる
Gaussian 3 -> 2 個のタイルに重なる

この配列に対して exclusive scan を行うと、次のようになります。

input:   [3, 1, 5, 2]
output:  [0, 3, 4, 9]

これにより、各 Gaussian が、タイルと Gaussian の対応配列のどの範囲を使えばいいかが分かります。

Gaussian 0 -> tile_gaussian_list[0..3)
Gaussian 1 -> tile_gaussian_list[3..4)
Gaussian 2 -> tile_gaussian_list[4..9)
Gaussian 3 -> tile_gaussian_list[9..11)

このように、Prefix Scan は 可変長の出力を1次元配列に並べるための開始位置を求める処理として使えます。

4.1.2 Up-Sweep

実装では、Blelloch Scan と呼ばれる方法で exclusive scan を行っています。

Blelloch Scan は大きく 2 段階に分かれます。

  1. Up-Sweep
  2. Down-Sweep

まず Up-Sweep では、配列を Binary Tree として考え、Leaf から Root に向かって部分和を作ります。

例えば、以下のような入力配列を考えます。

input:
[3, 1, 5, 0, 2, 4, 1, 6]

この配列を、Leaf が 8 個ある Binary Tree として見ます。

最初は、隣り合う 2 要素の合計を作ります。

1次元配列では以下のようになります。

[3, 4, 5, 5, 2, 6, 1, 7]
    ^     ^     ^     ^

次に、4 要素単位の合計を作ります。

木構造で見ると、さらに 1 段上のノードを作る段階です。

1次元配列では以下のようになります。

[3, 4, 5, 9, 2, 6, 1, 13]
          ^            ^

最後に、8 要素全体の合計を作ります。

木構造で見ると、Root を作る段階です。

1次元配列では以下のようになります

[3, 4, 5, 9, 2, 6, 1, 22]
                       ^

つまり、Up-Sweepでは、各ノードが担当する配列の範囲の合計を その範囲の右端の要素に畳み込んでいると考えると分かりやすいです。

4.1.3 Down-Sweep

Prefix Scan の目標は、各要素について 自分より前の要素の合計 を求めることです。

Down-Sweepでは、配列を Binary Tree として考え、Up-Sweep で作った部分和を使いながら、Root から Leaf に向かって値を構築していきます。

このとき、親ノード から 左右の子ノード を構築する規則は次のようになります。

left  = parent
right = parent + left が担当する範囲の合計

左の子ノード が担当する範囲の前にある要素は、親ノード が担当する範囲の前にある要素と同じです。
そのため、left には parent の値をそのまま渡します。

一方、右の子ノード が担当する範囲の前には、親ノードの前にある要素に加えて、左の子ノードが担当する範囲の要素があります。
そのため、right には parent + left の範囲の合計 を渡します。

この規則を Root から Leaf まで繰り返すことで、最終的に各 Leaf には 自分より前の要素の合計 = 累積和 が入ります。

前節では、Up-Sweep によって以下の配列を得ました。

[3, 4, 5, 9, 2, 6, 1, 22]

ここで、Root は 配列全体を担当します。
配列全体より前には要素が存在しないため、Root を 0 にします。

1次元配列では以下のようになります。

[3, 4, 5, 9, 2, 6, 1, 0]
                      ^

次に、全体を左右に分けます。

root: [0..8) = 0
left:   [0..4)
right:  [4..8)

Up-Sweep によって、左の子ノード [0..4) の合計は 9 だと分かっています。

これをDown-Sweepの規則に当てはめると、

left  = parent
      = root
      = 0

right = parent + left が担当する範囲の合計
      = root + 9
      = 0 + 9
      = 9

木構造で見ると以下のようになります

1次元配列では以下のようになります。

[3, 4, 5, 0, 2, 6, 1, 9]
          ^           ^

次は、2 要素ずつに分けて、同様に考えます。

1次元配列では以下のようになります。

[3, 0, 5, 4, 2, 9, 1, 15]
    ^     ^     ^      ^

最後に、1要素ずつに分けて同様の操作を行います。

1次元配列では以下のようになります。

[0, 3, 4, 9, 9, 11, 15, 16]

これが Prefix Scan (exclusive scan) の結果になります。

4.1.4 GPU 上での Prefix Scan と Shared Memory

前節までで見たように、Up-Sweep / Down-Sweep は各段階ごとに同じ深さのノードを並列に計算できます。

ただし、Up-Sweep / Down-Sweep は、各段階が完全に独立しているわけではありません。
例えば Up-Sweep では、2 要素の和を作ったあと、その結果を使って 4 要素の和を作る必要があります。
つまり、次の段階では、前の段階で計算した値を正しく使う必要があります。

そのため、workgroup 内の thread が共有できる一時領域として Shared Memory を使います。

WGSL では、以下のようにして Shared Memory を構築できます。

var<workgroup> sdata: array<u32, WORKGROUP_SIZE>;

各 thread は、自分が担当する値を Shared Memory に書き込みます。
その後、Up-Sweep / Down-Sweep の各段階で Shared Memory を更新していきます。

このとき重要になるのが workgroupBarrier() です。

ある段階で Shared Memory に書き込んだ値を、次の段階で別の thread が読みます。
そのため、段階の区切りごとに workgroupBarrier() を入れ、前の段階の書き込みが次の段階の読み取りから見えるようにします。

2 要素の和を Shared Memory に書き込む

workgroupBarrier()

Shared Memory から 2 要素の和を読み、4 要素の和を作る

workgroupBarrier()

Shared Memory から 4 要素の和を読み、8 要素の和を作る

workgroupBarrier()がないと、次の段階の thread が、前の段階でまだ反映されていない古い値を読んでしまう可能性があります。
その場合、部分和が壊れ、Prefix Scan の結果も正しくなりません。


もう一つ重要なのは、Shared Memory は workgroup ごとに独立しているという点です。
つまり、ある workgroup の Shared Memory は、その workgroup 内の thread だけが読み書きできます。
別の workgroup の Shared Memory を直接読むことはできません。

そのため、この段階で完成するのは、配列全体の Prefix Scan ではなく、workgroup が担当する範囲ごとの Prefix Scan です。

4.1.5 workgroup をまたぐ Prefix Scan

workgroup 内の scan だけでは、複数の workgroup にまたがる配列全体の Prefix Scan は完成しません。

例えば、入力配列が複数の workgroup に分かれているとします。

input:
[ block 0 ][ block 1 ][ block 2 ][ block 3 ]...

各 block は、1つの workgroup が担当する範囲です。

workgroup 0 -> block 0 を scan
workgroup 1 -> block 1 を scan
workgroup 2 -> block 2 を scan
workgroup 3 -> block 3 を scan
...

この時点では、各 block 内では正しい Prefix Scan が作られています。

[ scan(block 0) ][ scan(block 1) ][ scan(block 2) ][ scan(block 3) ]...

しかし、block 1 以降には、それより前の block の合計がまだ反映されていません。

つまり、配列全体として正しい Prefix Scan にするには、各 block に対して 自分より前の block の合計 を足す必要があります。

そのため、各 workgroup は、自分が担当した block 全体の合計を書き出しておくことが必要です。

block_sums:
[ sum(block 0), sum(block 1), sum(block 2), sum(block 3),... ]

この block_sums に対してさらに Prefix Scan を行うと、各 block に足すべき値が得られます。

block_sums:
[100, 80, 120, 50]

prefix scan:
[0, 100, 180, 300]

この結果は、各 block に足すべき補正値です。

block 0 -> +0
block 1 -> +100
block 2 -> +180
block 3 -> +300

この補正値を各 block 内の scan 結果に加えることで、workgroup をまたいだ Prefix Scan が完成します。

ただし、block_sums はサイズによっては 1 つの workgroup で処理できるとは限りません。

元の配列が大きい場合、block_sums も大きな配列となり、workgroup の数も多くなります。

この場合、1回の Prefix Scan で全ての補正値を計算できないことがあります。

そのため、今回の実装では Prefix Scan を階層的に実行します。

Level 0:
  入力配列 を Prefix Scan して offsets に書き出す
  各 workgroup の合計を block_sums0 に書き出す

Level 1:
  block_sums0 を Prefix Scan して block_offsets0 に書き出す
  各 workgroup の合計を block_sums1 に書き出す

Level 2:
  block_sums1 を Prefix Scan して block_offsets1 に書き出す

その後、上位 level で得た補正値を下位 level に足していきます。

block_offsets1
    |
    | Add Block Offsets
    v
block_offsets0 を補正

block_offsets0
    |
    | Add Block Offsets
    v
offsets を補正

これにより、最終的に offsets は、配列全体に対する正しい exclusive scan になります。


実装では、Prefix Scan を 3 段階で行っています。

各 workgroup が WORKGROUP_SIZE 個の要素を処理できるので、3 段階の階層で扱える入力サイズの上限は、WORKGROUP_SIZE^3 になります。

Level 0: N 個の要素を処理
Level 1: ceil(N / WORKGROUP_SIZE) 個の block_sums0 を処理
Level 2: ceil(N / WORKGROUP_SIZE^2) 個の block_sums1 を処理

Level 2 が 1 workgroup 内で処理できるためには、

ceil(N / WORKGROUP_SIZE^2) <= WORKGROUP_SIZE

である必要があります。

したがって

N <= WORKGROUP_SIZE^3

これが 3 段階構成で扱える上限になります。

実装では、WORKGROUP_SIZE = 256 としています。

256^3 = 16,777,216

なので、約 1,677 万要素まで扱えます。

Prefix Scan Pass の入力配列は、可視 Gaussian が 重なるタイル数の配列 です。
つまり、要素数は画面に映る Gaussian 数に対応します。
理論的には、約 1,677 万 Gaussian まで扱えることになります。
より大きな入力を扱いたい場合は、Prefix Scan の level 数を増やす必要があります。

4.1.6 Prefix Scan 用の Dispatch Indirect

実装では、Prefix Scan は次の 5 つの compute pass に分かれています。

  1. Scan Level 0
  2. Scan Level 1
  3. Scan Level 2
  4. Add Level 1
  5. Add Level 0

Scan Level 0 では、入力配列を Prefix Scan します。

Scan Level 1 では、Level 0 で作成した block_sums0 を Prefix Scan します。

Scan Level 2 では、Level 1 で作成した block_sums1 を Prefix Scan します。

その後、上位 Level で得られた補正値を下位 Level に足していきます。

Scan Level 0:
  tiles_touched -> offsets, block_sums0

Scan Level 1:
  block_sums0 -> block_offsets0, block_sums1

Scan Level 2:
  block_sums1 -> block_offsets1

Add Level 1:
  block_offsets1 を block_offsets0 に加算する

Add Level 0:
  block_offsets0 を offsets に加算する

Prefix Scan の入力サイズは、Preprocess Pass の結果である 画面に映るGaussian数 = visible_count によって決まります。

visible_count は GPU 上で計算された値なので、Prefix Scan の前に、dispatch 数を GPU 上で作成します。

まず、各 level の要素数を求めます。

n0 = visible_count
n1 = ceil(n0 / WORKGROUP_SIZE)
n2 = ceil(n1 / WORKGROUP_SIZE)

それぞれの意味は次の通りです。

n0:
  Scan Level 0 が処理する要素数
  tiles_touched の要素数

n1:
  Scan Level 1 が処理する要素数
  Level 0 の block 数
  block_sums0 の要素数

n2:
  Scan Level 2 が処理する要素数
  Level 1 の block 数
  block_sums1 の要素数

次に、それぞれの pass に必要な workgroup 数を計算します。

Scan Level 0:
  ceil(n0 / WORKGROUP_SIZE)

Scan Level 1:
  ceil(n1 / WORKGROUP_SIZE)

Scan Level 2:
  ceil(n2 / WORKGROUP_SIZE)

Add Level 1:
  ceil(n1 / WORKGROUP_SIZE)

Add Level 0:
  ceil(n0 / WORKGROUP_SIZE)

この 5 つの dispatch 数を、Dispatch Indirect 用の buffer に書き込んでおきます。

実行時には、各 pass がこの buffer を参照して dispatch されます。

4.1.7 Gaussian とタイルとの対応配列のサイズを求めておく

Prefix Scan によって、各 Gaussian が、タイルと Gaussian の対応配列のどこから書き込めばよいかを表す offsets が求まりました。

次に必要なのは、対応配列全体のサイズです。

後続の Duplicate Pass では、各 Gaussian を重なるタイルの数だけ対応配列へ書き込みます。
そのため、対応配列に必要な要素数は、各 Gaussian が重なるタイル数 の合計になります。

実装では、階層的な Prefix Scan の過程で作られた block_sums1 を合計することで求めています。

これは、後続の Duplicate Pass や Radix Sort Pass で、何個の対応情報を処理するかを決めるために使います。

4.1.8 実装

Prefix Scan Pass の実装は以下を参照してください。

4.2 Duplicate Pass

Prefix Scan Pass によって、各 Gaussian が対応配列のどこから書き込めばよいかを表す offsets が求まりました。

また、Preprocess pass では Gaussian がタイルと重なる矩形範囲 tile_rect を求めました。

Duplicate Pass では、これを使って、タイルと Gaussian の対応配列を構築します。

例えば、ある Gaussian A が次のタイルに重なるとします。

tile 5
tile 6
tile 13

この場合、次のような対応情報を作ります。

tile 5  -> Gaussian A
tile 6  -> Gaussian A
tile 13 -> Gaussian A

1 つの Gaussian から、重なるタイルの数だけ対応情報が作られます。

4.2.1 offsets を使って書き込み位置を決める

Duplicate Pass では、各 Gaussian が並列に処理されます。
そのため、タイルと Gaussian の対応配列へ書き込む位置が衝突しないようにする必要があります。

これは、Prefix Scan Pass で求めた offsets を使えば問題ありません。

例えば、各 Gaussian と対応するタイル数を次のように考えると、

tiles_touched = [3, 1, 5, 2]
offsets       = [0, 3, 4, 9]

であれば、各 Gaussian は次の範囲に書き込みます。

Gaussian 0 -> list[0..3)
Gaussian 1 -> list[3..4)
Gaussian 2 -> list[4..9)
Gaussian 3 -> list[9..11)

各 Gaussian は自分に割り当てられた範囲だけに書き込むため、並列に処理しても書き込み位置が衝突しません。

4.2.2 重なるタイルだけを書き込む

offsets によって書き込む範囲は特定できますが、書き込むべきタイルについては注意が必要です。

Duplicate Pass が受け取るのは、Preprocess Pass で求めた tile_rect です。

2.3.4 の説明とも重なりますが、tile_rect は Gaussian とタイルが重なる矩形範囲であり、その中のすべてのタイルが実際に Gaussian と重なるとは限りません。

Preprocess Pass では、tile_rect 内の各タイルに対して Gaussian と重なるかどうかの判定に通ったタイル数を数えています。

Prefix Scan Pass で得られた offsets は、このタイル数をもとに作られています。

そのため Duplicate Pass でも、同じ tile_rect を走査し、Preprocess Pass と同じように Gaussian と重なるかどうかの判定を再度行います。

判定に通ったタイルだけを書き込むことで、Preprocess Pass で数えたタイル数と Duplicate Pass で実際に書き込む数を一致させます。

4.2.3 pair_key / pair_value を作る

タイルと Gaussian の対応情報は、実装上は pair_keypair_value に分けています。
これらは同じサイズの 1 次元配列であり、同じ index の要素同士が対応します。

pair_key = sort 用の key 
pair_value = 対応する Gaussian 情報を参照するための index

pair_key は、後続の Radix Sort Pass で並び替えるための key です。

実装では、pair_key を次のような 64bit key として扱います。

上位 32bit: tile_id
下位 32bit: depth_key


pair_key

tile_id は、対応するタイルを表す ID です。

depth_key は、そのままの z 座標ではなく、sort した結果が手前から奥の順になるように作成した正の depth key です。
実装では、depth_keybitcast<u32>(depth) によって u32 として扱い、tile_id と組み合わせて pair_key に pack しています。

このような key にしておくと、pair_key を小さい順に sort するだけで、

  • tile_id ごとにまとまる
  • 同じ tile_id の中では手前から奥の順に並ぶ

という状態を作れます。

例えば、Gaussian A が tile 5, tile 6, tile 13 に重なる場合、次のような対応情報を作ります。

pair_key   = { tile_id: 5,  depth_key: A.depth_key }
pair_value = Gaussian A

pair_key   = { tile_id: 6,  depth_key: A.depth_key }
pair_value = Gaussian A

pair_key   = { tile_id: 13, depth_key: A.depth_key }
pair_value = Gaussian A

この時点では、対応情報はまだタイルごとにまとまっていません。
このままでは、レンダリング時に、特定のタイルに属する Gaussian 一覧をすぐに取り出せません。

そのため、次の Radix Sort Pass で pair_key を使って並び替えます。

4.2.4 実装

Duplicate Pass の実装は以下を参照してください。

4.3 Radix Sort Pass

Radix Sort Pass では、pair_key を使ってタイルと Gaussian の対応情報を並び替えます。

目指す並びは、次のような形です。

tile 0 の Gaussian
tile 0 の Gaussian
tile 0 の Gaussian
tile 1 の Gaussian
tile 1 の Gaussian
tile 2 の Gaussian
...

さらに、同じタイルの中では奥行き順に並びます。

tile 0 の Gaussian  // 手前
tile 0 の Gaussian
tile 0 の Gaussian  // 奥
tile 1 の Gaussian  // 手前
tile 1 の Gaussian  // 奥
tile 2 の Gaussian
...

これにより、後続のレンダリングでは、各タイルが対応配列のどの範囲に対応するかを求められるようになります。

4.3.1 Radix Sortとは

Radix Sort は、key を小さな単位に分けて、複数回に分けて並び替える sort 方法です。

実装では、64bit の pair_key を 8bit ずつ処理します。

そのため、Radix Sort Pass は 8 回 (= 64 bit / 8 bit) の pass に分かれています。

pass 0: depth の下位 8bit
pass 1: depth の次の 8bit
pass 2: depth の次の 8bit
pass 3: depth の上位 8bit
pass 4: tile_id の下位 8bit
pass 5: tile_id の次の 8bit
pass 6: tile_id の次の 8bit
pass 7: tile_id の上位 8bit

4.3.2 8 bit keyのカウント

Radix Sortでは、まず初めに、8bit の値、つまり 0 ~ 255 のそれぞれの要素数を数えます。

このとき、各要素がどの bin に入るかを数えます。

例えば、8 bit keys が以下のようになっているとします。

keys: 
[ 4, 1, 4, 0, 2, 1, 4, 0 ]

この場合、bin ごとの個数は次のようになります。

bin 0 -> 2 個
bin 1 -> 2 個
bin 2 -> 1 個
bin 3 -> 0 個
bin 4 -> 3 個

1次元配列として書くと、次のようにになります。

histgram:
[2, 2, 1, 0, 3]

各 index は bin に対応しています。

index 0 -> bin 0 の個数
index 1 -> bin 1 の個数
index 2 -> bin 2 の個数
index 3 -> bin 3 の個数
index 4 -> bin 4 の個数
...

この histgram が分かれば、ソート後の形は分かります。

sorted:
[0, 0, 1, 1, 2, 4, 4, 4]


Histgram


実装では、workgroupサイズを 256 としています。

8bit 値は 0 ~ 255 の 256 種類なので、workgroup 内の各 thread が 1 つの bin を担当できます。

thread 0   -> value 0 の個数を担当
thread 1   -> value 1 の個数を担当
thread 2   -> value 2 の個数を担当
...
thread 255 -> value 255 の個数を担当

つまり、1 つの workgroup 内で、現在見ている 8bit の全パターンについて個数を数えます。

4.3.3 Prefix Scan の応用

bin ごとの個数の配列からソート結果を出力先の1次元配列に書き込むには、書き込む位置を求める必要があります。

ここで、前述した Prefix Scan が使えます。

先ほどの結果に対して Prefix Scan (exclusive scan) を行うと、以下のようになります。

[2, 2, 1, 0, 3]



exclusive scan:
[0, 2, 4, 5, 5]

これは、各値の書き込み開始位置を表します。

bin 0 -> output[0..2)
bin 1 -> output[2..4)
bin 2 -> output[4..5)
bin 3 -> output[5..5)
bin 4 -> output[5..8)

つまり Radix Sort Pass でも、Prefix Scan Pass と同じような考え方を使います。


ただし、実際には workgroup ごとに histgram を作ります。

例えば、各 workgroup が持つ bin の個数が次のようになっているとします。

bin 0 bin 1 bin 2
workgroup 0 2 個 1 個 0 個
workgroup 1 3 個 0 個 2 個
workgroup 2 1 個 4 個 1 個

まず、bin ごとの合計を求めます。

bin 0: total = 2 + 3 + 1 = 6
bin 1: total = 1 + 0 + 4 = 5
bin 2: total = 0 + 2 + 1 = 3

この合計に対して Prefix Scan (exclusive scan) を行うと、bin 全体の開始位置が分かります。

bin totals:
[6, 5, 3]

exclusive scan:
[0, 6, 11]

つまり、出力配列全体では次の範囲になります。

bin 0 -> output[0..6)
bin 1 -> output[6..11)
bin 2 -> output[11..14)

次に、同じ bin の中で、workgroup ごとの書き込み開始位置を求めます。

例えば bin 0 について見ると、

bin 0:
  workgroup 0 に 2 個
  workgroup 1 に 3 個
  workgroup 2 に 1 個

なので、各 workgroup は次の範囲に書き込みます。

bin 0:
  workgroup 0 -> output[0..2)
  workgroup 1 -> output[2..5)
  workgroup 2 -> output[5..6)

つまり、各 workgroup は以下の 2 つを知る必要があります。

  1. bin 全体の開始位置
  2. 自分より前の workgroup にある同じ bin の個数

これらを足すことで、その workgroup がその bin をどこから書けばよいかが分かります。


各 workgroup の bin の書き込み位置


Prefix Scan Pass では、入力配列全体を処理する必要がありました。
そのため、入力が大きい場合は workgroup をまたぎ、block_sumsAdd Block Offsets による補正が必要でした。

一方、Radix Sort の Prefix Scan する対象は、各 bin の合計数です。

bin totals:
[bin 0 の合計, bin 1 の合計, ..., bin 255 の合計]

8bit の値は 0..255 なので、bin の数は 256 個です。

この実装では workgroupサイズ = 256 としているため、256 個の bin totals を 1 workgroup 内で Prefix Scan できます。

thread 0   -> bin 0
thread 1   -> bin 1
thread 2   -> bin 2
...
thread 255 -> bin 255

そのため、Radix Sort の bin totals に対しては、Prefix Scan Pass のような多段階の Add Block Offsets は必要ありません。


ここまでで、各 bin の書き込み開始位置が求まりました。

ただし、同じ workgroup 内にも、同じ bin に属する要素が複数存在します。

例えば、ある workgroup 内で、各 thread が次の bin に属しているとします。

thread:  0    1    2    3    4    5    6    7
bin:     3    1    3    0    3    1    0    3

このとき、bin 3 に属する要素は次の 4 つです。

thread 0
thread 2
thread 4
thread 7

そのため、同じ bin の中で自分が何番目かを求める必要があります。

実装では、各 thread が自分の属する bin に flag bit を立てます。
その後、自分より前に立っている bit の数を数えることで、同じ bin 内での順番を求めます。

例えば、先ほどの例でにおいて、bin 3の flag bit は以下のようになります。

bin 3 flags:
thread: 0 1 2 3 4 5 6 7
flag:   1 0 1 0 1 0 0 1

thread 4 より前にある bin 3 の要素は、thread 0 と thread 2 の 2 個です。

そのため、thread 4 は bin 3 の中では 2 番目の位置に書き込めばよいことが分かります。

これを用いて、pair_keypair_value を同じ位置へ書き込むことで、Radix Sortを実現しています。

4.3.4 Ping-Pong Buffer

Radix Sort は 8bit ずつ、複数回に分けて並び替えます。

そのため、ある pass の出力を、次の pass の入力として使う必要があります。

ただし、8 回の pass それぞれに専用のバッファを用意する必要はありません。

実装では、入力用と一時出力用のバッファを用意し、pass ごとに入れ替えて使います。

pass 0:
  pair_keys        -> pair_keys_tmp
  pair_values      -> pair_values_tmp

pass 1:
  pair_keys_tmp    -> pair_keys
  pair_values_tmp  -> pair_values

pass 2:
  pair_keys        -> pair_keys_tmp
  pair_values      -> pair_values_tmp

pass 3:
  pair_keys_tmp    -> pair_keys
  pair_values_tmp  -> pair_values
...

このように、読み込み元と書き込み先を交互に使うことで、前の pass の結果を次の pass の入力として使えます。

4.3.5 実装

Radix Sort の実装は、Mirco Werner 氏の VkRadixSort を参考にしています。

詳しくは以下を参照してください。

4.4 Tile Range Pass

Radix Sort Pass によって、タイルと Gaussian の対応配列は、tile_id 順に並びました。

この状態になっていれば、同じタイルに属する Gaussian は連続しています。

次に必要なのは、各タイルがこの配列のどの範囲を読めばよいかです。

そこで Tile Range Pass では、次のような tile_ranges を作ります。

tile_ranges[tile_id] = { start, end }

例えば、以下のような状態であれば、

tile 0 の Gaussian
tile 0 の Gaussian
tile 0 の Gaussian
tile 1 の Gaussian
tile 1 の Gaussian
tile 2 の Gaussian

tile_ranges は以下のようになります。

tile_ranges[0] = { start: 0, end: 3 }
tile_ranges[1] = { start: 3, end: 5 }
tile_ranges[2] = { start: 5, end: 6 }

これにより、後続のレンダリングでは、特定のタイルの Gaussian だけをすぐに取り出せます。

4.4.1 境界を見つける

Tile Range Pass では、ソート済みの対応配列を 1 要素ずつ確認します。

各要素について、

前の要素と tile_id が違うか
次の要素と tile_id が違うか

を調べます。

前の要素と tile_id が違う場合、その位置はタイルの開始位置です。

tile_id:
[0, 0, 0, 1, 1, 2, 2, 2]
          ^
    tile 1 の start

次の要素と tile_id が違う場合、その次の位置は tile の終了位置です。

tile_id:
[0, 0, 0, 1, 1, 2, 2, 2]
       ^
  tile 0 の end = 3

実装では、各 thread がソート済みの pair_key の各 key を担当します。

各 thread は、自分が担当する pair_key の前後を確認し、その位置がタイル範囲の開始位置または終了位置かどうかを判定します。

4.4.2 空のタイル

すべてのタイルに Gaussian が存在するとは限りません。

例えば、タイル 3 に対応する Gaussian が 1 つもない場合、ソート済みの対応配列にはタイル 3 の要素が現れません。

tile_id:
[0, 0, 0, 1, 1, 2, 2, 2, 4, 4]

この場合、Tile Range Pass では tile_ranges[3] は更新されません。

そのため、この実装では Tile Range Pass の前に tile_ranges を初期化する Clear Tile Ranges Pass を用意しています。

これにより、Gaussian が存在しないタイルは空の範囲として扱えます。

tile_ranges[3] = { start: 0, end: 0 }

4.4.3 実装

Tile Range Pass の実装は以下を参照してください。

4.5 Tile Render Pass

ここまでの pass によって、以下の情報が得られました。

  • タイルと Gaussian の対応配列(タイル/奥行き順にソート済み)
  • 各タイルが対応配列のどの範囲を読むかを表す tile_ranges

Tile Render Passではこれらの情報を使って効率的に各 pixel の色を計算してレンダリングします。

4.5.1 1 workgroup が 1 tile を担当する

この実装では、タイルサイズを 16 x 16 pixel としています。

Tile Render Pass では、1 workgroup が 1 tile を担当します。

workgroup (0, 0) -> tile (0, 0)
workgroup (1, 0) -> tile (1, 0)
workgroup (2, 0) -> tile (2, 0)

Rust 側でも、dispatch 数はタイル数に合わせており、画面上のタイルグリッド全体に対して compute shader を実行しています。

tile_render_pass.dispatch_workgroups(tiles_width, tiles_height, 1);

そして、workgroupサイズはタイルサイズと同じ 256 としており、各 thread が タイル内の 1 pixel を担当しています。

4.5.2 Gaussian 範囲を取得する

Tile Range Pass で作成した tile_ranges から、workgroupに対応するタイルに属する Gaussian 範囲を取得できます。

例えば タイル 5 であれば、

tile_ranges[5] = { start: 120, end: 145 }

そして、tile 5 で処理すべき Gaussian は次の範囲にあります。

sorted_ids[120..145]

これにより、Radix Sort 後の pair_value、つまり、タイル/奥行き順に並んだ Gaussian index の一覧を取得できます。

そのため、Tile Render Pass では、全 Gaussian を走査する必要がありません。

4.5.3 Shared Memory による batch 処理

Tile Render Pass では、Gaussian の一覧 (start..end) を一定数ずつ batch に分けて処理します。

実装では、1 batch のサイズは workgroup サイズ、つまり 1 タイル内の pixel 数と同じ 256 です。

同じ workgroup 内の 各 thread は、それぞれ 1つの Gaussian 情報を Shared Memory に書き込みます。

thread 0   -> Gaussian 0 を書き込む
thread 1   -> Gaussian 1 を書き込む
thread 2   -> Gaussian 2 を書き込む
...
thread 255 -> Gaussian 255 を書き込む

読み込む情報は、Preprocess Pass で作成したものです。

uv      // 画面上の中心座標
conic   // 2D共分散行列の逆行列
opacity // 透明度
color   // 色

その後、workgroupBarrier() によって、全 thread の読み込みが完了してから pixel 評価へ進みます。

この batch 処理により、workgroup 内 (= タイル内) の 各 pixel が同じ Gaussian 情報を共有できます。

また、Shared Memory は、一般にメモリアクセス速度が global memory よりも速いため、メモリアクセス負荷を軽減する恩恵も得られます。

4.5.4 Gaussian を評価する

各 pixel は、その tile に属する Gaussian を順に評価します。

ここで Preprocess Pass でも説明した 2D Gaussian の式を使います。

G_{\text{2D}}(p) = \exp\left( -\frac{1}{2} (p - \mu_{\text{2D}})^T \Sigma_{\text{2D}}^{-1} (p - \mu_{\text{2D}}) \right)
  • p は pixel 位置
  • \mu_{\text{2D}} は Gaussian の画面上の中心位置
  • \Sigma_{\text{2D}}^{-1} は 2D 共分散行列の逆行列

先ほど batch として Shared Memory に書き出した uv/conic を上式に適用することにより、Gaussian の寄与を得ることができます。

最終的な alpha は、opacity と Gaussian の寄与を掛け合わせて求めます。

// G: ガウシアンの寄与
alpha = opacity * G

4.5.5 alpha 合成

タイルと Gaussian の対応配列について、Radix Sort Pass によって、同じタイル内の Gaussian は奥行き順にソートされています。

そのため、Tile Render Pass では 属しているタイルの範囲を順に処理するだけで、奥行き順に alpha 合成ができます。

alpha 合成では、各 Gaussian の色の寄与を accum に累積していきます。
また、奥側にどれだけ色が届くかを表す透過率 T を使います。

// 疑似コード

var T = 1.0;                  // 奥側に残っている透過率
var accum = vec3<f32>(0.0);   // 累積する色

// この pixel が属するタイルの Gaussian を、手前から奥へ順に処理する
for (/* ソート済みの Gaussian 一覧 */) {
    let G = eval_gaussian(pixel, gaussian);
    let alpha = opacity * G;
    let color = gaussian.color;

    accum += color * alpha * T; 

    T *= 1.0 - alpha;

    if (T < 1e-2) {
        break;
    }
}

手前の Gaussian から順に合成していくため、奥にある Gaussian の寄与は、手前の Gaussian の alpha によって弱くなります。

また、T が十分小さくなった場合、それ以上奥の Gaussian はほとんど見えません。
実装では、T < 1e-2 になったら、その pixel の処理を終了します。

最終的に得られた accum を pixel の色として Render Texture に書き込みます。

4.5.6 実装

Tile Render Pass の実装は以下を参照してください

まとめ

この記事では、3DGS のラスタライズの流れを見てきました。

3DGS の描画では、各 Gaussian を画面上に投影し、pixel ごとに寄与を計算して alpha 合成します。
ただし、すべての pixel がすべての Gaussian を処理すると重いため、今回は Tile-Based Rendering を使って処理対象を絞り込みました。

実装では、まず Preprocess Pass で各 Gaussian の画面上の位置、形状、色、透明度を計算しました。
次に Prefix Scan / Duplicate / Radix Sort / Tile Range Pass によって、タイルごとに処理すべき Gaussian の範囲を作りました。

最後に Tile Render Pass で、各 pixel が自分のタイルに属する Gaussian だけを評価し、奥行き順に alpha 合成して描画しました。

これらは全て GPU 処理の組み合わせで構成されており、リアルタイムな描画を可能にしています。

参照

https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting

https://www.ppsloan.org/publications/StupidSH36.pdf

https://developer.nvidia.com/gpugems/gpugems3/part-vi-gpu-computing/chapter-39-parallel-prefix-sum-scan-cuda

https://github.com/MircoWerner/VkRadixSort

https://sotrh.github.io/learn-wgpu

脚注
  1. https://sotrh.github.io/learn-wgpu/beginner/tutorial4-buffer ↩︎

  2. Peter-Pike Sloan, “Stupid Spherical Harmonics (SH) Tricks”, 2008, Appendix A2 “Polynomial Forms of SH Basis”. https://www.ppsloan.org/publications/StupidSH36.pdf ↩︎

  3. 実装では、厳密な矩形と楕円の交差判定ではなく、タイル矩形を外接円で近似した判定を行なっています。 ↩︎

Discussion