poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/vec_rnx_vmp_ref.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/vec_rnx_vmp_ref.cRaw
1#include <assert.h>
2#include <string.h>
3
4#include "../coeffs/coeffs_arithmetic.h"
5#include "../reim/reim_fft.h"
6#include "../reim4/reim4_arithmetic.h"
7#include "vec_rnx_arithmetic_private.h"
8
9/** @brief number of bytes in a RNX_VMP_PMAT (for manual allocation) */
10EXPORT uint64_t fft64_bytes_of_rnx_vmp_pmat(const MOD_RNX* module, // N
11 uint64_t nrows, uint64_t ncols) { // dimensions
12 return nrows * ncols * module->n * sizeof(double);
13}
14
15/** @brief prepares a vmp matrix (contiguous row-major version) */
16EXPORT void fft64_rnx_vmp_prepare_contiguous_ref( //
17 const MOD_RNX* module, // N
18 RNX_VMP_PMAT* pmat, // output
19 const double* mat, uint64_t nrows, uint64_t ncols, // a
20 uint8_t* tmp_space // scratch space
21) {
22 // there is an edge case if nn < 8
23 const uint64_t nn = module->n;
24 const uint64_t m = module->m;
25
26 double* const dtmp = (double*)tmp_space;
27 double* const output_mat = (double*)pmat;
28 double* start_addr = (double*)pmat;
29 uint64_t offset = nrows * ncols * 8;
30
31 if (nn >= 8) {
32 for (uint64_t row_i = 0; row_i < nrows; row_i++) {
33 for (uint64_t col_i = 0; col_i < ncols; col_i++) {
34 rnx_divide_by_m_ref(nn, m, dtmp, mat + (row_i * ncols + col_i) * nn);
35 reim_fft(module->precomp.fft64.p_fft, dtmp);
36
37 if (col_i == (ncols - 1) && (ncols % 2 == 1)) {
38 // special case: last column out of an odd column number
39 start_addr = output_mat + col_i * nrows * 8 // col == ncols-1
40 + row_i * 8;
41 } else {
42 // general case: columns go by pair
43 start_addr = output_mat + (col_i / 2) * (2 * nrows) * 8 // second: col pair index
44 + row_i * 2 * 8 // third: row index
45 + (col_i % 2) * 8;
46 }
47
48 for (uint64_t blk_i = 0; blk_i < m / 4; blk_i++) {
49 // extract blk from tmp and save it
50 reim4_extract_1blk_from_reim_ref(m, blk_i, start_addr + blk_i * offset, dtmp);
51 }
52 }
53 }
54 } else {
55 for (uint64_t row_i = 0; row_i < nrows; row_i++) {
56 for (uint64_t col_i = 0; col_i < ncols; col_i++) {
57 double* res = output_mat + (col_i * nrows + row_i) * nn;
58 rnx_divide_by_m_ref(nn, m, res, mat + (row_i * ncols + col_i) * nn);
59 reim_fft(module->precomp.fft64.p_fft, res);
60 }
61 }
62 }
63}
64
65/** @brief prepares a vmp matrix (mat[row]+col*N points to the item) */
66EXPORT void fft64_rnx_vmp_prepare_dblptr_ref( //
67 const MOD_RNX* module, // N
68 RNX_VMP_PMAT* pmat, // output
69 const double** mat, uint64_t nrows, uint64_t ncols, // a
70 uint8_t* tmp_space // scratch space
71) {
72 for (uint64_t row_i = 0; row_i < nrows; row_i++) {
73 fft64_rnx_vmp_prepare_row_ref(module, pmat, mat[row_i], row_i, nrows, ncols, tmp_space);
74 }
75}
76
77/** @brief prepares the ith-row of a vmp matrix with nrows and ncols */
78EXPORT void fft64_rnx_vmp_prepare_row_ref( //
79 const MOD_RNX* module, // N
80 RNX_VMP_PMAT* pmat, // output
81 const double* row, uint64_t row_i, uint64_t nrows, uint64_t ncols, // a
82 uint8_t* tmp_space // scratch space
83) {
84 // there is an edge case if nn < 8
85 const uint64_t nn = module->n;
86 const uint64_t m = module->m;
87
88 double* const dtmp = (double*)tmp_space;
89 double* const output_mat = (double*)pmat;
90 double* start_addr = (double*)pmat;
91 uint64_t offset = nrows * ncols * 8;
92
93 if (nn >= 8) {
94 for (uint64_t col_i = 0; col_i < ncols; col_i++) {
95 rnx_divide_by_m_ref(nn, m, dtmp, row + col_i * nn);
96 reim_fft(module->precomp.fft64.p_fft, dtmp);
97
98 if (col_i == (ncols - 1) && (ncols % 2 == 1)) {
99 // special case: last column out of an odd column number
100 start_addr = output_mat + col_i * nrows * 8 // col == ncols-1
101 + row_i * 8;
102 } else {
103 // general case: columns go by pair
104 start_addr = output_mat + (col_i / 2) * (2 * nrows) * 8 // second: col pair index
105 + row_i * 2 * 8 // third: row index
106 + (col_i % 2) * 8;
107 }
108
109 for (uint64_t blk_i = 0; blk_i < m / 4; blk_i++) {
110 // extract blk from tmp and save it
111 reim4_extract_1blk_from_reim_ref(m, blk_i, start_addr + blk_i * offset, dtmp);
112 }
113 }
114 } else {
115 for (uint64_t col_i = 0; col_i < ncols; col_i++) {
116 double* res = output_mat + (col_i * nrows + row_i) * nn;
117 rnx_divide_by_m_ref(nn, m, res, row + col_i * nn);
118 reim_fft(module->precomp.fft64.p_fft, res);
119 }
120 }
121}
122
123/** @brief number of scratch bytes necessary to prepare a matrix */
124EXPORT uint64_t fft64_rnx_vmp_prepare_tmp_bytes_ref(const MOD_RNX* module) {
125 const uint64_t nn = module->n;
126 return nn * sizeof(int64_t);
127}
128
129/** @brief minimal size of the tmp_space */
130EXPORT void fft64_rnx_vmp_apply_dft_to_dft_ref( //
131 const MOD_RNX* module, // N
132 double* res, uint64_t res_size, uint64_t res_sl, // res
133 const double* a_dft, uint64_t a_size, uint64_t a_sl, // a
134 const RNX_VMP_PMAT* pmat, uint64_t nrows, uint64_t ncols, // prep matrix
135 uint8_t* tmp_space // scratch space (a_size*sizeof(reim4) bytes)
136) {
137 const uint64_t m = module->m;
138 const uint64_t nn = module->n;
139
140 double* mat2cols_output = (double*)tmp_space; // 128 bytes
141 double* extracted_blk = (double*)tmp_space + 16; // 64*min(nrows,a_size) bytes
142
143 double* mat_input = (double*)pmat;
144
145 const uint64_t row_max = nrows < a_size ? nrows : a_size;
146 const uint64_t col_max = ncols < res_size ? ncols : res_size;
147
148 if (row_max > 0 && col_max > 0) {
149 if (nn >= 8) {
150 for (uint64_t blk_i = 0; blk_i < m / 4; blk_i++) {
151 double* mat_blk_start = mat_input + blk_i * (8 * nrows * ncols);
152
153 reim4_extract_1blk_from_contiguous_reim_sl_ref(m, a_sl, row_max, blk_i, extracted_blk, a_dft);
154 // apply mat2cols
155 for (uint64_t col_i = 0; col_i < col_max - 1; col_i += 2) {
156 uint64_t col_offset = col_i * (8 * nrows);
157 reim4_vec_mat2cols_product_ref(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset);
158
159 reim4_save_1blk_to_reim_ref(m, blk_i, res + col_i * res_sl, mat2cols_output);
160 reim4_save_1blk_to_reim_ref(m, blk_i, res + (col_i + 1) * res_sl, mat2cols_output + 8);
161 }
162
163 // check if col_max is odd, then special case
164 if (col_max % 2 == 1) {
165 uint64_t last_col = col_max - 1;
166 uint64_t col_offset = last_col * (8 * nrows);
167
168 // the last column is alone in the pmat: vec_mat1col
169 if (ncols == col_max) {
170 reim4_vec_mat1col_product_ref(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset);
171 } else {
172 // the last column is part of a colpair in the pmat: vec_mat2cols and ignore the second position
173 reim4_vec_mat2cols_product_ref(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset);
174 }
175 reim4_save_1blk_to_reim_ref(m, blk_i, res + last_col * res_sl, mat2cols_output);
176 }
177 }
178 } else {
179 const double* in;
180 uint64_t in_sl;
181 if (res == a_dft) {
182 // it is in place: copy the input vector
183 in = (double*)tmp_space;
184 in_sl = nn;
185 // vec_rnx_copy(module, (double*)tmp_space, row_max, nn, a_dft, row_max, a_sl);
186 for (uint64_t row_i = 0; row_i < row_max; row_i++) {
187 memcpy((double*)tmp_space + row_i * nn, a_dft + row_i * a_sl, nn * sizeof(double));
188 }
189 } else {
190 // it is out of place: do the product directly
191 in = a_dft;
192 in_sl = a_sl;
193 }
194 for (uint64_t col_i = 0; col_i < col_max; col_i++) {
195 double* pmat_col = mat_input + col_i * nrows * nn;
196 {
197 reim_fftvec_mul(module->precomp.fft64.p_fftvec_mul, //
198 res + col_i * res_sl, //
199 in, //
200 pmat_col);
201 }
202 for (uint64_t row_i = 1; row_i < row_max; row_i++) {
203 reim_fftvec_addmul(module->precomp.fft64.p_fftvec_addmul, //
204 res + col_i * res_sl, //
205 in + row_i * in_sl, //
206 pmat_col + row_i * nn);
207 }
208 }
209 }
210 }
211 // zero out remaining bytes (if any)
212 for (uint64_t i = col_max; i < res_size; ++i) {
213 memset(res + i * res_sl, 0, nn * sizeof(double));
214 }
215}
216
217/** @brief applies a vmp product res = a x pmat */
218EXPORT void fft64_rnx_vmp_apply_tmp_a_ref( //
219 const MOD_RNX* module, // N
220 double* res, uint64_t res_size, uint64_t res_sl, // res (addr must be != a)
221 double* tmpa, uint64_t a_size, uint64_t a_sl, // a (will be overwritten)
222 const RNX_VMP_PMAT* pmat, uint64_t nrows, uint64_t ncols, // prep matrix
223 uint8_t* tmp_space // scratch space
224) {
225 const uint64_t nn = module->n;
226 const uint64_t rows = nrows < a_size ? nrows : a_size;
227 const uint64_t cols = ncols < res_size ? ncols : res_size;
228
229 // fft is done in place on the input (tmpa is destroyed)
230 for (uint64_t i = 0; i < rows; ++i) {
231 reim_fft(module->precomp.fft64.p_fft, tmpa + i * a_sl);
232 }
233 fft64_rnx_vmp_apply_dft_to_dft_ref(module, //
234 res, cols, res_sl, //
235 tmpa, rows, a_sl, //
236 pmat, nrows, ncols, //
237 tmp_space);
238 // ifft is done in place on the output
239 for (uint64_t i = 0; i < cols; ++i) {
240 reim_ifft(module->precomp.fft64.p_ifft, res + i * res_sl);
241 }
242 // zero out the remaining positions
243 for (uint64_t i = cols; i < res_size; ++i) {
244 memset(res + i * res_sl, 0, nn * sizeof(double));
245 }
246}
247
248/** @brief minimal size of the tmp_space */
249EXPORT uint64_t fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref( //
250 const MOD_RNX* module, // N
251 uint64_t res_size, // res
252 uint64_t a_size, // a
253 uint64_t nrows, uint64_t ncols // prep matrix
254) {
255 const uint64_t row_max = nrows < a_size ? nrows : a_size;
256
257 return (128) + (64 * row_max);
258}
259
260#ifdef __APPLE__
261EXPORT uint64_t fft64_rnx_vmp_apply_tmp_a_tmp_bytes_ref( //
262 const MOD_RNX* module, // N
263 uint64_t res_size, // res
264 uint64_t a_size, // a
265 uint64_t nrows, uint64_t ncols // prep matrix
266) {
267 return fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref(module, res_size, a_size, nrows, ncols);
268}
269#else
270EXPORT uint64_t fft64_rnx_vmp_apply_tmp_a_tmp_bytes_ref( //
271 const MOD_RNX* module, // N
272 uint64_t res_size, // res
273 uint64_t a_size, // a
274 uint64_t nrows, uint64_t ncols // prep matrix
275 ) __attribute((alias("fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref")));
276#endif
277// avx aliases that need to be defined in the same .c file
278
279/** @brief number of scratch bytes necessary to prepare a matrix */
280#ifdef __APPLE__
281#pragma weak fft64_rnx_vmp_prepare_tmp_bytes_avx = fft64_rnx_vmp_prepare_tmp_bytes_ref
282#else
283EXPORT uint64_t fft64_rnx_vmp_prepare_tmp_bytes_avx(const MOD_RNX* module)
284 __attribute((alias("fft64_rnx_vmp_prepare_tmp_bytes_ref")));
285#endif
286
287/** @brief minimal size of the tmp_space */
288#ifdef __APPLE__
289#pragma weak fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_avx = fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref
290#else
291EXPORT uint64_t fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_avx( //
292 const MOD_RNX* module, // N
293 uint64_t res_size, // res
294 uint64_t a_size, // a
295 uint64_t nrows, uint64_t ncols // prep matrix
296 ) __attribute((alias("fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref")));
297#endif
298
299#ifdef __APPLE__
300#pragma weak fft64_rnx_vmp_apply_tmp_a_tmp_bytes_avx = fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref
301#else
302EXPORT uint64_t fft64_rnx_vmp_apply_tmp_a_tmp_bytes_avx( //
303 const MOD_RNX* module, // N
304 uint64_t res_size, // res
305 uint64_t a_size, // a
306 uint64_t nrows, uint64_t ncols // prep matrix
307 ) __attribute((alias("fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref")));
308#endif
309// wrappers