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#[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}