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