From 7ec93d8b690e8b15992b12b7fadd4a7b71eb4ba0 Mon Sep 17 00:00:00 2001 From: Sergi Soler Arrufat Date: Wed, 10 Jun 2026 10:55:58 +0200 Subject: [PATCH 1/2] Add VMP prepared to DFT variant Intended to be a more performant alternative to vmp_apply_dft_to_dft when one can prepare the input vector (eg. PIR). The prepared vec is in reim4 dft vs the reim dft it would usually be, which can result in significant performance gains in certain memory-bound vmp workflows. --- spqlios/arithmetic/module_api.c | 5 + spqlios/arithmetic/vec_znx_arithmetic.h | 28 ++++- .../arithmetic/vec_znx_arithmetic_private.h | 39 ++++++ spqlios/arithmetic/vector_matrix_product.c | 113 ++++++++++++++++++ .../arithmetic/vector_matrix_product_avx.c | 55 +++++++++ test/spqlios_vmp_product_test.cpp | 90 ++++++++++++++ test/testlib/fft64_layouts.cpp | 44 +++++++ test/testlib/fft64_layouts.h | 17 +++ 8 files changed, 390 insertions(+), 1 deletion(-) diff --git a/spqlios/arithmetic/module_api.c b/spqlios/arithmetic/module_api.c index 05631ec..88e5f2c 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_contiguous_vec = fft64_vmp_prepare_contiguous_vec_ref; module->func.vmp_prepare_contiguous_tmp_bytes = fft64_vmp_prepare_contiguous_tmp_bytes; + module->func.vmp_prepare_contiguous_vec_tmp_bytes = fft64_vmp_prepare_contiguous_vec_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_prepared_to_dft = fft64_vmp_apply_prepared_to_dft_ref; module->func.vmp_apply_dft_to_dft_tmp_bytes = fft64_vmp_apply_dft_to_dft_tmp_bytes; + module->func.vmp_apply_prepared_to_dft_tmp_bytes = fft64_vmp_apply_prepared_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_prepared_to_dft = fft64_vmp_apply_prepared_to_dft_avx; } } diff --git a/spqlios/arithmetic/vec_znx_arithmetic.h b/spqlios/arithmetic/vec_znx_arithmetic.h index fde7df0..02aacce 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 vec_znx_dft_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_contiguous_vec(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_contiguous_vec_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_prepared_to_dft(const MODULE* module, // N + VEC_ZNX_DFT* res, const uint64_t res_size, // res + const VEC_ZNX_DFT* 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_prepared_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..24f5d59 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_contiguous_vec) VMP_PREPARE_CONTIGUOUS_VEC_F; typedef typeof(vmp_prepare_contiguous_tmp_bytes) VMP_PREPARE_CONTIGUOUS_TMP_BYTES_F; +typedef typeof(vmp_prepare_contiguous_vec_tmp_bytes) VMP_PREPARE_CONTIGUOUS_VEC_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_prepared_to_dft) VMP_APPLY_PREPARED_TO_DFT_F; typedef typeof(vmp_apply_dft_to_dft_tmp_bytes) VMP_APPLY_DFT_TO_DFT_TMP_BYTES_F; +typedef typeof(vmp_apply_prepared_to_dft_tmp_bytes) VMP_APPLY_PREPARED_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_CONTIGUOUS_VEC_F* vmp_prepare_contiguous_vec; VMP_PREPARE_CONTIGUOUS_TMP_BYTES_F* vmp_prepare_contiguous_tmp_bytes; + VMP_PREPARE_CONTIGUOUS_VEC_TMP_BYTES_F* vmp_prepare_contiguous_vec_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_PREPARED_TO_DFT_F* vmp_apply_prepared_to_dft; VMP_APPLY_DFT_TO_DFT_TMP_BYTES_F* vmp_apply_dft_to_dft_tmp_bytes; + VMP_APPLY_PREPARED_TO_DFT_TMP_BYTES_F* vmp_apply_prepared_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_contiguous_vec_tmp_bytes(const MODULE* module, // N + uint64_t nrows, uint64_t a_size); + +EXPORT void fft64_vmp_prepare_contiguous_vec_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_prepared_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_prepared_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_prepared_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..f0cb618 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_contiguous_vec(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_contiguous_vec(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_contiguous_vec_tmp_bytes(const MODULE* module, // N + uint64_t nrows, uint64_t a_size) { + return module->func.vmp_prepare_contiguous_vec_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_contiguous_vec_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_contiguous_vec_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_prepared_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_prepared_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_prepared_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_prepared_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_prepared_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_dft_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..7849a57 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_prepared_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..5a47c3c 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_contiguous_vec(VMP_PREPARE_CONTIGUOUS_VEC_F* prepare_contiguous, + VMP_PREPARE_CONTIGUOUS_VEC_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_contiguous_vec_ref) { + test_vmp_prepare_contiguous_vec(fft64_vmp_prepare_contiguous_vec_ref, fft64_vmp_prepare_contiguous_vec_tmp_bytes); +} + +static void test_vmp_pvec_apply(VMP_APPLY_PREPARED_TO_DFT_F* apply, VMP_APPLY_PREPARED_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_contiguous_vec_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_contiguous_vec_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_prepared_to_dft_ref) { + test_vmp_pvec_apply(fft64_vmp_apply_prepared_to_dft_ref, fft64_vmp_apply_prepared_to_dft_tmp_bytes); +} + +#ifdef __x86_64__ +TEST(vec_znx, fft64_vmp_apply_prepared_to_dft_avx) { + test_vmp_pvec_apply(fft64_vmp_apply_prepared_to_dft_avx, fft64_vmp_apply_prepared_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 From 083e1cd2fe62b476def6047d18386641fd5aab95 Mon Sep 17 00:00:00 2001 From: Sergi Soler Arrufat Date: Mon, 22 Jun 2026 10:52:03 +0200 Subject: [PATCH 2/2] Prepared vec for VMP renames --- spqlios/arithmetic/module_api.c | 10 +-- spqlios/arithmetic/vec_znx_arithmetic.h | 34 ++++----- .../arithmetic/vec_znx_arithmetic_private.h | 60 ++++++++-------- spqlios/arithmetic/vector_matrix_product.c | 72 +++++++++---------- .../arithmetic/vector_matrix_product_avx.c | 12 ++-- test/spqlios_vmp_product_test.cpp | 26 +++---- 6 files changed, 107 insertions(+), 107 deletions(-) diff --git a/spqlios/arithmetic/module_api.c b/spqlios/arithmetic/module_api.c index 88e5f2c..8303ac9 100644 --- a/spqlios/arithmetic/module_api.c +++ b/spqlios/arithmetic/module_api.c @@ -47,15 +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_contiguous_vec = fft64_vmp_prepare_contiguous_vec_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_contiguous_vec_tmp_bytes = fft64_vmp_prepare_contiguous_vec_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_prepared_to_dft = fft64_vmp_apply_prepared_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_prepared_to_dft_tmp_bytes = fft64_vmp_apply_prepared_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; @@ -74,7 +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_prepared_to_dft = fft64_vmp_apply_prepared_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 02aacce..9b6a2eb 100644 --- a/spqlios/arithmetic/vec_znx_arithmetic.h +++ b/spqlios/arithmetic/vec_znx_arithmetic.h @@ -26,7 +26,7 @@ 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 vec_znx_dft_t VMP_PVEC; +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 */ @@ -153,10 +153,10 @@ EXPORT void vmp_prepare_contiguous(const MODULE* module, ); /** @brief prepares a vmp vector */ -EXPORT void vmp_prepare_contiguous_vec(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 +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) */ @@ -340,8 +340,8 @@ 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_contiguous_vec_tmp_bytes(const MODULE* module, // N - uint64_t nrows, uint64_t a_size); +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 @@ -367,12 +367,12 @@ EXPORT void vmp_apply_dft_to_dft(const MODULE* module, // uint8_t* tmp_space // scratch space (a_size*sizeof(reim4) bytes) ); -EXPORT void vmp_apply_prepared_to_dft(const MODULE* module, // N - VEC_ZNX_DFT* res, const uint64_t res_size, // res - const VEC_ZNX_DFT* 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) +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 */ @@ -382,10 +382,10 @@ 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_prepared_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 +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 */ diff --git a/spqlios/arithmetic/vec_znx_arithmetic_private.h b/spqlios/arithmetic/vec_znx_arithmetic_private.h index 24f5d59..25e62ef 100644 --- a/spqlios/arithmetic/vec_znx_arithmetic_private.h +++ b/spqlios/arithmetic/vec_znx_arithmetic_private.h @@ -85,9 +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_contiguous_vec) VMP_PREPARE_CONTIGUOUS_VEC_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_contiguous_vec_tmp_bytes) VMP_PREPARE_CONTIGUOUS_VEC_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; @@ -97,9 +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_prepared_to_dft) VMP_APPLY_PREPARED_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_prepared_to_dft_tmp_bytes) VMP_APPLY_PREPARED_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; @@ -139,15 +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_CONTIGUOUS_VEC_F* vmp_prepare_contiguous_vec; + VMP_PREPARE_VECTOR_F* vmp_prepare_vector; VMP_PREPARE_CONTIGUOUS_TMP_BYTES_F* vmp_prepare_contiguous_tmp_bytes; - VMP_PREPARE_CONTIGUOUS_VEC_TMP_BYTES_F* vmp_prepare_contiguous_vec_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_PREPARED_TO_DFT_F* vmp_apply_prepared_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_PREPARED_TO_DFT_TMP_BYTES_F* vmp_apply_prepared_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; @@ -509,13 +509,13 @@ 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_contiguous_vec_tmp_bytes(const MODULE* module, // N - uint64_t nrows, uint64_t a_size); +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_contiguous_vec_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 +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) */ @@ -552,20 +552,20 @@ 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_prepared_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_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_prepared_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) +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 */ @@ -582,9 +582,9 @@ EXPORT uint64_t fft64_vmp_apply_dft_to_dft_tmp_bytes(const MODULE* module, uint64_t nrows, uint64_t ncols // prep matrix ); -EXPORT uint64_t fft64_vmp_apply_prepared_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 +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 f0cb618..a88e405 100644 --- a/spqlios/arithmetic/vector_matrix_product.c +++ b/spqlios/arithmetic/vector_matrix_product.c @@ -34,12 +34,12 @@ EXPORT void vmp_prepare_contiguous(const MODULE* module, } /** @brief prepares a vmp vector */ -EXPORT void vmp_prepare_contiguous_vec(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 +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_contiguous_vec(module, pvec, nrows, a, a_size, a_sl, tmp_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 */ @@ -49,9 +49,9 @@ EXPORT uint64_t vmp_prepare_contiguous_tmp_bytes(const MODULE* module, // N } /** @brief minimal scratch space byte-size required for the vmp_prepare_vec function */ -EXPORT uint64_t vmp_prepare_contiguous_vec_tmp_bytes(const MODULE* module, // N - uint64_t nrows, uint64_t a_size) { - return module->func.vmp_prepare_contiguous_vec_tmp_bytes(module, nrows, a_size); +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) */ @@ -110,18 +110,18 @@ EXPORT uint64_t fft64_vmp_prepare_contiguous_tmp_bytes(const MODULE* module, // } /** @brief prepares a vmp vector */ -EXPORT void fft64_vmp_prepare_contiguous_vec_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 +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_contiguous_vec_tmp_bytes(const MODULE* module, // N - uint64_t nrows, uint64_t a_size) { +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); } @@ -211,12 +211,12 @@ 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_prepared_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) +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; @@ -289,10 +289,10 @@ EXPORT uint64_t fft64_vmp_apply_dft_to_dft_tmp_bytes(const MODULE* module, } /** @brief minimal size of the tmp_space */ -EXPORT uint64_t fft64_vmp_apply_prepared_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 +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); } @@ -307,14 +307,14 @@ 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_prepared_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) +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_prepared_to_dft(module, res, res_size, a_prep, a_size, pmat, nrows, ncols, tmp_space); + 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 @@ -325,12 +325,12 @@ 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_prepared_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 +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_dft_to_dft_tmp_bytes(module, res_size, a_size, nrows, ncols); + 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) */ diff --git a/spqlios/arithmetic/vector_matrix_product_avx.c b/spqlios/arithmetic/vector_matrix_product_avx.c index 7849a57..6ed44c1 100644 --- a/spqlios/arithmetic/vector_matrix_product_avx.c +++ b/spqlios/arithmetic/vector_matrix_product_avx.c @@ -136,12 +136,12 @@ EXPORT void fft64_vmp_apply_dft_to_dft_avx(const MODULE* module, memset(vec_output + col_max * nn, 0, (res_size - col_max) * nn * sizeof(double)); } -EXPORT void fft64_vmp_apply_prepared_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) +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; diff --git a/test/spqlios_vmp_product_test.cpp b/test/spqlios_vmp_product_test.cpp index 5a47c3c..2268a4a 100644 --- a/test/spqlios_vmp_product_test.cpp +++ b/test/spqlios_vmp_product_test.cpp @@ -120,8 +120,8 @@ TEST(vec_znx, fft64_vmp_apply_dft_to_dft_avx) { } #endif -static void test_vmp_prepare_contiguous_vec(VMP_PREPARE_CONTIGUOUS_VEC_F* prepare_contiguous, - VMP_PREPARE_CONTIGUOUS_VEC_TMP_BYTES_F* tmp_bytes) { +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); @@ -151,11 +151,11 @@ static void test_vmp_prepare_contiguous_vec(VMP_PREPARE_CONTIGUOUS_VEC_F* prepar } } -TEST(vec_znx, fft64_vmp_prepare_contiguous_vec_ref) { - test_vmp_prepare_contiguous_vec(fft64_vmp_prepare_contiguous_vec_ref, fft64_vmp_prepare_contiguous_vec_tmp_bytes); +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_PREPARED_TO_DFT_F* apply, VMP_APPLY_PREPARED_TO_DFT_TMP_BYTES_F* 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}) { @@ -167,14 +167,14 @@ static void test_vmp_pvec_apply(VMP_APPLY_PREPARED_TO_DFT_F* apply, VMP_APPLY_PR 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_contiguous_vec_tmp_bytes(module, in_size, in_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_contiguous_vec_ref(module, // - pvec.data, in_size, // - in.data(), in_size, nn, tmp_space.data()); + 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) { @@ -200,12 +200,12 @@ static void test_vmp_pvec_apply(VMP_APPLY_PREPARED_TO_DFT_F* apply, VMP_APPLY_PR } } -TEST(vec_znx, fft64_vmp_apply_prepared_to_dft_ref) { - test_vmp_pvec_apply(fft64_vmp_apply_prepared_to_dft_ref, fft64_vmp_apply_prepared_to_dft_tmp_bytes); +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_prepared_to_dft_avx) { - test_vmp_pvec_apply(fft64_vmp_apply_prepared_to_dft_avx, fft64_vmp_apply_prepared_to_dft_tmp_bytes); +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