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