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