poulpy/poulpy-cpu-avx/src/reim/conversion.rs
poulpy-cpu-avx/src/reim/conversion.rsRaw
1/// # Correctness
2/// Ensured for inputs absolute value bounded by 2^50-1
3/// # Safety
4/// Caller must ensure the CPU supports FMA (e.g., via `is_x86_feature_detected!("fma")`);
5#[target_feature(enable = "fma")]
6pub fn reim_from_znx_i64_bnd50_fma(res: &mut [f64], a: &[i64]) {
7 #[cfg(debug_assertions)]
8 {
9 assert_eq!(res.len(), a.len())
10 }
11
12 let n: usize = res.len();
13
14 unsafe {
15 use std::arch::x86_64::{
16 __m256d, __m256i, _mm256_add_epi64, _mm256_castsi256_pd, _mm256_loadu_si256, _mm256_or_pd, _mm256_set1_epi64x,
17 _mm256_set1_pd, _mm256_storeu_pd, _mm256_sub_pd,
18 };
19
20 let expo: f64 = (1i64 << 52) as f64;
21 let add_cst: i64 = 1i64 << 51;
22 let sub_cst: f64 = (3i64 << 51) as f64;
23
24 let expo_256: __m256d = _mm256_set1_pd(expo);
25 let add_cst_256: __m256i = _mm256_set1_epi64x(add_cst);
26 let sub_cst_256: __m256d = _mm256_set1_pd(sub_cst);
27
28 let mut res_ptr: *mut f64 = res.as_mut_ptr();
29 let mut a_ptr: *const __m256i = a.as_ptr() as *const __m256i;
30
31 let span: usize = n >> 2;
32
33 for _ in 0..span {
34 let mut ai64_256: __m256i = _mm256_loadu_si256(a_ptr);
35
36 ai64_256 = _mm256_add_epi64(ai64_256, add_cst_256);
37
38 let mut af64_256: __m256d = _mm256_castsi256_pd(ai64_256);
39 af64_256 = _mm256_or_pd(af64_256, expo_256);
40 af64_256 = _mm256_sub_pd(af64_256, sub_cst_256);
41
42 _mm256_storeu_pd(res_ptr, af64_256);
43
44 res_ptr = res_ptr.add(4);
45 a_ptr = a_ptr.add(1);
46 }
47
48 if !res.len().is_multiple_of(4) {
49 use poulpy_hal::reference::fft64::reim::reim_from_znx_i64_ref;
50 reim_from_znx_i64_ref(&mut res[span << 2..], &a[span << 2..])
51 }
52 }
53}
54
55/// # Correctness
56/// Only ensured for inputs absoluate value bounded by 2^63-1
57/// # Safety
58/// Caller must ensure the CPU supports FMA (e.g., via `is_x86_feature_detected!("fma,avx2")`);
59#[allow(dead_code)]
60#[target_feature(enable = "avx2,fma")]
61pub fn reim_to_znx_i64_bnd63_avx2_fma(res: &mut [i64], divisor: f64, a: &[f64]) {
62 #[cfg(debug_assertions)]
63 {
64 assert_eq!(res.len(), a.len())
65 }
66
67 let sign_mask: u64 = 0x8000000000000000u64;
68 let expo_mask: u64 = 0x7FF0000000000000u64;
69 let mantissa_mask: u64 = (i64::MAX as u64) ^ expo_mask;
70 let mantissa_msb: u64 = 0x0010000000000000u64;
71 let divi_bits: f64 = divisor * (1i64 << 52) as f64;
72 let offset: f64 = divisor / 2.;
73
74 unsafe {
75 use std::arch::x86_64::{
76 __m256d, __m256i, _mm256_add_pd, _mm256_and_pd, _mm256_and_si256, _mm256_castpd_si256, _mm256_castsi256_pd,
77 _mm256_loadu_pd, _mm256_or_pd, _mm256_or_si256, _mm256_set1_epi64x, _mm256_set1_pd, _mm256_sllv_epi64,
78 _mm256_srli_epi64, _mm256_srlv_epi64, _mm256_sub_epi64, _mm256_xor_si256,
79 };
80
81 let sign_mask_256: __m256d = _mm256_castsi256_pd(_mm256_set1_epi64x(sign_mask as i64));
82 let expo_mask_256: __m256i = _mm256_set1_epi64x(expo_mask as i64);
83 let mantissa_mask_256: __m256i = _mm256_set1_epi64x(mantissa_mask as i64);
84 let mantissa_msb_256: __m256i = _mm256_set1_epi64x(mantissa_msb as i64);
85 let offset_256 = _mm256_set1_pd(offset);
86 let divi_bits_256 = _mm256_castpd_si256(_mm256_set1_pd(divi_bits));
87
88 let mut res_ptr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
89 let mut a_ptr: *const f64 = a.as_ptr();
90
91 let span: usize = res.len() >> 2;
92
93 for _ in 0..span {
94 // read the next value
95 use std::arch::x86_64::_mm256_storeu_si256;
96 let mut a: __m256d = _mm256_loadu_pd(a_ptr);
97
98 // a += sign(a) * m/2
99 let asign: __m256d = _mm256_and_pd(a, sign_mask_256);
100 a = _mm256_add_pd(a, _mm256_or_pd(asign, offset_256));
101
102 // sign: either 0 or -1
103 let mut sign_mask: __m256i = _mm256_castpd_si256(asign);
104 sign_mask = _mm256_sub_epi64(_mm256_set1_epi64x(0), _mm256_srli_epi64(sign_mask, 63));
105
106 // compute the exponents
107 let a0exp: __m256i = _mm256_and_si256(_mm256_castpd_si256(a), expo_mask_256);
108 let mut a0lsh: __m256i = _mm256_sub_epi64(a0exp, divi_bits_256);
109 let mut a0rsh: __m256i = _mm256_sub_epi64(divi_bits_256, a0exp);
110 a0lsh = _mm256_srli_epi64(a0lsh, 52);
111 a0rsh = _mm256_srli_epi64(a0rsh, 52);
112
113 // compute the new mantissa
114 let mut a0pos: __m256i = _mm256_and_si256(_mm256_castpd_si256(a), mantissa_mask_256);
115 a0pos = _mm256_or_si256(a0pos, mantissa_msb_256);
116 a0lsh = _mm256_sllv_epi64(a0pos, a0lsh);
117 a0rsh = _mm256_srlv_epi64(a0pos, a0rsh);
118 let mut out: __m256i = _mm256_or_si256(a0lsh, a0rsh);
119
120 // negate if the sign was negative
121 out = _mm256_xor_si256(out, sign_mask);
122 out = _mm256_sub_epi64(out, sign_mask);
123
124 // stores
125 _mm256_storeu_si256(res_ptr, out);
126
127 res_ptr = res_ptr.add(1);
128 a_ptr = a_ptr.add(4);
129 }
130
131 if !res.len().is_multiple_of(4) {
132 use poulpy_hal::reference::fft64::reim::reim_to_znx_i64_ref;
133 reim_to_znx_i64_ref(&mut res[span << 2..], divisor, &a[span << 2..])
134 }
135 }
136}
137
138/// # Correctness
139/// Only ensured for inputs absoluate value bounded by 2^63-1
140/// # Safety
141/// Caller must ensure the CPU supports FMA (e.g., via `is_x86_feature_detected!("fma,avx2")`);
142#[target_feature(enable = "avx2,fma")]
143pub fn reim_to_znx_i64_inplace_bnd63_avx2_fma(res: &mut [f64], divisor: f64) {
144 let sign_mask: u64 = 0x8000000000000000u64;
145 let expo_mask: u64 = 0x7FF0000000000000u64;
146 let mantissa_mask: u64 = (i64::MAX as u64) ^ expo_mask;
147 let mantissa_msb: u64 = 0x0010000000000000u64;
148 let divi_bits: f64 = divisor * (1i64 << 52) as f64;
149 let offset: f64 = divisor / 2.;
150
151 unsafe {
152 use std::arch::x86_64::{
153 __m256d, __m256i, _mm256_add_pd, _mm256_and_pd, _mm256_and_si256, _mm256_castpd_si256, _mm256_castsi256_pd,
154 _mm256_loadu_pd, _mm256_or_pd, _mm256_or_si256, _mm256_set1_epi64x, _mm256_set1_pd, _mm256_sllv_epi64,
155 _mm256_srli_epi64, _mm256_srlv_epi64, _mm256_sub_epi64, _mm256_xor_si256,
156 };
157
158 use poulpy_hal::reference::fft64::reim::reim_to_znx_i64_inplace_ref;
159
160 let sign_mask_256: __m256d = _mm256_castsi256_pd(_mm256_set1_epi64x(sign_mask as i64));
161 let expo_mask_256: __m256i = _mm256_set1_epi64x(expo_mask as i64);
162 let mantissa_mask_256: __m256i = _mm256_set1_epi64x(mantissa_mask as i64);
163 let mantissa_msb_256: __m256i = _mm256_set1_epi64x(mantissa_msb as i64);
164 let offset_256: __m256d = _mm256_set1_pd(offset);
165 let divi_bits_256: __m256i = _mm256_castpd_si256(_mm256_set1_pd(divi_bits));
166
167 let mut res_ptr_4xi64: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
168 let mut res_ptr_1xf64: *mut f64 = res.as_mut_ptr();
169
170 let span: usize = res.len() >> 2;
171
172 for _ in 0..span {
173 // read the next value
174 use std::arch::x86_64::_mm256_storeu_si256;
175 let mut a: __m256d = _mm256_loadu_pd(res_ptr_1xf64);
176
177 // a += sign(a) * m/2
178 let asign: __m256d = _mm256_and_pd(a, sign_mask_256);
179 a = _mm256_add_pd(a, _mm256_or_pd(asign, offset_256));
180
181 // sign: either 0 or -1
182 let mut sign_mask: __m256i = _mm256_castpd_si256(asign);
183 sign_mask = _mm256_sub_epi64(_mm256_set1_epi64x(0), _mm256_srli_epi64(sign_mask, 63));
184
185 // compute the exponents
186 let a0exp: __m256i = _mm256_and_si256(_mm256_castpd_si256(a), expo_mask_256);
187 let mut a0lsh: __m256i = _mm256_sub_epi64(a0exp, divi_bits_256);
188 let mut a0rsh: __m256i = _mm256_sub_epi64(divi_bits_256, a0exp);
189 a0lsh = _mm256_srli_epi64(a0lsh, 52);
190 a0rsh = _mm256_srli_epi64(a0rsh, 52);
191
192 // compute the new mantissa
193 let mut a0pos: __m256i = _mm256_and_si256(_mm256_castpd_si256(a), mantissa_mask_256);
194 a0pos = _mm256_or_si256(a0pos, mantissa_msb_256);
195 a0lsh = _mm256_sllv_epi64(a0pos, a0lsh);
196 a0rsh = _mm256_srlv_epi64(a0pos, a0rsh);
197 let mut out: __m256i = _mm256_or_si256(a0lsh, a0rsh);
198
199 // negate if the sign was negative
200 out = _mm256_xor_si256(out, sign_mask);
201 out = _mm256_sub_epi64(out, sign_mask);
202
203 // stores
204 _mm256_storeu_si256(res_ptr_4xi64, out);
205
206 res_ptr_4xi64 = res_ptr_4xi64.add(1);
207 res_ptr_1xf64 = res_ptr_1xf64.add(4);
208 }
209
210 if !res.len().is_multiple_of(4) {
211 reim_to_znx_i64_inplace_ref(&mut res[span << 2..], divisor)
212 }
213 }
214}
215
216/// # Correctness
217/// Only ensured for inputs absoluate value bounded by 2^50-1
218/// # Safety
219/// Caller must ensure the CPU supports FMA (e.g., via `is_x86_feature_detected!("fma")`);
220#[target_feature(enable = "fma")]
221#[allow(dead_code)]
222pub fn reim_to_znx_i64_avx2_bnd50_fma(res: &mut [i64], divisor: f64, a: &[f64]) {
223 #[cfg(debug_assertions)]
224 {
225 assert_eq!(res.len(), a.len())
226 }
227
228 unsafe {
229 use std::arch::x86_64::{
230 __m256d, __m256i, _mm256_add_pd, _mm256_and_si256, _mm256_castpd_si256, _mm256_loadu_pd, _mm256_set1_epi64x,
231 _mm256_set1_pd, _mm256_storeu_si256, _mm256_sub_epi64,
232 };
233
234 let mantissa_mask: u64 = 0x000FFFFFFFFFFFFFu64;
235 let sub_cst: i64 = 1i64 << 51;
236 let add_cst: f64 = divisor * (3i64 << 51) as f64;
237
238 let sub_cst_4: __m256i = _mm256_set1_epi64x(sub_cst);
239 let add_cst_4: std::arch::x86_64::__m256d = _mm256_set1_pd(add_cst);
240 let mantissa_mask_4: __m256i = _mm256_set1_epi64x(mantissa_mask as i64);
241
242 let mut res_ptr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
243 let mut a_ptr = a.as_ptr();
244
245 let span: usize = res.len() >> 2;
246
247 for _ in 0..span {
248 // read the next value
249 let mut a: __m256d = _mm256_loadu_pd(a_ptr);
250 a = _mm256_add_pd(a, add_cst_4);
251 let mut ai: __m256i = _mm256_castpd_si256(a);
252 ai = _mm256_and_si256(ai, mantissa_mask_4);
253 ai = _mm256_sub_epi64(ai, sub_cst_4);
254 // store the next value
255 _mm256_storeu_si256(res_ptr, ai);
256
257 res_ptr = res_ptr.add(1);
258 a_ptr = a_ptr.add(4);
259 }
260
261 if !res.len().is_multiple_of(4) {
262 use poulpy_hal::reference::fft64::reim::reim_to_znx_i64_ref;
263 reim_to_znx_i64_ref(&mut res[span << 2..], divisor, &a[span << 2..])
264 }
265 }
266}