Optimize sparse vec with byte packing into u32

Reduces memory usage for large symbol counts by ~40%
This commit is contained in:
Christopher Berner 2020-01-12 19:22:02 -08:00
parent 886075c1b0
commit fd93b712c8
3 changed files with 100 additions and 68 deletions

@ -213,8 +213,8 @@ impl<'a> OctetIter<'a> {
x.keys_values()
.map(|(physical_col, value)| {
(
self.sparse_physical_col_to_logical.unwrap()[*physical_col],
value.clone(),
self.sparse_physical_col_to_logical.unwrap()[physical_col],
value,
)
})
.filter(|(logical_col, _)| {
@ -248,7 +248,7 @@ impl<'a> Iterator for OctetIter<'a> {
self.sparse_index += 1;
let logical_col = self.sparse_physical_col_to_logical.unwrap()[entry.0];
if logical_col >= self.start_col && logical_col < self.end_col {
return Some((logical_col, entry.1.clone()));
return Some((logical_col, entry.1));
}
}
return None;

@ -46,8 +46,8 @@ impl SparseBinaryMatrix {
}
for row in 0..self.height {
for (col, value) in self.sparse_elements[row].keys_values() {
if *value != Octet::zero() {
debug_assert!(self.sparse_column_index[*col].exists(row));
if value != Octet::zero() {
debug_assert!(self.sparse_column_index[col].exists(row));
}
}
}
@ -143,8 +143,8 @@ impl BinaryMatrix for SparseBinaryMatrix {
let mut ones = 0;
let physical_row = self.logical_row_to_physical[row];
for (physical_col, value) in self.sparse_elements[physical_row].keys_values() {
let col = self.physical_col_to_logical[*physical_col];
if col >= start_col && col < end_col && *value == Octet::one() {
let col = self.physical_col_to_logical[physical_col];
if col >= start_col && col < end_col && value == Octet::one() {
ones += 1;
}
}
@ -175,8 +175,7 @@ impl BinaryMatrix for SparseBinaryMatrix {
} else {
return self.sparse_elements[physical_i]
.get(physical_j)
.unwrap_or(&Octet::zero())
.clone();
.unwrap_or_else(Octet::zero);
}
}
@ -228,7 +227,7 @@ impl BinaryMatrix for SparseBinaryMatrix {
self.sparse_column_index = vec![SparseValuelessVec::with_capacity(50); self.width];
for (physical_row, elements) in self.sparse_elements.iter().enumerate() {
for (physical_col, _) in elements.keys_values() {
self.sparse_column_index[*physical_col].insert_last(physical_row);
self.sparse_column_index[physical_col].insert_last(physical_row);
}
}
}
@ -301,7 +300,7 @@ impl BinaryMatrix for SparseBinaryMatrix {
}
if !self.column_index_disabled {
for (col, _) in self.sparse_elements[physical_row].keys_values() {
self.sparse_column_index[*col].insert(physical_row)
self.sparse_column_index[col].insert(physical_row)
}
}
}
@ -331,7 +330,7 @@ impl BinaryMatrix for SparseBinaryMatrix {
let new_columns = dest_row.add_assign(temp_row);
if !self.column_index_disabled {
for new_col in new_columns {
self.sparse_column_index[new_col].insert(physical_dest);
self.sparse_column_index[new_col as usize].insert(physical_dest);
}
}

@ -3,10 +3,38 @@ use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::mem::size_of;
#[derive(Clone, Debug, PartialEq, PartialOrd, Eq, Ord, Serialize, Deserialize, Hash)]
struct PackedEntry {
value: u32,
}
impl PackedEntry {
fn new(index: u32, binary_value: u8) -> PackedEntry {
debug_assert!(index < 16777216);
debug_assert!(binary_value < 2);
PackedEntry {
value: index << 8 | (binary_value as u32),
}
}
fn value(&self) -> Octet {
Octet::new((self.value & 0xFF) as u8)
}
fn add_value(&mut self, value: Octet) {
// Index is stored in upper 24-bits, but XOR'ing with zero won't change it.
self.value ^= value.byte() as u32
}
fn index(&self) -> u32 {
self.value >> 8
}
}
#[derive(Clone, Debug, PartialEq, PartialOrd, Eq, Ord, Serialize, Deserialize, Hash)]
pub struct SparseBinaryVec {
// Kept sorted by the usize (key)
elements: Vec<(usize, Octet)>,
elements: Vec<PackedEntry>,
}
impl SparseBinaryVec {
@ -18,42 +46,42 @@ impl SparseBinaryVec {
// Returns the internal index into self.elements matching key i, or the index
// at which it can be inserted (maintaining sorted order)
fn key_to_internal_index(&self, i: usize) -> Result<usize, usize> {
self.elements.binary_search_by_key(&i, |(index, _)| *index)
fn key_to_internal_index(&self, i: u32) -> Result<usize, usize> {
self.elements
.binary_search_by_key(&i, |entry| entry.index())
}
pub fn size_in_bytes(&self) -> usize {
size_of::<Self>() + size_of::<(usize, Octet)>() * self.elements.len()
size_of::<Self>() + size_of::<PackedEntry>() * self.elements.len()
}
pub fn len(&self) -> usize {
self.elements.len()
}
pub fn get_by_raw_index(&self, i: usize) -> &(usize, Octet) {
&self.elements[i]
pub fn get_by_raw_index(&self, i: usize) -> (usize, Octet) {
(self.elements[i].index() as usize, self.elements[i].value())
}
// Returns a vector of new column indices that this row contains
pub fn add_assign(&mut self, other: &SparseBinaryVec) -> Vec<usize> {
pub fn add_assign(&mut self, other: &SparseBinaryVec) -> Vec<u32> {
// Fast path for a single value that's being eliminated
// TODO: Probably wouldn't need this if we implemented "Furthermore, the row operations
// required for the HDPC rows may be performed for all such rows in one
// process, by using the algorithm described in Section 5.3.3.3."
if other.elements.len() == 1 {
let (other_col, other_value) = &other.elements[0];
match self.key_to_internal_index(*other_col) {
let other_entry = &other.elements[0];
match self.key_to_internal_index(other_entry.index()) {
Ok(index) => {
let self_value = &mut self.elements[index].1;
*self_value += other_value;
if *self_value == Octet::zero() {
let self_entry = &mut self.elements[index];
self_entry.add_value(other_entry.value());
if self_entry.value() == Octet::zero() {
self.elements.remove(index);
}
}
Err(index) => {
self.elements
.insert(index, (*other_col, other_value.clone()));
return vec![*other_col];
self.elements.insert(index, other_entry.clone());
return vec![other_entry.index()];
}
};
return vec![];
@ -62,48 +90,48 @@ impl SparseBinaryVec {
let mut result = Vec::with_capacity(self.elements.len() + other.elements.len());
let mut self_iter = self.elements.iter();
let mut other_iter = other.elements.iter();
let mut self_entry = self_iter.next();
let mut other_entry = other_iter.next();
let mut self_next = self_iter.next();
let mut other_next = other_iter.next();
let mut new_columns = Vec::with_capacity(10);
loop {
if let Some((self_col, self_value)) = self_entry {
if let Some((other_col, other_value)) = other_entry {
match self_col.cmp(&other_col) {
if let Some(self_entry) = self_next {
if let Some(other_entry) = other_next {
match self_entry.index().cmp(&other_entry.index()) {
Ordering::Less => {
if *self_value != Octet::zero() {
result.push((*self_col, self_value.clone()));
if self_entry.value() != Octet::zero() {
result.push(self_entry.clone());
}
self_entry = self_iter.next();
self_next = self_iter.next();
}
Ordering::Equal => {
let value = self_value + other_value;
let value = self_entry.value() + other_entry.value();
if value != Octet::zero() {
result.push((*self_col, value));
result.push(PackedEntry::new(self_entry.index(), value.byte()));
}
self_entry = self_iter.next();
other_entry = other_iter.next();
self_next = self_iter.next();
other_next = other_iter.next();
}
Ordering::Greater => {
if *other_value != Octet::zero() {
new_columns.push(*other_col);
result.push((*other_col, other_value.clone()));
if other_entry.value() != Octet::zero() {
new_columns.push(other_entry.index());
result.push(other_entry.clone());
}
other_entry = other_iter.next();
other_next = other_iter.next();
}
}
} else {
if *self_value != Octet::zero() {
result.push((*self_col, self_value.clone()));
if self_entry.value() != Octet::zero() {
result.push(self_entry.clone());
}
self_entry = self_iter.next();
self_next = self_iter.next();
}
} else if let Some((other_col, other_value)) = other_entry {
if *other_value != Octet::zero() {
new_columns.push(*other_col);
result.push((*other_col, other_value.clone()));
} else if let Some(other_entry) = other_next {
if other_entry.value() != Octet::zero() {
new_columns.push(other_entry.index());
result.push(other_entry.clone());
}
other_entry = other_iter.next();
other_next = other_iter.next();
} else {
break;
}
@ -114,31 +142,36 @@ impl SparseBinaryVec {
}
pub fn remove(&mut self, i: usize) -> Option<Octet> {
match self.key_to_internal_index(i) {
Ok(index) => Some(self.elements.remove(index).1),
match self.key_to_internal_index(i as u32) {
Ok(index) => Some(self.elements.remove(index).value()),
Err(_) => None,
}
}
pub fn retain<P: Fn(&(usize, Octet)) -> bool>(&mut self, predicate: P) {
self.elements.retain(predicate);
self.elements
.retain(|entry| predicate(&(entry.index() as usize, entry.value())));
}
pub fn get(&self, i: usize) -> Option<&Octet> {
match self.key_to_internal_index(i) {
Ok(index) => Some(&self.elements[index].1),
pub fn get(&self, i: usize) -> Option<Octet> {
match self.key_to_internal_index(i as u32) {
Ok(index) => Some(self.elements[index].value()),
Err(_) => None,
}
}
pub fn keys_values(&self) -> impl Iterator<Item = &(usize, Octet)> {
self.elements.iter()
pub fn keys_values(&self) -> impl Iterator<Item = (usize, Octet)> + '_ {
self.elements
.iter()
.map(|entry| (entry.index() as usize, entry.value()))
}
pub fn insert(&mut self, i: usize, value: Octet) {
match self.key_to_internal_index(i) {
Ok(index) => self.elements[index] = (i, value),
Err(index) => self.elements.insert(index, (i, value)),
match self.key_to_internal_index(i as u32) {
Ok(index) => self.elements[index] = PackedEntry::new(i as u32, value.byte()),
Err(index) => self
.elements
.insert(index, PackedEntry::new(i as u32, value.byte())),
}
}
}
@ -213,7 +246,7 @@ mod tests {
let mut sparse = SparseBinaryVec::with_capacity(size);
for _ in 0..size {
let i = rand::thread_rng().gen_range(0, size);
let value = rand::thread_rng().gen();
let value = rand::thread_rng().gen_range(0, 2);
dense[i] = value;
sparse.insert(i, Octet::new(value));
}
@ -227,13 +260,13 @@ mod tests {
let mut dense1 = vec![Octet::zero(); 8];
let mut sparse1 = SparseBinaryVec::with_capacity(8);
for i in 0..4 {
let value = rand::thread_rng().gen();
let value = rand::thread_rng().gen_range(0, 2);
dense1[i * 2] = Octet::new(value);
sparse1.insert(i * 2, Octet::new(value));
}
for i in 0..8 {
let actual = sparse1.get(i).map(|x| x.clone()).unwrap_or(Octet::zero());
let actual = sparse1.get(i).unwrap_or(Octet::zero());
let expected = dense1[i].clone();
assert_eq!(
actual, expected,
@ -245,13 +278,13 @@ mod tests {
let mut dense2 = vec![Octet::zero(); 8];
let mut sparse2 = SparseBinaryVec::with_capacity(8);
for i in 0..4 {
let value = rand::thread_rng().gen();
let value = rand::thread_rng().gen_range(0, 2);
dense2[i] = Octet::new(value);
sparse2.insert(i, Octet::new(value));
}
for i in 0..8 {
let actual = sparse2.get(i).map(|x| x.clone()).unwrap_or(Octet::zero());
let actual = sparse2.get(i).unwrap_or(Octet::zero());
let expected = dense2[i].clone();
assert_eq!(
actual, expected,
@ -263,7 +296,7 @@ mod tests {
sparse1.add_assign(&sparse2);
for i in 0..8 {
let actual = sparse1.get(i).map(|x| x.clone()).unwrap_or(Octet::zero());
let actual = sparse1.get(i).unwrap_or(Octet::zero());
let expected = &dense1[i] + &dense2[i];
assert_eq!(
actual, expected,