poulpy/spqlios/lib/test/spqlios_vec_rnx_vmp_test.cpp
spqlios/lib/test/spqlios_vec_rnx_vmp_test.cppRaw
1#include "gtest/gtest.h"
2#include "../spqlios/arithmetic/vec_rnx_arithmetic_private.h"
3#include "../spqlios/reim/reim_fft.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_CONTIGUOUS_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_contiguous_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_contiguous_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_contiguous_tmp_bytes_avx);
183}
184#endif
185
186/// rnx_vmp_apply_dft_to_dft
187
188static void test_vmp_apply_tmp_a_outplace( //
189 RNX_VMP_APPLY_TMP_A_F* apply, //
190 RNX_VMP_APPLY_TMP_A_TMP_BYTES_F* tmp_bytes) {
191 for (uint64_t nn : {2, 4, 8, 64}) {
192 MOD_RNX* module = new_rnx_module_info(nn, FFT64);
193 for (uint64_t mat_nrows : {1, 4, 7}) {
194 for (uint64_t mat_ncols : {1, 2, 5}) {
195 for (uint64_t in_size : {1, 4, 7}) {
196 for (uint64_t out_size : {1, 2, 5}) {
197 const uint64_t in_sl = nn + uniform_u64_bits(2);
198 const uint64_t out_sl = nn + uniform_u64_bits(2);
199 rnx_vec_f64_layout in(nn, in_size, in_sl);
200 fft64_rnx_vmp_pmat_layout pmat(nn, mat_nrows, mat_ncols);
201 rnx_vec_f64_layout out(nn, out_size, out_sl);
202 in.fill_random(0);
203 pmat.fill_random(0);
204 // naive computation of the product
205 std::vector<rnx_f64> expect(out_size);
206 for (uint64_t col = 0; col < out_size; ++col) {
207 reim_fft64vec ex = reim_fft64vec::zero(nn);
208 for (uint64_t row = 0; row < std::min(mat_nrows, in_size); ++row) {
209 ex += pmat.get_zext(row, col) * simple_fft64(in.get_copy(row));
210 }
211 expect[col] = simple_ifft64(ex);
212 }
213 // apply the product
214 std::vector<uint8_t> tmp(tmp_bytes(module, out_size, in_size, mat_nrows, mat_ncols));
215 apply(module, //
216 out.data(), out_size, out_sl, //
217 in.data(), in_size, in_sl, //
218 pmat.data, mat_nrows, mat_ncols, //
219 tmp.data());
220 // check that the output is close from the expectation
221 for (uint64_t col = 0; col < out_size; ++col) {
222 rnx_f64 actual = out.get_copy_zext(col);
223 ASSERT_LE(infty_dist(actual, expect[col]), 1e-10);
224 }
225 }
226 }
227 }
228 }
229 delete_rnx_module_info(module);
230 }
231}
232
233static void test_vmp_apply_tmp_a_inplace( //
234 RNX_VMP_APPLY_TMP_A_F* apply, //
235 RNX_VMP_APPLY_TMP_A_TMP_BYTES_F* tmp_bytes) {
236 for (uint64_t nn : {2, 4, 8, 64}) {
237 MOD_RNX* module = new_rnx_module_info(nn, FFT64);
238 for (uint64_t mat_nrows : {1, 4, 7}) {
239 for (uint64_t mat_ncols : {1, 2, 5}) {
240 for (uint64_t in_size : {1, 4, 7}) {
241 for (uint64_t out_size : {1, 2, 5}) {
242 const uint64_t in_out_sl = nn + uniform_u64_bits(2);
243 rnx_vec_f64_layout in_out(nn, std::max(in_size, out_size), in_out_sl);
244 fft64_rnx_vmp_pmat_layout pmat(nn, mat_nrows, mat_ncols);
245 in_out.fill_random(0);
246 pmat.fill_random(0);
247 // naive computation of the product
248 std::vector<rnx_f64> expect(out_size);
249 for (uint64_t col = 0; col < out_size; ++col) {
250 reim_fft64vec ex = reim_fft64vec::zero(nn);
251 for (uint64_t row = 0; row < std::min(mat_nrows, in_size); ++row) {
252 ex += pmat.get_zext(row, col) * simple_fft64(in_out.get_copy(row));
253 }
254 expect[col] = simple_ifft64(ex);
255 }
256 // apply the product
257 std::vector<uint8_t> tmp(tmp_bytes(module, out_size, in_size, mat_nrows, mat_ncols));
258 apply(module, //
259 in_out.data(), out_size, in_out_sl, //
260 in_out.data(), in_size, in_out_sl, //
261 pmat.data, mat_nrows, mat_ncols, //
262 tmp.data());
263 // check that the output is close from the expectation
264 for (uint64_t col = 0; col < out_size; ++col) {
265 rnx_f64 actual = in_out.get_copy_zext(col);
266 ASSERT_LE(infty_dist(actual, expect[col]), 1e-10);
267 }
268 }
269 }
270 }
271 }
272 delete_rnx_module_info(module);
273 }
274}
275
276static void test_vmp_apply_tmp_a( //
277 RNX_VMP_APPLY_TMP_A_F* apply, //
278 RNX_VMP_APPLY_TMP_A_TMP_BYTES_F* tmp_bytes) {
279 test_vmp_apply_tmp_a_outplace(apply, tmp_bytes);
280 test_vmp_apply_tmp_a_inplace(apply, tmp_bytes);
281}
282
283TEST(vec_znx, fft64_vmp_apply_tmp_a) { test_vmp_apply_tmp_a(rnx_vmp_apply_tmp_a, rnx_vmp_apply_tmp_a_tmp_bytes); }
284TEST(vec_znx, fft64_vmp_apply_tmp_a_ref) {
285 test_vmp_apply_tmp_a(fft64_rnx_vmp_apply_tmp_a_ref, fft64_rnx_vmp_apply_tmp_a_tmp_bytes_ref);
286}
287#ifdef __x86_64__
288TEST(vec_znx, fft64_vmp_apply_tmp_a_avx) {
289 test_vmp_apply_tmp_a(fft64_rnx_vmp_apply_tmp_a_avx, fft64_rnx_vmp_apply_tmp_a_tmp_bytes_avx);
290}
291#endif