arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/module_api.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/module_api.cRaw
  1#include <string.h>
  2
  3#include "vec_znx_arithmetic_private.h"
  4
  5static void fill_generic_virtual_table(MODULE* module) {
  6  // TODO add default ref handler here
  7  module->func.vec_znx_zero = vec_znx_zero_ref;
  8  module->func.vec_znx_copy = vec_znx_copy_ref;
  9  module->func.vec_znx_negate = vec_znx_negate_ref;
 10  module->func.vec_znx_add = vec_znx_add_ref;
 11  module->func.vec_znx_sub = vec_znx_sub_ref;
 12  module->func.vec_znx_rotate = vec_znx_rotate_ref;
 13  module->func.vec_znx_mul_xp_minus_one = vec_znx_mul_xp_minus_one_ref;
 14  module->func.vec_znx_automorphism = vec_znx_automorphism_ref;
 15  module->func.vec_znx_normalize_base2k = vec_znx_normalize_base2k_ref;
 16  module->func.vec_znx_normalize_base2k_tmp_bytes = vec_znx_normalize_base2k_tmp_bytes_ref;
 17  if (CPU_SUPPORTS("avx2")) {
 18    // TODO add avx handlers here
 19    module->func.vec_znx_negate = vec_znx_negate_avx;
 20    module->func.vec_znx_add = vec_znx_add_avx;
 21    module->func.vec_znx_sub = vec_znx_sub_avx;
 22  }
 23}
 24
 25static void fill_fft64_virtual_table(MODULE* module) {
 26  // TODO add default ref handler here
 27  // module->func.vec_znx_dft = ...;
 28  module->func.vec_znx_big_normalize_base2k = fft64_vec_znx_big_normalize_base2k;
 29  module->func.vec_znx_big_normalize_base2k_tmp_bytes = fft64_vec_znx_big_normalize_base2k_tmp_bytes;
 30  module->func.vec_znx_big_range_normalize_base2k = fft64_vec_znx_big_range_normalize_base2k;
 31  module->func.vec_znx_big_range_normalize_base2k_tmp_bytes = fft64_vec_znx_big_range_normalize_base2k_tmp_bytes;
 32  module->func.vec_znx_dft = fft64_vec_znx_dft;
 33  module->func.vec_znx_idft = fft64_vec_znx_idft;
 34  module->func.vec_dft_add = fft64_vec_dft_add;
 35  module->func.vec_dft_sub = fft64_vec_dft_sub;
 36  module->func.vec_znx_idft_tmp_bytes = fft64_vec_znx_idft_tmp_bytes;
 37  module->func.vec_znx_idft_tmp_a = fft64_vec_znx_idft_tmp_a;
 38  module->func.vec_znx_big_add = fft64_vec_znx_big_add;
 39  module->func.vec_znx_big_add_small = fft64_vec_znx_big_add_small;
 40  module->func.vec_znx_big_add_small2 = fft64_vec_znx_big_add_small2;
 41  module->func.vec_znx_big_sub = fft64_vec_znx_big_sub;
 42  module->func.vec_znx_big_sub_small_a = fft64_vec_znx_big_sub_small_a;
 43  module->func.vec_znx_big_sub_small_b = fft64_vec_znx_big_sub_small_b;
 44  module->func.vec_znx_big_sub_small2 = fft64_vec_znx_big_sub_small2;
 45  module->func.vec_znx_big_rotate = fft64_vec_znx_big_rotate;
 46  module->func.vec_znx_big_automorphism = fft64_vec_znx_big_automorphism;
 47  module->func.svp_prepare = fft64_svp_prepare_ref;
 48  module->func.svp_apply_dft = fft64_svp_apply_dft_ref;
 49  module->func.svp_apply_dft_to_dft = fft64_svp_apply_dft_to_dft_ref;
 50  module->func.znx_small_single_product = fft64_znx_small_single_product;
 51  module->func.znx_small_single_product_tmp_bytes = fft64_znx_small_single_product_tmp_bytes;
 52  module->func.vmp_prepare_contiguous = fft64_vmp_prepare_contiguous_ref;
 53  module->func.vmp_prepare_tmp_bytes = fft64_vmp_prepare_tmp_bytes;
 54  module->func.vmp_apply_dft = fft64_vmp_apply_dft_ref;
 55  module->func.vmp_apply_dft_add = fft64_vmp_apply_dft_add_ref;
 56  module->func.vmp_apply_dft_tmp_bytes = fft64_vmp_apply_dft_tmp_bytes;
 57  module->func.vmp_apply_dft_to_dft = fft64_vmp_apply_dft_to_dft_ref;
 58  module->func.vmp_apply_dft_to_dft_add = fft64_vmp_apply_dft_to_dft_add_ref;
 59  module->func.vmp_apply_dft_to_dft_tmp_bytes = fft64_vmp_apply_dft_to_dft_tmp_bytes;
 60  module->func.bytes_of_vec_znx_dft = fft64_bytes_of_vec_znx_dft;
 61  module->func.bytes_of_vec_znx_big = fft64_bytes_of_vec_znx_big;
 62  module->func.bytes_of_svp_ppol = fft64_bytes_of_svp_ppol;
 63  module->func.bytes_of_vmp_pmat = fft64_bytes_of_vmp_pmat;
 64  if (CPU_SUPPORTS("avx2")) {
 65    // TODO add avx handlers here
 66    // TODO: enable when avx implementation is done
 67    module->func.vmp_prepare_contiguous = fft64_vmp_prepare_contiguous_avx;
 68    module->func.vmp_apply_dft = fft64_vmp_apply_dft_avx;
 69    module->func.vmp_apply_dft_add = fft64_vmp_apply_dft_add_avx;
 70    module->func.vmp_apply_dft_to_dft = fft64_vmp_apply_dft_to_dft_avx;
 71    module->func.vmp_apply_dft_to_dft_add = fft64_vmp_apply_dft_to_dft_add_avx;
 72  }
 73}
 74
 75static void fill_ntt120_virtual_table(MODULE* module) {
 76  // TODO add default ref handler here
 77  // module->func.vec_znx_dft = ...;
 78  if (CPU_SUPPORTS("avx2")) {
 79    // TODO add avx handlers here
 80    module->func.vec_znx_dft = ntt120_vec_znx_dft_avx;
 81    module->func.vec_znx_idft = ntt120_vec_znx_idft_avx;
 82    module->func.vec_znx_idft_tmp_bytes = ntt120_vec_znx_idft_tmp_bytes_avx;
 83    module->func.vec_znx_idft_tmp_a = ntt120_vec_znx_idft_tmp_a_avx;
 84  }
 85}
 86
 87static void fill_virtual_table(MODULE* module) {
 88  fill_generic_virtual_table(module);
 89  switch (module->module_type) {
 90    case FFT64:
 91      fill_fft64_virtual_table(module);
 92      break;
 93    case NTT120:
 94      fill_ntt120_virtual_table(module);
 95      break;
 96    default:
 97      NOT_SUPPORTED();  // invalid type
 98  }
 99}
