poulpy/base2k/src/lib.rs
base2k/src/lib.rsRaw
1pub mod encoding;
2#[allow(non_camel_case_types, non_snake_case, non_upper_case_globals, dead_code, improper_ctypes)]
3// Other modules and exports
4pub mod ffi;
5pub mod mat_znx_dft;
6pub mod mat_znx_dft_ops;
7pub mod module;
8pub mod sampling;
9pub mod scalar_znx;
10pub mod scalar_znx_dft;
11pub mod scalar_znx_dft_ops;
12pub mod stats;
13pub mod vec_znx;
14pub mod vec_znx_big;
15pub mod vec_znx_big_ops;
16pub mod vec_znx_dft;
17pub mod vec_znx_dft_ops;
18pub mod vec_znx_ops;
19pub mod znx_base;
20
21use std::{
22 any::type_name,
23 ops::{DerefMut, Sub},
24};
25
26pub use encoding::*;
27pub use mat_znx_dft::*;
28pub use mat_znx_dft_ops::*;
29pub use module::*;
30use rand_core::le;
31use rand_distr::num_traits::sign;
32pub use sampling::*;
33pub use scalar_znx::*;
34pub use scalar_znx_dft::*;
35pub use scalar_znx_dft_ops::*;
36pub use stats::*;
37pub use vec_znx::*;
38pub use vec_znx_big::*;
39pub use vec_znx_big_ops::*;
40pub use vec_znx_dft::*;
41pub use vec_znx_dft_ops::*;
42pub use vec_znx_ops::*;
43pub use znx_base::*;
44
45pub const GALOISGENERATOR: u64 = 5;
46pub const DEFAULTALIGN: usize = 64;
47
48fn is_aligned_custom<T>(ptr: *const T, align: usize) -> bool {
49 (ptr as usize) % align == 0
50}
51
52pub fn is_aligned<T>(ptr: *const T) -> bool {
53 is_aligned_custom(ptr, DEFAULTALIGN)
54}
55
56pub fn assert_alignement<T>(ptr: *const T) {
57 assert!(
58 is_aligned(ptr),
59 "invalid alignement: ensure passed bytes have been allocated with [alloc_aligned_u8] or [alloc_aligned]"
60 )
61}
62
63pub fn cast<T, V>(data: &[T]) -> &[V] {
64 let ptr: *const V = data.as_ptr() as *const V;
65 let len: usize = data.len() / size_of::<V>();
66 unsafe { std::slice::from_raw_parts(ptr, len) }
67}
68
69pub fn cast_mut<T, V>(data: &[T]) -> &mut [V] {
70 let ptr: *mut V = data.as_ptr() as *mut V;
71 let len: usize = data.len() / size_of::<V>();
72 unsafe { std::slice::from_raw_parts_mut(ptr, len) }
73}
74
75/// Allocates a block of bytes with a custom alignement.
76/// Alignement must be a power of two and size a multiple of the alignement.
77/// Allocated memory is initialized to zero.
78fn alloc_aligned_custom_u8(size: usize, align: usize) -> Vec<u8> {
79 assert!(
80 align.is_power_of_two(),
81 "Alignment must be a power of two but is {}",
82 align
83 );
84 assert_eq!(
85 (size * size_of::<u8>()) % align,
86 0,
87 "size={} must be a multiple of align={}",
88 size,
89 align
90 );
91 unsafe {
92 let layout: std::alloc::Layout = std::alloc::Layout::from_size_align(size, align).expect("Invalid alignment");
93 let ptr: *mut u8 = std::alloc::alloc(layout);
94 if ptr.is_null() {
95 panic!("Memory allocation failed");
96 }
97 assert!(
98 is_aligned_custom(ptr, align),
99 "Memory allocation at {:p} is not aligned to {} bytes",
100 ptr,
101 align
102 );
103 // Init allocated memory to zero
104 std::ptr::write_bytes(ptr, 0, size);
105 Vec::from_raw_parts(ptr, size, size)
106 }
107}
108
109/// Allocates a block of T aligned with [DEFAULTALIGN].
110/// Size of T * size msut be a multiple of [DEFAULTALIGN].
111pub fn alloc_aligned_custom<T>(size: usize, align: usize) -> Vec<T> {
112 assert_eq!(
113 (size * size_of::<T>()) % align,
114 0,
115 "size={} must be a multiple of align={}",
116 size,
117 align
118 );
119 let mut vec_u8: Vec<u8> = alloc_aligned_custom_u8(size_of::<T>() * size, align);
120 let ptr: *mut T = vec_u8.as_mut_ptr() as *mut T;
121 let len: usize = vec_u8.len() / size_of::<T>();
122 let cap: usize = vec_u8.capacity() / size_of::<T>();
123 std::mem::forget(vec_u8);
124 unsafe { Vec::from_raw_parts(ptr, len, cap) }
125}
126
127/// Allocates an aligned vector of size equal to the smallest multiple
128/// of [DEFAULTALIGN]/size_of::<T>() that is equal or greater to `size`.
129pub fn alloc_aligned<T>(size: usize) -> Vec<T> {
130 alloc_aligned_custom::<T>(
131 size + (size % (DEFAULTALIGN / size_of::<T>())),
132 DEFAULTALIGN,
133 )
134}
135
136pub struct ScratchOwned(Vec<u8>);
137
138impl ScratchOwned {
139 pub fn new(byte_count: usize) -> Self {
140 let data: Vec<u8> = alloc_aligned(byte_count);
141 Self(data)
142 }
143
144 pub fn borrow(&mut self) -> &mut ScratchBorr {
145 ScratchBorr::new(&mut self.0)
146 }
147}
148
149pub struct ScratchBorr {
150 data: [u8],
151}
152
153impl ScratchBorr {
154 fn new(data: &mut [u8]) -> &mut Self {
155 unsafe { &mut *(data as *mut [u8] as *mut Self) }
156 }
157
158 fn take_slice_aligned(data: &mut [u8], take_len: usize) -> (&mut [u8], &mut [u8]) {
159 let ptr = data.as_mut_ptr();
160 let self_len = data.len();
161
162 let aligned_offset = ptr.align_offset(DEFAULTALIGN);
163 let aligned_len = self_len.saturating_sub(aligned_offset);
164
165 if let Some(rem_len) = aligned_len.checked_sub(take_len) {
166 unsafe {
167 let rem_ptr = ptr.add(aligned_offset).add(take_len);
168 let rem_slice = &mut *std::ptr::slice_from_raw_parts_mut(rem_ptr, rem_len);
169
170 let take_slice = &mut *std::ptr::slice_from_raw_parts_mut(ptr.add(aligned_offset), take_len);
171
172 return (take_slice, rem_slice);
173 }
174 } else {
175 panic!(
176 "Attempted to take {} from scratch with {} aligned bytes left",
177 take_len,
178 take_len,
179 // type_name::<T>(),
180 // aligned_len
181 );
182 }
183 }
184
185 fn tmp_scalar_slice<T>(&mut self, len: usize) -> (&mut [T], &mut Self) {
186 let (take_slice, rem_slice) = Self::take_slice_aligned(&mut self.data, len * std::mem::size_of::<T>());
187
188 unsafe {
189 (
190 &mut *(std::ptr::slice_from_raw_parts_mut(take_slice.as_mut_ptr() as *mut T, len)),
191 Self::new(rem_slice),
192 )
193 }
194 }
195
196 fn tmp_vec_znx_dft<B: Backend>(
197 &mut self,
198 module: &Module<B>,
199 cols: usize,
200 size: usize,
201 ) -> (VecZnxDft<&mut [u8], B>, &mut Self) {
202 let (take_slice, rem_slice) = Self::take_slice_aligned(&mut self.data, bytes_of_vec_znx_dft(module, cols, size));
203
204 (
205 VecZnxDft::from_data(take_slice, module.n(), cols, size),
206 Self::new(rem_slice),
207 )
208 }
209
210 fn tmp_vec_znx_big<D: for<'a> From<&'a mut [u8]>, B: Backend>(
211 &mut self,
212 module: &Module<B>,
213 cols: usize,
214 size: usize,
215 ) -> (VecZnxBig<D, B>, &mut Self) {
216 let (take_slice, rem_slice) = Self::take_slice_aligned(&mut self.data, bytes_of_vec_znx_big(module, cols, size));
217
218 (
219 VecZnxBig::from_data(D::from(take_slice), module.n(), cols, size),
220 Self::new(rem_slice),
221 )
222 }
223}
224
225// pub struct ScratchBorrowed<'a> {
226// data: &'a mut [u8],
227// }
228
229// impl<'a> ScratchBorrowed<'a> {
230// fn take_slice<T>(&mut self, take_len: usize) -> (&mut [T], ScratchBorrowed<'_>) {
231// let ptr = self.data.as_mut_ptr();
232// let self_len = self.data.len();
233
234// //TODO(Jay): print the offset sometimes, just to check
235// let aligned_offset = ptr.align_offset(DEFAULTALIGN);
236// let aligned_len = self_len.saturating_sub(aligned_offset);
237
238// let take_len_bytes = take_len * std::mem::size_of::<T>();
239
240// if let Some(rem_len) = aligned_len.checked_sub(take_len_bytes) {
241// unsafe {
242// let rem_ptr = ptr.add(aligned_offset).add(take_len_bytes);
243// let rem_slice = &mut *std::ptr::slice_from_raw_parts_mut(rem_ptr, rem_len);
244
245// let take_slice = &mut *std::ptr::slice_from_raw_parts_mut(ptr.add(aligned_offset) as *mut T, take_len_bytes);
246
247// return (take_slice, ScratchBorrowed { data: rem_slice });
248// }
249// } else {
250// panic!(
251// "Attempted to take {} (={} elements of {}) from scratch with {} aligned bytes left",
252// take_len_bytes,
253// take_len,
254// type_name::<T>(),
255// aligned_len
256// );
257// }
258// }
259
260// fn reborrow(&mut self) -> ScratchBorrowed<'a> {
261// //(Jay)TODO: `data: &mut *self.data` does not work because liftime of &mut self is different from 'a.
262// // But it feels that there should be a simpler impl. than the one below
263// Self {
264// data: unsafe { &mut *std::ptr::slice_from_raw_parts_mut(self.data.as_mut_ptr(), self.data.len()) },
265// }
266// }
267
268// fn tmp_vec_znx_dft<B: Backend>(&mut self, module: &Module<B>, cols: usize, size: usize) -> (VecZnxDft<&mut [u8], B>, Self) {
269// let (data, re_scratch) = self.take_slice::<u8>(vec_znx_dft::bytes_of_vec_znx_dft(module, cols, size));
270// (
271// VecZnxDft::from_data(data, module.n(), cols, size),
272// re_scratch,
273// )
274// }
275
276// pub(crate) fn len(&self) -> usize {
277// self.data.len()
278// }
279// }
280
281// pub trait Scratch<D> {
282// fn tmp_vec_znx_dft<B: Backend>(&mut self, module: &Module<B>, cols: usize, size: usize) -> (D, &mut Self);
283// }
284
285// impl<'a> Scratch<&'a mut [u8]> for ScratchBorr {
286// fn tmp_vec_znx_dft<B: Backend>(&mut self, module: &Module<B>, cols: usize, size: usize) -> (&'a mut [u8], &mut Self) {
287// let (data, rem_scratch) = self.tmp_scalar_slice(vec_znx_dft::bytes_of_vec_znx_dft(module, cols, size));
288// (
289// data
290// rem_scratch,
291// )
292// }
293
294// // fn tmp_vec_znx_big<B: Backend>(&mut self, module: &Module<B>, cols: usize, size: usize) -> (VecZnxBig<&mut [u8], B>, Self) {
295// // // let (data, re_scratch) = self.take_slice(vec_znx_big::bytes_of_vec_znx_big(module, cols, size));
296// // // (
297// // // VecZnxBig::from_data(data, module.n(), cols, size),
298// // // re_scratch,
299// // // )
300// // }
301
302// // fn scalar_slice<T>(&mut self, len: usize) -> (&mut [T], Self) {
303// // self.take_slice::<T>(len)
304// // }
305
306// // fn reborrow(&mut self) -> Self {
307// // self.reborrow()
308// // }
309// }