Skip to content

Commit

Permalink
checkin
Browse files Browse the repository at this point in the history
  • Loading branch information
@mosure
mosure committed Dec 28, 2023
1 parent d071f6d commit 166a7dc
Show file tree
Hide file tree
Showing 6 changed files with 191 additions and 95 deletions.
4 changes: 2 additions & 2 deletions Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,8 @@ default = [

# "packed",
"planar",
"f32",
# "f16",
# "f32",
"f16",

"query_select",
"query_sparse",
Expand Down
120 changes: 61 additions & 59 deletions src/gaussian/cloud.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,10 +56,10 @@ pub struct GaussianCloud {
#[cfg(feature = "f16")]
pub rotation_scale_opacity_packed128: Vec<RotationScaleOpacityPacked128>,

#[cfg(not(feature = "f16"))]
pub rotation: Vec<Rotation>,
#[cfg(not(feature = "f16"))]
pub scale_opacity: Vec<ScaleOpacity>,
// #[cfg(feature = "f32")]
// pub rotation: Vec<Rotation>,
// #[cfg(feature = "f32")]
// pub scale_opacity: Vec<ScaleOpacity>,
}

impl GaussianCloud {
Expand Down Expand Up @@ -100,55 +100,51 @@ impl GaussianCloud {
}


pub fn rotation(&self, index: usize) -> &[f32; 4] {
#[cfg(feature = "f16")]
return &self.rotation_scale_opacity_packed128[index].rotation;
// pub fn rotation(&self, index: usize) -> &[f32; 4] {
// #[cfg(feature = "f16")]
// return &self.rotation_scale_opacity_packed128[index].rotation;

#[cfg(not(feature = "f16"))]
return &self.rotation[index].rotation;
}
// #[cfg(feature = "f32")]
// return &self.rotation[index].rotation;
// }

pub fn rotation_mut(&mut self, index: usize) -> &mut [f32; 4] {
#[cfg(feature = "f16")]
return &mut self.rotation_scale_opacity_packed128[index].rotation;
// pub fn rotation_mut(&mut self, index: usize) -> &mut [f32; 4] {
// #[cfg(feature = "f16")]
// return &mut self.rotation_scale_opacity_packed128[index].rotation;

#[cfg(not(feature = "f16"))]
return &mut self.rotation[index].rotation;
}
// #[cfg(feature = "f32")]
// return &mut self.rotation[index].rotation;
// }


pub fn scale(&self, index: usize) -> &[f32; 3] {
#[cfg(feature = "f16")]
return &self.rotation_scale_opacity_packed128[index].scale;
// pub fn scale(&self, index: usize) -> &[f32; 3] {
// #[cfg(feature = "f16")]
// return &self.rotation_scale_opacity_packed128[index].scale;

#[cfg(not(feature = "f16"))]
return &self.scale_opacity[index].scale;
}
// #[cfg(feature = "f32")]
// return &self.scale_opacity[index].scale;
// }

pub fn scale_mut(&mut self, index: usize) -> &mut [f32; 3] {
#[cfg(feature = "f16")]
return &mut self.rotation_scale_opacity_packed128[index].scale;
// pub fn scale_mut(&mut self, index: usize) -> &mut [f32; 3] {
// #[cfg(feature = "f16")]
// return &mut self.rotation_scale_opacity_packed128[index].scale;

#[cfg(not(feature = "f16"))]
return &mut self.scale_opacity[index].scale;
}
// #[cfg(feature = "f32")]
// return &mut self.scale_opacity[index].scale;
// }


#[cfg(feature = "f32")]
pub fn gaussian(&self, index: usize) -> Gaussian {
Gaussian {
position_visibility: self.position_visibility[index],
spherical_harmonic: self.spherical_harmonic[index],

#[cfg(feature = "f16")]
rotation_scale_opacity_packed128: self.rotation_scale_opacity_packed128[index],

#[cfg(not(feature = "f16"))]
rotation: self.rotation[index],
#[cfg(not(feature = "f16"))]
scale_opacity: self.scale_opacity[index],
}
}

#[cfg(feature = "f32")]
pub fn gaussian_iter(&self) -> impl Iterator<Item=Gaussian> + '_ {
self.position_visibility.iter()
.zip(self.spherical_harmonic.iter())
Expand All @@ -159,62 +155,68 @@ impl GaussianCloud {
position_visibility: *position_visibility,
spherical_harmonic: *spherical_harmonic,

#[cfg(feature = "f16")]
rotation_scale_opacity_packed128: *rotation_scale_opacity_packed128,

#[cfg(not(feature = "f16"))]
rotation: *rotation,
#[cfg(not(feature = "f16"))]
scale_opacity: *scale_opacity,
}
})
}


