Skip to content

Add VMP prepared to DFT variant #60

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MGeorgie merged 2 commits into from
Jun 22, 2026
Merged

Add VMP prepared to DFT variant #60

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7ec93d8
Add VMP prepared to DFT variant
fraret Jun 10, 2026
083e1cd
Prepared vec for VMP renames
fraret Jun 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions spqlios/arithmetic/module_api.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
}
}

Expand Down
28 changes: 27 additions & 1 deletion spqlios/arithmetic/vec_znx_arithmetic.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 */
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
39 changes: 39 additions & 0 deletions spqlios/arithmetic/vec_znx_arithmetic_private.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
113 changes: 113 additions & 0 deletions spqlios/arithmetic/vector_matrix_product.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
Loading
Loading