mirror of
https://github.com/storytold/spark.git
synced 2026-10-12 23:49:01 +00:00
Implement float32 sorting option (#129)
* Add commented benchmarkSort to compare JS vs Wasm sorting. * Add 32-bit float sort via SparkViewpooint.sort32, using 2-pass radix-65536 sort. Turn on sort32 by default for examples/editor. Implemented Rust sort and JS sort, updated benchmarking code to include float16 and float32 sort. * Add sort32 to docs/spark-viewpoint.md.
This commit is contained in:
committed by
GitHub
parent
2ab9ee1e78
commit
0bfb6c8963
@@ -4,7 +4,7 @@ use js_sys::{Float32Array, Uint16Array, Uint32Array};
|
||||
use wasm_bindgen::prelude::*;
|
||||
|
||||
mod sort;
|
||||
use sort::{sort_internal, SortBuffers};
|
||||
use sort::{sort_internal, SortBuffers, sort32_internal, Sort32Buffers};
|
||||
|
||||
mod raycast;
|
||||
use raycast::{raycast_ellipsoids, raycast_spheres};
|
||||
@@ -13,6 +13,7 @@ const RAYCAST_BUFFER_COUNT: u32 = 65536;
|
||||
|
||||
thread_local! {
|
||||
static SORT_BUFFERS: RefCell<SortBuffers> = RefCell::new(SortBuffers::default());
|
||||
static SORT32_BUFFERS: RefCell<Sort32Buffers> = RefCell::new(Sort32Buffers::default());
|
||||
static RAYCAST_BUFFER: RefCell<Vec<u32>> = RefCell::new(vec![0; RAYCAST_BUFFER_COUNT as usize * 4]);
|
||||
}
|
||||
|
||||
@@ -45,6 +46,35 @@ pub fn sort_splats(
|
||||
active_splats
|
||||
}
|
||||
|
||||
#[wasm_bindgen]
|
||||
pub fn sort32_splats(
|
||||
num_splats: u32, readback: Uint32Array, ordering: Uint32Array,
|
||||
) -> u32 {
|
||||
let max_splats = readback.length() as usize;
|
||||
|
||||
let active_splats = SORT32_BUFFERS.with_borrow_mut(|buffers| {
|
||||
buffers.ensure_size(max_splats);
|
||||
let sub_readback = readback.subarray(0, num_splats);
|
||||
sub_readback.copy_to(&mut buffers.readback[..num_splats as usize]);
|
||||
|
||||
let active_splats = match sort32_internal(buffers, max_splats, num_splats as usize) {
|
||||
Ok(active_splats) => active_splats,
|
||||
Err(err) => {
|
||||
wasm_bindgen::throw_str(&format!("{}", err));
|
||||
}
|
||||
};
|
||||
|
||||
if active_splats > 0 {
|
||||
// Copy out ordering result
|
||||
let subarray = &buffers.ordering[..active_splats as usize];
|
||||
ordering.subarray(0, active_splats).copy_from(&subarray);
|
||||
}
|
||||
active_splats
|
||||
});
|
||||
|
||||
active_splats
|
||||
}
|
||||
|
||||
#[wasm_bindgen]
|
||||
pub fn raycast_splats(
|
||||
origin_x: f32, origin_y: f32, origin_z: f32,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use anyhow::anyhow;
|
||||
|
||||
const DEPTH_INFINITY: u32 = 0x7c00;
|
||||
const DEPTH_SIZE: usize = DEPTH_INFINITY as usize + 1;
|
||||
const DEPTH_INFINITY_F16: u32 = 0x7c00;
|
||||
const DEPTH_SIZE_F16: usize = DEPTH_INFINITY_F16 as usize + 1;
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct SortBuffers {
|
||||
@@ -18,8 +18,8 @@ impl SortBuffers {
|
||||
if self.ordering.len() < max_splats {
|
||||
self.ordering.resize(max_splats, 0);
|
||||
}
|
||||
if self.buckets.len() < DEPTH_SIZE {
|
||||
self.buckets.resize(DEPTH_SIZE, 0);
|
||||
if self.buckets.len() < DEPTH_SIZE_F16 {
|
||||
self.buckets.resize(DEPTH_SIZE_F16, 0);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -30,11 +30,11 @@ pub fn sort_internal(buffers: &mut SortBuffers, num_splats: usize) -> anyhow::Re
|
||||
|
||||
// Set the bucket counts to zero
|
||||
buckets.clear();
|
||||
buckets.resize(DEPTH_SIZE, 0);
|
||||
buckets.resize(DEPTH_SIZE_F16, 0);
|
||||
|
||||
// Count the number of splats in each bucket
|
||||
for &metric in readback.iter() {
|
||||
if (metric as u32) < DEPTH_INFINITY {
|
||||
if (metric as u32) < DEPTH_INFINITY_F16 {
|
||||
buckets[metric as usize] += 1;
|
||||
}
|
||||
}
|
||||
@@ -49,7 +49,7 @@ pub fn sort_internal(buffers: &mut SortBuffers, num_splats: usize) -> anyhow::Re
|
||||
|
||||
// Write out splat indices at the right location using bucket offsets
|
||||
for (index, &metric) in readback.iter().enumerate() {
|
||||
if (metric as u32) < DEPTH_INFINITY {
|
||||
if (metric as u32) < DEPTH_INFINITY_F16 {
|
||||
ordering[buckets[metric as usize] as usize] = index as u32;
|
||||
buckets[metric as usize] += 1;
|
||||
}
|
||||
@@ -65,3 +65,111 @@ pub fn sort_internal(buffers: &mut SortBuffers, num_splats: usize) -> anyhow::Re
|
||||
}
|
||||
Ok(active_splats)
|
||||
}
|
||||
|
||||
const DEPTH_INFINITY_F32: u32 = 0x7f800000;
|
||||
const RADIX_BASE: usize = 1 << 16; // 65536
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct Sort32Buffers {
|
||||
/// raw f32 bit‑patterns (one per splat)
|
||||
pub readback: Vec<u32>,
|
||||
/// output indices
|
||||
pub ordering: Vec<u32>,
|
||||
/// bucket counts / offsets (length == RADIX_BASE)
|
||||
pub buckets16: Vec<u32>,
|
||||
/// scratch space for indices
|
||||
pub scratch: Vec<u32>,
|
||||
}
|
||||
|
||||
impl Sort32Buffers {
|
||||
/// ensure all internal buffers are large enough for up to `max_splats`
|
||||
pub fn ensure_size(&mut self, max_splats: usize) {
|
||||
if self.readback.len() < max_splats {
|
||||
self.readback.resize(max_splats, 0);
|
||||
}
|
||||
if self.ordering.len() < max_splats {
|
||||
self.ordering.resize(max_splats, 0);
|
||||
}
|
||||
if self.scratch.len() < max_splats {
|
||||
self.scratch.resize(max_splats, 0);
|
||||
}
|
||||
if self.buckets16.len() < RADIX_BASE {
|
||||
self.buckets16.resize(RADIX_BASE, 0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Two‑pass radix sort (base 2¹⁶) of 32‑bit float bit‑patterns,
|
||||
/// descending order (largest keys first). Mirrors the JS `sort32Splats`.
|
||||
pub fn sort32_internal(
|
||||
buffers: &mut Sort32Buffers,
|
||||
max_splats: usize,
|
||||
num_splats: usize,
|
||||
) -> anyhow::Result<u32> {
|
||||
// make sure our buffers can hold `max_splats`
|
||||
buffers.ensure_size(max_splats);
|
||||
|
||||
let Sort32Buffers { readback, ordering, buckets16, scratch } = buffers;
|
||||
let keys = &readback[..num_splats];
|
||||
|
||||
// ——— Pass #1: bucket by inv(low 16 bits) ———
|
||||
buckets16.fill(0);
|
||||
for &key in keys.iter() {
|
||||
if key < DEPTH_INFINITY_F32 {
|
||||
let inv = !key;
|
||||
buckets16[(inv & 0xFFFF) as usize] += 1;
|
||||
}
|
||||
}
|
||||
// exclusive prefix‑sum → starting offsets
|
||||
let mut total: u32 = 0;
|
||||
for slot in buckets16.iter_mut() {
|
||||
let cnt = *slot;
|
||||
*slot = total;
|
||||
total = total.wrapping_add(cnt);
|
||||
}
|
||||
let active_splats = total;
|
||||
|
||||
// scatter into scratch by low bits of inv
|
||||
for (i, &key) in keys.iter().enumerate() {
|
||||
if key < DEPTH_INFINITY_F32 {
|
||||
let inv = !key;
|
||||
let lo = (inv & 0xFFFF) as usize;
|
||||
scratch[buckets16[lo] as usize] = i as u32;
|
||||
buckets16[lo] += 1;
|
||||
}
|
||||
}
|
||||
|
||||
// ——— Pass #2: bucket by inv(high 16 bits) ———
|
||||
buckets16.fill(0);
|
||||
for &idx in scratch.iter().take(active_splats as usize) {
|
||||
let key = keys[idx as usize];
|
||||
let inv = !key;
|
||||
buckets16[(inv >> 16) as usize] += 1;
|
||||
}
|
||||
// exclusive prefix‑sum again
|
||||
let mut sum: u32 = 0;
|
||||
for slot in buckets16.iter_mut() {
|
||||
let cnt = *slot;
|
||||
*slot = sum;
|
||||
sum = sum.wrapping_add(cnt);
|
||||
}
|
||||
// scatter into final ordering by high bits of inv
|
||||
for &idx in scratch.iter().take(active_splats as usize) {
|
||||
let key = keys[idx as usize];
|
||||
let inv = !key;
|
||||
let hi = (inv >> 16) as usize;
|
||||
ordering[buckets16[hi] as usize] = idx;
|
||||
buckets16[hi] += 1;
|
||||
}
|
||||
|
||||
// sanity‑check: last bucket should have consumed all entries
|
||||
if buckets16[RADIX_BASE - 1] != active_splats {
|
||||
return Err(anyhow!(
|
||||
"Expected {} active splats but got {}",
|
||||
active_splats,
|
||||
buckets16[RADIX_BASE - 1]
|
||||
));
|
||||
}
|
||||
|
||||
Ok(active_splats)
|
||||
}
|
||||
Reference in New Issue
Block a user