poulpy/poulpy-backend/src/cpu_fft64_avx/znx_avx/mul.rs
poulpy-backend/src/cpu_fft64_avx/znx_avx/mul.rsRaw
1/// Multiply/divide by a power of two with rounding matching [poulpy_hal::reference::znx::znx_mul_power_of_two_ref].
2///
3/// # Safety
4/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
5/// all inputs must have the same length and must not alias.
6#[cfg(target_arch = "x86_64")]
7#[target_feature(enable = "avx2")]
8pub unsafe fn znx_mul_power_of_two_avx(k: i64, res: &mut [i64], a: &[i64]) {
9 #[cfg(debug_assertions)]
10 {
11 assert_eq!(res.len(), a.len());
12 }
13
14 use core::arch::x86_64::{
15 __m128i, __m256i, _mm_cvtsi32_si128, _mm256_add_epi64, _mm256_and_si256, _mm256_cmpgt_epi64, _mm256_loadu_si256,
16 _mm256_or_si256, _mm256_set1_epi64x, _mm256_setzero_si256, _mm256_sll_epi64, _mm256_srl_epi64, _mm256_srli_epi64,
17 _mm256_storeu_si256, _mm256_sub_epi64,
18 };
19
20 let n: usize = res.len();
21
22 if n == 0 {
23 return;
24 }
25
26 if k == 0 {
27 use poulpy_hal::reference::znx::znx_copy_ref;
28 znx_copy_ref(res, a);
29 return;
30 }
31
32 let span: usize = n >> 2; // number of 256-bit chunks
33
34 unsafe {
35 let mut rr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
36 let mut aa: *const __m256i = a.as_ptr() as *const __m256i;
37
38 if k > 0 {
39 // Left shift by k (variable count).
40 #[cfg(debug_assertions)]
41 {
42 debug_assert!(k <= 63);
43 }
44 let cnt128: __m128i = _mm_cvtsi32_si128(k as i32);
45 for _ in 0..span {
46 let x: __m256i = _mm256_loadu_si256(aa);
47 let y: __m256i = _mm256_sll_epi64(x, cnt128);
48 _mm256_storeu_si256(rr, y);
49 rr = rr.add(1);
50 aa = aa.add(1);
51 }
52
53 // tail
54 if !n.is_multiple_of(4) {
55 use poulpy_hal::reference::znx::znx_mul_power_of_two_ref;
56
57 znx_mul_power_of_two_ref(k, &mut res[span << 2..], &a[span << 2..]);
58 }
59 return;
60 }
61
62 // k < 0 => arithmetic right shift with rounding:
63 // for each x:
64 // sign_bit = (x >> 63) & 1
65 // bias = (1<<(kp-1)) - sign_bit
66 // t = x + bias
67 // y = t >> kp (arithmetic)
68 let kp = -k;
69 #[cfg(debug_assertions)]
70 {
71 debug_assert!((1..=63).contains(&kp));
72 }
73
74 let cnt_right: __m128i = _mm_cvtsi32_si128(kp as i32);
75 let bias_base: __m256i = _mm256_set1_epi64x(1_i64 << (kp - 1));
76 let top_mask: __m256i = _mm256_set1_epi64x(-1_i64 << (64 - kp)); // high kp bits
77 let zero: __m256i = _mm256_setzero_si256();
78
79 for _ in 0..span {
80 let x = _mm256_loadu_si256(aa);
81
82 // bias = (1 << (kp-1)) - sign_bit
83 let sign_bit_x: __m256i = _mm256_srli_epi64(x, 63);
84 let bias: __m256i = _mm256_sub_epi64(bias_base, sign_bit_x);
85
86 // t = x + bias
87 let t: __m256i = _mm256_add_epi64(x, bias);
88
89 // logical shift
90 let lsr: __m256i = _mm256_srl_epi64(t, cnt_right);
91
92 // sign extension
93 let neg: __m256i = _mm256_cmpgt_epi64(zero, t);
94 let fill: __m256i = _mm256_and_si256(neg, top_mask);
95 let y: __m256i = _mm256_or_si256(lsr, fill);
96
97 _mm256_storeu_si256(rr, y);
98 rr = rr.add(1);
99 aa = aa.add(1);
100 }
101 }
102
103 // tail
104 if !n.is_multiple_of(4) {
105 use poulpy_hal::reference::znx::znx_mul_power_of_two_ref;
106
107 znx_mul_power_of_two_ref(k, &mut res[span << 2..], &a[span << 2..]);
108 }
109}
110
111/// Multiply/divide inplace by a power of two with rounding matching [poulpy_hal::reference::znx::znx_mul_power_of_two_inplace_ref].
112///
113/// # Safety
114/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
115/// all inputs must have the same length and must not alias.
116#[cfg(target_arch = "x86_64")]
117#[target_feature(enable = "avx2")]
118pub unsafe fn znx_mul_power_of_two_inplace_avx(k: i64, res: &mut [i64]) {
119 use core::arch::x86_64::{
120 __m128i, __m256i, _mm_cvtsi32_si128, _mm256_add_epi64, _mm256_and_si256, _mm256_cmpgt_epi64, _mm256_loadu_si256,
121 _mm256_or_si256, _mm256_set1_epi64x, _mm256_setzero_si256, _mm256_sll_epi64, _mm256_srl_epi64, _mm256_srli_epi64,
122 _mm256_storeu_si256, _mm256_sub_epi64,
123 };
124
125 let n: usize = res.len();
126
127 if n == 0 {
128 return;
129 }
130
131 if k == 0 {
132 return;
133 }
134
135 let span: usize = n >> 2; // number of 256-bit chunks
136
137 unsafe {
138 let mut rr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
139
140 if k > 0 {
141 // Left shift by k (variable count).
142 #[cfg(debug_assertions)]
143 {
144 debug_assert!(k <= 63);
145 }
146 let cnt128: __m128i = _mm_cvtsi32_si128(k as i32);
147 for _ in 0..span {
148 let x: __m256i = _mm256_loadu_si256(rr);
149 let y: __m256i = _mm256_sll_epi64(x, cnt128);
150 _mm256_storeu_si256(rr, y);
151 rr = rr.add(1);
152 }
153
154 // tail
155 if !n.is_multiple_of(4) {
156 use poulpy_hal::reference::znx::znx_mul_power_of_two_inplace_ref;
157 znx_mul_power_of_two_inplace_ref(k, &mut res[span << 2..]);
158 }
159 return;
160 }
161
162 // k < 0 => arithmetic right shift with rounding:
163 // for each x:
164 // sign_bit = (x >> 63) & 1
165 // bias = (1<<(kp-1)) - sign_bit
166 // t = x + bias
167 // y = t >> kp (arithmetic)
168 let kp = -k;
169 #[cfg(debug_assertions)]
170 {
171 debug_assert!((1..=63).contains(&kp));
172 }
173
174 let cnt_right: __m128i = _mm_cvtsi32_si128(kp as i32);
175 let bias_base: __m256i = _mm256_set1_epi64x(1_i64 << (kp - 1));
176 let top_mask: __m256i = _mm256_set1_epi64x(-1_i64 << (64 - kp)); // high kp bits
177 let zero: __m256i = _mm256_setzero_si256();
178
179 for _ in 0..span {
180 let x = _mm256_loadu_si256(rr);
181
182 // bias = (1 << (kp-1)) - sign_bit
183 let sign_bit_x: __m256i = _mm256_srli_epi64(x, 63);
184 let bias: __m256i = _mm256_sub_epi64(bias_base, sign_bit_x);
185
186 // t = x + bias
187 let t: __m256i = _mm256_add_epi64(x, bias);
188
189 // logical shift
190 let lsr: __m256i = _mm256_srl_epi64(t, cnt_right);
191
192 // sign extension
193 let neg: __m256i = _mm256_cmpgt_epi64(zero, t);
194 let fill: __m256i = _mm256_and_si256(neg, top_mask);
195 let y: __m256i = _mm256_or_si256(lsr, fill);
196
197 _mm256_storeu_si256(rr, y);
198 rr = rr.add(1);
199 }
200 }
201
202 // tail
203 if !n.is_multiple_of(4) {
204 use poulpy_hal::reference::znx::znx_mul_power_of_two_inplace_ref;
205 znx_mul_power_of_two_inplace_ref(k, &mut res[span << 2..]);
206 }
207}
208
209/// Multiply/divide by a power of two and add on the result with rounding matching [poulpy_hal::reference::znx::znx_mul_power_of_two_inplace_ref].
210///
211/// # Safety
212/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
213/// all inputs must have the same length and must not alias.
214#[cfg(target_arch = "x86_64")]
215#[target_feature(enable = "avx2")]
216pub unsafe fn znx_mul_add_power_of_two_avx(k: i64, res: &mut [i64], a: &[i64]) {
217 #[cfg(debug_assertions)]
218 {
219 assert_eq!(res.len(), a.len());
220 }
221
222 use core::arch::x86_64::{
223 __m128i, __m256i, _mm_cvtsi32_si128, _mm256_add_epi64, _mm256_and_si256, _mm256_cmpgt_epi64, _mm256_loadu_si256,
224 _mm256_or_si256, _mm256_set1_epi64x, _mm256_setzero_si256, _mm256_sll_epi64, _mm256_srl_epi64, _mm256_srli_epi64,
225 _mm256_storeu_si256, _mm256_sub_epi64,
226 };
227
228 let n: usize = res.len();
229
230 if n == 0 {
231 return;
232 }
233
234 if k == 0 {
235 use crate::cpu_fft64_avx::znx_avx::znx_add_inplace_avx;
236
237 znx_add_inplace_avx(res, a);
238 return;
239 }
240
241 let span: usize = n >> 2; // number of 256-bit chunks
242
243 unsafe {
244 let mut rr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
245 let mut aa: *const __m256i = a.as_ptr() as *const __m256i;
246
247 if k > 0 {
248 // Left shift by k (variable count).
249 #[cfg(debug_assertions)]
250 {
251 debug_assert!(k <= 63);
252 }
253 let cnt128: __m128i = _mm_cvtsi32_si128(k as i32);
254 for _ in 0..span {
255 let x: __m256i = _mm256_loadu_si256(aa);
256 let y: __m256i = _mm256_loadu_si256(rr);
257 _mm256_storeu_si256(rr, _mm256_add_epi64(y, _mm256_sll_epi64(x, cnt128)));
258 rr = rr.add(1);
259 aa = aa.add(1);
260 }
261
262 // tail
263 if !n.is_multiple_of(4) {
264 use poulpy_hal::reference::znx::znx_mul_add_power_of_two_ref;
265
266 znx_mul_add_power_of_two_ref(k, &mut res[span << 2..], &a[span << 2..]);
267 }
268 return;
269 }
270
271 // k < 0 => arithmetic right shift with rounding:
272 // for each x:
273 // sign_bit = (x >> 63) & 1
274 // bias = (1<<(kp-1)) - sign_bit
275 // t = x + bias
276 // y = t >> kp (arithmetic)
277 let kp = -k;
278 #[cfg(debug_assertions)]
279 {
280 debug_assert!((1..=63).contains(&kp));
281 }
282
283 let cnt_right: __m128i = _mm_cvtsi32_si128(kp as i32);
284 let bias_base: __m256i = _mm256_set1_epi64x(1_i64 << (kp - 1));
285 let top_mask: __m256i = _mm256_set1_epi64x(-1_i64 << (64 - kp)); // high kp bits
286 let zero: __m256i = _mm256_setzero_si256();
287
288 for _ in 0..span {
289 let x: __m256i = _mm256_loadu_si256(aa);
290 let y: __m256i = _mm256_loadu_si256(rr);
291
292 // bias = (1 << (kp-1)) - sign_bit
293 let sign_bit_x: __m256i = _mm256_srli_epi64(x, 63);
294 let bias: __m256i = _mm256_sub_epi64(bias_base, sign_bit_x);
295
296 // t = x + bias
297 let t: __m256i = _mm256_add_epi64(x, bias);
298
299 // logical shift
300 let lsr: __m256i = _mm256_srl_epi64(t, cnt_right);
301
302 // sign extension
303 let neg: __m256i = _mm256_cmpgt_epi64(zero, t);
304 let fill: __m256i = _mm256_and_si256(neg, top_mask);
305 let out: __m256i = _mm256_or_si256(lsr, fill);
306
307 _mm256_storeu_si256(rr, _mm256_add_epi64(y, out));
308 rr = rr.add(1);
309 aa = aa.add(1);
310 }
311 }
312
313 // tail
314 if !n.is_multiple_of(4) {
315 use poulpy_hal::reference::znx::znx_mul_add_power_of_two_ref;
316 znx_mul_add_power_of_two_ref(k, &mut res[span << 2..], &a[span << 2..]);
317 }
318}