poulpy/poulpy-cpu-avx/src/znx_avx/normalization.rs
poulpy-cpu-avx/src/znx_avx/normalization.rsRaw
1use std::arch::x86_64::__m256i;
2
3/// Vector forms of those constants (broadcast to all lanes)
4///
5/// # Safety
6/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
7/// all inputs must have the same length and must not alias.
8#[target_feature(enable = "avx2")]
9fn normalize_consts_avx(base2k: usize) -> (__m256i, __m256i, __m256i, __m256i) {
10 use std::arch::x86_64::_mm256_set1_epi64x;
11
12 assert!((1..=63).contains(&base2k));
13 let mask_k: i64 = ((1u64 << base2k) - 1) as i64; // 0..k-1 bits set
14 let sign_k: i64 = (1u64 << (base2k - 1)) as i64; // bit k-1
15 let topmask: i64 = (!0u64 << (64 - base2k)) as i64; // top k bits set
16 let sh_k: __m256i = _mm256_set1_epi64x(base2k as i64);
17 (
18 _mm256_set1_epi64x(mask_k), // mask_k_vec
19 _mm256_set1_epi64x(sign_k), // sign_k_vec
20 sh_k, // shift_k_vec
21 _mm256_set1_epi64x(topmask), // topmask_vec
22 )
23}
24
25/// AVX2 get_digit using masks (no arithmetic shift needed):
26/// digit = ((x & mask_k) ^ sign_k) - sign_k
27///
28/// # Safety
29/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
30/// all inputs must have the same length and must not alias.
31#[target_feature(enable = "avx2")]
32fn get_digit_avx(x: __m256i, mask_k: __m256i, sign_k: __m256i) -> __m256i {
33 use std::arch::x86_64::{_mm256_and_si256, _mm256_sub_epi64, _mm256_xor_si256};
34 let low: __m256i = _mm256_and_si256(x, mask_k);
35 let t: __m256i = _mm256_xor_si256(low, sign_k);
36 _mm256_sub_epi64(t, sign_k)
37}
38
39/// AVX2 get_carry using precomputed shift and topmask:
40/// carry = (x - digit) >>_arith k
41///
42/// # Safety
43/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
44/// all inputs must have the same length and must not alias.
45#[target_feature(enable = "avx2")]
46unsafe fn get_carry_avx(
47 x: __m256i,
48 digit: __m256i,
49 base2k: __m256i, // _mm256_set1_epi64x(k)
50 top_mask: __m256i, // (!0 << (64 - k)) broadcast
51) -> __m256i {
52 use std::arch::x86_64::{
53 __m256i, _mm256_and_si256, _mm256_cmpgt_epi64, _mm256_or_si256, _mm256_setzero_si256, _mm256_srlv_epi64, _mm256_sub_epi64,
54 };
55 let diff: __m256i = _mm256_sub_epi64(x, digit);
56 let lsr: __m256i = _mm256_srlv_epi64(diff, base2k); // logical >>
57 let neg: __m256i = _mm256_cmpgt_epi64(_mm256_setzero_si256(), diff); // 0xFFFF.. where v<0
58 let fill: __m256i = _mm256_and_si256(neg, top_mask); // top k bits if negative
59 _mm256_or_si256(lsr, fill)
60}
61
62/// # Safety
63/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
64/// `res` and `src` must have the same length and must not alias.
65#[cfg(target_arch = "x86_64")]
66#[target_feature(enable = "avx2")]
67pub fn znx_extract_digit_addmul_avx(base2k: usize, lsh: usize, res: &mut [i64], src: &mut [i64]) {
68 #[cfg(debug_assertions)]
69 {
70 assert_eq!(res.len(), src.len());
71 }
72
73 use std::arch::x86_64::{
74 __m256i, _mm256_add_epi64, _mm256_loadu_si256, _mm256_set1_epi64x, _mm256_sllv_epi64, _mm256_storeu_si256,
75 };
76
77 let n: usize = res.len();
78 let span: usize = n >> 2;
79
80 unsafe {
81 let mut rr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
82 let mut ss: *mut __m256i = src.as_mut_ptr() as *mut __m256i;
83
84 // constants for digit/carry extraction
85 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(base2k);
86 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
87
88 for _ in 0..span {
89 // load source & extract digit/carry
90 let sv: __m256i = _mm256_loadu_si256(ss);
91 let digit_256: __m256i = get_digit_avx(sv, mask, sign);
92 let carry_256: __m256i = get_carry_avx(sv, digit_256, basek_vec, top_mask);
93
94 // res += (digit << lsh)
95 let rv: __m256i = _mm256_loadu_si256(rr);
96 let madd: __m256i = _mm256_sllv_epi64(digit_256, lsh_v);
97 let sum: __m256i = _mm256_add_epi64(rv, madd);
98
99 _mm256_storeu_si256(rr, sum);
100 _mm256_storeu_si256(ss, carry_256);
101
102 rr = rr.add(1);
103 ss = ss.add(1);
104 }
105 }
106
107 // tail (scalar)
108 if !n.is_multiple_of(4) {
109 use poulpy_hal::reference::znx::znx_extract_digit_addmul_ref;
110
111 let off: usize = span << 2;
112 znx_extract_digit_addmul_ref(base2k, lsh, &mut res[off..], &mut src[off..]);
113 }
114}
115
116/// # Safety
117/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
118/// `res` and `src` must have the same length and must not alias.
119#[cfg(target_arch = "x86_64")]
120#[target_feature(enable = "avx2")]
121pub fn znx_normalize_digit_avx(base2k: usize, res: &mut [i64], src: &mut [i64]) {
122 #[cfg(debug_assertions)]
123 {
124 assert_eq!(res.len(), src.len());
125 }
126
127 use std::arch::x86_64::{__m256i, _mm256_add_epi64, _mm256_loadu_si256, _mm256_storeu_si256};
128
129 let n: usize = res.len();
130 let span: usize = n >> 2;
131
132 unsafe {
133 // Pointers to 256-bit lanes
134 let mut rr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
135 let mut ss: *mut __m256i = src.as_mut_ptr() as *mut __m256i;
136
137 // Constants for digit/carry extraction
138 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(base2k);
139
140 for _ in 0..span {
141 // Load res lane
142 let rv: __m256i = _mm256_loadu_si256(rr);
143
144 // Extract digit and carry from res
145 let digit_256: __m256i = get_digit_avx(rv, mask, sign);
146 let carry_256: __m256i = get_carry_avx(rv, digit_256, basek_vec, top_mask);
147
148 // src += carry
149 let sv: __m256i = _mm256_loadu_si256(ss);
150 let sum: __m256i = _mm256_add_epi64(sv, carry_256);
151
152 _mm256_storeu_si256(ss, sum);
153 _mm256_storeu_si256(rr, digit_256);
154
155 rr = rr.add(1);
156 ss = ss.add(1);
157 }
158 }
159
160 // scalar tail
161 if !n.is_multiple_of(4) {
162 use poulpy_hal::reference::znx::znx_normalize_digit_ref;
163
164 let off = span << 2;
165 znx_normalize_digit_ref(base2k, &mut res[off..], &mut src[off..]);
166 }
167}
168
169/// # Safety
170/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
171/// all inputs must have the same length and must not alias.
172#[target_feature(enable = "avx2")]
173pub fn znx_normalize_first_step_carry_only_avx(base2k: usize, lsh: usize, x: &[i64], carry: &mut [i64]) {
174 #[cfg(debug_assertions)]
175 {
176 assert!(x.len() <= carry.len());
177 assert!(lsh < base2k);
178 }
179
180 use std::arch::x86_64::{_mm256_loadu_si256, _mm256_storeu_si256};
181
182 let n: usize = x.len();
183
184 let span: usize = n >> 2;
185
186 unsafe {
187 let mut xx: *const __m256i = x.as_ptr() as *const __m256i;
188 let mut cc: *mut __m256i = carry.as_ptr() as *mut __m256i;
189
190 let (mask, sign, basek_vec, top_mask) = if lsh == 0 {
191 normalize_consts_avx(base2k)
192 } else {
193 normalize_consts_avx(base2k - lsh)
194 };
195
196 for _ in 0..span {
197 let xv: __m256i = _mm256_loadu_si256(xx);
198
199 // (x << (64 - base2k)) >> (64 - base2k)
200 let digit_256: __m256i = get_digit_avx(xv, mask, sign);
201
202 // (x - digit) >> base2k
203 let carry_256: __m256i = get_carry_avx(xv, digit_256, basek_vec, top_mask);
204
205 _mm256_storeu_si256(cc, carry_256);
206
207 xx = xx.add(1);
208 cc = cc.add(1);
209 }
210 }
211
212 // tail
213 if !x.len().is_multiple_of(4) {
214 use poulpy_hal::reference::znx::znx_normalize_first_step_carry_only_ref;
215
216 znx_normalize_first_step_carry_only_ref(base2k, lsh, &x[span << 2..], &mut carry[span << 2..]);
217 }
218}
219
220/// # Safety
221/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
222/// all inputs must have the same length and must not alias.
223#[target_feature(enable = "avx2")]
224pub fn znx_normalize_first_step_inplace_avx(base2k: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
225 #[cfg(debug_assertions)]
226 {
227 assert!(x.len() <= carry.len());
228 assert!(lsh < base2k);
229 }
230
231 use std::arch::x86_64::{_mm256_loadu_si256, _mm256_set1_epi64x, _mm256_sllv_epi64, _mm256_storeu_si256};
232
233 let n: usize = x.len();
234
235 let span: usize = n >> 2;
236
237 unsafe {
238 let mut xx: *mut __m256i = x.as_mut_ptr() as *mut __m256i;
239 let mut cc: *mut __m256i = carry.as_ptr() as *mut __m256i;
240
241 if lsh == 0 {
242 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(base2k);
243
244 for _ in 0..span {
245 let xv: __m256i = _mm256_loadu_si256(xx);
246
247 // (x << (64 - base2k)) >> (64 - base2k)
248 let digit_256: __m256i = get_digit_avx(xv, mask, sign);
249
250 // (x - digit) >> base2k
251 let carry_256: __m256i = get_carry_avx(xv, digit_256, basek_vec, top_mask);
252
253 _mm256_storeu_si256(xx, digit_256);
254 _mm256_storeu_si256(cc, carry_256);
255
256 xx = xx.add(1);
257 cc = cc.add(1);
258 }
259 } else {
260 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(base2k - lsh);
261
262 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
263
264 for _ in 0..span {
265 let xv: __m256i = _mm256_loadu_si256(xx);
266
267 // (x << (64 - base2k)) >> (64 - base2k)
268 let digit_256: __m256i = get_digit_avx(xv, mask, sign);
269
270 // (x - digit) >> base2k
271 let carry_256: __m256i = get_carry_avx(xv, digit_256, basek_vec, top_mask);
272
273 _mm256_storeu_si256(xx, _mm256_sllv_epi64(digit_256, lsh_v));
274 _mm256_storeu_si256(cc, carry_256);
275
276 xx = xx.add(1);
277 cc = cc.add(1);
278 }
279 }
280 }
281
282 // tail
283 if !x.len().is_multiple_of(4) {
284 use poulpy_hal::reference::znx::znx_normalize_first_step_inplace_ref;
285
286 znx_normalize_first_step_inplace_ref(base2k, lsh, &mut x[span << 2..], &mut carry[span << 2..]);
287 }
288}
289
290/// # Safety
291/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
292/// all inputs must have the same length and must not alias.
293#[target_feature(enable = "avx2")]
294pub fn znx_normalize_first_step_avx(base2k: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
295 #[cfg(debug_assertions)]
296 {
297 assert_eq!(x.len(), a.len());
298 assert!(x.len() <= carry.len());
299 assert!(lsh < base2k);
300 }
301
302 use std::arch::x86_64::{_mm256_loadu_si256, _mm256_sllv_epi64, _mm256_storeu_si256};
303
304 let n: usize = x.len();
305
306 let span: usize = n >> 2;
307
308 unsafe {
309 let mut xx: *mut __m256i = x.as_mut_ptr() as *mut __m256i;
310 let mut aa: *const __m256i = a.as_ptr() as *const __m256i;
311 let mut cc: *mut __m256i = carry.as_ptr() as *mut __m256i;
312
313 if lsh == 0 {
314 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(base2k);
315
316 for _ in 0..span {
317 let av: __m256i = _mm256_loadu_si256(aa);
318
319 // (x << (64 - base2k)) >> (64 - base2k)
320 let digit_256: __m256i = get_digit_avx(av, mask, sign);
321
322 // (x - digit) >> base2k
323 let carry_256: __m256i = get_carry_avx(av, digit_256, basek_vec, top_mask);
324
325 _mm256_storeu_si256(xx, digit_256);
326 _mm256_storeu_si256(cc, carry_256);
327
328 xx = xx.add(1);
329 aa = aa.add(1);
330 cc = cc.add(1);
331 }
332 } else {
333 use std::arch::x86_64::_mm256_set1_epi64x;
334
335 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(base2k - lsh);
336
337 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
338
339 for _ in 0..span {
340 let av: __m256i = _mm256_loadu_si256(aa);
341
342 // (x << (64 - base2k)) >> (64 - base2k)
343 let digit_256: __m256i = get_digit_avx(av, mask, sign);
344
345 // (x - digit) >> base2k
346 let carry_256: __m256i = get_carry_avx(av, digit_256, basek_vec, top_mask);
347
348 _mm256_storeu_si256(xx, _mm256_sllv_epi64(digit_256, lsh_v));
349 _mm256_storeu_si256(cc, carry_256);
350
351 xx = xx.add(1);
352 aa = aa.add(1);
353 cc = cc.add(1);
354 }
355 }
356 }
357
358 // tail
359 if !x.len().is_multiple_of(4) {
360 use poulpy_hal::reference::znx::znx_normalize_first_step_ref;
361
362 znx_normalize_first_step_ref(
363 base2k,
364 lsh,
365 &mut x[span << 2..],
366 &a[span << 2..],
367 &mut carry[span << 2..],
368 );
369 }
370}
371
372/// # Safety
373/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
374/// all inputs must have the same length and must not alias.
375#[target_feature(enable = "avx2")]
376pub fn znx_normalize_middle_step_inplace_avx(base2k: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
377 #[cfg(debug_assertions)]
378 {
379 assert!(x.len() <= carry.len());
380 assert!(lsh < base2k);
381 }
382
383 use std::arch::x86_64::{_mm256_add_epi64, _mm256_loadu_si256, _mm256_sllv_epi64, _mm256_storeu_si256};
384
385 let n: usize = x.len();
386
387 let span: usize = n >> 2;
388
389 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(base2k);
390
391 unsafe {
392 let mut xx: *mut __m256i = x.as_mut_ptr() as *mut __m256i;
393 let mut cc: *mut __m256i = carry.as_mut_ptr() as *mut __m256i;
394
395 if lsh == 0 {
396 for _ in 0..span {
397 let xv: __m256i = _mm256_loadu_si256(xx);
398 let cv: __m256i = _mm256_loadu_si256(cc);
399
400 let d0: __m256i = get_digit_avx(xv, mask, sign);
401 let c0: __m256i = get_carry_avx(xv, d0, basek_vec, top_mask);
402
403 let s: __m256i = _mm256_add_epi64(d0, cv);
404 let x1: __m256i = get_digit_avx(s, mask, sign);
405 let c1: __m256i = get_carry_avx(s, x1, basek_vec, top_mask);
406 let cout: __m256i = _mm256_add_epi64(c0, c1);
407
408 _mm256_storeu_si256(xx, x1);
409 _mm256_storeu_si256(cc, cout);
410
411 xx = xx.add(1);
412 cc = cc.add(1);
413 }
414 } else {
415 use std::arch::x86_64::_mm256_set1_epi64x;
416
417 let (mask_lsh, sign_lsh, basek_vec_lsh, top_mask_lsh) = normalize_consts_avx(base2k - lsh);
418
419 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
420
421 for _ in 0..span {
422 let xv: __m256i = _mm256_loadu_si256(xx);
423 let cv: __m256i = _mm256_loadu_si256(cc);
424
425 let d0: __m256i = get_digit_avx(xv, mask_lsh, sign_lsh);
426 let c0: __m256i = get_carry_avx(xv, d0, basek_vec_lsh, top_mask_lsh);
427
428 let d0_lsh: __m256i = _mm256_sllv_epi64(d0, lsh_v);
429
430 let s: __m256i = _mm256_add_epi64(d0_lsh, cv);
431 let x1: __m256i = get_digit_avx(s, mask, sign);
432 let c1: __m256i = get_carry_avx(s, x1, basek_vec, top_mask);
433 let cout: __m256i = _mm256_add_epi64(c0, c1);
434
435 _mm256_storeu_si256(xx, x1);
436 _mm256_storeu_si256(cc, cout);
437
438 xx = xx.add(1);
439 cc = cc.add(1);
440 }
441 }
442 }
443
444 if !x.len().is_multiple_of(4) {
445 use poulpy_hal::reference::znx::znx_normalize_middle_step_inplace_ref;
446
447 znx_normalize_middle_step_inplace_ref(base2k, lsh, &mut x[span << 2..], &mut carry[span << 2..]);
448 }
449}
450
451/// # Safety
452/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
453/// all inputs must have the same length and must not alias.
454#[target_feature(enable = "avx2")]
455pub fn znx_normalize_middle_step_carry_only_avx(base2k: usize, lsh: usize, x: &[i64], carry: &mut [i64]) {
456 #[cfg(debug_assertions)]
457 {
458 assert!(x.len() <= carry.len());
459 assert!(lsh < base2k);
460 }
461
462 use std::arch::x86_64::{_mm256_add_epi64, _mm256_loadu_si256, _mm256_sllv_epi64, _mm256_storeu_si256};
463
464 let n: usize = x.len();
465
466 let span: usize = n >> 2;
467
468 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(base2k);
469
470 unsafe {
471 let mut xx: *const __m256i = x.as_ptr() as *const __m256i;
472 let mut cc: *mut __m256i = carry.as_mut_ptr() as *mut __m256i;
473
474 if lsh == 0 {
475 for _ in 0..span {
476 let xv: __m256i = _mm256_loadu_si256(xx);
477 let cv: __m256i = _mm256_loadu_si256(cc);
478
479 let d0: __m256i = get_digit_avx(xv, mask, sign);
480 let c0: __m256i = get_carry_avx(xv, d0, basek_vec, top_mask);
481
482 let s: __m256i = _mm256_add_epi64(d0, cv);
483 let x1: __m256i = get_digit_avx(s, mask, sign);
484 let c1: __m256i = get_carry_avx(s, x1, basek_vec, top_mask);
485 let cout: __m256i = _mm256_add_epi64(c0, c1);
486
487 _mm256_storeu_si256(cc, cout);
488
489 xx = xx.add(1);
490 cc = cc.add(1);
491 }
492 } else {
493 use std::arch::x86_64::_mm256_set1_epi64x;
494
495 let (mask_lsh, sign_lsh, basek_vec_lsh, top_mask_lsh) = normalize_consts_avx(base2k - lsh);
496
497 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
498
499 for _ in 0..span {
500 let xv: __m256i = _mm256_loadu_si256(xx);
501 let cv: __m256i = _mm256_loadu_si256(cc);
502
503 let d0: __m256i = get_digit_avx(xv, mask_lsh, sign_lsh);
504 let c0: __m256i = get_carry_avx(xv, d0, basek_vec_lsh, top_mask_lsh);
505
506 let d0_lsh: __m256i = _mm256_sllv_epi64(d0, lsh_v);
507
508 let s: __m256i = _mm256_add_epi64(d0_lsh, cv);
509 let x1: __m256i = get_digit_avx(s, mask, sign);
510 let c1: __m256i = get_carry_avx(s, x1, basek_vec, top_mask);
511 let cout: __m256i = _mm256_add_epi64(c0, c1);
512
513 _mm256_storeu_si256(cc, cout);
514
515 xx = xx.add(1);
516 cc = cc.add(1);
517 }
518 }
519 }
520
521 if !x.len().is_multiple_of(4) {
522 use poulpy_hal::reference::znx::znx_normalize_middle_step_carry_only_ref;
523
524 znx_normalize_middle_step_carry_only_ref(base2k, lsh, &x[span << 2..], &mut carry[span << 2..]);
525 }
526}
527
528/// # Safety
529/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
530/// all inputs must have the same length and must not alias.
531#[target_feature(enable = "avx2")]
532pub fn znx_normalize_middle_step_avx(base2k: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
533 #[cfg(debug_assertions)]
534 {
535 assert_eq!(x.len(), a.len());
536 assert!(x.len() <= carry.len());
537 assert!(lsh < base2k);
538 }
539
540 use std::arch::x86_64::{_mm256_add_epi64, _mm256_loadu_si256, _mm256_sllv_epi64, _mm256_storeu_si256};
541
542 let n: usize = x.len();
543
544 let span: usize = n >> 2;
545
546 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(base2k);
547
548 unsafe {
549 let mut xx: *mut __m256i = x.as_mut_ptr() as *mut __m256i;
550 let mut aa: *const __m256i = a.as_ptr() as *const __m256i;
551 let mut cc: *mut __m256i = carry.as_ptr() as *mut __m256i;
552
553 if lsh == 0 {
554 for _ in 0..span {
555 let av: __m256i = _mm256_loadu_si256(aa);
556 let cv: __m256i = _mm256_loadu_si256(cc);
557
558 let d0: __m256i = get_digit_avx(av, mask, sign);
559 let c0: __m256i = get_carry_avx(av, d0, basek_vec, top_mask);
560
561 let s: __m256i = _mm256_add_epi64(d0, cv);
562 let x1: __m256i = get_digit_avx(s, mask, sign);
563 let c1: __m256i = get_carry_avx(s, x1, basek_vec, top_mask);
564 let cout: __m256i = _mm256_add_epi64(c0, c1);
565
566 _mm256_storeu_si256(xx, x1);
567 _mm256_storeu_si256(cc, cout);
568
569 xx = xx.add(1);
570 aa = aa.add(1);
571 cc = cc.add(1);
572 }
573 } else {
574 use std::arch::x86_64::_mm256_set1_epi64x;
575
576 let (mask_lsh, sign_lsh, basek_vec_lsh, top_mask_lsh) = normalize_consts_avx(base2k - lsh);
577
578 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
579
580 for _ in 0..span {
581 let av: __m256i = _mm256_loadu_si256(aa);
582 let cv: __m256i = _mm256_loadu_si256(cc);
583
584 let d0: __m256i = get_digit_avx(av, mask_lsh, sign_lsh);
585 let c0: __m256i = get_carry_avx(av, d0, basek_vec_lsh, top_mask_lsh);
586
587 let d0_lsh: __m256i = _mm256_sllv_epi64(d0, lsh_v);
588
589 let s: __m256i = _mm256_add_epi64(d0_lsh, cv);
590 let x1: __m256i = get_digit_avx(s, mask, sign);
591 let c1: __m256i = get_carry_avx(s, x1, basek_vec, top_mask);
592 let cout: __m256i = _mm256_add_epi64(c0, c1);
593
594 _mm256_storeu_si256(xx, x1);
595 _mm256_storeu_si256(cc, cout);
596
597 xx = xx.add(1);
598 aa = aa.add(1);
599 cc = cc.add(1);
600 }
601 }
602 }
603
604 if !x.len().is_multiple_of(4) {
605 use poulpy_hal::reference::znx::znx_normalize_middle_step_ref;
606
607 znx_normalize_middle_step_ref(
608 base2k,
609 lsh,
610 &mut x[span << 2..],
611 &a[span << 2..],
612 &mut carry[span << 2..],
613 );
614 }
615}
616
617/// # Safety
618/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
619/// all inputs must have the same length and must not alias.
620#[target_feature(enable = "avx2")]
621pub fn znx_normalize_final_step_inplace_avx(base2k: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
622 #[cfg(debug_assertions)]
623 {
624 assert!(x.len() <= carry.len());
625 assert!(lsh < base2k);
626 }
627
628 use std::arch::x86_64::{_mm256_add_epi64, _mm256_loadu_si256, _mm256_sllv_epi64, _mm256_storeu_si256};
629
630 let n: usize = x.len();
631
632 let span: usize = n >> 2;
633
634 let (mask, sign, _, _) = normalize_consts_avx(base2k);
635
636 unsafe {
637 let mut xx: *mut __m256i = x.as_mut_ptr() as *mut __m256i;
638 let mut cc: *mut __m256i = carry.as_ptr() as *mut __m256i;
639
640 if lsh == 0 {
641 for _ in 0..span {
642 let xv: __m256i = _mm256_loadu_si256(xx);
643 let cv: __m256i = _mm256_loadu_si256(cc);
644
645 let d0: __m256i = get_digit_avx(xv, mask, sign);
646 let s: __m256i = _mm256_add_epi64(d0, cv);
647 let x1: __m256i = get_digit_avx(s, mask, sign);
648
649 _mm256_storeu_si256(xx, x1);
650
651 xx = xx.add(1);
652 cc = cc.add(1);
653 }
654 } else {
655 use std::arch::x86_64::_mm256_set1_epi64x;
656
657 let (mask_lsh, sign_lsh, _, _) = normalize_consts_avx(base2k - lsh);
658
659 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
660
661 for _ in 0..span {
662 let xv: __m256i = _mm256_loadu_si256(xx);
663 let cv: __m256i = _mm256_loadu_si256(cc);
664
665 let d0: __m256i = get_digit_avx(xv, mask_lsh, sign_lsh);
666
667 let d0_lsh: __m256i = _mm256_sllv_epi64(d0, lsh_v);
668
669 let s: __m256i = _mm256_add_epi64(d0_lsh, cv);
670 let x1: __m256i = get_digit_avx(s, mask, sign);
671
672 _mm256_storeu_si256(xx, x1);
673
674 xx = xx.add(1);
675 cc = cc.add(1);
676 }
677 }
678 }
679
680 if !x.len().is_multiple_of(4) {
681 use poulpy_hal::reference::znx::znx_normalize_final_step_inplace_ref;
682
683 znx_normalize_final_step_inplace_ref(base2k, lsh, &mut x[span << 2..], &mut carry[span << 2..]);
684 }
685}
686
687/// # Safety
688/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
689/// all inputs must have the same length and must not alias.
690#[target_feature(enable = "avx2")]
691pub fn znx_normalize_final_step_avx(base2k: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
692 #[cfg(debug_assertions)]
693 {
694 assert_eq!(x.len(), a.len());
695 assert!(x.len() <= carry.len());
696 assert!(lsh < base2k);
697 }
698
699 use std::arch::x86_64::{_mm256_add_epi64, _mm256_loadu_si256, _mm256_sllv_epi64, _mm256_storeu_si256};
700
701 let n: usize = x.len();
702
703 let span: usize = n >> 2;
704
705 let (mask, sign, _, _) = normalize_consts_avx(base2k);
706
707 unsafe {
708 let mut xx: *mut __m256i = x.as_mut_ptr() as *mut __m256i;
709 let mut aa: *mut __m256i = a.as_ptr() as *mut __m256i;
710 let mut cc: *mut __m256i = carry.as_ptr() as *mut __m256i;
711
712 if lsh == 0 {
713 for _ in 0..span {
714 let av: __m256i = _mm256_loadu_si256(aa);
715 let cv: __m256i = _mm256_loadu_si256(cc);
716
717 let d0: __m256i = get_digit_avx(av, mask, sign);
718 let s: __m256i = _mm256_add_epi64(d0, cv);
719 let x1: __m256i = get_digit_avx(s, mask, sign);
720
721 _mm256_storeu_si256(xx, x1);
722
723 xx = xx.add(1);
724 aa = aa.add(1);
725 cc = cc.add(1);
726 }
727 } else {
728 use std::arch::x86_64::_mm256_set1_epi64x;
729
730 let (mask_lsh, sign_lsh, _, _) = normalize_consts_avx(base2k - lsh);
731
732 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
733
734 for _ in 0..span {
735 let av: __m256i = _mm256_loadu_si256(aa);
736 let cv: __m256i = _mm256_loadu_si256(cc);
737
738 let d0: __m256i = get_digit_avx(av, mask_lsh, sign_lsh);
739 let d0_lsh: __m256i = _mm256_sllv_epi64(d0, lsh_v);
740
741 let s: __m256i = _mm256_add_epi64(d0_lsh, cv);
742 let x1: __m256i = get_digit_avx(s, mask, sign);
743
744 _mm256_storeu_si256(xx, x1);
745
746 xx = xx.add(1);
747 aa = aa.add(1);
748 cc = cc.add(1);
749 }
750 }
751 }
752
753 if !x.len().is_multiple_of(4) {
754 use poulpy_hal::reference::znx::znx_normalize_final_step_ref;
755
756 znx_normalize_final_step_ref(
757 base2k,
758 lsh,
759 &mut x[span << 2..],
760 &a[span << 2..],
761 &mut carry[span << 2..],
762 );
763 }
764}
765
766mod tests {
767 use poulpy_hal::reference::znx::{
768 get_carry_i64, get_digit_i64, znx_extract_digit_addmul_ref, znx_normalize_digit_ref,
769 znx_normalize_final_step_inplace_ref, znx_normalize_final_step_ref, znx_normalize_first_step_inplace_ref,
770 znx_normalize_first_step_ref, znx_normalize_middle_step_inplace_ref, znx_normalize_middle_step_ref,
771 };
772
773 use super::*;
774
775 use std::arch::x86_64::{_mm256_loadu_si256, _mm256_storeu_si256};
776
777 #[allow(dead_code)]
778 #[target_feature(enable = "avx2")]
779 fn test_get_digit_avx_internal() {
780 let base2k: usize = 12;
781 let x: [i64; 4] = [
782 7638646372408325293,
783 -61440197422348985,
784 6835891051541717957,
785 -4835376105455195188,
786 ];
787 let y0: Vec<i64> = vec![
788 get_digit_i64(base2k, x[0]),
789 get_digit_i64(base2k, x[1]),
790 get_digit_i64(base2k, x[2]),
791 get_digit_i64(base2k, x[3]),
792 ];
793 let mut y1: Vec<i64> = vec![0i64; 4];
794 unsafe {
795 let x_256: __m256i = _mm256_loadu_si256(x.as_ptr() as *const __m256i);
796 let (mask, sign, _, _) = normalize_consts_avx(base2k);
797 let digit: __m256i = get_digit_avx(x_256, mask, sign);
798 _mm256_storeu_si256(y1.as_mut_ptr() as *mut __m256i, digit);
799 }
800 assert_eq!(y0, y1);
801 }
802
803 #[test]
804 fn test_get_digit_avx() {
805 if !std::is_x86_feature_detected!("avx2") {
806 eprintln!("skipping: CPU lacks avx2");
807 return;
808 };
809 unsafe {
810 test_get_digit_avx_internal();
811 }
812 }
813
814 #[allow(dead_code)]
815 #[target_feature(enable = "avx2")]
816 fn test_get_carry_avx_internal() {
817 let base2k: usize = 12;
818 let x: [i64; 4] = [
819 7638646372408325293,
820 -61440197422348985,
821 6835891051541717957,
822 -4835376105455195188,
823 ];
824 let carry: [i64; 4] = [1174467039, -144794816, -1466676977, 513122840];
825 let y0: Vec<i64> = vec![
826 get_carry_i64(base2k, x[0], carry[0]),
827 get_carry_i64(base2k, x[1], carry[1]),
828 get_carry_i64(base2k, x[2], carry[2]),
829 get_carry_i64(base2k, x[3], carry[3]),
830 ];
831 let mut y1: Vec<i64> = vec![0i64; 4];
832 unsafe {
833 let x_256: __m256i = _mm256_loadu_si256(x.as_ptr() as *const __m256i);
834 let d_256: __m256i = _mm256_loadu_si256(carry.as_ptr() as *const __m256i);
835 let (_, _, basek_vec, top_mask) = normalize_consts_avx(base2k);
836 let digit: __m256i = get_carry_avx(x_256, d_256, basek_vec, top_mask);
837 _mm256_storeu_si256(y1.as_mut_ptr() as *mut __m256i, digit);
838 }
839 assert_eq!(y0, y1);
840 }
841
842 #[test]
843 fn test_get_carry_avx() {
844 if !std::is_x86_feature_detected!("avx2") {
845 eprintln!("skipping: CPU lacks avx2");
846 return;
847 };
848 unsafe {
849 test_get_carry_avx_internal();
850 }
851 }
852
853 #[allow(dead_code)]
854 #[target_feature(enable = "avx2")]
855 fn test_znx_normalize_first_step_inplace_avx_internal() {
856 let mut y0: [i64; 4] = [
857 7638646372408325293,
858 -61440197422348985,
859 6835891051541717957,
860 -4835376105455195188,
861 ];
862 let mut y1: [i64; 4] = y0;
863
864 let mut c0: [i64; 4] = [
865 621182201135793202,
866 9000856573317006236,
867 5542252755421113668,
868 -6036847263131690631,
869 ];
870 let mut c1: [i64; 4] = c0;
871
872 let base2k = 12;
873
874 znx_normalize_first_step_inplace_ref(base2k, 0, &mut y0, &mut c0);
875 znx_normalize_first_step_inplace_avx(base2k, 0, &mut y1, &mut c1);
876
877 assert_eq!(y0, y1);
878 assert_eq!(c0, c1);
879
880 znx_normalize_first_step_inplace_ref(base2k, base2k - 1, &mut y0, &mut c0);
881 znx_normalize_first_step_inplace_avx(base2k, base2k - 1, &mut y1, &mut c1);
882
883 assert_eq!(y0, y1);
884 assert_eq!(c0, c1);
885 }
886
887 #[test]
888 fn test_znx_normalize_first_step_inplace_avx() {
889 if !std::is_x86_feature_detected!("avx2") {
890 eprintln!("skipping: CPU lacks avx2");
891 return;
892 };
893 unsafe {
894 test_znx_normalize_first_step_inplace_avx_internal();
895 }
896 }
897
898 #[allow(dead_code)]
899 #[target_feature(enable = "avx2")]
900 fn test_znx_normalize_middle_step_inplace_avx_internal() {
901 let mut y0: [i64; 4] = [
902 7638646372408325293,
903 -61440197422348985,
904 6835891051541717957,
905 -4835376105455195188,
906 ];
907 let mut y1: [i64; 4] = y0;
908
909 let mut c0: [i64; 4] = [
910 621182201135793202,
911 9000856573317006236,
912 5542252755421113668,
913 -6036847263131690631,
914 ];
915 let mut c1: [i64; 4] = c0;
916
917 let base2k = 12;
918
919 znx_normalize_middle_step_inplace_ref(base2k, 0, &mut y0, &mut c0);
920 znx_normalize_middle_step_inplace_avx(base2k, 0, &mut y1, &mut c1);
921
922 assert_eq!(y0, y1);
923 assert_eq!(c0, c1);
924
925 znx_normalize_middle_step_inplace_ref(base2k, base2k - 1, &mut y0, &mut c0);
926 znx_normalize_middle_step_inplace_avx(base2k, base2k - 1, &mut y1, &mut c1);
927
928 assert_eq!(y0, y1);
929 assert_eq!(c0, c1);
930 }
931
932 #[test]
933 fn test_znx_normalize_middle_step_inplace_avx() {
934 if !std::is_x86_feature_detected!("avx2") {
935 eprintln!("skipping: CPU lacks avx2");
936 return;
937 };
938 unsafe {
939 test_znx_normalize_middle_step_inplace_avx_internal();
940 }
941 }
942
943 #[allow(dead_code)]
944 #[target_feature(enable = "avx2")]
945 fn test_znx_normalize_final_step_inplace_avx_internal() {
946 let mut y0: [i64; 4] = [
947 7638646372408325293,
948 -61440197422348985,
949 6835891051541717957,
950 -4835376105455195188,
951 ];
952 let mut y1: [i64; 4] = y0;
953
954 let mut c0: [i64; 4] = [
955 621182201135793202,
956 9000856573317006236,
957 5542252755421113668,
958 -6036847263131690631,
959 ];
960 let mut c1: [i64; 4] = c0;
961
962 let base2k = 12;
963
964 znx_normalize_final_step_inplace_ref(base2k, 0, &mut y0, &mut c0);
965 znx_normalize_final_step_inplace_avx(base2k, 0, &mut y1, &mut c1);
966
967 assert_eq!(y0, y1);
968 assert_eq!(c0, c1);
969
970 znx_normalize_final_step_inplace_ref(base2k, base2k - 1, &mut y0, &mut c0);
971 znx_normalize_final_step_inplace_avx(base2k, base2k - 1, &mut y1, &mut c1);
972
973 assert_eq!(y0, y1);
974 assert_eq!(c0, c1);
975 }
976
977 #[test]
978 fn test_znx_normalize_final_step_inplace_avx() {
979 if !std::is_x86_feature_detected!("avx2") {
980 eprintln!("skipping: CPU lacks avx2");
981 return;
982 };
983 unsafe {
984 test_znx_normalize_final_step_inplace_avx_internal();
985 }
986 }
987
988 #[allow(dead_code)]
989 #[target_feature(enable = "avx2")]
990 fn test_znx_normalize_first_step_avx_internal() {
991 let mut y0: [i64; 4] = [
992 7638646372408325293,
993 -61440197422348985,
994 6835891051541717957,
995 -4835376105455195188,
996 ];
997 let mut y1: [i64; 4] = y0;
998 let a: [i64; 4] = y0;
999
1000 let mut c0: [i64; 4] = [
1001 621182201135793202,
1002 9000856573317006236,
1003 5542252755421113668,
1004 -6036847263131690631,
1005 ];
1006 let mut c1: [i64; 4] = c0;
1007
1008 let base2k = 12;
1009
1010 znx_normalize_first_step_ref(base2k, 0, &mut y0, &a, &mut c0);
1011 znx_normalize_first_step_avx(base2k, 0, &mut y1, &a, &mut c1);
1012
1013 assert_eq!(y0, y1);
1014 assert_eq!(c0, c1);
1015
1016 znx_normalize_first_step_ref(base2k, base2k - 1, &mut y0, &a, &mut c0);
1017 znx_normalize_first_step_avx(base2k, base2k - 1, &mut y1, &a, &mut c1);
1018
1019 assert_eq!(y0, y1);
1020 assert_eq!(c0, c1);
1021 }
1022
1023 #[test]
1024 fn test_znx_normalize_first_step_avx() {
1025 if !std::is_x86_feature_detected!("avx2") {
1026 eprintln!("skipping: CPU lacks avx2");
1027 return;
1028 };
1029 unsafe {
1030 test_znx_normalize_first_step_avx_internal();
1031 }
1032 }
1033
1034 #[allow(dead_code)]
1035 #[target_feature(enable = "avx2")]
1036 fn test_znx_normalize_middle_step_avx_internal() {
1037 let mut y0: [i64; 4] = [
1038 7638646372408325293,
1039 -61440197422348985,
1040 6835891051541717957,
1041 -4835376105455195188,
1042 ];
1043 let mut y1: [i64; 4] = y0;
1044 let a: [i64; 4] = y0;
1045
1046 let mut c0: [i64; 4] = [
1047 621182201135793202,
1048 9000856573317006236,
1049 5542252755421113668,
1050 -6036847263131690631,
1051 ];
1052 let mut c1: [i64; 4] = c0;
1053
1054 let base2k = 12;
1055
1056 znx_normalize_middle_step_ref(base2k, 0, &mut y0, &a, &mut c0);
1057 znx_normalize_middle_step_avx(base2k, 0, &mut y1, &a, &mut c1);
1058
1059 assert_eq!(y0, y1);
1060 assert_eq!(c0, c1);
1061
1062 znx_normalize_middle_step_ref(base2k, base2k - 1, &mut y0, &a, &mut c0);
1063 znx_normalize_middle_step_avx(base2k, base2k - 1, &mut y1, &a, &mut c1);
1064
1065 assert_eq!(y0, y1);
1066 assert_eq!(c0, c1);
1067 }
1068
1069 #[test]
1070 fn test_znx_normalize_middle_step_avx() {
1071 if !std::is_x86_feature_detected!("avx2") {
1072 eprintln!("skipping: CPU lacks avx2");
1073 return;
1074 };
1075 unsafe {
1076 test_znx_normalize_middle_step_avx_internal();
1077 }
1078 }
1079
1080 #[allow(dead_code)]
1081 #[target_feature(enable = "avx2")]
1082 fn test_znx_normalize_final_step_avx_internal() {
1083 let mut y0: [i64; 4] = [
1084 7638646372408325293,
1085 -61440197422348985,
1086 6835891051541717957,
1087 -4835376105455195188,
1088 ];
1089 let mut y1: [i64; 4] = y0;
1090 let a: [i64; 4] = y0;
1091
1092 let mut c0: [i64; 4] = [
1093 621182201135793202,
1094 9000856573317006236,
1095 5542252755421113668,
1096 -6036847263131690631,
1097 ];
1098 let mut c1: [i64; 4] = c0;
1099
1100 let base2k = 12;
1101
1102 znx_normalize_final_step_ref(base2k, 0, &mut y0, &a, &mut c0);
1103 znx_normalize_final_step_avx(base2k, 0, &mut y1, &a, &mut c1);
1104
1105 assert_eq!(y0, y1);
1106 assert_eq!(c0, c1);
1107
1108 znx_normalize_final_step_ref(base2k, base2k - 1, &mut y0, &a, &mut c0);
1109 znx_normalize_final_step_avx(base2k, base2k - 1, &mut y1, &a, &mut c1);
1110
1111 assert_eq!(y0, y1);
1112 assert_eq!(c0, c1);
1113 }
1114
1115 #[test]
1116 fn test_znx_normalize_final_step_avx() {
1117 if !std::is_x86_feature_detected!("avx2") {
1118 eprintln!("skipping: CPU lacks avx2");
1119 return;
1120 };
1121 unsafe {
1122 test_znx_normalize_final_step_avx_internal();
1123 }
1124 }
1125
1126 #[allow(dead_code)]
1127 #[target_feature(enable = "avx2")]
1128 fn znx_extract_digit_addmul_internal() {
1129 let mut y0: [i64; 4] = [
1130 7638646372408325293,
1131 -61440197422348985,
1132 6835891051541717957,
1133 -4835376105455195188,
1134 ];
1135 let mut y1: [i64; 4] = y0;
1136
1137 let mut c0: [i64; 4] = [
1138 621182201135793202,
1139 9000856573317006236,
1140 5542252755421113668,
1141 -6036847263131690631,
1142 ];
1143 let mut c1: [i64; 4] = c0;
1144
1145 let base2k: usize = 12;
1146
1147 znx_extract_digit_addmul_ref(base2k, 0, &mut y0, &mut c0);
1148 znx_extract_digit_addmul_avx(base2k, 0, &mut y1, &mut c1);
1149
1150 assert_eq!(y0, y1);
1151 assert_eq!(c0, c1);
1152
1153 znx_extract_digit_addmul_ref(base2k, base2k - 1, &mut y0, &mut c0);
1154 znx_extract_digit_addmul_avx(base2k, base2k - 1, &mut y1, &mut c1);
1155
1156 assert_eq!(y0, y1);
1157 assert_eq!(c0, c1);
1158 }
1159
1160 #[test]
1161 fn test_znx_extract_digit_addmul_avx() {
1162 if !std::is_x86_feature_detected!("avx2") {
1163 eprintln!("skipping: CPU lacks avx2");
1164 return;
1165 };
1166 unsafe {
1167 znx_extract_digit_addmul_internal();
1168 }
1169 }
1170
1171 #[allow(dead_code)]
1172 #[target_feature(enable = "avx2")]
1173 fn znx_normalize_digit_internal() {
1174 let mut y0: [i64; 4] = [
1175 7638646372408325293,
1176 -61440197422348985,
1177 6835891051541717957,
1178 -4835376105455195188,
1179 ];
1180 let mut y1: [i64; 4] = y0;
1181
1182 let mut c0: [i64; 4] = [
1183 621182201135793202,
1184 9000856573317006236,
1185 5542252755421113668,
1186 -6036847263131690631,
1187 ];
1188 let mut c1: [i64; 4] = c0;
1189
1190 let base2k: usize = 12;
1191
1192 znx_normalize_digit_ref(base2k, &mut y0, &mut c0);
1193 znx_normalize_digit_avx(base2k, &mut y1, &mut c1);
1194
1195 assert_eq!(y0, y1);
1196 assert_eq!(c0, c1);
1197 }
1198
1199 #[test]
1200 fn test_znx_normalize_digit_internal_avx() {
1201 if !std::is_x86_feature_detected!("avx2") {
1202 eprintln!("skipping: CPU lacks avx2");
1203 return;
1204 };
1205 unsafe {
1206 znx_normalize_digit_internal();
1207 }
1208 }
1209}