arnaucube
arnaucube/poulpyPublic Self-hosted repository
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#[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}