poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/vec_rnx_arithmetic.h
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/vec_rnx_arithmetic.hRaw
1#ifndef SPQLIOS_VEC_RNX_ARITHMETIC_H
2#define SPQLIOS_VEC_RNX_ARITHMETIC_H
3
4#include <stdint.h>
5
6#include "../commons.h"
7
8/**
9 * We support the following module families:
10 * - FFT64:
11 * the overall precision should fit at all times over 52 bits.
12 */
13typedef enum rnx_module_type_t { FFT64 } RNX_MODULE_TYPE;
14
15/** @brief opaque structure that describes the modules (RnX,ZnX,TnX) and the hardware */
16typedef struct rnx_module_info_t MOD_RNX;
17
18/**
19 * @brief obtain a module info for ring dimension N
20 * the module-info knows about:
21 * - the dimension N (or the complex dimension m=N/2)
22 * - any moduleuted fft or ntt items
23 * - the hardware (avx, arm64, x86, ...)
24 */
25EXPORT MOD_RNX* new_rnx_module_info(uint64_t N, RNX_MODULE_TYPE mode);
26EXPORT void delete_rnx_module_info(MOD_RNX* module_info);
27EXPORT uint64_t rnx_module_get_n(const MOD_RNX* module);
28
29// basic arithmetic
30
31/** @brief sets res = 0 */
32EXPORT void vec_rnx_zero( //
33 const MOD_RNX* module, // N
34 double* res, uint64_t res_size, uint64_t res_sl // res
35);
36
37/** @brief sets res = a */
38EXPORT void vec_rnx_copy( //
39 const MOD_RNX* module, // N
40 double* res, uint64_t res_size, uint64_t res_sl, // res
41 const double* a, uint64_t a_size, uint64_t a_sl // a
42);
43
44/** @brief sets res = -a */
45EXPORT void vec_rnx_negate( //
46 const MOD_RNX* module, // N
47 double* res, uint64_t res_size, uint64_t res_sl, // res
48 const double* a, uint64_t a_size, uint64_t a_sl // a
49);
50
51/** @brief sets res = a + b */
52EXPORT void vec_rnx_add( //
53 const MOD_RNX* module, // N
54 double* res, uint64_t res_size, uint64_t res_sl, // res
55 const double* a, uint64_t a_size, uint64_t a_sl, // a
56 const double* b, uint64_t b_size, uint64_t b_sl // b
57);
58
59/** @brief sets res = a - b */
60EXPORT void vec_rnx_sub( //
61 const MOD_RNX* module, // N
62 double* res, uint64_t res_size, uint64_t res_sl, // res
63 const double* a, uint64_t a_size, uint64_t a_sl, // a
64 const double* b, uint64_t b_size, uint64_t b_sl // b
65);
66
67/** @brief sets res = a . X^p */
68EXPORT void vec_rnx_rotate( //
69 const MOD_RNX* module, // N
70 const int64_t p, // rotation value
71 double* res, uint64_t res_size, uint64_t res_sl, // res
72 const double* a, uint64_t a_size, uint64_t a_sl // a
73);
74
75/** @brief sets res = a . (X^p - 1) */
76EXPORT void vec_rnx_mul_xp_minus_one( //
77 const MOD_RNX* module, // N
78 const int64_t p, // rotation value
79 double* res, uint64_t res_size, uint64_t res_sl, // res
80 const double* a, uint64_t a_size, uint64_t a_sl // a
81);
82
83/** @brief sets res = a(X^p) */
84EXPORT void vec_rnx_automorphism( //
85 const MOD_RNX* module, // N
86 int64_t p, // X -> X^p
87 double* res, uint64_t res_size, uint64_t res_sl, // res
88 const double* a, uint64_t a_size, uint64_t a_sl // a
89);
90
91///////////////////////////////////////////////////////////////////
92// conversions //
93///////////////////////////////////////////////////////////////////
94
95EXPORT void vec_rnx_to_znx32( //
96 const MOD_RNX* module, // N
97 int32_t* res, uint64_t res_size, uint64_t res_sl, // res
98 const double* a, uint64_t a_size, uint64_t a_sl // a
99);
100
101EXPORT void vec_rnx_from_znx32( //
102 const MOD_RNX* module, // N
103 double* res, uint64_t res_size, uint64_t res_sl, // res
104 const int32_t* a, uint64_t a_size, uint64_t a_sl // a
105);
106
107EXPORT void vec_rnx_to_tnx32( //
108 const MOD_RNX* module, // N
109 int32_t* res, uint64_t res_size, uint64_t res_sl, // res
110 const double* a, uint64_t a_size, uint64_t a_sl // a
111);
112
113EXPORT void vec_rnx_from_tnx32( //
114 const MOD_RNX* module, // N
115 double* res, uint64_t res_size, uint64_t res_sl, // res
116 const int32_t* a, uint64_t a_size, uint64_t a_sl // a
117);
118
119EXPORT void vec_rnx_to_tnx32x2( //
120 const MOD_RNX* module, // N
121 int32_t* res, uint64_t res_size, uint64_t res_sl, // res
122 const double* a, uint64_t a_size, uint64_t a_sl // a
123);
124
125EXPORT void vec_rnx_from_tnx32x2( //
126 const MOD_RNX* module, // N
127 double* res, uint64_t res_size, uint64_t res_sl, // res
128 const int32_t* a, uint64_t a_size, uint64_t a_sl // a
129);
130
131EXPORT void vec_rnx_to_tnxdbl( //
132 const MOD_RNX* module, // N
133 double* res, uint64_t res_size, uint64_t res_sl, // res
134 const double* a, uint64_t a_size, uint64_t a_sl // a
135);
136
137///////////////////////////////////////////////////////////////////
138// isolated products (n.log(n), but not particularly optimized //
139///////////////////////////////////////////////////////////////////
140
141/** @brief res = a * b : small polynomial product */
142EXPORT void rnx_small_single_product( //
143 const MOD_RNX* module, // N
144 double* res, // output
145 const double* a, // a
146 const double* b, // b
147 uint8_t* tmp); // scratch space
148
149EXPORT uint64_t rnx_small_single_product_tmp_bytes(const MOD_RNX* module);
150
151/** @brief res = a * b centermod 1: small polynomial product */
152EXPORT void tnxdbl_small_single_product( //
153 const MOD_RNX* module, // N
154 double* torus_res, // output
155 const double* int_a, // a
156 const double* torus_b, // b
157 uint8_t* tmp); // scratch space
158
159EXPORT uint64_t tnxdbl_small_single_product_tmp_bytes(const MOD_RNX* module);
160
161/** @brief res = a * b: small polynomial product */
162EXPORT void znx32_small_single_product( //
163 const MOD_RNX* module, // N
164 int32_t* int_res, // output
165 const int32_t* int_a, // a
166 const int32_t* int_b, // b
167 uint8_t* tmp); // scratch space
168
169EXPORT uint64_t znx32_small_single_product_tmp_bytes(const MOD_RNX* module);
170
171/** @brief res = a * b centermod 1: small polynomial product */
172EXPORT void tnx32_small_single_product( //
173 const MOD_RNX* module, // N
174 int32_t* torus_res, // output
175 const int32_t* int_a, // a
176 const int32_t* torus_b, // b
177 uint8_t* tmp); // scratch space
178
179EXPORT uint64_t tnx32_small_single_product_tmp_bytes(const MOD_RNX* module);
180
181///////////////////////////////////////////////////////////////////
182// prepared gadget decompositions (optimized) //
183///////////////////////////////////////////////////////////////////
184
185// decompose from tnx32
186
187typedef struct tnx32_approxdecomp_gadget_t TNX32_APPROXDECOMP_GADGET;
188
189/** @brief new gadget: delete with delete_tnx32_approxdecomp_gadget */
190EXPORT TNX32_APPROXDECOMP_GADGET* new_tnx32_approxdecomp_gadget( //
191 const MOD_RNX* module, // N
192 uint64_t k, uint64_t ell // base 2^K and size
193);
194EXPORT void delete_tnx32_approxdecomp_gadget(const MOD_RNX* module);
195
196/** @brief sets res = gadget_decompose(a) */
197EXPORT void rnx_approxdecomp_from_tnx32( //
198 const MOD_RNX* module, // N
199 const TNX32_APPROXDECOMP_GADGET* gadget, // output base 2^K
200 double* res, uint64_t res_size, uint64_t res_sl, // res
201 const int32_t* a // a
202);
203
204// decompose from tnx32x2
205
206typedef struct tnx32x2_approxdecomp_gadget_t TNX32X2_APPROXDECOMP_GADGET;
207
208/** @brief new gadget: delete with delete_tnx32x2_approxdecomp_gadget */
209EXPORT TNX32X2_APPROXDECOMP_GADGET* new_tnx32x2_approxdecomp_gadget(const MOD_RNX* module, uint64_t ka, uint64_t ella,
210 uint64_t kb, uint64_t ellb);
211EXPORT void delete_tnx32x2_approxdecomp_gadget(const MOD_RNX* module);
212
213/** @brief sets res = gadget_decompose(a) */
214EXPORT void rnx_approxdecomp_from_tnx32x2( //
215 const MOD_RNX* module, // N
216 const TNX32X2_APPROXDECOMP_GADGET* gadget, // output base 2^K
217 double* res, uint64_t res_size, uint64_t res_sl, // res
218 const int32_t* a // a
219);
220
221// decompose from tnxdbl
222
223typedef struct tnxdbl_approxdecomp_gadget_t TNXDBL_APPROXDECOMP_GADGET;
224
225/** @brief new gadget: delete with delete_tnxdbl_approxdecomp_gadget */
226EXPORT TNXDBL_APPROXDECOMP_GADGET* new_tnxdbl_approxdecomp_gadget( //
227 const MOD_RNX* module, // N
228 uint64_t k, uint64_t ell // base 2^K and size
229);
230EXPORT void delete_tnxdbl_approxdecomp_gadget(TNXDBL_APPROXDECOMP_GADGET* gadget);
231
232/** @brief sets res = gadget_decompose(a) */
233EXPORT void rnx_approxdecomp_from_tnxdbl( //
234 const MOD_RNX* module, // N
235 const TNXDBL_APPROXDECOMP_GADGET* gadget, // output base 2^K
236 double* res, uint64_t res_size, uint64_t res_sl, // res
237 const double* a); // a
238
239///////////////////////////////////////////////////////////////////
240// prepared scalar-vector product (optimized) //
241///////////////////////////////////////////////////////////////////
242
243/** @brief opaque type that represents a polynomial of RnX prepared for a scalar-vector product */
244typedef struct rnx_svp_ppol_t RNX_SVP_PPOL;
245
246/** @brief number of bytes in a RNX_VMP_PMAT (for manual allocation) */
247EXPORT uint64_t bytes_of_rnx_svp_ppol(const MOD_RNX* module); // N
248
249/** @brief allocates a prepared vector (release with delete_rnx_svp_ppol) */
250EXPORT RNX_SVP_PPOL* new_rnx_svp_ppol(const MOD_RNX* module); // N
251
252/** @brief frees memory for a prepared vector */
253EXPORT void delete_rnx_svp_ppol(RNX_SVP_PPOL* res);
254
255/** @brief prepares a svp polynomial */
256EXPORT void rnx_svp_prepare(const MOD_RNX* module, // N
257 RNX_SVP_PPOL* ppol, // output
258 const double* pol // a
259);
260
261/** @brief apply a svp product, result = ppol * a, presented in DFT space */
262EXPORT void rnx_svp_apply( //
263 const MOD_RNX* module, // N
264 double* res, uint64_t res_size, uint64_t res_sl, // output
265 const RNX_SVP_PPOL* ppol, // prepared pol
266 const double* a, uint64_t a_size, uint64_t a_sl // a
267);
268
269///////////////////////////////////////////////////////////////////
270// prepared vector-matrix product (optimized) //
271///////////////////////////////////////////////////////////////////
272
273typedef struct rnx_vmp_pmat_t RNX_VMP_PMAT;
274
275/** @brief number of bytes in a RNX_VMP_PMAT (for manual allocation) */
276EXPORT uint64_t bytes_of_rnx_vmp_pmat(const MOD_RNX* module, // N
277 uint64_t nrows, uint64_t ncols); // dimensions
278
279/** @brief allocates a prepared matrix (release with delete_rnx_vmp_pmat) */
280EXPORT RNX_VMP_PMAT* new_rnx_vmp_pmat(const MOD_RNX* module, // N
281 uint64_t nrows, uint64_t ncols); // dimensions
282EXPORT void delete_rnx_vmp_pmat(RNX_VMP_PMAT* ptr);
283
284/** @brief prepares a vmp matrix (contiguous row-major version) */
285EXPORT void rnx_vmp_prepare_contiguous( //
286 const MOD_RNX* module, // N
287 RNX_VMP_PMAT* pmat, // output
288 const double* a, uint64_t nrows, uint64_t ncols, // a
289 uint8_t* tmp_space // scratch space
290);
291
292/** @brief prepares a vmp matrix (mat[row]+col*N points to the item) */
293EXPORT void rnx_vmp_prepare_dblptr( //
294 const MOD_RNX* module, // N
295 RNX_VMP_PMAT* pmat, // output
296 const double** a, uint64_t nrows, uint64_t ncols, // a
297 uint8_t* tmp_space // scratch space
298);
299
300/** @brief prepares the ith-row of a vmp matrix with nrows and ncols */
301EXPORT void rnx_vmp_prepare_row( //
302 const MOD_RNX* module, // N
303 RNX_VMP_PMAT* pmat, // output
304 const double* a, uint64_t row_i, uint64_t nrows, uint64_t ncols, // a
305 uint8_t* tmp_space // scratch space
306);
307
308/** @brief number of scratch bytes necessary to prepare a matrix */
309EXPORT uint64_t rnx_vmp_prepare_tmp_bytes(const MOD_RNX* module);
310
311/** @brief applies a vmp product res = a x pmat */
312EXPORT void rnx_vmp_apply_tmp_a( //
313 const MOD_RNX* module, // N
314 double* res, uint64_t res_size, uint64_t res_sl, // res
315 double* tmpa, uint64_t a_size, uint64_t a_sl, // a (will be overwritten)
316 const RNX_VMP_PMAT* pmat, uint64_t nrows, uint64_t ncols, // prep matrix
317 uint8_t* tmp_space // scratch space
318);
319
320EXPORT uint64_t rnx_vmp_apply_tmp_a_tmp_bytes( //
321 const MOD_RNX* module, // N
322 uint64_t res_size, // res size
323 uint64_t a_size, // a size
324 uint64_t nrows, uint64_t ncols // prep matrix dims
325);
326
327/** @brief minimal size of the tmp_space */
328EXPORT void rnx_vmp_apply_dft_to_dft( //
329 const MOD_RNX* module, // N
330 double* res, uint64_t res_size, uint64_t res_sl, // res
331 const double* a_dft, uint64_t a_size, uint64_t a_sl, // a
332 const RNX_VMP_PMAT* pmat, uint64_t nrows, uint64_t ncols, // prep matrix
333 uint8_t* tmp_space // scratch space (a_size*sizeof(reim4) bytes)
334);
335
336/** @brief minimal size of the tmp_space */
337EXPORT uint64_t rnx_vmp_apply_dft_to_dft_tmp_bytes( //
338 const MOD_RNX* module, // N
339 uint64_t res_size, // res
340 uint64_t a_size, // a
341 uint64_t nrows, uint64_t ncols // prep matrix
342);
343
344/** @brief sets res = DFT(a) */
345EXPORT void vec_rnx_dft(const MOD_RNX* module, // N
346 double* res, uint64_t res_size, uint64_t res_sl, // res
347 const double* a, uint64_t a_size, uint64_t a_sl // a
348);
349
350/** @brief sets res = iDFT(a_dft) -- idft is not normalized */
351EXPORT void vec_rnx_idft(const MOD_RNX* module, // N
352 double* res, uint64_t res_size, uint64_t res_sl, // res
353 const double* a_dft, uint64_t a_size, uint64_t a_sl // a
354);
355
356#endif // SPQLIOS_VEC_RNX_ARITHMETIC_H