poulpy/spqlios/lib/spqlios/arithmetic/zn_vmp_int32_avx.c
spqlios/lib/spqlios/arithmetic/zn_vmp_int32_avx.cRaw
1// This file is actually a template: it will be compiled multiple times with
2// different INTTYPES
3#ifndef INTTYPE
4#define INTTYPE int32_t
5#define INTSN i32
6#endif
7
8#include <immintrin.h>
9#include <memory.h>
10
11#include "zn_arithmetic_private.h"
12
13#define concat_inner(aa, bb, cc) aa##_##bb##_##cc
14#define concat(aa, bb, cc) concat_inner(aa, bb, cc)
15#define zn32_vec_fn(cc) concat(zn32_vec, INTSN, cc)
16
17static void zn32_vec_mat32cols_avx_prefetch(uint64_t nrows, int32_t* res, const INTTYPE* a, const int32_t* b) {
18 if (nrows == 0) {
19 memset(res, 0, 32 * sizeof(int32_t));
20 return;
21 }
22 const int32_t* bb = b;
23 const int32_t* pref_bb = b;
24 const uint64_t pref_iters = 128;
25 const uint64_t pref_start = pref_iters < nrows ? pref_iters : nrows;
26 const uint64_t pref_last = pref_iters > nrows ? 0 : nrows - pref_iters;
27 // let's do some prefetching of the GSW key, since on some cpus,
28 // it helps
29 for (uint64_t i = 0; i < pref_start; ++i) {
30 __builtin_prefetch(pref_bb, 0, _MM_HINT_T0);
31 __builtin_prefetch(pref_bb + 16, 0, _MM_HINT_T0);
32 pref_bb += 32;
33 }
34 // we do the first iteration
35 __m256i x = _mm256_set1_epi32(a[0]);
36 __m256i r0 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb)));
37 __m256i r1 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 8)));
38 __m256i r2 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 16)));
39 __m256i r3 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 24)));
40 bb += 32;
41 uint64_t row = 1;
42 for (; //
43 row < pref_last; //
44 ++row, bb += 32) {
45 // prefetch the next iteration
46 __builtin_prefetch(pref_bb, 0, _MM_HINT_T0);
47 __builtin_prefetch(pref_bb + 16, 0, _MM_HINT_T0);
48 pref_bb += 32;
49 INTTYPE ai = a[row];
50 if (ai == 0) continue;
51 x = _mm256_set1_epi32(ai);
52 r0 = _mm256_add_epi32(r0, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb))));
53 r1 = _mm256_add_epi32(r1, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 8))));
54 r2 = _mm256_add_epi32(r2, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 16))));
55 r3 = _mm256_add_epi32(r3, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 24))));
56 }
57 for (; //
58 row < nrows; //
59 ++row, bb += 32) {
60 INTTYPE ai = a[row];
61 if (ai == 0) continue;
62 x = _mm256_set1_epi32(ai);
63 r0 = _mm256_add_epi32(r0, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb))));
64 r1 = _mm256_add_epi32(r1, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 8))));
65 r2 = _mm256_add_epi32(r2, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 16))));
66 r3 = _mm256_add_epi32(r3, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 24))));
67 }
68 _mm256_storeu_si256((__m256i*)(res), r0);
69 _mm256_storeu_si256((__m256i*)(res + 8), r1);
70 _mm256_storeu_si256((__m256i*)(res + 16), r2);
71 _mm256_storeu_si256((__m256i*)(res + 24), r3);
72}
73
74void zn32_vec_fn(mat32cols_avx)(uint64_t nrows, int32_t* res, const INTTYPE* a, const int32_t* b, uint64_t b_sl) {
75 if (nrows == 0) {
76 memset(res, 0, 32 * sizeof(int32_t));
77 return;
78 }
79 const INTTYPE* aa = a;
80 const INTTYPE* const aaend = a + nrows;
81 const int32_t* bb = b;
82 __m256i x = _mm256_set1_epi32(*aa);
83 __m256i r0 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb)));
84 __m256i r1 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 8)));
85 __m256i r2 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 16)));
86 __m256i r3 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 24)));
87 bb += b_sl;
88 ++aa;
89 for (; //
90 aa < aaend; //
91 bb += b_sl, ++aa) {
92 INTTYPE ai = *aa;
93 if (ai == 0) continue;
94 x = _mm256_set1_epi32(ai);
95 r0 = _mm256_add_epi32(r0, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb))));
96 r1 = _mm256_add_epi32(r1, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 8))));
97 r2 = _mm256_add_epi32(r2, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 16))));
98 r3 = _mm256_add_epi32(r3, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 24))));
99 }
100 _mm256_storeu_si256((__m256i*)(res), r0);
101 _mm256_storeu_si256((__m256i*)(res + 8), r1);
102 _mm256_storeu_si256((__m256i*)(res + 16), r2);
103 _mm256_storeu_si256((__m256i*)(res + 24), r3);
104}
105
106void zn32_vec_fn(mat24cols_avx)(uint64_t nrows, int32_t* res, const INTTYPE* a, const int32_t* b, uint64_t b_sl) {
107 if (nrows == 0) {
108 memset(res, 0, 24 * sizeof(int32_t));
109 return;
110 }
111 const INTTYPE* aa = a;
112 const INTTYPE* const aaend = a + nrows;
113 const int32_t* bb = b;
114 __m256i x = _mm256_set1_epi32(*aa);
115 __m256i r0 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb)));
116 __m256i r1 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 8)));
117 __m256i r2 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 16)));
118 bb += b_sl;
119 ++aa;
120 for (; //
121 aa < aaend; //
122 bb += b_sl, ++aa) {
123 INTTYPE ai = *aa;
124 if (ai == 0) continue;
125 x = _mm256_set1_epi32(ai);
126 r0 = _mm256_add_epi32(r0, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb))));
127 r1 = _mm256_add_epi32(r1, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 8))));
128 r2 = _mm256_add_epi32(r2, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 16))));
129 }
130 _mm256_storeu_si256((__m256i*)(res), r0);
131 _mm256_storeu_si256((__m256i*)(res + 8), r1);
132 _mm256_storeu_si256((__m256i*)(res + 16), r2);
133}
134void zn32_vec_fn(mat16cols_avx)(uint64_t nrows, int32_t* res, const INTTYPE* a, const int32_t* b, uint64_t b_sl) {
135 if (nrows == 0) {
136 memset(res, 0, 16 * sizeof(int32_t));
137 return;
138 }
139 const INTTYPE* aa = a;
140 const INTTYPE* const aaend = a + nrows;
141 const int32_t* bb = b;
142 __m256i x = _mm256_set1_epi32(*aa);
143 __m256i r0 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb)));
144 __m256i r1 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 8)));
145 bb += b_sl;
146 ++aa;
147 for (; //
148 aa < aaend; //
149 bb += b_sl, ++aa) {
150 INTTYPE ai = *aa;
151 if (ai == 0) continue;
152 x = _mm256_set1_epi32(ai);
153 r0 = _mm256_add_epi32(r0, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb))));
154 r1 = _mm256_add_epi32(r1, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb + 8))));
155 }
156 _mm256_storeu_si256((__m256i*)(res), r0);
157 _mm256_storeu_si256((__m256i*)(res + 8), r1);
158}
159
160void zn32_vec_fn(mat8cols_avx)(uint64_t nrows, int32_t* res, const INTTYPE* a, const int32_t* b, uint64_t b_sl) {
161 if (nrows == 0) {
162 memset(res, 0, 8 * sizeof(int32_t));
163 return;
164 }
165 const INTTYPE* aa = a;
166 const INTTYPE* const aaend = a + nrows;
167 const int32_t* bb = b;
168 __m256i x = _mm256_set1_epi32(*aa);
169 __m256i r0 = _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb)));
170 bb += b_sl;
171 ++aa;
172 for (; //
173 aa < aaend; //
174 bb += b_sl, ++aa) {
175 INTTYPE ai = *aa;
176 if (ai == 0) continue;
177 x = _mm256_set1_epi32(ai);
178 r0 = _mm256_add_epi32(r0, _mm256_mullo_epi32(x, _mm256_loadu_si256((__m256i*)(bb))));
179 }
180 _mm256_storeu_si256((__m256i*)(res), r0);
181}
182
183typedef void (*vm_f)(uint64_t nrows, //
184 int32_t* res, //
185 const INTTYPE* a, //
186 const int32_t* b, uint64_t b_sl //
187);
188static const vm_f zn32_vec_mat8kcols_avx[4] = { //
189 zn32_vec_fn(mat8cols_avx), //
190 zn32_vec_fn(mat16cols_avx), //
191 zn32_vec_fn(mat24cols_avx), //
192 zn32_vec_fn(mat32cols_avx)};
193
194/** @brief applies a vmp product (int32_t* input) */
195EXPORT void concat(default_zn32_vmp_apply, INTSN, avx)( //
196 const MOD_Z* module, //
197 int32_t* res, uint64_t res_size, //
198 const INTTYPE* a, uint64_t a_size, //
199 const ZN32_VMP_PMAT* pmat, uint64_t nrows, uint64_t ncols) {
200 const uint64_t rows = a_size < nrows ? a_size : nrows;
201 const uint64_t cols = res_size < ncols ? res_size : ncols;
202 const uint64_t ncolblk = cols >> 5;
203 const uint64_t ncolrem = cols & 31;
204 // copy the first full blocks
205 const uint64_t full_blk_size = nrows * 32;
206 const int32_t* mat = (int32_t*)pmat;
207 int32_t* rr = res;
208 for (uint64_t blk = 0; //
209 blk < ncolblk; //
210 ++blk, mat += full_blk_size, rr += 32) {
211 zn32_vec_mat32cols_avx_prefetch(rows, rr, a, mat);
212 }
213 // last block
214 if (ncolrem) {
215 uint64_t orig_rem = ncols - (ncolblk << 5);
216 uint64_t b_sl = orig_rem >= 32 ? 32 : orig_rem;
217 int32_t tmp[32];
218 zn32_vec_mat8kcols_avx[(ncolrem - 1) >> 3](rows, tmp, a, mat, b_sl);
219 memcpy(rr, tmp, ncolrem * sizeof(int32_t));
220 }
221 // trailing bytes
222 memset(res + cols, 0, (res_size - cols) * sizeof(int32_t));
223}