cluster_raster.wgsl

#import bevy_pbr::cluster::{
    Aabb, CLUSTERABLE_OBJECT_TYPE_DECAL, CLUSTERABLE_OBJECT_TYPE_IRRADIANCE_VOLUME,
    CLUSTERABLE_OBJECT_TYPE_POINT_LIGHT, CLUSTERABLE_OBJECT_TYPE_REFLECTION_PROBE,
    CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT, ClusterableObjectZSlice,
    calculate_sphere_cluster_bounds, compute_view_from_world_scale
}
#import bevy_pbr::clustered_forward
#import bevy_pbr::light_probes::transpose_affine_matrix
#import bevy_pbr::mesh_view_types::{
    ClusterOffsetsAndCounts, ClusterableObjectIndexLists, ClusteredDecals, ClusteredLights,
    LightProbes, Lights, POINT_LIGHT_FLAGS_SPOT_LIGHT_Y_NEGATIVE
}
#import bevy_render::view::View

// The shader that performs the cluster-object intersection tests and assigns
// objects to clusters as appropriate.
//
// Assuming the froxel grid has size WxHxD, this shader is expected to run on a
// viewport of size WxH. It draws each *Z slice* as an axis aligned quad such
// that each fragment shader invocation represents a single cluster-object pair
// in the froxel grid. The fragment shader performs a finer test to see if the
// object might intersects the cluster and, if it succeeds, records the result.
// Because the result is written to storage buffers, color writes are disabled
// for this shader invocation.
//
// This shader runs twice: once to accumulate the *count* of each object type in
// each cluster, and once to *populate* the actual IDs. Both invocations of the
// shader must compute the exact same visibility results.

struct Vertex {
    @builtin(instance_index) instance_id: u32,
    @location(0) position: vec2<f32>,
}

// Data output from the vertex shader and input to the fragment shader.
struct Varyings {
    @builtin(position) position: vec4<f32>,
    // The index of the Z slice we're rasterizing.
    @location(0) @interpolate(flat) instance_id: u32,
    // The view-space center of the bounding sphere of the object.
    @location(1) @interpolate(flat) sphere_position: vec3<f32>,
    // The view-space radius of the bounding sphere of the object.
    @location(2) @interpolate(flat) sphere_radius: f32,
}

// The same as the `ClusterOffsetsAndCounts` structure, but with atomic fields
// so that we can write to it.
struct ClusterOffsetsAndCountsAtomic {
    data: array<ClusterOffsetsAndCountsElementAtomic>,
}

// The same as the `ClusterOffsetsAndCountsElement` structure, but with atomic
// fields so that we can write to it.
struct ClusterOffsetsAndCountsElementAtomic {
    offset: atomic<u32>,
    point_lights: atomic<u32>,
    spot_lights: atomic<u32>,
    reflection_probes: atomic<u32>,
    irradiance_volumes: atomic<u32>,
    decals: atomic<u32>,
    pad_a: u32,
    pad_b: u32,
}

// The list of clusterable object Z slices that we read from to.
@group(0) @binding(0) var<storage> z_slices: array<ClusterableObjectZSlice>;
#ifndef VERTEX_SHADER
// The list of indices per cluster that we write to in the populate pass.
@group(0) @binding(1) var<storage, read_write> index_lists: ClusterableObjectIndexLists;
#endif  // VERTEX_SHADER
// Information about each light.
@group(0) @binding(2) var<storage> clustered_lights: ClusteredLights;
// Information about each light probe (reflection probe or irradiance volume).
@group(0) @binding(3) var<uniform> light_probes: LightProbes;
// Information about each clustered decal.
@group(0) @binding(4) var<storage> clustered_decals: ClusteredDecals;
// Information about the clusters as a whole, including the dimensions of the
// cluster grid.
@group(0) @binding(5) var<uniform> lights: Lights;
// Information about the view.
@group(0) @binding(6) var<uniform> view: View;
#ifndef VERTEX_SHADER
#ifdef POPULATE_PASS
// The number of objects in each cluster, and the offset of each list.
@group(0) @binding(7) var<storage> offsets_and_counts: ClusterOffsetsAndCounts;
// For each cluster, the counts of objects written *so far* to it.
//
// We use this during the populate phase in order to write the ID of each object
// to the correct spot.
@group(0) @binding(8) var<storage, read_write> scratchpad_offsets_and_counts:
    ClusterOffsetsAndCountsAtomic;
