diff --git a/spqlios/arithmetic/module_api.c b/spqlios/arithmetic/module_api.c index 05631ec..8303ac9 100644 --- a/spqlios/arithmetic/module_api.c +++ b/spqlios/arithmetic/module_api.c @@ -47,11 +47,15 @@ static void fill_fft64_virtual_table(MODULE* module) { module->func.znx_small_single_product = fft64_znx_small_single_product; module->func.znx_small_single_product_tmp_bytes = fft64_znx_small_single_product_tmp_bytes; module->func.vmp_prepare_contiguous = fft64_vmp_prepare_contiguous_ref; + module->func.vmp_prepare_vector = fft64_vmp_prepare_vector_ref; module->func.vmp_prepare_contiguous_tmp_bytes = fft64_vmp_prepare_contiguous_tmp_bytes; + module->func.vmp_prepare_vector_tmp_bytes = fft64_vmp_prepare_vector_tmp_bytes; module->func.vmp_apply_dft = fft64_vmp_apply_dft_ref; module->func.vmp_apply_dft_tmp_bytes = fft64_vmp_apply_dft_tmp_bytes; module->func.vmp_apply_dft_to_dft = fft64_vmp_apply_dft_to_dft_ref; + module->func.vmp_apply_pvec_to_dft = fft64_vmp_apply_pvec_to_dft_ref; module->func.vmp_apply_dft_to_dft_tmp_bytes = fft64_vmp_apply_dft_to_dft_tmp_bytes; + module->func.vmp_apply_pvec_to_dft_tmp_bytes = fft64_vmp_apply_pvec_to_dft_tmp_bytes; module->func.cnv_prepare_right_contiguous_tmp_bytes = fft64_convolution_prepare_right_contiguous_tmp_bytes; module->func.cnv_prepare_left_contiguous_tmp_bytes = fft64_convolution_prepare_left_contiguous_tmp_bytes; module->func.cnv_prepare_right_contiguous = fft64_convolution_prepare_right_contiguous_ref; @@ -70,6 +74,7 @@ static void fill_fft64_virtual_table(MODULE* module) { module->func.vmp_prepare_contiguous = fft64_vmp_prepare_contiguous_avx; module->func.vmp_apply_dft = fft64_vmp_apply_dft_avx; module->func.vmp_apply_dft_to_dft = fft64_vmp_apply_dft_to_dft_avx; + module->func.vmp_apply_pvec_to_dft = fft64_vmp_apply_pvec_to_dft_avx; } } diff --git a/spqlios/arithmetic/vec_znx_arithmetic.h b/spqlios/arithmetic/vec_znx_arithmetic.h index fde7df0..9b6a2eb 100644 --- a/spqlios/arithmetic/vec_znx_arithmetic.h +++ b/spqlios/arithmetic/vec_znx_arithmetic.h @@ -25,6 +25,8 @@ typedef struct module_info_t MODULE; typedef struct vmp_pmat_t VMP_PMAT; /** @brief opaque type that represents a vector of znx in DFT space */ typedef struct vec_znx_dft_t VEC_ZNX_DFT; +/** @brief opaque type that represents a vector of znx in prepared for vmp DFT space */ +typedef struct vmp_pvec_t VMP_PVEC; /** @brief opaque type that represents a vector of znx in large coeffs space */ typedef struct vec_znx_bigcoeff_t VEC_ZNX_BIG; /** @brief opaque type that represents a prepared scalar vector product */ @@ -150,6 +152,13 @@ EXPORT void vmp_prepare_contiguous(const MODULE* module, uint8_t* tmp_space // scratch space ); +/** @brief prepares a vmp vector */ +EXPORT void vmp_prepare_vector(const MODULE* module, // N + VMP_PVEC* pvec, uint64_t nrows, // output + const int64_t* a, uint64_t a_size, uint64_t a_sl, // a + uint8_t* tmp_space // scratch space +); + /** @brief prepares a vmp matrix (mat[row*ncols+col] points to the item) */ EXPORT void vmp_prepare_dblptr(const MODULE* module, // N VMP_PMAT* pmat, // output @@ -330,6 +339,10 @@ EXPORT void vmp_prepare_contiguous(const MODULE* module, EXPORT uint64_t vmp_prepare_contiguous_tmp_bytes(const MODULE* module, // N uint64_t nrows, uint64_t ncols); +/** @brief minimal scratch space byte-size required for the vmp_prepare_vec function */ +EXPORT uint64_t vmp_prepare_vector_tmp_bytes(const MODULE* module, // N + uint64_t nrows, uint64_t a_size); + /** @brief applies a vmp product (result in DFT space) */ EXPORT void vmp_apply_dft(const MODULE* module, // N VEC_ZNX_DFT* res, uint64_t res_size, // res @@ -353,7 +366,14 @@ EXPORT void vmp_apply_dft_to_dft(const MODULE* module, // const uint64_t ncols, // prep matrix uint8_t* tmp_space // scratch space (a_size*sizeof(reim4) bytes) ); -; + +EXPORT void vmp_apply_pvec_to_dft(const MODULE* module, // N + VEC_ZNX_DFT* res, const uint64_t res_size, // res + const VMP_PVEC* a_prep, uint64_t a_size, // a + const VMP_PMAT* pmat, const uint64_t nrows, + const uint64_t ncols, // prep matrix + uint8_t* tmp_space // scratch space (a_size*sizeof(reim4) bytes) +); /** @brief minimal size of the tmp_space */ EXPORT uint64_t vmp_apply_dft_to_dft_tmp_bytes(const MODULE* module, // N @@ -362,6 +382,12 @@ EXPORT uint64_t vmp_apply_dft_to_dft_tmp_bytes(const MODULE* module, / uint64_t nrows, uint64_t ncols // prep matrix ); +EXPORT uint64_t vmp_apply_pvec_to_dft_tmp_bytes(const MODULE* module, // N + uint64_t res_size, // res + uint64_t a_size, // a + uint64_t nrows, uint64_t ncols // prep matrix +); + /** @brief prepares the right vector for convolution */ EXPORT void cnv_prepare_right_contiguous(const MODULE* module, // N CNV_PVEC_R* pvec, uint64_t nrows, // output diff --git a/spqlios/arithmetic/vec_znx_arithmetic_private.h b/spqlios/arithmetic/vec_znx_arithmetic_private.h index 9b8551c..25e62ef 100644 --- a/spqlios/arithmetic/vec_znx_arithmetic_private.h +++ b/spqlios/arithmetic/vec_znx_arithmetic_private.h @@ -85,7 +85,9 @@ typedef typeof(svp_apply_dft_to_dft) SVP_APPLY_DFT_TO_DFT_F; typedef typeof(znx_small_single_product) ZNX_SMALL_SINGLE_PRODUCT_F; typedef typeof(znx_small_single_product_tmp_bytes) ZNX_SMALL_SINGLE_PRODUCT_TMP_BYTES_F; typedef typeof(vmp_prepare_contiguous) VMP_PREPARE_CONTIGUOUS_F; +typedef typeof(vmp_prepare_vector) VMP_PREPARE_VECTOR_F; typedef typeof(vmp_prepare_contiguous_tmp_bytes) VMP_PREPARE_CONTIGUOUS_TMP_BYTES_F; +typedef typeof(vmp_prepare_vector_tmp_bytes) VMP_PREPARE_VECTOR_TMP_BYTES_F; typedef typeof(cnv_apply_dft) CNV_APPLY_DFT_F; typedef typeof(cnv_apply_dft_tmp_bytes) CNV_APPLY_DFT_TMP_BYTES_F; typedef typeof(cnv_prepare_left_contiguous) CNV_PREPARE_LEFT_CONTIGUOUS_F; @@ -95,7 +97,9 @@ typedef typeof(cnv_prepare_right_contiguous_tmp_bytes) CNV_PREPARE_RIGHT_CONTIGU typedef typeof(vmp_apply_dft) VMP_APPLY_DFT_F; typedef typeof(vmp_apply_dft_tmp_bytes) VMP_APPLY_DFT_TMP_BYTES_F; typedef typeof(vmp_apply_dft_to_dft) VMP_APPLY_DFT_TO_DFT_F; +typedef typeof(vmp_apply_pvec_to_dft) VMP_APPLY_PVEC_TO_DFT_F; typedef typeof(vmp_apply_dft_to_dft_tmp_bytes) VMP_APPLY_DFT_TO_DFT_TMP_BYTES_F; +typedef typeof(vmp_apply_pvec_to_dft_tmp_bytes) VMP_APPLY_PVEC_TO_DFT_TMP_BYTES_F; typedef typeof(bytes_of_vec_znx_dft) BYTES_OF_VEC_ZNX_DFT_F; typedef typeof(bytes_of_vec_znx_big) BYTES_OF_VEC_ZNX_BIG_F; typedef typeof(bytes_of_svp_ppol) BYTES_OF_SVP_PPOL_F; @@ -135,11 +139,15 @@ struct module_virtual_functions_t { ZNX_SMALL_SINGLE_PRODUCT_F* znx_small_single_product; ZNX_SMALL_SINGLE_PRODUCT_TMP_BYTES_F* znx_small_single_product_tmp_bytes; VMP_PREPARE_CONTIGUOUS_F* vmp_prepare_contiguous; + VMP_PREPARE_VECTOR_F* vmp_prepare_vector; VMP_PREPARE_CONTIGUOUS_TMP_BYTES_F* vmp_prepare_contiguous_tmp_bytes; + VMP_PREPARE_VECTOR_TMP_BYTES_F* vmp_prepare_vector_tmp_bytes; VMP_APPLY_DFT_F* vmp_apply_dft; VMP_APPLY_DFT_TMP_BYTES_F* vmp_apply_dft_tmp_bytes; VMP_APPLY_DFT_TO_DFT_F* vmp_apply_dft_to_dft; + VMP_APPLY_PVEC_TO_DFT_F* vmp_apply_pvec_to_dft; VMP_APPLY_DFT_TO_DFT_TMP_BYTES_F* vmp_apply_dft_to_dft_tmp_bytes; + VMP_APPLY_PVEC_TO_DFT_TMP_BYTES_F* vmp_apply_pvec_to_dft_tmp_bytes; CNV_APPLY_DFT_F* cnv_apply_dft; CNV_APPLY_DFT_TMP_BYTES_F* cnv_apply_dft_tmp_bytes; CNV_PREPARE_LEFT_CONTIGUOUS_F* cnv_prepare_left_contiguous; @@ -501,6 +509,15 @@ EXPORT void fft64_vmp_prepare_contiguous_avx(const MODULE* module, EXPORT uint64_t fft64_vmp_prepare_contiguous_tmp_bytes(const MODULE* module, // N uint64_t nrows, uint64_t ncols); +EXPORT uint64_t fft64_vmp_prepare_vector_tmp_bytes(const MODULE* module, // N + uint64_t nrows, uint64_t a_size); + +EXPORT void fft64_vmp_prepare_vector_ref(const MODULE* module, // N + VMP_PVEC* pvec, uint64_t nrows, // output + const int64_t* a, uint64_t a_size, uint64_t a_sl, // a + uint8_t* tmp_space // scratch space +); + /** @brief applies a vmp product (result in DFT space) */ EXPORT void fft64_vmp_apply_dft_ref(const MODULE* module, // N VEC_ZNX_DFT* res, uint64_t res_size, // res @@ -535,6 +552,22 @@ EXPORT void fft64_vmp_apply_dft_to_dft_avx(const MODULE* module, uint8_t* tmp_space // scratch space (a_size*sizeof(reim4) bytes) ); +EXPORT void fft64_vmp_apply_pvec_to_dft_ref(const MODULE* module, // N + VEC_ZNX_DFT* res, const uint64_t res_size, // res + const VMP_PVEC* a_prep, uint64_t a_size, // a + const VMP_PMAT* pmat, const uint64_t nrows, + const uint64_t ncols, // prep matrix + uint8_t* tmp_space // scratch space (128 bytes) +); + +EXPORT void fft64_vmp_apply_pvec_to_dft_avx(const MODULE* module, // N + VEC_ZNX_DFT* res, const uint64_t res_size, // res + const VMP_PVEC* a_prep, uint64_t a_size, // a + const VMP_PMAT* pmat, const uint64_t nrows, + const uint64_t ncols, // prep matrix + uint8_t* tmp_space // scratch space (128 bytes) +); + /** @brief minimal size of the tmp_space */ EXPORT uint64_t fft64_vmp_apply_dft_tmp_bytes(const MODULE* module, // N uint64_t res_size, // res @@ -548,4 +581,10 @@ EXPORT uint64_t fft64_vmp_apply_dft_to_dft_tmp_bytes(const MODULE* module, uint64_t a_size, // a uint64_t nrows, uint64_t ncols // prep matrix ); + +EXPORT uint64_t fft64_vmp_apply_pvec_to_dft_tmp_bytes(const MODULE* module, // N + uint64_t res_size, // res + uint64_t a_size, // a + uint64_t nrows, uint64_t ncols // prep matrix +); #endif // SPQLIOS_VEC_ZNX_ARITHMETIC_PRIVATE_H diff --git a/spqlios/arithmetic/vector_matrix_product.c b/spqlios/arithmetic/vector_matrix_product.c index e1951a3..a88e405 100644 --- a/spqlios/arithmetic/vector_matrix_product.c +++ b/spqlios/arithmetic/vector_matrix_product.c @@ -33,12 +33,27 @@ EXPORT void vmp_prepare_contiguous(const MODULE* module, module->func.vmp_prepare_contiguous(module, pmat, mat, nrows, ncols, tmp_space); } +/** @brief prepares a vmp vector */ +EXPORT void vmp_prepare_vector(const MODULE* module, // N + VMP_PVEC* pvec, uint64_t nrows, // output + const int64_t* a, uint64_t a_size, uint64_t a_sl, // a + uint8_t* tmp_space // scratch space +) { + module->func.vmp_prepare_vector(module, pvec, nrows, a, a_size, a_sl, tmp_space); +} + /** @brief minimal scratch space byte-size required for the vmp_prepare function */ EXPORT uint64_t vmp_prepare_contiguous_tmp_bytes(const MODULE* module, // N uint64_t nrows, uint64_t ncols) { return module->func.vmp_prepare_contiguous_tmp_bytes(module, nrows, ncols); } +/** @brief minimal scratch space byte-size required for the vmp_prepare_vec function */ +EXPORT uint64_t vmp_prepare_vector_tmp_bytes(const MODULE* module, // N + uint64_t nrows, uint64_t a_size) { + return module->func.vmp_prepare_vector_tmp_bytes(module, nrows, a_size); +} + /** @brief prepares a vmp matrix (contiguous row-major version) */ EXPORT void fft64_vmp_prepare_contiguous_ref(const MODULE* module, // N VMP_PMAT* pmat, // output @@ -94,6 +109,23 @@ EXPORT uint64_t fft64_vmp_prepare_contiguous_tmp_bytes(const MODULE* module, // return nn * sizeof(int64_t); } +/** @brief prepares a vmp vector */ +EXPORT void fft64_vmp_prepare_vector_ref(const MODULE* module, // N + VMP_PVEC* pvec, uint64_t nrows, // output + const int64_t* a, uint64_t a_size, uint64_t a_sl, // a + uint8_t* tmp_space // scratch space +) { + // Same data format as fft64 prepared convolution + fft64_convolution_prepare_contiguous_ref(module, (double*)pvec, nrows, a, a_size, a_sl, tmp_space); +} + +/** @brief minimal scratch space byte-size required for the vmp_prepare_contiguous_vec function */ +EXPORT uint64_t fft64_vmp_prepare_vector_tmp_bytes(const MODULE* module, // N + uint64_t nrows, uint64_t a_size) { + // Same format than right/left convolution + return fft64_convolution_prepare_right_contiguous_tmp_bytes(module, nrows, a_size); +} + /** @brief applies a vmp product (result in DFT space) */ EXPORT void fft64_vmp_apply_dft_ref(const MODULE* module, // N VEC_ZNX_DFT* res, uint64_t res_size, // res @@ -179,6 +211,60 @@ EXPORT void fft64_vmp_apply_dft_to_dft_ref(const MODULE* module, memset(vec_output + col_max * nn, 0, (res_size - col_max) * nn * sizeof(double)); } +EXPORT void fft64_vmp_apply_pvec_to_dft_ref(const MODULE* module, // N + VEC_ZNX_DFT* res, const uint64_t res_size, // res + const VMP_PVEC* a_prep, uint64_t a_size, // a + const VMP_PMAT* pmat, const uint64_t nrows, + const uint64_t ncols, // prep matrix + uint8_t* tmp_space // scratch space (a_size*sizeof(reim4) bytes) +) { + const uint64_t m = module->m; + const uint64_t nn = module->nn; + + double* mat2cols_output = (double*)tmp_space; // 128 bytes + + double* mat_input = (double*)pmat; + double* vec_output = (double*)res; + + const uint64_t row_max = nrows < a_size ? nrows : a_size; + const uint64_t col_max = ncols < res_size ? ncols : res_size; + + if (nn >= 8) { + for (uint64_t blk_i = 0; blk_i < m / 4; blk_i++) { + double* mat_blk_start = mat_input + blk_i * (8 * nrows * ncols); + + double* extracted_blk = (double*)a_prep + 4l * 2 * a_size * blk_i; + // apply mat2cols + for (uint64_t col_i = 0; col_i < col_max - 1; col_i += 2) { + uint64_t col_offset = col_i * (8 * nrows); + reim4_vec_mat2cols_product_ref(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset); + + reim4_save_1blk_to_reim_ref(m, blk_i, vec_output + col_i * nn, mat2cols_output); + reim4_save_1blk_to_reim_ref(m, blk_i, vec_output + (col_i + 1) * nn, mat2cols_output + 8); + } + + // check if col_max is odd, then special case + if (col_max % 2 == 1) { + uint64_t last_col = col_max - 1; + uint64_t col_offset = last_col * (8 * nrows); + + // the last column is alone in the pmat: vec_mat1col + if (ncols == col_max) { + reim4_vec_mat1col_product_ref(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset); + } else { + // the last column is part of a colpair in the pmat: vec_mat2cols and ignore the second position + reim4_vec_mat2cols_product_ref(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset); + } + reim4_save_1blk_to_reim_ref(m, blk_i, vec_output + last_col * nn, mat2cols_output); + } + } + } else { + NOT_IMPLEMENTED() + } + + // zero out remaining bytes + memset(vec_output + col_max * nn, 0, (res_size - col_max) * nn * sizeof(double)); +} /** @brief minimal size of the tmp_space */ EXPORT uint64_t fft64_vmp_apply_dft_tmp_bytes(const MODULE* module, // N uint64_t res_size, // res @@ -202,6 +288,15 @@ EXPORT uint64_t fft64_vmp_apply_dft_to_dft_tmp_bytes(const MODULE* module, return (128) + (64 * row_max); } +/** @brief minimal size of the tmp_space */ +EXPORT uint64_t fft64_vmp_apply_pvec_to_dft_tmp_bytes(const MODULE* module, // N + uint64_t res_size, // res + uint64_t a_size, // a + uint64_t nrows, uint64_t ncols // prep matrix +) { + return (128); +} + EXPORT void vmp_apply_dft_to_dft(const MODULE* module, // N VEC_ZNX_DFT* res, const uint64_t res_size, // res const VEC_ZNX_DFT* a_dft, uint64_t a_size, // a @@ -212,6 +307,16 @@ EXPORT void vmp_apply_dft_to_dft(const MODULE* module, // module->func.vmp_apply_dft_to_dft(module, res, res_size, a_dft, a_size, pmat, nrows, ncols, tmp_space); } +EXPORT void vmp_apply_pvec_to_dft(const MODULE* module, // N + VEC_ZNX_DFT* res, const uint64_t res_size, // res + const VMP_PVEC* a_prep, uint64_t a_size, // a + const VMP_PMAT* pmat, const uint64_t nrows, + const uint64_t ncols, // prep matrix + uint8_t* tmp_space // scratch space (a_size*sizeof(reim4) bytes) +) { + module->func.vmp_apply_pvec_to_dft(module, res, res_size, a_prep, a_size, pmat, nrows, ncols, tmp_space); +} + EXPORT uint64_t vmp_apply_dft_to_dft_tmp_bytes(const MODULE* module, // N uint64_t res_size, // res uint64_t a_size, // a @@ -220,6 +325,14 @@ EXPORT uint64_t vmp_apply_dft_to_dft_tmp_bytes(const MODULE* module, / return module->func.vmp_apply_dft_to_dft_tmp_bytes(module, res_size, a_size, nrows, ncols); } +EXPORT uint64_t vmp_apply_pvec_to_dft_tmp_bytes(const MODULE* module, // N + uint64_t res_size, // res + uint64_t a_size, // a + uint64_t nrows, uint64_t ncols // prep matrix +) { + return module->func.vmp_apply_pvec_to_dft_tmp_bytes(module, res_size, a_size, nrows, ncols); +} + /** @brief applies a vmp product (result in DFT space) */ EXPORT void vmp_apply_dft(const MODULE* module, // N VEC_ZNX_DFT* res, uint64_t res_size, // res diff --git a/spqlios/arithmetic/vector_matrix_product_avx.c b/spqlios/arithmetic/vector_matrix_product_avx.c index f428650..6ed44c1 100644 --- a/spqlios/arithmetic/vector_matrix_product_avx.c +++ b/spqlios/arithmetic/vector_matrix_product_avx.c @@ -135,3 +135,58 @@ EXPORT void fft64_vmp_apply_dft_to_dft_avx(const MODULE* module, // zero out remaining bytes memset(vec_output + col_max * nn, 0, (res_size - col_max) * nn * sizeof(double)); } + +EXPORT void fft64_vmp_apply_pvec_to_dft_avx(const MODULE* module, // N + VEC_ZNX_DFT* res, const uint64_t res_size, // res + const VMP_PVEC* a_prep, uint64_t a_size, // a + const VMP_PMAT* pmat, const uint64_t nrows, + const uint64_t ncols, // prep matrix + uint8_t* tmp_space // scratch space (a_size*sizeof(reim4) bytes) +) { + const uint64_t m = module->m; + const uint64_t nn = module->nn; + + double* mat2cols_output = (double*)tmp_space; // 128 bytes + + double* mat_input = (double*)pmat; + double* vec_output = (double*)res; + + const uint64_t row_max = nrows < a_size ? nrows : a_size; + const uint64_t col_max = ncols < res_size ? ncols : res_size; + + if (nn >= 8) { + for (uint64_t blk_i = 0; blk_i < m / 4; blk_i++) { + double* mat_blk_start = mat_input + blk_i * (8 * nrows * ncols); + + double* extracted_blk = (double*)a_prep + 4l * 2 * a_size * blk_i; + // apply mat2cols + for (uint64_t col_i = 0; col_i < col_max - 1; col_i += 2) { + uint64_t col_offset = col_i * (8 * nrows); + reim4_vec_mat2cols_product_avx2(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset); + + reim4_save_1blk_to_reim_avx(m, blk_i, vec_output + col_i * nn, mat2cols_output); + reim4_save_1blk_to_reim_avx(m, blk_i, vec_output + (col_i + 1) * nn, mat2cols_output + 8); + } + + // check if col_max is odd, then special case + if (col_max % 2 == 1) { + uint64_t last_col = col_max - 1; + uint64_t col_offset = last_col * (8 * nrows); + + // the last column is alone in the pmat: vec_mat1col + if (ncols == col_max) + reim4_vec_mat1col_product_avx2(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset); + else { + // the last column is part of a colpair in the pmat: vec_mat2cols and ignore the second position + reim4_vec_mat2cols_product_avx2(row_max, mat2cols_output, extracted_blk, mat_blk_start + col_offset); + } + reim4_save_1blk_to_reim_avx(m, blk_i, vec_output + last_col * nn, mat2cols_output); + } + } + } else { + NOT_IMPLEMENTED() + } + + // zero out remaining bytes + memset(vec_output + col_max * nn, 0, (res_size - col_max) * nn * sizeof(double)); +} diff --git a/test/spqlios_vmp_product_test.cpp b/test/spqlios_vmp_product_test.cpp index cb55818..2268a4a 100644 --- a/test/spqlios_vmp_product_test.cpp +++ b/test/spqlios_vmp_product_test.cpp @@ -119,3 +119,93 @@ TEST(vec_znx, fft64_vmp_apply_dft_to_dft_avx) { test_vmp_apply(fft64_vmp_apply_dft_to_dft_avx, fft64_vmp_apply_dft_to_dft_tmp_bytes); } #endif + +static void test_vmp_prepare_vector(VMP_PREPARE_VECTOR_F* prepare_contiguous, + VMP_PREPARE_VECTOR_TMP_BYTES_F* tmp_bytes) { + // tests when n >= 8 + for (uint64_t nn : {8, 32}) { + MODULE* module = new_module_info(nn, FFT64); + uint64_t nblk = nn / 8; + for (uint64_t in_nrows : {0, 1, 2, 5}) { + for (uint64_t out_nrows : {0, 1, 3, 5}) { + znx_vec_i64_layout vec(nn, in_nrows, nn); + fft64_vmp_pvec_layout pvec(nn, out_nrows); + vec.fill_random(30); + std::vector tmp_space(tmp_bytes(module, out_nrows, in_nrows)); + thash hash_before = vec.content_hash(); + prepare_contiguous(module, // + pvec.data, out_nrows, // + vec.data(), in_nrows, nn, tmp_space.data()); + ASSERT_EQ(vec.content_hash(), hash_before); + for (uint64_t row = 0; row < out_nrows; ++row) { + reim_fft64vec tmp = simple_fft64(vec.get_copy_zext(row)); + for (uint64_t blk = 0; blk < nblk; ++blk) { + reim4_elem expect = tmp.get_blk(blk); + reim4_elem actual = pvec.get(row, blk); + ASSERT_LE(infty_dist(actual, expect), 1e-10); + } + } + } + } + delete_module_info(module); + } +} + +TEST(vec_znx, fft64_vmp_prepare_vector_ref) { + test_vmp_prepare_vector(fft64_vmp_prepare_vector_ref, fft64_vmp_prepare_vector_tmp_bytes); +} + +static void test_vmp_pvec_apply(VMP_APPLY_PVEC_TO_DFT_F* apply, VMP_APPLY_PVEC_TO_DFT_TMP_BYTES_F* tmp_bytes) { + for (uint64_t nn : {8, 64}) { // Preparation not possible for nn < 8 due to convolution not implementing it yet + MODULE* module = new_module_info(nn, FFT64); + for (uint64_t mat_nrows : {1, 4, 7}) { + for (uint64_t mat_ncols : {1, 2, 5}) { + for (uint64_t in_size : {1, 4, 7}) { + for (uint64_t out_size : {1, 2, 5}) { + znx_vec_i64_layout in(nn, in_size, nn); + fft64_vmp_pvec_layout pvec(nn, in_size); + fft64_vec_znx_dft_layout in_dft(nn, in_size); + fft64_vmp_pmat_layout pmat(nn, mat_nrows, mat_ncols); + fft64_vec_znx_dft_layout out(nn, out_size); + std::vector tmp_space(fft64_vmp_prepare_vector_tmp_bytes(module, in_size, in_size)); + + in.fill_random(0); + pmat.fill_random(0); + vec_znx_dft(module, in_dft.data, in_size, in.data(), in_size, nn); + fft64_vmp_prepare_vector_ref(module, // + pvec.data, in_size, // + in.data(), in_size, nn, tmp_space.data()); + // naive computation of the product + std::vector expect(out_size, reim_fft64vec(nn)); + for (uint64_t col = 0; col < out_size; ++col) { + reim_fft64vec ex = reim_fft64vec::zero(nn); + for (uint64_t row = 0; row < std::min(mat_nrows, in_size); ++row) { + ex += pmat.get_zext(row, col) * in_dft.get_copy_zext(row); + } + expect[col] = ex; + } + // apply the product + std::vector tmp(tmp_bytes(module, out_size, in_size, mat_nrows, mat_ncols)); + apply(module, out.data, out_size, pvec.data, in_size, pmat.data, mat_nrows, mat_ncols, tmp.data()); + // check that the output is close from the expectation + for (uint64_t col = 0; col < out_size; ++col) { + reim_fft64vec actual = out.get_copy_zext(col); + ASSERT_LE(infty_dist(actual, expect[col]), 1e-10); + } + } + } + } + } + delete_module_info(module); + } +} + +TEST(vec_znx, fft64_vmp_apply_pvec_to_dft_ref) { + test_vmp_pvec_apply(fft64_vmp_apply_pvec_to_dft_ref, fft64_vmp_apply_pvec_to_dft_tmp_bytes); +} + +#ifdef __x86_64__ +TEST(vec_znx, fft64_vmp_apply_pvec_to_dft_avx) { + test_vmp_pvec_apply(fft64_vmp_apply_pvec_to_dft_avx, fft64_vmp_apply_pvec_to_dft_tmp_bytes); +} +#endif diff --git a/test/testlib/fft64_layouts.cpp b/test/testlib/fft64_layouts.cpp index 2e42994..d6020d6 100644 --- a/test/testlib/fft64_layouts.cpp +++ b/test/testlib/fft64_layouts.cpp @@ -300,3 +300,47 @@ void fft64_cnv_right_layout::fill_random(double log2bound) { } fft64_cnv_right_layout::~fft64_cnv_right_layout() { spqlios_free(data); } + +fft64_vmp_pvec_layout::fft64_vmp_pvec_layout(uint64_t n, uint64_t size) + : nn(n), // + size(size), + data((VMP_PVEC*)alloc64(size * nn * 8)) {} + +reim4_elem fft64_vmp_pvec_layout::get(uint64_t idx, uint64_t blk) const { + REQUIRE_DRAMATICALLY(idx < size, "idx overflow: " << idx << " / " << size); + REQUIRE_DRAMATICALLY(blk < nn / 8, "block overflow: " << blk << " / " << (nn / 8)); + return reim4_elem(((double*)data) + 8 * (blk * size + idx)); +} + +void fft64_vmp_pvec_layout::set(uint64_t idx, uint64_t blk, const reim4_elem& v) { + REQUIRE_DRAMATICALLY(idx < size, "idx overflow: " << idx << " / " << size); + REQUIRE_DRAMATICALLY(blk < nn / 8, "block overflow: " << blk << " / " << (nn / 8)); + v.save_as(((double*)data) + 8 * (blk * size + idx)); +} + +reim_fft64vec fft64_vmp_pvec_layout::get_zext(uint64_t row) const { + if (row >= size) { + return reim_fft64vec::zero(nn); + } + reim_fft64vec res(nn); + for (uint64_t blk = 0; blk < nn / 8; ++blk) { + reim4_elem v = get(row, blk); + res.set_blk(blk, v); + } + return res; +} + +void fft64_vmp_pvec_layout::set(uint64_t idx, const reim_fft64vec& value) { + REQUIRE_DRAMATICALLY(idx < size, "idx overflow: " << idx << " / " << size); + for (uint64_t blk = 0; blk < nn / 8; ++blk) { + set(idx, blk, value.get_blk(blk)); + } +} + +void fft64_vmp_pvec_layout::fill_random(double log2bound) { + for (uint64_t row = 0; row < size; ++row) { + set(row, reim_fft64vec::random(nn, log2bound)); + } +} + +fft64_vmp_pvec_layout::~fft64_vmp_pvec_layout() { spqlios_free(data); } diff --git a/test/testlib/fft64_layouts.h b/test/testlib/fft64_layouts.h index fabd2e3..f6717c9 100644 --- a/test/testlib/fft64_layouts.h +++ b/test/testlib/fft64_layouts.h @@ -118,4 +118,21 @@ class fft64_cnv_right_layout { ~fft64_cnv_right_layout(); }; +/** @brief test layout for the VMP_PVEC */ +class fft64_vmp_pvec_layout { + public: + const uint64_t nn; + const uint64_t size; + VMP_PVEC* const data; + fft64_vmp_pvec_layout(uint64_t n, uint64_t size); + [[nodiscard]] reim4_elem get(uint64_t idx, uint64_t blk) const; + void set(uint64_t idx, uint64_t blk, const reim4_elem& v); + void set(uint64_t idx, const reim_fft64vec&); + [[nodiscard]] reim_fft64vec get_zext(uint64_t idx) const; + [[nodiscard]] thash content_hash() const; + /** @brief fill with random double values (unstructured) */ + void fill_random(double log2bound); + ~fft64_vmp_pvec_layout(); +}; + #endif // SPQLIOS_FFT64_LAYOUTS_H