poulpy/poulpy-backend/src/cpu_fft64_avx/znx_avx/normalization.rs
poulpy-backend/src/cpu_fft64_avx/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(basek: usize) -> (__m256i, __m256i, __m256i, __m256i) {
10 use std::arch::x86_64::_mm256_set1_epi64x;
11
12 assert!((1..=63).contains(&basek));
13 let mask_k: i64 = ((1u64 << basek) - 1) as i64; // 0..k-1 bits set
14 let sign_k: i64 = (1u64 << (basek - 1)) as i64; // bit k-1
15 let topmask: i64 = (!0u64 << (64 - basek)) as i64; // top k bits set
16 let sh_k: __m256i = _mm256_set1_epi64x(basek 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 basek: __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, basek); // 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/// all inputs must have the same length and must not alias.
65#[target_feature(enable = "avx2")]
66pub fn znx_normalize_first_step_carry_only_avx(basek: usize, lsh: usize, x: &[i64], carry: &mut [i64]) {
67 #[cfg(debug_assertions)]
68 {
69 assert_eq!(x.len(), carry.len());
70 assert!(lsh < basek);
71 }
72
73 use std::arch::x86_64::{_mm256_loadu_si256, _mm256_storeu_si256};
74
75 let n: usize = x.len();
76
77 let span: usize = n >> 2;
78
79 unsafe {
80 let mut xx: *const __m256i = x.as_ptr() as *const __m256i;
81 let mut cc: *mut __m256i = carry.as_ptr() as *mut __m256i;
82
83 let (mask, sign, basek_vec, top_mask) = if lsh == 0 {
84 normalize_consts_avx(basek)
85 } else {
86 normalize_consts_avx(basek - lsh)
87 };
88
89 for _ in 0..span {
90 let xx_256: __m256i = _mm256_loadu_si256(xx);
91
92 // (x << (64 - basek)) >> (64 - basek)
93 let digit_256: __m256i = get_digit_avx(xx_256, mask, sign);
94
95 // (x - digit) >> basek
96 let carry_256: __m256i = get_carry_avx(xx_256, digit_256, basek_vec, top_mask);
97
98 _mm256_storeu_si256(cc, carry_256);
99
100 xx = xx.add(1);
101 cc = cc.add(1);
102 }
103 }
104
105 // tail
106 if !x.len().is_multiple_of(4) {
107 use poulpy_hal::reference::znx::znx_normalize_first_step_carry_only_ref;
108
109 znx_normalize_first_step_carry_only_ref(basek, lsh, &x[span << 2..], &mut carry[span << 2..]);
110 }
111}
112
113/// # Safety
114/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
115/// all inputs must have the same length and must not alias.
116#[target_feature(enable = "avx2")]
117pub fn znx_normalize_first_step_inplace_avx(basek: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
118 #[cfg(debug_assertions)]
119 {
120 assert_eq!(x.len(), carry.len());
121 assert!(lsh < basek);
122 }
123
124 use std::arch::x86_64::{_mm256_loadu_si256, _mm256_set1_epi64x, _mm256_sllv_epi64, _mm256_storeu_si256};
125
126 let n: usize = x.len();
127
128 let span: usize = n >> 2;
129
130 unsafe {
131 let mut xx: *mut __m256i = x.as_mut_ptr() as *mut __m256i;
132 let mut cc: *mut __m256i = carry.as_ptr() as *mut __m256i;
133
134 if lsh == 0 {
135 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(basek);
136
137 for _ in 0..span {
138 let xx_256: __m256i = _mm256_loadu_si256(xx);
139
140 // (x << (64 - basek)) >> (64 - basek)
141 let digit_256: __m256i = get_digit_avx(xx_256, mask, sign);
142
143 // (x - digit) >> basek
144 let carry_256: __m256i = get_carry_avx(xx_256, digit_256, basek_vec, top_mask);
145
146 _mm256_storeu_si256(xx, digit_256);
147 _mm256_storeu_si256(cc, carry_256);
148
149 xx = xx.add(1);
150 cc = cc.add(1);
151 }
152 } else {
153 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(basek - lsh);
154
155 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
156
157 for _ in 0..span {
158 let xx_256: __m256i = _mm256_loadu_si256(xx);
159
160 // (x << (64 - basek)) >> (64 - basek)
161 let digit_256: __m256i = get_digit_avx(xx_256, mask, sign);
162
163 // (x - digit) >> basek
164 let carry_256: __m256i = get_carry_avx(xx_256, digit_256, basek_vec, top_mask);
165
166 _mm256_storeu_si256(xx, _mm256_sllv_epi64(digit_256, lsh_v));
167 _mm256_storeu_si256(cc, carry_256);
168
169 xx = xx.add(1);
170 cc = cc.add(1);
171 }
172 }
173 }
174
175 // tail
176 if !x.len().is_multiple_of(4) {
177 use poulpy_hal::reference::znx::znx_normalize_first_step_inplace_ref;
178
179 znx_normalize_first_step_inplace_ref(basek, lsh, &mut x[span << 2..], &mut carry[span << 2..]);
180 }
181}
182
183/// # Safety
184/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
185/// all inputs must have the same length and must not alias.
186#[target_feature(enable = "avx2")]
187pub fn znx_normalize_first_step_avx(basek: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
188 #[cfg(debug_assertions)]
189 {
190 assert_eq!(x.len(), carry.len());
191 assert_eq!(a.len(), carry.len());
192 assert!(lsh < basek);
193 }
194
195 use std::arch::x86_64::{_mm256_loadu_si256, _mm256_sllv_epi64, _mm256_storeu_si256};
196
197 let n: usize = x.len();
198
199 let span: usize = n >> 2;
200
201 unsafe {
202 let mut xx: *mut __m256i = x.as_mut_ptr() as *mut __m256i;
203 let mut aa: *const __m256i = a.as_ptr() as *const __m256i;
204 let mut cc: *mut __m256i = carry.as_ptr() as *mut __m256i;
205
206 if lsh == 0 {
207 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(basek);
208
209 for _ in 0..span {
210 let aa_256: __m256i = _mm256_loadu_si256(aa);
211
212 // (x << (64 - basek)) >> (64 - basek)
213 let digit_256: __m256i = get_digit_avx(aa_256, mask, sign);
214
215 // (x - digit) >> basek
216 let carry_256: __m256i = get_carry_avx(aa_256, digit_256, basek_vec, top_mask);
217
218 _mm256_storeu_si256(xx, digit_256);
219 _mm256_storeu_si256(cc, carry_256);
220
221 xx = xx.add(1);
222 aa = aa.add(1);
223 cc = cc.add(1);
224 }
225 } else {
226 use std::arch::x86_64::_mm256_set1_epi64x;
227
228 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(basek - lsh);
229
230 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
231
232 for _ in 0..span {
233 let aa_256: __m256i = _mm256_loadu_si256(aa);
234
235 // (x << (64 - basek)) >> (64 - basek)
236 let digit_256: __m256i = get_digit_avx(aa_256, mask, sign);
237
238 // (x - digit) >> basek
239 let carry_256: __m256i = get_carry_avx(aa_256, digit_256, basek_vec, top_mask);
240
241 _mm256_storeu_si256(xx, _mm256_sllv_epi64(digit_256, lsh_v));
242 _mm256_storeu_si256(cc, carry_256);
243
244 xx = xx.add(1);
245 aa = aa.add(1);
246 cc = cc.add(1);
247 }
248 }
249 }
250
251 // tail
252 if !x.len().is_multiple_of(4) {
253 use poulpy_hal::reference::znx::znx_normalize_first_step_ref;
254
255 znx_normalize_first_step_ref(
256 basek,
257 lsh,
258 &mut x[span << 2..],
259 &a[span << 2..],
260 &mut carry[span << 2..],
261 );
262 }
263}
264
265/// # Safety
266/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
267/// all inputs must have the same length and must not alias.
268#[target_feature(enable = "avx2")]
269pub fn znx_normalize_middle_step_inplace_avx(basek: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
270 #[cfg(debug_assertions)]
271 {
272 assert_eq!(x.len(), carry.len());
273 assert!(lsh < basek);
274 }
275
276 use std::arch::x86_64::{_mm256_add_epi64, _mm256_loadu_si256, _mm256_sllv_epi64, _mm256_storeu_si256};
277
278 let n: usize = x.len();
279
280 let span: usize = n >> 2;
281
282 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(basek);
283
284 unsafe {
285 let mut xx: *mut __m256i = x.as_mut_ptr() as *mut __m256i;
286 let mut cc: *mut __m256i = carry.as_mut_ptr() as *mut __m256i;
287
288 if lsh == 0 {
289 for _ in 0..span {
290 let xx_256: __m256i = _mm256_loadu_si256(xx);
291 let cc_256: __m256i = _mm256_loadu_si256(cc);
292
293 let d0: __m256i = get_digit_avx(xx_256, mask, sign);
294 let c0: __m256i = get_carry_avx(xx_256, d0, basek_vec, top_mask);
295
296 let s: __m256i = _mm256_add_epi64(d0, cc_256);
297 let x1: __m256i = get_digit_avx(s, mask, sign);
298 let c1: __m256i = get_carry_avx(s, x1, basek_vec, top_mask);
299 let cout: __m256i = _mm256_add_epi64(c0, c1);
300
301 _mm256_storeu_si256(xx, x1);
302 _mm256_storeu_si256(cc, cout);
303
304 xx = xx.add(1);
305 cc = cc.add(1);
306 }
307 } else {
308 use std::arch::x86_64::_mm256_set1_epi64x;
309
310 let (mask_lsh, sign_lsh, basek_vec_lsh, top_mask_lsh) = normalize_consts_avx(basek - lsh);
311
312 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
313
314 for _ in 0..span {
315 let xx_256: __m256i = _mm256_loadu_si256(xx);
316 let cc_256: __m256i = _mm256_loadu_si256(cc);
317
318 let d0: __m256i = get_digit_avx(xx_256, mask_lsh, sign_lsh);
319 let c0: __m256i = get_carry_avx(xx_256, d0, basek_vec_lsh, top_mask_lsh);
320
321 let d0_lsh: __m256i = _mm256_sllv_epi64(d0, lsh_v);
322
323 let s: __m256i = _mm256_add_epi64(d0_lsh, cc_256);
324 let x1: __m256i = get_digit_avx(s, mask, sign);
325 let c1: __m256i = get_carry_avx(s, x1, basek_vec, top_mask);
326 let cout: __m256i = _mm256_add_epi64(c0, c1);
327
328 _mm256_storeu_si256(xx, x1);
329 _mm256_storeu_si256(cc, cout);
330
331 xx = xx.add(1);
332 cc = cc.add(1);
333 }
334 }
335 }
336
337 if !x.len().is_multiple_of(4) {
338 use poulpy_hal::reference::znx::znx_normalize_middle_step_inplace_ref;
339
340 znx_normalize_middle_step_inplace_ref(basek, lsh, &mut x[span << 2..], &mut carry[span << 2..]);
341 }
342}
343
344/// # Safety
345/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
346/// all inputs must have the same length and must not alias.
347#[target_feature(enable = "avx2")]
348pub fn znx_normalize_middle_step_carry_only_avx(basek: usize, lsh: usize, x: &[i64], carry: &mut [i64]) {
349 #[cfg(debug_assertions)]
350 {
351 assert_eq!(x.len(), carry.len());
352 assert!(lsh < basek);
353 }
354
355 use std::arch::x86_64::{_mm256_add_epi64, _mm256_loadu_si256, _mm256_sllv_epi64, _mm256_storeu_si256};
356
357 let n: usize = x.len();
358
359 let span: usize = n >> 2;
360
361 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(basek);
362
363 unsafe {
364 let mut xx: *const __m256i = x.as_ptr() as *const __m256i;
365 let mut cc: *mut __m256i = carry.as_mut_ptr() as *mut __m256i;
366
367 if lsh == 0 {
368 for _ in 0..span {
369 let xx_256: __m256i = _mm256_loadu_si256(xx);
370 let cc_256: __m256i = _mm256_loadu_si256(cc);
371
372 let d0: __m256i = get_digit_avx(xx_256, mask, sign);
373 let c0: __m256i = get_carry_avx(xx_256, d0, basek_vec, top_mask);
374
375 let s: __m256i = _mm256_add_epi64(d0, cc_256);
376 let x1: __m256i = get_digit_avx(s, mask, sign);
377 let c1: __m256i = get_carry_avx(s, x1, basek_vec, top_mask);
378 let cout: __m256i = _mm256_add_epi64(c0, c1);
379
380 _mm256_storeu_si256(cc, cout);
381
382 xx = xx.add(1);
383 cc = cc.add(1);
384 }
385 } else {
386 use std::arch::x86_64::_mm256_set1_epi64x;
387
388 let (mask_lsh, sign_lsh, basek_vec_lsh, top_mask_lsh) = normalize_consts_avx(basek - lsh);
389
390 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
391
392 for _ in 0..span {
393 let xx_256: __m256i = _mm256_loadu_si256(xx);
394 let cc_256: __m256i = _mm256_loadu_si256(cc);
395
396 let d0: __m256i = get_digit_avx(xx_256, mask_lsh, sign_lsh);
397 let c0: __m256i = get_carry_avx(xx_256, d0, basek_vec_lsh, top_mask_lsh);
398
399 let d0_lsh: __m256i = _mm256_sllv_epi64(d0, lsh_v);
400
401 let s: __m256i = _mm256_add_epi64(d0_lsh, cc_256);
402 let x1: __m256i = get_digit_avx(s, mask, sign);
403 let c1: __m256i = get_carry_avx(s, x1, basek_vec, top_mask);
404 let cout: __m256i = _mm256_add_epi64(c0, c1);
405
406 _mm256_storeu_si256(cc, cout);
407
408 xx = xx.add(1);
409 cc = cc.add(1);
410 }
411 }
412 }
413
414 if !x.len().is_multiple_of(4) {
415 use poulpy_hal::reference::znx::znx_normalize_middle_step_carry_only_ref;
416
417 znx_normalize_middle_step_carry_only_ref(basek, lsh, &x[span << 2..], &mut carry[span << 2..]);
418 }
419}
420
421/// # Safety
422/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
423/// all inputs must have the same length and must not alias.
424#[target_feature(enable = "avx2")]
425pub fn znx_normalize_middle_step_avx(basek: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
426 #[cfg(debug_assertions)]
427 {
428 assert_eq!(x.len(), carry.len());
429 assert_eq!(a.len(), carry.len());
430 assert!(lsh < basek);
431 }
432
433 use std::arch::x86_64::{_mm256_add_epi64, _mm256_loadu_si256, _mm256_sllv_epi64, _mm256_storeu_si256};
434
435 let n: usize = x.len();
436
437 let span: usize = n >> 2;
438
439 let (mask, sign, basek_vec, top_mask) = normalize_consts_avx(basek);
440
441 unsafe {
442 let mut xx: *mut __m256i = x.as_mut_ptr() as *mut __m256i;
443 let mut aa: *const __m256i = a.as_ptr() as *const __m256i;
444 let mut cc: *mut __m256i = carry.as_ptr() as *mut __m256i;
445
446 if lsh == 0 {
447 for _ in 0..span {
448 let aa_256: __m256i = _mm256_loadu_si256(aa);
449 let cc_256: __m256i = _mm256_loadu_si256(cc);
450
451 let d0: __m256i = get_digit_avx(aa_256, mask, sign);
452 let c0: __m256i = get_carry_avx(aa_256, d0, basek_vec, top_mask);
453
454 let s: __m256i = _mm256_add_epi64(d0, cc_256);
455 let x1: __m256i = get_digit_avx(s, mask, sign);
456 let c1: __m256i = get_carry_avx(s, x1, basek_vec, top_mask);
457 let cout: __m256i = _mm256_add_epi64(c0, c1);
458
459 _mm256_storeu_si256(xx, x1);
460 _mm256_storeu_si256(cc, cout);
461
462 xx = xx.add(1);
463 aa = aa.add(1);
464 cc = cc.add(1);
465 }
466 } else {
467 use std::arch::x86_64::_mm256_set1_epi64x;
468
469 let (mask_lsh, sign_lsh, basek_vec_lsh, top_mask_lsh) = normalize_consts_avx(basek - lsh);
470
471 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
472
473 for _ in 0..span {
474 let aa_256: __m256i = _mm256_loadu_si256(aa);
475 let cc_256: __m256i = _mm256_loadu_si256(cc);
476
477 let d0: __m256i = get_digit_avx(aa_256, mask_lsh, sign_lsh);
478 let c0: __m256i = get_carry_avx(aa_256, d0, basek_vec_lsh, top_mask_lsh);
479
480 let d0_lsh: __m256i = _mm256_sllv_epi64(d0, lsh_v);
481
482 let s: __m256i = _mm256_add_epi64(d0_lsh, cc_256);
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(xx, x1);
488 _mm256_storeu_si256(cc, cout);
489
490 xx = xx.add(1);
491 aa = aa.add(1);
492 cc = cc.add(1);
493 }
494 }
495 }
496
497 if !x.len().is_multiple_of(4) {
498 use poulpy_hal::reference::znx::znx_normalize_middle_step_ref;
499
500 znx_normalize_middle_step_ref(
501 basek,
502 lsh,
503 &mut x[span << 2..],
504 &a[span << 2..],
505 &mut carry[span << 2..],
506 );
507 }
508}
509
510/// # Safety
511/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
512/// all inputs must have the same length and must not alias.
513#[target_feature(enable = "avx2")]
514pub fn znx_normalize_final_step_inplace_avx(basek: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
515 #[cfg(debug_assertions)]
516 {
517 assert_eq!(x.len(), carry.len());
518 assert!(lsh < basek);
519 }
520
521 use std::arch::x86_64::{_mm256_add_epi64, _mm256_loadu_si256, _mm256_sllv_epi64, _mm256_storeu_si256};
522
523 let n: usize = x.len();
524
525 let span: usize = n >> 2;
526
527 let (mask, sign, _, _) = normalize_consts_avx(basek);
528
529 unsafe {
530 let mut xx: *mut __m256i = x.as_mut_ptr() as *mut __m256i;
531 let mut cc: *mut __m256i = carry.as_ptr() as *mut __m256i;
532
533 if lsh == 0 {
534 for _ in 0..span {
535 let xv: __m256i = _mm256_loadu_si256(xx);
536 let cv: __m256i = _mm256_loadu_si256(cc);
537
538 let d0: __m256i = get_digit_avx(xv, mask, sign);
539 let s: __m256i = _mm256_add_epi64(d0, cv);
540 let x1: __m256i = get_digit_avx(s, mask, sign);
541
542 _mm256_storeu_si256(xx, x1);
543
544 xx = xx.add(1);
545 cc = cc.add(1);
546 }
547 } else {
548 use std::arch::x86_64::_mm256_set1_epi64x;
549
550 let (mask_lsh, sign_lsh, _, _) = normalize_consts_avx(basek - lsh);
551
552 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
553
554 for _ in 0..span {
555 let xv: __m256i = _mm256_loadu_si256(xx);
556 let cv: __m256i = _mm256_loadu_si256(cc);
557
558 let d0: __m256i = get_digit_avx(xv, mask_lsh, sign_lsh);
559
560 let d0_lsh: __m256i = _mm256_sllv_epi64(d0, lsh_v);
561
562 let s: __m256i = _mm256_add_epi64(d0_lsh, cv);
563 let x1: __m256i = get_digit_avx(s, mask, sign);
564
565 _mm256_storeu_si256(xx, x1);
566
567 xx = xx.add(1);
568 cc = cc.add(1);
569 }
570 }
571 }
572
573 if !x.len().is_multiple_of(4) {
574 use poulpy_hal::reference::znx::znx_normalize_final_step_inplace_ref;
575
576 znx_normalize_final_step_inplace_ref(basek, lsh, &mut x[span << 2..], &mut carry[span << 2..]);
577 }
578}
579
580/// # Safety
581/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
582/// all inputs must have the same length and must not alias.
583#[target_feature(enable = "avx2")]
584pub fn znx_normalize_final_step_avx(basek: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
585 #[cfg(debug_assertions)]
586 {
587 assert_eq!(x.len(), carry.len());
588 assert_eq!(a.len(), carry.len());
589 assert!(lsh < basek);
590 }
591
592 use std::arch::x86_64::{_mm256_add_epi64, _mm256_loadu_si256, _mm256_sllv_epi64, _mm256_storeu_si256};
593
594 let n: usize = x.len();
595
596 let span: usize = n >> 2;
597
598 let (mask, sign, _, _) = normalize_consts_avx(basek);
599
600 unsafe {
601 let mut xx: *mut __m256i = x.as_mut_ptr() as *mut __m256i;
602 let mut aa: *mut __m256i = a.as_ptr() as *mut __m256i;
603 let mut cc: *mut __m256i = carry.as_ptr() as *mut __m256i;
604
605 if lsh == 0 {
606 for _ in 0..span {
607 let av: __m256i = _mm256_loadu_si256(aa);
608 let cv: __m256i = _mm256_loadu_si256(cc);
609
610 let d0: __m256i = get_digit_avx(av, mask, sign);
611 let s: __m256i = _mm256_add_epi64(d0, cv);
612 let x1: __m256i = get_digit_avx(s, mask, sign);
613
614 _mm256_storeu_si256(xx, x1);
615
616 xx = xx.add(1);
617 aa = aa.add(1);
618 cc = cc.add(1);
619 }
620 } else {
621 use std::arch::x86_64::_mm256_set1_epi64x;
622
623 let (mask_lsh, sign_lsh, _, _) = normalize_consts_avx(basek - lsh);
624
625 let lsh_v: __m256i = _mm256_set1_epi64x(lsh as i64);
626
627 for _ in 0..span {
628 let av: __m256i = _mm256_loadu_si256(aa);
629 let cv: __m256i = _mm256_loadu_si256(cc);
630
631 let d0: __m256i = get_digit_avx(av, mask_lsh, sign_lsh);
632 let d0_lsh: __m256i = _mm256_sllv_epi64(d0, lsh_v);
633
634 let s: __m256i = _mm256_add_epi64(d0_lsh, cv);
635 let x1: __m256i = get_digit_avx(s, mask, sign);
636
637 _mm256_storeu_si256(xx, x1);
638
639 xx = xx.add(1);
640 aa = aa.add(1);
641 cc = cc.add(1);
642 }
643 }
644 }
645
646 if !x.len().is_multiple_of(4) {
647 use poulpy_hal::reference::znx::znx_normalize_final_step_ref;
648
649 znx_normalize_final_step_ref(
650 basek,
651 lsh,
652 &mut x[span << 2..],
653 &a[span << 2..],
654 &mut carry[span << 2..],
655 );
656 }
657}
658
659mod tests {
660 use poulpy_hal::reference::znx::{
661 get_carry, get_digit, znx_normalize_final_step_inplace_ref, znx_normalize_final_step_ref,
662 znx_normalize_first_step_inplace_ref, znx_normalize_first_step_ref, znx_normalize_middle_step_inplace_ref,
663 znx_normalize_middle_step_ref,
664 };
665
666 use super::*;
667
668 use std::arch::x86_64::{_mm256_loadu_si256, _mm256_storeu_si256};
669
670 #[allow(dead_code)]
671 #[target_feature(enable = "avx2")]
672 fn test_get_digit_avx_internal() {
673 let basek: usize = 12;
674 let x: [i64; 4] = [
675 7638646372408325293,
676 -61440197422348985,
677 6835891051541717957,
678 -4835376105455195188,
679 ];
680 let y0: Vec<i64> = vec![
681 get_digit(basek, x[0]),
682 get_digit(basek, x[1]),
683 get_digit(basek, x[2]),
684 get_digit(basek, x[3]),
685 ];
686 let mut y1: Vec<i64> = vec![0i64; 4];
687 unsafe {
688 let x_256: __m256i = _mm256_loadu_si256(x.as_ptr() as *const __m256i);
689 let (mask, sign, _, _) = normalize_consts_avx(basek);
690 let digit: __m256i = get_digit_avx(x_256, mask, sign);
691 _mm256_storeu_si256(y1.as_mut_ptr() as *mut __m256i, digit);
692 }
693 assert_eq!(y0, y1);
694 }
695
696 #[test]
697 fn test_get_digit_avx() {
698 if !std::is_x86_feature_detected!("avx2") {
699 eprintln!("skipping: CPU lacks avx2");
700 return;
701 };
702 unsafe {
703 test_get_digit_avx_internal();
704 }
705 }
706
707 #[allow(dead_code)]
708 #[target_feature(enable = "avx2")]
709 fn test_get_carry_avx_internal() {
710 let basek: usize = 12;
711 let x: [i64; 4] = [
712 7638646372408325293,
713 -61440197422348985,
714 6835891051541717957,
715 -4835376105455195188,
716 ];
717 let carry: [i64; 4] = [1174467039, -144794816, -1466676977, 513122840];
718 let y0: Vec<i64> = vec![
719 get_carry(basek, x[0], carry[0]),
720 get_carry(basek, x[1], carry[1]),
721 get_carry(basek, x[2], carry[2]),
722 get_carry(basek, x[3], carry[3]),
723 ];
724 let mut y1: Vec<i64> = vec![0i64; 4];
725 unsafe {
726 let x_256: __m256i = _mm256_loadu_si256(x.as_ptr() as *const __m256i);
727 let d_256: __m256i = _mm256_loadu_si256(carry.as_ptr() as *const __m256i);
728 let (_, _, basek_vec, top_mask) = normalize_consts_avx(basek);
729 let digit: __m256i = get_carry_avx(x_256, d_256, basek_vec, top_mask);
730 _mm256_storeu_si256(y1.as_mut_ptr() as *mut __m256i, digit);
731 }
732 assert_eq!(y0, y1);
733 }
734
735 #[test]
736 fn test_get_carry_avx() {
737 if !std::is_x86_feature_detected!("avx2") {
738 eprintln!("skipping: CPU lacks avx2");
739 return;
740 };
741 unsafe {
742 test_get_carry_avx_internal();
743 }
744 }
745
746 #[allow(dead_code)]
747 #[target_feature(enable = "avx2")]
748 fn test_znx_normalize_first_step_inplace_avx_internal() {
749 let mut y0: [i64; 4] = [
750 7638646372408325293,
751 -61440197422348985,
752 6835891051541717957,
753 -4835376105455195188,
754 ];
755 let mut y1: [i64; 4] = y0;
756
757 let mut c0: [i64; 4] = [
758 621182201135793202,
759 9000856573317006236,
760 5542252755421113668,
761 -6036847263131690631,
762 ];
763 let mut c1: [i64; 4] = c0;
764
765 let basek = 12;
766
767 znx_normalize_first_step_inplace_ref(basek, 0, &mut y0, &mut c0);
768 znx_normalize_first_step_inplace_avx(basek, 0, &mut y1, &mut c1);
769
770 assert_eq!(y0, y1);
771 assert_eq!(c0, c1);
772
773 znx_normalize_first_step_inplace_ref(basek, basek - 1, &mut y0, &mut c0);
774 znx_normalize_first_step_inplace_avx(basek, basek - 1, &mut y1, &mut c1);
775
776 assert_eq!(y0, y1);
777 assert_eq!(c0, c1);
778 }
779
780 #[test]
781 fn test_znx_normalize_first_step_inplace_avx() {
782 if !std::is_x86_feature_detected!("avx2") {
783 eprintln!("skipping: CPU lacks avx2");
784 return;
785 };
786 unsafe {
787 test_znx_normalize_first_step_inplace_avx_internal();
788 }
789 }
790
791 #[allow(dead_code)]
792 #[target_feature(enable = "avx2")]
793 fn test_znx_normalize_middle_step_inplace_avx_internal() {
794 let mut y0: [i64; 4] = [
795 7638646372408325293,
796 -61440197422348985,
797 6835891051541717957,
798 -4835376105455195188,
799 ];
800 let mut y1: [i64; 4] = y0;
801
802 let mut c0: [i64; 4] = [
803 621182201135793202,
804 9000856573317006236,
805 5542252755421113668,
806 -6036847263131690631,
807 ];
808 let mut c1: [i64; 4] = c0;
809
810 let basek = 12;
811
812 znx_normalize_middle_step_inplace_ref(basek, 0, &mut y0, &mut c0);
813 znx_normalize_middle_step_inplace_avx(basek, 0, &mut y1, &mut c1);
814
815 assert_eq!(y0, y1);
816 assert_eq!(c0, c1);
817
818 znx_normalize_middle_step_inplace_ref(basek, basek - 1, &mut y0, &mut c0);
819 znx_normalize_middle_step_inplace_avx(basek, basek - 1, &mut y1, &mut c1);
820
821 assert_eq!(y0, y1);
822 assert_eq!(c0, c1);
823 }
824
825 #[test]
826 fn test_znx_normalize_middle_step_inplace_avx() {
827 if !std::is_x86_feature_detected!("avx2") {
828 eprintln!("skipping: CPU lacks avx2");
829 return;
830 };
831 unsafe {
832 test_znx_normalize_middle_step_inplace_avx_internal();
833 }
834 }
835
836 #[allow(dead_code)]
837 #[target_feature(enable = "avx2")]
838 fn test_znx_normalize_final_step_inplace_avx_internal() {
839 let mut y0: [i64; 4] = [
840 7638646372408325293,
841 -61440197422348985,
842 6835891051541717957,
843 -4835376105455195188,
844 ];
845 let mut y1: [i64; 4] = y0;
846
847 let mut c0: [i64; 4] = [
848 621182201135793202,
849 9000856573317006236,
850 5542252755421113668,
851 -6036847263131690631,
852 ];
853 let mut c1: [i64; 4] = c0;
854
855 let basek = 12;
856
857 znx_normalize_final_step_inplace_ref(basek, 0, &mut y0, &mut c0);
858 znx_normalize_final_step_inplace_avx(basek, 0, &mut y1, &mut c1);
859
860 assert_eq!(y0, y1);
861 assert_eq!(c0, c1);
862
863 znx_normalize_final_step_inplace_ref(basek, basek - 1, &mut y0, &mut c0);
864 znx_normalize_final_step_inplace_avx(basek, basek - 1, &mut y1, &mut c1);
865
866 assert_eq!(y0, y1);
867 assert_eq!(c0, c1);
868 }
869
870 #[test]
871 fn test_znx_normalize_final_step_inplace_avx() {
872 if !std::is_x86_feature_detected!("avx2") {
873 eprintln!("skipping: CPU lacks avx2");
874 return;
875 };
876 unsafe {
877 test_znx_normalize_final_step_inplace_avx_internal();
878 }
879 }
880
881 #[allow(dead_code)]
882 #[target_feature(enable = "avx2")]
883 fn test_znx_normalize_first_step_avx_internal() {
884 let mut y0: [i64; 4] = [
885 7638646372408325293,
886 -61440197422348985,
887 6835891051541717957,
888 -4835376105455195188,
889 ];
890 let mut y1: [i64; 4] = y0;
891 let a: [i64; 4] = y0;
892
893 let mut c0: [i64; 4] = [
894 621182201135793202,
895 9000856573317006236,
896 5542252755421113668,
897 -6036847263131690631,
898 ];
899 let mut c1: [i64; 4] = c0;
900
901 let basek = 12;
902
903 znx_normalize_first_step_ref(basek, 0, &mut y0, &a, &mut c0);
904 znx_normalize_first_step_avx(basek, 0, &mut y1, &a, &mut c1);
905
906 assert_eq!(y0, y1);
907 assert_eq!(c0, c1);
908
909 znx_normalize_first_step_ref(basek, basek - 1, &mut y0, &a, &mut c0);
910 znx_normalize_first_step_avx(basek, basek - 1, &mut y1, &a, &mut c1);
911
912 assert_eq!(y0, y1);
913 assert_eq!(c0, c1);
914 }
915
916 #[test]
917 fn test_znx_normalize_first_step_avx() {
918 if !std::is_x86_feature_detected!("avx2") {
919 eprintln!("skipping: CPU lacks avx2");
920 return;
921 };
922 unsafe {
923 test_znx_normalize_first_step_avx_internal();
924 }
925 }
926
927 #[allow(dead_code)]
928 #[target_feature(enable = "avx2")]
929 fn test_znx_normalize_middle_step_avx_internal() {
930 let mut y0: [i64; 4] = [
931 7638646372408325293,
932 -61440197422348985,
933 6835891051541717957,
934 -4835376105455195188,
935 ];
936 let mut y1: [i64; 4] = y0;
937 let a: [i64; 4] = y0;
938
939 let mut c0: [i64; 4] = [
940 621182201135793202,
941 9000856573317006236,
942 5542252755421113668,
943 -6036847263131690631,
944 ];
945 let mut c1: [i64; 4] = c0;
946
947 let basek = 12;
948
949 znx_normalize_middle_step_ref(basek, 0, &mut y0, &a, &mut c0);
950 znx_normalize_middle_step_avx(basek, 0, &mut y1, &a, &mut c1);
951
952 assert_eq!(y0, y1);
953 assert_eq!(c0, c1);
954
955 znx_normalize_middle_step_ref(basek, basek - 1, &mut y0, &a, &mut c0);
956 znx_normalize_middle_step_avx(basek, basek - 1, &mut y1, &a, &mut c1);
957
958 assert_eq!(y0, y1);
959 assert_eq!(c0, c1);
960 }
961
962 #[test]
963 fn test_znx_normalize_middle_step_avx() {
964 if !std::is_x86_feature_detected!("avx2") {
965 eprintln!("skipping: CPU lacks avx2");
966 return;
967 };
968 unsafe {
969 test_znx_normalize_middle_step_avx_internal();
970 }
971 }
972
973 #[allow(dead_code)]
974 #[target_feature(enable = "avx2")]
975 fn test_znx_normalize_final_step_avx_internal() {
976 let mut y0: [i64; 4] = [
977 7638646372408325293,
978 -61440197422348985,
979 6835891051541717957,
980 -4835376105455195188,
981 ];
982 let mut y1: [i64; 4] = y0;
983 let a: [i64; 4] = y0;
984
985 let mut c0: [i64; 4] = [
986 621182201135793202,
987 9000856573317006236,
988 5542252755421113668,
989 -6036847263131690631,
990 ];
991 let mut c1: [i64; 4] = c0;
992
993 let basek = 12;
994
995 znx_normalize_final_step_ref(basek, 0, &mut y0, &a, &mut c0);
996 znx_normalize_final_step_avx(basek, 0, &mut y1, &a, &mut c1);
997
998 assert_eq!(y0, y1);
999 assert_eq!(c0, c1);
1000
1001 znx_normalize_final_step_ref(basek, basek - 1, &mut y0, &a, &mut c0);
1002 znx_normalize_final_step_avx(basek, basek - 1, &mut y1, &a, &mut c1);
1003
1004 assert_eq!(y0, y1);
1005 assert_eq!(c0, c1);
1006 }
1007
1008 #[test]
1009 fn test_znx_normalize_final_step_avx() {
1010 if !std::is_x86_feature_detected!("avx2") {
1011 eprintln!("skipping: CPU lacks avx2");
1012 return;
1013 };
1014 unsafe {
1015 test_znx_normalize_final_step_avx_internal();
1016 }
1017 }
1018}