#else   // POPULATE_PASS
// The number of objects in each cluster.
//
// During the count pass, we write to this.
@group(0) @binding(7) var<storage, read_write> offsets_and_counts: ClusterOffsetsAndCountsAtomic;
#endif  // POPULATE_PASS
#endif  // VERTEX_SHADER

// The vertex entry point.
@vertex
fn vertex_main(vertex: Vertex) -> Varyings {
    let instance_id = vertex.instance_id;
    let object_index = z_slices[instance_id].object_index;
    let object_type = z_slices[instance_id].object_type;

    // Look up the world space bounding sphere of the object.
    let bounding_sphere = get_object_bounding_sphere(object_index, object_type);
    let position = bounding_sphere.xyz;
    let radius = bounding_sphere.w;

    let view_from_world_scale = compute_view_from_world_scale(view.world_from_view);
    let max_view_from_world_scale = max(view_from_world_scale.x,
        max(view_from_world_scale.y, view_from_world_scale.z));
    let is_orthographic = view.clip_from_view[3].w == 1.0;

    // Calculate an approximate AABB of the cluster by computing its bounding
    // sphere and converting that to the AABB.
    // It's possible to do better, as the CPU version of
    // `assign_objects_to_clusters` does with its *iterative sphere refinement*
    // algorithm. However, that's sequential. I believe that this simple
    // approach is the standard way to do cluster assignment on GPU.
    let cluster_bounds = calculate_sphere_cluster_bounds(
        position,
        radius,
        view.view_from_world,
        view.clip_from_view,
        view_from_world_scale,
        lights.cluster_dimensions.xyz,
        lights.cluster_factors.zw,
        is_orthographic
    );
    let cluster_bounds_xy = vec4<u32>(cluster_bounds.min.xy, cluster_bounds.max.xy + vec2<u32>(1u));

    // Calculate the bounding sphere's center and radius in view space.
    let view_position = (view.view_from_world * vec4(position, 1.0)).xyz;
    let view_radius = max_view_from_world_scale * radius;

    return Varyings(
        calculate_vertex_position(vertex, cluster_bounds_xy),
        instance_id,
        view_position,
        view_radius
    );
}

// Returns the position of the quad vertex necessary to enclose all the
// fragments that represent the cluster AABB.
// The cluster bounds are supplied as `vec4(min X, min Y, max X, max Y)`.
fn calculate_vertex_position(vertex: Vertex, cluster_bounds: vec4<u32>) -> vec4<f32> {
    let framebuffer_position = vec2<u32>(
        select(cluster_bounds.x, cluster_bounds.z, vertex.position.x == 1.0),
        select(cluster_bounds.y, cluster_bounds.w, vertex.position.y == 1.0)
    );
    let vertex_position =
        vec2<f32>(framebuffer_position) / vec2<f32>(lights.cluster_dimensions.xy);
    return vec4(mix(vec2(-1.0, 1.0), vec2(1.0, -1.0), vertex_position), 0.0, 1.0);
}

#ifndef VERTEX_SHADER

