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