1#define_import_path bevy_gaussian_splatting::bindings
2
3#import bevy_pbr::prepass_bindings::PreviousViewUniforms
4#import bevy_render::globals::Globals
5#import bevy_render::view::View
6
7@group(0) @binding(0) var<uniform> view: View;
8@group(0) @binding(1) var<uniform> globals: Globals;
9@group(0) @binding(2) var<uniform> previous_view_uniforms: PreviousViewUniforms;
10
11@group(0) @binding(14) var<storage> visibility_ranges: array<vec4<f32>>;
12
13struct GaussianUniforms {
14    transform: mat4x4<f32>,
15    global_opacity: f32,
16    global_scale: f32,
17    count: u32,
18    count_root_ceil: u32,
19    time: f32,
20    time_start: f32,
21    time_stop: f32,
22    num_classes: u32,
23    min: vec4<f32>,
24    max: vec4<f32>,
25};
26@group(1) @binding(0) var<uniform> gaussian_uniforms: GaussianUniforms;
27
28#ifdef GAUSSIAN_3D_STRUCTURE
29    #ifdef PACKED_F32
30        struct Gaussian {
31            @location(0) rotation: vec4<f32>,
32            @location(1) position_visibility: vec4<f32>,
33            @location(2) scale_opacity: vec4<f32>,
34            sh: array<f32, #{SH_COEFF_COUNT}>,
35        };
36
37        #ifdef READ_WRITE_POINTS
38            @group(2) @binding(0) var<storage, read_write> points: array<Gaussian>;
39        #else
40            @group(2) @binding(0) var<storage, read> points: array<Gaussian>;
41        #endif
42
43        #ifdef BINARY_GAUSSIAN_OP
44            #ifdef READ_WRITE_POINTS
45                @group(3) @binding(0) var<storage, read_write> rhs_points: array<Gaussian>;
46            #else
47                @group(3) @binding(0) var<storage, read> rhs_points: array<Gaussian>;
48            #endif
49
50            @group(4) @binding(0) var<storage, read_write> out_points: array<Gaussian>;
51        #endif
52    #endif
53
54    #ifdef PLANAR_F32
55        #ifdef READ_WRITE_POINTS
56            @group(2) @binding(0) var<storage, read_write> position_visibility: array<vec4<f32>>;
57        #else
58            @group(2) @binding(0) var<storage, read> position_visibility: array<vec4<f32>>;
59        #endif
60
61        @group(2) @binding(1) var<storage, read> spherical_harmonics: array<array<f32, #{SH_COEFF_COUNT}>>;
62
63        #ifdef BINARY_GAUSSIAN_OP
64            #ifdef READ_WRITE_POINTS
65                @group(3) @binding(0) var<storage, read_write> rhs_position_visibility: array<vec4<f32>>;
66            #else
67                @group(3) @binding(0) var<storage, read> rhs_position_visibility: array<vec4<f32>>;
68            #endif
69
70            @group(3) @binding(1) var<storage, read> rhs_spherical_harmonics: array<array<f32, #{SH_COEFF_COUNT}>>;
71
72            #ifdef PRECOMPUTE_COVARIANCE_3D
73                @group(3) @binding(2) var<storage, read> rhs_covariance_3d_opacity: array<array<f32, 8>>;
74            #else
75                @group(3) @binding(2) var<storage, read> rhs_rotation: array<vec4<f32>>;
76                @group(3) @binding(3) var<storage, read> rhs_scale_opacity: array<vec4<f32>>;
77            #endif
78        #endif
79
80
81        #ifdef BINARY_GAUSSIAN_OP
82            @group(4) @binding(0) var<storage, read_write> out_position_visibility: array<vec4<f32>>;
83            @group(4) @binding(1) var<storage, read_write> out_spherical_harmonics: array<array<f32, #{SH_COEFF_COUNT}>>;
84
85            #ifdef PRECOMPUTE_COVARIANCE_3D
86                @group(4) @binding(2) var<storage, read_write> out_covariance_3d_opacity: array<array<f32, 8>>;
87            #else
88                @group(4) @binding(2) var<storage, read_write> out_rotation: array<vec4<f32>>;
89                @group(4) @binding(3) var<storage, read_write> out_scale_opacity: array<vec4<f32>>;
90            #endif
91        #endif
92
93        #ifdef PRECOMPUTE_COVARIANCE_3D
94            @group(2) @binding(2) var<storage, read> covariance_3d_opacity: array<array<f32, 8>>;
95        #else
96            @group(2) @binding(2) var<storage, read> rotation: array<vec4<f32>>;
97            @group(2) @binding(3) var<storage, read> scale_opacity: array<vec4<f32>>;
98        #endif
99    #endif
100
101    #ifdef PLANAR_F16
102        #ifdef READ_WRITE_POINTS
103            @group(2) @binding(0) var<storage, read_write> position_visibility: array<vec4<f32>>;
104        #else
105            @group(2) @binding(0) var<storage, read> position_visibility: array<vec4<f32>>;
106        #endif
107
108        @group(2) @binding(1) var<storage, read> spherical_harmonics: array<array<u32, #{HALF_SH_COEFF_COUNT}>>;
109
110        #ifdef BINARY_GAUSSIAN_OP
111            #ifdef READ_WRITE_POINTS
112                @group(3) @binding(0) var<storage, read_write> rhs_position_visibility: array<vec4<f32>>;
113            #else
114                @group(3) @binding(0) var<storage, read> rhs_position_visibility: array<vec4<f32>>;
115            #endif
116
117            @group(3) @binding(1) var<storage, read> rhs_spherical_harmonics: array<array<u32, #{HALF_SH_COEFF_COUNT}>>;
118
119            #ifdef PRECOMPUTE_COVARIANCE_3D
120                @group(3) @binding(2) var<storage, read> rhs_covariance_3d_opacity: array<vec4<u32>>;
121            #else
122                @group(3) @binding(2) var<storage, read> rhs_rotation_scale_opacity: array<vec4<u32>>;
123            #endif
124
125            @group(4) @binding(0) var<storage, read_write> out_position_visibility: array<vec4<f32>>;
126            @group(4) @binding(1) var<storage, read_write> out_spherical_harmonics: array<array<u32, #{HALF_SH_COEFF_COUNT}>>;
127
128            #ifdef PRECOMPUTE_COVARIANCE_3D
129                @group(4) @binding(2) var<storage, read_write> out_covariance_3d_opacity: array<vec4<u32>>;
130            #else
131                @group(4) @binding(2) var<storage, read_write> out_rotation_scale_opacity: array<vec4<u32>>;
132            #endif
133        #endif
134
135        #ifdef PRECOMPUTE_COVARIANCE_3D
136            @group(2) @binding(2) var<storage, read> covariance_3d_opacity: array<vec4<u32>>;
137        #else
138            @group(2) @binding(2) var<storage, read> rotation_scale_opacity: array<vec4<u32>>;
139        #endif
140    #endif
141
142    #ifdef PLANAR_TEXTURE_F16
143        @group(2) @binding(0) var position_visibility: texture_2d<f32>;
144
145        #if SH_VEC4_PLANES == 1
146            @group(2) @binding(1) var spherical_harmonics: texture_2d<u32>;
147        #else
148            @group(2) @binding(1) var spherical_harmonics: texture_2d_array<u32>;
149        #endif
150
151        #ifdef PRECOMPUTE_COVARIANCE_3D
152            @group(2) @binding(2) var covariance_3d_opacity: texture_2d<u32>;
153        #else
154            @group(2) @binding(2) var rotation_scale_opacity: texture_2d<u32>;
155        #endif
156
157        #ifdef BINARY_GAUSSIAN_OP
158            @group(3) @binding(0) var rhs_position_visibility: texture_2d<f32>;
159
160            #if SH_VEC4_PLANES == 1
161                @group(3) @binding(1) var rhs_spherical_harmonics: texture_2d<u32>;
162            #else
163                @group(3) @binding(1) var rhs_spherical_harmonics: texture_2d_array<u32>;
164            #endif
165
166            #ifdef PRECOMPUTE_COVARIANCE_3D
167                @group(3) @binding(2) var rhs_covariance_3d_opacity: texture_2d<u32>;
168            #else
169                @group(3) @binding(2) var rhs_rotation_scale_opacity: texture_2d<u32>;
170            #endif
171        #endif
172    #endif
173
174    #ifdef PLANAR_TEXTURE_F32
175        @group(2) @binding(0) var position_visibility: texture_2d<f32>;
176
177        #if SH_VEC4_PLANES == 1
178            @group(2) @binding(1) var spherical_harmonics: texture_2d<f32>;
179        #else
180            @group(2) @binding(1) var spherical_harmonics: texture_2d_array<f32>;
181        #endif
182
183        // TODO: support f32_cov3d_opacity texture
184
185        @group(2) @binding(2) var rotation_scale_opacity: texture_2d<f32>;
186
187        #ifdef BINARY_GAUSSIAN_OP
188            @group(3) @binding(0) var rhs_position_visibility: texture_2d<f32>;
189
190            #if SH_VEC4_PLANES == 1
191                @group(3) @binding(1) var rhs_spherical_harmonics: texture_2d<f32>;
192            #else
193                @group(3) @binding(1) var rhs_spherical_harmonics: texture_2d_array<f32>;
194            #endif
195
196            @group(3) @binding(2) var rhs_rotation_scale_opacity: texture_2d<f32>;
197        #endif
198    #endif
199#else ifdef GAUSSIAN_4D
200    #ifdef PLANAR_F32
201        #ifdef READ_WRITE_POINTS
202            @group(2) @binding(0) var<storage, read_write> position_visibility: array<vec4<f32>>;
203        #else
204            @group(2) @binding(0) var<storage, read> position_visibility: array<vec4<f32>>;
205        #endif
206
207        @group(2) @binding(1) var<storage, read> spherindrical_harmonics: array<array<f32, #{SH_4D_COEFF_COUNT}>>;
208
209        @group(2) @binding(2) var<storage, read> isotropic_rotations: array<array<f32, 8>>;
210        @group(2) @binding(3) var<storage, read> scale_opacity: array<vec4<f32>>;
211        @group(2) @binding(4) var<storage, read> timestamp_timescale: array<vec4<f32>>;
212
213        #ifdef BINARY_GAUSSIAN_OP
214            @group(3) @binding(0) var<storage, read> rhs_position_visibility: array<vec4<f32>>;
215            @group(3) @binding(1) var<storage, read> rhs_spherindrical_harmonics: array<array<f32, #{SH_4D_COEFF_COUNT}>>;
216            @group(3) @binding(2) var<storage, read> rhs_isotropic_rotations: array<array<f32, 8>>;
217            @group(3) @binding(3) var<storage, read> rhs_scale_opacity: array<vec4<f32>>;
218            @group(3) @binding(4) var<storage, read> rhs_timestamp_timescale: array<vec4<f32>>;
219        #endif
220    #endif
221#endif
222
223struct DrawIndirect {
224    vertex_count: u32,
225    instance_count: atomic<u32>,
226    base_vertex: u32,
227    base_instance: u32,
228}
229
230struct Entry {
231    key: u32,
232    value: u32,
233}