arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/vec_znx_avx.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/vec_znx_avx.cRaw
  1#include <string.h>
  2
  3#include "../coeffs/coeffs_arithmetic.h"
  4#include "../reim4/reim4_arithmetic.h"
  5#include "vec_znx_arithmetic_private.h"
  6
  7// specialized function (ref)
  8
  9// Note: these functions do not have an avx variant.
 10#define znx_copy_i64_avx znx_copy_i64_ref
 11#define znx_zero_i64_avx znx_zero_i64_ref
 12
 13EXPORT void vec_znx_add_avx(const MODULE* module,                              // N
 14                            int64_t* res, uint64_t res_size, uint64_t res_sl,  // res
 15                            const int64_t* a, uint64_t a_size, uint64_t a_sl,  // a
 16                            const int64_t* b, uint64_t b_size, uint64_t b_sl   // b
 17) {
 18  const uint64_t nn = module->nn;
 19  if (a_size <= b_size) {
 20    const uint64_t sum_idx = res_size < a_size ? res_size : a_size;
 21    const uint64_t copy_idx = res_size < b_size ? res_size : b_size;
 22    // add up to the smallest dimension
 23    for (uint64_t i = 0; i < sum_idx; ++i) {
 24      znx_add_i64_avx(nn, res + i * res_sl, a + i * a_sl, b + i * b_sl);
 25    }
 26    // then copy to the largest dimension
 27    for (uint64_t i = sum_idx; i < copy_idx; ++i) {
 28      znx_copy_i64_avx(nn, res + i * res_sl, b + i * b_sl);
 29    }
 30    // then extend with zeros
 31    for (uint64_t i = copy_idx; i < res_size; ++i) {
 32      znx_zero_i64_avx(nn, res + i * res_sl);
 33    }
 34  } else {
 35    const uint64_t sum_idx = res_size < b_size ? res_size : b_size;
 36    const uint64_t copy_idx = res_size < a_size ? res_size : a_size;
 37    // add up to the smallest dimension
 38    for (uint64_t i = 0; i < sum_idx; ++i) {
 39      znx_add_i64_avx(nn, res + i * res_sl, a + i * a_sl, b + i * b_sl);
 40    }
 41    // then copy to the largest dimension
 42    for (uint64_t i = sum_idx; i < copy_idx; ++i) {
 43      znx_copy_i64_avx(nn, res + i * res_sl, a + i * a_sl);
 44    }
 45    // then extend with zeros
 46    for (uint64_t i = copy_idx; i < res_size; ++i) {
 47      znx_zero_i64_avx(nn, res + i * res_sl);
 48    }
 49  }
 50}
 51
 52EXPORT void vec_znx_sub_avx(const MODULE* module,                              // N
 53                            int64_t* res, uint64_t res_size, uint64_t res_sl,  // res
 54                            const int64_t* a, uint64_t a_size, uint64_t a_sl,  // a
 55                            const int64_t* b, uint64_t b_size, uint64_t b_sl   // b
 56) {
 57  const uint64_t nn = module->nn;
 58  if (a_size <= b_size) {
 59    const uint64_t sub_idx = res_size < a_size ? res_size : a_size;
 60    const uint64_t copy_idx = res_size < b_size ? res_size : b_size;
 61    // subtract up to the smallest dimension
 62    for (uint64_t i = 0; i < sub_idx; ++i) {
 63      znx_sub_i64_avx(nn, res + i * res_sl, a + i * a_sl, b + i * b_sl);
 64    }
 65    // then negate to the largest dimension
 66    for (uint64_t i = sub_idx; i < copy_idx; ++i) {
 67      znx_negate_i64_avx(nn, res + i * res_sl, b + i * b_sl);
 68    }
 69    // then extend with zeros
 70    for (uint64_t i = copy_idx; i < res_size; ++i) {
 71      znx_zero_i64_avx(nn, res + i * res_sl);
 72    }
 73  } else {
 74    const uint64_t sub_idx = res_size < b_size ? res_size : b_size;
 75    const uint64_t copy_idx = res_size < a_size ? res_size : a_size;
 76    // subtract up to the smallest dimension
 77    for (uint64_t i = 0; i < sub_idx; ++i) {
 78      znx_sub_i64_avx(nn, res + i * res_sl, a + i * a_sl, b + i * b_sl);
 79    }
 80    // then copy to the largest dimension
 81    for (uint64_t i = sub_idx; i < copy_idx; ++i) {
 82      znx_copy_i64_avx(nn, res + i * res_sl, a + i * a_sl);
 83    }
 84    // then extend with zeros
 85    for (uint64_t i = copy_idx; i < res_size; ++i) {
 86      znx_zero_i64_avx(nn, res + i * res_sl);
 87    }
 88  }
 89}
 90
 91EXPORT void vec_znx_negate_avx(const MODULE* module,                              // N
 92                               int64_t* res, uint64_t res_size, uint64_t res_sl,  // res
 93                               const int64_t* a, uint64_t a_size, uint64_t a_sl   // a
 94) {
 95  uint64_t nn = module->nn;
 96  uint64_t smin = res_size < a_size ? res_size : a_size;
 97  for (uint64_t i = 0; i < smin; ++i) {
 98    znx_negate_i64_avx(nn, res + i * res_sl, a + i * a_sl);
 99  }
100  for (uint64_t i = smin; i < res_size; ++i) {
101    znx_zero_i64_ref(nn, res + i * res_sl);
102  }
103}