poulpy/spqlios/lib/spqlios/arithmetic/vec_rnx_vmp_ref.c
spqlios/lib/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 number of scratch bytes necessary to prepare a matrix */
66EXPORT uint64_t fft64_rnx_vmp_prepare_contiguous_tmp_bytes_ref(const MOD_RNX* module) {
67 const uint64_t nn = module->n;
68 return nn * sizeof(int64_t);
69}
70
71/** @brief minimal size of the tmp_space */
72EXPORT void fft64_rnx_vmp_apply_dft_to_dft_ref( //
73 const MOD_RNX* module, // N
74 double* res, uint64_t res_size, uint64_t res_sl, // res
75 const double* a_dft, uint64_t a_size, uint64_t a_sl, // a
76 const RNX_VMP_PMAT* pmat, uint64_t nrows, uint64_t ncols, // prep matrix
77 uint8_t* tmp_space // scratch space (a_size*sizeof(reim4) bytes)
78) {
79 const uint64_t m = module->m;
80 const uint64_t nn = module->n;
81
82 double* mat2cols_output = (double*)tmp_space; // 128 bytes
83 double* extracted_blk = (double*)tmp_space + 16; // 64*min(nrows,a_size) bytes
84
85 double* mat_input = (double*)pmat;
86
87 const uint64_t row_max = nrows < a_size ? nrows : a_size;
88 const uint64_t col_max = ncols < res_size ? ncols : res_size;
89
90 if (row_max > 0 && col_max > 0) {
91 if (nn >= 8) {
92 for (uint64_t blk_i = 0; blk_i < m / 4; blk_i++) {
93 double* mat_blk_start = mat_input + blk_i * (8 * nrows * ncols);
94
95 reim4_extract_1blk_from_contiguous_reim_sl_ref(m, a_sl, row_max, blk_i, extracted_blk, a_dft);
96 // apply mat2cols
97 for (uint64_t col_i = 0; col_i < col_max - 1; col_i += 2) {
98 uint64_t col_offset = col_i * (8 * nrows);
99 reim4_vec_mat2cols_product_ref(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset);
100
101 reim4_save_1blk_to_reim_ref(m, blk_i, res + col_i * res_sl, mat2cols_output);
102 reim4_save_1blk_to_reim_ref(m, blk_i, res + (col_i + 1) * res_sl, mat2cols_output + 8);
103 }
104
105 // check if col_max is odd, then special case
106 if (col_max % 2 == 1) {
107 uint64_t last_col = col_max - 1;
108 uint64_t col_offset = last_col * (8 * nrows);
109
110 // the last column is alone in the pmat: vec_mat1col
111 if (ncols == col_max) {
112 reim4_vec_mat1col_product_ref(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset);
113 } else {
114 // the last column is part of a colpair in the pmat: vec_mat2cols and ignore the second position
115 reim4_vec_mat2cols_product_ref(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset);
116 }
117 reim4_save_1blk_to_reim_ref(m, blk_i, res + last_col * res_sl, mat2cols_output);
118 }
119 }
120 } else {
121 const double* in;
122 uint64_t in_sl;
123 if (res == a_dft) {
124 // it is in place: copy the input vector
125 in = (double*)tmp_space;
126 in_sl = nn;
127 // vec_rnx_copy(module, (double*)tmp_space, row_max, nn, a_dft, row_max, a_sl);
128 for (uint64_t row_i = 0; row_i < row_max; row_i++) {
129 memcpy((double*)tmp_space + row_i * nn, a_dft + row_i * a_sl, nn * sizeof(double));
130 }
131 } else {
132 // it is out of place: do the product directly
133 in = a_dft;
134 in_sl = a_sl;
135 }
136 for (uint64_t col_i = 0; col_i < col_max; col_i++) {
137 double* pmat_col = mat_input + col_i * nrows * nn;
138 {
139 reim_fftvec_mul(module->precomp.fft64.p_fftvec_mul, //
140 res + col_i * res_sl, //
141 in, //
142 pmat_col);
143 }
144 for (uint64_t row_i = 1; row_i < row_max; row_i++) {
145 reim_fftvec_addmul(module->precomp.fft64.p_fftvec_addmul, //
146 res + col_i * res_sl, //
147 in + row_i * in_sl, //
148 pmat_col + row_i * nn);
149 }
150 }
151 }
152 }
153 // zero out remaining bytes (if any)
154 for (uint64_t i = col_max; i < res_size; ++i) {
155 memset(res + i * res_sl, 0, nn * sizeof(double));
156 }
157}
158
159/** @brief applies a vmp product res = a x pmat */
160EXPORT void fft64_rnx_vmp_apply_tmp_a_ref( //
161 const MOD_RNX* module, // N
162 double* res, uint64_t res_size, uint64_t res_sl, // res (addr must be != a)
163 double* tmpa, uint64_t a_size, uint64_t a_sl, // a (will be overwritten)
164 const RNX_VMP_PMAT* pmat, uint64_t nrows, uint64_t ncols, // prep matrix
165 uint8_t* tmp_space // scratch space
166) {
167 const uint64_t nn = module->n;
168 const uint64_t rows = nrows < a_size ? nrows : a_size;
169 const uint64_t cols = ncols < res_size ? ncols : res_size;
170
171 // fft is done in place on the input (tmpa is destroyed)
172 for (uint64_t i = 0; i < rows; ++i) {
173 reim_fft(module->precomp.fft64.p_fft, tmpa + i * a_sl);
174 }
175 fft64_rnx_vmp_apply_dft_to_dft_ref(module, //
176 res, cols, res_sl, //
177 tmpa, rows, a_sl, //
178 pmat, nrows, ncols, //
179 tmp_space);
180 // ifft is done in place on the output
181 for (uint64_t i = 0; i < cols; ++i) {
182 reim_ifft(module->precomp.fft64.p_ifft, res + i * res_sl);
183 }
184 // zero out the remaining positions
185 for (uint64_t i = cols; i < res_size; ++i) {
186 memset(res + i * res_sl, 0, nn * sizeof(double));
187 }
188}
189
190/** @brief minimal size of the tmp_space */
191EXPORT uint64_t fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref( //
192 const MOD_RNX* module, // N
193 uint64_t res_size, // res
194 uint64_t a_size, // a
195 uint64_t nrows, uint64_t ncols // prep matrix
196) {
197 const uint64_t row_max = nrows < a_size ? nrows : a_size;
198
199 return (128) + (64 * row_max);
200}
201
202#ifdef __APPLE__
203EXPORT uint64_t fft64_rnx_vmp_apply_tmp_a_tmp_bytes_ref( //
204 const MOD_RNX* module, // N
205 uint64_t res_size, // res
206 uint64_t a_size, // a
207 uint64_t nrows, uint64_t ncols // prep matrix
208) {
209 return fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref(module, res_size, a_size, nrows, ncols);
210}
211#else
212EXPORT uint64_t fft64_rnx_vmp_apply_tmp_a_tmp_bytes_ref( //
213 const MOD_RNX* module, // N
214 uint64_t res_size, // res
215 uint64_t a_size, // a
216 uint64_t nrows, uint64_t ncols // prep matrix
217 ) __attribute((alias("fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref")));
218#endif
219// avx aliases that need to be defined in the same .c file
220
221/** @brief number of scratch bytes necessary to prepare a matrix */
222#ifdef __APPLE__
223#pragma weak fft64_rnx_vmp_prepare_contiguous_tmp_bytes_avx = fft64_rnx_vmp_prepare_contiguous_tmp_bytes_ref
224#else
225EXPORT uint64_t fft64_rnx_vmp_prepare_contiguous_tmp_bytes_avx(const MOD_RNX* module)
226 __attribute((alias("fft64_rnx_vmp_prepare_contiguous_tmp_bytes_ref")));
227#endif
228
229/** @brief minimal size of the tmp_space */
230#ifdef __APPLE__
231#pragma weak fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_avx = fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref
232#else
233EXPORT uint64_t fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_avx( //
234 const MOD_RNX* module, // N
235 uint64_t res_size, // res
236 uint64_t a_size, // a
237 uint64_t nrows, uint64_t ncols // prep matrix
238 ) __attribute((alias("fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref")));
239#endif
240
241#ifdef __APPLE__
242#pragma weak fft64_rnx_vmp_apply_tmp_a_tmp_bytes_avx = fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref
243#else
244EXPORT uint64_t fft64_rnx_vmp_apply_tmp_a_tmp_bytes_avx( //
245 const MOD_RNX* module, // N
246 uint64_t res_size, // res
247 uint64_t a_size, // a
248 uint64_t nrows, uint64_t ncols // prep matrix
249 ) __attribute((alias("fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref")));
250#endif
251// wrappers