// Performs a fine-grained test to ensure that the object intersects a single
// froxel and records the result.
@fragment
fn fragment_main(varyings: Varyings) -> @location(0) vec4<f32> {
    let instance_id = varyings.instance_id;
    let object_index = z_slices[instance_id].object_index;
    let object_type = z_slices[instance_id].object_type;
    let z_slice = z_slices[instance_id].z_slice;

    let is_orthographic = view.clip_from_view[3].w == 1.0;
    let screen_size = view.viewport.zw;
    let tile_size = screen_size / vec2<f32>(lights.cluster_dimensions.xy);

    let z_near_far = compute_z_near_and_z_far(is_orthographic);
    let z_near = z_near_far.x;
    let z_far = z_near_far.y;

    // Determine the AABB of the cluster this fragment represents.
    let cluster_position = vec3<u32>(vec2<u32>(floor(varyings.position.xy)), z_slice);
    let cluster_aabb = compute_aabb_for_cluster(
        z_near,
        z_far,
        tile_size,
        screen_size,
        view.view_from_clip,
        is_orthographic,
        lights.cluster_dimensions.xyz,
        cluster_position
    );
    let cluster_aabb_center = (cluster_aabb.max + cluster_aabb.min) * 0.5;
    let cluster_aabb_half_size = (cluster_aabb.max - cluster_aabb.min) * 0.5;

    // See if the object sphere intersects the AABB. If it doesn't, cull the
    // object.
    let object_intersects_cluster_aabb = sphere_intersects_aabb(
        varyings.sphere_position,
        varyings.sphere_radius,
        cluster_aabb_center,
        cluster_aabb_half_size
    );
    if (!object_intersects_cluster_aabb) {
        return vec4<f32>(0.0);
    }

    // Do further, more precise culling for spot lights.
    if (object_type == CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT && cull_spot_light(
        object_index,
        cluster_aabb_center,
        length(cluster_aabb_half_size),
        varyings.sphere_position,
        varyings.sphere_radius
    )) {
        return vec4<f32>(0.0);
    }

    let cluster_index =
        clustered_forward::fragment_cluster_index(cluster_position, lights.cluster_dimensions);

    // If this is the populate pass, reserve a slot and write in the actual
    // object index. Otherwise, if this is the count pass, just bump the
    // appropriate counter.
#ifdef POPULATE_PASS
    let output_index = allocate_list_entry(cluster_index, object_type);
    if (output_index < arrayLength(&index_lists.data)) {
        index_lists.data[output_index] = object_index;
    }
#else   // POPULATE_PASS
    increment_object_count(cluster_index, object_type);
#endif  // POPULATE_PASS

    return vec4<f32>(0.0);
}

#endif  // VERTEX_SHADER

// Returns true if the given sphere intersects the AABB with the given
// boundaries.
fn sphere_intersects_aabb(
    sphere_center: vec3<f32>,
    sphere_radius: f32,
    aabb_center: vec3<f32>,
    aabb_half_size: vec3<f32>
) -> bool {
    let delta = max(vec3(0.0), abs(aabb_center - sphere_center) - aabb_half_size);
    let dist_sq = dot(delta, delta);
    return dist_sq <= sphere_radius * sphere_radius;
}

