Optimize transpose

This commit is contained in:
Emanuele Cesena 2023-05-10 12:28:44 -05:00
parent 779456c2c9
commit f1cc284d42

View File

@ -1,6 +1,6 @@
use alloc::vec;
use alloc::vec::Vec;
use plonky2_maybe_rayon::*;
#[doc(inline)]
pub use plonky2_util::*;
@ -20,35 +20,11 @@ pub(crate) fn transpose_poly_values<F: Field>(polys: Vec<PolynomialValues<F>>) -
}
pub fn transpose<F: Field>(matrix: &[Vec<F>]) -> Vec<Vec<F>> {
let l = matrix.len();
let w = matrix[0].len();
let mut transposed = vec![vec![]; w];
for i in 0..w {
transposed[i].reserve_exact(l);
unsafe {
// After .reserve_exact(l), transposed[i] will have capacity at least l. Hence, set_len
// will not cause the buffer to overrun.
transposed[i].set_len(l);
}
}
// Optimization: ensure the larger loop is outside.
if w >= l {
for i in 0..w {
for j in 0..l {
transposed[i][j] = matrix[j][i];
}
}
} else {
for j in 0..l {
for i in 0..w {
transposed[i][j] = matrix[j][i];
}
}
}
transposed
let len = matrix[0].len();
(0..len)
.into_par_iter()
.map(|i| matrix.iter().map(|row| row[i]).collect())
.collect()
}
pub(crate) fn reverse_bits(n: usize, num_bits: usize) -> usize {