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