pub fn spherical_harmonic(&self, index: usize) -> &[f32; SH_COEFF_COUNT] {
&self.spherical_harmonic[index].coefficients
pub fn spherical_harmonic(&self, index: usize) -> &SphericalHarmonicCoefficients {
&self.spherical_harmonic[index]
}

pub fn spherical_harmonic_mut(&mut self, index: usize) -> &mut [f32; SH_COEFF_COUNT] {
&mut self.spherical_harmonic[index].coefficients
pub fn spherical_harmonic_mut(&mut self, index: usize) -> &mut SphericalHarmonicCoefficients {
&mut self.spherical_harmonic[index]
}
}


impl GaussianCloud {
#[cfg(feature = "f16")]
pub fn subset(&self, indicies: &[usize]) -> Self {
let mut position_visibility = Vec::with_capacity(indicies.len());
let mut spherical_harmonic = Vec::with_capacity(indicies.len());
let mut rotation = Vec::with_capacity(indicies.len());
let mut scale_opacity = Vec::with_capacity(indicies.len());
let mut rotation_scale_opacity_packed128 = Vec::with_capacity(indicies.len());

for &index in indicies.iter() {
position_visibility.push(self.position_visibility[index]);
spherical_harmonic.push(self.spherical_harmonic[index]);

#[cfg(feature = "f16")]
rotation_scale_opacity_packed128.push(self.rotation_scale_opacity_packed128[index]);
}

#[cfg(not(feature = "f16"))]
Self {
position_visibility,
spherical_harmonic,
rotation_scale_opacity_packed128,
}
}

#[cfg(feature = "f32")]
pub fn subset(&self, indicies: &[usize]) -> Self {
let mut position_visibility = Vec::with_capacity(indicies.len());
let mut spherical_harmonic = Vec::with_capacity(indicies.len());
let mut rotation = Vec::with_capacity(indicies.len());
let mut scale_opacity = Vec::with_capacity(indicies.len());

let mut rotation_scale_opacity_packed128 = Vec::with_capacity(indicies.len());

for &index in indicies.iter() {
position_visibility.push(self.position_visibility[index]);
spherical_harmonic.push(self.spherical_harmonic[index]);
rotation.push(self.rotation[index]);
#[cfg(not(feature = "f16"))]
scale_opacity.push(self.scale_opacity[index]);
}

Self {
position_visibility,
spherical_harmonic,

#[cfg(feature = "f16")]
rotation_scale_opacity_packed128,

#[cfg(not(feature = "f16"))]
rotation,
#[cfg(not(feature = "f16"))]
scale_opacity,
}
}

#[cfg(feature = "f32")]
pub fn to_packed(&self) -> Vec<Gaussian> {
let mut gaussians = Vec::with_capacity(self.len());

Expand All @@ -241,9 +243,9 @@ impl GaussianCloud {
#[cfg(feature = "f16")]
rotation_scale_opacity_packed128.push(gaussian.rotation_scale_opacity_packed128);

#[cfg(not(feature = "f16"))]
#[cfg(feature = "f32")]
rotation.push(gaussian.rotation);
#[cfg(not(feature = "f16"))]
#[cfg(feature = "f32")]
scale_opacity.push(gaussian.scale_opacity);
}

Expand All @@ -254,9 +256,9 @@ impl GaussianCloud {
#[cfg(feature = "f16")]
rotation_scale_opacity_packed128,

#[cfg(not(feature = "f16"))]
#[cfg(feature = "f32")]
rotation,
#[cfg(not(feature = "f16"))]
#[cfg(feature = "f32")]
scale_opacity,
}
}
Expand Down
65 changes: 40 additions & 25 deletions src/gaussian/f16.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
use rand::prelude::Distribution;
use std::marker::Copy;

use half::f16;
Expand All @@ -17,15 +16,14 @@ use serde::{
};


type f16_pod_t = [u8; 2];

#[derive(
Clone,
Debug,
Default,
Copy,
PartialEq,
Reflect,
ShaderType,
Pod,
Zeroable,
Serialize,
Expand All @@ -34,28 +32,22 @@ type f16_pod_t = [u8; 2];
#[repr(C)]
pub struct RotationScaleOpacityPacked128 {
#[reflect(ignore)]
pub rotation: [f16_pod_t; 4],
#[reflect(ignore)]
pub scale: [f16_pod_t; 3],
pub rotation: [u32; 2],
#[reflect(ignore)]
pub opacity: f16_pod_t,
pub scale_opacity: [u32; 2],
}

impl From<[f32; 8]> for RotationScaleOpacityPacked128 {
fn from(rotation_scale_opacity: [f32; 8]) -> Self {
Self {
rotation: [
f16::from_f32(rotation_scale_opacity[0]).to_bits().to_be_bytes(),
f16::from_f32(rotation_scale_opacity[1]).to_bits().to_be_bytes(),
f16::from_f32(rotation_scale_opacity[2]).to_bits().to_be_bytes(),
f16::from_f32(rotation_scale_opacity[3]).to_bits().to_be_bytes(),
pack_f32s_to_u32(rotation_scale_opacity[0], rotation_scale_opacity[1]),
pack_f32s_to_u32(rotation_scale_opacity[2], rotation_scale_opacity[3]),
],
scale: [
f16::from_f32(rotation_scale_opacity[4]).to_bits().to_be_bytes(),
f16::from_f32(rotation_scale_opacity[5]).to_bits().to_be_bytes(),
f16::from_f32(rotation_scale_opacity[6]).to_bits().to_be_bytes(),
scale_opacity: [
pack_f32s_to_u32(rotation_scale_opacity[4], rotation_scale_opacity[5]),
pack_f32s_to_u32(rotation_scale_opacity[6], rotation_scale_opacity[7]),
],
opacity: f16::from_f32(rotation_scale_opacity[7]).to_bits().to_be_bytes(),
}
}
}
Expand All @@ -64,17 +56,40 @@ impl From<[f16; 8]> for RotationScaleOpacityPacked128 {
fn from(rotation_scale_opacity: [f16; 8]) -> Self {
Self {
rotation: [
rotation_scale_opacity[0].to_bits().to_be_bytes(),
rotation_scale_opacity[1].to_bits().to_be_bytes(),
rotation_scale_opacity[2].to_bits().to_be_bytes(),
rotation_scale_opacity[3].to_bits().to_be_bytes(),
pack_f16s_to_u32(rotation_scale_opacity[0], rotation_scale_opacity[1]),
pack_f16s_to_u32(rotation_scale_opacity[2], rotation_scale_opacity[3]),
],
scale: [
rotation_scale_opacity[4].to_bits().to_be_bytes(),
rotation_scale_opacity[5].to_bits().to_be_bytes(),
rotation_scale_opacity[6].to_bits().to_be_bytes(),
scale_opacity: [
pack_f16s_to_u32(rotation_scale_opacity[4], rotation_scale_opacity[5]),
pack_f16s_to_u32(rotation_scale_opacity[6], rotation_scale_opacity[7]),
],
opacity: rotation_scale_opacity[7].to_bits().to_be_bytes(),
}
}
}

impl From<[u32; 4]> for RotationScaleOpacityPacked128 {
fn from(rotation_scale_opacity: [u32; 4]) -> Self {
Self {
rotation: [rotation_scale_opacity[0], rotation_scale_opacity[1]],
scale_opacity: [rotation_scale_opacity[2], rotation_scale_opacity[3]],
}
}
}


pub fn pack_f32s_to_u32(upper: f32, lower: f32) -> u32 {
let upper_bits = (f16::from_f32(upper).to_bits() as u32) << 16;
let lower_bits = f16::from_f32(lower).to_bits() as u32;
upper_bits | lower_bits
}

pub fn pack_f16s_to_u32(upper: f16, lower: f16) -> u32 {
let upper_bits = (upper.to_bits() as u32) << 16;
let lower_bits = lower.to_bits() as u32;
upper_bits | lower_bits
}




// TODO: add conversions from f32
26 changes: 22 additions & 4 deletions src/gaussian/rand.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,9 @@ use rand::{
Rng,
};

#[cfg(feature = "f16")]
use crate::gaussian::f16::pack_f32s_to_u32;

use crate::{
gaussian::{
cloud::GaussianCloud,
Expand Down Expand Up @@ -38,11 +41,26 @@ impl Distribution<Gaussian> for rand::distributions::Standard {
].into(),
spherical_harmonic: SphericalHarmonicCoefficients {
coefficients: {
let mut coefficients = [0.0; SH_COEFF_COUNT];
for coefficient in coefficients.iter_mut() {
*coefficient = rng.gen_range(-1.0..1.0);
#[cfg(feature = "f16")]
{
let mut coefficients: [u32; SH_COEFF_COUNT / 2];
for coefficient in coefficients.iter_mut() {
let upper = rng.gen_range(-1.0..1.0);
let lower = rng.gen_range(-1.0..1.0);

*coefficient = pack_f32s_to_u32(upper, lower);
}
coefficients
}

#[cfg(feature = "f32")]
{
let mut coefficients = [0.0; SH_COEFF_COUNT];
for coefficient in coefficients.iter_mut() {
*coefficient = rng.gen_range(-1.0..1.0);
}
coefficients
}
coefficients
},
},
}
Expand Down
2 changes: 2 additions & 0 deletions src/io/ply.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,8 @@ impl PropertyAccess for Gaussian {

if interleaved_idx < SH_COEFF_COUNT {
self.spherical_harmonic.coefficients[interleaved_idx] = v;
} else {
// TODO: convert higher degree SH to lower degree SH
}
}
(_, _) => {},
Expand Down
Loading

0 comments on commit 166a7dc

Please sign in to comment.