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