arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/poulpy-backend/src/cpu_fft64_avx/reim/fft_vec_avx2_fma.rs
poulpy-backend/src/cpu_fft64_avx/reim/fft_vec_avx2_fma.rsRaw
  1/// # Safety
  2/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
  3#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
  4#[target_feature(enable = "avx2,fma")]
  5pub fn reim_add_avx2_fma(res: &mut [f64], a: &[f64], b: &[f64]) {
  6    #[cfg(debug_assertions)]
  7    {
  8        assert_eq!(a.len(), res.len());
  9        assert_eq!(b.len(), res.len());
 10    }
 11
 12    use std::arch::x86_64::{__m256d, _mm256_add_pd, _mm256_loadu_pd, _mm256_storeu_pd};
 13
 14    let span: usize = res.len() >> 2;
 15
 16    unsafe {
 17        let mut rr: *mut f64 = res.as_mut_ptr();
 18        let mut aa: *const f64 = a.as_ptr();
 19        let mut bb: *const f64 = b.as_ptr();
 20
 21        for _ in 0..span {
 22            let a_256: __m256d = _mm256_loadu_pd(aa);
 23            let b_256: __m256d = _mm256_loadu_pd(bb);
 24            _mm256_storeu_pd(rr, _mm256_add_pd(a_256, b_256));
 25            rr = rr.add(4);
 26            aa = aa.add(4);
 27            bb = bb.add(4);
 28        }
 29    }
 30}
 31
 32/// # Safety
 33/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
 34#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
 35#[target_feature(enable = "avx2,fma")]
 36pub fn reim_add_inplace_avx2_fma(res: &mut [f64], a: &[f64]) {
 37    #[cfg(debug_assertions)]
 38    {
 39        assert_eq!(a.len(), res.len());
 40    }
 41
 42    use std::arch::x86_64::{__m256d, _mm256_add_pd, _mm256_loadu_pd, _mm256_storeu_pd};
 43
 44    let span: usize = res.len() >> 2;
 45
 46    unsafe {
 47        let mut rr: *mut f64 = res.as_mut_ptr();
 48        let mut aa: *const f64 = a.as_ptr();
 49
 50        for _ in 0..span {
 51            let a_256: __m256d = _mm256_loadu_pd(aa);
 52            let r_256: __m256d = _mm256_loadu_pd(rr);
 53            _mm256_storeu_pd(rr, _mm256_add_pd(r_256, a_256));
 54            rr = rr.add(4);
 55            aa = aa.add(4);
 56        }
 57    }
 58}
 59
 60/// # Safety
 61/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
 62#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
 63#[target_feature(enable = "avx2,fma")]
 64pub fn reim_sub_avx2_fma(res: &mut [f64], a: &[f64], b: &[f64]) {
 65    #[cfg(debug_assertions)]
 66    {
 67        assert_eq!(a.len(), res.len());
 68        assert_eq!(b.len(), res.len());
 69    }
 70
 71    use std::arch::x86_64::{__m256d, _mm256_loadu_pd, _mm256_storeu_pd, _mm256_sub_pd};
 72
 73    let span: usize = res.len() >> 2;
 74
 75    unsafe {
 76        let mut rr: *mut f64 = res.as_mut_ptr();
 77        let mut aa: *const f64 = a.as_ptr();
 78        let mut bb: *const f64 = b.as_ptr();
 79
 80        for _ in 0..span {
 81            let a_256: __m256d = _mm256_loadu_pd(aa);
 82            let b_256: __m256d = _mm256_loadu_pd(bb);
 83            _mm256_storeu_pd(rr, _mm256_sub_pd(a_256, b_256));
 84            rr = rr.add(4);
 85            aa = aa.add(4);
 86            bb = bb.add(4);
 87        }
 88    }
 89}
 90
 91/// # Safety
 92/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
 93#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
 94#[target_feature(enable = "avx2,fma")]
 95pub fn reim_sub_ab_inplace_avx2_fma(res: &mut [f64], a: &[f64]) {
 96    #[cfg(debug_assertions)]
 97    {
 98        assert_eq!(a.len(), res.len());
 99    }
100
101    use std::arch::x86_64::{__m256d, _mm256_loadu_pd, _mm256_storeu_pd, _mm256_sub_pd};
102
103    let span: usize = res.len() >> 2;
104
105    unsafe {
106        let mut rr: *mut f64 = res.as_mut_ptr();
107        let mut aa: *const f64 = a.as_ptr();
108
109        for _ in 0..span {
110            let a_256: __m256d = _mm256_loadu_pd(aa);
111            let r_256: __m256d = _mm256_loadu_pd(rr);
112            _mm256_storeu_pd(rr, _mm256_sub_pd(r_256, a_256));
113            rr = rr.add(4);
114            aa = aa.add(4);
115        }
116    }
117}
118
119/// # Safety
120/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
121#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
122#[target_feature(enable = "avx2,fma")]
123pub fn reim_sub_ba_inplace_avx2_fma(res: &mut [f64], a: &[f64]) {
124    #[cfg(debug_assertions)]
125    {
126        assert_eq!(a.len(), res.len());
127    }
128
129    use std::arch::x86_64::{__m256d, _mm256_loadu_pd, _mm256_storeu_pd, _mm256_sub_pd};
130
131    let span: usize = res.len() >> 2;
132
133    unsafe {
134        let mut rr: *mut f64 = res.as_mut_ptr();
135        let mut aa: *const f64 = a.as_ptr();
136
137        for _ in 0..span {
138            let a_256: __m256d = _mm256_loadu_pd(aa);
139            let r_256: __m256d = _mm256_loadu_pd(rr);
140            _mm256_storeu_pd(rr, _mm256_sub_pd(a_256, r_256));
141            rr = rr.add(4);
142            aa = aa.add(4);
143        }
144    }
145}
146
147/// # Safety
148/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
149#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
150#[target_feature(enable = "avx2,fma")]
151pub fn reim_negate_avx2_fma(res: &mut [f64], a: &[f64]) {
152    #[cfg(debug_assertions)]
153    {
154        assert_eq!(a.len(), res.len());
155    }
156
157    use std::arch::x86_64::{__m256d, _mm256_loadu_pd, _mm256_storeu_pd, _mm256_xor_pd};
158
159    let span: usize = res.len() >> 2;
160
161    unsafe {
162        use std::arch::x86_64::_mm256_set1_pd;
163
164        let mut rr: *mut f64 = res.as_mut_ptr();
165        let mut aa: *const f64 = a.as_ptr();
166
167        let neg0: __m256d = _mm256_set1_pd(-0.0);
168
169        for _ in 0..span {
170            let a_256: __m256d = _mm256_loadu_pd(aa);
171            _mm256_storeu_pd(rr, _mm256_xor_pd(a_256, neg0));
172            rr = rr.add(4);
173            aa = aa.add(4);
174        }
175    }
176}
177
178/// # Safety
179/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
180#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
181#[target_feature(enable = "avx2,fma")]
182pub fn reim_negate_inplace_avx2_fma(res: &mut [f64]) {
183    use std::arch::x86_64::{__m256d, _mm256_loadu_pd, _mm256_storeu_pd, _mm256_xor_pd};
184
185    let span: usize = res.len() >> 2;
186
187    unsafe {
188        use std::arch::x86_64::_mm256_set1_pd;
189
190        let mut rr: *mut f64 = res.as_mut_ptr();
191        let neg0: __m256d = _mm256_set1_pd(-0.0);
192
193        for _ in 0..span {
194            let r_256: __m256d = _mm256_loadu_pd(rr);
195            _mm256_storeu_pd(rr, _mm256_xor_pd(r_256, neg0));
196            rr = rr.add(4);
197        }
198    }
199}
200
201/// # Safety
202/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
203#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
204#[target_feature(enable = "avx2,fma")]
205pub fn reim_addmul_avx2_fma(res: &mut [f64], a: &[f64], b: &[f64]) {
206    #[cfg(debug_assertions)]
207    {
208        assert_eq!(a.len(), res.len());
209        assert_eq!(b.len(), res.len());
210    }
211
212    let m: usize = res.len() >> 1;
213
214    let (rr, ri) = res.split_at_mut(m);
215    let (ar, ai) = a.split_at(m);
216    let (br, bi) = b.split_at(m);
217
218    unsafe {
219        let mut rr_ptr: *mut f64 = rr.as_mut_ptr();
220        let mut ri_ptr: *mut f64 = ri.as_mut_ptr();
221        let mut ar_ptr: *const f64 = ar.as_ptr();
222        let mut ai_ptr: *const f64 = ai.as_ptr();
223        let mut br_ptr: *const f64 = br.as_ptr();
224        let mut bi_ptr: *const f64 = bi.as_ptr();
225
226        use std::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_fmsub_pd, _mm256_loadu_pd, _mm256_storeu_pd};
227
228        for _ in 0..(m >> 2) {
229            let mut rr: __m256d = _mm256_loadu_pd(rr_ptr);
230            let mut ri: __m256d = _mm256_loadu_pd(ri_ptr);
231            let ar: __m256d = _mm256_loadu_pd(ar_ptr);
232            let ai: __m256d = _mm256_loadu_pd(ai_ptr);
233            let br: __m256d = _mm256_loadu_pd(br_ptr);
234            let bi: __m256d = _mm256_loadu_pd(bi_ptr);
235
236            rr = _mm256_fmsub_pd(ai, bi, rr);
237            rr = _mm256_fmsub_pd(ar, br, rr);
238            ri = _mm256_fmadd_pd(ar, bi, ri);
239            ri = _mm256_fmadd_pd(ai, br, ri);
240
241            _mm256_storeu_pd(rr_ptr, rr);
242            _mm256_storeu_pd(ri_ptr, ri);
243
244            rr_ptr = rr_ptr.add(4);
245            ri_ptr = ri_ptr.add(4);
246            ar_ptr = ar_ptr.add(4);
247            ai_ptr = ai_ptr.add(4);
248            br_ptr = br_ptr.add(4);
249            bi_ptr = bi_ptr.add(4);
250        }
251    }
252}
253
254/// # Safety
255/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
256#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
257#[target_feature(enable = "avx2,fma")]
258pub fn reim_mul_avx2_fma(res: &mut [f64], a: &[f64], b: &[f64]) {
259    #[cfg(debug_assertions)]
260    {
261        assert_eq!(a.len(), res.len());
262        assert_eq!(b.len(), res.len());
263    }
264
265    let m: usize = res.len() >> 1;
266
267    let (rr, ri) = res.split_at_mut(m);
268    let (ar, ai) = a.split_at(m);
269    let (br, bi) = b.split_at(m);
270
271    unsafe {
272        let mut rr_ptr: *mut f64 = rr.as_mut_ptr();
273        let mut ri_ptr: *mut f64 = ri.as_mut_ptr();
274        let mut ar_ptr: *const f64 = ar.as_ptr();
275        let mut ai_ptr: *const f64 = ai.as_ptr();
276        let mut br_ptr: *const f64 = br.as_ptr();
277        let mut bi_ptr: *const f64 = bi.as_ptr();
278
279        use std::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_fmsub_pd, _mm256_loadu_pd, _mm256_mul_pd, _mm256_storeu_pd};
280
281        for _ in 0..(m >> 2) {
282            let ar: __m256d = _mm256_loadu_pd(ar_ptr);
283            let ai: __m256d = _mm256_loadu_pd(ai_ptr);
284            let br: __m256d = _mm256_loadu_pd(br_ptr);
285            let bi: __m256d = _mm256_loadu_pd(bi_ptr);
286
287            let t1: __m256d = _mm256_mul_pd(ai, bi);
288            let t2: __m256d = _mm256_mul_pd(ar, bi);
289
290            let rr: __m256d = _mm256_fmsub_pd(ar, br, t1);
291            let ri: __m256d = _mm256_fmadd_pd(ai, br, t2);
292
293            _mm256_storeu_pd(rr_ptr, rr);
294            _mm256_storeu_pd(ri_ptr, ri);
295
296            rr_ptr = rr_ptr.add(4);
297            ri_ptr = ri_ptr.add(4);
298            ar_ptr = ar_ptr.add(4);
299            ai_ptr = ai_ptr.add(4);
300            br_ptr = br_ptr.add(4);
301            bi_ptr = bi_ptr.add(4);
302        }
303    }
304}
305
306/// # Safety
307/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
308#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
309#[target_feature(enable = "avx2,fma")]
310pub fn reim_mul_inplace_avx2_fma(res: &mut [f64], a: &[f64]) {
311    #[cfg(debug_assertions)]
312    {
313        assert_eq!(a.len(), res.len());
314    }
315
316    let m: usize = res.len() >> 1;
317
318    let (rr, ri) = res.split_at_mut(m);
319    let (ar, ai) = a.split_at(m);
320
321    unsafe {
322        let mut rr_ptr: *mut f64 = rr.as_mut_ptr();
323        let mut ri_ptr: *mut f64 = ri.as_mut_ptr();
324        let mut ar_ptr: *const f64 = ar.as_ptr();
325        let mut ai_ptr: *const f64 = ai.as_ptr();
326
327        use std::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_fmsub_pd, _mm256_loadu_pd, _mm256_mul_pd, _mm256_storeu_pd};
328
329        for _ in 0..(m >> 2) {
330            let ar: __m256d = _mm256_loadu_pd(ar_ptr);
331            let ai: __m256d = _mm256_loadu_pd(ai_ptr);
332            let br: __m256d = _mm256_loadu_pd(rr_ptr);
333            let bi: __m256d = _mm256_loadu_pd(ri_ptr);
334
335            let t1: __m256d = _mm256_mul_pd(ai, bi);
336            let t2: __m256d = _mm256_mul_pd(ar, bi);
337
338            let rr = _mm256_fmsub_pd(ar, br, t1);
339            let ri = _mm256_fmadd_pd(ai, br, t2);
340
341            _mm256_storeu_pd(rr_ptr, rr);
342            _mm256_storeu_pd(ri_ptr, ri);
343
344            rr_ptr = rr_ptr.add(4);
345            ri_ptr = ri_ptr.add(4);
346            ar_ptr = ar_ptr.add(4);
347            ai_ptr = ai_ptr.add(4);
348        }
349    }
350}