diff --git a/CHANGELOG.md b/CHANGELOG.md index 276bc2cb..bc0cf0f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,14 @@ signals a backwards-compatible change. ## [Unreleased] +### Changed + +- **Count Sketch estimates no longer allocate.** The per-row estimates a + median is taken over are collected on the stack instead of in a fresh + `Vec`: up to 20 rows for `Vector2D` (the fast and regular paths), + `BitMatrix`, `CountSketchWithL2Norm` (estimates and `l2_squared`) and + `CountSketchFold`, and a `[f64; ROWS]` array for the fixed-size matrices. + About 8 to 14 ns less per estimate on small shapes. Results are unchanged. ### Fixed - **`Kll::update` regains most of its pre-0.4.0 speed.** With the lazily grown diff --git a/src/common/structures/bit_matrix.rs b/src/common/structures/bit_matrix.rs index 5e232596..dd86cd60 100644 --- a/src/common/structures/bit_matrix.rs +++ b/src/common/structures/bit_matrix.rs @@ -256,7 +256,7 @@ impl MatrixStorage for BitMatrix { Hash: MatrixFastHash, F: Fn(&Self::Counter, usize, &Hash) -> f64, { - let mut values: Vec = (0..self.rows) + let mut values: crate::common::structures::matrix_storage::RowEstimates = (0..self.rows) .map(|row| { let col = self.col_for_row(hashed_val, row); op(&self.get(row, col), row, hashed_val) diff --git a/src/common/structures/fixed_structure.rs b/src/common/structures/fixed_structure.rs index 5cebaa3e..c0a77156 100644 --- a/src/common/structures/fixed_structure.rs +++ b/src/common/structures/fixed_structure.rs @@ -294,11 +294,11 @@ macro_rules! impl_fixed_matrix { Hash: $crate::MatrixFastHash, F: Fn(&Self::Counter, usize, &Hash) -> f64, { - let mut estimates = Vec::with_capacity($rows); - for row in 0..$rows { + let mut estimates = [0.0_f64; $rows]; + for (row, estimate) in estimates.iter_mut().enumerate() { let col = hashed_val.col_for_row(row, $cols); let idx = row * $cols + col; - estimates.push(op(&self.data[idx], row, hashed_val)); + *estimate = op(&self.data[idx], row, hashed_val); } $crate::compute_median_inline_f64(&mut estimates) } diff --git a/src/common/structures/matrix_storage.rs b/src/common/structures/matrix_storage.rs index ea886f90..19885d0a 100644 --- a/src/common/structures/matrix_storage.rs +++ b/src/common/structures/matrix_storage.rs @@ -4,6 +4,10 @@ use smallvec::SmallVec; use crate::DataInput; +/// Per-row estimates collected for a median, held inline up to +/// [`MATRIX_MAX_ROWS`](crate::common::MATRIX_MAX_ROWS) rows. +pub(crate) type RowEstimates = SmallVec<[f64; crate::common::MATRIX_MAX_ROWS]>; + /// Fast-path hash container for matrix-backed sketches. #[derive(Clone, Debug)] pub enum MatrixHash { diff --git a/src/common/structures/vector2d.rs b/src/common/structures/vector2d.rs index fea3c5a3..d36370e1 100644 --- a/src/common/structures/vector2d.rs +++ b/src/common/structures/vector2d.rs @@ -1,3 +1,4 @@ +use crate::common::structures::matrix_storage::RowEstimates; use std::ops::{Index, IndexMut}; use crate::{ @@ -394,7 +395,7 @@ impl Vector2D { Hash: MatrixFastHash, F: Fn(&T, usize, &Hash) -> f64, { - let mut estimates = Vec::with_capacity(self.rows); + let mut estimates = RowEstimates::with_capacity(self.rows); for row in 0..self.rows { let col = self.col_for_row(hashed_val, row); let idx = row * self.cols + col; @@ -566,7 +567,7 @@ impl Vector2D { where F: Fn(&T, &Q, usize, &MatrixHash) -> f64, { - let mut estimates = Vec::with_capacity(self.rows); + let mut estimates = RowEstimates::with_capacity(self.rows); for row in 0..self.rows { let col = self.col_for_row(hashed_val, row); let idx = row * self.cols + col; diff --git a/src/sketches/count_sketch_fold.rs b/src/sketches/count_sketch_fold.rs index 4ec43449..6c3b12cd 100644 --- a/src/sketches/count_sketch_fold.rs +++ b/src/sketches/count_sketch_fold.rs @@ -14,6 +14,7 @@ //! additional approximation error — the accuracy is identical to a full-width //! Count Sketch with W columns. +use crate::common::structures::matrix_storage::RowEstimates; use crate::count_min_fold::FoldCell; use crate::{DataInput, DefaultXxHasher, SketchHasher, TopKHeap, compute_median_inline_f64}; use std::marker::PhantomData; @@ -198,7 +199,7 @@ impl CountSketchFold { /// Returns the Count Sketch frequency estimate for `key` (median of /// sign-corrected row estimates). pub fn estimate(&self, key: &DataInput) -> i64 { - let mut estimates = Vec::with_capacity(self.rows); + let mut estimates = RowEstimates::with_capacity(self.rows); for r in 0..self.rows { let (full_col, sign) = self.hash_for(r, key); let fc = self.fold_col_of(full_col); diff --git a/src/sketches/countsketch.rs b/src/sketches/countsketch.rs index a389a5e4..cb627632 100644 --- a/src/sketches/countsketch.rs +++ b/src/sketches/countsketch.rs @@ -10,6 +10,7 @@ use crate::asapv1::admission_rate::admit_all; use crate::common::hash::check_matrix_shape; use crate::common::structure_utils::AdmittedRows; +use crate::common::structures::matrix_storage::RowEstimates; use crate::{ DataInput, DefaultMatrixI32, DefaultMatrixI64, DefaultMatrixI128, DefaultXxHasher, FastPath, FastPathHasher, FixedMatrix, MATRIX_MAX_ROWS, MatrixFastHash, MatrixStorage, NitroTarget, @@ -382,7 +383,7 @@ where pub fn estimate(&self, value: &DataInput) -> f64 { let rows = self.counts.rows(); let cols = self.counts.cols(); - let mut estimates = Vec::with_capacity(rows); + let mut estimates = RowEstimates::with_capacity(rows); for r in 0..rows { let hashed = H::hash64_seeded(r, value); let col = ((hashed & LOWER_32_MASK) as usize) % cols; diff --git a/src/sketches/countsketch_l2_norm.rs b/src/sketches/countsketch_l2_norm.rs index a4137966..e7f7029d 100644 --- a/src/sketches/countsketch_l2_norm.rs +++ b/src/sketches/countsketch_l2_norm.rs @@ -3,6 +3,7 @@ //! [`CountSketchWithL2Norm`] is the per-layer sketch of `UnivMon` and //! `UnivMonTwoTier`; its L2 norm sets their terminal-mode heavy-hitter threshold. +use crate::common::structures::matrix_storage::RowEstimates; use std::marker::PhantomData; use crate::common::structures::matrix_storage::cols_mask_bits; @@ -229,7 +230,7 @@ impl CountSketchWithL2Norm { /// Returns the median estimate of the squared L2 norm across rows. pub fn l2_squared(&self) -> f64 { - let mut values: Vec = self.row_l2_squared.as_slice()[..self.rows] + let mut values: RowEstimates = self.row_l2_squared.as_slice()[..self.rows] .iter() .map(|&v| v as f64) .collect(); @@ -254,7 +255,7 @@ impl CountSketchWithL2Norm { pub fn estimate_with_hash(&self, hashed_val: u128) -> f64 { let mask_bits = self.counts.mask_bits() as usize; let mask = (1u128 << mask_bits) - 1; - let mut lst = Vec::with_capacity(self.rows); + let mut lst = RowEstimates::with_capacity(self.rows); let mut shift_amount = 0; let mut sign_bit_pos = 127;