poulpy/spqlios/lib/spqlios/arithmetic/vec_rnx_vmp_avx.c
spqlios/lib/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 minimal size of the tmp_space */
61EXPORT void fft64_rnx_vmp_apply_dft_to_dft_avx( //
62 const MOD_RNX* module, // N
63 double* res, uint64_t res_size, uint64_t res_sl, // res
64 const double* a_dft, uint64_t a_size, uint64_t a_sl, // a
65 const RNX_VMP_PMAT* pmat, uint64_t nrows, uint64_t ncols, // prep matrix
66 uint8_t* tmp_space // scratch space (a_size*sizeof(reim4) bytes)
67) {
68 const uint64_t m = module->m;
69 const uint64_t nn = module->n;
70
71 double* mat2cols_output = (double*)tmp_space; // 128 bytes
72 double* extracted_blk = (double*)tmp_space + 16; // 64*min(nrows,a_size) bytes
73
74 double* mat_input = (double*)pmat;
75
76 const uint64_t row_max = nrows < a_size ? nrows : a_size;
77 const uint64_t col_max = ncols < res_size ? ncols : res_size;
78
79 if (row_max > 0 && col_max > 0) {
80 if (nn >= 8) {
81 // let's do some prefetching of the GSW key, since on some cpus,
82 // it helps
83 const uint64_t ms4 = m >> 2; // m/4
84 const uint64_t gsw_iter_doubles = 8 * nrows * ncols;
85 const uint64_t pref_doubles = 1200;
86 const double* gsw_pref_ptr = mat_input;
87 const double* const gsw_ptr_end = mat_input + ms4 * gsw_iter_doubles;
88 const double* gsw_pref_ptr_target = mat_input + pref_doubles;
89 for (; gsw_pref_ptr < gsw_pref_ptr_target; gsw_pref_ptr += 8) {
90 __builtin_prefetch(gsw_pref_ptr, 0, _MM_HINT_T0);
91 }
92 const double* mat_blk_start;
93 uint64_t blk_i;
94 for (blk_i = 0, mat_blk_start = mat_input; blk_i < ms4; blk_i++, mat_blk_start += gsw_iter_doubles) {
95 // prefetch the next iteration
96 if (gsw_pref_ptr_target < gsw_ptr_end) {
97 gsw_pref_ptr_target += gsw_iter_doubles;
98 if (gsw_pref_ptr_target > gsw_ptr_end) gsw_pref_ptr_target = gsw_ptr_end;
99 for (; gsw_pref_ptr < gsw_pref_ptr_target; gsw_pref_ptr += 8) {
100 __builtin_prefetch(gsw_pref_ptr, 0, _MM_HINT_T0);
101 }
102 }
103 reim4_extract_1blk_from_contiguous_reim_sl_avx(m, a_sl, row_max, blk_i, extracted_blk, a_dft);
104 // apply mat2cols
105 for (uint64_t col_i = 0; col_i < col_max - 1; col_i += 2) {
106 uint64_t col_offset = col_i * (8 * nrows);
107 reim4_vec_mat2cols_product_avx2(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset);
108
109 reim4_save_1blk_to_reim_avx(m, blk_i, res + col_i * res_sl, mat2cols_output);
110 reim4_save_1blk_to_reim_avx(m, blk_i, res + (col_i + 1) * res_sl, mat2cols_output + 8);
111 }
112
113 // check if col_max is odd, then special case
114 if (col_max % 2 == 1) {
115 uint64_t last_col = col_max - 1;
116 uint64_t col_offset = last_col * (8 * nrows);
117
118 // the last column is alone in the pmat: vec_mat1col
119 if (ncols == col_max) {
120 reim4_vec_mat1col_product_avx2(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset);
121 } else {
122 // the last column is part of a colpair in the pmat: vec_mat2cols and ignore the second position
123 reim4_vec_mat2cols_product_avx2(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset);
124 }
125 reim4_save_1blk_to_reim_avx(m, blk_i, res + last_col * res_sl, mat2cols_output);
126 }
127 }
128 } else {
129 const double* in;
130 uint64_t in_sl;
131 if (res == a_dft) {
132 // it is in place: copy the input vector
133 in = (double*)tmp_space;
134 in_sl = nn;
135 // vec_rnx_copy(module, (double*)tmp_space, row_max, nn, a_dft, row_max, a_sl);
136 for (uint64_t row_i = 0; row_i < row_max; row_i++) {
137 memcpy((double*)tmp_space + row_i * nn, a_dft + row_i * a_sl, nn * sizeof(double));
138 }
139 } else {
140 // it is out of place: do the product directly
141 in = a_dft;
142 in_sl = a_sl;
143 }
144 for (uint64_t col_i = 0; col_i < col_max; col_i++) {
145 double* pmat_col = mat_input + col_i * nrows * nn;
146 {
147 reim_fftvec_mul(module->precomp.fft64.p_fftvec_mul, //
148 res + col_i * res_sl, //
149 in, //
150 pmat_col);
151 }
152 for (uint64_t row_i = 1; row_i < row_max; row_i++) {
153 reim_fftvec_addmul(module->precomp.fft64.p_fftvec_addmul, //
154 res + col_i * res_sl, //
155 in + row_i * in_sl, //
156 pmat_col + row_i * nn);
157 }
158 }
159 }
160 }
161 // zero out remaining bytes (if any)
162 for (uint64_t i = col_max; i < res_size; ++i) {
163 memset(res + i * res_sl, 0, nn * sizeof(double));
164 }
165}
166
167/** @brief applies a vmp product res = a x pmat */
168EXPORT void fft64_rnx_vmp_apply_tmp_a_avx( //
169 const MOD_RNX* module, // N
170 double* res, uint64_t res_size, uint64_t res_sl, // res (addr must be != a)
171 double* tmpa, uint64_t a_size, uint64_t a_sl, // a (will be overwritten)
172 const RNX_VMP_PMAT* pmat, uint64_t nrows, uint64_t ncols, // prep matrix
173 uint8_t* tmp_space // scratch space
174) {
175 const uint64_t nn = module->n;
176 const uint64_t rows = nrows < a_size ? nrows : a_size;
177 const uint64_t cols = ncols < res_size ? ncols : res_size;
178
179 // fft is done in place on the input (tmpa is destroyed)
180 for (uint64_t i = 0; i < rows; ++i) {
181 reim_fft(module->precomp.fft64.p_fft, tmpa + i * a_sl);
182 }
183 fft64_rnx_vmp_apply_dft_to_dft_avx(module, //
184 res, cols, res_sl, //
185 tmpa, rows, a_sl, //
186 pmat, nrows, ncols, //
187 tmp_space);
188 // ifft is done in place on the output
189 for (uint64_t i = 0; i < cols; ++i) {
190 reim_ifft(module->precomp.fft64.p_ifft, res + i * res_sl);
191 }
192 // zero out the remaining positions
193 for (uint64_t i = cols; i < res_size; ++i) {
194 memset(res + i * res_sl, 0, nn * sizeof(double));
195 }
196}