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}