// See `bevy_light::cluster::assign::compute_aabb_for_cluster`.
fn compute_aabb_for_cluster(
    z_near: f32,
    z_far: f32,
    tile_size: vec2<f32>,
    screen_size: vec2<f32>,
    view_from_clip: mat4x4<f32>,
    is_orthographic: bool,
    cluster_dimensions: vec3<u32>,
    ijk_u: vec3<u32>
) -> Aabb {
    let ijk = vec3<f32>(ijk_u);

    // Calculate the minimum and maximum points in screen space
    let p_min = ijk.xy * tile_size;
    let p_max = p_min + tile_size;

    var cluster_min: vec3<f32>;
    var cluster_max: vec3<f32>;
    if (is_orthographic) {
        // Use linear depth slicing for orthographic

        // Convert to view space at the cluster near and far planes
        // NOTE: 1.0 is the near plane due to using reverse z projections
        var p_min = screen_to_view(screen_size, view_from_clip, p_min, 0.0).xyz;
        var p_max = screen_to_view(screen_size, view_from_clip, p_max, 0.0).xyz;

        // calculate cluster depth using z_near and z_far
        p_min.z = -z_near + (z_near - z_far) * ijk.z / f32(cluster_dimensions.z);
        p_max.z = -z_near + (z_near - z_far) * (ijk.z + 1.0) / f32(cluster_dimensions.z);

        cluster_min = min(p_min, p_max);
        cluster_max = max(p_min, p_max);
    } else {
        // Convert to view space at the near plane
        // NOTE: 1.0 is the near plane due to using reverse z projections
        let p_min = screen_to_view(screen_size, view_from_clip, p_min, 1.0);
        let p_max = screen_to_view(screen_size, view_from_clip, p_max, 1.0);

        let z_far_over_z_near = -z_far / -z_near;
        var cluster_near = 0.0;
        if (ijk.z != 0.0) {
            cluster_near = -z_near *
                pow(z_far_over_z_near, (ijk.z - 1.0) / f32(cluster_dimensions.z - 1u));
        }
        // NOTE: This could be simplified to:
        // cluster_far = cluster_near * z_far_over_z_near;
        var cluster_far: f32;
        if (cluster_dimensions.z == 1u) {
            cluster_far = -z_far;
        } else {
            cluster_far = -z_near * pow(z_far_over_z_near, ijk.z / f32(cluster_dimensions.z - 1u));
        }

        // Calculate the four intersection points of the min and max points with the cluster near and far planes
        let p_min_near = line_intersection_to_z_plane(vec3(0.0), p_min.xyz, cluster_near);
        let p_min_far = line_intersection_to_z_plane(vec3(0.0), p_min.xyz, cluster_far);
        let p_max_near = line_intersection_to_z_plane(vec3(0.0), p_max.xyz, cluster_near);
        let p_max_far = line_intersection_to_z_plane(vec3(0.0), p_max.xyz, cluster_far);

        cluster_min = min(min(p_min_near, p_min_far), min(p_max_near, p_max_far));
        cluster_max = max(max(p_min_near, p_min_far), max(p_max_near, p_max_far));
    }

    return Aabb(cluster_min, cluster_max);
}

// Converts a screen-space position to a view-space position.
// See `bevy_light::cluster::assign::screen_to_view`.
fn screen_to_view(
    screen_size: vec2<f32>,
    view_from_clip: mat4x4<f32>,
    screen: vec2<f32>,
    ndc_z: f32
) -> vec4<f32> {
    let tex_coord = screen / screen_size;
    let clip = vec4(
        tex_coord.x * 2.0 - 1.0,
        (1.0 - tex_coord.y) * 2.0 - 1.0,
        ndc_z,
        1.0
    );
    return clip_to_view(view_from_clip, clip);
}

// Converts a clip-space position to a view-space position.
// See `bevy_light::cluster::assign::clip_to_view`.
fn clip_to_view(view_from_clip: mat4x4<f32>, clip: vec4<f32>) -> vec4<f32> {
    let view = view_from_clip * clip;
    return view / view.w;
}

// Calculate the intersection of a ray from the eye through the view space
// position to a z plane
// See `bevy_light::cluster::assign::line_intersection_to_z_plane`.
fn line_intersection_to_z_plane(origin: vec3<f32>, p: vec3<f32>, z: f32) -> vec3<f32> {
    let v = p - origin;
    let t = (z - dot(vec3(0.0, 0.0, 1.0), origin)) / dot(vec3(0.0, 0.0, 1.0), v);
    return origin + t * v;
}

// Computes the near and far extents of the cluster grid.
fn compute_z_near_and_z_far(is_orthographic: bool) -> vec2<f32> {
    let cluster_dimensions = vec3<f32>(lights.cluster_dimensions.xyz);
    let cluster_factors = lights.cluster_factors;
    var z_near: f32;
    var z_far: f32;
    if (is_orthographic) {
        z_near = -cluster_factors.z;
        z_far = -(cluster_dimensions.z + cluster_factors.w * cluster_factors.z) / cluster_factors.w;
    } else {
        z_near = exp(cluster_factors.w / cluster_factors.z);
        z_far = exp((cluster_dimensions.z + cluster_factors.w - 1.0) / cluster_factors.z);
    }
    return vec2<f32>(z_near, z_far);
}

