Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Better variance, debug target, and more #45

Merged
merged 2 commits into from
Oct 26, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions blade-render/code/fill-gbuf.wgsl
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,6 @@
#include "debug.inc.wgsl"
#include "debug-param.inc.wgsl"

const DEBUG_CONSISTENCY: bool = false;

// Has to match the host!
struct Vertex {
pos: vec3<f32>,
Expand Down Expand Up @@ -46,6 +44,7 @@ var out_depth: texture_storage_2d<r32float, write>;
var out_flat_normal: texture_storage_2d<rgba8snorm, write>;
var out_basis: texture_storage_2d<rgba8snorm, write>;
var out_albedo: texture_storage_2d<rgba8unorm, write>;
var out_debug: texture_storage_2d<rgba8unorm, write>;

fn decode_normal(raw: u32) -> vec3<f32> {
return unpack4x8snorm(raw).xyz;
Expand Down Expand Up @@ -149,16 +148,20 @@ fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
let base_color_sample = textureSampleLevel(textures[entry.base_color_texture], sampler_linear, tex_coords, lod);
albedo = (base_color_factor * base_color_sample).xyz;
}
if (DEBUG_CONSISTENCY) {
if (debug.view_mode == DebugMode_HitConsistency) {
let reprojected = get_projected_pixel(camera, hit_position);
let barycentrics_pos_diff = positions * barycentrics - hit_position;
let camera_projection_diff = vec2<f32>(global_id.xy) - vec2<f32>(reprojected);
albedo = vec3<f32>(length(barycentrics_pos_diff), length(camera_projection_diff), 0.0);
let consistency = vec4<f32>(length(barycentrics_pos_diff), length(camera_projection_diff), 0.0, 0.0);
textureStore(out_debug, global_id.xy, consistency);
}
} else {
if (enable_debug) {
debug_buf.entry = DebugEntry();
}
if (debug.view_mode != DebugMode_Final) {
textureStore(out_debug, global_id.xy, vec4<f32>(0.0));
}
}

textureStore(out_depth, global_id.xy, vec4<f32>(depth, 0.0, 0.0, 0.0));
Expand Down
31 changes: 19 additions & 12 deletions blade-render/code/post-proc.wgsl
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
#use PostProcMode

struct ToneMapParams {
enabled: u32,
mode: u32,
average_lum: f32,
key_value: f32,
// minimum value of the pixels mapped to white brightness
Expand All @@ -8,6 +10,7 @@ struct ToneMapParams {

var t_albedo: texture_2d<f32>;
var light_diffuse: texture_2d<f32>;
var t_debug: texture_2d<f32>;
var<uniform> tone_map_params: ToneMapParams;

struct VertexOutput {
Expand All @@ -25,17 +28,21 @@ fn blit_vs(@builtin(vertex_index) vi: u32) -> VertexOutput {

@fragment
fn blit_fs(vo: VertexOutput) -> @location(0) vec4<f32> {
let tc = vec2<i32>(i32(vo.clip_pos.x), i32(vo.input_size.y) - i32(vo.clip_pos.y));
let albedo = textureLoad(t_albedo, tc, 0);
let radiance = textureLoad(light_diffuse, tc, 0);
let color = albedo * radiance;
if (tone_map_params.enabled != 0u) {
// Following https://blog.en.uwa4d.com/2022/07/19/physically-based-renderingg-hdr-tone-mapping/
let l_adjusted = tone_map_params.key_value / tone_map_params.average_lum * color.xyz;
let l_white = tone_map_params.white_level;
let l_ldr = l_adjusted * (1.0 + l_adjusted / (l_white*l_white)) / (1.0 + l_adjusted);
return vec4<f32>(l_ldr, 1.0);
let tc = vec2<i32>(i32(vo.clip_pos.x), i32(vo.input_size.y) - i32(vo.clip_pos.y) - 1);
if (tone_map_params.mode == PostProcMode_Debug) {
return textureLoad(t_debug, tc, 0);
} else {
return color;
let albedo = textureLoad(t_albedo, tc, 0);
let radiance = textureLoad(light_diffuse, tc, 0);
let color = albedo * radiance;
if (tone_map_params.mode == PostProcMode_Tonemap) {
// Following https://blog.en.uwa4d.com/2022/07/19/physically-based-renderingg-hdr-tone-mapping/
let l_adjusted = tone_map_params.key_value / tone_map_params.average_lum * color.xyz;
let l_white = tone_map_params.white_level;
let l_ldr = l_adjusted * (1.0 + l_adjusted / (l_white*l_white)) / (1.0 + l_adjusted);
return vec4<f32>(l_ldr, 1.0);
} else {
return color;
}
}
}
70 changes: 43 additions & 27 deletions blade-render/code/ray-trace.wgsl
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,9 @@ const PAIRWISE_MIS: bool = true;
// Bitterli's pseudocode (where it's 1) and NVidia's RTXDI implementation (where it's 0).
// With Bitterli's 1 we have MIS not respecting the prior history enough.
const BASE_CANONICAL_MIS: f32 = 0.05;
// See "DECOUPLING SHADING AND REUSE" in
// "Rearchitecting Spatiotemporal Resampling for Production"
const DECOUPLED_SHADING: bool = false;

struct MainParams {
frame_index: u32,
Expand Down Expand Up @@ -54,6 +57,7 @@ struct LiveReservoir {
selected_uv: vec2<f32>,
selected_light_index: u32,
selected_target_score: f32,
radiance: vec3<f32>,
weight_sum: f32,
history: f32,
}
Expand Down Expand Up @@ -81,9 +85,10 @@ fn bump_reservoir(r: ptr<function, LiveReservoir>, history: f32) {
}
fn make_reservoir(ls: LightSample, light_index: u32, brdf: vec3<f32>) -> LiveReservoir {
var r: LiveReservoir;
r.radiance = ls.radiance * brdf;
r.selected_uv = ls.uv;
r.selected_light_index = light_index;
r.selected_target_score = compute_target_score(ls.radiance * brdf);
r.selected_target_score = compute_target_score(r.radiance);
r.weight_sum = r.selected_target_score / ls.pdf;
r.history = 1.0;
return r;
Expand All @@ -95,6 +100,7 @@ fn merge_reservoir(r: ptr<function, LiveReservoir>, other: LiveReservoir, random
(*r).selected_light_index = other.selected_light_index;
(*r).selected_uv = other.selected_uv;
(*r).selected_target_score = other.selected_target_score;
(*r).radiance = other.radiance;
return true;
} else {
return false;
Expand All @@ -105,6 +111,7 @@ fn unpack_reservoir(f: StoredReservoir, max_history: u32) -> LiveReservoir {
r.selected_light_index = f.light_index;
r.selected_uv = f.light_uv;
r.selected_target_score = f.target_score;
r.radiance = vec3<f32>(0.0); // to be continued...
let history = min(f.confidence, f32(max_history));
r.weight_sum = f.contribution_weight * f.target_score * history;
r.history = history;
Expand All @@ -131,6 +138,7 @@ var t_prev_basis: texture_2d<f32>;
var t_flat_normal: texture_2d<f32>;
var t_prev_flat_normal: texture_2d<f32>;
var out_diffuse: texture_storage_2d<rgba16float, write>;
var out_debug: texture_storage_2d<rgba8unorm, write>;

fn sample_circle(random: f32) -> vec2<f32> {
let angle = 2.0 * PI * random;
Expand Down Expand Up @@ -237,34 +245,36 @@ fn evaluate_reflected_light(surface: Surface, light_index: u32, light_uv: vec2<f
return radiance * brdf;
}

struct TargetPdf {
struct TargetScore {
color: vec3<f32>,
score: f32,
}

fn make_target_pdf(color: vec3<f32>) -> TargetPdf {
return TargetPdf(color, compute_target_score(color));
fn make_target_score(color: vec3<f32>) -> TargetScore {
return TargetScore(color, compute_target_score(color));
}

fn estimate_target_pdf_with_occlusion(surface: Surface, position: vec3<f32>, light_index: u32, light_uv: vec2<f32>, debug_len: f32) -> TargetPdf {
fn estimate_target_score_with_occlusion(
surface: Surface, position: vec3<f32>, light_index: u32, light_uv: vec2<f32>, debug_len: f32
) -> TargetScore {
if (light_index != 0u) {
return TargetPdf();
return TargetScore();
}
let direction = map_equirect_uv_to_dir(light_uv);
if (dot(direction, surface.flat_normal) <= 0.0) {
return TargetPdf();
return TargetScore();
}
let brdf = evaluate_brdf(surface, direction);
if (brdf <= 0.0) {
return TargetPdf();
return TargetScore();
}

if (check_ray_occluded(position, direction, debug_len)) {
return TargetPdf();
return TargetScore();
} else {
//Note: same as `evaluate_reflected_light`
let radiance = textureSampleLevel(env_map, sampler_nearest, light_uv, 0.0).xyz;
return make_target_pdf(brdf * radiance);
return make_target_score(brdf * radiance);
}
}

Expand Down Expand Up @@ -312,8 +322,7 @@ struct RestirOutput {

fn compute_restir(surface: Surface, pixel: vec2<i32>, rng: ptr<function, RandomState>, enable_debug: bool) -> RestirOutput {
if (debug.view_mode == DebugMode_Depth) {
let depth = vec3<f32>(surface.depth / camera.depth);
return RestirOutput(depth);
textureStore(out_debug, pixel, vec4<f32>(surface.depth / camera.depth));
}
let ray_dir = get_ray_direction(camera, pixel);
let pixel_index = get_reservoir_index(pixel, camera);
Expand All @@ -327,11 +336,10 @@ fn compute_restir(surface: Surface, pixel: vec2<i32>, rng: ptr<function, RandomS
let position = camera.position + surface.depth * ray_dir;
let normal = qrot(surface.basis, vec3<f32>(0.0, 0.0, 1.0));
if (debug.view_mode == DebugMode_Normal) {
return RestirOutput(normal);
textureStore(out_debug, pixel, vec4<f32>(normal, 0.0));
}

var canonical = LiveReservoir();
var canonical_radiance = vec3<f32>(0.0);
for (var i = 0u; i < parameters.num_environment_samples; i += 1u) {
var ls: LightSample;
if (parameters.environment_importance_sampling != 0u) {
Expand All @@ -343,9 +351,7 @@ fn compute_restir(surface: Surface, pixel: vec2<i32>, rng: ptr<function, RandomS
let brdf = evaluate_sample(ls, surface, position, debug_len);
if (brdf > 0.0) {
let other = make_reservoir(ls, 0u, vec3<f32>(brdf));
if (merge_reservoir(&canonical, other, random_gen(rng))) {
canonical_radiance = ls.radiance * brdf;
}
merge_reservoir(&canonical, other, random_gen(rng));
} else {
bump_reservoir(&canonical, 1.0);
}
Expand Down Expand Up @@ -396,13 +402,13 @@ fn compute_restir(surface: Surface, pixel: vec2<i32>, rng: ptr<function, RandomS
var reservoir = LiveReservoir();
var shaded_color = vec3<f32>(0.0);
var mis_canonical = BASE_CANONICAL_MIS;
var color_and_weight = vec4<f32>(0.0);
for (var rid = 0u; rid < accepted_count; rid += 1u) {
let neighbor_index = accepted_reservoir_indices[rid];
let neighbor = prev_reservoirs[neighbor_index];

let max_history = select(parameters.spatial_tap_history, parameters.temporal_history, rid == temporal_index);
var other: LiveReservoir;
var other_color: vec3<f32>;
if (PAIRWISE_MIS) {
let neighbor_pixel = get_pixel_from_reservoir_index(neighbor_index, prev_camera);
let neighbor_history = min(neighbor.confidence, f32(max_history));
Expand All @@ -411,7 +417,7 @@ fn compute_restir(surface: Surface, pixel: vec2<i32>, rng: ptr<function, RandomS
let neighbor_dir = get_ray_direction(prev_camera, neighbor_pixel);
let neighbor_position = prev_camera.position + neighbor_surface.depth * neighbor_dir;

let t_canonical_at_neighbor = estimate_target_pdf_with_occlusion(
let t_canonical_at_neighbor = estimate_target_score_with_occlusion(
neighbor_surface, neighbor_position, canonical.selected_light_index, canonical.selected_uv, debug_len);
let mis_sub_canonical = balance_heuristic(
t_canonical_at_neighbor.score, canonical.selected_target_score,
Expand All @@ -424,7 +430,7 @@ fn compute_restir(surface: Surface, pixel: vec2<i32>, rng: ptr<function, RandomS
// target light has moved, and re-evaluate the occlusion.
// 2. we can use the cached target score, and there is no use of the target color
//let t_neighbor_at_neighbor = estimate_target_pdf(neighbor_surface, neighbor_position, neighbor.selected_dir);
let t_neighbor_at_canonical = estimate_target_pdf_with_occlusion(
let t_neighbor_at_canonical = estimate_target_score_with_occlusion(
surface, position, neighbor.light_index, neighbor.light_uv, debug_len);
let mis_neighbor = balance_heuristic(
neighbor.target_score, t_neighbor_at_canonical.score,
Expand All @@ -437,32 +443,42 @@ fn compute_restir(surface: Surface, pixel: vec2<i32>, rng: ptr<function, RandomS
other.weight_sum = t_neighbor_at_canonical.score * neighbor.contribution_weight * mis_neighbor.weight;
//Note: should be needed according to the paper
// other.history *= min(mis_neighbor.history, mis_sub_canonical.history);
other_color = t_neighbor_at_canonical.color;
other.radiance = t_neighbor_at_canonical.color;
} else {
other = unpack_reservoir(neighbor, max_history);
other_color = evaluate_reflected_light(surface, other.selected_light_index, other.selected_uv);
other.radiance = evaluate_reflected_light(surface, other.selected_light_index, other.selected_uv);
}

if (DECOUPLED_SHADING) {
color_and_weight += other.weight_sum * vec4<f32>(neighbor.contribution_weight * other.radiance, 1.0);
}
if (other.weight_sum <= 0.0) {
bump_reservoir(&reservoir, other.history);
} else if (merge_reservoir(&reservoir, other, random_gen(rng))) {
shaded_color = other_color;
} else {
merge_reservoir(&reservoir, other, random_gen(rng));
}
}

// Finally, merge in the canonical sample
if (PAIRWISE_MIS) {
canonical.weight_sum *= mis_canonical / canonical.history;
}
if (merge_reservoir(&reservoir, canonical, random_gen(rng))) {
shaded_color = canonical_radiance;
if (DECOUPLED_SHADING) {
//FIXME: issue with near zero denominator. Do we need do use BASE_CANONICAL_MIS?
let cw = canonical.weight_sum / max(canonical.selected_target_score * mis_canonical, 0.1);
color_and_weight += canonical.weight_sum * vec4<f32>(cw * canonical.radiance, 1.0);
}
merge_reservoir(&reservoir, canonical, random_gen(rng));

let effective_history = select(reservoir.history, BASE_CANONICAL_MIS + f32(accepted_count), PAIRWISE_MIS);
let stored = pack_reservoir_detail(reservoir, effective_history);
reservoirs[pixel_index] = stored;
var ro = RestirOutput();
ro.radiance = stored.contribution_weight * shaded_color;
if (DECOUPLED_SHADING) {
ro.radiance = color_and_weight.xyz / max(color_and_weight.w, 0.001);
} else {
ro.radiance = stored.contribution_weight * reservoir.radiance;
}
return ro;
}

Expand Down
1 change: 1 addition & 0 deletions blade-render/src/asset_hub.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ impl AssetHub {
sh_baker.register_enum::<crate::render::DebugMode>();
sh_baker.register_bitflags::<crate::render::DebugDrawFlags>();
sh_baker.register_bitflags::<crate::render::DebugTextureFlags>();
sh_baker.register_enum::<crate::render::PostProcMode>();
let shaders = AssetManager::new(target, choir, sh_baker);

Self {
Expand Down
Loading
Loading