Initial commit
This commit is contained in:
Máté Homolya
parent
8255e6cda9
commit
ffc0613e18
BIN
assets/environment_maps/ballawley_park_1k.ktx2
Normal file
BIN
assets/environment_maps/ballawley_park_1k.ktx2
Normal file
Binary file not shown.
@ -197,6 +197,22 @@ impl Image {
|
||||
})
|
||||
.map(DynamicImage::ImageRgba8)
|
||||
}
|
||||
TextureFormat::Rgba16Float => {
|
||||
use half::f16;
|
||||
let pixel_count = (width * height) as usize;
|
||||
let mut rgba32f_data = Vec::<f32>::with_capacity(pixel_count * 4);
|
||||
for rgba16f in data.chunks_exact(8) {
|
||||
let r = f16::from_bits(u16::from_le_bytes([rgba16f[0], rgba16f[1]])).to_f32();
|
||||
let g = f16::from_bits(u16::from_le_bytes([rgba16f[2], rgba16f[3]])).to_f32();
|
||||
let b = f16::from_bits(u16::from_le_bytes([rgba16f[4], rgba16f[5]])).to_f32();
|
||||
let a = f16::from_bits(u16::from_le_bytes([rgba16f[6], rgba16f[7]])).to_f32();
|
||||
rgba32f_data.push(r);
|
||||
rgba32f_data.push(g);
|
||||
rgba32f_data.push(b);
|
||||
rgba32f_data.push(a);
|
||||
}
|
||||
ImageBuffer::from_raw(width, height, rgba32f_data).map(DynamicImage::ImageRgba32F)
|
||||
}
|
||||
// Throw and error if conversion isn't supported
|
||||
texture_format => return Err(IntoDynamicImageError::UnsupportedFormat(texture_format)),
|
||||
}
|
||||
|
||||
BIN
crates/bevy_pbr/src/atmosphere/bluenoise.ktx2
Normal file
BIN
crates/bevy_pbr/src/atmosphere/bluenoise.ktx2
Normal file
Binary file not shown.
@ -45,6 +45,7 @@ use bevy_ecs::{
|
||||
schedule::IntoScheduleConfigs,
|
||||
system::{lifetimeless::Read, Query},
|
||||
};
|
||||
use bevy_image::Image;
|
||||
use bevy_math::{UVec2, UVec3, Vec3};
|
||||
use bevy_reflect::{std_traits::ReflectDefault, Reflect};
|
||||
use bevy_render::{
|
||||
|
||||
425
crates/bevy_pbr/src/light_probe/environment_filter.wgsl
Normal file
425
crates/bevy_pbr/src/light_probe/environment_filter.wgsl
Normal file
@ -0,0 +1,425 @@
|
||||
#import bevy_render::maths::{PI, PI_2};
|
||||
#import bevy_pbr::lighting::perceptualRoughnessToRoughness;
|
||||
|
||||
struct PrefilterConstants {
|
||||
mip_level: f32,
|
||||
sample_count: u32,
|
||||
roughness: f32,
|
||||
blue_noise_size: vec2f,
|
||||
}
|
||||
|
||||
@group(0) @binding(0) var input_texture: texture_2d_array<f32>;
|
||||
@group(0) @binding(1) var input_sampler: sampler;
|
||||
@group(0) @binding(2) var output_texture: texture_storage_2d_array<rgba16float, write>;
|
||||
@group(0) @binding(3) var<uniform> constants: PrefilterConstants;
|
||||
@group(0) @binding(4) var blue_noise_texture: texture_2d<f32>;
|
||||
|
||||
// Tonemapping functions to reduce fireflies
|
||||
fn rcp(x: f32) -> f32 { return 1.0 / x; }
|
||||
fn max3(x: vec3f) -> f32 { return max(x.r, max(x.g, x.b)); }
|
||||
fn tonemap(color: vec3f) -> vec3f {
|
||||
return color / (color + vec3(5000.0));
|
||||
}
|
||||
fn reverse_tonemap(color: vec3f) -> vec3f {
|
||||
return 5000.0 * color / (vec3(1.0) - color);
|
||||
}
|
||||
|
||||
// Predefined set of uniform directions
|
||||
fn get_uniform_direction(index: u32) -> vec3f {
|
||||
var dir = vec3f(0.0);
|
||||
|
||||
switch(index % 64u) {
|
||||
case 0u: { dir = vec3f(0.91593, -0.347884, 0.200123); }
|
||||
case 1u: { dir = vec3f(-0.244493, -0.710186, -0.660196); }
|
||||
case 2u: { dir = vec3f(-0.838322, 0.259442, 0.479484); }
|
||||
case 3u: { dir = vec3f(0.245473, 0.891464, -0.380835); }
|
||||
case 4u: { dir = vec3f(0.632533, -0.155099, 0.758846); }
|
||||
case 5u: { dir = vec3f(-0.20644, -0.973183, -0.101474); }
|
||||
case 6u: { dir = vec3f(-0.269471, 0.0483681, -0.961793); }
|
||||
case 7u: { dir = vec3f(0.143331, 0.973557, 0.177887); }
|
||||
case 8u: { dir = vec3f(0.725872, -0.086002, -0.682432); }
|
||||
case 9u: { dir = vec3f(-0.076835, -0.886014, 0.457249); }
|
||||
case 10u: { dir = vec3f(-0.913781, 0.0503775, -0.403071); }
|
||||
case 11u: { dir = vec3f(0.0159914, 0.676129, 0.73661); }
|
||||
case 12u: { dir = vec3f(0.992288, 0.00772121, -0.12371); }
|
||||
case 13u: { dir = vec3f(0.00641109, -0.177892, -0.984029); }
|
||||
case 14u: { dir = vec3f(-0.985566, -0.0665794, 0.155651); }
|
||||
case 15u: { dir = vec3f(-0.0700448, 0.706071, -0.704668); }
|
||||
case 16u: { dir = vec3f(0.89279, 0.117001, 0.435013); }
|
||||
case 17u: { dir = vec3f(0.142896, -0.893697, -0.425307); }
|
||||
case 18u: { dir = vec3f(-0.687174, -0.132142, 0.714374); }
|
||||
case 19u: { dir = vec3f(-0.217251, 0.965143, -0.145946); }
|
||||
case 20u: { dir = vec3f(0.108209, 0.0279573, 0.993735); }
|
||||
case 21u: { dir = vec3f(0.274912, -0.952168, 0.133416); }
|
||||
case 22u: { dir = vec3f(-0.653478, -0.211134, -0.726904); }
|
||||
case 23u: { dir = vec3f(-0.307126, 0.85749, 0.412777); }
|
||||
case 24u: { dir = vec3f(0.831999, 0.327845, -0.447543); }
|
||||
case 25u: { dir = vec3f(0.283463, -0.663772, 0.692138); }
|
||||
case 26u: { dir = vec3f(-0.893939, -0.415437, -0.168182); }
|
||||
case 27u: { dir = vec3f(-0.106605, 0.211719, 0.971499); }
|
||||
case 28u: { dir = vec3f(0.873146, 0.474611, 0.11118); }
|
||||
case 29u: { dir = vec3f(0.332658, -0.572825, -0.74914); }
|
||||
case 30u: { dir = vec3f(-0.781162, -0.487098, 0.390541); }
|
||||
case 31u: { dir = vec3f(-0.490404, 0.734038, -0.469779); }
|
||||
case 32u: { dir = vec3f(0.604084, 0.431641, 0.669902); }
|
||||
case 33u: { dir = vec3f(0.593065, -0.782314, -0.190417); }
|
||||
case 34u: { dir = vec3f(-0.244516, -0.197766, 0.949263); }
|
||||
case 35u: { dir = vec3f(-0.650394, 0.754372, 0.0889437); }
|
||||
case 36u: { dir = vec3f(0.468682, 0.430484, -0.771376); }
|
||||
case 37u: { dir = vec3f(0.647992, -0.666677, 0.368305); }
|
||||
case 38u: { dir = vec3f(-0.604909, -0.626104, -0.492015); }
|
||||
case 39u: { dir = vec3f(-0.564322, 0.511928, 0.647666); }
|
||||
case 40u: { dir = vec3f(0.633455, 0.743985, -0.212653); }
|
||||
case 41u: { dir = vec3f(0.292272, -0.234942, 0.927027); }
|
||||
case 42u: { dir = vec3f(-0.600382, -0.796926, 0.0667077); }
|
||||
case 43u: { dir = vec3f(-0.497216, 0.350652, -0.793612); }
|
||||
case 44u: { dir = vec3f(0.516356, 0.783334, 0.346069); }
|
||||
case 45u: { dir = vec3f(0.729109, -0.451604, -0.514251); }
|
||||
case 46u: { dir = vec3f(-0.389822, -0.675926, 0.62543); }
|
||||
case 47u: { dir = vec3f(-0.856868, 0.458917, -0.234889); }
|
||||
case 48u: { dir = vec3f(0.189162, 0.381537, 0.904791); }
|
||||
case 49u: { dir = vec3f(0.907219, -0.4183, 0.0444719); }
|
||||
case 50u: { dir = vec3f(-0.225508, -0.532484, -0.815848); }
|
||||
case 51u: { dir = vec3f(-0.882371, 0.341401, 0.323833); }
|
||||
case 52u: { dir = vec3f(0.279638, 0.796232, -0.536486); }
|
||||
case 53u: { dir = vec3f(0.759697, -0.242935, 0.603194); }
|
||||
case 54u: { dir = vec3f(-0.265275, -0.929255, -0.257125); }
|
||||
case 55u: { dir = vec3f(-0.455978, 0.114803, 0.882555); }
|
||||
case 56u: { dir = vec3f(0.213508, 0.976688, 0.0222359); }
|
||||
case 57u: { dir = vec3f(0.536034, -0.10141, -0.838084); }
|
||||
case 58u: { dir = vec3f(-0.147707, -0.941925, 0.301597); }
|
||||
case 59u: { dir = vec3f(-0.822975, 0.102675, -0.558722); }
|
||||
case 60u: { dir = vec3f(0.0753368, 0.810439, 0.580958); }
|
||||
case 61u: { dir = vec3f(0.958193, -0.0618399, -0.279361); }
|
||||
case 62u: { dir = vec3f(-0.0168296, -0.509477, 0.860319); }
|
||||
case 63u: { dir = vec3f(-0.999999, 0.00159255, 0.0); }
|
||||
default: { dir = vec3f(0.0, 0.0, 1.0); }
|
||||
}
|
||||
|
||||
return normalize(dir);
|
||||
}
|
||||
|
||||
// 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 blue_noise_offset(pixel_coords: vec2u) -> vec2f {
|
||||
// Get a stable random offset for this pixel
|
||||
let noise_size = vec2u(u32(constants.blue_noise_size.x), u32(constants.blue_noise_size.y));
|
||||
let noise_coords = pixel_coords % noise_size;
|
||||
return textureSampleLevel(blue_noise_texture, input_sampler, vec2f(noise_coords) / constants.blue_noise_size, 0.0).rg;
|
||||
}
|
||||
|
||||
// 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 = sqrt((1.0 - xi.y) / (1.0 + (a * a - 1.0) * xi.y));
|
||||
let sin_theta = 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 blue_noise = blue_noise_offset(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 + blue_noise); // 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));
|
||||
}
|
||||
|
||||
@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);
|
||||
|
||||
// Create tangent space matrix
|
||||
let tangent_frame = calculate_tangent_frame(normal);
|
||||
|
||||
// Get blue noise offset for stratification
|
||||
let blue_noise = blue_noise_offset(coords);
|
||||
|
||||
var irradiance = vec3f(0.0);
|
||||
var total_weight = 0.0;
|
||||
|
||||
let sample_count = min(constants.sample_count, 64u);
|
||||
|
||||
for (var i = 0u; i < sample_count; i++) {
|
||||
// Using a predefined set of directions provides good hemisphere coverage for diffuse
|
||||
var dir = get_uniform_direction((i + u32(coords.x * 7u + coords.y * 11u + face * 5u)) % 64u);
|
||||
|
||||
// Ensure the direction is in the hemisphere defined by the normal
|
||||
let NoL = dot(normal, dir);
|
||||
|
||||
// Flip the direction if it's in the wrong hemisphere
|
||||
if (NoL < 0.0) {
|
||||
dir = -dir;
|
||||
}
|
||||
|
||||
// Recalculate NoL after possible flipping
|
||||
let weight = max(dot(normal, dir), 0.0);
|
||||
|
||||
if (weight > 0.0) {
|
||||
// Lambert PDF
|
||||
let pdf = weight / PI;
|
||||
let width = f32(size.x);
|
||||
|
||||
// Filtered importance sampling
|
||||
let mip_level = clamp(
|
||||
calculate_environment_map_lod(pdf, width, f32(sample_count)),
|
||||
1.0,
|
||||
constants.roughness * 3.0
|
||||
);
|
||||
|
||||
// Sample environment with the calculated mip level
|
||||
let sample_color = sample_environment(dir, mip_level).rgb;
|
||||
|
||||
// Accumulate the sample
|
||||
irradiance += sample_color * weight;
|
||||
total_weight += weight;
|
||||
}
|
||||
}
|
||||
|
||||
// Normalize and scale by PI for diffuse BRDF
|
||||
if (total_weight > 0.0) {
|
||||
irradiance = irradiance / total_weight * PI;
|
||||
}
|
||||
|
||||
// Add some low-frequency ambient term to avoid completely dark areas
|
||||
irradiance = max(irradiance, vec3f(0.01, 0.01, 0.01));
|
||||
|
||||
// Write result to output texture
|
||||
textureStore(output_texture, coords, face, vec4f(irradiance, 1.0));
|
||||
}
|
||||
@ -31,6 +31,9 @@ use bevy_render::{
|
||||
Extract, ExtractSchedule, Render, RenderApp, RenderSystems,
|
||||
};
|
||||
use bevy_transform::{components::Transform, prelude::GlobalTransform};
|
||||
use prefilter::{
|
||||
create_environment_map_from_prefilter, extract_prefilter_entities, prepare_prefilter_bind_groups, prepare_prefilter_textures, FilteredEnvironmentMapLight, PrefilterPipelines, SpdNode
|
||||
};
|
||||
use tracing::error;
|
||||
|
||||
use core::{hash::Hash, ops::Deref};
|
||||
@ -41,6 +44,7 @@ use self::irradiance_volume::IrradianceVolume;
|
||||
|
||||
pub mod environment_map;
|
||||
pub mod irradiance_volume;
|
||||
pub mod prefilter;
|
||||
|
||||
/// The maximum number of each type of light probe that each view will consider.
|
||||
///
|
||||
@ -343,7 +347,9 @@ impl Plugin for LightProbePlugin {
|
||||
|
||||
app.register_type::<LightProbe>()
|
||||
.register_type::<EnvironmentMapLight>()
|
||||
.register_type::<IrradianceVolume>();
|
||||
.register_type::<IrradianceVolume>()
|
||||
.add_plugins(ExtractComponentPlugin::<FilteredEnvironmentMapLight>::default())
|
||||
.add_systems(Update, create_environment_map_from_prefilter);
|
||||
}
|
||||
|
||||
fn finish(&self, app: &mut App) {
|
||||
@ -355,9 +361,26 @@ impl Plugin for LightProbePlugin {
|
||||
.add_plugins(ExtractInstancesPlugin::<EnvironmentMapIds>::new())
|
||||
.init_resource::<LightProbesBuffer>()
|
||||
.init_resource::<EnvironmentMapUniformBuffer>()
|
||||
.init_resource::<PrefilterBindGroupLayouts>()
|
||||
.init_resource::<PrefilterSamplers>()
|
||||
.init_resource::<PrefilterPipelines>()
|
||||
.add_render_graph_node::<SpdNode>(Core3d, PrefilterNode::GenerateMipmap)
|
||||
.add_render_graph_node::<RadianceMapNode>(Core3d, PrefilterNode::RadianceMap)
|
||||
.add_render_graph_node::<IrradianceMapNode>(Core3d, PrefilterNode::IrradianceMap)
|
||||
.add_render_graph_edges(
|
||||
Core3d,
|
||||
(
|
||||
Node3d::EndPrepasses,
|
||||
PrefilterNode::GenerateMipmap,
|
||||
PrefilterNode::RadianceMap,
|
||||
PrefilterNode::IrradianceMap,
|
||||
Node3d::StartMainPass,
|
||||
),
|
||||
)
|
||||
.add_systems(ExtractSchedule, gather_environment_map_uniform)
|
||||
.add_systems(ExtractSchedule, gather_light_probes::<EnvironmentMapLight>)
|
||||
.add_systems(ExtractSchedule, gather_light_probes::<IrradianceVolume>)
|
||||
.add_systems(ExtractSchedule, extract_prefilter_entities.after(create_environment_map_from_prefilter),)
|
||||
.add_systems(
|
||||
Render,
|
||||
(upload_light_probes, prepare_environment_uniform_buffer)
|
||||
|
||||
938
crates/bevy_pbr/src/light_probe/prefilter.rs
Normal file
938
crates/bevy_pbr/src/light_probe/prefilter.rs
Normal file
@ -0,0 +1,938 @@
|
||||
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::atmosphere;
|
||||
use crate::light_probe::environment_map::EnvironmentMapLight;
|
||||
|
||||
/// A handle to the SPD (Single Pass Downsampling) shader.
|
||||
pub const SPD_SHADER_HANDLE: Handle<Shader> = weak_handle!("5dcf400c-bcb3-49b9-8b7e-80f4117eaf82");
|
||||
|
||||
/// A handle to the environment filter shader.
|
||||
pub const ENVIRONMENT_FILTER_SHADER_HANDLE: Handle<Shader> =
|
||||
weak_handle!("3110b545-78e0-48fc-b86e-8bc0ea50fc67");
|
||||
|
||||
/// Labels for the prefiltering nodes
|
||||
#[derive(PartialEq, Eq, Debug, Copy, Clone, Hash, RenderLabel)]
|
||||
pub enum PrefilterNode {
|
||||
GenerateMipmap,
|
||||
RadianceMap,
|
||||
IrradianceMap,
|
||||
}
|
||||
|
||||
/// Stores the bind group layouts for the prefiltering process
|
||||
#[derive(Resource)]
|
||||
pub struct PrefilterBindGroupLayouts {
|
||||
pub spd: BindGroupLayout,
|
||||
pub radiance: BindGroupLayout,
|
||||
pub irradiance: BindGroupLayout,
|
||||
}
|
||||
|
||||
impl FromWorld for PrefilterBindGroupLayouts {
|
||||
fn from_world(world: &mut World) -> Self {
|
||||
let render_device = world.resource::<RenderDevice>();
|
||||
|
||||
// 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::<SpdConstants>(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::<PrefilterConstants>(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::<PrefilterConstants>(false)), // Uniforms
|
||||
(4, texture_2d(TextureSampleType::Float { filterable: true })), // Blue noise texture
|
||||
),
|
||||
),
|
||||
);
|
||||
|
||||
Self {
|
||||
spd,
|
||||
radiance,
|
||||
irradiance,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Samplers for the prefiltering process
|
||||
#[derive(Resource)]
|
||||
pub struct PrefilterSamplers {
|
||||
pub linear: Sampler,
|
||||
}
|
||||
|
||||
impl FromWorld for PrefilterSamplers {
|
||||
fn from_world(world: &mut World) -> Self {
|
||||
let render_device = world.resource::<RenderDevice>();
|
||||
|
||||
let linear = render_device.create_sampler(&SamplerDescriptor {
|
||||
label: Some("prefilter_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 prefiltering process
|
||||
#[derive(Resource)]
|
||||
pub struct PrefilterPipelines {
|
||||
pub spd_first: CachedComputePipelineId,
|
||||
pub spd_second: CachedComputePipelineId,
|
||||
pub radiance: CachedComputePipelineId,
|
||||
pub irradiance: CachedComputePipelineId,
|
||||
}
|
||||
|
||||
impl FromWorld for PrefilterPipelines {
|
||||
fn from_world(world: &mut World) -> Self {
|
||||
let pipeline_cache = world.resource::<PipelineCache>();
|
||||
let layouts = world.resource::<PrefilterBindGroupLayouts>();
|
||||
|
||||
let render_device = world.resource::<RenderDevice>();
|
||||
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 FilteredEnvironmentMapLight {
|
||||
pub environment_map: Handle<Image>,
|
||||
pub intensity: f32,
|
||||
pub rotation: Quat,
|
||||
pub affects_lightmapped_mesh_diffuse: bool,
|
||||
}
|
||||
|
||||
impl Default for FilteredEnvironmentMapLight {
|
||||
fn default() -> Self {
|
||||
FilteredEnvironmentMapLight {
|
||||
environment_map: Handle::default(),
|
||||
intensity: 0.0,
|
||||
rotation: Quat::IDENTITY,
|
||||
affects_lightmapped_mesh_diffuse: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn extract_prefilter_entities(
|
||||
prefilter_query: Extract<
|
||||
Query<(
|
||||
RenderEntity,
|
||||
&FilteredEnvironmentMapLight,
|
||||
&EnvironmentMapLight,
|
||||
)>,
|
||||
>,
|
||||
mut commands: Commands,
|
||||
render_images: Res<RenderAssets<GpuImage>>,
|
||||
) {
|
||||
for (entity, filtered_env_map, env_map_light) in prefilter_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,
|
||||
};
|
||||
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,
|
||||
}
|
||||
|
||||
#[derive(Component)]
|
||||
pub struct PrefilterTextures {
|
||||
pub environment_map: CachedTexture,
|
||||
}
|
||||
|
||||
/// Prepares textures needed for prefiltering
|
||||
pub fn prepare_prefilter_textures(
|
||||
light_probes: Query<Entity, With<RenderEnvironmentMap>>,
|
||||
render_device: Res<RenderDevice>,
|
||||
mut texture_cache: ResMut<TextureCache>,
|
||||
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("prefilter_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(PrefilterTextures { 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 prefiltering
|
||||
#[derive(Clone, Copy, ShaderType)]
|
||||
#[repr(C)]
|
||||
pub struct PrefilterConstants {
|
||||
mip_level: f32,
|
||||
sample_count: u32,
|
||||
roughness: f32,
|
||||
blue_noise_size: Vec2,
|
||||
}
|
||||
|
||||
/// Stores bind groups for the prefiltering process
|
||||
#[derive(Component)]
|
||||
pub struct PrefilterBindGroups {
|
||||
pub spd: BindGroup,
|
||||
pub radiance: Vec<BindGroup>, // One per mip level
|
||||
pub irradiance: BindGroup,
|
||||
}
|
||||
|
||||
/// Prepares bind groups for prefiltering
|
||||
pub fn prepare_prefilter_bind_groups(
|
||||
light_probes: Query<
|
||||
(Entity, &PrefilterTextures, &RenderEnvironmentMap),
|
||||
With<RenderEnvironmentMap>,
|
||||
>,
|
||||
render_device: Res<RenderDevice>,
|
||||
queue: Res<RenderQueue>,
|
||||
layouts: Res<PrefilterBindGroupLayouts>,
|
||||
samplers: Res<PrefilterSamplers>,
|
||||
render_images: Res<RenderAssets<GpuImage>>,
|
||||
mut commands: Commands,
|
||||
) {
|
||||
// Get blue noise texture
|
||||
let blue_noise_texture = render_images
|
||||
.get(&atmosphere::shaders::BLUENOISE_TEXTURE)
|
||||
.expect("Blue noise 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 = PrefilterConstants {
|
||||
mip_level: mip as f32,
|
||||
sample_count,
|
||||
roughness,
|
||||
blue_noise_size: Vec2::new(
|
||||
blue_noise_texture.size.width as f32,
|
||||
blue_noise_texture.size.height as f32,
|
||||
),
|
||||
};
|
||||
|
||||
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, &blue_noise_texture.texture_view),
|
||||
)),
|
||||
);
|
||||
|
||||
radiance_bind_groups.push(bind_group);
|
||||
}
|
||||
|
||||
// Create irradiance bind group
|
||||
let irradiance_constants = PrefilterConstants {
|
||||
mip_level: 0.0,
|
||||
sample_count: 64,
|
||||
roughness: 1.0,
|
||||
blue_noise_size: Vec2::new(
|
||||
blue_noise_texture.size.width as f32,
|
||||
blue_noise_texture.size.height as f32,
|
||||
),
|
||||
};
|
||||
|
||||
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, &blue_noise_texture.texture_view),
|
||||
)),
|
||||
);
|
||||
|
||||
commands.entity(entity).insert(PrefilterBindGroups {
|
||||
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<PrefilterBindGroups>,
|
||||
Read<RenderEnvironmentMap>,
|
||||
)>,
|
||||
}
|
||||
|
||||
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::<PipelineCache>();
|
||||
let pipelines = world.resource::<PrefilterPipelines>();
|
||||
|
||||
// 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::<PrefilterTextures>(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<PrefilterBindGroups>)>,
|
||||
}
|
||||
|
||||
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::<PipelineCache>();
|
||||
let pipelines = world.resource::<PrefilterPipelines>();
|
||||
|
||||
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<PrefilterBindGroups>)>,
|
||||
}
|
||||
|
||||
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::<PipelineCache>();
|
||||
let pipelines = world.resource::<PrefilterPipelines>();
|
||||
|
||||
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 creates an `EnvironmentMapLight` component from the prefiltered textures
|
||||
pub fn create_environment_map_from_prefilter(
|
||||
mut commands: Commands,
|
||||
mut images: ResMut<Assets<Image>>,
|
||||
query: Query<(Entity, &FilteredEnvironmentMapLight), Without<EnvironmentMapLight>>,
|
||||
) {
|
||||
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,
|
||||
});
|
||||
}
|
||||
}
|
||||
398
crates/bevy_pbr/src/light_probe/spd.wgsl
Normal file
398
crates/bevy_pbr/src/light_probe/spd.wgsl
Normal file
@ -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<f32>;
|
||||
@group(0) @binding(1) var mip_1: texture_storage_2d_array<rgba16float, write>;
|
||||
@group(0) @binding(2) var mip_2: texture_storage_2d_array<rgba16float, write>;
|
||||
@group(0) @binding(3) var mip_3: texture_storage_2d_array<rgba16float, write>;
|
||||
@group(0) @binding(4) var mip_4: texture_storage_2d_array<rgba16float, write>;
|
||||
@group(0) @binding(5) var mip_5: texture_storage_2d_array<rgba16float, write>;
|
||||
@group(0) @binding(6) var mip_6: texture_storage_2d_array<rgba16float, read_write>;
|
||||
@group(0) @binding(7) var mip_7: texture_storage_2d_array<rgba16float, write>;
|
||||
@group(0) @binding(8) var mip_8: texture_storage_2d_array<rgba16float, write>;
|
||||
@group(0) @binding(9) var mip_9: texture_storage_2d_array<rgba16float, write>;
|
||||
@group(0) @binding(10) var mip_10: texture_storage_2d_array<rgba16float, write>;
|
||||
@group(0) @binding(11) var mip_11: texture_storage_2d_array<rgba16float, write>;
|
||||
@group(0) @binding(12) var mip_12: texture_storage_2d_array<rgba16float, write>;
|
||||
@group(0) @binding(13) var sampler_linear_clamp: sampler;
|
||||
@group(0) @binding(14) var<uniform> constants: Constants;
|
||||
struct Constants { mips: u32, inverse_input_size: vec2f }
|
||||
|
||||
var<workgroup> spd_intermediate_r: array<array<f32, 16>, 16>;
|
||||
var<workgroup> spd_intermediate_g: array<array<f32, 16>, 16>;
|
||||
var<workgroup> spd_intermediate_b: array<array<f32, 16>, 16>;
|
||||
var<workgroup> spd_intermediate_a: array<array<f32, 16>, 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<vec4f, 4>;
|
||||
|
||||
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);
|
||||
}
|
||||
@ -37,6 +37,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 prefiltered environment map is shown.
|
||||
PrefilteredEnvironmentMap = 3,
|
||||
}
|
||||
|
||||
// The various reflection maps.
|
||||
@ -53,6 +55,9 @@ struct Cubemaps {
|
||||
// The specular cubemap that reflects both the world and the cubes.
|
||||
specular_reflection_probe: Handle<Image>,
|
||||
|
||||
// Unfiltered environment map
|
||||
unfiltered_environment_map: Handle<Image>,
|
||||
|
||||
// The skybox cubemap image. This is almost the same as
|
||||
// `specular_environment_map`.
|
||||
skybox: Handle<Image>,
|
||||
@ -146,6 +151,21 @@ fn spawn_reflection_probe(commands: &mut Commands, cubemaps: &Cubemaps) {
|
||||
));
|
||||
}
|
||||
|
||||
fn spawn_prefiltered_environment_map(commands: &mut Commands, cubemaps: &Cubemaps) {
|
||||
println!("spawn_prefiltered_environment_map");
|
||||
commands.spawn((
|
||||
LightProbe,
|
||||
FilteredEnvironmentMapLight {
|
||||
environment_map: cubemaps.unfiltered_environment_map.clone(),
|
||||
intensity: 5000.0,
|
||||
..default()
|
||||
},
|
||||
// 2.0 because the sphere's radius is 1.0 and we want to fully enclose it.
|
||||
Transform::from_scale(Vec3::splat(2.0)),
|
||||
));
|
||||
}
|
||||
|
||||
|
||||
// Spawns the help text.
|
||||
fn spawn_text(commands: &mut Commands, app_status: &AppStatus) {
|
||||
// Create the text.
|
||||
@ -196,7 +216,7 @@ 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() {
|
||||
@ -205,6 +225,7 @@ fn change_reflection_type(
|
||||
match app_status.reflection_mode {
|
||||
ReflectionMode::None | ReflectionMode::EnvironmentMap => {}
|
||||
ReflectionMode::ReflectionProbe => spawn_reflection_probe(&mut commands, &cubemaps),
|
||||
ReflectionMode::PrefilteredEnvironmentMap => spawn_prefiltered_environment_map(&mut commands, &cubemaps),
|
||||
}
|
||||
|
||||
// Add or remove the environment map from the camera.
|
||||
@ -213,7 +234,7 @@ fn change_reflection_type(
|
||||
ReflectionMode::None => {
|
||||
commands.entity(camera).remove::<EnvironmentMapLight>();
|
||||
}
|
||||
ReflectionMode::EnvironmentMap | ReflectionMode::ReflectionProbe => {
|
||||
ReflectionMode::EnvironmentMap | ReflectionMode::ReflectionProbe | ReflectionMode::PrefilteredEnvironmentMap => {
|
||||
commands
|
||||
.entity(camera)
|
||||
.insert(create_camera_environment_map_light(&cubemaps));
|
||||
@ -244,6 +265,7 @@ impl TryFrom<u32> for ReflectionMode {
|
||||
0 => Ok(ReflectionMode::None),
|
||||
1 => Ok(ReflectionMode::EnvironmentMap),
|
||||
2 => Ok(ReflectionMode::ReflectionProbe),
|
||||
3 => Ok(ReflectionMode::PrefilteredEnvironmentMap),
|
||||
_ => Err(()),
|
||||
}
|
||||
}
|
||||
@ -255,6 +277,7 @@ impl Display for ReflectionMode {
|
||||
ReflectionMode::None => "No reflections",
|
||||
ReflectionMode::EnvironmentMap => "Environment map",
|
||||
ReflectionMode::ReflectionProbe => "Reflection probe",
|
||||
ReflectionMode::PrefilteredEnvironmentMap => "Prefiltered environment map",
|
||||
};
|
||||
formatter.write_str(text)
|
||||
}
|
||||
@ -321,6 +344,7 @@ 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/ballawley_park_1k.ktx2"),
|
||||
skybox: specular_map,
|
||||
}
|
||||
}
|
||||
@ -329,8 +353,8 @@ impl FromWorld for Cubemaps {
|
||||
impl Default for AppStatus {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
reflection_mode: ReflectionMode::ReflectionProbe,
|
||||
rotating: true,
|
||||
reflection_mode: ReflectionMode::PrefilteredEnvironmentMap,
|
||||
rotating: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user