1#define_import_path bevy_gaussian_splatting::texture
2
3#ifdef PRECOMPUTE_COVARIANCE_3D
4#import bevy_gaussian_splatting::bindings::{
5    gaussian_uniforms,
6    position_visibility,
7    spherical_harmonics,
8    covariance_3d_opacity,
9};
10#else
11#import bevy_gaussian_splatting::bindings::{
12    gaussian_uniforms,
13    position_visibility,
14    spherical_harmonics,
15    rotation,
16    rotation_scale_opacity,
17    scale_opacity,
18};
19#endif
20
21#import bevy_gaussian_splatting::spherical_harmonics::{
22    shc,
23    spherical_harmonics_lookup,
24    srgb_to_linear,
25}
26
27fn location(index: u32) -> vec2<i32> {
28    return vec2<i32>(
29        i32(index) % i32(gaussian_uniforms.count_root_ceil),
30        i32(index) / i32(gaussian_uniforms.count_root_ceil),
31    );
32}
33
34#ifdef PLANAR_TEXTURE_F16
35
36fn get_position(index: u32) -> vec3<f32> {
37    let sample = textureLoad(
38        position_visibility,
39        location(index),
40        0,
41    );
42
43    return sample.xyz;
44}
45
46fn get_sh_vec(
47    index: u32,
48    plane: i32,
49) -> vec4<u32> {
50#if SH_VEC4_PLANES == 1
51    return textureLoad(
52        spherical_harmonics,
53        location(index),
54        plane,
55    );
56#else
57    return textureLoad(
58        spherical_harmonics,
59        location(index),
60        plane,
61        0,
62    );
63#endif
64}
65
66#ifdef WEBGL2
67fn get_color(
68    index: u32,
69    ray_direction: vec3<f32>,
70) -> vec3<f32> {
71    let s0 = get_sh_vec(index, 0);
72
73    let v0 = unpack2x16float(s0.x);
74    let v1 = unpack2x16float(s0.y);
75    let v2 = unpack2x16float(s0.z);
76    let v3 = unpack2x16float(s0.w);
77
78    let rds = ray_direction * ray_direction;
79    var color = vec3<f32>(0.5);
80
81    color += shc[ 0] * vec3<f32>(
82        v0.x,
83        v0.y,
84        v1.x,
85    );
86
87#if SH_COEFF_COUNT > 11
88    let r1 = vec3<f32>(
89        v1.y,
90        v2.x,
91        v2.y,
92    );
93
94    let s1 = get_sh_vec(index, 1);
95    let v4 = unpack2x16float(s1.x);
96    let v5 = unpack2x16float(s1.y);
97    let v6 = unpack2x16float(s1.z);
98    let v7 = unpack2x16float(s1.w);
99
100    let r2 = vec3<f32>(
101        v3.x,
102        v3.y,
103        v4.x,
104    );
105
106    let r3 = vec3<f32>(
107        v4.y,
108        v5.x,
109        v5.y,
110    );
111
112    color += shc[ 1] * r1 * ray_direction.y;
113    color += shc[ 2] * r2 * ray_direction.z;
114    color += shc[ 3] * r3 * ray_direction.x;
115#endif
116
117#if SH_COEFF_COUNT > 26
118    let r4 = vec3<f32>(
119        v6.x,
120        v6.y,
121        v7.x,
122    );
123
124    let s2 = get_sh_vec(index, 2);
125    let v8 = unpack2x16float(s2.x);
126    let v9 = unpack2x16float(s2.y);
127    let v10 = unpack2x16float(s2.z);
128    let v11 = unpack2x16float(s2.w);
129
130    let r5 = vec3<f32>(
131        v7.y,
132        v8.x,
133        v8.y,
134    );
135
136    let r6 = vec3<f32>(
137        v9.x,
138        v9.y,
139        v10.x,
140    );
141
142    let r7 = vec3<f32>(
143        v10.y,
144        v11.x,
145        v11.y,
146    );
147
148    let s3 = get_sh_vec(index, 3);
149    let v12 = unpack2x16float(s3.x);
150    let v13 = unpack2x16float(s3.y);
151    let v14 = unpack2x16float(s3.z);
152    let v15 = unpack2x16float(s3.w);
153
154    let r8 = vec3<f32>(
155        v12.x,
156        v12.y,
157        v13.x,
158    );
159
160    color += shc[ 4] * r4 * ray_direction.x * ray_direction.y;
161    color += shc[ 5] * r5 * ray_direction.y * ray_direction.z;
162    color += shc[ 6] * r6 * (2.0 * rds.z - rds.x - rds.y);
163    color += shc[ 7] * r7 * ray_direction.x * ray_direction.z;
164    color += shc[ 8] * r8 * (rds.x - rds.y);
165#endif
166
167#if SH_COEFF_COUNT > 47
168    let r9 = vec3<f32>(
169        v13.y,
170        v14.x,
171        v14.y,
172    );
173
174    let s4 = get_sh_vec(index, 4);
175    let v16 = unpack2x16float(s4.x);
176    let v17 = unpack2x16float(s4.y);
177    let v18 = unpack2x16float(s4.z);
178    let v19 = unpack2x16float(s4.w);
179
180    let r10 = vec3<f32>(
181        v15.x,
182        v15.y,
183        v16.x,
184    );
185
186    let r11 = vec3<f32>(
187        v16.y,
188        v17.x,
189        v17.y,
190    );
191
192    let r12 = vec3<f32>(
193        v18.x,
194        v18.y,
195        v19.x,
196    );
197
198    let s5 = get_sh_vec(index, 5);
199    let v20 = unpack2x16float(s5.x);
200    let v21 = unpack2x16float(s5.y);
201    let v22 = unpack2x16float(s5.z);
202    let v23 = unpack2x16float(s5.w);
203
204    let r13 = vec3<f32>(
205        v19.y,
206        v20.x,
207        v20.y,
208    );
209
210    let r14 = vec3<f32>(
211        v21.x,
212        v21.y,
213        v22.x,
214    );
215
216    let r15 = vec3<f32>(
217        v22.y,
218        v23.x,
219        v23.y,
220    );
221
222    color += shc[ 9] * r9 * ray_direction.y * (3.0 * rds.x - rds.y);
223    color += shc[10] * r10 * ray_direction.x * ray_direction.y * ray_direction.z;
224    color += shc[11] * r11 * ray_direction.y * (4.0 * rds.z - rds.x - rds.y);
225    color += shc[12] * r12 * ray_direction.z * (2.0 * rds.z - 3.0 * rds.x - 3.0 * rds.y);
226    color += shc[13] * r13 * ray_direction.x * (4.0 * rds.z - rds.x - rds.y);
227    color += shc[14] * r14 * ray_direction.z * (rds.x - rds.y);
228    color += shc[15] * r15 * ray_direction.x * (rds.x - 3.0 * rds.y);
229#endif
230
231    return srgb_to_linear(color);
232}
233#else
234fn get_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
235    var coefficients: array<f32, #{SH_COEFF_COUNT}>;
236
237    for (var i = 0u; i < #{SH_VEC4_PLANES}u; i = i + 1u) {
238        let sample = get_sh_vec(index, i32(i));
239
240        let v0 = unpack2x16float(sample.x);
241        let v1 = unpack2x16float(sample.y);
242        let v2 = unpack2x16float(sample.z);
243        let v3 = unpack2x16float(sample.w);
244
245        let base_index = i * 8u;
246        coefficients[base_index     ] = v0.x;
247        coefficients[base_index + 1u] = v0.y;
248
249        coefficients[base_index + 2u] = v1.x;
250        coefficients[base_index + 3u] = v1.y;
251
252        coefficients[base_index + 4u] = v2.x;
253        coefficients[base_index + 5u] = v2.y;
254
255        coefficients[base_index + 6u] = v3.x;
256        coefficients[base_index + 7u] = v3.y;
257    }
258
259    return coefficients;
260}
261
262fn get_color(
263    index: u32,
264    ray_direction: vec3<f32>,
265) -> vec3<f32> {
266    let sh = get_spherical_harmonics(index);
267    let color = spherical_harmonics_lookup(ray_direction, sh);
268    return srgb_to_linear(color);
269}
270#endif
271
272#ifdef PRECOMPUTE_COVARIANCE_3D
273    fn get_cov3d(index: u32) -> array<f32, 6> {
274        let sample = textureLoad(
275            covariance_3d_opacity,
276            location(index),
277            0,
278        );
279
280        let c0 = unpack2x16float(sample.x);
281        let c1 = unpack2x16float(sample.y);
282        let c2 = unpack2x16float(sample.z);
283
284        var cov3d: array<f32, 6>;
285
286        cov3d[0] = c0.y;
287        cov3d[1] = c0.x;
288        cov3d[2] = c1.y;
289        cov3d[3] = c1.x;
290        cov3d[4] = c2.y;
291        cov3d[5] = c2.x;
292
293        return cov3d;
294    }
295#else
296    fn get_rotation(index: u32) -> vec4<f32> {
297        let sample = textureLoad(
298            rotation_scale_opacity,
299            location(index),
300            0,
301        );
302
303        let q0 = unpack2x16float(sample.x);
304        let q1 = unpack2x16float(sample.y);
305
306        return vec4<f32>(
307            q0.yx,
308            q1.yx,
309        );
310    }
311
312    fn get_scale(index: u32) -> vec3<f32> {
313        let sample = textureLoad(
314            rotation_scale_opacity,
315            location(index),
316            0,
317        );
318
319        let s0 = unpack2x16float(sample.z);
320        let s1 = unpack2x16float(sample.w);
321
322        return vec3<f32>(
323            s0.yx,
324            s1.y,
325        );
326    }
327#endif
328
329fn get_opacity(index: u32) -> f32 {
330#ifdef PRECOMPUTE_COVARIANCE_3D
331    let sample = textureLoad(
332        covariance_3d_opacity,
333        location(index),
334        0,
335    );
336
337    return unpack2x16float(sample.w).y;
338#else
339    let sample = textureLoad(
340        rotation_scale_opacity,
341        location(index),
342        0,
343    );
344
345    return unpack2x16float(sample.w).x;
346#endif
347}
348
349fn get_visibility(index: u32) -> f32 {
350    let sample = textureLoad(
351        position_visibility,
352        location(index),
353        0,
354    );
355
356    return sample.w;
357}
358#endif
359
360// TODO: support f32
361#ifdef PLANAR_TEXTURE_F32
362fn get_position(index: u32) -> vec3<f32> {
363    return position_visibility[index].xyz;
364}
365
366fn get_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
367    return spherical_harmonics[index];
368}
369
370fn get_rotation(index: u32) -> vec4<f32> {
371    return rotation[index];
372}
373
374fn get_scale(index: u32) -> vec3<f32> {
375    return scale_opacity[index].xyz;
376}
377
378fn get_opacity(index: u32) -> f32 {
379    return scale_opacity[index].w;
380}
381
382fn get_visibility(index: u32) -> f32 {
383    return position_visibility[index].w;
384}
385#endif