1#import bevy_pbr::cluster::{2 Aabb, CLUSTERABLE_OBJECT_TYPE_DECAL, CLUSTERABLE_OBJECT_TYPE_IRRADIANCE_VOLUME,3 CLUSTERABLE_OBJECT_TYPE_POINT_LIGHT, CLUSTERABLE_OBJECT_TYPE_REFLECTION_PROBE,4 CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT, ClusterableObjectZSlice,5 calculate_sphere_cluster_bounds, compute_view_from_world_scale6}7#import bevy_pbr::clustered_forward8#import bevy_pbr::light_probes::transpose_affine_matrix9#import bevy_pbr::mesh_view_types::{10 ClusterOffsetsAndCounts, ClusterableObjectIndexLists, ClusteredDecals, ClusteredLights,11 LightProbes, Lights, POINT_LIGHT_FLAGS_SPOT_LIGHT_Y_NEGATIVE12}13#import bevy_render::view::View15// The shader that performs the cluster-object intersection tests and assigns16// objects to clusters as appropriate.17//18// Assuming the froxel grid has size WxHxD, this shader is expected to run on a19// viewport of size WxH. It draws each *Z slice* as an axis aligned quad such20// that each fragment shader invocation represents a single cluster-object pair21// in the froxel grid. The fragment shader performs a finer test to see if the22// object might intersects the cluster and, if it succeeds, records the result.23// Because the result is written to storage buffers, color writes are disabled24// for this shader invocation.25//26// This shader runs twice: once to accumulate the *count* of each object type in27// each cluster, and once to *populate* the actual IDs. Both invocations of the28// shader must compute the exact same visibility results.30struct Vertex {31 @builtin(instance_index) instance_id: u32,32 @location(0) position: vec2<f32>,33}35// Data output from the vertex shader and input to the fragment shader.36struct Varyings {37 @builtin(position) position: vec4<f32>,38 // The index of the Z slice we're rasterizing.39 @location(0) @interpolate(flat) instance_id: u32,40 // The view-space center of the bounding sphere of the object.41 @location(1) @interpolate(flat) sphere_position: vec3<f32>,42 // The view-space radius of the bounding sphere of the object.43 @location(2) @interpolate(flat) sphere_radius: f32,44}46// The same as the `ClusterOffsetsAndCounts` structure, but with atomic fields47// so that we can write to it.48struct ClusterOffsetsAndCountsAtomic {49 data: array<ClusterOffsetsAndCountsElementAtomic>,50}52// The same as the `ClusterOffsetsAndCountsElement` structure, but with atomic53// fields so that we can write to it.54struct ClusterOffsetsAndCountsElementAtomic {55 offset: atomic<u32>,56 point_lights: atomic<u32>,57 spot_lights: atomic<u32>,58 reflection_probes: atomic<u32>,59 irradiance_volumes: atomic<u32>,60 decals: atomic<u32>,61 pad_a: u32,62 pad_b: u32,63}65// The list of clusterable object Z slices that we read from to.66@group(0) @binding(0) var<storage> z_slices: array<ClusterableObjectZSlice>;67#ifndef VERTEX_SHADER68// The list of indices per cluster that we write to in the populate pass.69@group(0) @binding(1) var<storage, read_write> index_lists: ClusterableObjectIndexLists;70#endif // VERTEX_SHADER71// Information about each light.72@group(0) @binding(2) var<storage> clustered_lights: ClusteredLights;73// Information about each light probe (reflection probe or irradiance volume).74@group(0) @binding(3) var<uniform> light_probes: LightProbes;75// Information about each clustered decal.76@group(0) @binding(4) var<storage> clustered_decals: ClusteredDecals;77// Information about the clusters as a whole, including the dimensions of the78// cluster grid.79@group(0) @binding(5) var<uniform> lights: Lights;80// Information about the view.81@group(0) @binding(6) var<uniform> view: View;82#ifndef VERTEX_SHADER83#ifdef POPULATE_PASS84// The number of objects in each cluster, and the offset of each list.85@group(0) @binding(7) var<storage> offsets_and_counts: ClusterOffsetsAndCounts;86// For each cluster, the counts of objects written *so far* to it.87//88// We use this during the populate phase in order to write the ID of each object89// to the correct spot.90@group(0) @binding(8) var<storage, read_write> scratchpad_offsets_and_counts:91 ClusterOffsetsAndCountsAtomic;92#else // POPULATE_PASS93// The number of objects in each cluster.94//95// During the count pass, we write to this.96@group(0) @binding(7) var<storage, read_write> offsets_and_counts: ClusterOffsetsAndCountsAtomic;97#endif // POPULATE_PASS98#endif // VERTEX_SHADER100// The vertex entry point.101@vertex102fn vertex_main(vertex: Vertex) -> Varyings {103 let instance_id = vertex.instance_id;104 let object_index = z_slices[instance_id].object_index;105 let object_type = z_slices[instance_id].object_type;107 // Look up the world space bounding sphere of the object.108 let bounding_sphere = get_object_bounding_sphere(object_index, object_type);109 let position = bounding_sphere.xyz;110 let radius = bounding_sphere.w;112 let view_from_world_scale = compute_view_from_world_scale(view.world_from_view);113 let max_view_from_world_scale = max(view_from_world_scale.x,114 max(view_from_world_scale.y, view_from_world_scale.z));115 let is_orthographic = view.clip_from_view[3].w == 1.0;117 // Calculate an approximate AABB of the cluster by computing its bounding118 // sphere and converting that to the AABB.119 // It's possible to do better, as the CPU version of120 // `assign_objects_to_clusters` does with its *iterative sphere refinement*121 // algorithm. However, that's sequential. I believe that this simple122 // approach is the standard way to do cluster assignment on GPU.123 let cluster_bounds = calculate_sphere_cluster_bounds(124 position,125 radius,126 view.view_from_world,127 view.clip_from_view,128 view_from_world_scale,129 lights.cluster_dimensions.xyz,130 lights.cluster_factors.zw,131 is_orthographic132 );133 let cluster_bounds_xy = vec4<u32>(cluster_bounds.min.xy, cluster_bounds.max.xy + vec2<u32>(1u));135 // Calculate the bounding sphere's center and radius in view space.136 let view_position = (view.view_from_world * vec4(position, 1.0)).xyz;137 let view_radius = max_view_from_world_scale * radius;139 return Varyings(140 calculate_vertex_position(vertex, cluster_bounds_xy),141 instance_id,142 view_position,143 view_radius144 );145}147// Returns the position of the quad vertex necessary to enclose all the148// fragments that represent the cluster AABB.149// The cluster bounds are supplied as `vec4(min X, min Y, max X, max Y)`.150fn calculate_vertex_position(vertex: Vertex, cluster_bounds: vec4<u32>) -> vec4<f32> {151 let framebuffer_position = vec2<u32>(152 select(cluster_bounds.x, cluster_bounds.z, vertex.position.x == 1.0),153 select(cluster_bounds.y, cluster_bounds.w, vertex.position.y == 1.0)154 );155 let vertex_position =156 vec2<f32>(framebuffer_position) / vec2<f32>(lights.cluster_dimensions.xy);157 return vec4(mix(vec2(-1.0, 1.0), vec2(1.0, -1.0), vertex_position), 0.0, 1.0);158}160#ifndef VERTEX_SHADER162// Performs a fine-grained test to ensure that the object intersects a single163// froxel and records the result.164@fragment165fn fragment_main(varyings: Varyings) -> @location(0) vec4<f32> {166 let instance_id = varyings.instance_id;167 let object_index = z_slices[instance_id].object_index;168 let object_type = z_slices[instance_id].object_type;169 let z_slice = z_slices[instance_id].z_slice;171 let is_orthographic = view.clip_from_view[3].w == 1.0;172 let screen_size = view.viewport.zw;173 let tile_size = screen_size / vec2<f32>(lights.cluster_dimensions.xy);175 let z_near_far = compute_z_near_and_z_far(is_orthographic);176 let z_near = z_near_far.x;177 let z_far = z_near_far.y;179 // Determine the AABB of the cluster this fragment represents.180 let cluster_position = vec3<u32>(vec2<u32>(floor(varyings.position.xy)), z_slice);181 let cluster_aabb = compute_aabb_for_cluster(182 z_near,183 z_far,184 tile_size,185 screen_size,186 view.view_from_clip,187 is_orthographic,188 lights.cluster_dimensions.xyz,189 cluster_position190 );191 let cluster_aabb_center = (cluster_aabb.max + cluster_aabb.min) * 0.5;192 let cluster_aabb_half_size = (cluster_aabb.max - cluster_aabb.min) * 0.5;194 // See if the object sphere intersects the AABB. If it doesn't, cull the195 // object.196 let object_intersects_cluster_aabb = sphere_intersects_aabb(197 varyings.sphere_position,198 varyings.sphere_radius,199 cluster_aabb_center,200 cluster_aabb_half_size201 );202 if (!object_intersects_cluster_aabb) {203 return vec4<f32>(0.0);204 }206 // Do further, more precise culling for spot lights.207 if (object_type == CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT && cull_spot_light(208 object_index,209 cluster_aabb_center,210 length(cluster_aabb_half_size),211 varyings.sphere_position,212 varyings.sphere_radius213 )) {214 return vec4<f32>(0.0);215 }217 let cluster_index =218 clustered_forward::fragment_cluster_index(cluster_position, lights.cluster_dimensions);220 // If this is the populate pass, reserve a slot and write in the actual221 // object index. Otherwise, if this is the count pass, just bump the222 // appropriate counter.223#ifdef POPULATE_PASS224 let output_index = allocate_list_entry(cluster_index, object_type);225 if (output_index < arrayLength(&index_lists.data)) {226 index_lists.data[output_index] = object_index;227 }228#else // POPULATE_PASS229 increment_object_count(cluster_index, object_type);230#endif // POPULATE_PASS232 return vec4<f32>(0.0);233}235#endif // VERTEX_SHADER237// Returns true if the given sphere intersects the AABB with the given238// boundaries.239fn sphere_intersects_aabb(240 sphere_center: vec3<f32>,241 sphere_radius: f32,242 aabb_center: vec3<f32>,243 aabb_half_size: vec3<f32>244) -> bool {245 let delta = max(vec3(0.0), abs(aabb_center - sphere_center) - aabb_half_size);246 let dist_sq = dot(delta, delta);247 return dist_sq <= sphere_radius * sphere_radius;248}250// See `bevy_light::cluster::assign::compute_aabb_for_cluster`.251fn compute_aabb_for_cluster(252 z_near: f32,253 z_far: f32,254 tile_size: vec2<f32>,255 screen_size: vec2<f32>,256 view_from_clip: mat4x4<f32>,257 is_orthographic: bool,258 cluster_dimensions: vec3<u32>,259 ijk_u: vec3<u32>260) -> Aabb {261 let ijk = vec3<f32>(ijk_u);263 // Calculate the minimum and maximum points in screen space264 let p_min = ijk.xy * tile_size;265 let p_max = p_min + tile_size;267 var cluster_min: vec3<f32>;268 var cluster_max: vec3<f32>;269 if (is_orthographic) {270 // Use linear depth slicing for orthographic272 // Convert to view space at the cluster near and far planes273 // NOTE: 1.0 is the near plane due to using reverse z projections274 var p_min = screen_to_view(screen_size, view_from_clip, p_min, 0.0).xyz;275 var p_max = screen_to_view(screen_size, view_from_clip, p_max, 0.0).xyz;277 // calculate cluster depth using z_near and z_far278 p_min.z = -z_near + (z_near - z_far) * ijk.z / f32(cluster_dimensions.z);279 p_max.z = -z_near + (z_near - z_far) * (ijk.z + 1.0) / f32(cluster_dimensions.z);281 cluster_min = min(p_min, p_max);282 cluster_max = max(p_min, p_max);283 } else {284 // Convert to view space at the near plane285 // NOTE: 1.0 is the near plane due to using reverse z projections286 let p_min = screen_to_view(screen_size, view_from_clip, p_min, 1.0);287 let p_max = screen_to_view(screen_size, view_from_clip, p_max, 1.0);289 let z_far_over_z_near = -z_far / -z_near;290 var cluster_near = 0.0;291 if (ijk.z != 0.0) {292 cluster_near = -z_near *293 pow(z_far_over_z_near, (ijk.z - 1.0) / f32(cluster_dimensions.z - 1u));294 }295 // NOTE: This could be simplified to:296 // cluster_far = cluster_near * z_far_over_z_near;297 var cluster_far: f32;298 if (cluster_dimensions.z == 1u) {299 cluster_far = -z_far;300 } else {301 cluster_far = -z_near * pow(z_far_over_z_near, ijk.z / f32(cluster_dimensions.z - 1u));302 }304 // Calculate the four intersection points of the min and max points with the cluster near and far planes305 let p_min_near = line_intersection_to_z_plane(vec3(0.0), p_min.xyz, cluster_near);306 let p_min_far = line_intersection_to_z_plane(vec3(0.0), p_min.xyz, cluster_far);307 let p_max_near = line_intersection_to_z_plane(vec3(0.0), p_max.xyz, cluster_near);308 let p_max_far = line_intersection_to_z_plane(vec3(0.0), p_max.xyz, cluster_far);310 cluster_min = min(min(p_min_near, p_min_far), min(p_max_near, p_max_far));311 cluster_max = max(max(p_min_near, p_min_far), max(p_max_near, p_max_far));312 }314 return Aabb(cluster_min, cluster_max);315}317// Converts a screen-space position to a view-space position.318// See `bevy_light::cluster::assign::screen_to_view`.319fn screen_to_view(320 screen_size: vec2<f32>,321 view_from_clip: mat4x4<f32>,322 screen: vec2<f32>,323 ndc_z: f32324) -> vec4<f32> {325 let tex_coord = screen / screen_size;326 let clip = vec4(327 tex_coord.x * 2.0 - 1.0,328 (1.0 - tex_coord.y) * 2.0 - 1.0,329 ndc_z,330 1.0331 );332 return clip_to_view(view_from_clip, clip);333}335// Converts a clip-space position to a view-space position.336// See `bevy_light::cluster::assign::clip_to_view`.337fn clip_to_view(view_from_clip: mat4x4<f32>, clip: vec4<f32>) -> vec4<f32> {338 let view = view_from_clip * clip;339 return view / view.w;340}342// Calculate the intersection of a ray from the eye through the view space343// position to a z plane344// See `bevy_light::cluster::assign::line_intersection_to_z_plane`.345fn line_intersection_to_z_plane(origin: vec3<f32>, p: vec3<f32>, z: f32) -> vec3<f32> {346 let v = p - origin;347 let t = (z - dot(vec3(0.0, 0.0, 1.0), origin)) / dot(vec3(0.0, 0.0, 1.0), v);348 return origin + t * v;349}351// Computes the near and far extents of the cluster grid.352fn compute_z_near_and_z_far(is_orthographic: bool) -> vec2<f32> {353 let cluster_dimensions = vec3<f32>(lights.cluster_dimensions.xyz);354 let cluster_factors = lights.cluster_factors;355 var z_near: f32;356 var z_far: f32;357 if (is_orthographic) {358 z_near = -cluster_factors.z;359 z_far = -(cluster_dimensions.z + cluster_factors.w * cluster_factors.z) / cluster_factors.w;360 } else {361 z_near = exp(cluster_factors.w / cluster_factors.z);362 z_far = exp((cluster_dimensions.z + cluster_factors.w - 1.0) / cluster_factors.z);363 }364 return vec2<f32>(z_near, z_far);365}367// Returns true if a spot light should be culled.368// See `assign_objects_to_clusters` in `bevy_light/src/cluster/assign.rs`.369fn cull_spot_light(370 object_index: u32,371 cluster_aabb_sphere_center: vec3<f32>,372 cluster_aabb_sphere_radius: f32,373 sphere_position: vec3<f32>,374 sphere_radius: f32375) -> bool {376 let light_custom_data = clustered_lights.data[object_index].light_custom_data;377 let light_flags = clustered_lights.data[object_index].flags;378 let light_tan_angle = clustered_lights.data[object_index].spot_light_tan_angle;380 // `spot_light_dir_sin_cos` in `assign_objects_to_clusters` uses381 // `normalize(view_from_world *382 // vec4(clusterable_object.transform.back(), 0.0).xyz)`.383 // What we have is the XZ value of384 // `clusterable_object.transform.forward()`. So we have to first385 // calculate the missing Y value, then flip it to go from forward to386 // back, then transform by the view-from-world matrix to get387 // `view_light_direction`.388 let world_light_direction_rev_xz = light_custom_data.xy;389 let world_light_direction_rev_y_sign = select(390 1.0,391 -1.0,392 (light_flags & POINT_LIGHT_FLAGS_SPOT_LIGHT_Y_NEGATIVE) != 0u393 );394 let world_light_direction_rev = vec3(395 world_light_direction_rev_xz.x,396 world_light_direction_rev_y_sign * sqrt(397 1.0 -398 world_light_direction_rev_xz.x * world_light_direction_rev_xz.x -399 world_light_direction_rev_xz.y * world_light_direction_rev_xz.y400 ),401 world_light_direction_rev_xz.y402 );403 let world_light_direction = -world_light_direction_rev;404 let view_light_direction = normalize((view.view_from_world *405 vec4(world_light_direction, 0.0)).xyz);407 let angle_cos = cos_atan(light_tan_angle);408 let angle_sin = sin_atan(light_tan_angle);410 // test -- based on https://bartwronski.com/2017/04/13/cull-that-cone/411 let spot_light_offset = sphere_position - cluster_aabb_sphere_center;412 let spot_light_dist_sq = dot(spot_light_offset, spot_light_offset);413 let v1_len = dot(spot_light_offset, view_light_direction);415 let distance_closest_point = angle_cos * sqrt(spot_light_dist_sq - v1_len * v1_len) -416 v1_len * angle_sin;417 let angle_cull = distance_closest_point > cluster_aabb_sphere_radius;419 let front_cull = v1_len > cluster_aabb_sphere_radius + sphere_radius;420 let back_cull = v1_len < -cluster_aabb_sphere_radius;422 return angle_cull || front_cull || back_cull;423}425// Computes `cos(atan(x))` cheaply.426// See https://en.wikipedia.org/wiki/List_of_trigonometric_identities427fn cos_atan(tan_theta: f32) -> f32 {428 return inverseSqrt(1.0 + tan_theta * tan_theta);429}431// Computes `sin(atan(x))` cheaply.432// See https://en.wikipedia.org/wiki/List_of_trigonometric_identities433fn sin_atan(tan_theta: f32) -> f32 {434 return tan_theta * inverseSqrt(1.0 + tan_theta * tan_theta);435}437#ifndef VERTEX_SHADER438#ifdef POPULATE_PASS439// Allocates space in the appropriate list and returns the global index that the440// object index should be written to.441fn allocate_list_entry(cluster_index: u32, object_type: u32) -> u32 {442 switch (object_type) {443 case CLUSTERABLE_OBJECT_TYPE_POINT_LIGHT: {444 return offsets_and_counts.data[cluster_index][0u].x +445 atomicAdd(&scratchpad_offsets_and_counts.data[cluster_index].point_lights, 1u);446 }447 case CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT: {448 return offsets_and_counts.data[cluster_index][0u].x +449 offsets_and_counts.data[cluster_index][0u].y +450 atomicAdd(&scratchpad_offsets_and_counts.data[cluster_index].spot_lights, 1u);451 }452 case CLUSTERABLE_OBJECT_TYPE_REFLECTION_PROBE: {453 return offsets_and_counts.data[cluster_index][0u].x +454 offsets_and_counts.data[cluster_index][0u].y +455 offsets_and_counts.data[cluster_index][0u].z +456 atomicAdd(&scratchpad_offsets_and_counts.data[cluster_index].reflection_probes, 1u);457 }458 case CLUSTERABLE_OBJECT_TYPE_IRRADIANCE_VOLUME: {459 return offsets_and_counts.data[cluster_index][0u].x +460 offsets_and_counts.data[cluster_index][0u].y +461 offsets_and_counts.data[cluster_index][0u].z +462 offsets_and_counts.data[cluster_index][0u].w +463 atomicAdd(464 &scratchpad_offsets_and_counts.data[cluster_index].irradiance_volumes,465 1u466 );467 }468 case CLUSTERABLE_OBJECT_TYPE_DECAL: {469 return offsets_and_counts.data[cluster_index][0u].x +470 offsets_and_counts.data[cluster_index][0u].y +471 offsets_and_counts.data[cluster_index][0u].z +472 offsets_and_counts.data[cluster_index][0u].w +473 offsets_and_counts.data[cluster_index][1u].x +474 atomicAdd(&scratchpad_offsets_and_counts.data[cluster_index].decals, 1u);475 }476 default: {}477 }478 return 0xffffffffu;479}480#else // POPULATE_PASS481// Increments the count of objects of the given type for the given cluster.482fn increment_object_count(cluster_index: u32, object_type: u32) {483 switch (object_type) {484 case CLUSTERABLE_OBJECT_TYPE_POINT_LIGHT: {485 atomicAdd(&offsets_and_counts.data[cluster_index].point_lights, 1u);486 }487 case CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT: {488 atomicAdd(&offsets_and_counts.data[cluster_index].spot_lights, 1u);489 }490 case CLUSTERABLE_OBJECT_TYPE_REFLECTION_PROBE: {491 atomicAdd(&offsets_and_counts.data[cluster_index].reflection_probes, 1u);492 }493 case CLUSTERABLE_OBJECT_TYPE_IRRADIANCE_VOLUME: {494 atomicAdd(&offsets_and_counts.data[cluster_index].irradiance_volumes, 1u);495 }496 case CLUSTERABLE_OBJECT_TYPE_DECAL: {497 atomicAdd(&offsets_and_counts.data[cluster_index].decals, 1u);498 }499 default: {}500 }501}502#endif // POPULATE_PASS503#endif // VERTEX_SHADER505// Looks up and returns the world-space center and radius of the bounding sphere506// for the object with the given index and type.507// Returns a 4-vector with the fields `vec4(center X, center Y, center Z, radius)`.508fn get_object_bounding_sphere(object_index: u32, object_type: u32) -> vec4<f32> {509 var position = vec3<f32>(0.0);510 var radius = 0.0;511 switch (object_type) {512 case CLUSTERABLE_OBJECT_TYPE_POINT_LIGHT: {513 position = clustered_lights.data[object_index].position_radius.xyz;514 radius = clustered_lights.data[object_index].range;515 }516 case CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT: {517 position = clustered_lights.data[object_index].position_radius.xyz;518 radius = clustered_lights.data[object_index].range;519 }520 case CLUSTERABLE_OBJECT_TYPE_REFLECTION_PROBE: {521 position = light_probes.reflection_probes[object_index].world_position;522 radius = light_probes.reflection_probes[object_index].bounding_sphere_radius;523 }524 case CLUSTERABLE_OBJECT_TYPE_IRRADIANCE_VOLUME: {525 position = light_probes.irradiance_volumes[object_index].world_position;526 radius = light_probes.irradiance_volumes[object_index].bounding_sphere_radius;527 }528 case CLUSTERABLE_OBJECT_TYPE_DECAL: {529 position = clustered_decals.decals[object_index].world_position;530 radius = clustered_decals.decals[object_index].bounding_sphere_radius;531 }532 default: {}533 }534 return vec4(position, radius);535}
Search across packages, versions, files, declarations, and comments.
Use multiple terms to narrow results (AND), and add @vertex,
@fragment, or @compute to filter by shader stage.