Skip to content

Commit d343e4d

Browse files
mvaligurskyMartin Valigursky
mvaligursky
and
Martin Valigursky
authored
First few gsplat chunks converted to WGSL (#7648)
Co-authored-by: Martin Valigursky <mvaligursky@snapchat.com>
1 parent 026afc1 commit d343e4d

File tree

5 files changed

+360
-81
lines changed

5 files changed

+360
-81
lines changed

‎src/scene/shader-lib/chunks-wgsl/chunks-wgsl.js

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,16 +46,16 @@ import gammaPS from './common/frag/gamma.js';
4646
import glossPS from './standard/frag/gloss.js';
4747
// import gsplatCenterVS from './gsplat/vert/gsplatCenter.js';
4848
// import gsplatColorVS from './gsplat/vert/gsplatColor.js';
49-
// import gsplatCommonVS from './gsplat/vert/gsplatCommon.js';
49+
import gsplatCommonVS from './gsplat/vert/gsplatCommon.js';
5050
// import gsplatCompressedDataVS from './gsplat/vert/gsplatCompressedData.js';
5151
// import gsplatCompressedSHVS from './gsplat/vert/gsplatCompressedSH.js';
5252
// import gsplatCornerVS from './gsplat/vert/gsplatCorner.js';
5353
// import gsplatDataVS from './gsplat/vert/gsplatData.js';
5454
// import gsplatOutputVS from './gsplat/vert/gsplatOutput.js';
55-
// import gsplatPS from './gsplat/frag/gsplat.js';
55+
import gsplatPS from './gsplat/frag/gsplat.js';
5656
// import gsplatSHVS from './gsplat/vert/gsplatSH.js';
5757
// import gsplatSourceVS from './gsplat/vert/gsplatSource.js';
58-
// import gsplatVS from './gsplat/vert/gsplat.js';
58+
import gsplatVS from './gsplat/vert/gsplat.js';
5959
import immediateLinePS from './internal/frag/immediateLine.js';
6060
import immediateLineVS from './internal/vert/immediateLine.js';
6161
import iridescenceDiffractionPS from './lit/frag/iridescenceDiffraction.js';
@@ -262,15 +262,15 @@ const shaderChunksWGSL = {
262262
// gsplatCenterVS,
263263
// gsplatCornerVS,
264264
// gsplatColorVS,
265-
// gsplatCommonVS,
265+
gsplatCommonVS,
266266
// gsplatCompressedDataVS,
267267
// gsplatCompressedSHVS,
268268
// gsplatDataVS,
269269
// gsplatOutputVS,
270-
// gsplatPS,
270+
gsplatPS,
271271
// gsplatSHVS,
272272
// gsplatSourceVS,
273-
// gsplatVS,
273+
gsplatVS,
274274
immediateLinePS,
275275
immediateLineVS,
276276
iridescenceDiffractionPS,

‎src/scene/shader-lib/chunks-wgsl/gsplat/frag/gsplat.js

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
export default /* wgsl */`
2+
3+
#ifndef DITHER_NONE
4+
#include "bayerPS"
5+
#include "opacityDitherPS"
6+
varying id: f32;
7+
#endif
8+
9+
#ifdef PICK_PASS
10+
#include "pickPS"
11+
#endif
12+
13+
#if defined(SHADOW_PASS) || defined(PICK_PASS) || defined(PREPASS_PASS)
14+
uniform alphaClip: f32;
15+
#endif
16+
17+
#ifdef PREPASS_PASS
18+
varying vLinearDepth: f32;
19+
#include "floatAsUintPS"
20+
#endif
21+
22+
varying gaussianUV: vec2f;
23+
varying gaussianColor: vec4f;
24+
25+
@fragment
26+
fn fragmentMain(input: FragmentInput) -> FragmentOutput {
27+
var output: FragmentOutput;
28+
29+
let A: f32 = dot(gaussianUV, gaussianUV);
30+
if (A > 1.0) {
31+
discard;
32+
return output;
33+
}
34+
35+
// evaluate alpha
36+
var alpha: f32 = exp(-A * 4.0) * gaussianColor.a;
37+
38+
#if defined(SHADOW_PASS) || defined(PICK_PASS) || defined(PREPASS_PASS)
39+
if (alpha < uniform.alphaClip) {
40+
discard;
41+
return output;
42+
}
43+
#endif
44+
45+
#ifdef PICK_PASS
46+
47+
output.color = getPickOutput();
48+
49+
#elif SHADOW_PASS
50+
51+
output.color = vec4f(0.0, 0.0, 0.0, 1.0);
52+
53+
#elif PREPASS_PASS
54+
55+
output.color = float2vec4(vLinearDepth);
56+
57+
#else
58+
59+
if (alpha < (1.0 / 255.0)) {
60+
discard;
61+
return output;
62+
}
63+
64+
#ifndef DITHER_NONE
65+
opacityDither(&alpha, id * 0.013);
66+
#endif
67+
68+
output.color = vec4f(input.gaussianColor.xyz * alpha, alpha);
69+
#endif
70+
71+
return output;
72+
}`;

‎src/scene/shader-lib/chunks-wgsl/gsplat/vert/gsplat.js

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
export default /* wgsl */`
2+
#include "gsplatCommonVS"
3+
4+
varying gaussianUV: vec2f;
5+
varying gaussianColor: vec4f;
6+
7+
#ifndef DITHER_NONE
8+
varying id: f32;
9+
#endif
10+
11+
const discardVec: vec4f = vec4f(0.0, 0.0, 2.0, 1.0);
12+
13+
#ifdef PREPASS_PASS
14+
varying vLinearDepth: f32;
15+
#endif
16+
17+
@vertex
18+
fn vertexMain(input: VertexInput) -> VertexOutput {
19+
var output: VertexOutput;
20+
21+
// read gaussian details
22+
var source: SplatSource;
23+
if (!initSource(&source)) {
24+
output.position = discardVec;
25+
return output;
26+
}
27+
28+
let modelCenter: vec3f = readCenter(&source);
29+
30+
var center: SplatCenter;
31+
if (!initCenter(modelCenter, &center)) {
32+
output.position = discardVec;
33+
return output;
34+
}
35+
36+
// project center to screen space
37+
var corner: SplatCorner;
38+
if (!initCorner(&source, &center, &corner)) {
39+
output.position = discardVec;
40+
return output;
41+
}
42+
43+
// read color
44+
var clr: vec4f = readColor(&source);
45+
46+
#if GSPLAT_AA
47+
// apply AA compensation
48+
clr.a = clr.a * corner.aaFactor;
49+
#endif
50+
51+
// evaluate spherical harmonics
52+
#if SH_BANDS > 0
53+
// calculate the model-space view direction
54+
let modelView3x3 = mat3x3f(center.modelView[0].xyz, center.modelView[1].xyz, center.modelView[2].xyz);
55+
let dir = normalize(modelView3x3 * center.view);
56+
clr = vec4f(clr.xyz + evalSH(&source, dir), clr.a);
57+
#endif
58+
59+
clipCorner(&corner, clr.w);
60+
61+
// write output
62+
output.position = center.proj + vec4f(corner.offset, 0.0, 0.0);
63+
output.gaussianUV = corner.uv;
64+
output.gaussianColor = vec4f(prepareOutputFromGamma(max(clr.xyz, vec3f(0.0))), clr.w);
65+
66+
#ifndef DITHER_NONE
67+
output.id = f32(source.id);
68+
#endif
69+
70+
#ifdef PREPASS_PASS
71+
output.vLinearDepth = -center.view.z;
72+
#endif
73+
74+
return output;
75+
}
76+
`;

‎src/scene/shader-lib/chunks-wgsl/gsplat/vert/gsplatCommon.js

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
export default /* wgsl */`
2+
3+
// stores the source UV and order of the splat
4+
struct SplatSource {
5+
order: u32, // render order
6+
id: u32, // splat id
7+
uv: vec2<i32>, // splat uv
8+
cornerUV: vec2f, // corner coordinates for this vertex of the gaussian (-1, -1)..(1, 1)
9+
}
10+
11+
// stores the camera and clip space position of the gaussian center
12+
struct SplatCenter {
13+
view: vec3f, // center in view space
14+
proj: vec4f, // center in clip space
15+
modelView: mat4x4f, // model-view matrix
16+
projMat00: f32, // elememt [0][0] of the projection matrix
17+
}
18+
19+
// stores the offset from center for the current gaussian
20+
struct SplatCorner {
21+
offset: vec2f, // corner offset from center in clip space
22+
uv: vec2f, // corner uv
23+
#if GSPLAT_AA
24+
aaFactor: f32, // for scenes generated with antialiasing
25+
#endif
26+
}
27+
28+
#if GSPLAT_COMPRESSED_DATA == true
29+
#include "gsplatCompressedDataVS"
30+
#include "gsplatCompressedSHVS"
31+
#elif GSPLAT_SOGS_DATA == true
32+
#include "gsplatSogsDataVS"
33+
#include "gsplatSogsColorVS"
34+
#include "gsplatSogsSHVS"
35+
#else
36+
#include "gsplatDataVS"
37+
#include "gsplatColorVS"
38+
#include "gsplatSHVS"
39+
#endif
40+
41+
#include "gsplatSourceVS"
42+
#include "gsplatCenterVS"
43+
#include "gsplatCornerVS"
44+
#include "gsplatOutputVS"
45+
46+
// modify the gaussian corner so it excludes gaussian regions with alpha less than 1/255
47+
fn clipCorner(corner: ptr<function, SplatCorner>, alpha: f32) {
48+
let clip: f32 = min(1.0, sqrt(-log(1.0 / (255.0 * alpha))) / 2.0);
49+
corner.offset = corner.offset * clip;
50+
corner.uv = corner.uv * clip;
51+
}
52+
53+
// spherical Harmonics
54+
55+
#if SH_BANDS > 0
56+
const SH_C1: f32 = 0.4886025119029199;
57+
58+
#if SH_BANDS > 1
59+
const SH_C2_0: f32 = 1.0925484305920792;
60+
const SH_C2_1: f32 = -1.0925484305920792;
61+
const SH_C2_2: f32 = 0.31539156525252005;
62+
const SH_C2_3: f32 = -1.0925484305920792;
63+
const SH_C2_4: f32 = 0.5462742152960396;
64+
#endif
65+
66+
#if SH_BANDS > 2
67+
const SH_C3_0: f32 = -0.5900435899266435;
68+
const SH_C3_1: f32 = 2.890611442640554;
69+
const SH_C3_2: f32 = -0.4570457994644658;
70+
const SH_C3_3: f32 = 0.3731763325901154;
71+
const SH_C3_4: f32 = -0.4570457994644658;
72+
const SH_C3_5: f32 = 1.445305721320277;
73+
const SH_C3_6: f32 = -0.5900435899266435;
74+
#endif
75+
76+
// see https://github.com/graphdeco-inria/gaussian-splatting/blob/main/utils/sh_utils.py
77+
fn evalSH(source: ptr<function, SplatSource>, dir: vec3f) -> vec3f {
78+
79+
#if SH_BANDS > 0
80+
#if SH_BANDS == 1
81+
var sh: array<vec3f, 3>;
82+
#elif SH_BANDS == 2
83+
var sh: array<vec3f, 8>;
84+
#elif SH_BANDS == 3
85+
var sh: array<vec3f, 15>;
86+
#endif
87+
#endif
88+
89+
var scale: f32;
90+
readSHData(source, &sh, &scale);
91+
92+
let x = dir.x;
93+
let y = dir.y;
94+
let z = dir.z;
95+
96+
// 1st degree
97+
var result = SH_C1 * (-sh[0] * y + sh[1] * z - sh[2] * x);
98+
99+
#if SH_BANDS > 1
100+
// 2nd degree
101+
let xx = x * x;
102+
let yy = y * y;
103+
let zz = z * z;
104+
let xy = x * y;
105+
let yz = y * z;
106+
let xz = x * z;
107+
108+
result = result + (
109+
sh[3] * (SH_C2_0 * xy) +
110+
sh[4] * (SH_C2_1 * yz) +
111+
sh[5] * (SH_C2_2 * (2.0 * zz - xx - yy)) +
112+
sh[6] * (SH_C2_3 * xz) +
113+
sh[7] * (SH_C2_4 * (xx - yy))
114+
);
115+
#endif
116+
117+
#if SH_BANDS > 2
118+
// 3rd degree
119+
result = result + (
120+
sh[8] * (SH_C3_0 * y * (3.0 * xx - yy)) +
121+
sh[9] * (SH_C3_1 * xy * z) +
122+
sh[10] * (SH_C3_2 * y * (4.0 * zz - xx - yy)) +
123+
sh[11] * (SH_C3_3 * z * (2.0 * zz - 3.0 * xx - 3.0 * yy)) +
124+
sh[12] * (SH_C3_4 * x * (4.0 * zz - xx - yy)) +
125+
sh[13] * (SH_C3_5 * z * (xx - yy)) +
126+
sh[14] * (SH_C3_6 * x * (xx - 3.0 * yy))
127+
);
128+
#endif
129+
130+
return result * scale;
131+
}
132+
#endif
133+
`;

0 commit comments

Comments
 (0)