diff --git a/crates/bevy_solari/src/pathtracer/pathtracer.wgsl b/crates/bevy_solari/src/pathtracer/pathtracer.wgsl index 45c5e519594f0..da216da959201 100644 --- a/crates/bevy_solari/src/pathtracer/pathtracer.wgsl +++ b/crates/bevy_solari/src/pathtracer/pathtracer.wgsl @@ -1,9 +1,11 @@ #import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance +#import bevy_pbr::pbr_functions::calculate_tbn_mikktspace #import bevy_pbr::utils::{rand_f, rand_vec2f, sample_cosine_hemisphere} #import bevy_render::maths::PI #import bevy_render::view::View -#import bevy_solari::sampling::sample_random_light -#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX} +#import bevy_solari::brdf::evaluate_brdf +#import bevy_solari::sampling::{sample_random_light, random_light_pdf, sample_ggx_vndf, ggx_vndf_pdf, balance_heuristic, power_heuristic} +#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, ResolvedRayHitFull, RAY_T_MIN, RAY_T_MAX} @group(1) @binding(0) var accumulation_texture: texture_storage_2d; @group(1) @binding(1) var view_output: texture_storage_2d; @@ -35,32 +37,45 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3) { // Path trace var radiance = vec3(0.0); var throughput = vec3(1.0); + var p_bounce = 0.0; + var bounce_was_perfect_reflection = true; + var previous_normal = vec3(0.0); loop { let ray_hit = trace_ray(ray_origin, ray_direction, ray_t_min, RAY_T_MAX, RAY_FLAG_NONE); if ray_hit.kind != RAY_QUERY_INTERSECTION_NONE { let ray_hit = resolve_ray_hit_full(ray_hit); + let wo = -ray_direction; - // Evaluate material BRDF - let diffuse_brdf = ray_hit.material.base_color / PI; + var mis_weight = 1.0; + if !bounce_was_perfect_reflection { + let p_light = random_light_pdf(ray_hit); + mis_weight = power_heuristic(p_bounce, p_light); + } + radiance += mis_weight * throughput * ray_hit.material.emissive; - // Use emissive only on the first ray (coming from the camera) - if ray_t_min == 0.0 { radiance = ray_hit.material.emissive; } - - // Sample direct lighting - let direct_lighting = sample_random_light(ray_hit.world_position, ray_hit.world_normal, &rng); - radiance += throughput * diffuse_brdf * direct_lighting.radiance * direct_lighting.inverse_pdf; + // Sample direct lighting, but only if the surface is not mirror-like + let is_perfectly_specular = ray_hit.material.roughness < 0.0001 && ray_hit.material.metallic > 0.9999; + if !is_perfectly_specular { + let direct_lighting = sample_random_light(ray_hit.world_position, ray_hit.world_normal, &rng); + let pdf_of_bounce = brdf_pdf(wo, direct_lighting.wi, ray_hit); + mis_weight = power_heuristic(1.0 / direct_lighting.inverse_pdf, pdf_of_bounce); + let direct_lighting_brdf = evaluate_brdf(ray_hit.world_normal, wo, direct_lighting.wi, ray_hit.material); + radiance += mis_weight * throughput * direct_lighting.radiance * direct_lighting.inverse_pdf * direct_lighting_brdf; + } // Sample new ray direction from the material BRDF for next bounce - ray_direction = sample_cosine_hemisphere(ray_hit.world_normal, &rng); - - // Update other variables for next bounce + let next_bounce = importance_sample_next_bounce(wo, ray_hit, &rng); + ray_direction = next_bounce.wi; ray_origin = ray_hit.world_position; ray_t_min = RAY_T_MIN; + p_bounce = next_bounce.pdf; + bounce_was_perfect_reflection = next_bounce.perfectly_specular_bounce; + previous_normal = ray_hit.world_normal; // Update throughput for next bounce - let cos_theta = dot(-ray_direction, ray_hit.world_normal); - let cosine_hemisphere_pdf = cos_theta / PI; // Weight for the next bounce because we importance sampled the diffuse BRDF for the next ray direction - throughput *= (diffuse_brdf * cos_theta) / cosine_hemisphere_pdf; + let brdf = evaluate_brdf(ray_hit.world_normal, wo, next_bounce.wi, ray_hit.material); + let cos_theta = dot(next_bounce.wi, ray_hit.world_normal); + throughput *= (brdf * cos_theta) / next_bounce.pdf; // Russian roulette for early termination let p = luminance(throughput); @@ -77,3 +92,60 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3) { textureStore(accumulation_texture, global_id.xy, vec4(new_color, old_color.a + 1.0)); textureStore(view_output, global_id.xy, vec4(new_color, 1.0)); } + +struct NextBounce { + wi: vec3, + pdf: f32, + perfectly_specular_bounce: bool, +} + +fn importance_sample_next_bounce(wo: vec3, ray_hit: ResolvedRayHitFull, rng: ptr) -> NextBounce { + let is_perfectly_specular = ray_hit.material.roughness < 0.0001 && ray_hit.material.metallic > 0.9999; + if is_perfectly_specular { + return NextBounce(reflect(-wo, ray_hit.world_normal), 1.0, true); + } + let diffuse_weight = mix(mix(0.4f, 0.9f, ray_hit.material.perceptual_roughness), 0.f, ray_hit.material.metallic); + let specular_weight = 1.0 - diffuse_weight; + + let TBN = calculate_tbn_mikktspace(ray_hit.world_normal, ray_hit.world_tangent); + let T = TBN[0]; + let B = TBN[1]; + let N = TBN[2]; + + let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N)); + + var wi: vec3; + var wi_tangent: vec3; + let diffuse_selected = rand_f(rng) < diffuse_weight; + if diffuse_selected { + wi = sample_cosine_hemisphere(ray_hit.world_normal, rng); + wi_tangent = vec3(dot(wi, T), dot(wi, B), dot(wi, N)); + } else { + wi_tangent = sample_ggx_vndf(wo_tangent, ray_hit.material.roughness, rng); + wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N; + } + + let diffuse_pdf = dot(wi, ray_hit.world_normal) / PI; + let specular_pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness); + let pdf = (diffuse_weight * diffuse_pdf) + (specular_weight * specular_pdf); + + return NextBounce(wi, pdf, false); +} + +fn brdf_pdf(wo: vec3, wi: vec3, ray_hit: ResolvedRayHitFull) -> f32 { + let diffuse_weight = mix(mix(0.4f, 0.9f, ray_hit.material.roughness), 0.f, ray_hit.material.metallic); + let specular_weight = 1.0 - diffuse_weight; + + let TBN = calculate_tbn_mikktspace(ray_hit.world_normal, ray_hit.world_tangent); + let T = TBN[0]; + let B = TBN[1]; + let N = TBN[2]; + + let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N)); + let wi_tangent = vec3(dot(wi, T), dot(wi, B), dot(wi, N)); + + let diffuse_pdf = wi_tangent.z / PI; + let specular_pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness); + let pdf = (diffuse_weight * diffuse_pdf) + (specular_weight * specular_pdf); + return pdf; +} diff --git a/crates/bevy_solari/src/scene/binder.rs b/crates/bevy_solari/src/scene/binder.rs index 0685c1dcf1acf..81923592cde18 100644 --- a/crates/bevy_solari/src/scene/binder.rs +++ b/crates/bevy_solari/src/scene/binder.rs @@ -115,13 +115,23 @@ pub fn prepare_raytracing_scene_bindings( let Some(emissive_texture_id) = process_texture(&material.emissive_texture) else { continue; }; + let Some(metallic_roughness_texture_id) = + process_texture(&material.metallic_roughness_texture) + else { + continue; + }; materials.get_mut().push(GpuMaterial { - base_color: material.base_color.to_linear(), - emissive: material.emissive, - base_color_texture_id, normal_map_texture_id, + base_color_texture_id, emissive_texture_id, + metallic_roughness_texture_id, + + base_color: LinearRgba::from(material.base_color).to_vec3(), + perceptual_roughness: material.perceptual_roughness, + emissive: material.emissive.to_vec3(), + metallic: material.metallic, + reflectance: LinearRgba::from(material.specular_tint).to_vec3() * material.reflectance, _padding: Default::default(), }); @@ -180,11 +190,12 @@ pub fn prepare_raytracing_scene_bindings( vertex_buffer_offset: vertex_slice.range.start, index_buffer_id, index_buffer_offset: index_slice.range.start, + triangle_count: (index_slice.range.len() / 3) as u32, }); material_ids.get_mut().push(material_id); - if material.emissive != LinearRgba::BLACK { + if material.emissive != Vec3::ZERO { light_sources .get_mut() .push(GpuLightSource::new_emissive_mesh_light( @@ -342,16 +353,22 @@ struct GpuInstanceGeometryIds { vertex_buffer_offset: u32, index_buffer_id: u32, index_buffer_offset: u32, + triangle_count: u32, } #[derive(ShaderType)] struct GpuMaterial { - base_color: LinearRgba, - emissive: LinearRgba, - base_color_texture_id: u32, normal_map_texture_id: u32, + base_color_texture_id: u32, emissive_texture_id: u32, - _padding: u32, + metallic_roughness_texture_id: u32, + + base_color: Vec3, + perceptual_roughness: f32, + emissive: Vec3, + metallic: f32, + reflectance: Vec3, + _padding: f32, } #[derive(ShaderType)] diff --git a/crates/bevy_solari/src/scene/brdf.wgsl b/crates/bevy_solari/src/scene/brdf.wgsl new file mode 100644 index 0000000000000..bc42203481928 --- /dev/null +++ b/crates/bevy_solari/src/scene/brdf.wgsl @@ -0,0 +1,56 @@ +#define_import_path bevy_solari::brdf + +#import bevy_pbr::lighting::{F_AB, D_GGX, V_SmithGGXCorrelated, fresnel, specular_multiscatter} +#import bevy_pbr::pbr_functions::{calculate_diffuse_color, calculate_F0} +#import bevy_render::maths::PI +#import bevy_solari::scene_bindings::ResolvedMaterial + +fn evaluate_brdf( + world_normal: vec3, + wo: vec3, + wi: vec3, + material: ResolvedMaterial, +) -> vec3 { + let diffuse_brdf = diffuse_brdf(material.base_color, material.metallic); + let specular_brdf = specular_brdf( + world_normal, + wo, + wi, + material.base_color, + material.metallic, + material.reflectance, + material.perceptual_roughness, + material.roughness, + ); + return diffuse_brdf + specular_brdf; +} + +fn diffuse_brdf(base_color: vec3, metallic: f32) -> vec3 { + let diffuse_color = calculate_diffuse_color(base_color, metallic, 0.0, 0.0); + return diffuse_color / PI; +} + +fn specular_brdf( + N: vec3, + V: vec3, + L: vec3, + base_color: vec3, + metallic: f32, + reflectance: vec3, + perceptual_roughness: f32, + roughness: f32, +) -> vec3 { + let H = normalize(L + V); + let NdotL = saturate(dot(N, L)); + let NdotH = saturate(dot(N, H)); + let LdotH = saturate(dot(L, H)); + let NdotV = max(dot(N, V), 0.0001); + + let F0 = calculate_F0(base_color, metallic, reflectance); + let F_ab = F_AB(perceptual_roughness, NdotV); + + let D = D_GGX(roughness, NdotH); + let Vs = V_SmithGGXCorrelated(roughness, NdotV, NdotL); + let F = fresnel(F0, LdotH); + return specular_multiscatter(D, Vs, F, F0, F_ab, 1.0); +} diff --git a/crates/bevy_solari/src/scene/mod.rs b/crates/bevy_solari/src/scene/mod.rs index f1af566c1c8ab..b0aa7e512bf71 100644 --- a/crates/bevy_solari/src/scene/mod.rs +++ b/crates/bevy_solari/src/scene/mod.rs @@ -31,6 +31,7 @@ pub struct RaytracingScenePlugin; impl Plugin for RaytracingScenePlugin { fn build(&self, app: &mut App) { + load_shader_library!(app, "brdf.wgsl"); load_shader_library!(app, "raytracing_scene_bindings.wgsl"); load_shader_library!(app, "sampling.wgsl"); diff --git a/crates/bevy_solari/src/scene/raytracing_scene_bindings.wgsl b/crates/bevy_solari/src/scene/raytracing_scene_bindings.wgsl index eeed96ad8e818..7359ad9063e2d 100644 --- a/crates/bevy_solari/src/scene/raytracing_scene_bindings.wgsl +++ b/crates/bevy_solari/src/scene/raytracing_scene_bindings.wgsl @@ -1,10 +1,14 @@ #define_import_path bevy_solari::scene_bindings +#import bevy_pbr::lighting::perceptualRoughnessToRoughness +#import bevy_pbr::pbr_functions::calculate_tbn_mikktspace + struct InstanceGeometryIds { vertex_buffer_id: u32, vertex_buffer_offset: u32, index_buffer_id: u32, index_buffer_offset: u32, + triangle_count: u32, } struct VertexBuffer { vertices: array } @@ -34,12 +38,17 @@ fn unpack_vertex(packed: PackedVertex) -> Vertex { } struct Material { - base_color: vec4, - emissive: vec4, - base_color_texture_id: u32, normal_map_texture_id: u32, + base_color_texture_id: u32, emissive_texture_id: u32, - _padding: u32, + metallic_roughness_texture_id: u32, + + base_color: vec3, + perceptual_roughness: f32, + emissive: vec3, + metallic: f32, + reflectance: vec3, + _padding: f32, } const TEXTURE_MAP_NONE = 0xFFFFFFFFu; @@ -94,14 +103,20 @@ fn sample_texture(id: u32, uv: vec2) -> vec3 { struct ResolvedMaterial { base_color: vec3, emissive: vec3, + reflectance: vec3, + perceptual_roughness: f32, + roughness: f32, + metallic: f32, } struct ResolvedRayHitFull { world_position: vec3, world_normal: vec3, geometric_world_normal: vec3, + world_tangent: vec4, uv: vec2, triangle_area: f32, + triangle_count: u32, material: ResolvedMaterial, } @@ -118,6 +133,17 @@ fn resolve_material(material: Material, uv: vec2) -> ResolvedMaterial { m.emissive *= sample_texture(material.emissive_texture_id, uv); } + m.reflectance = material.reflectance; + + m.perceptual_roughness = material.perceptual_roughness; + m.metallic = material.metallic; + if material.metallic_roughness_texture_id != TEXTURE_MAP_NONE { + let metallic_roughness = sample_texture(material.metallic_roughness_texture_id, uv); + m.perceptual_roughness *= metallic_roughness.g; + m.metallic *= metallic_roughness.b; + } + m.roughness = clamp(m.perceptual_roughness * m.perceptual_roughness, 0.001, 1.0); + return m; } @@ -144,15 +170,20 @@ fn resolve_triangle_data_full(instance_id: u32, triangle_id: u32, barycentrics: let uv = mat3x2(vertices[0].uv, vertices[1].uv, vertices[2].uv) * barycentrics; + let local_tangent = mat3x3(vertices[0].tangent.xyz, vertices[1].tangent.xyz, vertices[2].tangent.xyz) * barycentrics; + let world_tangent = vec4( + normalize(mat3x3(transform[0].xyz, transform[1].xyz, transform[2].xyz) * local_tangent), + vertices[0].tangent.w, + ); + let local_normal = mat3x3(vertices[0].normal, vertices[1].normal, vertices[2].normal) * barycentrics; // TODO: Use barycentric lerp, ray_hit.object_to_world, cross product geo normal var world_normal = normalize(mat3x3(transform[0].xyz, transform[1].xyz, transform[2].xyz) * local_normal); let geometric_world_normal = world_normal; if material.normal_map_texture_id != TEXTURE_MAP_NONE { - let local_tangent = mat3x3(vertices[0].tangent.xyz, vertices[1].tangent.xyz, vertices[2].tangent.xyz) * barycentrics; - let world_tangent = normalize(mat3x3(transform[0].xyz, transform[1].xyz, transform[2].xyz) * local_tangent); - let N = world_normal; - let T = world_tangent; - let B = vertices[0].tangent.w * cross(N, T); + let TBN = calculate_tbn_mikktspace(world_normal, world_tangent); + let T = TBN[0]; + let B = TBN[1]; + let N = TBN[2]; let Nt = sample_texture(material.normal_map_texture_id, uv); world_normal = normalize(Nt.x * T + Nt.y * B + Nt.z * N); } @@ -163,5 +194,5 @@ fn resolve_triangle_data_full(instance_id: u32, triangle_id: u32, barycentrics: let resolved_material = resolve_material(material, uv); - return ResolvedRayHitFull(world_position, world_normal, geometric_world_normal, uv, triangle_area, resolved_material); + return ResolvedRayHitFull(world_position, world_normal, geometric_world_normal, world_tangent, uv, triangle_area, instance_geometry_ids.triangle_count, resolved_material); } diff --git a/crates/bevy_solari/src/scene/sampling.wgsl b/crates/bevy_solari/src/scene/sampling.wgsl index 298cf8ad679a0..8385f9b3999e6 100644 --- a/crates/bevy_solari/src/scene/sampling.wgsl +++ b/crates/bevy_solari/src/scene/sampling.wgsl @@ -1,8 +1,57 @@ #define_import_path bevy_solari::sampling +#import bevy_pbr::lighting::D_GGX #import bevy_pbr::utils::{rand_f, rand_vec2f, rand_u, rand_range_u} -#import bevy_render::maths::PI_2 -#import bevy_solari::scene_bindings::{trace_ray, RAY_T_MIN, RAY_T_MAX, light_sources, directional_lights, LightSource, LIGHT_SOURCE_KIND_DIRECTIONAL, resolve_triangle_data_full} +#import bevy_render::maths::{PI_2, orthonormalize} +#import bevy_solari::scene_bindings::{trace_ray, RAY_T_MIN, RAY_T_MAX, light_sources, directional_lights, LightSource, LIGHT_SOURCE_KIND_DIRECTIONAL, resolve_triangle_data_full, ResolvedRayHitFull} + +fn power_heuristic(f: f32, g: f32) -> f32 { + return f * f / (f * f + g * g); +} + +fn balance_heuristic(f: f32, g: f32) -> f32 { + return f / (f + g); +} + +// https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf (Listing 1) +fn sample_ggx_vndf(wi_tangent: vec3, roughness: f32, rng: ptr) -> vec3 { + let i = wi_tangent; + let rand = rand_vec2f(rng); + let i_std = normalize(vec3(i.xy * roughness, i.z)); + let phi = PI_2 * rand.x; + let a = roughness; + let s = 1.0 + length(vec2(i.xy)); + let a2 = a * a; + let s2 = s * s; + let k = (1.0 - a2) * s2 / (s2 + a2 * i.z * i.z); + let b = select(i_std.z, k * i_std.z, i.z > 0.0); + let z = fma(1.0 - rand.y, 1.0 + b, -b); + let sin_theta = sqrt(saturate(1.0 - z * z)); + let o_std = vec3(sin_theta * cos(phi), sin_theta * sin(phi), z); + let m_std = i_std + o_std; + let m = normalize(vec3(m_std.xy * roughness, m_std.z)); + return 2.0 * dot(i, m) * m - i; +} + +// https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf (Listing 2) +fn ggx_vndf_pdf(wi_tangent: vec3, wo_tangent: vec3, roughness: f32) -> f32 { + let i = wi_tangent; + let o = wo_tangent; + let m = normalize(i + o); + let ndf = D_GGX(roughness, saturate(m.z)); + let ai = roughness * i.xy; + let len2 = dot(ai, ai); + let t = sqrt(len2 + i.z * i.z); + if i.z >= 0.0 { + let a = roughness; + let s = 1.0 + length(i.xy); + let a2 = a * a; + let s2 = s * s; + let k = (1.0 - a2) * s2 / (s2 + a2 * i.z * i.z); + return ndf / (2.0 * (k * i.z + t)); + } + return ndf * (t - i.z) / (2.0 * len2); +} struct LightSample { light_id: u32, @@ -39,6 +88,12 @@ fn sample_random_light(ray_origin: vec3, origin_world_normal: vec3, rn return light_contribution; } +fn random_light_pdf(hit: ResolvedRayHitFull) -> f32 { + let light_count = arrayLength(&light_sources); + let p_light = 1.0 / f32(light_count); + return p_light / (hit.triangle_area * f32(hit.triangle_count)); +} + fn generate_random_light_sample(rng: ptr) -> GenerateRandomLightSampleResult { let light_count = arrayLength(&light_sources); let light_id = rand_range_u(light_count, rng); diff --git a/examples/3d/solari.rs b/examples/3d/solari.rs index 417cfe4f30b32..88422d093dcca 100644 --- a/examples/3d/solari.rs +++ b/examples/3d/solari.rs @@ -87,7 +87,7 @@ fn add_raytracing_meshes_on_scene_load( mut commands: Commands, args: Res, ) { - // Ensure meshes are bevy_solari compatible + // Ensure meshes are Solari compatible for (_, mesh) in meshes.iter_mut() { mesh.remove_attribute(Mesh::ATTRIBUTE_UV_1.id); mesh.remove_attribute(Mesh::ATTRIBUTE_COLOR.id); @@ -113,8 +113,13 @@ fn add_raytracing_meshes_on_scene_load( } } - // Increase material emissive intensity to make it prettier for the example + // Adjust scene materials to better demo Solari features for (_, material) in materials.iter_mut() { material.emissive *= 200.0; + + if material.base_color.to_linear() == LinearRgba::new(0.5, 0.5, 0.5, 1.0) { + material.metallic = 1.0; + material.perceptual_roughness = 0.15; + } } } diff --git a/release-content/release-notes/bevy_solari.md b/release-content/release-notes/bevy_solari.md index 22417faa320e4..7a32ce80cf95d 100644 --- a/release-content/release-notes/bevy_solari.md +++ b/release-content/release-notes/bevy_solari.md @@ -1,7 +1,7 @@ --- title: Initial raytraced lighting progress (bevy_solari) -authors: ["@JMS55"] -pull_requests: [19058, 19620, 19790, 20020, 20113, 20213, 20259] +authors: ["@JMS55", "@SparkyPotato"] +pull_requests: [19058, 19620, 19790, 20020, 20113, 20213, 20242, 20259] --- (TODO: Embed solari example screenshot here)