poulpy/base2k/src/vec_znx_dft_ops.rs
base2k/src/vec_znx_dft_ops.rsRaw
1use crate::ffi::{vec_znx_big, vec_znx_dft};
2use crate::vec_znx_dft::bytes_of_vec_znx_dft;
3use crate::znx_base::ZnxInfos;
4use crate::{
5 Backend, Scratch, VecZnxBig, VecZnxBigToMut, VecZnxDft, VecZnxDftOwned, VecZnxDftToMut, VecZnxDftToRef, VecZnxToRef,
6 ZnxSliceSize,
7};
8use crate::{FFT64, Module, ZnxView, ZnxViewMut, ZnxZero};
9use std::cmp::min;
10
11pub trait VecZnxDftAlloc<B: Backend> {
12 /// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
13 fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDftOwned<B>;
14
15 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
16 ///
17 /// Behavior: takes ownership of the backing array.
18 ///
19 /// # Arguments
20 ///
21 /// * `cols`: the number of cols of the [VecZnxDft].
22 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
23 ///
24 /// # Panics
25 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
26 fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<B>;
27
28 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
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 bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize;
38}
39
40pub trait VecZnxDftOps<B: Backend> {
41 /// Returns the minimum number of bytes necessary to allocate
42 /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
43 fn vec_znx_idft_tmp_bytes(&self) -> usize;
44
45 fn vec_znx_dft_copy<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
46 where
47 R: VecZnxDftToMut<B>,
48 A: VecZnxDftToRef<B>;
49
50 /// b <- IDFT(a), uses a as scratch space.
51 fn vec_znx_idft_tmp_a<R, A>(&self, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
52 where
53 R: VecZnxBigToMut<B>,
54 A: VecZnxDftToMut<B>;
55
56 /// Consumes a to return IDFT(a) in big coeff space.
57 fn vec_znx_idft_consume<D>(&self, a: VecZnxDft<D, B>) -> VecZnxBig<D, FFT64>
58 where
59 VecZnxDft<D, FFT64>: VecZnxDftToMut<FFT64>;
60
61 fn vec_znx_idft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
62 where
63 R: VecZnxBigToMut<B>,
64 A: VecZnxDftToRef<B>;
65
66 fn vec_znx_dft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
67 where
68 R: VecZnxDftToMut<B>,
69 A: VecZnxToRef;
70}
71
72impl<B: Backend> VecZnxDftAlloc<B> for Module<B> {
73 fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDftOwned<B> {
74 VecZnxDftOwned::new(&self, cols, size)
75 }
76
77 fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<B> {
78 VecZnxDftOwned::new_from_bytes(self, cols, size, bytes)
79 }
80
81 fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize {
82 bytes_of_vec_znx_dft(self, cols, size)
83 }
84}
85
86impl VecZnxDftOps<FFT64> for Module<FFT64> {
87 fn vec_znx_dft_copy<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
88 where
89 R: VecZnxDftToMut<FFT64>,
90 A: VecZnxDftToRef<FFT64>,
91 {
92 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
93 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
94
95 let min_size: usize = min(res_mut.size(), a_ref.size());
96
97 (0..min_size).for_each(|j| {
98 res_mut
99 .at_mut(res_col, j)
100 .copy_from_slice(a_ref.at(a_col, j));
101 });
102 (min_size..res_mut.size()).for_each(|j| {
103 res_mut.zero_at(res_col, j);
104 })
105 }
106
107 fn vec_znx_idft_tmp_a<R, A>(&self, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
108 where
109 R: VecZnxBigToMut<FFT64>,
110 A: VecZnxDftToMut<FFT64>,
111 {
112 let mut res_mut: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
113 let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
114
115 let min_size: usize = min(res_mut.size(), a_mut.size());
116
117 unsafe {
118 (0..min_size).for_each(|j| {
119 vec_znx_dft::vec_znx_idft_tmp_a(
120 self.ptr,
121 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
122 1 as u64,
123 a_mut.at_mut_ptr(a_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
124 1 as u64,
125 )
126 });
127 (min_size..res_mut.size()).for_each(|j| {
128 res_mut.zero_at(res_col, j);
129 })
130 }
131 }
132
133 fn vec_znx_idft_consume<D>(&self, mut a: VecZnxDft<D, FFT64>) -> VecZnxBig<D, FFT64>
134 where
135 VecZnxDft<D, FFT64>: VecZnxDftToMut<FFT64>,
136 {
137 let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
138
139 unsafe {
140 // Rev col and rows because ZnxDft.sl() >= ZnxBig.sl()
141 (0..a_mut.size()).for_each(|j| {
142 (0..a_mut.cols()).for_each(|i| {
143 vec_znx_dft::vec_znx_idft_tmp_a(
144 self.ptr,
145 a_mut.at_mut_ptr(i, j) as *mut vec_znx_big::vec_znx_big_t,
146 1 as u64,
147 a_mut.at_mut_ptr(i, j) as *mut vec_znx_dft::vec_znx_dft_t,
148 1 as u64,
149 )
150 });
151 });
152 }
153
154 a.into_big()
155 }
156
157 fn vec_znx_idft_tmp_bytes(&self) -> usize {
158 unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.ptr) as usize }
159 }
160
161 /// b <- DFT(a)
162 ///
163 /// # Panics
164 /// If b.cols < a_col
165 fn vec_znx_dft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
166 where
167 R: VecZnxDftToMut<FFT64>,
168 A: VecZnxToRef,
169 {
170 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
171 let a_ref: crate::VecZnx<&[u8]> = a.to_ref();
172
173 let min_size: usize = min(res_mut.size(), a_ref.size());
174
175 unsafe {
176 (0..min_size).for_each(|j| {
177 vec_znx_dft::vec_znx_dft(
178 self.ptr,
179 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
180 1 as u64,
181 a_ref.at_ptr(a_col, j),
182 1 as u64,
183 a_ref.sl() as u64,
184 )
185 });
186 (min_size..res_mut.size()).for_each(|j| {
187 res_mut.zero_at(res_col, j);
188 });
189 }
190 }
191
192 // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
193 fn vec_znx_idft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
194 where
195 R: VecZnxBigToMut<FFT64>,
196 A: VecZnxDftToRef<FFT64>,
197 {
198 let mut res_mut: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
199 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
200
201 let (tmp_bytes, _) = scratch.tmp_slice(self.vec_znx_idft_tmp_bytes());
202
203 let min_size: usize = min(res_mut.size(), a_ref.size());
204
205 unsafe {
206 (0..min_size).for_each(|j| {
207 vec_znx_dft::vec_znx_idft(
208 self.ptr,
209 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
210 1 as u64,
211 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
212 1 as u64,
213 tmp_bytes.as_mut_ptr(),
214 )
215 });
216 (min_size..res_mut.size()).for_each(|j| {
217 res_mut.zero_at(res_col, j);
218 });
219 }
220 }
221}