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