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