poulpy/base2k/src/lib.rs
base2k/src/lib.rsRaw
1pub mod commons;
2pub mod encoding;
3#[allow(non_camel_case_types, non_snake_case, non_upper_case_globals, dead_code, improper_ctypes)]
4// Other modules and exports
5pub mod ffi;
6pub mod mat_znx_dft;
7pub mod module;
8pub mod sampling;
9pub mod scalar_znx_dft;
10pub mod stats;
11pub mod vec_znx;
12pub mod vec_znx_big;
13pub mod vec_znx_big_ops;
14pub mod vec_znx_dft;
15pub mod vec_znx_ops;
16
17pub use commons::*;
18pub use encoding::*;
19pub use mat_znx_dft::*;
20pub use module::*;
21pub use sampling::*;
22pub use scalar_znx_dft::*;
23#[allow(unused_imports)]
24pub use stats::*;
25pub use vec_znx::*;
26pub use vec_znx_big::*;
27pub use vec_znx_big_ops::*;
28pub use vec_znx_dft::*;
29pub use vec_znx_ops::*;
30
31pub const GALOISGENERATOR: u64 = 5;
32pub const DEFAULTALIGN: usize = 64;
33
34fn is_aligned_custom<T>(ptr: *const T, align: usize) -> bool {
35 (ptr as usize) % align == 0
36}
37
38pub fn is_aligned<T>(ptr: *const T) -> bool {
39 is_aligned_custom(ptr, DEFAULTALIGN)
40}
41
42pub fn assert_alignement<T>(ptr: *const T) {
43 assert!(
44 is_aligned(ptr),
45 "invalid alignement: ensure passed bytes have been allocated with [alloc_aligned_u8] or [alloc_aligned]"
46 )
47}
48
49pub fn cast<T, V>(data: &[T]) -> &[V] {
50 let ptr: *const V = data.as_ptr() as *const V;
51 let len: usize = data.len() / std::mem::size_of::<V>();
52 unsafe { std::slice::from_raw_parts(ptr, len) }
53}
54
55pub fn cast_mut<T, V>(data: &[T]) -> &mut [V] {
56 let ptr: *mut V = data.as_ptr() as *mut V;
57 let len: usize = data.len() / std::mem::size_of::<V>();
58 unsafe { std::slice::from_raw_parts_mut(ptr, len) }
59}
60
61/// Allocates a block of bytes with a custom alignement.
62/// Alignement must be a power of two and size a multiple of the alignement.
63/// Allocated memory is initialized to zero.
64fn alloc_aligned_custom_u8(size: usize, align: usize) -> Vec<u8> {
65 assert!(
66 align.is_power_of_two(),
67 "Alignment must be a power of two but is {}",
68 align
69 );
70 assert_eq!(
71 (size * std::mem::size_of::<u8>()) % align,
72 0,
73 "size={} must be a multiple of align={}",
74 size,
75 align
76 );
77 unsafe {
78 let layout: std::alloc::Layout = std::alloc::Layout::from_size_align(size, align).expect("Invalid alignment");
79 let ptr: *mut u8 = std::alloc::alloc(layout);
80 if ptr.is_null() {
81 panic!("Memory allocation failed");
82 }
83 assert!(
84 is_aligned_custom(ptr, align),
85 "Memory allocation at {:p} is not aligned to {} bytes",
86 ptr,
87 align
88 );
89 // Init allocated memory to zero
90 std::ptr::write_bytes(ptr, 0, size);
91 Vec::from_raw_parts(ptr, size, size)
92 }
93}
94
95/// Allocates a block of T aligned with [DEFAULTALIGN].
96/// Size of T * size msut be a multiple of [DEFAULTALIGN].
97pub fn alloc_aligned_custom<T>(size: usize, align: usize) -> Vec<T> {
98 assert_eq!(
99 (size * std::mem::size_of::<T>()) % align,
100 0,
101 "size={} must be a multiple of align={}",
102 size,
103 align
104 );
105 let mut vec_u8: Vec<u8> = alloc_aligned_custom_u8(std::mem::size_of::<T>() * size, align);
106 let ptr: *mut T = vec_u8.as_mut_ptr() as *mut T;
107 let len: usize = vec_u8.len() / std::mem::size_of::<T>();
108 let cap: usize = vec_u8.capacity() / std::mem::size_of::<T>();
109 std::mem::forget(vec_u8);
110 unsafe { Vec::from_raw_parts(ptr, len, cap) }
111}
112
113// Allocates an aligned of size equal to the smallest power of two equal or greater to `size` that is
114// at least as bit as DEFAULTALIGN / std::mem::size_of::<T>().
115pub fn alloc_aligned<T>(size: usize) -> Vec<T> {
116 alloc_aligned_custom::<T>(
117 std::cmp::max(
118 size.next_power_of_two(),
119 DEFAULTALIGN / std::mem::size_of::<T>(),
120 ),
121 DEFAULTALIGN,
122 )
123}