diff --git a/assets/environment_maps/goegap_road_2k.ktx2 b/assets/environment_maps/goegap_road_2k.ktx2 new file mode 100644 index 0000000000000..50bf6201d66b4 Binary files /dev/null and b/assets/environment_maps/goegap_road_2k.ktx2 differ diff --git a/crates/bevy_pbr/src/atmosphere/mod.rs b/crates/bevy_pbr/src/atmosphere/mod.rs index e7f17f0e1e855..39c7b227d56f2 100644 --- a/crates/bevy_pbr/src/atmosphere/mod.rs +++ b/crates/bevy_pbr/src/atmosphere/mod.rs @@ -74,7 +74,7 @@ use self::{ }, }; -mod shaders { +pub mod shaders { use bevy_asset::{weak_handle, Handle}; use bevy_render::render_resource::Shader; diff --git a/crates/bevy_pbr/src/light_probe/environment_filter.wgsl b/crates/bevy_pbr/src/light_probe/environment_filter.wgsl new file mode 100644 index 0000000000000..0b576b6305670 --- /dev/null +++ b/crates/bevy_pbr/src/light_probe/environment_filter.wgsl @@ -0,0 +1,354 @@ +#import bevy_render::maths::{PI, PI_2, fast_sqrt}; +#import bevy_pbr::lighting::perceptualRoughnessToRoughness; + +struct FilteringConstants { + mip_level: f32, + sample_count: u32, + roughness: f32, + blue_noise_size: vec2f, + white_point: f32, +} + +@group(0) @binding(0) var input_texture: texture_2d_array; +@group(0) @binding(1) var input_sampler: sampler; +@group(0) @binding(2) var output_texture: texture_storage_2d_array; +@group(0) @binding(3) var constants: FilteringConstants; +@group(0) @binding(4) var blue_noise_texture: texture_2d; + +// Tonemapping functions to reduce fireflies +fn tonemap(color: vec3f) -> vec3f { + return color / (color + vec3(constants.white_point)); +} +fn reverse_tonemap(color: vec3f) -> vec3f { + return constants.white_point * color / (vec3(1.0) - color); +} + +// Convert UV and face index to direction vector +fn sample_cube_dir(uv: vec2f, face: u32) -> vec3f { + // Convert from [0,1] to [-1,1] + let uvc = 2.0 * uv - 1.0; + + // Generate direction based on the cube face + var dir: vec3f; + switch(face) { + case 0u: { dir = vec3f( 1.0, -uvc.y, -uvc.x); } // +X + case 1u: { dir = vec3f(-1.0, -uvc.y, uvc.x); } // -X + case 2u: { dir = vec3f( uvc.x, 1.0, uvc.y); } // +Y + case 3u: { dir = vec3f( uvc.x, -1.0, -uvc.y); } // -Y + case 4u: { dir = vec3f( uvc.x, -uvc.y, 1.0); } // +Z + case 5u: { dir = vec3f(-uvc.x, -uvc.y, -1.0); } // -Z + default: { dir = vec3f(0.0); } + } + return normalize(dir); +} + +// Convert direction vector to cube face UV +struct CubeUV { + uv: vec2f, + face: u32, +} +fn dir_to_cube_uv(dir: vec3f) -> CubeUV { + let abs_dir = abs(dir); + var face: u32 = 0u; + var uv: vec2f = vec2f(0.0); + + // Find the dominant axis to determine face + if (abs_dir.x >= abs_dir.y && abs_dir.x >= abs_dir.z) { + // X axis is dominant + if (dir.x > 0.0) { + face = 0u; // +X + uv = vec2f(-dir.z, -dir.y) / dir.x; + } else { + face = 1u; // -X + uv = vec2f(dir.z, -dir.y) / abs_dir.x; + } + } else if (abs_dir.y >= abs_dir.x && abs_dir.y >= abs_dir.z) { + // Y axis is dominant + if (dir.y > 0.0) { + face = 2u; // +Y + uv = vec2f(dir.x, dir.z) / dir.y; + } else { + face = 3u; // -Y + uv = vec2f(dir.x, -dir.z) / abs_dir.y; + } + } else { + // Z axis is dominant + if (dir.z > 0.0) { + face = 4u; // +Z + uv = vec2f(dir.x, -dir.y) / dir.z; + } else { + face = 5u; // -Z + uv = vec2f(-dir.x, -dir.y) / abs_dir.z; + } + } + + // Convert from [-1,1] to [0,1] + return CubeUV(uv * 0.5 + 0.5, face); +} + +// Sample an environment map with a specific LOD +fn sample_environment(dir: vec3f, level: f32) -> vec4f { + let cube_uv = dir_to_cube_uv(dir); + return textureSampleLevel(input_texture, input_sampler, cube_uv.uv, cube_uv.face, level); +} + +// Calculate tangent space for the given normal +fn calculate_tangent_frame(normal: vec3f) -> mat3x3f { + // Use a robust method to pick a tangent + var up = vec3f(1.0, 0.0, 0.0); + if abs(normal.z) < 0.999 { + up = vec3f(0.0, 0.0, 1.0); + } + let tangent = normalize(cross(up, normal)); + let bitangent = cross(normal, tangent); + return mat3x3f(tangent, bitangent, normal); +} + +// Hammersley sequence for quasi-random points +fn hammersley_2d(i: u32, n: u32) -> vec2f { + // Van der Corput sequence + var bits = i; + bits = (bits << 16u) | (bits >> 16u); + bits = ((bits & 0x55555555u) << 1u) | ((bits & 0xAAAAAAAAu) >> 1u); + bits = ((bits & 0x33333333u) << 2u) | ((bits & 0xCCCCCCCCu) >> 2u); + bits = ((bits & 0x0F0F0F0Fu) << 4u) | ((bits & 0xF0F0F0F0u) >> 4u); + bits = ((bits & 0x00FF00FFu) << 8u) | ((bits & 0xFF00FF00u) >> 8u); + let vdc = f32(bits) * 2.3283064365386963e-10; // 1 / 0x100000000 + return vec2f(f32(i) / f32(n), vdc); +} + +// Blue noise randomization +fn sample_noise(pixel_coords: vec2u) -> vec4f { + let noise_size = vec2u(u32(constants.blue_noise_size.x), u32(constants.blue_noise_size.y)); + let noise_coords = pixel_coords % noise_size; + let uv = vec2f(noise_coords) / constants.blue_noise_size; + return textureSampleLevel(blue_noise_texture, input_sampler, uv, 0.0); +} + +// GGX/Trowbridge-Reitz normal distribution function (D term) +fn D_GGX(roughness: f32, NdotH: f32) -> f32 { + let oneMinusNdotHSquared = 1.0 - NdotH * NdotH; + let a = NdotH * roughness; + let k = roughness / (oneMinusNdotHSquared + a * a); + let d = k * k * (1.0 / PI); + return d; +} + +// Importance sample GGX normal distribution function for a given roughness +fn importance_sample_ggx(xi: vec2f, roughness: f32, normal: vec3f) -> vec3f { + // Use roughness^2 to ensure correct specular highlights + let a = roughness * roughness; + + // Sample in spherical coordinates + let phi = 2.0 * PI * xi.x; + + // GGX mapping from uniform random to GGX distribution + let cos_theta = fast_sqrt((1.0 - xi.y) / (1.0 + (a * a - 1.0) * xi.y)); + let sin_theta = fast_sqrt(1.0 - cos_theta * cos_theta); + + // Convert to cartesian + let h = vec3f( + sin_theta * cos(phi), + sin_theta * sin(phi), + cos_theta + ); + + // Transform from tangent to world space + return calculate_tangent_frame(normal) * h; +} + +// Calculate LOD for environment map lookup using filtered importance sampling +fn calculate_environment_map_lod(pdf: f32, width: f32, samples: f32) -> f32 { + // Solid angle of current sample + let omega_s = 1.0 / (samples * pdf); + + // Solid angle of a texel in the environment map + let omega_p = 4.0 * PI / (6.0 * width * width); + + // Filtered importance sampling: compute the correct LOD + return 0.5 * log2(omega_s / omega_p); +} + +// Smith geometric shadowing function +fn G_Smith(NoV: f32, NoL: f32, roughness: f32) -> f32 { + let k = (roughness * roughness) / 2.0; + let GGXL = NoL / (NoL * (1.0 - k) + k); + let GGXV = NoV / (NoV * (1.0 - k) + k); + return GGXL * GGXV; +} + +@compute +@workgroup_size(8, 8, 1) +fn generate_radiance_map(@builtin(global_invocation_id) global_id: vec3u) { + let size = textureDimensions(output_texture).xy; + let invSize = 1.0 / vec2f(size); + + let coords = vec2u(global_id.xy); + let face = global_id.z; + + if (any(coords >= size)) { + return; + } + + // Convert texture coordinates to direction vector + let uv = (vec2f(coords) + 0.5) * invSize; + let normal = sample_cube_dir(uv, face); + + // For radiance map, view direction = normal for perfect reflection + let view = normal; + + // Get the roughness parameter + let roughness = constants.roughness; + + // Get blue noise offset for stratification + let vector_noise = sample_noise(coords); + + var radiance = vec3f(0.0); + var total_weight = 0.0; + + // Skip sampling for mirror reflection (roughness = 0) + if (roughness < 0.01) { + radiance = sample_environment(normal, 0.0).rgb; + textureStore(output_texture, coords, face, vec4f(radiance, 1.0)); + return; + } + + // For higher roughness values, use importance sampling + let sample_count = constants.sample_count; + + for (var i = 0u; i < sample_count; i++) { + // Get sample coordinates from Hammersley sequence with blue noise offset + var xi = hammersley_2d(i, sample_count); + xi = fract(xi + vector_noise.rg); // Apply Cranley-Patterson rotation + + // Sample the GGX distribution to get a half vector + let half_vector = importance_sample_ggx(xi, roughness, normal); + + // Calculate reflection vector from half vector + let light_dir = reflect(-view, half_vector); + + // Calculate weight (N·L) + let NoL = dot(normal, light_dir); + + if (NoL > 0.0) { + // Calculate values needed for PDF + let NoH = dot(normal, half_vector); + let VoH = dot(view, half_vector); + let NoV = dot(normal, view); + + // Get the geometric shadowing term + let G = G_Smith(NoV, NoL, roughness); + + // Probability Distribution Function + let pdf = D_GGX(roughness, NoH) * NoH / (4.0 * VoH); + + // Calculate LOD using filtered importance sampling + // This is crucial to avoid fireflies and improve quality + let width = f32(size.x); + let lod = calculate_environment_map_lod(pdf, width, f32(sample_count)); + + // Get source mip level - ensure we don't go negative + let source_mip = max(0.0, lod); + + // Sample environment map with the light direction + var sample_color = sample_environment(light_dir, source_mip).rgb; + sample_color = tonemap(sample_color); + + // Accumulate weighted sample, including geometric term + radiance += sample_color * NoL * G; + total_weight += NoL * G; + } + } + + // Normalize by total weight + if (total_weight > 0.0) { + radiance = radiance / total_weight; + } + + // Reverse tonemap + radiance = reverse_tonemap(radiance); + + // Write result to output texture + textureStore(output_texture, coords, face, vec4f(radiance, 1.0)); +} + +// Calculate spherical coordinates using spiral pattern +// and golden angle to get a uniform distribution +fn uniform_sample_sphere(i: u32, normal: vec3f) -> vec3f { + // Get stratified sample index + let strat_i = i % constants.sample_count; + + let golden_angle = 2.4; + let full_sphere = f32(constants.sample_count) * 2.0; + let z = 1.0 - (2.0 * f32(strat_i) + 1.0) / full_sphere; + let r = fast_sqrt(1.0 - z * z); + + let phi = f32(strat_i) * golden_angle; + + // Create the direction vector + let dir_uniform = vec3f( + r * cos(phi), + r * sin(phi), + z + ); + + let tangent_frame = calculate_tangent_frame(normal); + return normalize(tangent_frame * dir_uniform); +} + +@compute +@workgroup_size(8, 8, 1) +fn generate_irradiance_map(@builtin(global_invocation_id) global_id: vec3u) { + let size = textureDimensions(output_texture).xy; + let invSize = 1.0 / vec2f(size); + + let coords = vec2u(global_id.xy); + let face = global_id.z; + + if (any(coords >= size)) { + return; + } + + // Convert texture coordinates to direction vector + let uv = (vec2f(coords) + 0.5) * invSize; + let normal = sample_cube_dir(uv, face); + + var irradiance = vec3f(0.0); + var total_weight = 0.0; + + // Use uniform sampling on a hemisphere + for (var i = 0u; i < constants.sample_count; i++) { + // Get a uniform direction on unit sphere + var sample_dir = uniform_sample_sphere(i, normal); + + // Calculate the cosine weight (N·L) + let weight = max(dot(normal, sample_dir), 0.0); + + // Skip samples below horizon or at grazing angles + if (weight <= 0.001) { + continue; + } + + // Sample environment with level 0 (no mip) + var sample_color = sample_environment(sample_dir, 0.0).rgb; + + // Apply tonemapping to reduce fireflies + sample_color = tonemap(sample_color); + + // Accumulate the contribution + irradiance += sample_color * weight; + total_weight += weight; + } + + // Normalize by total weight + irradiance = irradiance / total_weight; + + // Scale by PI to account for the Lambert BRDF normalization factor + irradiance *= PI; + + // Reverse tonemap to restore HDR range + irradiance = reverse_tonemap(irradiance); + + // Write result to output texture + textureStore(output_texture, coords, face, vec4f(irradiance, 1.0)); +} diff --git a/crates/bevy_pbr/src/light_probe/generate.rs b/crates/bevy_pbr/src/light_probe/generate.rs new file mode 100644 index 0000000000000..71f6ab2c401a3 --- /dev/null +++ b/crates/bevy_pbr/src/light_probe/generate.rs @@ -0,0 +1,961 @@ +use bevy_asset::{weak_handle, Assets, Handle}; +use bevy_ecs::{ + component::Component, + entity::Entity, + query::{QueryState, With, Without}, + resource::Resource, + system::{lifetimeless::Read, Commands, Query, Res, ResMut}, + world::{FromWorld, World}, +}; +use bevy_image::Image; +use bevy_math::{Quat, Vec2}; +use bevy_reflect::Reflect; +use bevy_render::{ + extract_component::ExtractComponent, + render_asset::{RenderAssetUsages, RenderAssets}, + render_graph::{Node, NodeRunError, RenderGraphContext, RenderLabel}, + render_resource::{ + binding_types::*, AddressMode, BindGroup, BindGroupEntries, BindGroupLayout, + BindGroupLayoutEntries, CachedComputePipelineId, ComputePassDescriptor, + ComputePipelineDescriptor, Extent3d, FilterMode, PipelineCache, Sampler, + SamplerBindingType, SamplerDescriptor, Shader, ShaderDefVal, ShaderStages, ShaderType, + StorageTextureAccess, Texture, TextureAspect, TextureDescriptor, TextureDimension, + TextureFormat, TextureSampleType, TextureUsages, TextureView, TextureViewDescriptor, + TextureViewDimension, UniformBuffer, + }, + renderer::{RenderContext, RenderDevice, RenderQueue}, + settings::WgpuFeatures, + sync_world::RenderEntity, + texture::{CachedTexture, GpuImage, TextureCache}, + Extract, +}; + +use crate::light_probe::environment_map::EnvironmentMapLight; + +/// Single Pass Downsampling (SPD) shader handle +pub const SPD_SHADER_HANDLE: Handle = weak_handle!("5dcf400c-bcb3-49b9-8b7e-80f4117eaf82"); + +/// Environment Filter shader handle +pub const ENVIRONMENT_FILTER_SHADER_HANDLE: Handle = + weak_handle!("3110b545-78e0-48fc-b86e-8bc0ea50fc67"); + +/// Sphere Cosine Weighted Irradiance shader handle +pub const STBN_SPHERE: Handle = weak_handle!("3110b545-78e0-48fc-b86e-8bc0ea50fc67"); +pub const STBN_VEC2: Handle = weak_handle!("3110b545-78e0-48fc-b86e-8bc0ea50fc67"); + +/// Labels for the environment map generation nodes +#[derive(PartialEq, Eq, Debug, Copy, Clone, Hash, RenderLabel)] +pub enum GeneratorNode { + Mipmap, + Radiance, + Irradiance, +} + +/// Stores the bind group layouts for the environment map generation pipelines +#[derive(Resource)] +pub struct GeneratorBindGroupLayouts { + pub spd: BindGroupLayout, + pub radiance: BindGroupLayout, + pub irradiance: BindGroupLayout, +} + +impl FromWorld for GeneratorBindGroupLayouts { + fn from_world(world: &mut World) -> Self { + let render_device = world.resource::(); + + // SPD (Single Pass Downsampling) bind group layout + let spd = render_device.create_bind_group_layout( + "spd_bind_group_layout", + &BindGroupLayoutEntries::with_indices( + ShaderStages::COMPUTE, + ( + ( + 0, + texture_2d_array(TextureSampleType::Float { filterable: true }), + ), // Source texture + ( + 1, + texture_storage_2d_array( + TextureFormat::Rgba16Float, + StorageTextureAccess::WriteOnly, + ), + ), // Output mip 1 + ( + 2, + texture_storage_2d_array( + TextureFormat::Rgba16Float, + StorageTextureAccess::WriteOnly, + ), + ), // Output mip 2 + ( + 3, + texture_storage_2d_array( + TextureFormat::Rgba16Float, + StorageTextureAccess::WriteOnly, + ), + ), // Output mip 3 + ( + 4, + texture_storage_2d_array( + TextureFormat::Rgba16Float, + StorageTextureAccess::WriteOnly, + ), + ), // Output mip 4 + ( + 5, + texture_storage_2d_array( + TextureFormat::Rgba16Float, + StorageTextureAccess::WriteOnly, + ), + ), // Output mip 5 + ( + 6, + texture_storage_2d_array( + TextureFormat::Rgba16Float, + StorageTextureAccess::ReadWrite, + ), + ), // Output mip 6 + ( + 7, + texture_storage_2d_array( + TextureFormat::Rgba16Float, + StorageTextureAccess::WriteOnly, + ), + ), // Output mip 7 + ( + 8, + texture_storage_2d_array( + TextureFormat::Rgba16Float, + StorageTextureAccess::WriteOnly, + ), + ), // Output mip 8 + // ( + // 9, + // texture_storage_2d_array( + // TextureFormat::Rgba16Float, + // StorageTextureAccess::WriteOnly, + // ), + // ), // Output mip 9 + // ( + // 10, + // texture_storage_2d_array( + // TextureFormat::Rgba16Float, + // StorageTextureAccess::WriteOnly, + // ), + // ), // Output mip 10 + // ( + // 11, + // texture_storage_2d_array( + // TextureFormat::Rgba16Float, + // StorageTextureAccess::WriteOnly, + // ), + // ), // Output mip 11 + // ( + // 12, + // texture_storage_2d_array( + // TextureFormat::Rgba16Float, + // StorageTextureAccess::WriteOnly, + // ), + // ), // Output mip 12 + (13, sampler(SamplerBindingType::Filtering)), // Linear sampler + (14, uniform_buffer::(false)), // Uniforms + ), + ), + ); + + // Radiance map bind group layout + let radiance = render_device.create_bind_group_layout( + "radiance_bind_group_layout", + &BindGroupLayoutEntries::with_indices( + ShaderStages::COMPUTE, + ( + ( + 0, + texture_2d_array(TextureSampleType::Float { filterable: true }), + ), // Source environment cubemap + (1, sampler(SamplerBindingType::Filtering)), // Source sampler + ( + 2, + texture_storage_2d_array( + TextureFormat::Rgba16Float, + StorageTextureAccess::WriteOnly, + ), + ), // Output specular map + (3, uniform_buffer::(false)), // Uniforms + (4, texture_2d(TextureSampleType::Float { filterable: true })), // Blue noise texture + ), + ), + ); + + // Irradiance convolution bind group layout + let irradiance = render_device.create_bind_group_layout( + "irradiance_bind_group_layout", + &BindGroupLayoutEntries::with_indices( + ShaderStages::COMPUTE, + ( + ( + 0, + texture_2d_array(TextureSampleType::Float { filterable: true }), + ), // Source environment cubemap + (1, sampler(SamplerBindingType::Filtering)), // Source sampler + ( + 2, + texture_storage_2d_array( + TextureFormat::Rgba16Float, + StorageTextureAccess::WriteOnly, + ), + ), // Output irradiance map + (3, uniform_buffer::(false)), // Uniforms + (4, texture_2d(TextureSampleType::Float { filterable: true })), // Blue noise texture + ), + ), + ); + + Self { + spd, + radiance, + irradiance, + } + } +} + +/// Samplers for the environment map generation pipelines +#[derive(Resource)] +pub struct GeneratorSamplers { + pub linear: Sampler, +} + +impl FromWorld for GeneratorSamplers { + fn from_world(world: &mut World) -> Self { + let render_device = world.resource::(); + + let linear = render_device.create_sampler(&SamplerDescriptor { + label: Some("generator_linear_sampler"), + address_mode_u: AddressMode::ClampToEdge, + address_mode_v: AddressMode::ClampToEdge, + address_mode_w: AddressMode::ClampToEdge, + mag_filter: FilterMode::Linear, + min_filter: FilterMode::Linear, + mipmap_filter: FilterMode::Linear, + ..Default::default() + }); + + Self { linear } + } +} + +/// Pipelines for the environment map generation pipelines +#[derive(Resource)] +pub struct GeneratorPipelines { + pub spd_first: CachedComputePipelineId, + pub spd_second: CachedComputePipelineId, + pub radiance: CachedComputePipelineId, + pub irradiance: CachedComputePipelineId, +} + +impl FromWorld for GeneratorPipelines { + fn from_world(world: &mut World) -> Self { + let pipeline_cache = world.resource::(); + let layouts = world.resource::(); + + let render_device = world.resource::(); + let features = render_device.features(); + let shader_defs = if features.contains(WgpuFeatures::SUBGROUP) { + vec![ShaderDefVal::Int("SUBGROUP_SUPPORT".into(), 1)] + } else { + vec![] + }; + + // Single Pass Downsampling for Base Mip Levels (0-5) + let spd_first = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { + label: Some("spd_first_pipeline".into()), + layout: vec![layouts.spd.clone()], + push_constant_ranges: vec![], + shader: SPD_SHADER_HANDLE, + shader_defs: shader_defs.clone(), + entry_point: "spd_downsample_first".into(), + zero_initialize_workgroup_memory: false, + }); + + // Single Pass Downsampling for Remaining Mip Levels (6-12) + let spd_second = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { + label: Some("spd_second_pipeline".into()), + layout: vec![layouts.spd.clone()], + push_constant_ranges: vec![], + shader: SPD_SHADER_HANDLE, + shader_defs, + entry_point: "spd_downsample_second".into(), + zero_initialize_workgroup_memory: false, + }); + + // Radiance map for Specular Environment Maps + let radiance = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { + label: Some("radiance_pipeline".into()), + layout: vec![layouts.radiance.clone()], + push_constant_ranges: vec![], + shader: ENVIRONMENT_FILTER_SHADER_HANDLE, + shader_defs: vec![], + entry_point: "generate_radiance_map".into(), + zero_initialize_workgroup_memory: false, + }); + + // Irradiance map for Diffuse Environment Maps + let irradiance = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { + label: Some("irradiance_pipeline".into()), + layout: vec![layouts.irradiance.clone()], + push_constant_ranges: vec![], + shader: ENVIRONMENT_FILTER_SHADER_HANDLE, + shader_defs: vec![], + entry_point: "generate_irradiance_map".into(), + zero_initialize_workgroup_memory: false, + }); + + Self { + spd_first, + spd_second, + radiance, + irradiance, + } + } +} + +#[derive(Component, Clone, Reflect, ExtractComponent)] +pub struct GeneratedEnvironmentMapLight { + pub environment_map: Handle, + pub intensity: f32, + pub rotation: Quat, + pub affects_lightmapped_mesh_diffuse: bool, + pub white_point: f32, +} + +impl Default for GeneratedEnvironmentMapLight { + fn default() -> Self { + GeneratedEnvironmentMapLight { + environment_map: Handle::default(), + intensity: 0.0, + rotation: Quat::IDENTITY, + affects_lightmapped_mesh_diffuse: true, + white_point: 2.0, + } + } +} + +pub fn extract_generator_entities( + query: Extract< + Query<( + RenderEntity, + &GeneratedEnvironmentMapLight, + &EnvironmentMapLight, + )>, + >, + mut commands: Commands, + render_images: Res>, +) { + for (entity, filtered_env_map, env_map_light) in query.iter() { + let env_map = render_images + .get(&filtered_env_map.environment_map) + .expect("Environment map not found"); + + let diffuse_map = render_images.get(&env_map_light.diffuse_map); + let specular_map = render_images.get(&env_map_light.specular_map); + + // continue if the diffuse map is not found + if diffuse_map.is_none() || specular_map.is_none() { + continue; + } + + let diffuse_map = diffuse_map.unwrap(); + let specular_map = specular_map.unwrap(); + + let render_filtered_env_map = RenderEnvironmentMap { + environment_map: env_map.clone(), + diffuse_map: diffuse_map.clone(), + specular_map: specular_map.clone(), + intensity: filtered_env_map.intensity, + rotation: filtered_env_map.rotation, + affects_lightmapped_mesh_diffuse: filtered_env_map.affects_lightmapped_mesh_diffuse, + white_point: filtered_env_map.white_point, + }; + commands + .get_entity(entity) + .expect("Entity not synced to render world") + .insert(render_filtered_env_map); + } +} + +// A render-world specific version of FilteredEnvironmentMapLight that uses CachedTexture +#[derive(Component, Clone)] +pub struct RenderEnvironmentMap { + pub environment_map: GpuImage, + pub diffuse_map: GpuImage, + pub specular_map: GpuImage, + pub intensity: f32, + pub rotation: Quat, + pub affects_lightmapped_mesh_diffuse: bool, + pub white_point: f32, +} + +#[derive(Component)] +pub struct IntermediateTextures { + pub environment_map: CachedTexture, +} + +/// Prepares textures needed for single pass downsampling +pub fn prepare_intermediate_textures( + light_probes: Query>, + render_device: Res, + mut texture_cache: ResMut, + mut commands: Commands, +) { + for entity in &light_probes { + // Create environment map with 8 mip levels (512x512 -> 1x1) + let environment_map = texture_cache.get( + &render_device, + TextureDescriptor { + label: Some("intermediate_environment_map"), + size: Extent3d { + width: 512, + height: 512, + depth_or_array_layers: 6, // Cubemap faces + }, + mip_level_count: 9, // 512, 256, 128, 64, 32, 16, 8, 4, 2, 1 + sample_count: 1, + dimension: TextureDimension::D2, + format: TextureFormat::Rgba16Float, + usage: TextureUsages::TEXTURE_BINDING + | TextureUsages::STORAGE_BINDING + | TextureUsages::COPY_DST, + view_formats: &[], + }, + ); + + commands + .entity(entity) + .insert(IntermediateTextures { environment_map }); + } +} + +/// Shader constants for SPD algorithm +#[derive(Clone, Copy, ShaderType)] +#[repr(C)] +pub struct SpdConstants { + mips: u32, + inverse_input_size: Vec2, + _padding: u32, +} + +/// Constants for filtering +#[derive(Clone, Copy, ShaderType)] +#[repr(C)] +pub struct FilteringConstants { + mip_level: f32, + sample_count: u32, + roughness: f32, + blue_noise_size: Vec2, + white_point: f32, +} + +/// Stores bind groups for the environment map generation pipelines +#[derive(Component)] +pub struct GeneratorBindGroups { + pub spd: BindGroup, + pub radiance: Vec, // One per mip level + pub irradiance: BindGroup, +} + +/// Prepares bind groups for environment map generation pipelines +pub fn prepare_generator_bind_groups( + light_probes: Query< + (Entity, &IntermediateTextures, &RenderEnvironmentMap), + With, + >, + render_device: Res, + queue: Res, + layouts: Res, + samplers: Res, + render_images: Res>, + mut commands: Commands, +) { + // Get blue noise texture + let sphere_cosine_weights = render_images + .get(&STBN_SPHERE) + .expect("Sphere cosine weights texture not loaded"); + + let vector2_uniform = render_images + .get(&STBN_VEC2) + .expect("Vector2 uniform texture not loaded"); + + for (entity, textures, env_map_light) in &light_probes { + // Create SPD bind group + let spd_constants = SpdConstants { + mips: 8, // Number of mip levels + inverse_input_size: Vec2::new(1.0 / 512.0, 1.0 / 512.0), // 1.0 / input size + _padding: 0, + }; + + let mut spd_constants_buffer = UniformBuffer::from(spd_constants); + spd_constants_buffer.write_buffer(&render_device, &queue); + + let input_env_map = + env_map_light + .environment_map + .texture + .create_view(&TextureViewDescriptor { + dimension: Some(TextureViewDimension::D2Array), + ..Default::default() + }); + + let spd_bind_group = render_device.create_bind_group( + "spd_bind_group", + &layouts.spd, + &BindGroupEntries::with_indices(( + (0, &input_env_map), + ( + 1, + &create_storage_view(&textures.environment_map.texture, 1, &render_device), + ), + ( + 2, + &create_storage_view(&textures.environment_map.texture, 2, &render_device), + ), + ( + 3, + &create_storage_view(&textures.environment_map.texture, 3, &render_device), + ), + ( + 4, + &create_storage_view(&textures.environment_map.texture, 4, &render_device), + ), + ( + 5, + &create_storage_view(&textures.environment_map.texture, 5, &render_device), + ), + ( + 6, + &create_storage_view(&textures.environment_map.texture, 6, &render_device), + ), + ( + 7, + &create_storage_view(&textures.environment_map.texture, 7, &render_device), + ), + ( + 8, + &create_storage_view(&textures.environment_map.texture, 8, &render_device), + ), + // ( + // 9, + // &create_storage_view(&textures.environment_map.texture, 9, &render_device), + // ), + // ( + // 10, + // &create_storage_view(&textures.environment_map.texture, 10, &render_device), + // ), + // ( + // 11, + // &create_storage_view(&textures.environment_map.texture, 11, &render_device), + // ), + // ( + // 12, + // &create_storage_view(&textures.environment_map.texture, 12, &render_device), + // ), + (13, &samplers.linear), + (14, &spd_constants_buffer), + )), + ); + + // Create radiance map bind groups for each mip level + let num_mips = 9; + let mut radiance_bind_groups = Vec::with_capacity(num_mips); + + for mip in 0..num_mips { + let roughness = mip as f32 / (num_mips - 1) as f32; + + // For higher roughness values, use importance sampling with optimized sample count + let sample_count = if roughness < 0.01 { + 1 // Mirror reflection + } else if roughness < 0.25 { + 16 + } else if roughness < 0.5 { + 32 + } else if roughness < 0.75 { + 64 + } else { + 128 + }; + + let radiance_constants = FilteringConstants { + mip_level: mip as f32, + sample_count, + roughness, + blue_noise_size: Vec2::new( + vector2_uniform.size.width as f32, + vector2_uniform.size.height as f32, + ), + white_point: env_map_light.white_point, + }; + + let mut radiance_constants_buffer = UniformBuffer::from(radiance_constants); + radiance_constants_buffer.write_buffer(&render_device, &queue); + + let mip_storage_view = create_storage_view( + &env_map_light.specular_map.texture, + mip as u32, + &render_device, + ); + let bind_group = render_device.create_bind_group( + Some(format!("radiance_bind_group_mip_{}", mip).as_str()), + &layouts.radiance, + &BindGroupEntries::with_indices(( + (0, &textures.environment_map.default_view), + (1, &samplers.linear), + (2, &mip_storage_view), + (3, &radiance_constants_buffer), + (4, &vector2_uniform.texture_view), + )), + ); + + radiance_bind_groups.push(bind_group); + } + + // Create irradiance bind group + let irradiance_constants = FilteringConstants { + mip_level: 0.0, + // 32 phi, 32 theta = 1024 samples total + sample_count: 1024, + roughness: 1.0, + blue_noise_size: Vec2::new( + sphere_cosine_weights.size.width as f32, + sphere_cosine_weights.size.height as f32, + ), + white_point: env_map_light.white_point, + }; + + let mut irradiance_constants_buffer = UniformBuffer::from(irradiance_constants); + irradiance_constants_buffer.write_buffer(&render_device, &queue); + + // create a 2d array view + let irradiance_map = + env_map_light + .diffuse_map + .texture + .create_view(&TextureViewDescriptor { + dimension: Some(TextureViewDimension::D2Array), + ..Default::default() + }); + + let irradiance_bind_group = render_device.create_bind_group( + "irradiance_bind_group", + &layouts.irradiance, + &BindGroupEntries::with_indices(( + (0, &textures.environment_map.default_view), + (1, &samplers.linear), + (2, &irradiance_map), + (3, &irradiance_constants_buffer), + (4, &sphere_cosine_weights.texture_view), + )), + ); + + commands.entity(entity).insert(GeneratorBindGroups { + spd: spd_bind_group, + radiance: radiance_bind_groups, + irradiance: irradiance_bind_group, + }); + } +} + +/// Helper function to create a storage texture view for a specific mip level +fn create_storage_view(texture: &Texture, mip: u32, _render_device: &RenderDevice) -> TextureView { + texture.create_view(&TextureViewDescriptor { + label: Some(format!("storage_view_mip_{}", mip).as_str()), + format: Some(texture.format()), + dimension: Some(TextureViewDimension::D2Array), + aspect: TextureAspect::All, + base_mip_level: mip, + mip_level_count: Some(1), + base_array_layer: 0, + array_layer_count: Some(texture.depth_or_array_layers()), + usage: Some(TextureUsages::STORAGE_BINDING), + }) +} + +/// SPD Node implementation that handles both parts of the downsampling (mips 0-12) +pub struct SpdNode { + query: QueryState<( + Entity, + Read, + Read, + )>, +} + +impl FromWorld for SpdNode { + fn from_world(world: &mut World) -> Self { + Self { + query: QueryState::new(world), + } + } +} + +impl Node for SpdNode { + fn update(&mut self, world: &mut World) { + self.query.update_archetypes(world); + } + + fn run( + &self, + _graph: &mut RenderGraphContext, + render_context: &mut RenderContext, + world: &World, + ) -> Result<(), NodeRunError> { + let pipeline_cache = world.resource::(); + let pipelines = world.resource::(); + + // First pass (mips 0-5) + let Some(spd_first_pipeline) = pipeline_cache.get_compute_pipeline(pipelines.spd_first) + else { + return Ok(()); + }; + + // Second pass (mips 6-12) + let Some(spd_second_pipeline) = pipeline_cache.get_compute_pipeline(pipelines.spd_second) + else { + return Ok(()); + }; + + for (entity, bind_groups, env_map_light) in self.query.iter_manual(world) { + // Copy original environment map to mip 0 of the intermediate environment map + let textures = world.get::(entity).unwrap(); + + render_context.command_encoder().copy_texture_to_texture( + env_map_light.environment_map.texture.as_image_copy(), + textures.environment_map.texture.as_image_copy(), + Extent3d { + width: 512, + height: 512, + depth_or_array_layers: 6, + }, + ); + + // First pass - process mips 0-5 + { + let mut compute_pass = + render_context + .command_encoder() + .begin_compute_pass(&ComputePassDescriptor { + label: Some("spd_first_pass"), + timestamp_writes: None, + }); + + compute_pass.set_pipeline(spd_first_pipeline); + compute_pass.set_bind_group(0, &bind_groups.spd, &[]); + + // Calculate the optimal dispatch size based on our shader's workgroup size and thread mapping + // The workgroup size is 256x1x1, and our remap_for_wave_reduction maps these threads to a 8x8 block + // For a 512x512 texture, we need 512/64 = 8 workgroups in X and 512/64 = 8 workgroups in Y + // Each workgroup processes 64x64 pixels (256 threads each handling 16 pixels) + compute_pass.dispatch_workgroups(8, 8, 6); // 6 faces of cubemap + } + + // Second pass - process mips 6-12 + { + let mut compute_pass = + render_context + .command_encoder() + .begin_compute_pass(&ComputePassDescriptor { + label: Some("spd_second_pass"), + timestamp_writes: None, + }); + + compute_pass.set_pipeline(spd_second_pipeline); + compute_pass.set_bind_group(0, &bind_groups.spd, &[]); + + // Dispatch workgroups - for each face + compute_pass.dispatch_workgroups(2, 2, 6); + } + } + + Ok(()) + } +} + +/// Radiance map node for generating specular environment maps +pub struct RadianceMapNode { + query: QueryState<(Entity, Read)>, +} + +impl FromWorld for RadianceMapNode { + fn from_world(world: &mut World) -> Self { + Self { + query: QueryState::new(world), + } + } +} + +impl Node for RadianceMapNode { + fn update(&mut self, world: &mut World) { + self.query.update_archetypes(world); + } + + fn run( + &self, + _graph: &mut RenderGraphContext, + render_context: &mut RenderContext, + world: &World, + ) -> Result<(), NodeRunError> { + let pipeline_cache = world.resource::(); + let pipelines = world.resource::(); + + let Some(radiance_pipeline) = pipeline_cache.get_compute_pipeline(pipelines.radiance) + else { + return Ok(()); + }; + + for (_, bind_groups) in self.query.iter_manual(world) { + let mut compute_pass = + render_context + .command_encoder() + .begin_compute_pass(&ComputePassDescriptor { + label: Some("radiance_map_pass"), + timestamp_writes: None, + }); + + compute_pass.set_pipeline(radiance_pipeline); + + // Process each mip level + for (mip, bind_group) in bind_groups.radiance.iter().enumerate() { + compute_pass.set_bind_group(0, bind_group, &[]); + + // Calculate dispatch size based on mip level + let mip_size = 512u32 >> mip; + let workgroup_count = mip_size.max(8) / 8; + + // Dispatch for all 6 faces + compute_pass.dispatch_workgroups(workgroup_count, workgroup_count, 6); + } + } + + Ok(()) + } +} + +/// Irradiance Convolution Node +pub struct IrradianceMapNode { + query: QueryState<(Entity, Read)>, +} + +impl FromWorld for IrradianceMapNode { + fn from_world(world: &mut World) -> Self { + Self { + query: QueryState::new(world), + } + } +} + +impl Node for IrradianceMapNode { + fn update(&mut self, world: &mut World) { + self.query.update_archetypes(world); + } + + fn run( + &self, + _graph: &mut RenderGraphContext, + render_context: &mut RenderContext, + world: &World, + ) -> Result<(), NodeRunError> { + let pipeline_cache = world.resource::(); + let pipelines = world.resource::(); + + let Some(irradiance_pipeline) = pipeline_cache.get_compute_pipeline(pipelines.irradiance) + else { + return Ok(()); + }; + + for (_, bind_groups) in self.query.iter_manual(world) { + let mut compute_pass = + render_context + .command_encoder() + .begin_compute_pass(&ComputePassDescriptor { + label: Some("irradiance_map_pass"), + timestamp_writes: None, + }); + + compute_pass.set_pipeline(irradiance_pipeline); + compute_pass.set_bind_group(0, &bind_groups.irradiance, &[]); + + // Dispatch workgroups - 32x32 texture with 8x8 workgroups + compute_pass.dispatch_workgroups(4, 4, 6); // 6 faces + } + + Ok(()) + } +} + +/// System that generates an `EnvironmentMapLight` component based on the `GeneratedEnvironmentMapLight` component +pub fn generate_environment_map_light( + mut commands: Commands, + mut images: ResMut>, + query: Query<(Entity, &GeneratedEnvironmentMapLight), Without>, +) { + for (entity, filtered_env_map) in &query { + // Create a placeholder for the irradiance map + let mut diffuse = Image::new_fill( + Extent3d { + width: 32, + height: 32, + depth_or_array_layers: 6, + }, + TextureDimension::D2, + &[0; 8], + TextureFormat::Rgba16Float, + RenderAssetUsages::all(), + ); + + diffuse.texture_descriptor.usage = + TextureUsages::TEXTURE_BINDING | TextureUsages::STORAGE_BINDING; + + diffuse.texture_view_descriptor = Some(TextureViewDescriptor { + dimension: Some(TextureViewDimension::Cube), + ..Default::default() + }); + + let diffuse_handle = images.add(diffuse); + + // Create a placeholder for the specular map + let mut specular = Image::new_fill( + Extent3d { + width: 512, + height: 512, + depth_or_array_layers: 6, + }, + TextureDimension::D2, + &[0; 8], + TextureFormat::Rgba16Float, + RenderAssetUsages::all(), + ); + + // Set up for mipmaps + specular.texture_descriptor.usage = + TextureUsages::TEXTURE_BINDING | TextureUsages::STORAGE_BINDING; + specular.texture_descriptor.mip_level_count = 9; + + // When setting mip_level_count, we need to allocate appropriate data size + // For GPU-generated mipmaps, we can set data to None since the GPU will generate the data + specular.data = None; + + specular.texture_view_descriptor = Some(TextureViewDescriptor { + dimension: Some(TextureViewDimension::Cube), + mip_level_count: Some(9), + ..Default::default() + }); + + let specular_handle = images.add(specular); + + // Add the EnvironmentMapLight component with the placeholder handles + commands.entity(entity).insert(EnvironmentMapLight { + diffuse_map: diffuse_handle, + specular_map: specular_handle, + intensity: filtered_env_map.intensity, + rotation: filtered_env_map.rotation, + affects_lightmapped_mesh_diffuse: filtered_env_map.affects_lightmapped_mesh_diffuse, + }); + } +} diff --git a/crates/bevy_pbr/src/light_probe/irradiance_volume.rs b/crates/bevy_pbr/src/light_probe/irradiance_volume.rs index 05dd51c3795bc..479e05501c971 100644 --- a/crates/bevy_pbr/src/light_probe/irradiance_volume.rs +++ b/crates/bevy_pbr/src/light_probe/irradiance_volume.rs @@ -126,7 +126,7 @@ //! //! [Blender]: http://blender.org/ //! -//! [baking tool]: https://docs.blender.org/manual/en/latest/render/eevee/render_settings/indirect_lighting.html +//! [baking tool]: https://docs.blender.org/manual/en/latest/render/eevee/light_probes/volume.html //! //! [`bevy-baked-gi`]: https://github.com/pcwalton/bevy-baked-gi //! diff --git a/crates/bevy_pbr/src/light_probe/mod.rs b/crates/bevy_pbr/src/light_probe/mod.rs index ebfc7c7e7c614..81b23bda87ee2 100644 --- a/crates/bevy_pbr/src/light_probe/mod.rs +++ b/crates/bevy_pbr/src/light_probe/mod.rs @@ -1,8 +1,11 @@ //! Light probes for baked global illumination. -use bevy_app::{App, Plugin}; -use bevy_asset::{load_internal_asset, weak_handle, AssetId, Handle}; -use bevy_core_pipeline::core_3d::Camera3d; +use bevy_app::{App, Plugin, Update}; +use bevy_asset::{load_internal_asset, load_internal_binary_asset, weak_handle, AssetId, Handle}; +use bevy_core_pipeline::core_3d::{ + graph::{Core3d, Node3d}, + Camera3d, +}; use bevy_derive::{Deref, DerefMut}; use bevy_ecs::{ component::Component, @@ -18,9 +21,11 @@ use bevy_math::{Affine3A, FloatOrd, Mat4, Vec3A, Vec4}; use bevy_platform::collections::HashMap; use bevy_reflect::{std_traits::ReflectDefault, Reflect}; use bevy_render::{ + extract_component::ExtractComponentPlugin, extract_instances::ExtractInstancesPlugin, primitives::{Aabb, Frustum}, render_asset::RenderAssets, + render_graph::RenderGraphApp, render_resource::{DynamicUniformBuffer, Sampler, Shader, ShaderType, TextureView}, renderer::{RenderAdapter, RenderDevice, RenderQueue}, settings::WgpuFeatures, @@ -30,14 +35,23 @@ use bevy_render::{ Extract, ExtractSchedule, Render, RenderApp, RenderSet, }; use bevy_transform::{components::Transform, prelude::GlobalTransform}; +use generate::{ + extract_generator_entities, generate_environment_map_light, prepare_generator_bind_groups, + prepare_intermediate_textures, GeneratedEnvironmentMapLight, GeneratorPipelines, SpdNode, + STBN_SPHERE, STBN_VEC2, +}; use tracing::error; use core::{hash::Hash, ops::Deref}; use crate::{ irradiance_volume::IRRADIANCE_VOLUME_SHADER_HANDLE, - light_probe::environment_map::{ - EnvironmentMapIds, EnvironmentMapLight, ENVIRONMENT_MAP_SHADER_HANDLE, + light_probe::{ + environment_map::{EnvironmentMapIds, EnvironmentMapLight, ENVIRONMENT_MAP_SHADER_HANDLE}, + generate::{ + GeneratorBindGroupLayouts, GeneratorNode, GeneratorSamplers, IrradianceMapNode, + RadianceMapNode, ENVIRONMENT_FILTER_SHADER_HANDLE, SPD_SHADER_HANDLE, + }, }, }; @@ -47,6 +61,7 @@ pub const LIGHT_PROBE_SHADER_HANDLE: Handle = weak_handle!("e80a2ae6-1c5a-4d9a-a852-d66ff0e6bf7f"); pub mod environment_map; +pub mod generate; pub mod irradiance_volume; /// The maximum number of each type of light probe that each view will consider. @@ -362,10 +377,47 @@ impl Plugin for LightProbePlugin { "irradiance_volume.wgsl", Shader::from_wgsl ); + load_internal_asset!(app, SPD_SHADER_HANDLE, "spd.wgsl", Shader::from_wgsl); + load_internal_asset!( + app, + ENVIRONMENT_FILTER_SHADER_HANDLE, + "environment_filter.wgsl", + Shader::from_wgsl + ); + load_internal_binary_asset!( + app, + STBN_SPHERE, + "noise/sphere_coshemi_gauss1_0.png", + |bytes, _: String| Image::from_buffer( + bytes, + bevy_image::ImageType::Format(bevy_image::ImageFormat::Png), + bevy_image::CompressedImageFormats::NONE, + false, + bevy_image::ImageSampler::Default, + bevy_asset::RenderAssetUsages::RENDER_WORLD, + ) + .expect("Failed to load sphere cosine weighted blue noise texture") + ); + load_internal_binary_asset!( + app, + STBN_VEC2, + "noise/vector2_uniform_gauss1_0.png", + |bytes, _: String| Image::from_buffer( + bytes, + bevy_image::ImageType::Format(bevy_image::ImageFormat::Png), + bevy_image::CompressedImageFormats::NONE, + false, + bevy_image::ImageSampler::Default, + bevy_asset::RenderAssetUsages::RENDER_WORLD, + ) + .expect("Failed to load vector2 uniform blue noise texture") + ); app.register_type::() .register_type::() - .register_type::(); + .register_type::() + .add_plugins(ExtractComponentPlugin::::default()) + .add_systems(Update, generate_environment_map_light); } fn finish(&self, app: &mut App) { @@ -377,13 +429,40 @@ impl Plugin for LightProbePlugin { .add_plugins(ExtractInstancesPlugin::::new()) .init_resource::() .init_resource::() + .init_resource::() + .init_resource::() + .init_resource::() + .add_render_graph_node::(Core3d, GeneratorNode::Mipmap) + .add_render_graph_node::(Core3d, GeneratorNode::Radiance) + .add_render_graph_node::(Core3d, GeneratorNode::Irradiance) + .add_render_graph_edges( + Core3d, + ( + Node3d::EndPrepasses, + GeneratorNode::Mipmap, + GeneratorNode::Radiance, + GeneratorNode::Irradiance, + Node3d::StartMainPass, + ), + ) .add_systems(ExtractSchedule, gather_environment_map_uniform) .add_systems(ExtractSchedule, gather_light_probes::) .add_systems(ExtractSchedule, gather_light_probes::) + .add_systems( + ExtractSchedule, + extract_generator_entities.after(generate_environment_map_light), + ) .add_systems( Render, - (upload_light_probes, prepare_environment_uniform_buffer) - .in_set(RenderSet::PrepareResources), + ( + prepare_generator_bind_groups.in_set(RenderSet::PrepareBindGroups), + ( + upload_light_probes, + prepare_environment_uniform_buffer, + prepare_intermediate_textures, + ) + .in_set(RenderSet::PrepareResources), + ), ); } } diff --git a/crates/bevy_pbr/src/light_probe/noise/sphere_coshemi_gauss1_0.png b/crates/bevy_pbr/src/light_probe/noise/sphere_coshemi_gauss1_0.png new file mode 100644 index 0000000000000..08199965dc832 Binary files /dev/null and b/crates/bevy_pbr/src/light_probe/noise/sphere_coshemi_gauss1_0.png differ diff --git a/crates/bevy_pbr/src/light_probe/noise/vector2_uniform_gauss1_0.png b/crates/bevy_pbr/src/light_probe/noise/vector2_uniform_gauss1_0.png new file mode 100644 index 0000000000000..c9b61dcae5c2c Binary files /dev/null and b/crates/bevy_pbr/src/light_probe/noise/vector2_uniform_gauss1_0.png differ diff --git a/crates/bevy_pbr/src/light_probe/spd.wgsl b/crates/bevy_pbr/src/light_probe/spd.wgsl new file mode 100644 index 0000000000000..2fd63ade1d786 --- /dev/null +++ b/crates/bevy_pbr/src/light_probe/spd.wgsl @@ -0,0 +1,398 @@ +// Ported from https://github.com/GPUOpen-LibrariesAndSDKs/FidelityFX-SDK/blob/c16b1d286b5b438b75da159ab51ff426bacea3d1/sdk/include/FidelityFX/gpu/spd/ffx_spd.h + +@group(0) @binding(0) var mip_0: texture_2d_array; +@group(0) @binding(1) var mip_1: texture_storage_2d_array; +@group(0) @binding(2) var mip_2: texture_storage_2d_array; +@group(0) @binding(3) var mip_3: texture_storage_2d_array; +@group(0) @binding(4) var mip_4: texture_storage_2d_array; +@group(0) @binding(5) var mip_5: texture_storage_2d_array; +@group(0) @binding(6) var mip_6: texture_storage_2d_array; +@group(0) @binding(7) var mip_7: texture_storage_2d_array; +@group(0) @binding(8) var mip_8: texture_storage_2d_array; +@group(0) @binding(9) var mip_9: texture_storage_2d_array; +@group(0) @binding(10) var mip_10: texture_storage_2d_array; +@group(0) @binding(11) var mip_11: texture_storage_2d_array; +@group(0) @binding(12) var mip_12: texture_storage_2d_array; +@group(0) @binding(13) var sampler_linear_clamp: sampler; +@group(0) @binding(14) var constants: Constants; +struct Constants { mips: u32, inverse_input_size: vec2f } + +var spd_intermediate_r: array, 16>; +var spd_intermediate_g: array, 16>; +var spd_intermediate_b: array, 16>; +var spd_intermediate_a: array, 16>; + +@compute +@workgroup_size(256, 1, 1) +fn spd_downsample_first( + @builtin(workgroup_id) workgroup_id: vec3u, + @builtin(local_invocation_index) local_invocation_index: u32, +#ifdef SUBGROUP_SUPPORT + @builtin(subgroup_invocation_id) subgroup_invocation_id: u32, +#endif +) { +#ifndef SUBGROUP_SUPPORT + let subgroup_invocation_id = 0u; +#endif + + let sub_xy = remap_for_wave_reduction(local_invocation_index % 64u); + let x = sub_xy.x + 8u * ((local_invocation_index >> 6u) % 2u); + let y = sub_xy.y + 8u * (local_invocation_index >> 7u); + + spd_downsample_mips_0_1(x, y, workgroup_id.xy, local_invocation_index, constants.mips, workgroup_id.z, subgroup_invocation_id); + + spd_downsample_next_four(x, y, workgroup_id.xy, local_invocation_index, 2u, constants.mips, workgroup_id.z, subgroup_invocation_id); +} + +// TODO: Once wgpu supports globallycoherent buffers, make it actually a single pass +@compute +@workgroup_size(256, 1, 1) +fn spd_downsample_second( + @builtin(workgroup_id) workgroup_id: vec3u, + @builtin(local_invocation_index) local_invocation_index: u32, +#ifdef SUBGROUP_SUPPORT + @builtin(subgroup_invocation_id) subgroup_invocation_id: u32, +#endif +) { +#ifndef SUBGROUP_SUPPORT + let subgroup_invocation_id = 0u; +#endif + + let sub_xy = remap_for_wave_reduction(local_invocation_index % 64u); + let x = sub_xy.x + 8u * ((local_invocation_index >> 6u) % 2u); + let y = sub_xy.y + 8u * (local_invocation_index >> 7u); + + spd_downsample_mips_6_7(x, y, constants.mips, workgroup_id.z); + + spd_downsample_next_four(x, y, vec2(0u), local_invocation_index, 8u, constants.mips, workgroup_id.z, subgroup_invocation_id); +} + +fn spd_downsample_mips_0_1(x: u32, y: u32, workgroup_id: vec2u, local_invocation_index: u32, mips: u32, slice: u32, subgroup_invocation_id: u32) { + var v: array; + + var tex = (workgroup_id * 64u) + vec2(x * 2u, y * 2u); + var pix = (workgroup_id * 32u) + vec2(x, y); + v[0] = spd_reduce_load_source_image(tex, slice); + spd_store(pix, v[0], 0u, slice); + + tex = (workgroup_id * 64u) + vec2(x * 2u + 32u, y * 2u); + pix = (workgroup_id * 32u) + vec2(x + 16u, y); + v[1] = spd_reduce_load_source_image(tex, slice); + spd_store(pix, v[1], 0u, slice); + + tex = (workgroup_id * 64u) + vec2(x * 2u, y * 2u + 32u); + pix = (workgroup_id * 32u) + vec2(x, y + 16u); + v[2] = spd_reduce_load_source_image(tex, slice); + spd_store(pix, v[2], 0u, slice); + + tex = (workgroup_id * 64u) + vec2(x * 2u + 32u, y * 2u + 32u); + pix = (workgroup_id * 32u) + vec2(x + 16u, y + 16u); + v[3] = spd_reduce_load_source_image(tex, slice); + spd_store(pix, v[3], 0u, slice); + + if mips <= 1u { return; } + +#ifdef SUBGROUP_SUPPORT + v[0] = spd_reduce_quad(v[0], subgroup_invocation_id); + v[1] = spd_reduce_quad(v[1], subgroup_invocation_id); + v[2] = spd_reduce_quad(v[2], subgroup_invocation_id); + v[3] = spd_reduce_quad(v[3], subgroup_invocation_id); + + if local_invocation_index % 4u == 0u { + spd_store((workgroup_id * 16u) + vec2(x / 2u, y / 2u), v[0], 1u, slice); + spd_store_intermediate(x / 2u, y / 2u, v[0]); + + spd_store((workgroup_id * 16u) + vec2(x / 2u + 8u, y / 2u), v[1], 1u, slice); + spd_store_intermediate(x / 2u + 8u, y / 2u, v[1]); + + spd_store((workgroup_id * 16u) + vec2(x / 2u, y / 2u + 8u), v[2], 1u, slice); + spd_store_intermediate(x / 2u, y / 2u + 8u, v[2]); + + spd_store((workgroup_id * 16u) + vec2(x / 2u + 8u, y / 2u + 8u), v[3], 1u, slice); + spd_store_intermediate(x / 2u + 8u, y / 2u + 8u, v[3]); + } +#else + for (var i = 0u; i < 4u; i++) { + spd_store_intermediate(x, y, v[i]); + workgroupBarrier(); + if local_invocation_index < 64u { + v[i] = spd_reduce_intermediate( + vec2(x * 2u + 0u, y * 2u + 0u), + vec2(x * 2u + 1u, y * 2u + 0u), + vec2(x * 2u + 0u, y * 2u + 1u), + vec2(x * 2u + 1u, y * 2u + 1u), + ); + spd_store(vec2(workgroup_id * 16) + vec2(x + (i % 2u) * 8u, y + (i / 2u) * 8u), v[i], 1u, slice); + } + workgroupBarrier(); + } + + if local_invocation_index < 64u { + spd_store_intermediate(x + 0u, y + 0u, v[0]); + spd_store_intermediate(x + 8u, y + 0u, v[1]); + spd_store_intermediate(x + 0u, y + 8u, v[2]); + spd_store_intermediate(x + 8u, y + 8u, v[3]); + } +#endif +} + +fn spd_downsample_next_four(x: u32, y: u32, workgroup_id: vec2u, local_invocation_index: u32, base_mip: u32, mips: u32, slice: u32, subgroup_invocation_id: u32) { + if mips <= base_mip { return; } + workgroupBarrier(); + spd_downsample_mip_2(x, y, workgroup_id, local_invocation_index, base_mip, slice, subgroup_invocation_id); + + if mips <= base_mip + 1u { return; } + workgroupBarrier(); + spd_downsample_mip_3(x, y, workgroup_id, local_invocation_index, base_mip + 1u, slice, subgroup_invocation_id); + + if mips <= base_mip + 2u { return; } + workgroupBarrier(); + spd_downsample_mip_4(x, y, workgroup_id, local_invocation_index, base_mip + 2u, slice, subgroup_invocation_id); + + if mips <= base_mip + 3u { return; } + workgroupBarrier(); + spd_downsample_mip_5(x, y, workgroup_id, local_invocation_index, base_mip + 3u, slice, subgroup_invocation_id); +} + +fn spd_downsample_mip_2(x: u32, y: u32, workgroup_id: vec2u, local_invocation_index: u32, base_mip: u32, slice: u32, subgroup_invocation_id: u32) { +#ifdef SUBGROUP_SUPPORT + var v = spd_load_intermediate(x, y); + v = spd_reduce_quad(v, subgroup_invocation_id); + if local_invocation_index % 4u == 0u { + spd_store((workgroup_id * 8u) + vec2(x / 2u, y / 2u), v, base_mip, slice); + spd_store_intermediate(x + (y / 2u) % 2u, y, v); + } +#else + if local_invocation_index < 64u { + let v = spd_reduce_intermediate( + vec2(x * 2u + 0u, y * 2u + 0u), + vec2(x * 2u + 1u, y * 2u + 0u), + vec2(x * 2u + 0u, y * 2u + 1u), + vec2(x * 2u + 1u, y * 2u + 1u), + ); + spd_store((workgroup_id * 8u) + vec2(x, y), v, base_mip, slice); + spd_store_intermediate(x * 2u + y % 2u, y * 2u, v); + } +#endif +} + +fn spd_downsample_mip_3(x: u32, y: u32, workgroup_id: vec2u, local_invocation_index: u32, base_mip: u32, slice: u32, subgroup_invocation_id: u32) { +#ifdef SUBGROUP_SUPPORT + if local_invocation_index < 64u { + var v = spd_load_intermediate(x * 2u + y % 2u, y * 2u); + v = spd_reduce_quad(v, subgroup_invocation_id); + if local_invocation_index % 4u == 0u { + spd_store((workgroup_id * 4u) + vec2(x / 2u, y / 2u), v, base_mip, slice); + spd_store_intermediate(x * 2u + y / 2u, y * 2u, v); + } + } +#else + if local_invocation_index < 16u { + let v = spd_reduce_intermediate( + vec2(x * 4u + 0u + 0u, y * 4u + 0u), + vec2(x * 4u + 2u + 0u, y * 4u + 0u), + vec2(x * 4u + 0u + 1u, y * 4u + 2u), + vec2(x * 4u + 2u + 1u, y * 4u + 2u), + ); + spd_store((workgroup_id * 4u) + vec2(x, y), v, base_mip, slice); + spd_store_intermediate(x * 4u + y, y * 4u, v); + } +#endif +} + +fn spd_downsample_mip_4(x: u32, y: u32, workgroup_id: vec2u, local_invocation_index: u32, base_mip: u32, slice: u32, subgroup_invocation_id: u32) { +#ifdef SUBGROUP_SUPPORT + if local_invocation_index < 16u { + var v = spd_load_intermediate(x * 4u + y, y * 4u); + v = spd_reduce_quad(v, subgroup_invocation_id); + if local_invocation_index % 4u == 0u { + spd_store((workgroup_id * 2u) + vec2(x / 2u, y / 2u), v, base_mip, slice); + spd_store_intermediate(x / 2u + y, 0u, v); + } + } +#else + if local_invocation_index < 4u { + let v = spd_reduce_intermediate( + vec2(x * 8u + 0u + 0u + y * 2u, y * 8u + 0u), + vec2(x * 8u + 4u + 0u + y * 2u, y * 8u + 0u), + vec2(x * 8u + 0u + 1u + y * 2u, y * 8u + 4u), + vec2(x * 8u + 4u + 1u + y * 2u, y * 8u + 4u), + ); + spd_store((workgroup_id * 2u) + vec2(x, y), v, base_mip, slice); + spd_store_intermediate(x + y * 2u, 0u, v); + } +#endif +} + +fn spd_downsample_mip_5(x: u32, y: u32, workgroup_id: vec2u, local_invocation_index: u32, base_mip: u32, slice: u32, subgroup_invocation_id: u32) { +#ifdef SUBGROUP_SUPPORT + if local_invocation_index < 4u { + var v = spd_load_intermediate(local_invocation_index, 0u); + v = spd_reduce_quad(v, subgroup_invocation_id); + if local_invocation_index % 4u == 0u { + spd_store(workgroup_id, v, base_mip, slice); + } + } +#else + if local_invocation_index < 1u { + let v = spd_reduce_intermediate(vec2(0u, 0u), vec2(1u, 0u), vec2(2u, 0u), vec2(3u, 0u)); + spd_store(workgroup_id, v, base_mip, slice); + } +#endif +} + +fn spd_downsample_mips_6_7(x: u32, y: u32, mips: u32, slice: u32) { + var tex = vec2(x * 4u + 0u, y * 4u + 0u); + var pix = vec2(x * 2u + 0u, y * 2u + 0u); + let v0 = spd_reduce_load_4( + vec2(x * 4u + 0u, y * 4u + 0u), + vec2(x * 4u + 1u, y * 4u + 0u), + vec2(x * 4u + 0u, y * 4u + 1u), + vec2(x * 4u + 1u, y * 4u + 1u), + slice + ); + spd_store(pix, v0, 6u, slice); + + tex = vec2(x * 4u + 2u, y * 4u + 0u); + pix = vec2(x * 2u + 1u, y * 2u + 0u); + let v1 = spd_reduce_load_4( + vec2(x * 4u + 2u, y * 4u + 0u), + vec2(x * 4u + 3u, y * 4u + 0u), + vec2(x * 4u + 2u, y * 4u + 1u), + vec2(x * 4u + 3u, y * 4u + 1u), + slice + ); + spd_store(pix, v1, 6u, slice); + + tex = vec2(x * 4u + 0u, y * 4u + 2u); + pix = vec2(x * 2u + 0u, y * 2u + 1u); + let v2 = spd_reduce_load_4( + vec2(x * 4u + 0u, y * 4u + 2u), + vec2(x * 4u + 1u, y * 4u + 2u), + vec2(x * 4u + 0u, y * 4u + 3u), + vec2(x * 4u + 1u, y * 4u + 3u), + slice + ); + spd_store(pix, v2, 6u, slice); + + tex = vec2(x * 4u + 2u, y * 4u + 2u); + pix = vec2(x * 2u + 1u, y * 2u + 1u); + let v3 = spd_reduce_load_4( + vec2(x * 4u + 2u, y * 4u + 2u), + vec2(x * 4u + 3u, y * 4u + 2u), + vec2(x * 4u + 2u, y * 4u + 3u), + vec2(x * 4u + 3u, y * 4u + 3u), + slice + ); + spd_store(pix, v3, 6u, slice); + + if mips < 7u { return; } + + let v = spd_reduce_4(v0, v1, v2, v3); + spd_store(vec2(x, y), v, 7u, slice); + spd_store_intermediate(x, y, v); +} + +fn remap_for_wave_reduction(a: u32) -> vec2u { + // This function maps linear thread IDs to 2D coordinates in a special pattern + // to ensure that neighboring threads process neighboring pixels + // For example, this transforms linear thread IDs 0,1,2,3 into a 2×2 square + + // Extract bits to form the X and Y coordinates + let x = insertBits(extractBits(a, 2u, 3u), a, 0u, 1u); + let y = insertBits(extractBits(a, 3u, 3u), extractBits(a, 1u, 2u), 0u, 2u); + + return vec2u(x, y); +} + +fn spd_reduce_load_source_image(uv: vec2u, slice: u32) -> vec4f { + let texture_coord = (vec2f(uv) + 0.5) * constants.inverse_input_size; + + let result = textureSampleLevel(mip_0, sampler_linear_clamp, texture_coord, slice, 0.0); + +#ifdef SRGB_CONVERSION + return vec4( + srgb_from_linear(result.r), + srgb_from_linear(result.g), + srgb_from_linear(result.b), + result.a + ); +#else + return result; +#endif + +} + +fn spd_store(pix: vec2u, value: vec4f, mip: u32, slice: u32) { + if mip >= constants.mips { return; } + switch mip { + case 0u: { textureStore(mip_1, pix, slice, value); } + case 1u: { textureStore(mip_2, pix, slice, value); } + case 2u: { textureStore(mip_3, pix, slice, value); } + case 3u: { textureStore(mip_4, pix, slice, value); } + case 4u: { textureStore(mip_5, pix, slice, value); } + case 5u: { textureStore(mip_6, pix, slice, value); } + case 6u: { textureStore(mip_7, pix, slice, value); } + case 7u: { textureStore(mip_8, pix, slice, value); } + // case 8u: { textureStore(mip_9, pix, slice, value); } + // case 9u: { textureStore(mip_10, pix, slice, value); } + // case 10u: { textureStore(mip_11, pix, slice, value); } + // case 11u: { textureStore(mip_12, pix, slice, value); } + default: {} + } +} + +fn spd_store_intermediate(x: u32, y: u32, value: vec4f) { + spd_intermediate_r[x][y] = value.x; + spd_intermediate_g[x][y] = value.y; + spd_intermediate_b[x][y] = value.z; + spd_intermediate_a[x][y] = value.w; +} + +fn spd_load_intermediate(x: u32, y: u32) -> vec4f { + return vec4(spd_intermediate_r[x][y], spd_intermediate_g[x][y], spd_intermediate_b[x][y], spd_intermediate_a[x][y]); +} + +fn spd_reduce_intermediate(i0: vec2u, i1: vec2u, i2: vec2u, i3: vec2u) -> vec4f { + let v0 = spd_load_intermediate(i0.x, i0.y); + let v1 = spd_load_intermediate(i1.x, i1.y); + let v2 = spd_load_intermediate(i2.x, i2.y); + let v3 = spd_load_intermediate(i3.x, i3.y); + return spd_reduce_4(v0, v1, v2, v3); +} + +fn spd_reduce_load_4(i0: vec2u, i1: vec2u, i2: vec2u, i3: vec2u, slice: u32) -> vec4f { + let v0 = textureLoad(mip_6, i0, slice); + let v1 = textureLoad(mip_6, i1, slice); + let v2 = textureLoad(mip_6, i2, slice); + let v3 = textureLoad(mip_6, i3, slice); + return spd_reduce_4(v0, v1, v2, v3); +} + +fn spd_reduce_4(v0: vec4f, v1: vec4f, v2: vec4f, v3: vec4f) -> vec4f { + return (v0 + v1 + v2 + v3) * 0.25; +} + +#ifdef SUBGROUP_SUPPORT +fn spd_reduce_quad(v: vec4f, subgroup_invocation_id: u32) -> vec4f { + let quad = subgroup_invocation_id & (~0x3u); + let v0 = v; + let v1 = subgroupBroadcast(v, quad | 1u); + let v2 = subgroupBroadcast(v, quad | 2u); + let v3 = subgroupBroadcast(v, quad | 3u); + return spd_reduce_4(v0, v1, v2, v3); + + // TODO: Use subgroup quad operations once wgpu supports them + // let v0 = v; + // let v1 = quadSwapX(v); + // let v2 = quadSwapY(v); + // let v3 = quadSwapDiagonal(v); + // return spd_reduce_4(v0, v1, v2, v3); +} +#endif + +fn srgb_from_linear(value: f32) -> f32 { + let j = vec3(0.0031308 * 12.92, 12.92, 1.0 / 2.4); + let k = vec2(1.055, -0.055); + return clamp(j.x, value * j.y, pow(value, j.z) * k.x + k.y); +} \ No newline at end of file diff --git a/examples/3d/reflection_probes.rs b/examples/3d/reflection_probes.rs index e2e3ea60b4f99..b66932dd2a140 100644 --- a/examples/3d/reflection_probes.rs +++ b/examples/3d/reflection_probes.rs @@ -6,7 +6,12 @@ //! //! Reflection probes don't work on WebGL 2 or WebGPU. -use bevy::{core_pipeline::Skybox, prelude::*}; +use bevy::{ + core_pipeline::Skybox, + pbr::generate::GeneratedEnvironmentMapLight, + prelude::*, + render::{camera::Exposure, render_resource::TextureUsages}, +}; use std::{ f32::consts::PI, @@ -25,6 +30,8 @@ struct AppStatus { reflection_mode: ReflectionMode, // Whether the user has requested the scene to rotate. rotating: bool, + // The current roughness of the central sphere + sphere_roughness: f32, } // Which environment maps the user has requested to display. @@ -37,6 +44,8 @@ enum ReflectionMode { // Both a world environment map and a reflection probe are present. The // reflection probe is shown in the sphere. ReflectionProbe = 2, + // A generated environment map is shown. + GeneratedEnvironmentMap = 3, } // The various reflection maps. @@ -53,6 +62,9 @@ struct Cubemaps { // The specular cubemap that reflects both the world and the cubes. specular_reflection_probe: Handle, + // Unfiltered environment map + unfiltered_environment_map: Handle, + // The skybox cubemap image. This is almost the same as // `specular_environment_map`. skybox: Handle, @@ -68,6 +80,7 @@ fn main() { .add_systems(PreUpdate, add_environment_map_to_camera) .add_systems(Update, change_reflection_type) .add_systems(Update, toggle_rotation) + .add_systems(Update, change_sphere_roughness) .add_systems( Update, rotate_camera @@ -75,6 +88,7 @@ fn main() { .after(change_reflection_type), ) .add_systems(Update, update_text.after(rotate_camera)) + .add_systems(Update, setup_environment_map_usage) .run(); } @@ -89,7 +103,7 @@ fn setup( ) { spawn_scene(&mut commands, &asset_server); spawn_camera(&mut commands); - spawn_sphere(&mut commands, &mut meshes, &mut materials); + spawn_sphere(&mut commands, &mut meshes, &mut materials, &app_status); spawn_reflection_probe(&mut commands, &cubemaps); spawn_text(&mut commands, &app_status); } @@ -109,6 +123,7 @@ fn spawn_camera(commands: &mut Commands) { hdr: true, ..default() }, + Exposure::OVERCAST, Transform::from_xyz(-6.483, 0.325, 4.381).looking_at(Vec3::ZERO, Vec3::Y), )); } @@ -118,6 +133,7 @@ fn spawn_sphere( commands: &mut Commands, meshes: &mut Assets, materials: &mut Assets, + app_status: &AppStatus, ) { // Create a sphere mesh. let sphere_mesh = meshes.add(Sphere::new(1.0).mesh().ico(7).unwrap()); @@ -126,11 +142,12 @@ fn spawn_sphere( commands.spawn(( Mesh3d(sphere_mesh.clone()), MeshMaterial3d(materials.add(StandardMaterial { - base_color: Srgba::hex("#ffd891").unwrap().into(), + base_color: Srgba::hex("#ffffff").unwrap().into(), metallic: 1.0, - perceptual_roughness: 0.0, + perceptual_roughness: app_status.sphere_roughness, ..StandardMaterial::default() })), + SphereMaterial, )); } @@ -149,6 +166,18 @@ fn spawn_reflection_probe(commands: &mut Commands, cubemaps: &Cubemaps) { )); } +fn spawn_generated_environment_map(commands: &mut Commands, cubemaps: &Cubemaps) { + commands.spawn(( + LightProbe, + GeneratedEnvironmentMapLight { + environment_map: cubemaps.unfiltered_environment_map.clone(), + intensity: 5000.0, + ..default() + }, + Transform::from_scale(Vec3::splat(2.0)), + )); +} + // Spawns the help text. fn spawn_text(commands: &mut Commands, app_status: &AppStatus) { // Create the text. @@ -187,6 +216,7 @@ fn add_environment_map_to_camera( fn change_reflection_type( mut commands: Commands, light_probe_query: Query>, + sky_box_query: Query>, camera_query: Query>, keyboard: Res>, mut app_status: ResMut, @@ -199,15 +229,21 @@ fn change_reflection_type( // Switch reflection mode. app_status.reflection_mode = - ReflectionMode::try_from((app_status.reflection_mode as u32 + 1) % 3).unwrap(); + ReflectionMode::try_from((app_status.reflection_mode as u32 + 1) % 4).unwrap(); // Add or remove the light probe. for light_probe in light_probe_query.iter() { commands.entity(light_probe).despawn(); } + for skybox in sky_box_query.iter() { + commands.entity(skybox).remove::(); + } match app_status.reflection_mode { ReflectionMode::None | ReflectionMode::EnvironmentMap => {} ReflectionMode::ReflectionProbe => spawn_reflection_probe(&mut commands, &cubemaps), + ReflectionMode::GeneratedEnvironmentMap => { + spawn_generated_environment_map(&mut commands, &cubemaps); + } } // Add or remove the environment map from the camera. @@ -216,10 +252,23 @@ fn change_reflection_type( ReflectionMode::None => { commands.entity(camera).remove::(); } - ReflectionMode::EnvironmentMap | ReflectionMode::ReflectionProbe => { + ReflectionMode::EnvironmentMap + | ReflectionMode::ReflectionProbe + | ReflectionMode::GeneratedEnvironmentMap => { + let image = match app_status.reflection_mode { + ReflectionMode::GeneratedEnvironmentMap => { + cubemaps.unfiltered_environment_map.clone() + } + _ => cubemaps.skybox.clone(), + }; commands .entity(camera) - .insert(create_camera_environment_map_light(&cubemaps)); + .insert(create_camera_environment_map_light(&cubemaps)) + .insert(Skybox { + image, + brightness: 5000.0, + ..default() + }); } } } @@ -247,6 +296,7 @@ impl TryFrom for ReflectionMode { 0 => Ok(ReflectionMode::None), 1 => Ok(ReflectionMode::EnvironmentMap), 2 => Ok(ReflectionMode::ReflectionProbe), + 3 => Ok(ReflectionMode::GeneratedEnvironmentMap), _ => Err(()), } } @@ -258,6 +308,7 @@ impl Display for ReflectionMode { ReflectionMode::None => "No reflections", ReflectionMode::EnvironmentMap => "Environment map", ReflectionMode::ReflectionProbe => "Reflection probe", + ReflectionMode::GeneratedEnvironmentMap => "Generated environment map", }; formatter.write_str(text) } @@ -274,8 +325,11 @@ impl AppStatus { }; format!( - "{}\n{}\n{}", - self.reflection_mode, rotation_help_text, REFLECTION_MODE_HELP_TEXT + "{}\n{}\nRoughness: {:.2}\n{}\nUp/Down arrows to change roughness", + self.reflection_mode, + rotation_help_text, + self.sphere_roughness, + REFLECTION_MODE_HELP_TEXT ) .into() } @@ -324,16 +378,62 @@ impl FromWorld for Cubemaps { specular_reflection_probe: world .load_asset("environment_maps/cubes_reflection_probe_specular_rgb9e5_zstd.ktx2"), specular_environment_map: specular_map.clone(), + unfiltered_environment_map: world.load_asset("environment_maps/goegap_road_2k.ktx2"), skybox: specular_map, } } } +fn setup_environment_map_usage(cubemaps: Res, mut images: ResMut>) { + if let Some(image) = images.get_mut(&cubemaps.unfiltered_environment_map) { + if !image + .texture_descriptor + .usage + .contains(TextureUsages::COPY_SRC) + { + image.texture_descriptor.usage |= TextureUsages::COPY_SRC; + } + } +} + impl Default for AppStatus { fn default() -> Self { Self { reflection_mode: ReflectionMode::ReflectionProbe, rotating: true, + sphere_roughness: 0.0, + } + } +} + +#[derive(Component)] +struct SphereMaterial; + +// A system that changes the sphere's roughness with up/down arrow keys +fn change_sphere_roughness( + keyboard: Res>, + mut app_status: ResMut, + mut materials: ResMut>, + sphere_query: Query<&MeshMaterial3d, With>, +) { + let roughness_delta = if keyboard.pressed(KeyCode::ArrowUp) { + 0.01 // Decrease roughness + } else if keyboard.pressed(KeyCode::ArrowDown) { + -0.01 // Increase roughness + } else { + 0.0 // No change + }; + + if roughness_delta != 0.0 { + // Update the app status + app_status.sphere_roughness = + (app_status.sphere_roughness + roughness_delta).clamp(0.0, 1.0); + + // Update the sphere material + for material_handle in sphere_query.iter() { + if let Some(material) = materials.get_mut(&material_handle.0) { + material.perceptual_roughness = app_status.sphere_roughness; + } } } }