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