// Returns true if a spot light should be culled.
// See `assign_objects_to_clusters` in `bevy_light/src/cluster/assign.rs`.
fn cull_spot_light(
    object_index: u32,
    cluster_aabb_sphere_center: vec3<f32>,
    cluster_aabb_sphere_radius: f32,
    sphere_position: vec3<f32>,
    sphere_radius: f32
) -> bool {
    let light_custom_data = clustered_lights.data[object_index].light_custom_data;
    let light_flags = clustered_lights.data[object_index].flags;
    let light_tan_angle = clustered_lights.data[object_index].spot_light_tan_angle;

    // `spot_light_dir_sin_cos` in `assign_objects_to_clusters` uses
    // `normalize(view_from_world *
    // vec4(clusterable_object.transform.back(), 0.0).xyz)`.
    // What we have is the XZ value of
    // `clusterable_object.transform.forward()`. So we have to first
    // calculate the missing Y value, then flip it to go from forward to
    // back, then transform by the view-from-world matrix to get
    // `view_light_direction`.
    let world_light_direction_rev_xz = light_custom_data.xy;
    let world_light_direction_rev_y_sign = select(
        1.0,
        -1.0,
        (light_flags & POINT_LIGHT_FLAGS_SPOT_LIGHT_Y_NEGATIVE) != 0u
    );
    let world_light_direction_rev = vec3(
        world_light_direction_rev_xz.x,
        world_light_direction_rev_y_sign * sqrt(
            1.0 -
                world_light_direction_rev_xz.x * world_light_direction_rev_xz.x -
                world_light_direction_rev_xz.y * world_light_direction_rev_xz.y
        ),
        world_light_direction_rev_xz.y
    );
    let world_light_direction = -world_light_direction_rev;
    let view_light_direction = normalize((view.view_from_world *
        vec4(world_light_direction, 0.0)).xyz);

    let angle_cos = cos_atan(light_tan_angle);
    let angle_sin = sin_atan(light_tan_angle);

    // test -- based on https://bartwronski.com/2017/04/13/cull-that-cone/
    let spot_light_offset = sphere_position - cluster_aabb_sphere_center;
    let spot_light_dist_sq = dot(spot_light_offset, spot_light_offset);
    let v1_len = dot(spot_light_offset, view_light_direction);

    let distance_closest_point = angle_cos * sqrt(spot_light_dist_sq - v1_len * v1_len) -
        v1_len * angle_sin;
    let angle_cull = distance_closest_point > cluster_aabb_sphere_radius;

    let front_cull = v1_len > cluster_aabb_sphere_radius + sphere_radius;
    let back_cull = v1_len < -cluster_aabb_sphere_radius;

    return angle_cull || front_cull || back_cull;
}

// Computes `cos(atan(x))` cheaply.
// See https://en.wikipedia.org/wiki/List_of_trigonometric_identities
fn cos_atan(tan_theta: f32) -> f32 {
    return inverseSqrt(1.0 + tan_theta * tan_theta);
}

// Computes `sin(atan(x))` cheaply.
// See https://en.wikipedia.org/wiki/List_of_trigonometric_identities
fn sin_atan(tan_theta: f32) -> f32 {
    return tan_theta * inverseSqrt(1.0 + tan_theta * tan_theta);
}

