12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091 |
- use spiral_rs::aligned_memory::*;
- use spiral_rs::arith::*;
- use spiral_rs::params::*;
- use spiral_rs::poly::*;
- use spiral_rs::server::*;
- use spiral_rs::util::*;
- use crossbeam::thread;
- pub fn load_item_from_slice<'a>(
- params: &'a Params,
- slice: &[u8],
- instance: usize,
- trial: usize,
- item_idx: usize,
- ) -> PolyMatrixRaw<'a> {
- let db_item_size = params.db_item_size;
- let instances = params.instances;
- let trials = params.n * params.n;
- let chunks = instances * trials;
- let bytes_per_chunk = f64::ceil(db_item_size as f64 / chunks as f64) as usize;
- let logp = f64::ceil(f64::log2(params.pt_modulus as f64)) as usize;
- let modp_words_per_chunk = f64::ceil((bytes_per_chunk * 8) as f64 / logp as f64) as usize;
- assert!(modp_words_per_chunk <= params.poly_len);
- let idx_item_in_file = item_idx * db_item_size;
- let idx_chunk = instance * trials + trial;
- let idx_poly_in_file = idx_item_in_file + idx_chunk * bytes_per_chunk;
- let mut out = PolyMatrixRaw::zero(params, 1, 1);
- let modp_words_read = f64::ceil((bytes_per_chunk * 8) as f64 / logp as f64) as usize;
- assert!(modp_words_read <= params.poly_len);
- for i in 0..modp_words_read {
- out.data[i] = read_arbitrary_bits(
- &slice[idx_poly_in_file..idx_poly_in_file + bytes_per_chunk],
- i * logp,
- logp,
- );
- assert!(out.data[i] <= params.pt_modulus);
- }
- out
- }
- pub fn load_db_from_slice_mt(params: &Params, slice: &[u8], numthreads: usize) -> AlignedMemory64 {
- let instances = params.instances;
- let trials = params.n * params.n;
- let dim0 = 1 << params.db_dim_1;
- let num_per = 1 << params.db_dim_2;
- let num_items = dim0 * num_per;
- let db_size_words = instances * trials * num_items * params.poly_len;
- let mut v = AlignedMemory64::new(db_size_words);
- for instance in 0..instances {
- for trial in 0..trials {
- let vslice = v.as_mut_slice();
- thread::scope(|s| {
- s.spawn(|_| {
- for i in 0..num_items {
- let ii = i % num_per;
- let j = i / num_per;
- let mut db_item = load_item_from_slice(¶ms, slice, instance, trial, i);
- // db_item.reduce_mod(params.pt_modulus);
- for z in 0..params.poly_len {
- db_item.data[z] =
- recenter_mod(db_item.data[z], params.pt_modulus, params.modulus);
- }
- let db_item_ntt = db_item.ntt();
- for z in 0..params.poly_len {
- let idx_dst = calc_index(
- &[instance, trial, z, ii, j],
- &[instances, trials, params.poly_len, num_per, dim0],
- );
- vslice[idx_dst] = db_item_ntt.data[z]
- | (db_item_ntt.data[params.poly_len + z] << PACKED_OFFSET_2);
- }
- }
- });
- })
- .unwrap();
- }
- }
- v
- }
|