feat: use cache BR for bit_reverse

Author: SuccinctPaul

starky/src/fft_p.rs

@@ -3,18 +3,67 @@ use crate::constant::{get_max_workers, MAX_OPS_PER_THREAD, MIN_OPS_PER_THREAD, S
 use crate::fft_worker::{fft_block, interpolate_prepare_block};
 use crate::helper::log2_any;
 use crate::traits::FieldExtension;
+use crate::utils::parallells::parallelize;
 use core::cmp::min;
+use lazy_static::lazy_static;
 use rayon::prelude::*;
-use rayon::{current_num_threads, scope};
+use std::collections::HashMap;
+use std::sync::Mutex;
 
+lazy_static! {
+    static ref BR_CACHE: Mutex<HashMap<usize, Vec<usize>>> = Mutex::new(HashMap::new());
+}
 pub fn BR(x: usize, domain_pow: usize) -> usize {
     assert!(domain_pow <= 32);
-    let mut x = x;
-    x = (x >> 16) | (x << 16);
-    x = ((x & 0xFF00FF00) >> 8) | ((x & 0x00FF00FF) << 8);
-    x = ((x & 0xF0F0F0F0) >> 4) | ((x & 0x0F0F0F0F) << 4);
-    x = ((x & 0xCCCCCCCC) >> 2) | ((x & 0x33333333) << 2);
-    (((x & 0xAAAAAAAA) >> 1) | ((x & 0x55555555) << 1)) >> (32 - domain_pow)
+    let cal = |x: usize, domain_pow: usize| -> usize {
+        let mut x = x;
+        x = (x >> 16) | (x << 16);
+        x = ((x & 0xFF00FF00) >> 8) | ((x & 0x00FF00FF) << 8);
+        x = ((x & 0xF0F0F0F0) >> 4) | ((x & 0x0F0F0F0F) << 4);
+        x = ((x & 0xCCCCCCCC) >> 2) | ((x & 0x33333333) << 2);
+        (((x & 0xAAAAAAAA) >> 1) | ((x & 0x55555555) << 1)) >> (32 - domain_pow)
+    };
+
+    // get cache by domain_pow
+    let mut map = BR_CACHE.lock().unwrap();
+    let mut cache = if map.contains_key(&domain_pow) {
+        map.remove(&domain_pow).unwrap() // get and remove the old values.
+    } else {
+        vec![]
+    };
+    // check if need append more to cache
+    let cache_len = cache.len();
+    let n = 1 << domain_pow;
+    if cache_len <= n || cache_len < x {
+        let end = if n >= x { n } else { x };
+        // todo parallel
+        for i in cache_len..=end {
+            let a = cal(i, domain_pow);
+            cache.push(a);
+        }
+    }
+    let res = cache[x];
+    // update map with cache
+    map.insert(domain_pow, cache);
+    res
+}
+fn BRs(start: usize, end: usize, domain_pow: usize) -> Vec<usize> {
+    assert!(end > start);
+    // 1. obtain a useless one to precompute the cache.
+    //      to make sure the cache existed and its len >= end.
+    BR(end, domain_pow);
+
+    // 2. get cache by domain_pow
+    let map = BR_CACHE.lock().unwrap();
+    let cache = if map.contains_key(&domain_pow) {
+        map.get(&domain_pow).unwrap()
+    } else {
+        // double check
+        BR(end, domain_pow);
+        map.get(&domain_pow).unwrap()
+    };
+
+    (start..end).map(|i| cache[i]).collect()
 }
 
 pub fn transpose<F: FieldExtension>(
@@ -45,26 +94,16 @@ pub fn bit_reverse<F: FieldExtension>(
     nbits: usize,
 ) {
     let n = 1 << nbits;
-    let num_threads = current_num_threads();
-    let mut chunk_size = n / num_threads;
-    if chunk_size < num_threads {
-        chunk_size = 1;
+    let ris = BRs(0, n, nbits); // move it outside the loop. obtain it from cache.
+
+    // todo parallel. Does the buffdst
+    let len = n * n_pols;
+    assert_eq!(len, buffdst.len()); // is ok?
+    for j in 0..len {
+        let i = j / n_pols;
+        let k = j % n_pols;
+        buffdst[j] = buffsrc[ris[i] * n_pols + k];
     }
-
-    scope(|scope| {
-        // for (chunk_num, v) in buffdst.chunks_mut(chunk).enumerate() {
-        for chunk_num in 0..num_threads {
-            scope.spawn(move |_| {
-                let start = chunk_num * chunk_size;
-                for i in start..(start + chunk_size) {
-                    let ri = BR(i, nbits);
-                    for k in 0..n_pols {
-                        buffdst[i * n_pols + k] = buffsrc[ri * n_pols + k];
-                    }
-                }
-            });
-        }
-    });
 }
 
 pub fn interpolate_bit_reverse<F: FieldExtension>(
@@ -91,27 +130,13 @@ pub fn inv_bit_reverse<F: FieldExtension>(
 ) {
     let n = 1 << nbits;
     let n_inv = F::inv(&F::from(n));
-    let num_threads = current_num_threads();
-    let mut chunk_size = n / num_threads;
-    if chunk_size < num_threads {
-        chunk_size = 1;
-    }
-
-    scope(|scope| {
-        // for (chunk_num, v) in buffdst.chunks_mut(chunk).enumerate() {
-        for chunk_num in 0..num_threads {
-            scope.spawn(move |_| {
-                let start = chunk_num * chunk_size;
-                for i in start..(start + chunk_size) {
-                    let ri = BR(i, nbits);
-                    let rii = (n - ri) % n;
-                    for p in 0..n_pols {
-                        buffdst[i * n_pols + p] = buffsrc[rii * n_pols + p] * n_inv;
-                    }
-                }
-            });
+    for i in 0..n {
+        let ri = BR(i, nbits);
+        let rii = (n - ri) % n;
+        for p in 0..n_pols {
+            buffdst[i * n_pols + p] = buffsrc[rii * n_pols + p] * n_inv;
         }
-    });
+    }
 }
 
 pub fn interpolate_prepare<F: FieldExtension>(buff: &mut Vec<F>, n_pols: usize, nbits: usize) {