#ifndef VERTEX_SHADER
#ifdef POPULATE_PASS
// Allocates space in the appropriate list and returns the global index that the
// object index should be written to.
fn allocate_list_entry(cluster_index: u32, object_type: u32) -> u32 {
    switch (object_type) {
        case CLUSTERABLE_OBJECT_TYPE_POINT_LIGHT: {
            return offsets_and_counts.data[cluster_index][0u].x +
                atomicAdd(&scratchpad_offsets_and_counts.data[cluster_index].point_lights, 1u);
        }
        case CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT: {
            return offsets_and_counts.data[cluster_index][0u].x +
                offsets_and_counts.data[cluster_index][0u].y +
                atomicAdd(&scratchpad_offsets_and_counts.data[cluster_index].spot_lights, 1u);
        }
        case CLUSTERABLE_OBJECT_TYPE_REFLECTION_PROBE: {
            return offsets_and_counts.data[cluster_index][0u].x +
                offsets_and_counts.data[cluster_index][0u].y +
                offsets_and_counts.data[cluster_index][0u].z +
                atomicAdd(&scratchpad_offsets_and_counts.data[cluster_index].reflection_probes, 1u);
        }
        case CLUSTERABLE_OBJECT_TYPE_IRRADIANCE_VOLUME: {
            return offsets_and_counts.data[cluster_index][0u].x +
                offsets_and_counts.data[cluster_index][0u].y +
                offsets_and_counts.data[cluster_index][0u].z +
                offsets_and_counts.data[cluster_index][0u].w +
                atomicAdd(
                    &scratchpad_offsets_and_counts.data[cluster_index].irradiance_volumes,
                    1u
                );
        }
        case CLUSTERABLE_OBJECT_TYPE_DECAL: {
            return offsets_and_counts.data[cluster_index][0u].x +
                offsets_and_counts.data[cluster_index][0u].y +
                offsets_and_counts.data[cluster_index][0u].z +
                offsets_and_counts.data[cluster_index][0u].w +
                offsets_and_counts.data[cluster_index][1u].x +
                atomicAdd(&scratchpad_offsets_and_counts.data[cluster_index].decals, 1u);
        }
        default: {}
    }
    return 0xffffffffu;
}
#else   // POPULATE_PASS
// Increments the count of objects of the given type for the given cluster.
fn increment_object_count(cluster_index: u32, object_type: u32) {
    switch (object_type) {
        case CLUSTERABLE_OBJECT_TYPE_POINT_LIGHT: {
            atomicAdd(&offsets_and_counts.data[cluster_index].point_lights, 1u);
        }
        case CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT: {
            atomicAdd(&offsets_and_counts.data[cluster_index].spot_lights, 1u);
        }
        case CLUSTERABLE_OBJECT_TYPE_REFLECTION_PROBE: {
            atomicAdd(&offsets_and_counts.data[cluster_index].reflection_probes, 1u);
        }
        case CLUSTERABLE_OBJECT_TYPE_IRRADIANCE_VOLUME: {
            atomicAdd(&offsets_and_counts.data[cluster_index].irradiance_volumes, 1u);
        }
        case CLUSTERABLE_OBJECT_TYPE_DECAL: {
            atomicAdd(&offsets_and_counts.data[cluster_index].decals, 1u);
        }
        default: {}
    }
}
#endif  // POPULATE_PASS
#endif  // VERTEX_SHADER

// Looks up and returns the world-space center and radius of the bounding sphere
// for the object with the given index and type.
// Returns a 4-vector with the fields `vec4(center X, center Y, center Z, radius)`.
fn get_object_bounding_sphere(object_index: u32, object_type: u32) -> vec4<f32> {
    var position = vec3<f32>(0.0);
    var radius = 0.0;
    switch (object_type) {
        case CLUSTERABLE_OBJECT_TYPE_POINT_LIGHT: {
            position = clustered_lights.data[object_index].position_radius.xyz;
            radius = clustered_lights.data[object_index].range;
        }
        case CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT: {
            position = clustered_lights.data[object_index].position_radius.xyz;
            radius = clustered_lights.data[object_index].range;
        }
        case CLUSTERABLE_OBJECT_TYPE_REFLECTION_PROBE: {
            position = light_probes.reflection_probes[object_index].world_position;
            radius = light_probes.reflection_probes[object_index].bounding_sphere_radius;
        }
        case CLUSTERABLE_OBJECT_TYPE_IRRADIANCE_VOLUME: {
            position = light_probes.irradiance_volumes[object_index].world_position;
            radius = light_probes.irradiance_volumes[object_index].bounding_sphere_radius;
        }
        case CLUSTERABLE_OBJECT_TYPE_DECAL: {
            position = clustered_decals.decals[object_index].world_position;
            radius = clustered_decals.decals[object_index].bounding_sphere_radius;
        }
        default: {}
    }
    return vec4(position, radius);
}