From ad85b4f044a33a138223fb566824a31a69a49190 Mon Sep 17 00:00:00 2001 From: wires Date: Wed, 28 May 2025 13:36:19 -0400 Subject: [PATCH] switch to sobol sampling --- Cargo.lock | 117 +++++++++++++++- core/Cargo.toml | 1 + core/build.rs | 17 --- core/src/lib.rs | 9 +- core/src/sampling.rs | 325 +++++++++++++++++++------------------------ examples/basic.py | 2 +- 6 files changed, 263 insertions(+), 208 deletions(-) delete mode 100644 core/build.rs diff --git a/Cargo.lock b/Cargo.lock index d0942c6..ceb24ad 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -14,6 +14,12 @@ version = "1.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6b1fc10dbac614ebc03540c9dbd60e83887fda27794998c6528f1782047d540" +[[package]] +name = "cfg-if" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4785bdd1c96b2a846b2bd7cc02e86b6b3dbf14e7e53446c4f54c92a361040822" + [[package]] name = "cfg-if" version = "1.0.0" @@ -32,6 +38,40 @@ dependencies = [ "syn", ] +[[package]] +name = "fasthash" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "032213946b4eaae09117ec63f020322b78ca7a31d8aa2cf64df3032e1579690f" +dependencies = [ + "cfg-if 0.1.10", + "fasthash-sys", + "num-traits", + "seahash", + "xoroshiro128", +] + +[[package]] +name = "fasthash-sys" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6de941abfe2e715cdd34009d90546f850597eb69ca628ddfbf616e53dda28f8" +dependencies = [ + "gcc", +] + +[[package]] +name = "fuchsia-cprng" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06f77d526c1a601b7c4cdd98f54b5eaabffc14d5f2f0296febdc7f357c6d3ba" + +[[package]] +name = "gcc" +version = "0.3.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f5f3913fa0bfe7ee1fd8248b6b9f42a5af4b9d65ec2dd2c3c26132b950ecfc2" + [[package]] name = "glam" version = "0.30.2" @@ -71,6 +111,7 @@ version = "0.1.0" dependencies = [ "bytemuck", "enum_dispatch", + "fasthash", "glam", "primal", ] @@ -184,7 +225,7 @@ version = "0.24.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5203598f366b11a02b13aa20cab591229ff0a89fd121a308a5df751d5fc9219" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", "indoc", "libc", "memoffset", @@ -250,6 +291,49 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "552840b97013b1a26992c11eac34bdd778e464601a4c2054b5f0bff7c6761293" +dependencies = [ + "fuchsia-cprng", + "libc", + "rand_core 0.3.1", + "rdrand", + "winapi", +] + +[[package]] +name = "rand_core" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a6fdeb83b075e8266dcc8762c22776f6877a63111121f5f8c7411e5be7eed4b" +dependencies = [ + "rand_core 0.4.2", +] + +[[package]] +name = "rand_core" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c33a3c44ca05fa6f1807d8e6743f3824e8509beca625669633be0acbdf509dc" + +[[package]] +name = "rdrand" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "678054eb77286b51581ba43620cc911abf02758c91f93f479767aed0f90458b2" +dependencies = [ + "rand_core 0.3.1", +] + +[[package]] +name = "seahash" +version = "3.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "58f57ca1d128a43733fd71d583e837b1f22239a37ebea09cde11d8d9a9080f47" + [[package]] name = "smallvec" version = "1.15.0" @@ -284,3 +368,34 @@ name = "unindent" version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "xoroshiro128" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0eeda34baec49c4f1eb2c04d59b761582fd6330010f9330ca696ca1a355dfcd" +dependencies = [ + "rand", +] diff --git a/core/Cargo.toml b/core/Cargo.toml index 3cc25f1..0cb7aef 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -9,6 +9,7 @@ repository.workspace = true [dependencies] bytemuck = { version = "1.22.0", optional = true } enum_dispatch = "0.3.13" +fasthash = "0.4.0" glam = "0.30.2" [features] diff --git a/core/build.rs b/core/build.rs deleted file mode 100644 index 97e239f..0000000 --- a/core/build.rs +++ /dev/null @@ -1,17 +0,0 @@ -use std::{env, fs::File, io::Write, path::PathBuf}; - -fn main() { - let dest: PathBuf = env::var("OUT_DIR").unwrap().into(); - let mut file = File::create(dest.join("primes.rs")).unwrap(); - - write_primes(&mut file).unwrap(); -} - -fn write_primes(file: &mut File) -> std::io::Result<()> { - write!(file, "[")?; - let sieve = primal::Sieve::new(7919); - for p in sieve.primes_from(2).take_while(|x| *x <= 7919) { - write!(file, "{p},")?; - } - write!(file, "]") -} diff --git a/core/src/lib.rs b/core/src/lib.rs index 6f92a66..bce4906 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -30,7 +30,7 @@ impl Scene { } pub fn render(&self, width: u32, height: u32, samples: u32) -> Box<[Rgba]> { - let mut sampler = Sampler::new(width, height); + let mut sampler = Sampler::new(width, height, samples, 0); let mut data = Box::new_uninit_slice(width as usize * height as usize); let w = width as f32; let h = height as f32; @@ -42,9 +42,10 @@ impl Scene { let mut rgb = Vector::ZERO; - for (offset_x, offset_y, _) in sampler.pixel(bx, by, samples) { - let x = bx as f32 + offset_x; - let y = by as f32 + offset_y; + for mut pixel_samples in sampler.pixel(bx, by) { + let (x_offset, y_offset) = pixel_samples.get_2d(); + let x = bx as f32 + x_offset; + let y = by as f32 + y_offset; let cx = 2.0 * x / w - 1.0; let cy = (1.0 - 2.0 * y / h) * aspect_ratio; diff --git a/core/src/sampling.rs b/core/src/sampling.rs index 9cfcbab..fa68689 100644 --- a/core/src/sampling.rs +++ b/core/src/sampling.rs @@ -1,220 +1,175 @@ -//! Halton sampling +use core::f32; -const PRIMES: [u32; 1000] = include!(concat!(env!("OUT_DIR"), "/primes.rs")); -const MAX_HALTON_RESOLUTION: u32 = 128; +use fasthash::murmur2::hash64_with_seed; + +// rust needs to add hex literals. this is 0x1p-32 +const ONE_OVER_MAX: f32 = 2.3283064e-10; + +const SOBOL_MATRIX: [u32; 32] = [ + 0x80000000, 0xc0000000, 0xa0000000, 0xf0000000, 0x88000000, 0xcc000000, 0xaa000000, 0xff000000, + 0x80800000, 0xc0c00000, 0xa0a00000, 0xf0f00000, 0x88880000, 0xcccc0000, 0xaaaa0000, 0xffff0000, + 0x80008000, 0xc000c000, 0xa000a000, 0xf000f000, 0x88008800, 0xcc00cc00, 0xaa00aa00, 0xff00ff00, + 0x80808080, 0xc0c0c0c0, 0xa0a0a0a0, 0xf0f0f0f0, 0x88888888, 0xcccccccc, 0xaaaaaaaa, 0xffffffff, +]; + +#[rustfmt::skip] +const PERMUTATIONS: [[u8; 4]; 24] = [ + [0, 1, 2, 3], [0, 1, 3, 2], [0, 2, 1, 3], [0, 2, 3, 1], [0, 3, 1, 2], [0, 3, 2, 1], + [1, 0, 2, 3], [1, 0, 3, 2], [1, 2, 0, 3], [1, 2, 3, 0], [1, 3, 0, 2], [1, 3, 2, 0], + [2, 0, 1, 3], [2, 0, 3, 1], [2, 1, 0, 3], [2, 1, 3, 0], [2, 3, 0, 1], [2, 3, 1, 0], + [3, 0, 1, 2], [3, 0, 2, 1], [3, 1, 0, 2], [3, 1, 2, 0], [3, 2, 0, 1], [3, 2, 1, 0], +]; pub struct Sampler { - x_scale: u8, - x_exponent: u8, - x_inverse: u8, - y_scale: u8, - y_exponent: u8, - y_inverse: u8, - index: u32, - dimension: u16, + seed: u32, + samples: u32, + log2_samples: u32, + base4_digits: u32, } impl Sampler { - pub fn new(width: u32, height: u32) -> Self { - let mut x_scale = 1u8; - let mut x_exponent = 0u8; - let x_res = u32::min(width, MAX_HALTON_RESOLUTION) as u8; - while x_scale < x_res { - x_scale *= 2; - x_exponent += 1; - } - - let mut y_scale = 1u8; - let mut y_exponent = 0u8; - let y_res = u32::min(height, MAX_HALTON_RESOLUTION) as u8; - while y_scale < y_res { - y_scale *= 3; - y_exponent += 1; - } - - let x_inverse = mod_inverse(y_scale, x_scale); - let y_inverse = mod_inverse(x_scale, y_scale); + pub fn new(width: u32, height: u32, samples: u32, seed: u32) -> Self { + let res = u32::next_power_of_two(u32::max(width, height)); + let log2_samples = u32::ilog2(samples); + let base4_digits = u32::ilog2(res) + log2_samples.div_ceil(2); Self { - dimension: 2, - index: 0, - x_scale, - x_exponent, - x_inverse, - y_scale, - y_exponent, - y_inverse, + seed, + samples, + log2_samples, + base4_digits, } } - pub fn pixel( - &mut self, - x: u32, - y: u32, - samples: u32, - ) -> impl Iterator { - let mut base = 0; - let sample_stride = self.y_scale as u32 * self.x_scale as u32; - - if sample_stride > 1 { - let mx = x % MAX_HALTON_RESOLUTION; - let my = y % MAX_HALTON_RESOLUTION; - - let x_offset = inverse_radical_inverse(mx, 2, self.x_exponent as u32); - let y_offset = inverse_radical_inverse(my, 3, self.y_exponent as u32); - - base += x_offset * (sample_stride / self.x_scale as u32) * self.x_inverse as u32; - base += y_offset * (sample_stride / self.y_scale as u32) * self.y_inverse as u32; - - base %= sample_stride; - } + pub fn pixel(&mut self, x: u32, y: u32) -> impl Iterator { + let &mut Self { + seed, + samples, + log2_samples, + base4_digits, + .. + } = self; (0..samples).map(move |i| { - let index = base + i * sample_stride; - let x_offset = radical_inverse(0, index >> self.x_exponent); - let y_offset = radical_inverse(1, index / self.y_scale as u32); - (x_offset, y_offset, Sample::new(index)) + let morton_index = (morton_2d(x, y) << log2_samples) | i; + + PixelSample { + seed, + log2_samples, + base4_digits, + morton_index, + dimension: 0, + } }) } } -pub struct Sample { - index: u32, +pub struct PixelSample { + seed: u32, + log2_samples: u32, + base4_digits: u32, + morton_index: u32, dimension: u32, } -impl Sample { - fn new(index: u32) -> Self { - Self { - index, - dimension: 2, - } - } - +impl PixelSample { pub fn get_1d(&mut self) -> f32 { - if self.dimension as usize >= PRIMES.len() { - self.dimension = 2; - } - let dim = self.dimension as u64; + let index = self.get_sample_index(); self.dimension += 1; - scrambled_radical_inverse(dim as usize, self.index, mix_bits(1 + (dim << 4))) + + let hash = hash64_with_seed(self.dimension.to_ne_bytes(), self.seed as u64); + let seed = hash as u32; + + sobol_sample0(index, seed) } pub fn get_2d(&mut self) -> (f32, f32) { - if self.dimension as usize + 1 >= PRIMES.len() { - self.dimension = 2; - } - let mut dim = self.dimension as u64; + let index = self.get_sample_index(); self.dimension += 2; - let x = scrambled_radical_inverse(dim as usize, self.index, mix_bits(1 + (dim << 4))); - dim += 1; - let y = scrambled_radical_inverse(dim as usize, self.index, mix_bits(1 + (dim << 4))); - (x, y) + + let hash = hash64_with_seed(self.dimension.to_ne_bytes(), self.seed as u64); + + let lo = hash as u32; + let hi = (hash >> 32) as u32; + + (sobol_sample0(index, lo), sobol_sample1(index, hi)) + } + + fn get_sample_index(&self) -> u32 { + let mut index: u32 = 0; + let last_digit = self.log2_samples & 1; + + for i in (last_digit..self.base4_digits).rev() { + let digit_shift = 2 * i - last_digit; + let digit = (self.morton_index >> digit_shift) & 3; + let higher_bits = self.morton_index >> (digit_shift + 2); + let mixed = mix_bits((higher_bits ^ (0x55555555 * self.dimension)) as u64); + let p = (mixed >> 24) as usize % PERMUTATIONS.len(); + let digit = PERMUTATIONS[p][digit as usize]; + index |= (digit as u32) << digit_shift; + } + + if last_digit != 0 { + let digit = self.morton_index & 1; + let higher_bits = self.morton_index >> 1; + let mixed = mix_bits((higher_bits ^ (0x55555555 * self.dimension)) as u64) as u32; + index |= (digit ^ mixed) & 1; + } + + index } } -fn egcd(a: i16, b: i16) -> (i16, i16) { - if a == 0 { - (0, 1) - } else { - let (x, y) = egcd(b % a, a); - (y - (b / a) * x, x) - } +fn sobol_sample0(a: u32, seed: u32) -> f32 { + let v = a.reverse_bits(); + f32::min(fast_owen(v, seed) as f32 * ONE_OVER_MAX, 1.0 - f32::EPSILON) } -fn mod_inverse(a: u8, m: u8) -> u8 { - let m = m as i16; - let (x, _) = egcd(a as i16, m); - ((x % m + m) % m) as u8 -} - -fn radical_inverse(base_index: usize, mut i: u32) -> f32 { - let base = PRIMES[base_index]; - let inv_base = 1.0 / base as f32; - let mut inv_base_m = 1.0; - let mut reversed = 0; - - while i > 0 { - let next = i / base; - let digit = i - next * base; - reversed = reversed * base + digit; - inv_base_m *= inv_base; - i = next - } - - f32::min(reversed as f32 * inv_base_m, 1.0 - f32::EPSILON) -} - -fn inverse_radical_inverse(mut inverse: u32, base: u32, digits: u32) -> u32 { - let mut i = 0; - for _ in 0..digits { - let digit = inverse % base; - inverse /= base; - i = i * base + digit; - } - i -} - -fn scrambled_radical_inverse(base_index: usize, mut i: u32, hash: u64) -> f32 { - let base = PRIMES[base_index]; - let inv_base = 1.0 / base as f32; - let mut inv_base_m = 1.0; - let mut reversed = 0; - - while 1.0 - inv_base_m < 1.0 { - let next = i / base; - let digit_value = i - next * base; - let digit_hash = mix_bits(hash ^ digit_value as u64); - let digit = permutation_element(digit_value, base, digit_hash as u32); - reversed = reversed * base + digit; - inv_base_m *= inv_base; - i = next; - } - - f32::min(reversed as f32 * inv_base_m, 1.0 - f32::EPSILON) -} - -// magic. Taken from pbrt/util/hash.h in pbrt -#[inline(always)] -fn mix_bits(mut val: u64) -> u64 { - val ^= val >> 31; - val *= 0x7fb5d329728ea185; - val ^= val >> 27; - val *= 0x81dadef4bc2dd44d; - val ^ (val >> 33) -} - -// more magic. Taken from pbrt/util/math.h in pbrt -#[inline(always)] -fn permutation_element(mut i: u32, n: u32, seed: u32) -> u32 { - let mut w = n - 1; - w |= w >> 1; - w |= w >> 2; - w |= w >> 4; - w |= w >> 8; - w |= w >> 16; - loop { - i ^= seed; - i *= 0xe170893d; - i ^= seed >> 16; - i ^= (i & w) >> 4; - i ^= seed >> 8; - i *= 0x0929eb3f; - i ^= seed >> 23; - i ^= (i & w) >> 1; - i *= 1 | seed >> 27; - i *= 0x6935fa69; - i ^= (i & w) >> 11; - i *= 0x74dcb303; - i ^= (i & w) >> 2; - i *= 0x9e501cc3; - i ^= (i & w) >> 2; - i *= 0xc860a3df; - i &= w; - i ^= i >> 5; - if i < n { +fn sobol_sample1(mut a: u32, seed: u32) -> f32 { + let mut v: u32 = 0; + for m in SOBOL_MATRIX { + if a & 1 != 0 { + v ^= m; + } + a >>= 1; + if a == 0 { break; } } - (i + seed) % n + f32::min(fast_owen(v, seed) as f32 * ONE_OVER_MAX, 1.0 - f32::EPSILON) +} + +// magic. taken from pbrt. +#[inline(always)] +fn fast_owen(mut v: u32, seed: u32) -> u32 { + v = v.reverse_bits(); + v ^= v.wrapping_mul(0x3d20adea); + v = v.wrapping_add(seed); + v = v.wrapping_mul((seed >> 16) | 1); + v ^= v.wrapping_mul(0x05526c56); + v ^= v.wrapping_mul(0x53a22864); + v.reverse_bits() +} + +// more magic. also taken from pbrt. +#[inline(always)] +fn mix_bits(mut v: u64) -> u64 { + v ^= v >> 31; + v = v.wrapping_mul(0x7fb5d329728ea185); + v ^= v >> 27; + v = v.wrapping_mul(0x81dadef4bc2dd44d); + v ^ (v >> 33) +} + +fn separate(mut v: u32) -> u32 { + v &= 0x0000ffff; + v = (v | (v << 8)) & 0x00ff00ff; + v = (v | (v << 4)) & 0x0f0f0f0f; + v = (v | (v << 2)) & 0x33333333; + v = (v | (v << 1)) & 0x55555555; + v +} + +fn morton_2d(x: u32, y: u32) -> u32 { + (separate(y) << 1) | separate(x) } diff --git a/examples/basic.py b/examples/basic.py index 11e6bdd..4810bb0 100755 --- a/examples/basic.py +++ b/examples/basic.py @@ -4,7 +4,7 @@ from mechthild import Scene, Shape, Camera WIDTH = 800 HEIGHT = 600 -SAMPLES = 20 +SAMPLES = 32 s1 = Shape.sphere((0.75, 1.0, 1.0), 1.0) s2 = Shape.sphere((-1.5, 2.0, -2.0), 2.0)