100
101static void fill_fft64_precomp(MODULE* module) {
102  // fill any necessary precomp stuff
103  module->mod.fft64.p_conv = new_reim_from_znx64_precomp(module->m, 50);
104  module->mod.fft64.p_fft = new_reim_fft_precomp(module->m, 0);
105  module->mod.fft64.p_reim_to_znx = new_reim_to_znx64_precomp(module->m, module->m, 63);
106  module->mod.fft64.p_ifft = new_reim_ifft_precomp(module->m, 0);
107  module->mod.fft64.p_addmul = new_reim_fftvec_addmul_precomp(module->m);
108  module->mod.fft64.mul_fft = new_reim_fftvec_mul_precomp(module->m);
109  module->mod.fft64.add_fft = new_reim_fftvec_add_precomp(module->m);
110  module->mod.fft64.sub_fft = new_reim_fftvec_sub_precomp(module->m);
111}
112static void fill_ntt120_precomp(MODULE* module) {
113  // fill any necessary precomp stuff
114  if (CPU_SUPPORTS("avx2")) {
115    module->mod.q120.p_ntt = q120_new_ntt_bb_precomp(module->nn);
116    module->mod.q120.p_intt = q120_new_intt_bb_precomp(module->nn);
117  }
118}
119
120static void fill_module_precomp(MODULE* module) {
121  switch (module->module_type) {
122    case FFT64:
123      fill_fft64_precomp(module);
124      break;
125    case NTT120:
126      fill_ntt120_precomp(module);
127      break;
128    default:
129      NOT_SUPPORTED();  // invalid type
130  }
131}
132
133static void fill_module(MODULE* module, uint64_t nn, MODULE_TYPE mtype) {
134  // init to zero to ensure that any non-initialized field bug is detected
135  // by at least a "proper" segfault
136  memset(module, 0, sizeof(MODULE));
137  module->module_type = mtype;
138  module->nn = nn;
139  module->m = nn >> 1;
140  fill_module_precomp(module);
141  fill_virtual_table(module);
142}
143
144EXPORT MODULE* new_module_info(uint64_t N, MODULE_TYPE mtype) {
145  MODULE* m = (MODULE*)malloc(sizeof(MODULE));
146  fill_module(m, N, mtype);
147  return m;
148}
149
150EXPORT void delete_module_info(MODULE* mod) {
151  switch (mod->module_type) {
152    case FFT64:
153      free(mod->mod.fft64.p_conv);
154      free(mod->mod.fft64.p_fft);
155      free(mod->mod.fft64.p_ifft);
156      free(mod->mod.fft64.p_reim_to_znx);
157      free(mod->mod.fft64.mul_fft);
158      free(mod->mod.fft64.p_addmul);
159      break;
160    case NTT120:
161      if (CPU_SUPPORTS("avx2")) {
162        q120_del_ntt_bb_precomp(mod->mod.q120.p_ntt);
163        q120_del_intt_bb_precomp(mod->mod.q120.p_intt);
164      }
165      break;
166    default:
167      break;
168  }
169  free(mod);
170}
171
172EXPORT uint64_t module_get_n(const MODULE* module) { return module->nn; }