Merge pull request #1033 from 0x0ece/transpose

Optimize transpose
This commit is contained in:
Nicholas Ward 2023-07-11 12:41:28 -07:00 committed by GitHub
commit ee5d1aa6c8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

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 {