poulpy/backend/src/lib.rs
backend/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
21pub use encoding::*;
22pub use mat_znx_dft::*;
23pub use mat_znx_dft_ops::*;
24pub use module::*;
25pub use sampling::*;
26pub use scalar_znx::*;
27pub use scalar_znx_dft::*;
28pub use scalar_znx_dft_ops::*;
29pub use stats::*;
30pub use vec_znx::*;
31pub use vec_znx_big::*;
32pub use vec_znx_big_ops::*;
33pub use vec_znx_dft::*;
34pub use vec_znx_dft_ops::*;
35pub use vec_znx_ops::*;
36pub use znx_base::*;
37
38pub const GALOISGENERATOR: u64 = 5;
39pub const DEFAULTALIGN: usize = 64;
40
41fn is_aligned_custom<T>(ptr: *const T, align: usize) -> bool {
42 (ptr as usize) % align == 0
43}
44
45pub fn is_aligned<T>(ptr: *const T) -> bool {
46 is_aligned_custom(ptr, DEFAULTALIGN)
47}
48
49pub fn assert_alignement<T>(ptr: *const T) {
50 assert!(
51 is_aligned(ptr),
52 "invalid alignement: ensure passed bytes have been allocated with [alloc_aligned_u8] or [alloc_aligned]"
53 )
54}
55
56pub fn cast<T, V>(data: &[T]) -> &[V] {
57 let ptr: *const V = data.as_ptr() as *const V;
58 let len: usize = data.len() / size_of::<V>();
59 unsafe { std::slice::from_raw_parts(ptr, len) }
60}
61
62pub fn cast_mut<T, V>(data: &[T]) -> &mut [V] {
63 let ptr: *mut V = data.as_ptr() as *mut V;
64 let len: usize = data.len() / size_of::<V>();
65 unsafe { std::slice::from_raw_parts_mut(ptr, len) }
66}
67
68/// Allocates a block of bytes with a custom alignement.
69/// Alignement must be a power of two and size a multiple of the alignement.
70/// Allocated memory is initialized to zero.
71fn alloc_aligned_custom_u8(size: usize, align: usize) -> Vec<u8> {
72 assert!(
73 align.is_power_of_two(),
74 "Alignment must be a power of two but is {}",
75 align
76 );
77 assert_eq!(
78 (size * size_of::<u8>()) % align,
79 0,
80 "size={} must be a multiple of align={}",
81 size,
82 align
83 );
84 unsafe {
85 let layout: std::alloc::Layout = std::alloc::Layout::from_size_align(size, align).expect("Invalid alignment");
86 let ptr: *mut u8 = std::alloc::alloc(layout);
87 if ptr.is_null() {
88 panic!("Memory allocation failed");
89 }
90 assert!(
91 is_aligned_custom(ptr, align),
92 "Memory allocation at {:p} is not aligned to {} bytes",
93 ptr,
94 align
95 );
96 // Init allocated memory to zero
97 std::ptr::write_bytes(ptr, 0, size);
98 Vec::from_raw_parts(ptr, size, size)
99 }
100}
101
102/// Allocates a block of T aligned with [DEFAULTALIGN].
103/// Size of T * size msut be a multiple of [DEFAULTALIGN].
104pub fn alloc_aligned_custom<T>(size: usize, align: usize) -> Vec<T> {
105 assert_eq!(
106 (size * size_of::<T>()) % align,
107 0,
108 "size={} must be a multiple of align={}",
109 size,
110 align
111 );
112 let mut vec_u8: Vec<u8> = alloc_aligned_custom_u8(size_of::<T>() * size, align);
113 let ptr: *mut T = vec_u8.as_mut_ptr() as *mut T;
114 let len: usize = vec_u8.len() / size_of::<T>();
115 let cap: usize = vec_u8.capacity() / size_of::<T>();
116 std::mem::forget(vec_u8);
117 unsafe { Vec::from_raw_parts(ptr, len, cap) }
118}
119
120/// Allocates an aligned vector of size equal to the smallest multiple
121/// of [DEFAULTALIGN]/size_of::<T>() that is equal or greater to `size`.
122pub fn alloc_aligned<T>(size: usize) -> Vec<T> {
123 alloc_aligned_custom::<T>(
124 size + (size % (DEFAULTALIGN / size_of::<T>())),
125 DEFAULTALIGN,
126 )
127}
128
129// Scratch implementation below
130
131pub struct ScratchOwned(Vec<u8>);
132
133impl ScratchOwned {
134 pub fn new(byte_count: usize) -> Self {
135 let data: Vec<u8> = alloc_aligned(byte_count);
136 Self(data)
137 }
138
139 pub fn borrow(&mut self) -> &mut Scratch {
140 Scratch::new(&mut self.0)
141 }
142}
143
144pub struct Scratch {
145 data: [u8],
146}
147
148impl Scratch {
149 fn new(data: &mut [u8]) -> &mut Self {
150 unsafe { &mut *(data as *mut [u8] as *mut Self) }
151 }
152
153 pub fn zero(&mut self) {
154 self.data.fill(0);
155 }
156
157 pub fn available(&self) -> usize {
158 let ptr: *const u8 = self.data.as_ptr();
159 let self_len: usize = self.data.len();
160 let aligned_offset: usize = ptr.align_offset(DEFAULTALIGN);
161 self_len.saturating_sub(aligned_offset)
162 }
163
164 fn take_slice_aligned(data: &mut [u8], take_len: usize) -> (&mut [u8], &mut [u8]) {
165 let ptr: *mut u8 = data.as_mut_ptr();
166 let self_len: usize = data.len();
167
168 let aligned_offset: usize = ptr.align_offset(DEFAULTALIGN);
169 let aligned_len: usize = self_len.saturating_sub(aligned_offset);
170
171 if let Some(rem_len) = aligned_len.checked_sub(take_len) {
172 unsafe {
173 let rem_ptr: *mut u8 = ptr.add(aligned_offset).add(take_len);
174 let rem_slice: &mut [u8] = &mut *std::ptr::slice_from_raw_parts_mut(rem_ptr, rem_len);
175
176 let take_slice: &mut [u8] = &mut *std::ptr::slice_from_raw_parts_mut(ptr.add(aligned_offset), take_len);
177
178 return (take_slice, rem_slice);
179 }
180 } else {
181 panic!(
182 "Attempted to take {} from scratch with {} aligned bytes left",
183 take_len,
184 aligned_len,
185 // type_name::<T>(),
186 // aligned_len
187 );
188 }
189 }
190
191 pub fn tmp_slice<T>(&mut self, len: usize) -> (&mut [T], &mut Self) {
192 let (take_slice, rem_slice) = Self::take_slice_aligned(&mut self.data, len * std::mem::size_of::<T>());
193
194 unsafe {
195 (
196 &mut *(std::ptr::slice_from_raw_parts_mut(take_slice.as_mut_ptr() as *mut T, len)),
197 Self::new(rem_slice),
198 )
199 }
200 }
201
202 pub fn tmp_scalar_znx<B: Backend>(&mut self, module: &Module<B>, cols: usize) -> (ScalarZnx<&mut [u8]>, &mut Self) {
203 let (take_slice, rem_slice) = Self::take_slice_aligned(&mut self.data, bytes_of_scalar_znx(module, cols));
204
205 (
206 ScalarZnx::from_data(take_slice, module.n(), cols),
207 Self::new(rem_slice),
208 )
209 }
210
211 pub fn tmp_scalar_znx_dft<B: Backend>(&mut self, module: &Module<B>, cols: usize) -> (ScalarZnxDft<&mut [u8], B>, &mut Self) {
212 let (take_slice, rem_slice) = Self::take_slice_aligned(&mut self.data, bytes_of_scalar_znx_dft(module, cols));
213
214 (
215 ScalarZnxDft::from_data(take_slice, module.n(), cols),
216 Self::new(rem_slice),
217 )
218 }
219
220 pub fn tmp_vec_znx_dft<B: Backend>(
221 &mut self,
222 module: &Module<B>,
223 cols: usize,
224 size: usize,
225 ) -> (VecZnxDft<&mut [u8], B>, &mut Self) {
226 let (take_slice, rem_slice) = Self::take_slice_aligned(&mut self.data, bytes_of_vec_znx_dft(module, cols, size));
227
228 (
229 VecZnxDft::from_data(take_slice, module.n(), cols, size),
230 Self::new(rem_slice),
231 )
232 }
233
234 pub fn tmp_vec_znx_big<B: Backend>(
235 &mut self,
236 module: &Module<B>,
237 cols: usize,
238 size: usize,
239 ) -> (VecZnxBig<&mut [u8], B>, &mut Self) {
240 let (take_slice, rem_slice) = Self::take_slice_aligned(&mut self.data, bytes_of_vec_znx_big(module, cols, size));
241
242 (
243 VecZnxBig::from_data(take_slice, module.n(), cols, size),
244 Self::new(rem_slice),
245 )
246 }
247
248 pub fn tmp_vec_znx<B: Backend>(&mut self, module: &Module<B>, cols: usize, size: usize) -> (VecZnx<&mut [u8]>, &mut Self) {
249 let (take_slice, rem_slice) = Self::take_slice_aligned(&mut self.data, module.bytes_of_vec_znx(cols, size));
250 (
251 VecZnx::from_data(take_slice, module.n(), cols, size),
252 Self::new(rem_slice),
253 )
254 }
255
256 pub fn tmp_mat_znx_dft<B: Backend>(
257 &mut self,
258 module: &Module<B>,
259 rows: usize,
260 cols_in: usize,
261 cols_out: usize,
262 size: usize,
263 ) -> (MatZnxDft<&mut [u8], B>, &mut Self) {
264 let (take_slice, rem_slice) = Self::take_slice_aligned(
265 &mut self.data,
266 module.bytes_of_mat_znx_dft(rows, cols_in, cols_out, size),
267 );
268 (
269 MatZnxDft::from_data(take_slice, module.n(), rows, cols_in, cols_out, size),
270 Self::new(rem_slice),
271 )
272 }
273}