poulpy/base2k/src/vec_znx_dft_ops.rs
base2k/src/vec_znx_dft_ops.rsRaw
1use crate::ffi::vec_znx_big;
2use crate::ffi::vec_znx_dft;
3use crate::znx_base::ZnxAlloc;
4use crate::znx_base::ZnxInfos;
5use crate::znx_base::ZnxLayout;
6use crate::znx_base::ZnxSliceSize;
7use crate::{Backend, FFT64, Module, VecZnx, VecZnxBig, VecZnxDft, ZnxZero, assert_alignement};
8use std::cmp::min;
9
10pub trait VecZnxDftOps<B: Backend> {
11 /// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
12 fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDft<B>;
13
14 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
15 ///
16 /// Behavior: takes ownership of the backing array.
17 ///
18 /// # Arguments
19 ///
20 /// * `cols`: the number of cols of the [VecZnxDft].
21 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
22 ///
23 /// # Panics
24 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
25 fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDft<B>;
26
27 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
28 ///
29 /// Behavior: the backing array is only borrowed.
30 ///
31 /// # Arguments
32 ///
33 /// * `cols`: the number of cols of the [VecZnxDft].
34 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
35 ///
36 /// # Panics
37 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
38 fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnxDft<B>;
39
40 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
41 ///
42 /// # Arguments
43 ///
44 /// * `cols`: the number of cols of the [VecZnxDft].
45 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
46 ///
47 /// # Panics
48 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
49 fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize;
50
51 /// Returns the minimum number of bytes necessary to allocate
52 /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
53 fn vec_znx_idft_tmp_bytes(&self) -> usize;
54
55 /// b <- IDFT(a), uses a as scratch space.
56 fn vec_znx_idft_tmp_a(&self, res: &mut VecZnxBig<B>, res_col: usize, a: &mut VecZnxDft<B>, a_cols: usize);
57
58 fn vec_znx_idft(&self, res: &mut VecZnxBig<B>, res_col: usize, a: &VecZnxDft<B>, a_col: usize, tmp_bytes: &mut [u8]);
59
60 fn vec_znx_dft(&self, res: &mut VecZnxDft<B>, res_col: usize, a: &VecZnx, a_col: usize);
61}
62
63impl VecZnxDftOps<FFT64> for Module<FFT64> {
64 fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDft<FFT64> {
65 VecZnxDft::<FFT64>::new(&self, 1, cols, size)
66 }
67
68 fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDft<FFT64> {
69 VecZnxDft::from_bytes(self, 1, cols, size, bytes)
70 }
71
72 fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnxDft<FFT64> {
73 VecZnxDft::from_bytes_borrow(self, 1, cols, size, bytes)
74 }
75
76 fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize {
77 VecZnxDft::bytes_of(&self, 1, cols, size)
78 }
79
80 fn vec_znx_idft_tmp_a(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &mut VecZnxDft<FFT64>, a_col: usize) {
81 let min_size: usize = min(res.size(), a.size());
82
83 unsafe {
84 (0..min_size).for_each(|j| {
85 vec_znx_dft::vec_znx_idft_tmp_a(
86 self.ptr,
87 res.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
88 1 as u64,
89 a.at_ptr(a_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
90 1 as u64,
91 )
92 });
93 (min_size..res.size()).for_each(|j| {
94 res.zero_at(res_col, j);
95 })
96 }
97 }
98
99 fn vec_znx_idft_tmp_bytes(&self) -> usize {
100 unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.ptr) as usize }
101 }
102
103 /// b <- DFT(a)
104 ///
105 /// # Panics
106 /// If b.cols < a_cols
107 fn vec_znx_dft(&self, res: &mut VecZnxDft<FFT64>, res_col: usize, a: &VecZnx, a_col: usize) {
108 let min_size: usize = min(res.size(), a.size());
109
110 unsafe {
111 (0..min_size).for_each(|j| {
112 vec_znx_dft::vec_znx_dft(
113 self.ptr,
114 res.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
115 1 as u64,
116 a.at_ptr(a_col, j),
117 1 as u64,
118 a.sl() as u64,
119 )
120 });
121 (min_size..res.size()).for_each(|j| {
122 res.zero_at(res_col, j);
123 });
124 }
125 }
126
127 // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
128 fn vec_znx_idft(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnxDft<FFT64>, a_col: usize, tmp_bytes: &mut [u8]) {
129 #[cfg(debug_assertions)]
130 {
131 assert!(
132 tmp_bytes.len() >= Self::vec_znx_idft_tmp_bytes(self),
133 "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
134 tmp_bytes.len(),
135 Self::vec_znx_idft_tmp_bytes(self)
136 );
137 assert_alignement(tmp_bytes.as_ptr())
138 }
139
140 let min_size: usize = min(res.size(), a.size());
141
142 unsafe {
143 (0..min_size).for_each(|j| {
144 vec_znx_dft::vec_znx_idft(
145 self.ptr,
146 res.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
147 1 as u64,
148 a.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
149 1 as u64,
150 tmp_bytes.as_mut_ptr(),
151 )
152 });
153 (min_size..res.size()).for_each(|j| {
154 res.zero_at(res_col, j);
155 });
156 }
157 }
158}