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 /// b <- IDFT(a), uses a as scratch space.
46 fn vec_znx_idft_tmp_a<R, A>(&self, res: &mut R, res_col: usize, a: &mut A, a_cols: usize)
47 where
48 R: VecZnxBigToMut<B>,
49 A: VecZnxDftToMut<B>;
50 fn vec_znx_idft_consume<D>(&self, a: VecZnxDft<D, B>, a_cols: usize) -> VecZnxBig<D, FFT64>
51 where
52 VecZnxDft<D, FFT64>: VecZnxDftToMut<FFT64>;
53
54 fn vec_znx_idft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
55 where
56 R: VecZnxBigToMut<B>,
57 A: VecZnxDftToRef<B>;
58
59 fn vec_znx_dft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
60 where
61 R: VecZnxDftToMut<B>,
62 A: VecZnxToRef;
63}
64
65impl<B: Backend> VecZnxDftAlloc<B> for Module<B> {
66 fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDftOwned<B> {
67 VecZnxDftOwned::new(&self, cols, size)
68 }
69
70 fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<B> {
71 VecZnxDftOwned::new_from_bytes(self, cols, size, bytes)
72 }
73
74 fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize {
75 bytes_of_vec_znx_dft(self, cols, size)
76 }
77}
78
79impl VecZnxDftOps<FFT64> for Module<FFT64> {
80 fn vec_znx_idft_tmp_a<R, A>(&self, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
81 where
82 R: VecZnxBigToMut<FFT64>,
83 A: VecZnxDftToMut<FFT64>,
84 {
85 let mut res_mut = res.to_mut();
86 let mut a_mut = a.to_mut();
87
88 let min_size: usize = min(res_mut.size(), a_mut.size());
89
90 unsafe {
91 (0..min_size).for_each(|j| {
92 vec_znx_dft::vec_znx_idft_tmp_a(
93 self.ptr,
94 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
95 1 as u64,
96 a_mut.at_mut_ptr(a_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
97 1 as u64,
98 )
99 });
100 (min_size..res_mut.size()).for_each(|j| {
101 res_mut.zero_at(res_col, j);
102 })
103 }
104 }
105
106 fn vec_znx_idft_consume<D>(&self, mut a: VecZnxDft<D, FFT64>, a_col: usize) -> VecZnxBig<D, FFT64>
107 where
108 VecZnxDft<D, FFT64>: VecZnxDftToMut<FFT64>,
109 {
110 let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
111
112 unsafe {
113 (0..a_mut.size()).for_each(|j| {
114 vec_znx_dft::vec_znx_idft_tmp_a(
115 self.ptr,
116 a_mut.at_mut_ptr(a_col, j) as *mut vec_znx_big::vec_znx_big_t,
117 1 as u64,
118 a_mut.at_mut_ptr(a_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
119 1 as u64,
120 )
121 });
122
123 a.into_big()
124 }
125 }
126
127 fn vec_znx_idft_tmp_bytes(&self) -> usize {
128 unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.ptr) as usize }
129 }
130
131 /// b <- DFT(a)
132 ///
133 /// # Panics
134 /// If b.cols < a_cols
135 fn vec_znx_dft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
136 where
137 R: VecZnxDftToMut<FFT64>,
138 A: VecZnxToRef,
139 {
140 let mut res_mut = res.to_mut();
141 let a_ref = a.to_ref();
142
143 let min_size: usize = min(res_mut.size(), a_ref.size());
144
145 unsafe {
146 (0..min_size).for_each(|j| {
147 vec_znx_dft::vec_znx_dft(
148 self.ptr,
149 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
150 1 as u64,
151 a_ref.at_ptr(a_col, j),
152 1 as u64,
153 a_ref.sl() as u64,
154 )
155 });
156 (min_size..res_mut.size()).for_each(|j| {
157 res_mut.zero_at(res_col, j);
158 });
159 }
160 }
161
162 // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
163 fn vec_znx_idft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
164 where
165 R: VecZnxBigToMut<FFT64>,
166 A: VecZnxDftToRef<FFT64>,
167 {
168 let mut res_mut = res.to_mut();
169 let a_ref = a.to_ref();
170
171 let (tmp_bytes, _) = scratch.tmp_slice(self.vec_znx_idft_tmp_bytes());
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_idft(
178 self.ptr,
179 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
180 1 as u64,
181 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
182 1 as u64,
183 tmp_bytes.as_mut_ptr(),
184 )
185 });
186 (min_size..res_mut.size()).for_each(|j| {
187 res_mut.zero_at(res_col, j);
188 });
189 }
190 }
191}