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