poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/test/spqlios_vec_rnx_vmp_test.cpp
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/test/spqlios_vec_rnx_vmp_test.cppRaw
1#include "../spqlios/arithmetic/vec_rnx_arithmetic_private.h"
2#include "../spqlios/reim/reim_fft.h"
3#include "gtest/gtest.h"
4#include "testlib/vec_rnx_layout.h"
5
6static void test_vmp_apply_dft_to_dft_outplace( //
7 RNX_VMP_APPLY_DFT_TO_DFT_F* apply, //
8 RNX_VMP_APPLY_DFT_TO_DFT_TMP_BYTES_F* tmp_bytes) {
9 for (uint64_t nn : {2, 4, 8, 64}) {
10 MOD_RNX* module = new_rnx_module_info(nn, FFT64);
11 for (uint64_t mat_nrows : {1, 4, 7}) {
12 for (uint64_t mat_ncols : {1, 2, 5}) {
13 for (uint64_t in_size : {1, 4, 7}) {
14 for (uint64_t out_size : {1, 2, 5}) {
15 const uint64_t in_sl = nn + uniform_u64_bits(2);
16 const uint64_t out_sl = nn + uniform_u64_bits(2);
17 rnx_vec_f64_layout in(nn, in_size, in_sl);
18 fft64_rnx_vmp_pmat_layout pmat(nn, mat_nrows, mat_ncols);
19 rnx_vec_f64_layout out(nn, out_size, out_sl);
20 in.fill_random(0);
21 pmat.fill_random(0);
22 // naive computation of the product
23 std::vector<reim_fft64vec> expect(out_size, reim_fft64vec(nn));
24 for (uint64_t col = 0; col < out_size; ++col) {
25 reim_fft64vec ex = reim_fft64vec::zero(nn);
26 for (uint64_t row = 0; row < std::min(mat_nrows, in_size); ++row) {
27 ex += pmat.get_zext(row, col) * in.get_dft_copy(row);
28 }
29 expect[col] = ex;
30 }
31 // apply the product
32 std::vector<uint8_t> tmp(tmp_bytes(module, out_size, in_size, mat_nrows, mat_ncols));
33 apply(module, //
34 out.data(), out_size, out_sl, //
35 in.data(), in_size, in_sl, //
36 pmat.data, mat_nrows, mat_ncols, //
37 tmp.data());
38 // check that the output is close from the expectation
39 for (uint64_t col = 0; col < out_size; ++col) {
40 reim_fft64vec actual = out.get_dft_copy_zext(col);
41 ASSERT_LE(infty_dist(actual, expect[col]), 1e-10);
42 }
43 }
44 }
45 }
46 }
47 delete_rnx_module_info(module);
48 }
49}
50
51static void test_vmp_apply_dft_to_dft_inplace( //
52 RNX_VMP_APPLY_DFT_TO_DFT_F* apply, //
53 RNX_VMP_APPLY_DFT_TO_DFT_TMP_BYTES_F* tmp_bytes) {
54 for (uint64_t nn : {2, 4, 8, 64}) {
55 MOD_RNX* module = new_rnx_module_info(nn, FFT64);
56 for (uint64_t mat_nrows : {1, 2, 6}) {
57 for (uint64_t mat_ncols : {1, 2, 7, 8}) {
58 for (uint64_t in_size : {1, 3, 6}) {
59 for (uint64_t out_size : {1, 3, 6}) {
60 const uint64_t in_out_sl = nn + uniform_u64_bits(2);
61 rnx_vec_f64_layout in_out(nn, std::max(in_size, out_size), in_out_sl);
62 fft64_rnx_vmp_pmat_layout pmat(nn, mat_nrows, mat_ncols);
63 in_out.fill_random(0);
64 pmat.fill_random(0);
65 // naive computation of the product
66 std::vector<reim_fft64vec> expect(out_size, reim_fft64vec(nn));
67 for (uint64_t col = 0; col < out_size; ++col) {
68 reim_fft64vec ex = reim_fft64vec::zero(nn);
69 for (uint64_t row = 0; row < std::min(mat_nrows, in_size); ++row) {
70 ex += pmat.get_zext(row, col) * in_out.get_dft_copy(row);
71 }
72 expect[col] = ex;
73 }
74 // apply the product
75 std::vector<uint8_t> tmp(tmp_bytes(module, out_size, in_size, mat_nrows, mat_ncols));
76 apply(module, //
77 in_out.data(), out_size, in_out_sl, //
78 in_out.data(), in_size, in_out_sl, //
79 pmat.data, mat_nrows, mat_ncols, //
80 tmp.data());
81 // check that the output is close from the expectation
82 for (uint64_t col = 0; col < out_size; ++col) {
83 reim_fft64vec actual = in_out.get_dft_copy_zext(col);
84 ASSERT_LE(infty_dist(actual, expect[col]), 1e-10);
85 }
86 }
87 }
88 }
89 }
90 delete_rnx_module_info(module);
91 }
92}
93
94static void test_vmp_apply_dft_to_dft( //
95 RNX_VMP_APPLY_DFT_TO_DFT_F* apply, //
96 RNX_VMP_APPLY_DFT_TO_DFT_TMP_BYTES_F* tmp_bytes) {
97 test_vmp_apply_dft_to_dft_outplace(apply, tmp_bytes);
98 test_vmp_apply_dft_to_dft_inplace(apply, tmp_bytes);
99}
100
101TEST(vec_rnx, vmp_apply_to_dft) {
102 test_vmp_apply_dft_to_dft(rnx_vmp_apply_dft_to_dft, rnx_vmp_apply_dft_to_dft_tmp_bytes);
103}
104TEST(vec_rnx, fft64_vmp_apply_dft_to_dft_ref) {
105 test_vmp_apply_dft_to_dft(fft64_rnx_vmp_apply_dft_to_dft_ref, fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_ref);
106}
107#ifdef __x86_64__
108TEST(vec_rnx, fft64_vmp_apply_dft_to_dft_avx) {
109 test_vmp_apply_dft_to_dft(fft64_rnx_vmp_apply_dft_to_dft_avx, fft64_rnx_vmp_apply_dft_to_dft_tmp_bytes_avx);
110}
111#endif
112
113/// rnx_vmp_prepare
114
115static void test_vmp_prepare_contiguous(RNX_VMP_PREPARE_CONTIGUOUS_F* prepare_contiguous,
116 RNX_VMP_PREPARE_TMP_BYTES_F* tmp_bytes) {
117 // tests when n < 8
118 for (uint64_t nn : {2, 4}) {
119 const double one_over_m = 2. / nn;
120 MOD_RNX* module = new_rnx_module_info(nn, FFT64);
121 for (uint64_t nrows : {1, 2, 5}) {
122 for (uint64_t ncols : {2, 6, 7}) {
123 rnx_vec_f64_layout mat(nn, nrows * ncols, nn);
124 fft64_rnx_vmp_pmat_layout pmat(nn, nrows, ncols);
125 mat.fill_random(0);
126 std::vector<uint8_t> tmp_space(tmp_bytes(module));
127 thash hash_before = mat.content_hash();
128 prepare_contiguous(module, pmat.data, mat.data(), nrows, ncols, tmp_space.data());
129 ASSERT_EQ(mat.content_hash(), hash_before);
130 for (uint64_t row = 0; row < nrows; ++row) {
131 for (uint64_t col = 0; col < ncols; ++col) {
132 const double* pmatv = (double*)pmat.data + (col * nrows + row) * nn;
133 reim_fft64vec tmp = one_over_m * simple_fft64(mat.get_copy(row * ncols + col));
134 const double* tmpv = tmp.data();
135 for (uint64_t i = 0; i < nn; ++i) {
136 ASSERT_LE(abs(pmatv[i] - tmpv[i]), 1e-10);
137 }
138 }
139 }
140 }
141 }
142 delete_rnx_module_info(module);
143 }
144 // tests when n >= 8
145 for (uint64_t nn : {8, 32}) {
146 const double one_over_m = 2. / nn;
147 MOD_RNX* module = new_rnx_module_info(nn, FFT64);
148 uint64_t nblk = nn / 8;
149 for (uint64_t nrows : {1, 2, 5}) {
150 for (uint64_t ncols : {2, 6, 7}) {
151 rnx_vec_f64_layout mat(nn, nrows * ncols, nn);
152 fft64_rnx_vmp_pmat_layout pmat(nn, nrows, ncols);
153 mat.fill_random(0);
154 std::vector<uint8_t> tmp_space(tmp_bytes(module));
155 thash hash_before = mat.content_hash();
156 prepare_contiguous(module, pmat.data, mat.data(), nrows, ncols, tmp_space.data());
157 ASSERT_EQ(mat.content_hash(), hash_before);
158 for (uint64_t row = 0; row < nrows; ++row) {
159 for (uint64_t col = 0; col < ncols; ++col) {
160 reim_fft64vec tmp = one_over_m * simple_fft64(mat.get_copy(row * ncols + col));
161 for (uint64_t blk = 0; blk < nblk; ++blk) {
162 reim4_elem expect = tmp.get_blk(blk);
163 reim4_elem actual = pmat.get(row, col, blk);
164 ASSERT_LE(infty_dist(actual, expect), 1e-10);
165 }
166 }
167 }
168 }
169 }
170 delete_rnx_module_info(module);
171 }
172}
173
174TEST(vec_rnx, vmp_prepare_contiguous) {
175 test_vmp_prepare_contiguous(rnx_vmp_prepare_contiguous, rnx_vmp_prepare_tmp_bytes);
176}
177TEST(vec_rnx, fft64_vmp_prepare_contiguous_ref) {
178 test_vmp_prepare_contiguous(fft64_rnx_vmp_prepare_contiguous_ref, fft64_rnx_vmp_prepare_tmp_bytes_ref);
179}
180#ifdef __x86_64__
181TEST(vec_rnx, fft64_vmp_prepare_contiguous_avx) {
182 test_vmp_prepare_contiguous(fft64_rnx_vmp_prepare_contiguous_avx, fft64_rnx_vmp_prepare_tmp_bytes_avx);
183}
184#endif
185
186/// rnx_vmp_prepare_dblptr
187
188static void test_vmp_prepare_dblptr(RNX_VMP_PREPARE_DBLPTR_F* prepare_dblptr, RNX_VMP_PREPARE_TMP_BYTES_F* tmp_bytes) {
189 // tests when n < 8
190 for (uint64_t nn : {2, 4}) {
191 const double one_over_m = 2. / nn;
192 MOD_RNX* module = new_rnx_module_info(nn, FFT64);
193 for (uint64_t nrows : {1, 2, 5}) {
194 for (uint64_t ncols : {2, 6, 7}) {
195 rnx_vec_f64_layout mat(nn, nrows * ncols, nn);
196 fft64_rnx_vmp_pmat_layout pmat(nn, nrows, ncols);
197 mat.fill_random(0);
198 std::vector<uint8_t> tmp_space(tmp_bytes(module));
199 thash hash_before = mat.content_hash();
200 const double** mat_dblptr = (const double**)malloc(nrows * sizeof(double*));
201 for (size_t row_i = 0; row_i < nrows; row_i++) {
202 mat_dblptr[row_i] = &mat.data()[row_i * ncols * nn];
203 };
204 prepare_dblptr(module, pmat.data, mat_dblptr, nrows, ncols, tmp_space.data());
205 ASSERT_EQ(mat.content_hash(), hash_before);
206 for (uint64_t row = 0; row < nrows; ++row) {
207 for (uint64_t col = 0; col < ncols; ++col) {
208 const double* pmatv = (double*)pmat.data + (col * nrows + row) * nn;
209 reim_fft64vec tmp = one_over_m * simple_fft64(mat.get_copy(row * ncols + col));
210 const double* tmpv = tmp.data();
211 for (uint64_t i = 0; i < nn; ++i) {
212 ASSERT_LE(abs(pmatv[i] - tmpv[i]), 1e-10);
213 }
214 }
215 }
216 }
217 }
218 delete_rnx_module_info(module);
219 }
220 // tests when n >= 8
221 for (uint64_t nn : {8, 32}) {
222 const double one_over_m = 2. / nn;
223 MOD_RNX* module = new_rnx_module_info(nn, FFT64);
224 uint64_t nblk = nn / 8;
225 for (uint64_t nrows : {1, 2, 5}) {
226 for (uint64_t ncols : {2, 6, 7}) {
227 rnx_vec_f64_layout mat(nn, nrows * ncols, nn);
228 fft64_rnx_vmp_pmat_layout pmat(nn, nrows, ncols);
229 mat.fill_random(0);
230 std::vector<uint8_t> tmp_space(tmp_bytes(module));
231 thash hash_before = mat.content_hash();
232 const double** mat_dblptr = (const double**)malloc(nrows * sizeof(double*));
233 for (size_t row_i = 0; row_i < nrows; row_i++) {
234 mat_dblptr[row_i] = &mat.data()[row_i * ncols * nn];
235 };
236 prepare_dblptr(module, pmat.data, mat_dblptr, nrows, ncols, tmp_space.data());
237 ASSERT_EQ(mat.content_hash(), hash_before);
238 for (uint64_t row = 0; row < nrows; ++row) {
239 for (uint64_t col = 0; col < ncols; ++col) {
240 reim_fft64vec tmp = one_over_m * simple_fft64(mat.get_copy(row * ncols + col));
241 for (uint64_t blk = 0; blk < nblk; ++blk) {
242 reim4_elem expect = tmp.get_blk(blk);
243 reim4_elem actual = pmat.get(row, col, blk);
244 ASSERT_LE(infty_dist(actual, expect), 1e-10);
245 }
246 }
247 }
248 }
249 }
250 delete_rnx_module_info(module);
251 }
252}
253
254TEST(vec_rnx, vmp_prepare_dblptr) { test_vmp_prepare_dblptr(rnx_vmp_prepare_dblptr, rnx_vmp_prepare_tmp_bytes); }
255TEST(vec_rnx, fft64_vmp_prepare_dblptr_ref) {
256 test_vmp_prepare_dblptr(fft64_rnx_vmp_prepare_dblptr_ref, fft64_rnx_vmp_prepare_tmp_bytes_ref);
257}
258#ifdef __x86_64__
259TEST(vec_rnx, fft64_vmp_prepare_dblptr_avx) {
260 test_vmp_prepare_dblptr(fft64_rnx_vmp_prepare_dblptr_avx, fft64_rnx_vmp_prepare_tmp_bytes_avx);
261}
262#endif
263
264/// rnx_vmp_apply_dft_to_dft
265
266static void test_vmp_apply_tmp_a_outplace( //
267 RNX_VMP_APPLY_TMP_A_F* apply, //
268 RNX_VMP_APPLY_TMP_A_TMP_BYTES_F* tmp_bytes) {
269 for (uint64_t nn : {2, 4, 8, 64}) {
270 MOD_RNX* module = new_rnx_module_info(nn, FFT64);
271 for (uint64_t mat_nrows : {1, 4, 7}) {
272 for (uint64_t mat_ncols : {1, 2, 5}) {
273 for (uint64_t in_size : {1, 4, 7}) {
274 for (uint64_t out_size : {1, 2, 5}) {
275 const uint64_t in_sl = nn + uniform_u64_bits(2);
276 const uint64_t out_sl = nn + uniform_u64_bits(2);
277 rnx_vec_f64_layout in(nn, in_size, in_sl);
278 fft64_rnx_vmp_pmat_layout pmat(nn, mat_nrows, mat_ncols);
279 rnx_vec_f64_layout out(nn, out_size, out_sl);
280 in.fill_random(0);
281 pmat.fill_random(0);
282 // naive computation of the product
283 std::vector<rnx_f64> expect(out_size);
284 for (uint64_t col = 0; col < out_size; ++col) {
285 reim_fft64vec ex = reim_fft64vec::zero(nn);
286 for (uint64_t row = 0; row < std::min(mat_nrows, in_size); ++row) {
287 ex += pmat.get_zext(row, col) * simple_fft64(in.get_copy(row));
288 }
289 expect[col] = simple_ifft64(ex);
290 }
291 // apply the product
292 std::vector<uint8_t> tmp(tmp_bytes(module, out_size, in_size, mat_nrows, mat_ncols));
293 apply(module, //
294 out.data(), out_size, out_sl, //
295 in.data(), in_size, in_sl, //
296 pmat.data, mat_nrows, mat_ncols, //
297 tmp.data());
298 // check that the output is close from the expectation
299 for (uint64_t col = 0; col < out_size; ++col) {
300 rnx_f64 actual = out.get_copy_zext(col);
301 ASSERT_LE(infty_dist(actual, expect[col]), 1e-10);
302 }
303 }
304 }
305 }
306 }
307 delete_rnx_module_info(module);
308 }
309}
310
311static void test_vmp_apply_tmp_a_inplace( //
312 RNX_VMP_APPLY_TMP_A_F* apply, //
313 RNX_VMP_APPLY_TMP_A_TMP_BYTES_F* tmp_bytes) {
314 for (uint64_t nn : {2, 4, 8, 64}) {
315 MOD_RNX* module = new_rnx_module_info(nn, FFT64);
316 for (uint64_t mat_nrows : {1, 4, 7}) {
317 for (uint64_t mat_ncols : {1, 2, 5}) {
318 for (uint64_t in_size : {1, 4, 7}) {
319 for (uint64_t out_size : {1, 2, 5}) {
320 const uint64_t in_out_sl = nn + uniform_u64_bits(2);
321 rnx_vec_f64_layout in_out(nn, std::max(in_size, out_size), in_out_sl);
322 fft64_rnx_vmp_pmat_layout pmat(nn, mat_nrows, mat_ncols);
323 in_out.fill_random(0);
324 pmat.fill_random(0);
325 // naive computation of the product
326 std::vector<rnx_f64> expect(out_size);
327 for (uint64_t col = 0; col < out_size; ++col) {
328 reim_fft64vec ex = reim_fft64vec::zero(nn);
329 for (uint64_t row = 0; row < std::min(mat_nrows, in_size); ++row) {
330 ex += pmat.get_zext(row, col) * simple_fft64(in_out.get_copy(row));
331 }
332 expect[col] = simple_ifft64(ex);
333 }
334 // apply the product
335 std::vector<uint8_t> tmp(tmp_bytes(module, out_size, in_size, mat_nrows, mat_ncols));
336 apply(module, //
337 in_out.data(), out_size, in_out_sl, //
338 in_out.data(), in_size, in_out_sl, //
339 pmat.data, mat_nrows, mat_ncols, //
340 tmp.data());
341 // check that the output is close from the expectation
342 for (uint64_t col = 0; col < out_size; ++col) {
343 rnx_f64 actual = in_out.get_copy_zext(col);
344 ASSERT_LE(infty_dist(actual, expect[col]), 1e-10);
345 }
346 }
347 }
348 }
349 }
350 delete_rnx_module_info(module);
351 }
352}
353
354static void test_vmp_apply_tmp_a( //
355 RNX_VMP_APPLY_TMP_A_F* apply, //
356 RNX_VMP_APPLY_TMP_A_TMP_BYTES_F* tmp_bytes) {
357 test_vmp_apply_tmp_a_outplace(apply, tmp_bytes);
358 test_vmp_apply_tmp_a_inplace(apply, tmp_bytes);
359}
360
361TEST(vec_znx, fft64_vmp_apply_tmp_a) { test_vmp_apply_tmp_a(rnx_vmp_apply_tmp_a, rnx_vmp_apply_tmp_a_tmp_bytes); }
362TEST(vec_znx, fft64_vmp_apply_tmp_a_ref) {
363 test_vmp_apply_tmp_a(fft64_rnx_vmp_apply_tmp_a_ref, fft64_rnx_vmp_apply_tmp_a_tmp_bytes_ref);
364}
365#ifdef __x86_64__
366TEST(vec_znx, fft64_vmp_apply_tmp_a_avx) {
367 test_vmp_apply_tmp_a(fft64_rnx_vmp_apply_tmp_a_avx, fft64_rnx_vmp_apply_tmp_a_tmp_bytes_avx);
368}
369#endif