diff --git a/crates/bevy_core_pipeline/src/skybox/mod.rs b/crates/bevy_core_pipeline/src/skybox/mod.rs index 2d31cea138732..5a500ef590e60 100644 --- a/crates/bevy_core_pipeline/src/skybox/mod.rs +++ b/crates/bevy_core_pipeline/src/skybox/mod.rs @@ -64,8 +64,9 @@ impl Plugin for SkyboxPlugin { /// Adds a skybox to a 3D camera, based on a cubemap texture. /// -/// Note that this component does not (currently) affect the scene's lighting. -/// To do so, use `EnvironmentMapLight` alongside this component. +/// Note that this component does not affect the scene's lighting. +/// To do so, use `EnvironmentMapLight` or `GenerateEnvironmentMapLight` +/// alongside this component. /// /// See also . #[derive(Component, ExtractComponent, Clone)] diff --git a/crates/bevy_pbr/src/environment_map/diffuse_convolution.wgsl b/crates/bevy_pbr/src/environment_map/diffuse_convolution.wgsl new file mode 100644 index 0000000000000..ccf85043f7410 --- /dev/null +++ b/crates/bevy_pbr/src/environment_map/diffuse_convolution.wgsl @@ -0,0 +1,68 @@ +// Importance samples (Lambertian distribution) a skybox to produce a diffuse lighting cubemap +// Based on https://github.com/KhronosGroup/glTF-IBL-Sampler/blob/master/lib/source/shaders/filter.frag + +#import bevy_pbr::utils::PI + +@group(0) @binding(0) var skybox: texture_cube; +#ifdef RG11B10FLOAT +@group(0) @binding(1) var diffuse_map: texture_storage_2d_array; +#else +@group(0) @binding(1) var diffuse_map: texture_storage_2d_array; +#endif +@group(0) @binding(2) var bilinear: sampler; + +fn get_dir(u: f32, v: f32, face: u32) -> vec3 { + switch face { + case 0u: { return vec3(1.0, v, -u); } + case 1u: { return vec3(-1.0, v, u); } + case 2u: { return vec3(u, 1.0, -v); } + case 3u: { return vec3(u, -1.0, v); } + case 4u: { return vec3(u, v, 1.0); } + default { return vec3(-u, v, -1.0); } + } +} + +fn generate_tbn(normal: vec3) -> mat3x3 { + var bitangent = vec3(0.0, 1.0, 0.0); + + let n_dot_up = dot(normal, bitangent); + if abs(n_dot_up) >= 0.9999999 { + bitangent = vec3(0.0, 0.0, sign(n_dot_up)); + } + + let tangent = normalize(cross(bitangent, normal)); + bitangent = cross(normal, tangent); + + return mat3x3(tangent, bitangent, normal); +} + +@compute +@workgroup_size(8, 8, 1) +fn main(@builtin(global_invocation_id) id: vec3) { + let u = (f32(id.x) * 2.0 + 1.0) / 64.0 - 1.0; + let v = -(f32(id.y) * 2.0 + 1.0) / 64.0 + 1.0; + + let normal = normalize(get_dir(u, v, id.z)); + + var color = vec3(0.0); + for (var sample_i = 0u; sample_i < 32u; sample_i++) { + // R2 sequence - http://extremelearning.com.au/unreasonable-effectiveness-of-quasirandom-sequences + let r = fract(0.5 + f32(sample_i) * vec2(0.75487766624669276005, 0.5698402909980532659114)); + + let cos_theta = sqrt(1.0 - f32(r.y)); + let sin_theta = sqrt(r.y); + let phi = 2.0 * PI * r.x; + + let local_space_direction = normalize(vec3( + sin_theta * cos(phi), + sin_theta * sin(phi), + cos_theta, + )); + let direction = generate_tbn(normal) * local_space_direction; + + color += textureSampleLevel(skybox, bilinear, direction, 0.0).rgb; + } + color /= 32.0; + + textureStore(diffuse_map, id.xy, id.z, vec4(color, 1.0)); +} diff --git a/crates/bevy_pbr/src/environment_map/downsample.wgsl b/crates/bevy_pbr/src/environment_map/downsample.wgsl new file mode 100644 index 0000000000000..3bf40dc1725fe --- /dev/null +++ b/crates/bevy_pbr/src/environment_map/downsample.wgsl @@ -0,0 +1,73 @@ +// Step 1/2 in generating a specular lighting cubemap from a skybox: Downsamples a skybox into multiple mips +// Original source: https://www.activision.com/cdn/research/downsample_cubemap.txt + +// Copyright 2016 Activision Publishing, Inc. +// +// Permission is hereby granted, free of charge, to any person obtaining +// a copy of this software and associated documentation files (the "Software"), +// to deal in the Software without restriction, including without limitation +// the rights to use, copy, modify, merge, publish, distribute, sublicense, +// and/or sell copies of the Software, and to permit persons to whom the Software +// is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +@group(0) @binding(0) var tex_hi_res: texture_cube; +#ifdef RG11B10FLOAT +@group(0) @binding(1) var tex_los_res: texture_storage_2d_array; +#else +@group(0) @binding(1) var tex_los_res: texture_storage_2d_array; +#endif +@group(0) @binding(2) var bilinear: sampler; + +fn get_dir(u: f32, v: f32, face: u32) -> vec3 { + switch face { + case 0u: { return vec3(1.0, v, -u); } + case 1u: { return vec3(-1.0, v, u); } + case 2u: { return vec3(u, 1.0, -v); } + case 3u: { return vec3(u, -1.0, v); } + case 4u: { return vec3(u, v, 1.0); } + default { return vec3(-u, v, -1.0); } + } +} + +fn calc_weight(u: f32, v: f32) -> f32 { + let val = u * u + v * v + 1.0; + return val * sqrt(val); +} + +@compute +@workgroup_size(8, 8, 1) +fn main(@builtin(global_invocation_id) id: vec3) { + let res_lo = textureDimensions(tex_los_res).x; + + if all(vec2u(id.xy) < vec2u(res_lo)) { + let inv_res_lo = 1.0 / f32(res_lo); + + let u0 = (f32(id.x) * 2.0 + 1.0 - 0.75) * inv_res_lo - 1.0; + let u1 = (f32(id.x) * 2.0 + 1.0 + 0.75) * inv_res_lo - 1.0; + + let v0 = (f32(id.y) * 2.0 + 1.0 - 0.75) * -inv_res_lo + 1.0; + let v1 = (f32(id.y) * 2.0 + 1.0 + 0.75) * -inv_res_lo + 1.0; + + var weights = vec4(calc_weight(u0, v0), calc_weight(u1, v0), calc_weight(u0, v1), calc_weight(u1, v1)); + let wsum = 0.5 / dot(vec4(1.0), weights); + weights = weights * wsum + 0.125; + + var color = textureSampleLevel(tex_hi_res, bilinear, get_dir(u0, v0, id.z), 0.0) * weights.x; + color += textureSampleLevel(tex_hi_res, bilinear, get_dir(u1, v0, id.z), 0.0) * weights.y; + color += textureSampleLevel(tex_hi_res, bilinear, get_dir(u0, v1, id.z), 0.0) * weights.z; + color += textureSampleLevel(tex_hi_res, bilinear, get_dir(u1, v1, id.z), 0.0) * weights.w; + + textureStore(tex_los_res, id.xy, id.z, color); + } +} diff --git a/crates/bevy_pbr/src/environment_map/filter.wgsl b/crates/bevy_pbr/src/environment_map/filter.wgsl new file mode 100644 index 0000000000000..38effba0c4845 --- /dev/null +++ b/crates/bevy_pbr/src/environment_map/filter.wgsl @@ -0,0 +1,183 @@ +// Step 2/2 in generating a specular lighting cubemap from a skybox: Importance sample the GGX distribution based on the downsampled cubemap +// Original source: https://www.activision.com/cdn/research/filter_using_table_128.txt + +// Copyright 2016 Activision Publishing, Inc. +// +// Permission is hereby granted, free of charge, to any person obtaining +// a copy of this software and associated documentation files (the "Software"), +// to deal in the Software without restriction, including without limitation +// the rights to use, copy, modify, merge, publish, distribute, sublicense, +// and/or sell copies of the Software, and to permit persons to whom the Software +// is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +@group(0) @binding(0) var tex_in: texture_cube; +#ifdef RG11B10FLOAT +@group(0) @binding(1) var text_out0: texture_storage_2d_array; +@group(0) @binding(2) var text_out1: texture_storage_2d_array; +@group(0) @binding(3) var text_out2: texture_storage_2d_array; +@group(0) @binding(4) var text_out3: texture_storage_2d_array; +@group(0) @binding(5) var text_out4: texture_storage_2d_array; +@group(0) @binding(6) var text_out5: texture_storage_2d_array; +@group(0) @binding(7) var text_out6: texture_storage_2d_array; +#else +@group(0) @binding(1) var text_out0: texture_storage_2d_array; +@group(0) @binding(2) var text_out1: texture_storage_2d_array; +@group(0) @binding(3) var text_out2: texture_storage_2d_array; +@group(0) @binding(4) var text_out3: texture_storage_2d_array; +@group(0) @binding(5) var text_out4: texture_storage_2d_array; +@group(0) @binding(6) var text_out5: texture_storage_2d_array; +@group(0) @binding(7) var text_out6: texture_storage_2d_array; +#endif +@group(0) @binding(8) var trilinear: sampler; +@group(0) @binding(9) var coeffs: array, 24>, 3>, 5>, 7>; + +fn get_dir(u: f32, v: f32, face: u32) -> vec3 { + switch face { + case 0u: { return vec3(1.0, v, -u); } + case 1u: { return vec3(-1.0, v, u); } + case 2u: { return vec3(u, 1.0, -v); } + case 3u: { return vec3(u, -1.0, v); } + case 4u: { return vec3(u, v, 1.0); } + default { return vec3(-u, v, -1.0); } + } +} + +@compute +@workgroup_size(64, 1, 1) +fn main(@builtin(global_invocation_id) global_id: vec3) { + var id = global_id; + var level = 0u; + if id.x < 128u * 128u { + level = 0u; + } else if id.x < 128u * 128u + 64u * 64u { + level = 1u; + id.x -= 128u * 128u; + } else if id.x < 128u * 128u + 64u * 64u + 32u * 32u { + level = 2u; + id.x -= 128u * 128u + 64u * 64u; + } else if id.x < 128u * 128u + 64u * 64u + 32u * 32u + 16u * 16u { + level = 3u; + id.x -= 128u * 128u + 64u * 64u + 32u * 32u; + } else if id.x < 128u * 128u + 64u * 64u + 32u * 32u + 16u * 16u + 8u * 8u { + level = 4u; + id.x -= 128u * 128u + 64u * 64u + 32u * 32u + 16u * 16u; + } else if id.x < 128u * 128u + 64u * 64u + 32u * 32u + 16u * 16u + 8u * 8u + 4u * 4u { + level = 5u; + id.x -= 128u * 128u + 64u * 64u + 32u * 32u + 16u * 16u + 8u * 8u; + } else if id.x < 128u * 128u + 64u * 64u + 32u * 32u + 16u * 16u + 8u * 8u + 4u * 4u + 2u * 2u { + level = 6u; + id.x -= 128u * 128u + 64u * 64u + 32u * 32u + 16u * 16u + 8u * 8u + 4u * 4u; + } else { + return; + } + + id.z = id.y; + let res = 128u >> level; + id.y = id.x / res; + id.x -= id.y * res; + + let u = (f32(id.x) * 2.0 + 1.0) / f32(res) - 1.0; + let v = -(f32(id.y) * 2.0 + 1.0) / f32(res) + 1.0; + + let dir = get_dir(u, v, id.z); + let frame_z = normalize(dir); + let adir = abs(dir); + + var color = vec4(0.0); + for (var axis = 0u; axis < 3u; axis++) { + let other_axis0 = 1u - (axis & 1u) - (axis >> 1u); + let other_axis1 = 2u - (axis >> 1u); + + let frame_weight = (max(adir[other_axis0], adir[other_axis1]) - 0.75) / 0.25; + if frame_weight > 0.0 { + var up_vector = vec3(0.0); + switch axis { + case 0u: { up_vector = vec3(1.0, 0.0, 0.0); } + case 1u: { up_vector = vec3(0.0, 1.0, 0.0); } + default { up_vector = vec3(0.0, 0.0, 1.0); } + } + let frame_x = normalize(cross(up_vector, frame_z)); + let frame_y = cross(frame_z, frame_x); + + var nx = dir[other_axis0]; + var ny = dir[other_axis1]; + let nz = adir[axis]; + + let nmax_xy = max(abs(ny), abs(nx)); + nx /= nmax_xy; + ny /= nmax_xy; + + var theta = 0.0; + if ny < nx { + if ny <= -0.999 { theta = nx; } else { theta = ny; } + } else { + if ny >= 0.999 { theta = -nx; } else { theta = -ny; } + } + + var phi = 0.0; + if nz <= -0.999 { + phi = -nmax_xy; + } else if nz >= 0.999 { + phi = nmax_xy; + } else { + phi = nz; + } + let theta2 = theta * theta; + let phi2 = phi * phi; + + for (var i_super_tap = 0u; i_super_tap < 8u; i_super_tap++) { + let index = 8u * axis + i_super_tap; + var coeffs_dir0 = array(vec4(0.0), vec4(0.0), vec4(0.0)); + var coeffs_dir1 = array(vec4(0.0), vec4(0.0), vec4(0.0)); + var coeffs_dir2 = array(vec4(0.0), vec4(0.0), vec4(0.0)); + var coeffs_level = array(vec4(0.0), vec4(0.0), vec4(0.0)); + var coeffs_weight = array(vec4(0.0), vec4(0.0), vec4(0.0)); + + for (var i_coeff = 0u; i_coeff < 3u; i_coeff++) { + coeffs_dir0[i_coeff] = coeffs[level][0u][i_coeff][index]; + coeffs_dir1[i_coeff] = coeffs[level][1u][i_coeff][index]; + coeffs_dir2[i_coeff] = coeffs[level][2u][i_coeff][index]; + coeffs_level[i_coeff] = coeffs[level][3u][i_coeff][index]; + coeffs_weight[i_coeff] = coeffs[level][4u][i_coeff][index]; + } + + for (var i_sub_tap = 0u; i_sub_tap < 4u; i_sub_tap++) { + var sample_dir = frame_x * (coeffs_dir0[0u][i_sub_tap] + coeffs_dir0[1u][i_sub_tap] * theta2 + coeffs_dir0[2u][i_sub_tap] * phi2) + frame_y * (coeffs_dir1[0u][i_sub_tap] + coeffs_dir1[1u][i_sub_tap] * theta2 + coeffs_dir1[2u][i_sub_tap] * phi2) + frame_z * (coeffs_dir2[0u][i_sub_tap] + coeffs_dir2[1u][i_sub_tap] * theta2 + coeffs_dir2[2u][i_sub_tap] * phi2); + + var sample_level = coeffs_level[0u][i_sub_tap] + coeffs_level[1u][i_sub_tap] * theta2 + coeffs_level[2u][i_sub_tap] * phi2; + + var sample_weight = coeffs_weight[0u][i_sub_tap] + coeffs_weight[1u][i_sub_tap] * theta2 + coeffs_weight[2u][i_sub_tap] * phi2; + sample_weight *= frame_weight; + + sample_dir /= max(abs(sample_dir[0u]), max(abs(sample_dir[1u]), abs(sample_dir[2u]))); + sample_level += 0.75 * log2(dot(sample_dir, sample_dir)); + + color += vec4(textureSampleLevel(tex_in, trilinear, sample_dir, sample_level).rgb * sample_weight, sample_weight); + } + } + } + } + color /= color.a; + color = vec4(max(color.rgb, vec3(0.0)), 1.0); + + switch level { + case 0u: { textureStore(text_out0, id.xy, id.z, color); } + case 1u: { textureStore(text_out1, id.xy, id.z, color); } + case 2u: { textureStore(text_out2, id.xy, id.z, color); } + case 3u: { textureStore(text_out3, id.xy, id.z, color); } + case 4u: { textureStore(text_out4, id.xy, id.z, color); } + case 5u: { textureStore(text_out5, id.xy, id.z, color); } + default { textureStore(text_out6, id.xy, id.z, color); } + } +} diff --git a/crates/bevy_pbr/src/environment_map/filter_coefficients.bin b/crates/bevy_pbr/src/environment_map/filter_coefficients.bin new file mode 100644 index 0000000000000..724a61589702e Binary files /dev/null and b/crates/bevy_pbr/src/environment_map/filter_coefficients.bin differ diff --git a/crates/bevy_pbr/src/environment_map/generate_from_skybox.rs b/crates/bevy_pbr/src/environment_map/generate_from_skybox.rs new file mode 100644 index 0000000000000..7f50f77f1d265 --- /dev/null +++ b/crates/bevy_pbr/src/environment_map/generate_from_skybox.rs @@ -0,0 +1,565 @@ +// https://research.activision.com/publications/archives/fast-filtering-of-reflection-probes + +use super::{ + EnvironmentMapLight, DIFFUSE_CONVOLUTION_SHADER_HANDLE, DOWNSAMPLE_SHADER_HANDLE, + FILTER_SHADER_HANDLE, +}; +use bevy_asset::{Assets, Handle}; +use bevy_core_pipeline::Skybox; +use bevy_ecs::{ + prelude::{Component, Entity}, + query::{QueryItem, Without}, + system::{Commands, Query, Res, ResMut, Resource}, + world::{FromWorld, World}, +}; +use bevy_reflect::Reflect; +use bevy_render::{ + extract_component::ExtractComponent, + render_asset::RenderAssets, + render_graph::{NodeRunError, RenderGraphContext, ViewNode}, + render_resource::{binding_types::*, *}, + renderer::{RenderContext, RenderDevice}, + texture::{GpuImage, Image, ImageFilterMode, ImageSampler, ImageSamplerDescriptor, Volume}, +}; +use bevy_utils::default; +use std::num::NonZeroU64; + +// WARNING: Do not adjust these constants without adjusting the shader code +const DIFFUSE_MAP_SIZE: u32 = 64; +const SPECULAR_MAP_SIZE: u32 = 128; +const SPECULAR_MAP_MIP_COUNT: u32 = 7; + +const FILTER_COEFFICIENTS_SIZE: Option = NonZeroU64::new(40320); + +/// Automatically generate an [`EnvironmentMapLight`] from a [`Skybox`]. +/// +/// Usage: +/// * Add this component via `GenerateEnvironmentMapLight::default()` to an entity with a [`Skybox`] component. +/// * The first frame this component is added to the skybox entity, an [`EnvironmentMapLight`] +/// component will be generated and added to the skybox entity. +/// * For static (non-changing) skyboxes, remove this component 1 frame after adding it to the skybox entity to save performance. +/// +/// This component does not work on `WebGL2`, and must use [`GenerateEnvironmentMapLightTextureFormat::Rgba16Float`] on `WebGPU`. +#[derive(Component, ExtractComponent, Reflect, Default, Clone)] +pub struct GenerateEnvironmentMapLight { + texture_format: GenerateEnvironmentMapLightTextureFormat, + downsampled_cubemap: Option>, +} + +impl GenerateEnvironmentMapLight { + pub fn new_with_texture_format( + texture_format: GenerateEnvironmentMapLightTextureFormat, + ) -> Self { + Self { + texture_format, + downsampled_cubemap: None, + } + } +} + +#[derive(Reflect, Clone, Copy, Default)] +pub enum GenerateEnvironmentMapLightTextureFormat { + /// 4 bytes per pixel (smaller and faster), but may not be able to represent as wide a range of lighting values. + /// This is the [`Default`] on non-WASM platforms. + #[cfg_attr(not(target_arch = "wasm32"), default)] + Rg11b10Float, + /// 8 bytes per pixel. This is the [`Default`], and only supported option for `WebGPU`. + #[cfg_attr(target_arch = "wasm32", default)] + Rgba16Float, +} + +impl GenerateEnvironmentMapLightTextureFormat { + pub fn as_wgpu(&self) -> TextureFormat { + match self { + Self::Rg11b10Float => TextureFormat::Rg11b10Float, + Self::Rgba16Float => TextureFormat::Rgba16Float, + } + } +} + +#[derive(Resource)] +pub struct GenerateEnvironmentMapLightResources { + rg11b10float: GenerateEnvironmentMapLightResourcesSpecialized, + rgba16float: GenerateEnvironmentMapLightResourcesSpecialized, + filter_coefficients: Buffer, +} + +impl FromWorld for GenerateEnvironmentMapLightResources { + fn from_world(world: &mut World) -> Self { + let render_device = world.resource::(); + let pipeline_cache = world.resource::(); + + Self { + rg11b10float: GenerateEnvironmentMapLightResourcesSpecialized::new( + TextureFormat::Rg11b10Float, + render_device, + pipeline_cache, + ), + rgba16float: GenerateEnvironmentMapLightResourcesSpecialized::new( + TextureFormat::Rgba16Float, + render_device, + pipeline_cache, + ), + // Original source: https://www.activision.com/cdn/research/coeffs_quad_32.txt + filter_coefficients: render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("generate_environment_map_light_filter_coefficients"), + contents: include_bytes!("filter_coefficients.bin"), + usage: BufferUsages::UNIFORM, + }), + } + } +} + +struct GenerateEnvironmentMapLightResourcesSpecialized { + downsample_and_diffuse_convolution_layout: BindGroupLayout, + filter_layout: BindGroupLayout, + downsample_pipeline: CachedComputePipelineId, + filter_pipeline: CachedComputePipelineId, + diffuse_convolution_pipeline: CachedComputePipelineId, +} + +impl GenerateEnvironmentMapLightResourcesSpecialized { + fn new( + texture_format: TextureFormat, + render_device: &RenderDevice, + pipeline_cache: &PipelineCache, + ) -> Self { + let downsample_and_diffuse_convolution_layout = render_device.create_bind_group_layout( + "generate_environment_map_light_downsample_and_diffuse_convolution_bind_group_layout", + &BindGroupLayoutEntries::sequential( + ShaderStages::COMPUTE, + ( + texture_cube(TextureSampleType::Float { filterable: true }), + texture_storage_2d_array(texture_format, StorageTextureAccess::WriteOnly), + sampler(SamplerBindingType::Filtering), + ), + ), + ); + + let filter_layout = render_device.create_bind_group_layout( + "generate_environment_map_light_filter_bind_group_layout", + &BindGroupLayoutEntries::sequential( + ShaderStages::COMPUTE, + ( + texture_cube(TextureSampleType::Float { filterable: true }), + texture_storage_2d_array(texture_format, StorageTextureAccess::WriteOnly), + texture_storage_2d_array(texture_format, StorageTextureAccess::WriteOnly), + texture_storage_2d_array(texture_format, StorageTextureAccess::WriteOnly), + texture_storage_2d_array(texture_format, StorageTextureAccess::WriteOnly), + texture_storage_2d_array(texture_format, StorageTextureAccess::WriteOnly), + texture_storage_2d_array(texture_format, StorageTextureAccess::WriteOnly), + texture_storage_2d_array(texture_format, StorageTextureAccess::WriteOnly), + sampler(SamplerBindingType::Filtering), + uniform_buffer_sized(false, FILTER_COEFFICIENTS_SIZE), + ), + ), + ); + + let shader_defs = match texture_format { + TextureFormat::Rg11b10Float => vec!["RG11B10FLOAT".into()], + TextureFormat::Rgba16Float => vec![], + _ => unreachable!(), + }; + + let downsample_pipeline = + pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { + label: Some("generate_environment_map_light_downsample_pipeline".into()), + layout: vec![downsample_and_diffuse_convolution_layout.clone()], + push_constant_ranges: vec![], + shader: DOWNSAMPLE_SHADER_HANDLE, + shader_defs: shader_defs.clone(), + entry_point: "main".into(), + }); + + let filter_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { + label: Some("generate_environment_map_light_filter_pipeline".into()), + layout: vec![filter_layout.clone()], + push_constant_ranges: vec![], + shader: FILTER_SHADER_HANDLE, + shader_defs: shader_defs.clone(), + entry_point: "main".into(), + }); + + let diffuse_convolution_pipeline = + pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { + label: Some("generate_environment_map_light_diffuse_convolution_pipeline".into()), + layout: vec![downsample_and_diffuse_convolution_layout.clone()], + push_constant_ranges: vec![], + shader: DIFFUSE_CONVOLUTION_SHADER_HANDLE, + shader_defs, + entry_point: "main".into(), + }); + + Self { + downsample_and_diffuse_convolution_layout, + filter_layout, + downsample_pipeline, + filter_pipeline, + diffuse_convolution_pipeline, + } + } +} + +pub fn generate_dummy_environment_map_lights_for_skyboxes( + mut skyboxes: Query< + (Entity, &Skybox, &mut GenerateEnvironmentMapLight), + Without, + >, + mut commands: Commands, + mut images: ResMut>, +) { + for (entity, skybox, mut gen_env_map_light) in &mut skyboxes { + let skybox_size = match images.get(&skybox.0) { + Some(skybox) => skybox.texture_descriptor.size, + None => continue, + }; + + let texture_format = gen_env_map_light.texture_format.as_wgpu(); + + let diffuse_size = Extent3d { + width: DIFFUSE_MAP_SIZE, + height: DIFFUSE_MAP_SIZE, + depth_or_array_layers: 6, + }; + let diffuse_map = Image { + data: vec![0; texture_byte_count(diffuse_size, 1, texture_format)], + texture_descriptor: TextureDescriptor { + label: Some("generate_environment_map_light_diffuse_map_texture"), + size: diffuse_size, + mip_level_count: 1, + sample_count: 1, + dimension: TextureDimension::D2, + format: texture_format, + usage: TextureUsages::TEXTURE_BINDING | TextureUsages::STORAGE_BINDING, + view_formats: &[], + }, + sampler: ImageSampler::Descriptor(ImageSamplerDescriptor { + label: Some("generate_environment_map_light_downsample_sampler".to_owned()), + mag_filter: ImageFilterMode::Linear, + min_filter: ImageFilterMode::Linear, + mipmap_filter: ImageFilterMode::Nearest, + ..default() + }), + texture_view_descriptor: Some(TextureViewDescriptor { + label: Some("generate_environment_map_light_diffuse_map_texture_view"), + format: Some(texture_format), + dimension: Some(TextureViewDimension::Cube), + aspect: TextureAspect::All, + base_mip_level: 0, + mip_level_count: Some(1), + base_array_layer: 0, + array_layer_count: Some(6), + }), + }; + + let specular_size = Extent3d { + width: SPECULAR_MAP_SIZE, + height: SPECULAR_MAP_SIZE, + depth_or_array_layers: 6, + }; + let specular_map = Image { + data: vec![ + 0; + texture_byte_count(specular_size, SPECULAR_MAP_MIP_COUNT, texture_format) + ], + texture_descriptor: TextureDescriptor { + label: Some("generate_environment_map_light_specular_map_texture"), + size: specular_size, + mip_level_count: SPECULAR_MAP_MIP_COUNT, + sample_count: 1, + dimension: TextureDimension::D2, + format: texture_format, + usage: TextureUsages::TEXTURE_BINDING | TextureUsages::STORAGE_BINDING, + view_formats: &[], + }, + sampler: ImageSampler::Descriptor(ImageSamplerDescriptor { + label: Some("generate_environment_map_light_filter_sampler".to_owned()), + mag_filter: ImageFilterMode::Linear, + min_filter: ImageFilterMode::Linear, + mipmap_filter: ImageFilterMode::Linear, + ..default() + }), + texture_view_descriptor: Some(TextureViewDescriptor { + label: Some("generate_environment_map_light_specular_map_texture_view"), + format: Some(texture_format), + dimension: Some(TextureViewDimension::Cube), + aspect: TextureAspect::All, + base_mip_level: 0, + mip_level_count: Some(SPECULAR_MAP_MIP_COUNT), + base_array_layer: 0, + array_layer_count: Some(6), + }), + }; + + let downsampled_size = Extent3d { + width: skybox_size.width / 2, + height: skybox_size.height / 2, + depth_or_array_layers: 6, + }; + let downsampled_cubemap = Image { + data: vec![0; texture_byte_count(downsampled_size, 6, texture_format)], + texture_descriptor: TextureDescriptor { + label: Some("generate_environment_map_light_downsampled_cubemap"), + size: downsampled_size, + mip_level_count: 6, + sample_count: 1, + dimension: TextureDimension::D2, + format: texture_format, + usage: TextureUsages::TEXTURE_BINDING | TextureUsages::STORAGE_BINDING, + view_formats: &[], + }, + ..default() + }; + gen_env_map_light.downsampled_cubemap = Some(images.add(downsampled_cubemap)); + + commands.entity(entity).insert(EnvironmentMapLight { + diffuse_map: images.add(diffuse_map), + specular_map: images.add(specular_map), + }); + } +} + +#[derive(Component)] +pub struct GenerateEnvironmentMapLightBindGroups { + downsample: [BindGroup; 6], + filter: BindGroup, + diffuse_convolution: BindGroup, + downsampled_cubemap_size: u32, + texture_format: GenerateEnvironmentMapLightTextureFormat, +} + +// PERF: Cache bind groups +pub fn prepare_generate_environment_map_lights_for_skyboxes_bind_groups( + environment_map_lights: Query<( + Entity, + &Skybox, + &EnvironmentMapLight, + &GenerateEnvironmentMapLight, + )>, + resources: Res, + render_device: Res, + mut commands: Commands, + images: Res>, +) { + for (entity, skybox, environment_map_light, gen_env_map) in &environment_map_lights { + let (Some(skybox), Some(diffuse_map), Some(specular_map), Some(downsampled_cubemap)) = ( + images.get(&skybox.0), + images.get(&environment_map_light.diffuse_map), + images.get(&environment_map_light.specular_map), + gen_env_map + .downsampled_cubemap + .as_ref() + .and_then(|t| images.get(t)), + ) else { + continue; + }; + + let filter_coefficients = &resources.filter_coefficients; + let resources = match gen_env_map.texture_format { + GenerateEnvironmentMapLightTextureFormat::Rg11b10Float => &resources.rg11b10float, + GenerateEnvironmentMapLightTextureFormat::Rgba16Float => &resources.rgba16float, + }; + + let downsample1 = render_device.create_bind_group( + "generate_environment_map_light_downsample1_bind_group", + &resources.downsample_and_diffuse_convolution_layout, + &BindGroupEntries::sequential(( + &skybox.texture_view, + &d2array_view(0, downsampled_cubemap), + &diffuse_map.sampler, + )), + ); + let downsample2 = render_device.create_bind_group( + "generate_environment_map_light_downsample2_bind_group", + &resources.downsample_and_diffuse_convolution_layout, + &BindGroupEntries::sequential(( + &cube_view(0, downsampled_cubemap), + &d2array_view(1, downsampled_cubemap), + &diffuse_map.sampler, + )), + ); + let downsample3 = render_device.create_bind_group( + "generate_environment_map_light_downsample3_bind_group", + &resources.downsample_and_diffuse_convolution_layout, + &BindGroupEntries::sequential(( + &cube_view(1, downsampled_cubemap), + &d2array_view(2, downsampled_cubemap), + &diffuse_map.sampler, + )), + ); + let downsample4 = render_device.create_bind_group( + "generate_environment_map_light_downsample4_bind_group", + &resources.downsample_and_diffuse_convolution_layout, + &BindGroupEntries::sequential(( + &cube_view(2, downsampled_cubemap), + &d2array_view(3, downsampled_cubemap), + &diffuse_map.sampler, + )), + ); + let downsample5 = render_device.create_bind_group( + "generate_environment_map_light_downsample5_bind_group", + &resources.downsample_and_diffuse_convolution_layout, + &BindGroupEntries::sequential(( + &cube_view(3, downsampled_cubemap), + &d2array_view(4, downsampled_cubemap), + &diffuse_map.sampler, + )), + ); + let downsample6 = render_device.create_bind_group( + "generate_environment_map_light_downsample6_bind_group", + &resources.downsample_and_diffuse_convolution_layout, + &BindGroupEntries::sequential(( + &cube_view(4, downsampled_cubemap), + &d2array_view(5, downsampled_cubemap), + &diffuse_map.sampler, + )), + ); + + let filter = render_device.create_bind_group( + "generate_environment_map_light_filter_bind_group", + &resources.filter_layout, + &BindGroupEntries::sequential(( + &cube_view(0, downsampled_cubemap), + &d2array_view(0, specular_map), + &d2array_view(1, specular_map), + &d2array_view(2, specular_map), + &d2array_view(3, specular_map), + &d2array_view(4, specular_map), + &d2array_view(5, specular_map), + &d2array_view(6, specular_map), + &specular_map.sampler, + filter_coefficients.as_entire_binding(), + )), + ); + + let diffuse_convolution = render_device.create_bind_group( + "generate_environment_map_light_diffuse_convolution_bind_group", + &resources.downsample_and_diffuse_convolution_layout, + &BindGroupEntries::sequential(( + &skybox.texture_view, + &d2array_view(0, diffuse_map), + &diffuse_map.sampler, + )), + ); + + commands + .entity(entity) + .insert(GenerateEnvironmentMapLightBindGroups { + downsample: [ + downsample1, + downsample2, + downsample3, + downsample4, + downsample5, + downsample6, + ], + filter, + diffuse_convolution, + downsampled_cubemap_size: downsampled_cubemap.size.x as u32, + texture_format: gen_env_map.texture_format, + }); + } +} + +#[derive(Default)] +pub struct GenerateEnvironmentMapLightNode; + +impl ViewNode for GenerateEnvironmentMapLightNode { + type ViewData = &'static GenerateEnvironmentMapLightBindGroups; + + fn run( + &self, + _: &mut RenderGraphContext, + render_context: &mut RenderContext, + bind_groups: QueryItem, + world: &World, + ) -> Result<(), NodeRunError> { + let pipeline_cache = world.resource::(); + let resources = world.resource::(); + let resources = match bind_groups.texture_format { + GenerateEnvironmentMapLightTextureFormat::Rg11b10Float => &resources.rg11b10float, + GenerateEnvironmentMapLightTextureFormat::Rgba16Float => &resources.rgba16float, + }; + + let (Some(downsample_pipeline), Some(filter_pipeline), Some(diffuse_convolution_pipeline)) = ( + pipeline_cache.get_compute_pipeline(resources.downsample_pipeline), + pipeline_cache.get_compute_pipeline(resources.filter_pipeline), + pipeline_cache.get_compute_pipeline(resources.diffuse_convolution_pipeline), + ) else { + return Ok(()); + }; + + let command_encoder = render_context.command_encoder(); + let mut pass = command_encoder.begin_compute_pass(&ComputePassDescriptor { + label: Some("generate_environment_map_light_pass"), + timestamp_writes: None, + }); + + pass.set_pipeline(downsample_pipeline); + let mut texture_size = bind_groups.downsampled_cubemap_size; + for bind_group in &bind_groups.downsample { + let workgroup_count = div_ceil(texture_size, 8); + pass.set_bind_group(0, bind_group, &[]); + pass.dispatch_workgroups(workgroup_count, workgroup_count, 6); + texture_size /= 2; + } + + // PERF: Don't filter to generate the first mip level, just downsample and copy the skybox texture directly + pass.set_pipeline(filter_pipeline); + pass.set_bind_group(0, &bind_groups.filter, &[]); + // 342 * 64 (workgroup size) = 21888 + // 128 * 128 + 64 * 64 + 32 * 32 + 16 * 16 + 8 * 8 + 4 * 4 = 21840 + // 128 is SPECULAR_MAP_SIZE, and then halved SPECULAR_MAP_MIP_COUNT times + pass.dispatch_workgroups(342, 6, 1); + + // PERF: At this point, may want to copy the specular map to a compressed texture format + + pass.set_pipeline(diffuse_convolution_pipeline); + pass.set_bind_group(0, &bind_groups.diffuse_convolution, &[]); + pass.dispatch_workgroups(DIFFUSE_MAP_SIZE / 8, DIFFUSE_MAP_SIZE / 8, 6); + + Ok(()) + } +} + +fn texture_byte_count(mut size: Extent3d, mip_count: u32, texture_format: TextureFormat) -> usize { + let mut total_size = 0; + for _ in 0..mip_count { + total_size += size.volume(); + size.width /= 2; + size.height /= 2; + } + let block_size = texture_format.block_size(None).unwrap() as usize; + total_size * block_size +} + +fn cube_view(mip_level: u32, cubemap: &GpuImage) -> TextureView { + cubemap.texture.create_view(&TextureViewDescriptor { + label: Some("generate_environment_map_light_texture_view"), + format: Some(cubemap.texture_format), + dimension: Some(TextureViewDimension::Cube), + aspect: TextureAspect::All, + base_mip_level: mip_level, + mip_level_count: Some(1), + base_array_layer: 0, + array_layer_count: Some(6), + }) +} + +fn d2array_view(mip_level: u32, cubemap: &GpuImage) -> TextureView { + cubemap.texture.create_view(&TextureViewDescriptor { + label: Some("generate_environment_map_light_texture_view"), + format: Some(cubemap.texture_format), + dimension: Some(TextureViewDimension::D2Array), + aspect: TextureAspect::All, + base_mip_level: mip_level, + mip_level_count: Some(1), + base_array_layer: 0, + array_layer_count: Some(6), + }) +} + +/// Divide `numerator` by `denominator`, rounded up to the nearest multiple of `denominator`. +fn div_ceil(numerator: u32, denominator: u32) -> u32 { + (numerator + denominator - 1) / denominator +} diff --git a/crates/bevy_pbr/src/environment_map/mod.rs b/crates/bevy_pbr/src/environment_map/mod.rs index 823f264a17923..fb19b4a0db2f0 100644 --- a/crates/bevy_pbr/src/environment_map/mod.rs +++ b/crates/bevy_pbr/src/environment_map/mod.rs @@ -1,24 +1,53 @@ -use bevy_app::{App, Plugin}; +mod generate_from_skybox; + +use bevy_app::{App, Last, Plugin}; use bevy_asset::{load_internal_asset, Handle}; -use bevy_core_pipeline::prelude::Camera3d; -use bevy_ecs::{prelude::Component, query::With}; +use bevy_core_pipeline::{ + core_3d::{self, CORE_3D}, + prelude::Camera3d, +}; +use bevy_ecs::{ + prelude::Component, + query::{QueryItem, With}, + schedule::IntoSystemConfigs, +}; use bevy_reflect::Reflect; use bevy_render::{ extract_component::{ExtractComponent, ExtractComponentPlugin}, render_asset::RenderAssets, + render_graph::{RenderGraphApp, ViewNodeRunner}, render_resource::{ binding_types::{sampler, texture_cube}, *, }, texture::{FallbackImageCubemap, Image}, + Render, RenderApp, RenderSet, +}; +use generate_from_skybox::{ + generate_dummy_environment_map_lights_for_skyboxes, + prepare_generate_environment_map_lights_for_skyboxes_bind_groups, + GenerateEnvironmentMapLightNode, GenerateEnvironmentMapLightResources, }; +pub use generate_from_skybox::{ + GenerateEnvironmentMapLight, GenerateEnvironmentMapLightTextureFormat, +}; + +pub mod draw_3d_graph { + pub mod node { + /// Label for the generate environment map light render node. + pub const GENERATE_ENVIRONMENT_MAP_LIGHT: &str = "generate_environment_map_light"; + } +} pub const ENVIRONMENT_MAP_SHADER_HANDLE: Handle = Handle::weak_from_u128(154476556247605696); +pub const DOWNSAMPLE_SHADER_HANDLE: Handle = Handle::weak_from_u128(254476556247605696); +pub const FILTER_SHADER_HANDLE: Handle = Handle::weak_from_u128(354476556247605696); +pub const DIFFUSE_CONVOLUTION_SHADER_HANDLE: Handle = + Handle::weak_from_u128(454476556247605696); +pub struct EnvironmentMapLightPlugin; -pub struct EnvironmentMapPlugin; - -impl Plugin for EnvironmentMapPlugin { +impl Plugin for EnvironmentMapLightPlugin { fn build(&self, app: &mut App) { load_internal_asset!( app, @@ -26,9 +55,52 @@ impl Plugin for EnvironmentMapPlugin { "environment_map.wgsl", Shader::from_wgsl ); + load_internal_asset!( + app, + DOWNSAMPLE_SHADER_HANDLE, + "downsample.wgsl", + Shader::from_wgsl + ); + load_internal_asset!(app, FILTER_SHADER_HANDLE, "filter.wgsl", Shader::from_wgsl); + load_internal_asset!( + app, + DIFFUSE_CONVOLUTION_SHADER_HANDLE, + "diffuse_convolution.wgsl", + Shader::from_wgsl + ); app.register_type::() - .add_plugins(ExtractComponentPlugin::::default()); + .register_type::() + .register_type::() + .add_plugins(ExtractComponentPlugin::::default()) + .add_plugins(ExtractComponentPlugin::::default()); + } + + fn finish(&self, app: &mut App) { + if app.get_sub_app(RenderApp).is_err() { + return; + } + + app.add_systems(Last, generate_dummy_environment_map_lights_for_skyboxes); + + app.sub_app_mut(RenderApp) + .add_render_graph_node::>( + CORE_3D, + draw_3d_graph::node::GENERATE_ENVIRONMENT_MAP_LIGHT, + ) + .add_render_graph_edges( + CORE_3D, + &[ + draw_3d_graph::node::GENERATE_ENVIRONMENT_MAP_LIGHT, + core_3d::graph::node::START_MAIN_PASS, + ], + ) + .init_resource::() + .add_systems( + Render, + prepare_generate_environment_map_lights_for_skyboxes_bind_groups + .in_set(RenderSet::PrepareBindGroups), + ); } } @@ -42,12 +114,11 @@ impl Plugin for EnvironmentMapPlugin { /// The environment map must be prefiltered into a diffuse and specular cubemap based on the /// [split-sum approximation](https://cdn2.unrealengine.com/Resources/files/2013SiggraphPresentationsNotes-26915738.pdf). /// -/// To prefilter your environment map, you can use `KhronosGroup`'s [glTF-IBL-Sampler](https://github.com/KhronosGroup/glTF-IBL-Sampler). +/// To prefilter your environment map, you can either: +/// * Use the [`GenerateEnvironmentMapLight`] component at runtime. +/// * Use `KhronosGroup`'s [glTF-IBL-Sampler](https://github.com/KhronosGroup/glTF-IBL-Sampler) to prefilter it offline. /// The diffuse map uses the Lambertian distribution, and the specular map uses the GGX distribution. -/// -/// `KhronosGroup` also has several prefiltered environment maps that can be found [here](https://github.com/KhronosGroup/glTF-Sample-Environments). -#[derive(Component, Reflect, Clone, ExtractComponent)] -#[extract_component_filter(With)] +#[derive(Component, Reflect, Clone)] pub struct EnvironmentMapLight { pub diffuse_map: Handle, pub specular_map: Handle, @@ -61,25 +132,38 @@ impl EnvironmentMapLight { } } +impl ExtractComponent for EnvironmentMapLight { + type Data = &'static Self; + type Filter = With; + type Out = Self; + + fn extract_component(item: QueryItem<'_, Self::Data>) -> Option { + Some(item.clone()) + } +} + pub fn get_bindings<'a>( environment_map_light: Option<&EnvironmentMapLight>, images: &'a RenderAssets, fallback_image_cubemap: &'a FallbackImageCubemap, ) -> (&'a TextureView, &'a TextureView, &'a Sampler) { - let (diffuse_map, specular_map) = match ( + let (diffuse_map, specular_map, sampler) = match ( environment_map_light.and_then(|env_map| images.get(&env_map.diffuse_map)), environment_map_light.and_then(|env_map| images.get(&env_map.specular_map)), ) { - (Some(diffuse_map), Some(specular_map)) => { - (&diffuse_map.texture_view, &specular_map.texture_view) - } + (Some(diffuse_map), Some(specular_map)) => ( + &diffuse_map.texture_view, + &specular_map.texture_view, + &specular_map.sampler, + ), _ => ( &fallback_image_cubemap.texture_view, &fallback_image_cubemap.texture_view, + &fallback_image_cubemap.sampler, ), }; - (diffuse_map, specular_map, &fallback_image_cubemap.sampler) + (diffuse_map, specular_map, sampler) } pub fn get_bind_group_layout_entries() -> [BindGroupLayoutEntryBuilder; 3] { diff --git a/crates/bevy_pbr/src/lib.rs b/crates/bevy_pbr/src/lib.rs index ce104502c8b10..b6407181382a2 100644 --- a/crates/bevy_pbr/src/lib.rs +++ b/crates/bevy_pbr/src/lib.rs @@ -16,7 +16,9 @@ mod ssao; pub use alpha::*; pub use bundle::*; -pub use environment_map::EnvironmentMapLight; +pub use environment_map::{ + EnvironmentMapLight, GenerateEnvironmentMapLight, GenerateEnvironmentMapLightTextureFormat, +}; pub use extended_material::*; pub use fog::*; pub use light::*; @@ -69,7 +71,7 @@ use bevy_render::{ ExtractSchedule, Render, RenderApp, RenderSet, }; use bevy_transform::TransformSystem; -use environment_map::EnvironmentMapPlugin; +use environment_map::EnvironmentMapLightPlugin; use crate::deferred::DeferredPbrLightingPlugin; @@ -253,7 +255,7 @@ impl Plugin for PbrPlugin { ..Default::default() }, ScreenSpaceAmbientOcclusionPlugin, - EnvironmentMapPlugin, + EnvironmentMapLightPlugin, ExtractResourcePlugin::::default(), FogPlugin, ExtractResourcePlugin::::default(),