1#define_import_path bevy_gaussian_splatting::planar
2
3#ifdef GAUSSIAN_3D_STRUCTURE
4    #ifdef PRECOMPUTE_COVARIANCE_3D
5        #import bevy_gaussian_splatting::bindings::{
6            position_visibility,
7            spherical_harmonics,
8            covariance_3d_opacity,
9        }
10
11        #ifdef BINARY_GAUSSIAN_OP
12            #import bevy_gaussian_splatting::bindings::{
13                rhs_position_visibility,
14                rhs_spherical_harmonics,
15                rhs_covariance_3d_opacity,
16                out_position_visibility,
17                out_spherical_harmonics,
18                out_covariance_3d_opacity,
19            }
20        #endif
21    #else
22        #import bevy_gaussian_splatting::bindings::{
23            position_visibility,
24            spherical_harmonics,
25            rotation,
26            rotation_scale_opacity,
27            scale_opacity,
28        }
29
30        #ifdef BINARY_GAUSSIAN_OP
31            #import bevy_gaussian_splatting::bindings::{
32                rhs_position_visibility,
33                rhs_spherical_harmonics,
34                out_position_visibility,
35                out_spherical_harmonics,
36            }
37
38            #ifdef PLANAR_F16
39                #import bevy_gaussian_splatting::bindings::rhs_rotation_scale_opacity
40                #import bevy_gaussian_splatting::bindings::out_rotation_scale_opacity
41            #endif
42
43            #ifdef PLANAR_F32
44                #import bevy_gaussian_splatting::bindings::{
45                    rhs_rotation,
46                    rhs_scale_opacity,
47                    out_rotation,
48                    out_scale_opacity,
49                }
50            #endif
51        #endif
52    #endif
53
54    #import bevy_gaussian_splatting::spherical_harmonics::{
55        spherical_harmonics_lookup,
56        srgb_to_linear,
57    }
58#else ifdef GAUSSIAN_4D
59    #import bevy_gaussian_splatting::bindings::{
60        position_visibility,
61        spherindrical_harmonics,
62        isotropic_rotations,
63        scale_opacity,
64        timestamp_timescale,
65    }
66
67    #ifdef BINARY_GAUSSIAN_OP
68        #import bevy_gaussian_splatting::bindings::{
69            rhs_position_visibility,
70            rhs_spherindrical_harmonics,
71            rhs_isotropic_rotations,
72            rhs_scale_opacity,
73            rhs_timestamp_timescale,
74        }
75    #endif
76
77    #import bevy_gaussian_splatting::spherical_harmonics::srgb_to_linear
78    #import bevy_gaussian_splatting::spherindrical_harmonics::spherindrical_harmonics_lookup
79#endif
80
81fn planar_position(value: vec4<f32>) -> vec3<f32> {
82    return value.xyz;
83}
84
85fn planar_visibility(value: vec4<f32>) -> f32 {
86    return value.w;
87}
88
89#ifdef GAUSSIAN_3D_STRUCTURE
90    fn planar_color_from_sh(
91        ray_direction: vec3<f32>,
92        sh: array<f32, #{SH_COEFF_COUNT}>,
93    ) -> vec3<f32> {
94        let color = spherical_harmonics_lookup(ray_direction, sh);
95        return srgb_to_linear(color);
96    }
97
98    fn planar_scale_from(scale_opacity: vec4<f32>) -> vec3<f32> {
99        return scale_opacity.xyz;
100    }
101
102    fn planar_opacity_from(scale_opacity: vec4<f32>) -> f32 {
103        return scale_opacity.w;
104    }
105
106    #ifdef PLANAR_F16
107        fn planar_f16_decode_sh(
108            raw: array<u32, #{HALF_SH_COEFF_COUNT}>,
109        ) -> array<f32, #{SH_COEFF_COUNT}> {
110            var coefficients: array<f32, #{SH_COEFF_COUNT}>;
111
112            for (var i = 0u; i < #{HALF_SH_COEFF_COUNT}u; i = i + 1u) {
113                let values = unpack2x16float(raw[i]);
114
115                coefficients[i * 2u] = values[0];
116                coefficients[i * 2u + 1u] = values[1];
117            }
118
119            return coefficients;
120        }
121
122        #ifdef PRECOMPUTE_COVARIANCE_3D
123            fn planar_f16_decode_covariance(raw: vec4<u32>) -> array<f32, 6> {
124                let c0 = unpack2x16float(raw.x);
125                let c1 = unpack2x16float(raw.y);
126                let c2 = unpack2x16float(raw.z);
127
128                var cov3d: array<f32, 6>;
129
130                cov3d[0] = c0.y;
131                cov3d[1] = c0.x;
132                cov3d[2] = c1.y;
133                cov3d[3] = c1.x;
134                cov3d[4] = c2.y;
135                cov3d[5] = c2.x;
136
137                return cov3d;
138            }
139
140            fn planar_f16_opacity(raw: vec4<u32>) -> f32 {
141                return unpack2x16float(raw.w).y;
142            }
143        #else
144            fn planar_f16_decode_rotation(raw: vec4<u32>) -> vec4<f32> {
145                let q0 = unpack2x16float(raw.x);
146                let q1 = unpack2x16float(raw.y);
147
148                return vec4<f32>(
149                    q0.yx,
150                    q1.yx,
151                );
152            }
153
154            fn planar_f16_decode_scale(raw: vec4<u32>) -> vec3<f32> {
155                let s0 = unpack2x16float(raw.z);
156                let s1 = unpack2x16float(raw.w);
157
158                return vec3<f32>(
159                    s0.yx,
160                    s1.y,
161                );
162            }
163
164            fn planar_f16_opacity(raw: vec4<u32>) -> f32 {
165                return unpack2x16float(raw.w).x;
166            }
167        #endif
168
169        fn get_color(
170            index: u32,
171            ray_direction: vec3<f32>,
172        ) -> vec3<f32> {
173            let sh = planar_f16_decode_sh(spherical_harmonics[index]);
174            return planar_color_from_sh(ray_direction, sh);
175        }
176
177        fn get_position(index: u32) -> vec3<f32> {
178            return planar_position(position_visibility[index]);
179        }
180
181        fn get_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
182            return planar_f16_decode_sh(spherical_harmonics[index]);
183        }
184
185        #ifdef PRECOMPUTE_COVARIANCE_3D
186            fn get_cov3d(index: u32) -> array<f32, 6> {
187                return planar_f16_decode_covariance(covariance_3d_opacity[index]);
188            }
189        #else
190            fn get_rotation(index: u32) -> vec4<f32> {
191                return planar_f16_decode_rotation(rotation_scale_opacity[index]);
192            }
193
194            fn get_scale(index: u32) -> vec3<f32> {
195                return planar_f16_decode_scale(rotation_scale_opacity[index]);
196            }
197        #endif
198
199        fn get_opacity(index: u32) -> f32 {
200            #ifdef PRECOMPUTE_COVARIANCE_3D
201                return planar_f16_opacity(covariance_3d_opacity[index]);
202            #else
203                return planar_f16_opacity(rotation_scale_opacity[index]);
204            #endif
205        }
206
207        fn get_visibility(index: u32) -> f32 {
208            return planar_visibility(position_visibility[index]);
209        }
210
211        #ifdef BINARY_GAUSSIAN_OP
212            fn get_rhs_color(
213                index: u32,
214                ray_direction: vec3<f32>,
215            ) -> vec3<f32> {
216                let sh = planar_f16_decode_sh(rhs_spherical_harmonics[index]);
217                return planar_color_from_sh(ray_direction, sh);
218            }
219
220            fn get_rhs_position(index: u32) -> vec3<f32> {
221                return planar_position(rhs_position_visibility[index]);
222            }
223
224            fn get_rhs_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
225                return planar_f16_decode_sh(rhs_spherical_harmonics[index]);
226            }
227
228            #ifdef PRECOMPUTE_COVARIANCE_3D
229                fn get_rhs_cov3d(index: u32) -> array<f32, 6> {
230                    return planar_f16_decode_covariance(rhs_covariance_3d_opacity[index]);
231                }
232            #else
233                fn get_rhs_rotation(index: u32) -> vec4<f32> {
234                    return planar_f16_decode_rotation(rhs_rotation_scale_opacity[index]);
235                }
236
237                fn get_rhs_scale(index: u32) -> vec3<f32> {
238                    return planar_f16_decode_scale(rhs_rotation_scale_opacity[index]);
239                }
240            #endif
241
242            fn get_rhs_opacity(index: u32) -> f32 {
243                #ifdef PRECOMPUTE_COVARIANCE_3D
244                    return planar_f16_opacity(rhs_covariance_3d_opacity[index]);
245                #else
246                    return planar_f16_opacity(rhs_rotation_scale_opacity[index]);
247                #endif
248            }
249
250            fn get_rhs_visibility(index: u32) -> f32 {
251                return planar_visibility(rhs_position_visibility[index]);
252            }
253
254            fn set_output_position_visibility(
255                index: u32,
256                position: vec3<f32>,
257                visibility: f32,
258            ) {
259                out_position_visibility[index] = vec4<f32>(position, visibility);
260            }
261
262            fn set_output_spherical_harmonics(
263                index: u32,
264                sh: array<f32, #{SH_COEFF_COUNT}>,
265            ) {
266                for (var i = 0u; i < #{HALF_SH_COEFF_COUNT}; i = i + 1u) {
267                    let base = i * 2u;
268                    out_spherical_harmonics[index][i] = pack2x16float(vec2<f32>(
269                        sh[base],
270                        sh[base + 1u],
271                    ));
272                }
273            }
274
275            #ifdef PRECOMPUTE_COVARIANCE_3D
276                fn planar_f16_encode_covariance(
277                    cov: array<f32, 6>,
278                    opacity: f32,
279                ) -> vec4<u32> {
280                    return vec4<u32>(
281                        pack2x16float(vec2<f32>(cov[1], cov[0])),
282                        pack2x16float(vec2<f32>(cov[3], cov[2])),
283                        pack2x16float(vec2<f32>(cov[5], cov[4])),
284                        pack2x16float(vec2<f32>(0.0, opacity)),
285                    );
286                }
287
288                fn set_output_covariance(
289                    index: u32,
290                    cov: array<f32, 6>,
291                    opacity: f32,
292                ) {
293                    out_covariance_3d_opacity[index] = planar_f16_encode_covariance(cov, opacity);
294                }
295            #else
296                fn planar_f16_encode_rotation_scale_opacity(
297                    rotation: vec4<f32>,
298                    scale: vec3<f32>,
299                    opacity: f32,
300                ) -> vec4<u32> {
301                    return vec4<u32>(
302                        pack2x16float(vec2<f32>(rotation.y, rotation.x)),
303                        pack2x16float(vec2<f32>(rotation.w, rotation.z)),
304                        pack2x16float(vec2<f32>(scale.y, scale.x)),
305                        pack2x16float(vec2<f32>(opacity, scale.z)),
306                    );
307                }
308
309                fn set_output_transform(
310                    index: u32,
311                    rotation: vec4<f32>,
312                    scale: vec3<f32>,
313                    opacity: f32,
314                ) {
315                    out_rotation_scale_opacity[index] = planar_f16_encode_rotation_scale_opacity(
316                        rotation,
317                        scale,
318                        opacity,
319                    );
320                }
321            #endif
322
323        #endif
324    #else ifdef PLANAR_F32
325        fn get_color(
326            index: u32,
327            ray_direction: vec3<f32>,
328        ) -> vec3<f32> {
329            return planar_color_from_sh(ray_direction, spherical_harmonics[index]);
330        }
331
332        fn get_position(index: u32) -> vec3<f32> {
333            return planar_position(position_visibility[index]);
334        }
335
336        fn get_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
337            return spherical_harmonics[index];
338        }
339
340        fn get_rotation(index: u32) -> vec4<f32> {
341            return rotation[index];
342        }
343
344        fn get_scale(index: u32) -> vec3<f32> {
345            return planar_scale_from(scale_opacity[index]);
346        }
347
348        fn get_opacity(index: u32) -> f32 {
349            return planar_opacity_from(scale_opacity[index]);
350        }
351
352        fn get_visibility(index: u32) -> f32 {
353            return planar_visibility(position_visibility[index]);
354        }
355
356        #ifdef BINARY_GAUSSIAN_OP
357            fn get_rhs_color(
358                index: u32,
359                ray_direction: vec3<f32>,
360            ) -> vec3<f32> {
361                return planar_color_from_sh(ray_direction, rhs_spherical_harmonics[index]);
362            }
363
364            fn get_rhs_position(index: u32) -> vec3<f32> {
365                return planar_position(rhs_position_visibility[index]);
366            }
367
368            fn get_rhs_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
369                return rhs_spherical_harmonics[index];
370            }
371
372            fn get_rhs_rotation(index: u32) -> vec4<f32> {
373                return rhs_rotation[index];
374            }
375
376            fn get_rhs_scale(index: u32) -> vec3<f32> {
377                return planar_scale_from(rhs_scale_opacity[index]);
378            }
379
380            fn get_rhs_opacity(index: u32) -> f32 {
381                return planar_opacity_from(rhs_scale_opacity[index]);
382            }
383
384            fn get_rhs_visibility(index: u32) -> f32 {
385                return planar_visibility(rhs_position_visibility[index]);
386            }
387
388            fn set_output_position_visibility(
389                index: u32,
390                position: vec3<f32>,
391                visibility: f32,
392            ) {
393                out_position_visibility[index] = vec4<f32>(position, visibility);
394            }
395
396            fn set_output_spherical_harmonics(
397                index: u32,
398                sh: array<f32, #{SH_COEFF_COUNT}>,
399            ) {
400                for (var i = 0u; i < #{SH_COEFF_COUNT}; i = i + 1u) {
401                    out_spherical_harmonics[index][i] = sh[i];
402                }
403            }
404
405            #ifdef PRECOMPUTE_COVARIANCE_3D
406                fn set_output_covariance(
407                    index: u32,
408                    cov: array<f32, 6>,
409                    opacity: f32,
410                ) {
411                    out_covariance_3d_opacity[index][0] = cov[0];
412                    out_covariance_3d_opacity[index][1] = cov[1];
413                    out_covariance_3d_opacity[index][2] = cov[2];
414                    out_covariance_3d_opacity[index][3] = cov[3];
415                    out_covariance_3d_opacity[index][4] = cov[4];
416                    out_covariance_3d_opacity[index][5] = cov[5];
417                    out_covariance_3d_opacity[index][6] = 0.0;
418                    out_covariance_3d_opacity[index][7] = opacity;
419                }
420            #else
421                fn set_output_transform(
422                    index: u32,
423                    rotation: vec4<f32>,
424                    scale: vec3<f32>,
425                    opacity: f32,
426                ) {
427                    out_rotation[index] = rotation;
428                    out_scale_opacity[index] = vec4<f32>(scale, opacity);
429                }
430            #endif
431
432        #endif
433    #endif
434#else ifdef GAUSSIAN_4D
435    fn planar4d_color_from_sh(
436        ray_direction: vec3<f32>,
437        dir_t: f32,
438        sh: array<f32, #{SH_4D_COEFF_COUNT}>,
439    ) -> vec3<f32> {
440        let color = spherindrical_harmonics_lookup(ray_direction, dir_t, sh);
441        return srgb_to_linear(color);
442    }
443
444    fn planar4d_isotropic_rotations(raw: array<f32, 8>) -> mat2x4<f32> {
445        let r1x = raw[0];
446        let r1y = raw[1];
447        let r1z = raw[2];
448        let r1w = raw[3];
449
450        let r2x = raw[4];
451        let r2y = raw[5];
452        let r2z = raw[6];
453        let r2w = raw[7];
454
455        return mat2x4<f32>(
456            r1x, r1y, r1z, r1w,
457            r2x, r2y, r2z, r2w,
458        );
459    }
460
461    fn planar4d_scale_from(scale_opacity: vec4<f32>) -> vec3<f32> {
462        return scale_opacity.xyz;
463    }
464
465    fn planar4d_opacity_from(scale_opacity: vec4<f32>) -> f32 {
466        return scale_opacity.w;
467    }
468
469    fn planar4d_timestamp_from(timestamp_timescale: vec4<f32>) -> f32 {
470        return timestamp_timescale.x;
471    }
472
473    fn planar4d_time_scale_from(timestamp_timescale: vec4<f32>) -> f32 {
474        return timestamp_timescale.y;
475    }
476
477    #ifdef PLANAR_F32
478        fn get_color(
479            index: u32,
480            dir_t: f32,
481            ray_direction: vec3<f32>,
482        ) -> vec3<f32> {
483            return planar4d_color_from_sh(ray_direction, dir_t, spherindrical_harmonics[index]);
484        }
485
486        fn get_isotropic_rotations(index: u32) -> mat2x4<f32> {
487            return planar4d_isotropic_rotations(isotropic_rotations[index]);
488        }
489
490        fn get_scale(index: u32) -> vec3<f32> {
491            return planar4d_scale_from(scale_opacity[index]);
492        }
493
494        fn get_opacity(index: u32) -> f32 {
495            return planar4d_opacity_from(scale_opacity[index]);
496        }
497
498        fn get_position(index: u32) -> vec3<f32> {
499            return planar_position(position_visibility[index]);
500        }
501
502        fn get_visibility(index: u32) -> f32 {
503            return planar_visibility(position_visibility[index]);
504        }
505
506        fn get_spherindrical_harmonics(index: u32) -> array<f32, #{SH_4D_COEFF_COUNT}> {
507            return spherindrical_harmonics[index];
508        }
509
510        fn get_timestamp(index: u32) -> f32 {
511            return planar4d_timestamp_from(timestamp_timescale[index]);
512        }
513
514        fn get_time_scale(index: u32) -> f32 {
515            return planar4d_time_scale_from(timestamp_timescale[index]);
516        }
517
518        #ifdef BINARY_GAUSSIAN_OP
519            fn get_rhs_color(
520                index: u32,
521                dir_t: f32,
522                ray_direction: vec3<f32>,
523            ) -> vec3<f32> {
524                return planar4d_color_from_sh(ray_direction, dir_t, rhs_spherindrical_harmonics[index]);
525            }
526
527            fn get_rhs_isotropic_rotations(index: u32) -> mat2x4<f32> {
528                return planar4d_isotropic_rotations(rhs_isotropic_rotations[index]);
529            }
530
531            fn get_rhs_scale(index: u32) -> vec3<f32> {
532                return planar4d_scale_from(rhs_scale_opacity[index]);
533            }
534
535            fn get_rhs_opacity(index: u32) -> f32 {
536                return planar4d_opacity_from(rhs_scale_opacity[index]);
537            }
538
539            fn get_rhs_position(index: u32) -> vec3<f32> {
540                return planar_position(rhs_position_visibility[index]);
541            }
542
543            fn get_rhs_visibility(index: u32) -> f32 {
544                return planar_visibility(rhs_position_visibility[index]);
545            }
546
547            fn get_rhs_spherindrical_harmonics(index: u32) -> array<f32, #{SH_4D_COEFF_COUNT}> {
548                return rhs_spherindrical_harmonics[index];
549            }
550
551            fn get_rhs_timestamp(index: u32) -> f32 {
552                return planar4d_timestamp_from(rhs_timestamp_timescale[index]);
553            }
554
555            fn get_rhs_time_scale(index: u32) -> f32 {
556                return planar4d_time_scale_from(rhs_timestamp_timescale[index]);
557            }
558        #endif
559    #endif
560
561    // TODO: PLANAR_F16 for GAUSSIAN_4D
562#endif