Initial commit

rust/forge-internal-rs/src/lib.rs

@@ -0,0 +1,154 @@
+
+use std::cell::RefCell;
+use js_sys::{Array, ArrayBuffer, Float32Array, Object, Reflect, Uint16Array, Uint32Array, Uint8Array};
+use wasm_bindgen::prelude::*;
+
+use wlg::decode12_packed;
+
+mod sort;
+use sort::{old_sort_internal, OldSortBuffers, sort_internal, SortBuffers};
+
+mod raycast;
+use raycast::{raycast_ellipsoids, raycast_spheres};
+
+const RAYCAST_BUFFER_COUNT: u32 = 65536;
+
+thread_local! {
+    static SORT_BUFFERS: RefCell<SortBuffers> = RefCell::new(SortBuffers::default());
+    static OLD_SORT_BUFFERS: RefCell<OldSortBuffers> = RefCell::new(OldSortBuffers::default());
+    static RAYCAST_BUFFER: RefCell<Vec<u32>> = RefCell::new(vec![0; RAYCAST_BUFFER_COUNT as usize * 4]);
+}
+
+#[wasm_bindgen]
+pub fn decode_wlg(bytes: ArrayBuffer, tex_width: u32, tex_height: u32) -> Object {
+    let mut bytes = Uint8Array::new(&bytes).to_vec();
+    let (_settings, packed_splats) = match decode12_packed(&mut bytes) {
+        Ok(result) => result,
+        Err(err) => {
+            wasm_bindgen::throw_str(&format!("{}", err));
+        }
+    };
+
+    let num_splats = packed_splats.0.len() / 4;
+    let max_splats = if tex_width != 0 && tex_height != 0 {            
+        let width = tex_width as usize;
+        let height = num_splats.div_ceil(width).min(tex_height as usize);
+        let depth = num_splats.div_ceil(width * height);
+        width * height * depth
+    } else {
+        num_splats
+    };
+
+    let packed = Uint32Array::new_with_length(max_splats as u32 * 4);
+    let packed_slice = packed.subarray(0, packed_splats.0.len() as u32);
+    packed_slice.copy_from(&packed_splats.0);
+
+    let result = Object::new();
+    Reflect::set(&result, &JsValue::from_str("numSplats"), &JsValue::from_f64(num_splats as f64)).unwrap();
+    Reflect::set(&result, &JsValue::from_str("packedSplats"), &JsValue::from(packed)).unwrap();
+    result
+}
+
+#[wasm_bindgen]
+pub fn old_sort_splats(
+    max_splats: u32, total_splats: u32, readback: Array, ordering: Uint32Array,
+) -> u32 {
+    let max_splats = max_splats as usize;
+    let total_splats = total_splats as usize;
+    let num_layers = readback.length() as usize;
+    let layer_size = readback.get(0).dyn_into::<Uint8Array>().unwrap().length() as usize / 4;
+
+    let active_splats = OLD_SORT_BUFFERS.with_borrow_mut(|buffers| {
+        buffers.ensure_size(max_splats);
+
+        // Copy the readback data layers into a contiguous buffer
+        let mut layer_base = 0;
+        for layer in 0..num_layers {
+            let layer_count = layer_size.min(total_splats - layer_base);
+            if layer_count > 0 {
+                let layer_buffer = readback.get(layer as u32).dyn_into::<Uint8Array>().unwrap().buffer();
+                let layer_uint32 = Uint32Array::new_with_byte_offset_and_length(&layer_buffer, 0, layer_count as u32);
+                layer_uint32.copy_to(&mut buffers.readback[layer_base..layer_base + layer_count]);
+            }
+            layer_base += layer_count;
+        }
+
+        let active_splats = match old_sort_internal(buffers, total_splats) {
+            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 sort_splats(
+    num_splats: u32, readback: Uint16Array, ordering: Uint32Array,
+) -> u32 {
+    let max_splats = readback.length() as usize;
+
+    let active_splats = SORT_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 sort_internal(buffers, 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,
+    dir_x: f32, dir_y: f32, dir_z: f32,
+    near: f32, far: f32,
+    num_splats: u32, packed_splats: Uint32Array,
+    raycast_ellipsoid: bool,
+) -> Float32Array {
+    let mut distances = Vec::<f32>::new();
+
+    _ = RAYCAST_BUFFER.with_borrow_mut(|buffer| {
+        let mut base = 0;
+        while base < num_splats {
+            let chunk_size = RAYCAST_BUFFER_COUNT.min(num_splats - base);
+            let subarray = packed_splats.subarray(4 * base, 4 * (base + chunk_size));
+            let subbuffer = &mut buffer[0..(4 * chunk_size as usize)];
+            subarray.copy_to(subbuffer);
+
+            if raycast_ellipsoid {
+                raycast_ellipsoids(subbuffer, &mut distances, [origin_x, origin_y, origin_z], [dir_x, dir_y, dir_z], near, far);
+            } else {
+                raycast_spheres(subbuffer, &mut distances, [origin_x, origin_y, origin_z], [dir_x, dir_y, dir_z], near, far);
+            }
+
+            base += chunk_size;
+        }
+    });
+
+    let output = Float32Array::new_with_length(distances.len() as u32);
+    output.copy_from(&distances);
+    output
+}

rust/forge-internal-rs/src/raycast.rs

@@ -0,0 +1,170 @@
+use half::f16;
+use wlg::decode_scale;
+
+const MIN_OPACITY: f32 = 0.1;
+
+pub fn raycast_spheres(
+    buffer: &[u32], distances: &mut Vec<f32>, 
+    origin: [f32; 3], dir: [f32; 3], near: f32, far: f32,
+) {
+    let quad_a = vec3_dot(dir, dir);
+
+    for packed in buffer.chunks(4) {
+        let opacity = ((packed[0] >> 24) as u8) as f32 / 255.0;
+        if opacity < MIN_OPACITY {
+            continue;
+        }
+
+        let origin = vec3_sub(origin, extract_center(packed));
+        let scale = extract_scale(packed);
+        
+        // Model the Gsplat as a sphere for faster approximate raycasting
+        let radius = (scale[0] + scale[1] + scale[2]) / 3.0;
+
+        let quad_b = vec3_dot(dir, origin);
+        let quad_c = vec3_dot(origin, origin) - radius * radius;
+        let discriminant = quad_b * quad_b - quad_a * quad_c;
+        if discriminant < 0.0 {
+            continue;
+        }
+
+        let t = (-quad_b - discriminant.sqrt()) / quad_a;
+        if t >= near && t <= far {
+            distances.push(t);
+        }
+    }
+}
+
+pub fn raycast_ellipsoids(
+    buffer: &[u32], distances: &mut Vec<f32>, 
+    origin: [f32; 3], dir: [f32; 3], near: f32, far: f32,
+) {
+    for packed in buffer.chunks(4) {
+        let opacity = ((packed[0] >> 24) as u8) as f32 / 255.0;
+        if opacity < MIN_OPACITY {
+            continue;
+        }
+    
+        let origin = vec3_sub(origin, extract_center(packed));
+        let scale = extract_scale(packed);
+        let quat = extract_quat(packed);
+        let inv_quat = [-quat[0], -quat[1], -quat[2], quat[3]];
+
+        // Model the Gsplat as an ellipsoid for higher quality raycasting
+        let local_origin = quat_vec(inv_quat, origin);
+        let local_dir = quat_vec(inv_quat, dir);
+
+        let min_scale = scale[0].max(scale[1]).max(scale[2]) * 0.01;
+        let t = if scale[2] < min_scale {
+            // Treat it as a flat elliptical disk
+            if local_dir[2].abs() < 1e-6 {
+                continue;
+            }
+            let t = -local_origin[2] / local_dir[2];
+            let p_x = local_origin[0] + t * local_dir[0];
+            let p_y = local_origin[1] + t * local_dir[1];
+            if sqr(p_x / scale[0]) + sqr(p_y / scale[1]) > 1.0 {
+                continue;
+            }
+            t
+        } else if scale[1] < min_scale {
+            // Treat it as a flat elliptical disk
+            if local_dir[1].abs() < 1e-6 {
+                continue;
+            }
+            let t = -local_origin[1] / local_dir[1];
+            let p_x = local_origin[0] + t * local_dir[0];
+            let p_z = local_origin[2] + t * local_dir[2];
+            if sqr(p_x / scale[0]) + sqr(p_z / scale[2]) > 1.0 {
+                continue;
+            }
+            t
+        } else if scale[0] < min_scale {
+            // Treat it as a flat elliptical disk
+            if local_dir[0].abs() < 1e-6 {
+                continue;
+            }
+            let t = -local_origin[0] / local_dir[0];
+            let p_y = local_origin[1] + t * local_dir[1];
+            let p_z = local_origin[2] + t * local_dir[2];
+            if sqr(p_y / scale[1]) + sqr(p_z / scale[2]) > 1.0 {
+                continue;
+            }
+            t
+        } else {
+            let inv_scale = [1.0 / scale[0], 1.0 / scale[1], 1.0 / scale[2]];
+            let local_origin = vec3_mul(local_origin, inv_scale);
+            let local_dir = vec3_mul(local_dir, inv_scale);
+
+            let a = vec3_dot(local_dir, local_dir);
+            let b = vec3_dot(local_origin, local_dir);
+            let c = vec3_dot(local_origin, local_origin) - 1.0;
+            let discriminant = b * b - a * c;
+            if discriminant < 0.0 {
+                continue;
+            }
+
+            (-b - discriminant.sqrt()) / a
+        };
+        if t >= near && t <= far {
+            distances.push(t);
+        }
+    }
+}
+
+fn extract_center(packed: &[u32]) -> [f32; 3] {
+    let x = f16::from_bits(packed[1] as u16).to_f32();
+    let y = f16::from_bits((packed[1] >> 16) as u16).to_f32();
+    let z = f16::from_bits(packed[2] as u16).to_f32();
+    [x, y, z]
+}
+
+fn extract_scale(packed: &[u32]) -> [f32; 3] {
+    let scale_x = decode_scale(packed[3] as u8);
+    let scale_y = decode_scale((packed[3] >> 8) as u8);
+    let scale_z = decode_scale((packed[3] >> 16) as u8);
+    [scale_x, scale_y, scale_z]
+}
+
+fn extract_quat(packed: &[u32]) -> [f32; 4] {
+    let quat_x = ((packed[2] >> 16) as i8) as f32 / 127.0;
+    let quat_y = ((packed[2] >> 24) as i8) as f32 / 127.0;
+    let quat_z = ((packed[3] >> 24) as i8) as f32 / 127.0;
+    let quat_w = (1.0 - quat_x * quat_x - quat_y * quat_y - quat_z * quat_z).max(0.0).sqrt();
+    [quat_x, quat_y, quat_z, quat_w]
+}
+
+fn sqr(x: f32) -> f32 {
+    x * x
+}
+
+fn vec3_sub(a: [f32; 3], b: [f32; 3]) -> [f32; 3] {
+    [a[0] - b[0], a[1] - b[1], a[2] - b[2]]
+}
+
+fn vec3_mul(a: [f32; 3], b: [f32; 3]) -> [f32; 3] {
+    [a[0] * b[0], a[1] * b[1], a[2] * b[2]]
+}
+
+fn vec3_dot(a: [f32; 3], b: [f32; 3]) -> f32 {
+    a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
+}
+
+fn vec3_cross(a: [f32; 3], b: [f32; 3]) -> [f32; 3] {
+    [
+        a[1] * b[2] - a[2] * b[1],
+        a[2] * b[0] - a[0] * b[2],
+        a[0] * b[1] - a[1] * b[0],
+    ]
+}
+
+fn quat_vec(q: [f32; 4], v: [f32; 3]) -> [f32; 3] {
+    let q_vec = [q[0], q[1], q[2]];
+    let uv = vec3_cross(q_vec, v);
+    let uuv = vec3_cross(q_vec, uv);
+    [
+        v[0] + 2.0 * (q[3] * uv[0] + uuv[0]),
+        v[1] + 2.0 * (q[3] * uv[1] + uuv[1]),
+        v[2] + 2.0 * (q[3] * uv[2] + uuv[2]),
+    ]
+}

rust/forge-internal-rs/src/sort.rs

@@ -0,0 +1,133 @@
+use anyhow::anyhow;
+
+const DEPTH_INFINITY: u32 = 0x7c00;
+const DEPTH_SIZE: usize = DEPTH_INFINITY as usize + 1;
+
+#[derive(Default)]
+pub struct OldSortBuffers {
+    pub readback: Vec<u32>,
+    pub ordering: Vec<u32>,
+    pub buckets: Vec<u32>,
+}
+
+impl OldSortBuffers {
+    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.buckets.len() < DEPTH_SIZE {
+            self.buckets.resize(DEPTH_SIZE, 0);
+        }
+    }
+}
+
+pub fn old_sort_internal(buffers: &mut OldSortBuffers, total_splats: usize) -> anyhow::Result<u32> {
+    let OldSortBuffers { readback, ordering, buckets } = buffers;
+    let readback = &readback[..total_splats];
+
+    // Set the bucket counts to zero
+    buckets.clear();
+    buckets.resize(DEPTH_SIZE, 0);
+
+    // Count the number of splats in each bucket
+    for &readback_u32 in readback.iter() {
+        let priority = readback_u32 & 0x7FFF;
+        if priority < DEPTH_INFINITY {
+            buckets[priority as usize] += 1;
+        }
+    }
+
+    // Compute bucket starting offset
+    let mut active_splats = 0;
+    for count in buckets.iter_mut() {
+        let new_total = active_splats + *count;
+        *count = active_splats;
+        active_splats = new_total;
+    }
+
+    // Write out splat indices at the right location using bucket offsets
+    for (index, &readback_u32) in readback.iter().enumerate() {
+        let priority = readback_u32 & 0x7FFF;
+        if priority < DEPTH_INFINITY {
+            ordering[buckets[priority as usize] as usize] = index as u32;
+            buckets[priority as usize] += 1;
+        }
+    }
+
+    // Sanity check
+    if buckets[DEPTH_SIZE - 1] != active_splats {
+        return Err(anyhow!(
+            "Expected {} active splats but got {}",
+            active_splats,
+            buckets[DEPTH_SIZE - 1]
+        ));
+    }
+    Ok(active_splats)
+}
+
+
+#[derive(Default)]
+pub struct SortBuffers {
+    pub readback: Vec<u16>,
+    pub ordering: Vec<u32>,
+    pub buckets: Vec<u32>,
+}
+
+impl SortBuffers {
+    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.buckets.len() < DEPTH_SIZE {
+            self.buckets.resize(DEPTH_SIZE, 0);
+        }
+    }
+}
+
+pub fn sort_internal(buffers: &mut SortBuffers, num_splats: usize) -> anyhow::Result<u32> {
+    let SortBuffers { readback, ordering, buckets } = buffers;
+    let readback = &readback[..num_splats];
+
+    // Set the bucket counts to zero
+    buckets.clear();
+    buckets.resize(DEPTH_SIZE, 0);
+
+    // Count the number of splats in each bucket
+    for &metric in readback.iter() {
+        if (metric as u32) < DEPTH_INFINITY {
+            buckets[metric as usize] += 1;
+        }
+    }
+
+    // Compute bucket starting offset
+    let mut active_splats = 0;
+    for count in buckets.iter_mut().rev().skip(1) {
+        let new_total = active_splats + *count;
+        *count = active_splats;
+        active_splats = new_total;
+    }
+
+    // Write out splat indices at the right location using bucket offsets
+    for (index, &metric) in readback.iter().enumerate() {
+        if (metric as u32) < DEPTH_INFINITY {
+            ordering[buckets[metric as usize] as usize] = index as u32;
+            buckets[metric as usize] += 1;
+        }
+    }
+
+    // Sanity check
+    if buckets[0] != active_splats {
+        return Err(anyhow!(
+            "Expected {} active splats but got {}",
+            active_splats,
+            buckets[0]
+        ));
+    }
+    Ok(active_splats)
+}