poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/vec_rnx_api.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/vec_rnx_api.cRaw
1#include <string.h>
2
3#include "vec_rnx_arithmetic_private.h"
4
5void fft64_init_rnx_module_precomp(MOD_RNX* module) {
6 // Add here initialization of items that are in the precomp
7 const uint64_t m = module->m;
8 module->precomp.fft64.p_fft = new_reim_fft_precomp(m, 0);
9 module->precomp.fft64.p_ifft = new_reim_ifft_precomp(m, 0);
10 module->precomp.fft64.p_fftvec_add = new_reim_fftvec_add_precomp(m);
11 module->precomp.fft64.p_fftvec_mul = new_reim_fftvec_mul_precomp(m);
12 module->precomp.fft64.p_fftvec_addmul = new_reim_fftvec_addmul_precomp(m);
13}
14
15void fft64_finalize_rnx_module_precomp(MOD_RNX* module) {
16 // Add here deleters for items that are in the precomp
17 delete_reim_fft_precomp(module->precomp.fft64.p_fft);
18 delete_reim_ifft_precomp(module->precomp.fft64.p_ifft);
19 delete_reim_fftvec_add_precomp(module->precomp.fft64.p_fftvec_add);
20 delete_reim_fftvec_mul_precomp(module->precomp.fft64.p_fftvec_mul);
21 delete_reim_fftvec_addmul_precomp(module->precomp.fft64.p_fftvec_addmul);
22}
23
24void fft64_init_rnx_module_vtable(MOD_RNX* module) {
25 // Add function pointers here
26 module->vtable.vec_rnx_add = vec_rnx_add_ref;
27 module->vtable.vec_rnx_zero = vec_rnx_zero_ref;
28 module->vtable.vec_rnx_copy = vec_rnx_copy_ref;
29 module->vtable.vec_rnx_negate = vec_rnx_negate_ref;
30 module->vtable.vec_rnx_sub = vec_rnx_sub_ref;
31 module->vtable.vec_rnx_rotate = vec_rnx_rotate_ref;
32 module->vtable.vec_rnx_automorphism = vec_rnx_automorphism_ref;
33 module->vtable.vec_rnx_mul_xp_minus_one = vec_rnx_mul_xp_minus_one_ref;
34 module->vtable.rnx_vmp_apply_dft_to_dft_tmp_bytes = fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref;
35 module->vtable.rnx_vmp_apply_dft_to_dft = fft64_rnx_vmp_apply_dft_to_dft_ref;
36 module->vtable.rnx_vmp_apply_tmp_a_tmp_bytes = fft64_rnx_vmp_apply_tmp_a_tmp_bytes_ref;
37 module->vtable.rnx_vmp_apply_tmp_a = fft64_rnx_vmp_apply_tmp_a_ref;
38 module->vtable.rnx_vmp_prepare_tmp_bytes = fft64_rnx_vmp_prepare_tmp_bytes_ref;
39 module->vtable.rnx_vmp_prepare_contiguous = fft64_rnx_vmp_prepare_contiguous_ref;
40 module->vtable.rnx_vmp_prepare_dblptr = fft64_rnx_vmp_prepare_dblptr_ref;
41 module->vtable.rnx_vmp_prepare_row = fft64_rnx_vmp_prepare_row_ref;
42 module->vtable.bytes_of_rnx_vmp_pmat = fft64_bytes_of_rnx_vmp_pmat;
43 module->vtable.rnx_approxdecomp_from_tnxdbl = rnx_approxdecomp_from_tnxdbl_ref;
44 module->vtable.vec_rnx_to_znx32 = vec_rnx_to_znx32_ref;
45 module->vtable.vec_rnx_from_znx32 = vec_rnx_from_znx32_ref;
46 module->vtable.vec_rnx_to_tnx32 = vec_rnx_to_tnx32_ref;
47 module->vtable.vec_rnx_from_tnx32 = vec_rnx_from_tnx32_ref;
48 module->vtable.vec_rnx_to_tnxdbl = vec_rnx_to_tnxdbl_ref;
49 module->vtable.bytes_of_rnx_svp_ppol = fft64_bytes_of_rnx_svp_ppol;
50 module->vtable.rnx_svp_prepare = fft64_rnx_svp_prepare_ref;
51 module->vtable.rnx_svp_apply = fft64_rnx_svp_apply_ref;
52
53 // Add optimized function pointers here
54 if (CPU_SUPPORTS("avx")) {
55 module->vtable.vec_rnx_add = vec_rnx_add_avx;
56 module->vtable.vec_rnx_sub = vec_rnx_sub_avx;
57 module->vtable.vec_rnx_negate = vec_rnx_negate_avx;
58 module->vtable.rnx_vmp_apply_dft_to_dft_tmp_bytes = fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_avx;
59 module->vtable.rnx_vmp_apply_dft_to_dft = fft64_rnx_vmp_apply_dft_to_dft_avx;
60 module->vtable.rnx_vmp_apply_tmp_a_tmp_bytes = fft64_rnx_vmp_apply_tmp_a_tmp_bytes_avx;
61 module->vtable.rnx_vmp_apply_tmp_a = fft64_rnx_vmp_apply_tmp_a_avx;
62 module->vtable.rnx_vmp_prepare_tmp_bytes = fft64_rnx_vmp_prepare_tmp_bytes_avx;
63 module->vtable.rnx_vmp_prepare_contiguous = fft64_rnx_vmp_prepare_contiguous_avx;
64 module->vtable.rnx_vmp_prepare_dblptr = fft64_rnx_vmp_prepare_dblptr_avx;
65 module->vtable.rnx_vmp_prepare_row = fft64_rnx_vmp_prepare_row_avx;
66 module->vtable.rnx_approxdecomp_from_tnxdbl = rnx_approxdecomp_from_tnxdbl_avx;
67 }
68}
69
70void init_rnx_module_info(MOD_RNX* module, //
71 uint64_t n, RNX_MODULE_TYPE mtype) {
72 memset(module, 0, sizeof(MOD_RNX));
73 module->n = n;
74 module->m = n >> 1;
75 module->mtype = mtype;
76 switch (mtype) {
77 case FFT64:
78 fft64_init_rnx_module_precomp(module);
79 fft64_init_rnx_module_vtable(module);
80 break;
81 default:
82 NOT_SUPPORTED(); // unknown mtype
83 }
84}
85
86void finalize_rnx_module_info(MOD_RNX* module) {
87 if (module->custom) module->custom_deleter(module->custom);
88 switch (module->mtype) {
89 case FFT64:
90 fft64_finalize_rnx_module_precomp(module);
91 // fft64_finalize_rnx_module_vtable(module); // nothing to finalize
92 break;
93 default:
94 NOT_SUPPORTED(); // unknown mtype
95 }
96}
97
98EXPORT MOD_RNX* new_rnx_module_info(uint64_t nn, RNX_MODULE_TYPE mtype) {
99 MOD_RNX* res = (MOD_RNX*)malloc(sizeof(MOD_RNX));
100 init_rnx_module_info(res, nn, mtype);
101 return res;
102}
103
104EXPORT void delete_rnx_module_info(MOD_RNX* module_info) {
105 finalize_rnx_module_info(module_info);
106 free(module_info);
107}
108
109EXPORT uint64_t rnx_module_get_n(const MOD_RNX* module) { return module->n; }
110
111/** @brief allocates a prepared matrix (release with delete_rnx_vmp_pmat) */
112EXPORT RNX_VMP_PMAT* new_rnx_vmp_pmat(const MOD_RNX* module, // N
113 uint64_t nrows, uint64_t ncols) { // dimensions
114 return (RNX_VMP_PMAT*)spqlios_alloc(bytes_of_rnx_vmp_pmat(module, nrows, ncols));
115}
116EXPORT void delete_rnx_vmp_pmat(RNX_VMP_PMAT* ptr) { spqlios_free(ptr); }
117
118//////////////// wrappers //////////////////
119
120/** @brief sets res = a + b */
121EXPORT void vec_rnx_add( //
122 const MOD_RNX* module, // N
123 double* res, uint64_t res_size, uint64_t res_sl, // res
124 const double* a, uint64_t a_size, uint64_t a_sl, // a
125 const double* b, uint64_t b_size, uint64_t b_sl // b
126) {
127 module->vtable.vec_rnx_add(module, res, res_size, res_sl, a, a_size, a_sl, b, b_size, b_sl);
128}
129
130/** @brief sets res = 0 */
131EXPORT void vec_rnx_zero( //
132 const MOD_RNX* module, // N
133 double* res, uint64_t res_size, uint64_t res_sl // res
134) {
135 module->vtable.vec_rnx_zero(module, res, res_size, res_sl);
136}
137
138/** @brief sets res = a */
139EXPORT void vec_rnx_copy( //
140 const MOD_RNX* module, // N
141 double* res, uint64_t res_size, uint64_t res_sl, // res
142 const double* a, uint64_t a_size, uint64_t a_sl // a
143) {
144 module->vtable.vec_rnx_copy(module, res, res_size, res_sl, a, a_size, a_sl);
145}
146
147/** @brief sets res = -a */
148EXPORT void vec_rnx_negate( //
149 const MOD_RNX* module, // N
150 double* res, uint64_t res_size, uint64_t res_sl, // res
151 const double* a, uint64_t a_size, uint64_t a_sl // a
152) {
153 module->vtable.vec_rnx_negate(module, res, res_size, res_sl, a, a_size, a_sl);
154}
155
156/** @brief sets res = a - b */
157EXPORT void vec_rnx_sub( //
158 const MOD_RNX* module, // N
159 double* res, uint64_t res_size, uint64_t res_sl, // res
160 const double* a, uint64_t a_size, uint64_t a_sl, // a
161 const double* b, uint64_t b_size, uint64_t b_sl // b
162) {
163 module->vtable.vec_rnx_sub(module, res, res_size, res_sl, a, a_size, a_sl, b, b_size, b_sl);
164}
165
166/** @brief sets res = a . X^p */
167EXPORT void vec_rnx_rotate( //
168 const MOD_RNX* module, // N
169 const int64_t p, // rotation value
170 double* res, uint64_t res_size, uint64_t res_sl, // res
171 const double* a, uint64_t a_size, uint64_t a_sl // a
172) {
173 module->vtable.vec_rnx_rotate(module, p, res, res_size, res_sl, a, a_size, a_sl);
174}
175
176/** @brief sets res = a(X^p) */
177EXPORT void vec_rnx_automorphism( //
178 const MOD_RNX* module, // N
179 int64_t p, // X -> X^p
180 double* res, uint64_t res_size, uint64_t res_sl, // res
181 const double* a, uint64_t a_size, uint64_t a_sl // a
182) {
183 module->vtable.vec_rnx_automorphism(module, p, res, res_size, res_sl, a, a_size, a_sl);
184}
185
186EXPORT void vec_rnx_mul_xp_minus_one( //
187 const MOD_RNX* module, // N
188 const int64_t p, // rotation value
189 double* res, uint64_t res_size, uint64_t res_sl, // res
190 const double* a, uint64_t a_size, uint64_t a_sl // a
191) {
192 module->vtable.vec_rnx_mul_xp_minus_one(module, p, res, res_size, res_sl, a, a_size, a_sl);
193}
194/** @brief number of bytes in a RNX_VMP_PMAT (for manual allocation) */
195EXPORT uint64_t bytes_of_rnx_vmp_pmat(const MOD_RNX* module, // N
196 uint64_t nrows, uint64_t ncols) { // dimensions
197 return module->vtable.bytes_of_rnx_vmp_pmat(module, nrows, ncols);
198}
199
200/** @brief prepares a vmp matrix (contiguous row-major version) */
201EXPORT void rnx_vmp_prepare_contiguous( //
202 const MOD_RNX* module, // N
203 RNX_VMP_PMAT* pmat, // output
204 const double* a, uint64_t nrows, uint64_t ncols, // a
205 uint8_t* tmp_space // scratch space
206) {
207 module->vtable.rnx_vmp_prepare_contiguous(module, pmat, a, nrows, ncols, tmp_space);
208}
209
210/** @brief prepares a vmp matrix (mat[row]+col*N points to the item) */
211EXPORT void rnx_vmp_prepare_dblptr( //
212 const MOD_RNX* module, // N
213 RNX_VMP_PMAT* pmat, // output
214 const double** a, uint64_t nrows, uint64_t ncols, // a
215 uint8_t* tmp_space // scratch space
216) {
217 module->vtable.rnx_vmp_prepare_dblptr(module, pmat, a, nrows, ncols, tmp_space);
218}
219
220/** @brief prepares the ith-row of a vmp matrix with nrows and ncols */
221EXPORT void rnx_vmp_prepare_row( //
222 const MOD_RNX* module, // N
223 RNX_VMP_PMAT* pmat, // output
224 const double* a, uint64_t row_i, uint64_t nrows, uint64_t ncols, // a
225 uint8_t* tmp_space // scratch space
226) {
227 module->vtable.rnx_vmp_prepare_row(module, pmat, a, row_i, nrows, ncols, tmp_space);
228}
229
230/** @brief number of scratch bytes necessary to prepare a matrix */
231EXPORT uint64_t rnx_vmp_prepare_tmp_bytes(const MOD_RNX* module) {
232 return module->vtable.rnx_vmp_prepare_tmp_bytes(module);
233}
234
235/** @brief applies a vmp product res = a x pmat */
236EXPORT void rnx_vmp_apply_tmp_a( //
237 const MOD_RNX* module, // N
238 double* res, uint64_t res_size, uint64_t res_sl, // res
239 double* tmpa, uint64_t a_size, uint64_t a_sl, // a (will be overwritten)
240 const RNX_VMP_PMAT* pmat, uint64_t nrows, uint64_t ncols, // prep matrix
241 uint8_t* tmp_space // scratch space
242) {
243 module->vtable.rnx_vmp_apply_tmp_a(module, res, res_size, res_sl, tmpa, a_size, a_sl, pmat, nrows, ncols, tmp_space);
244}
245
246EXPORT uint64_t rnx_vmp_apply_tmp_a_tmp_bytes( //
247 const MOD_RNX* module, // N
248 uint64_t res_size, // res size
249 uint64_t a_size, // a size
250 uint64_t nrows, uint64_t ncols // prep matrix dims
251) {
252 return module->vtable.rnx_vmp_apply_tmp_a_tmp_bytes(module, res_size, a_size, nrows, ncols);
253}
254
255/** @brief minimal size of the tmp_space */
256EXPORT void rnx_vmp_apply_dft_to_dft( //
257 const MOD_RNX* module, // N
258 double* res, uint64_t res_size, uint64_t res_sl, // res
259 const double* a_dft, uint64_t a_size, uint64_t a_sl, // a
260 const RNX_VMP_PMAT* pmat, uint64_t nrows, uint64_t ncols, // prep matrix
261 uint8_t* tmp_space // scratch space (a_size*sizeof(reim4) bytes)
262) {
263 module->vtable.rnx_vmp_apply_dft_to_dft(module, res, res_size, res_sl, a_dft, a_size, a_sl, pmat, nrows, ncols,
264 tmp_space);
265}
266
267/** @brief minimal size of the tmp_space */
268EXPORT uint64_t rnx_vmp_apply_dft_to_dft_tmp_bytes( //
269 const MOD_RNX* module, // N
270 uint64_t res_size, // res
271 uint64_t a_size, // a
272 uint64_t nrows, uint64_t ncols // prep matrix
273) {
274 return module->vtable.rnx_vmp_apply_dft_to_dft_tmp_bytes(module, res_size, a_size, nrows, ncols);
275}
276
277EXPORT uint64_t bytes_of_rnx_svp_ppol(const MOD_RNX* module) { return module->vtable.bytes_of_rnx_svp_ppol(module); }
278
279EXPORT void rnx_svp_prepare(const MOD_RNX* module, // N
280 RNX_SVP_PPOL* ppol, // output
281 const double* pol // a
282) {
283 module->vtable.rnx_svp_prepare(module, ppol, pol);
284}
285
286EXPORT void rnx_svp_apply( //
287 const MOD_RNX* module, // N
288 double* res, uint64_t res_size, uint64_t res_sl, // output
289 const RNX_SVP_PPOL* ppol, // prepared pol
290 const double* a, uint64_t a_size, uint64_t a_sl // a
291) {
292 module->vtable.rnx_svp_apply(module, // N
293 res, res_size, res_sl, // output
294 ppol, // prepared pol
295 a, a_size, a_sl);
296}
297
298EXPORT void rnx_approxdecomp_from_tnxdbl( //
299 const MOD_RNX* module, // N
300 const TNXDBL_APPROXDECOMP_GADGET* gadget, // output base 2^K
301 double* res, uint64_t res_size, uint64_t res_sl, // res
302 const double* a) { // a
303 module->vtable.rnx_approxdecomp_from_tnxdbl(module, gadget, res, res_size, res_sl, a);
304}
305
306EXPORT void vec_rnx_to_znx32( //
307 const MOD_RNX* module, // N
308 int32_t* res, uint64_t res_size, uint64_t res_sl, // res
309 const double* a, uint64_t a_size, uint64_t a_sl // a
310) {
311 module->vtable.vec_rnx_to_znx32(module, res, res_size, res_sl, a, a_size, a_sl);
312}
313
314EXPORT void vec_rnx_from_znx32( //
315 const MOD_RNX* module, // N
316 double* res, uint64_t res_size, uint64_t res_sl, // res
317 const int32_t* a, uint64_t a_size, uint64_t a_sl // a
318) {
319 module->vtable.vec_rnx_from_znx32(module, res, res_size, res_sl, a, a_size, a_sl);
320}
321
322EXPORT void vec_rnx_to_tnx32( //
323 const MOD_RNX* module, // N
324 int32_t* res, uint64_t res_size, uint64_t res_sl, // res
325 const double* a, uint64_t a_size, uint64_t a_sl // a
326) {
327 module->vtable.vec_rnx_to_tnx32(module, res, res_size, res_sl, a, a_size, a_sl);
328}
329
330EXPORT void vec_rnx_from_tnx32( //
331 const MOD_RNX* module, // N
332 double* res, uint64_t res_size, uint64_t res_sl, // res
333 const int32_t* a, uint64_t a_size, uint64_t a_sl // a
334) {
335 module->vtable.vec_rnx_from_tnx32(module, res, res_size, res_sl, a, a_size, a_sl);
336}
337
338EXPORT void vec_rnx_to_tnxdbl( //
339 const MOD_RNX* module, // N
340 double* res, uint64_t res_size, uint64_t res_sl, // res
341 const double* a, uint64_t a_size, uint64_t a_sl // a
342) {
343 module->vtable.vec_rnx_to_tnxdbl(module, res, res_size, res_sl, a, a_size, a_sl);
344}