poulpy/poulpy-backend/src/cpu_fft64_avx/reim4/arithmetic_avx.rs
poulpy-backend/src/cpu_fft64_avx/reim4/arithmetic_avx.rsRaw
1/// # Safety
2/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
3#[cfg(target_arch = "x86_64")]
4#[target_feature(enable = "avx")]
5pub fn reim4_extract_1blk_from_reim_avx(m: usize, rows: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
6 use core::arch::x86_64::{__m256d, _mm256_loadu_pd, _mm256_storeu_pd};
7
8 unsafe {
9 let mut src_ptr: *const __m256d = src.as_ptr().add(blk << 2) as *const __m256d; // src + 4*blk
10 let mut dst_ptr: *mut __m256d = dst.as_mut_ptr() as *mut __m256d;
11
12 let step: usize = m >> 2;
13
14 // Each iteration copies 4 doubles; advance src by m doubles each row
15 for _ in 0..2 * rows {
16 let v: __m256d = _mm256_loadu_pd(src_ptr as *const f64);
17 _mm256_storeu_pd(dst_ptr as *mut f64, v);
18 dst_ptr = dst_ptr.add(1);
19 src_ptr = src_ptr.add(step);
20 }
21 }
22}
23
24/// # Safety
25/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
26#[cfg(target_arch = "x86_64")]
27#[target_feature(enable = "avx2,fma")]
28pub fn reim4_save_1blk_to_reim_avx<const OVERWRITE: bool>(m: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
29 use core::arch::x86_64::{__m256d, _mm256_add_pd, _mm256_loadu_pd, _mm256_storeu_pd};
30 unsafe {
31 let off: usize = blk * 4;
32 let src_ptr: *const f64 = src.as_ptr();
33
34 let s0: __m256d = _mm256_loadu_pd(src_ptr);
35 let s1: __m256d = _mm256_loadu_pd(src_ptr.add(4));
36
37 let d0_ptr: *mut f64 = dst.as_mut_ptr().add(off);
38 let d1_ptr: *mut f64 = d0_ptr.add(m);
39
40 if OVERWRITE {
41 _mm256_storeu_pd(d0_ptr, s0);
42 _mm256_storeu_pd(d1_ptr, s1);
43 } else {
44 let d0: __m256d = _mm256_loadu_pd(d0_ptr);
45 let d1: __m256d = _mm256_loadu_pd(d1_ptr);
46 _mm256_storeu_pd(d0_ptr, _mm256_add_pd(d0, s0));
47 _mm256_storeu_pd(d1_ptr, _mm256_add_pd(d1, s1));
48 }
49 }
50}
51
52/// # Safety
53/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
54#[cfg(target_arch = "x86_64")]
55#[target_feature(enable = "avx2,fma")]
56pub fn reim4_save_2blk_to_reim_avx<const OVERWRITE: bool>(
57 m: usize, //
58 blk: usize, // block index
59 dst: &mut [f64], //
60 src: &[f64], // 16 doubles [re1(4), im1(4), re2(4), im2(4)]
61) {
62 use core::arch::x86_64::{__m256d, _mm256_add_pd, _mm256_loadu_pd, _mm256_storeu_pd};
63 unsafe {
64 let off: usize = blk * 4;
65 let src_ptr: *const f64 = src.as_ptr();
66
67 let d0_ptr: *mut f64 = dst.as_mut_ptr().add(off);
68 let d1_ptr: *mut f64 = d0_ptr.add(m);
69 let d2_ptr: *mut f64 = d1_ptr.add(m);
70 let d3_ptr: *mut f64 = d2_ptr.add(m);
71
72 let s0: __m256d = _mm256_loadu_pd(src_ptr);
73 let s1: __m256d = _mm256_loadu_pd(src_ptr.add(4));
74 let s2: __m256d = _mm256_loadu_pd(src_ptr.add(8));
75 let s3: __m256d = _mm256_loadu_pd(src_ptr.add(12));
76
77 if OVERWRITE {
78 _mm256_storeu_pd(d0_ptr, s0);
79 _mm256_storeu_pd(d1_ptr, s1);
80 _mm256_storeu_pd(d2_ptr, s2);
81 _mm256_storeu_pd(d3_ptr, s3);
82 } else {
83 let d0: __m256d = _mm256_loadu_pd(d0_ptr);
84 let d1: __m256d = _mm256_loadu_pd(d1_ptr);
85 let d2: __m256d = _mm256_loadu_pd(d2_ptr);
86 let d3: __m256d = _mm256_loadu_pd(d3_ptr);
87 _mm256_storeu_pd(d0_ptr, _mm256_add_pd(d0, s0));
88 _mm256_storeu_pd(d1_ptr, _mm256_add_pd(d1, s1));
89 _mm256_storeu_pd(d2_ptr, _mm256_add_pd(d2, s2));
90 _mm256_storeu_pd(d3_ptr, _mm256_add_pd(d3, s3));
91 }
92 }
93}
94
95/// # Safety
96/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
97#[cfg(target_arch = "x86_64")]
98#[target_feature(enable = "avx2", enable = "fma")]
99pub fn reim4_vec_mat1col_product_avx(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
100 use core::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_loadu_pd, _mm256_setzero_pd, _mm256_storeu_pd};
101
102 #[cfg(debug_assertions)]
103 {
104 assert!(dst.len() >= 8, "dst must have at least 8 doubles");
105 assert!(u.len() >= nrows * 8, "u must be at least nrows * 8 doubles");
106 assert!(v.len() >= nrows * 8, "v must be at least nrows * 8 doubles");
107 }
108
109 unsafe {
110 use std::arch::x86_64::{_mm256_add_pd, _mm256_sub_pd};
111
112 let mut re1: __m256d = _mm256_setzero_pd();
113 let mut im1: __m256d = _mm256_setzero_pd();
114 let mut re2: __m256d = _mm256_setzero_pd();
115 let mut im2: __m256d = _mm256_setzero_pd();
116
117 let mut u_ptr: *const f64 = u.as_ptr();
118 let mut v_ptr: *const f64 = v.as_ptr();
119
120 for _ in 0..nrows {
121 let ur: __m256d = _mm256_loadu_pd(u_ptr);
122 let ui: __m256d = _mm256_loadu_pd(u_ptr.add(4));
123 let vr: __m256d = _mm256_loadu_pd(v_ptr);
124 let vi: __m256d = _mm256_loadu_pd(v_ptr.add(4));
125
126 // re1 = re1 + ur*vr;
127 re1 = _mm256_fmadd_pd(ur, vr, re1);
128 // im1 = im1 + ur*d;
129 im1 = _mm256_fmadd_pd(ur, vi, im1);
130 // re2 = re2 + ui*d;
131 re2 = _mm256_fmadd_pd(ui, vi, re2);
132 // im2 = im2 + ui*vr;
133 im2 = _mm256_fmadd_pd(ui, vr, im2);
134
135 u_ptr = u_ptr.add(8);
136 v_ptr = v_ptr.add(8);
137 }
138
139 // re1 - re2
140 _mm256_storeu_pd(dst.as_mut_ptr(), _mm256_sub_pd(re1, re2));
141
142 // im1 + im2
143 _mm256_storeu_pd(dst.as_mut_ptr().add(4), _mm256_add_pd(im1, im2));
144 }
145}
146
147/// # Safety
148/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
149#[cfg(target_arch = "x86_64")]
150#[target_feature(enable = "avx2", enable = "fma")]
151pub fn reim4_vec_mat2cols_product_avx(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
152 use core::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_fmsub_pd, _mm256_loadu_pd, _mm256_setzero_pd, _mm256_storeu_pd};
153
154 #[cfg(debug_assertions)]
155 {
156 assert!(
157 dst.len() >= 8,
158 "dst must be at least 8 doubles but is {}",
159 dst.len()
160 );
161 assert!(
162 u.len() >= nrows * 8,
163 "u must be at least nrows={} * 8 doubles but is {}",
164 nrows,
165 u.len()
166 );
167 assert!(
168 v.len() >= nrows * 16,
169 "v must be at least nrows={} * 16 doubles but is {}",
170 nrows,
171 v.len()
172 );
173 }
174
175 unsafe {
176 let mut re1: __m256d = _mm256_setzero_pd();
177 let mut im1: __m256d = _mm256_setzero_pd();
178 let mut re2: __m256d = _mm256_setzero_pd();
179 let mut im2: __m256d = _mm256_setzero_pd();
180
181 let mut u_ptr: *const f64 = u.as_ptr();
182 let mut v_ptr: *const f64 = v.as_ptr();
183
184 for _ in 0..nrows {
185 let ur: __m256d = _mm256_loadu_pd(u_ptr);
186 let ui: __m256d = _mm256_loadu_pd(u_ptr.add(4));
187
188 let ar: __m256d = _mm256_loadu_pd(v_ptr);
189 let ai: __m256d = _mm256_loadu_pd(v_ptr.add(4));
190 let br: __m256d = _mm256_loadu_pd(v_ptr.add(8));
191 let bi: __m256d = _mm256_loadu_pd(v_ptr.add(12));
192
193 // re1 = re1 - ui*ai; re2 = re2 - ui*bi;
194 re1 = _mm256_fmsub_pd(ui, ai, re1);
195 re2 = _mm256_fmsub_pd(ui, bi, re2);
196 // im1 = im1 + ur*ai; im2 = im2 + ur*bi;
197 im1 = _mm256_fmadd_pd(ur, ai, im1);
198 im2 = _mm256_fmadd_pd(ur, bi, im2);
199 // re1 = re1 - ur*ar; re2 = re2 - ur*br;
200 re1 = _mm256_fmsub_pd(ur, ar, re1);
201 re2 = _mm256_fmsub_pd(ur, br, re2);
202 // im1 = im1 + ui*ar; im2 = im2 + ui*br;
203 im1 = _mm256_fmadd_pd(ui, ar, im1);
204 im2 = _mm256_fmadd_pd(ui, br, im2);
205
206 u_ptr = u_ptr.add(8);
207 v_ptr = v_ptr.add(16);
208 }
209
210 _mm256_storeu_pd(dst.as_mut_ptr(), re1);
211 _mm256_storeu_pd(dst.as_mut_ptr().add(4), im1);
212 _mm256_storeu_pd(dst.as_mut_ptr().add(8), re2);
213 _mm256_storeu_pd(dst.as_mut_ptr().add(12), im2);
214 }
215}
216
217/// # Safety
218/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
219#[cfg(target_arch = "x86_64")]
220#[target_feature(enable = "avx2", enable = "fma")]
221pub fn reim4_vec_mat2cols_2ndcol_product_avx(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
222 use core::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_fmsub_pd, _mm256_loadu_pd, _mm256_setzero_pd, _mm256_storeu_pd};
223
224 #[cfg(debug_assertions)]
225 {
226 assert_eq!(dst.len(), 16, "dst must have 16 doubles");
227 assert!(u.len() >= nrows * 8, "u must be at least nrows * 8 doubles");
228 assert!(
229 v.len() >= nrows * 16,
230 "v must be at least nrows * 16 doubles"
231 );
232 }
233
234 unsafe {
235 let mut re1: __m256d = _mm256_setzero_pd();
236 let mut im1: __m256d = _mm256_setzero_pd();
237
238 let mut u_ptr: *const f64 = u.as_ptr();
239 let mut v_ptr: *const f64 = v.as_ptr().add(8); // Offset to 2nd column
240
241 for _ in 0..nrows {
242 let ur: __m256d = _mm256_loadu_pd(u_ptr);
243 let ui: __m256d = _mm256_loadu_pd(u_ptr.add(4));
244
245 let ar: __m256d = _mm256_loadu_pd(v_ptr);
246 let ai: __m256d = _mm256_loadu_pd(v_ptr.add(4));
247
248 // re1 = re1 - ui*ai; re2 = re2 - ui*bi;
249 re1 = _mm256_fmsub_pd(ui, ai, re1);
250 // im1 = im1 + ur*ai; im2 = im2 + ur*bi;
251 im1 = _mm256_fmadd_pd(ur, ai, im1);
252 // re1 = re1 - ur*ar; re2 = re2 - ur*br;
253 re1 = _mm256_fmsub_pd(ur, ar, re1);
254 // im1 = im1 + ui*ar; im2 = im2 + ui*br;
255 im1 = _mm256_fmadd_pd(ui, ar, im1);
256
257 u_ptr = u_ptr.add(8);
258 v_ptr = v_ptr.add(16);
259 }
260
261 _mm256_storeu_pd(dst.as_mut_ptr(), re1);
262 _mm256_storeu_pd(dst.as_mut_ptr().add(4), im1);
263 }
264}