poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/vec_rnx_vmp_avx.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/vec_rnx_vmp_avx.cRaw
1#include <assert.h>
2#include <immintrin.h>
3#include <string.h>
4
5#include "../coeffs/coeffs_arithmetic.h"
6#include "../reim/reim_fft.h"
7#include "../reim4/reim4_arithmetic.h"
8#include "vec_rnx_arithmetic_private.h"
9
10/** @brief prepares a vmp matrix (contiguous row-major version) */
11EXPORT void fft64_rnx_vmp_prepare_contiguous_avx( //
12 const MOD_RNX* module, // N
13 RNX_VMP_PMAT* pmat, // output
14 const double* mat, uint64_t nrows, uint64_t ncols, // a
15 uint8_t* tmp_space // scratch space
16) {
17 // there is an edge case if nn < 8
18 const uint64_t nn = module->n;
19 const uint64_t m = module->m;
20
21 double* const dtmp = (double*)tmp_space;
22 double* const output_mat = (double*)pmat;
23 double* start_addr = (double*)pmat;
24 uint64_t offset = nrows * ncols * 8;
25
26 if (nn >= 8) {
27 for (uint64_t row_i = 0; row_i < nrows; row_i++) {
28 for (uint64_t col_i = 0; col_i < ncols; col_i++) {
29 rnx_divide_by_m_avx(nn, m, dtmp, mat + (row_i * ncols + col_i) * nn);
30 reim_fft(module->precomp.fft64.p_fft, dtmp);
31
32 if (col_i == (ncols - 1) && (ncols % 2 == 1)) {
33 // special case: last column out of an odd column number
34 start_addr = output_mat + col_i * nrows * 8 // col == ncols-1
35 + row_i * 8;
36 } else {
37 // general case: columns go by pair
38 start_addr = output_mat + (col_i / 2) * (2 * nrows) * 8 // second: col pair index
39 + row_i * 2 * 8 // third: row index
40 + (col_i % 2) * 8;
41 }
42
43 for (uint64_t blk_i = 0; blk_i < m / 4; blk_i++) {
44 // extract blk from tmp and save it
45 reim4_extract_1blk_from_reim_avx(m, blk_i, start_addr + blk_i * offset, dtmp);
46 }
47 }
48 }
49 } else {
50 for (uint64_t row_i = 0; row_i < nrows; row_i++) {
51 for (uint64_t col_i = 0; col_i < ncols; col_i++) {
52 double* res = output_mat + (col_i * nrows + row_i) * nn;
53 rnx_divide_by_m_avx(nn, m, res, mat + (row_i * ncols + col_i) * nn);
54 reim_fft(module->precomp.fft64.p_fft, res);
55 }
56 }
57 }
58}
59
60/** @brief prepares a vmp matrix (mat[row]+col*N points to the item) */
61EXPORT void fft64_rnx_vmp_prepare_dblptr_avx( //
62 const MOD_RNX* module, // N
63 RNX_VMP_PMAT* pmat, // output
64 const double** mat, uint64_t nrows, uint64_t ncols, // a
65 uint8_t* tmp_space // scratch space
66) {
67 for (uint64_t row_i = 0; row_i < nrows; row_i++) {
68 fft64_rnx_vmp_prepare_row_avx(module, pmat, mat[row_i], row_i, nrows, ncols, tmp_space);
69 }
70}
71
72/** @brief prepares the ith-row of a vmp matrix with nrows and ncols */
73EXPORT void fft64_rnx_vmp_prepare_row_avx( //
74 const MOD_RNX* module, // N
75 RNX_VMP_PMAT* pmat, // output
76 const double* row, uint64_t row_i, uint64_t nrows, uint64_t ncols, // a
77 uint8_t* tmp_space // scratch space
78) {
79 // there is an edge case if nn < 8
80 const uint64_t nn = module->n;
81 const uint64_t m = module->m;
82
83 double* const dtmp = (double*)tmp_space;
84 double* const output_mat = (double*)pmat;
85 double* start_addr = (double*)pmat;
86 uint64_t offset = nrows * ncols * 8;
87
88 if (nn >= 8) {
89 for (uint64_t col_i = 0; col_i < ncols; col_i++) {
90 rnx_divide_by_m_avx(nn, m, dtmp, row + col_i * nn);
91 reim_fft(module->precomp.fft64.p_fft, dtmp);
92
93 if (col_i == (ncols - 1) && (ncols % 2 == 1)) {
94 // special case: last column out of an odd column number
95 start_addr = output_mat + col_i * nrows * 8 // col == ncols-1
96 + row_i * 8;
97 } else {
98 // general case: columns go by pair
99 start_addr = output_mat + (col_i / 2) * (2 * nrows) * 8 // second: col pair index
100 + row_i * 2 * 8 // third: row index
101 + (col_i % 2) * 8;
102 }
103
104 for (uint64_t blk_i = 0; blk_i < m / 4; blk_i++) {
105 // extract blk from tmp and save it
106 reim4_extract_1blk_from_reim_avx(m, blk_i, start_addr + blk_i * offset, dtmp);
107 }
108 }
109 } else {
110 for (uint64_t col_i = 0; col_i < ncols; col_i++) {
111 double* res = output_mat + (col_i * nrows + row_i) * nn;
112 rnx_divide_by_m_avx(nn, m, res, row + col_i * nn);
113 reim_fft(module->precomp.fft64.p_fft, res);
114 }
115 }
116}
117
118/** @brief minimal size of the tmp_space */
119EXPORT void fft64_rnx_vmp_apply_dft_to_dft_avx( //
120 const MOD_RNX* module, // N
121 double* res, uint64_t res_size, uint64_t res_sl, // res
122 const double* a_dft, uint64_t a_size, uint64_t a_sl, // a
123 const RNX_VMP_PMAT* pmat, uint64_t nrows, uint64_t ncols, // prep matrix
124 uint8_t* tmp_space // scratch space (a_size*sizeof(reim4) bytes)
125) {
126 const uint64_t m = module->m;
127 const uint64_t nn = module->n;
128
129 double* mat2cols_output = (double*)tmp_space; // 128 bytes
130 double* extracted_blk = (double*)tmp_space + 16; // 64*min(nrows,a_size) bytes
131
132 double* mat_input = (double*)pmat;
133
134 const uint64_t row_max = nrows < a_size ? nrows : a_size;
135 const uint64_t col_max = ncols < res_size ? ncols : res_size;
136
137 if (row_max > 0 && col_max > 0) {
138 if (nn >= 8) {
139 // let's do some prefetching of the GSW key, since on some cpus,
140 // it helps
141 const uint64_t ms4 = m >> 2; // m/4
142 const uint64_t gsw_iter_doubles = 8 * nrows * ncols;
143 const uint64_t pref_doubles = 1200;
144 const double* gsw_pref_ptr = mat_input;
145 const double* const gsw_ptr_end = mat_input + ms4 * gsw_iter_doubles;
146 const double* gsw_pref_ptr_target = mat_input + pref_doubles;
147 for (; gsw_pref_ptr < gsw_pref_ptr_target; gsw_pref_ptr += 8) {
148 __builtin_prefetch(gsw_pref_ptr, 0, _MM_HINT_T0);
149 }
150 const double* mat_blk_start;
151 uint64_t blk_i;
152 for (blk_i = 0, mat_blk_start = mat_input; blk_i < ms4; blk_i++, mat_blk_start += gsw_iter_doubles) {
153 // prefetch the next iteration
154 if (gsw_pref_ptr_target < gsw_ptr_end) {
155 gsw_pref_ptr_target += gsw_iter_doubles;
156 if (gsw_pref_ptr_target > gsw_ptr_end) gsw_pref_ptr_target = gsw_ptr_end;
157 for (; gsw_pref_ptr < gsw_pref_ptr_target; gsw_pref_ptr += 8) {
158 __builtin_prefetch(gsw_pref_ptr, 0, _MM_HINT_T0);
159 }
160 }
161 reim4_extract_1blk_from_contiguous_reim_sl_avx(m, a_sl, row_max, blk_i, extracted_blk, a_dft);
162 // apply mat2cols
163 for (uint64_t col_i = 0; col_i < col_max - 1; col_i += 2) {
164 uint64_t col_offset = col_i * (8 * nrows);
165 reim4_vec_mat2cols_product_avx2(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset);
166
167 reim4_save_1blk_to_reim_avx(m, blk_i, res + col_i * res_sl, mat2cols_output);
168 reim4_save_1blk_to_reim_avx(m, blk_i, res + (col_i + 1) * res_sl, mat2cols_output + 8);
169 }
170
171 // check if col_max is odd, then special case
172 if (col_max % 2 == 1) {
173 uint64_t last_col = col_max - 1;
174 uint64_t col_offset = last_col * (8 * nrows);
175
176 // the last column is alone in the pmat: vec_mat1col
177 if (ncols == col_max) {
178 reim4_vec_mat1col_product_avx2(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset);
179 } else {
180 // the last column is part of a colpair in the pmat: vec_mat2cols and ignore the second position
181 reim4_vec_mat2cols_product_avx2(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset);
182 }
183 reim4_save_1blk_to_reim_avx(m, blk_i, res + last_col * res_sl, mat2cols_output);
184 }
185 }
186 } else {
187 const double* in;
188 uint64_t in_sl;
189 if (res == a_dft) {
190 // it is in place: copy the input vector
191 in = (double*)tmp_space;
192 in_sl = nn;
193 // vec_rnx_copy(module, (double*)tmp_space, row_max, nn, a_dft, row_max, a_sl);
194 for (uint64_t row_i = 0; row_i < row_max; row_i++) {
195 memcpy((double*)tmp_space + row_i * nn, a_dft + row_i * a_sl, nn * sizeof(double));
196 }
197 } else {
198 // it is out of place: do the product directly
199 in = a_dft;
200 in_sl = a_sl;
201 }
202 for (uint64_t col_i = 0; col_i < col_max; col_i++) {
203 double* pmat_col = mat_input + col_i * nrows * nn;
204 {
205 reim_fftvec_mul(module->precomp.fft64.p_fftvec_mul, //
206 res + col_i * res_sl, //
207 in, //
208 pmat_col);
209 }
210 for (uint64_t row_i = 1; row_i < row_max; row_i++) {
211 reim_fftvec_addmul(module->precomp.fft64.p_fftvec_addmul, //
212 res + col_i * res_sl, //
213 in + row_i * in_sl, //
214 pmat_col + row_i * nn);
215 }
216 }
217 }
218 }
219 // zero out remaining bytes (if any)
220 for (uint64_t i = col_max; i < res_size; ++i) {
221 memset(res + i * res_sl, 0, nn * sizeof(double));
222 }
223}
224
225/** @brief applies a vmp product res = a x pmat */
226EXPORT void fft64_rnx_vmp_apply_tmp_a_avx( //
227 const MOD_RNX* module, // N
228 double* res, uint64_t res_size, uint64_t res_sl, // res (addr must be != a)
229 double* tmpa, uint64_t a_size, uint64_t a_sl, // a (will be overwritten)
230 const RNX_VMP_PMAT* pmat, uint64_t nrows, uint64_t ncols, // prep matrix
231 uint8_t* tmp_space // scratch space
232) {
233 const uint64_t nn = module->n;
234 const uint64_t rows = nrows < a_size ? nrows : a_size;
235 const uint64_t cols = ncols < res_size ? ncols : res_size;
236
237 // fft is done in place on the input (tmpa is destroyed)
238 for (uint64_t i = 0; i < rows; ++i) {
239 reim_fft(module->precomp.fft64.p_fft, tmpa + i * a_sl);
240 }
241 fft64_rnx_vmp_apply_dft_to_dft_avx(module, //
242 res, cols, res_sl, //
243 tmpa, rows, a_sl, //
244 pmat, nrows, ncols, //
245 tmp_space);
246 // ifft is done in place on the output
247 for (uint64_t i = 0; i < cols; ++i) {
248 reim_ifft(module->precomp.fft64.p_ifft, res + i * res_sl);
249 }
250 // zero out the remaining positions
251 for (uint64_t i = cols; i < res_size; ++i) {
252 memset(res + i * res_sl, 0, nn * sizeof(double));
253 }
254}