poulpy/poulpy-cpu-avx/src/convolution.rs
poulpy-cpu-avx/src/convolution.rsRaw
1use poulpy_hal::{
2 api::{Convolution, ModuleN, ScratchTakeBasic, TakeSlice, VecZnxDftApply, VecZnxDftBytesOf},
3 layouts::{
4 Backend, CnvPVecL, CnvPVecLToMut, CnvPVecLToRef, CnvPVecR, CnvPVecRToMut, CnvPVecRToRef, Module, Scratch, VecZnx,
5 VecZnxBig, VecZnxBigToMut, VecZnxDft, VecZnxDftToMut, VecZnxToRef, ZnxInfos,
6 },
7 oep::{CnvPVecBytesOfImpl, CnvPVecLAllocImpl, ConvolutionImpl},
8 reference::fft64::convolution::{
9 convolution_apply_dft, convolution_apply_dft_tmp_bytes, convolution_by_const_apply, convolution_by_const_apply_tmp_bytes,
10 convolution_pairwise_apply_dft, convolution_pairwise_apply_dft_tmp_bytes, convolution_prepare_left,
11 convolution_prepare_right,
12 },
13};
14
15use crate::{FFT64Avx, module::FFT64ModuleHandle};
16
17unsafe impl CnvPVecLAllocImpl<Self> for FFT64Avx {
18 fn cnv_pvec_left_alloc_impl(n: usize, cols: usize, size: usize) -> CnvPVecL<Vec<u8>, Self> {
19 CnvPVecL::alloc(n, cols, size)
20 }
21
22 fn cnv_pvec_right_alloc_impl(n: usize, cols: usize, size: usize) -> CnvPVecR<Vec<u8>, Self> {
23 CnvPVecR::alloc(n, cols, size)
24 }
25}
26
27unsafe impl CnvPVecBytesOfImpl for FFT64Avx {
28 fn bytes_of_cnv_pvec_left_impl(n: usize, cols: usize, size: usize) -> usize {
29 Self::layout_prep_word_count() * n * cols * size * size_of::<<FFT64Avx as Backend>::ScalarPrep>()
30 }
31
32 fn bytes_of_cnv_pvec_right_impl(n: usize, cols: usize, size: usize) -> usize {
33 Self::layout_prep_word_count() * n * cols * size * size_of::<<FFT64Avx as Backend>::ScalarPrep>()
34 }
35}
36
37unsafe impl ConvolutionImpl<Self> for FFT64Avx
38where
39 Module<Self>: ModuleN + VecZnxDftBytesOf + VecZnxDftApply<Self>,
40{
41 fn cnv_prepare_left_tmp_bytes_impl(module: &Module<Self>, res_size: usize, a_size: usize) -> usize {
42 module.bytes_of_vec_znx_dft(1, res_size.min(a_size))
43 }
44
45 fn cnv_prepare_left_impl<R, A>(module: &Module<Self>, res: &mut R, a: &A, scratch: &mut Scratch<Self>)
46 where
47 R: CnvPVecLToMut<Self>,
48 A: VecZnxToRef,
49 {
50 let res: &mut CnvPVecL<&mut [u8], FFT64Avx> = &mut res.to_mut();
51 let a: &VecZnx<&[u8]> = &a.to_ref();
52 let (mut tmp, _) = scratch.take_vec_znx_dft(module, 1, res.size().min(a.size()));
53 convolution_prepare_left(module.get_fft_table(), res, a, &mut tmp);
54 }
55
56 fn cnv_prepare_right_tmp_bytes_impl(module: &Module<Self>, res_size: usize, a_size: usize) -> usize {
57 module.bytes_of_vec_znx_dft(1, res_size.min(a_size))
58 }
59
60 fn cnv_prepare_right_impl<R, A>(module: &Module<Self>, res: &mut R, a: &A, scratch: &mut Scratch<Self>)
61 where
62 R: CnvPVecRToMut<Self>,
63 A: VecZnxToRef,
64 {
65 let res: &mut CnvPVecR<&mut [u8], FFT64Avx> = &mut res.to_mut();
66 let a: &VecZnx<&[u8]> = &a.to_ref();
67 let (mut tmp, _) = scratch.take_vec_znx_dft(module, 1, res.size().min(a.size()));
68 convolution_prepare_right(module.get_fft_table(), res, a, &mut tmp);
69 }
70 fn cnv_apply_dft_tmp_bytes_impl(
71 _module: &Module<Self>,
72 res_size: usize,
73 _res_offset: usize,
74 a_size: usize,
75 b_size: usize,
76 ) -> usize {
77 convolution_apply_dft_tmp_bytes(res_size, a_size, b_size)
78 }
79
80 fn cnv_by_const_apply_tmp_bytes_impl(
81 _module: &Module<Self>,
82 res_size: usize,
83 _res_offset: usize,
84 a_size: usize,
85 b_size: usize,
86 ) -> usize {
87 convolution_by_const_apply_tmp_bytes(res_size, a_size, b_size)
88 }
89
90 fn cnv_by_const_apply_impl<R, A>(
91 module: &Module<Self>,
92 res: &mut R,
93 res_offset: usize,
94 res_col: usize,
95 a: &A,
96 a_col: usize,
97 b: &[i64],
98 scratch: &mut Scratch<Self>,
99 ) where
100 R: VecZnxBigToMut<Self>,
101 A: VecZnxToRef,
102 {
103 let res: &mut VecZnxBig<&mut [u8], Self> = &mut res.to_mut();
104 let a: &VecZnx<&[u8]> = &a.to_ref();
105 let (tmp, _) =
106 scratch.take_slice(module.cnv_by_const_apply_tmp_bytes(res.size(), res_offset, a.size(), b.len()) / size_of::<i64>());
107 convolution_by_const_apply(res, res_offset, res_col, a, a_col, b, tmp);
108 }
109
110 fn cnv_apply_dft_impl<R, A, B>(
111 module: &Module<Self>,
112 res: &mut R,
113 res_offset: usize,
114 res_col: usize,
115 a: &A,
116 a_col: usize,
117 b: &B,
118 b_col: usize,
119 scratch: &mut Scratch<Self>,
120 ) where
121 R: VecZnxDftToMut<Self>,
122 A: CnvPVecLToRef<Self>,
123 B: CnvPVecRToRef<Self>,
124 {
125 let res: &mut VecZnxDft<&mut [u8], FFT64Avx> = &mut res.to_mut();
126 let a: &CnvPVecL<&[u8], FFT64Avx> = &a.to_ref();
127 let b: &CnvPVecR<&[u8], FFT64Avx> = &b.to_ref();
128 let (tmp, _) =
129 scratch.take_slice(module.cnv_apply_dft_tmp_bytes(res.size(), res_offset, a.size(), b.size()) / size_of::<f64>());
130 convolution_apply_dft(res, res_offset, res_col, a, a_col, b, b_col, tmp);
131 }
132
133 fn cnv_pairwise_apply_dft_tmp_bytes(
134 _module: &Module<Self>,
135 res_size: usize,
136 _res_offset: usize,
137 a_size: usize,
138 b_size: usize,
139 ) -> usize {
140 convolution_pairwise_apply_dft_tmp_bytes(res_size, a_size, b_size)
141 }
142
143 fn cnv_pairwise_apply_dft_impl<R, A, B>(
144 module: &Module<Self>,
145 res: &mut R,
146 res_offset: usize,
147 res_col: usize,
148 a: &A,
149 b: &B,
150 i: usize,
151 j: usize,
152 scratch: &mut Scratch<Self>,
153 ) where
154 R: VecZnxDftToMut<Self>,
155 A: CnvPVecLToRef<Self>,
156 B: CnvPVecRToRef<Self>,
157 {
158 let res: &mut VecZnxDft<&mut [u8], FFT64Avx> = &mut res.to_mut();
159 let a: &CnvPVecL<&[u8], FFT64Avx> = &a.to_ref();
160 let b: &CnvPVecR<&[u8], FFT64Avx> = &b.to_ref();
161 let (tmp, _) = scratch
162 .take_slice(module.cnv_pairwise_apply_dft_tmp_bytes(res.size(), res_offset, a.size(), b.size()) / size_of::<f64>());
163 convolution_pairwise_apply_dft(res, res_offset, res_col, a, b, i, j, tmp);
164 }
165}
166
167/// # Safety
168/// Caller must ensure the CPU supports AVX2.
169/// Assumes all inputs fit in i32 (so i32×i32→i64 is exact).
170#[target_feature(enable = "avx2")]
171pub unsafe fn i64_convolution_by_const_1coeff_avx(k: usize, dst: &mut [i64; 8], a: &[i64], a_size: usize, b: &[i64]) {
172 use core::arch::x86_64::{
173 __m256i, _mm256_add_epi64, _mm256_loadu_si256, _mm256_mul_epi32, _mm256_set1_epi32, _mm256_setzero_si256,
174 _mm256_storeu_si256,
175 };
176
177 dst.fill(0);
178
179 let b_size = b.len();
180 if k >= a_size + b_size {
181 return;
182 }
183
184 let j_min = k.saturating_sub(a_size - 1);
185 let j_max = (k + 1).min(b_size);
186
187 unsafe {
188 // Two accumulators = 8 outputs total
189 let mut acc_lo: __m256i = _mm256_setzero_si256(); // dst[0..4)
190 let mut acc_hi: __m256i = _mm256_setzero_si256(); // dst[4..8)
191
192 let mut a_ptr: *const i64 = a.as_ptr().add(8 * (k - j_min));
193 let mut b_ptr: *const i64 = b.as_ptr().add(j_min);
194
195 for _ in 0..(j_max - j_min) {
196 // Broadcast scalar b[j] as i32
197 let br: __m256i = _mm256_set1_epi32(*b_ptr as i32);
198
199 // ---- lower half: a[0..4) ----
200 let a_lo: __m256i = _mm256_loadu_si256(a_ptr as *const __m256i);
201
202 let prod_lo: __m256i = _mm256_mul_epi32(a_lo, br);
203
204 acc_lo = _mm256_add_epi64(acc_lo, prod_lo);
205
206 // ---- upper half: a[4..8) ----
207 let a_hi: __m256i = _mm256_loadu_si256(a_ptr.add(4) as *const __m256i);
208
209 let prod_hi: __m256i = _mm256_mul_epi32(a_hi, br);
210
211 acc_hi = _mm256_add_epi64(acc_hi, prod_hi);
212
213 a_ptr = a_ptr.sub(8);
214 b_ptr = b_ptr.add(1);
215 }
216
217 // Store final result
218 _mm256_storeu_si256(dst.as_mut_ptr() as *mut __m256i, acc_lo);
219 _mm256_storeu_si256(dst.as_mut_ptr().add(4) as *mut __m256i, acc_hi);
220 }
221}
222
223/// # Safety
224/// Caller must ensure the CPU supports AVX2.
225/// Assumes all values in `a` and `b` fit in i32 (so i32×i32→i64 is exact).
226#[target_feature(enable = "avx2")]
227pub unsafe fn i64_convolution_by_real_const_2coeffs_avx(
228 k: usize,
229 dst: &mut [i64; 16],
230 a: &[i64],
231 a_size: usize,
232 b: &[i64], // real scalars, stride-1
233) {
234 use core::arch::x86_64::{
235 __m256i, _mm256_add_epi64, _mm256_loadu_si256, _mm256_mul_epi32, _mm256_set1_epi32, _mm256_setzero_si256,
236 _mm256_storeu_si256,
237 };
238
239 let b_size: usize = b.len();
240
241 debug_assert!(a.len() >= 8 * a_size);
242
243 let k0: usize = k;
244 let k1: usize = k + 1;
245 let bound: usize = a_size + b_size;
246
247 if k0 >= bound {
248 unsafe {
249 let zero: __m256i = _mm256_setzero_si256();
250 let dst_ptr: *mut i64 = dst.as_mut_ptr();
251 _mm256_storeu_si256(dst_ptr as *mut __m256i, zero);
252 _mm256_storeu_si256(dst_ptr.add(4) as *mut __m256i, zero);
253 _mm256_storeu_si256(dst_ptr.add(8) as *mut __m256i, zero);
254 _mm256_storeu_si256(dst_ptr.add(12) as *mut __m256i, zero);
255 }
256 return;
257 }
258
259 unsafe {
260 let mut acc_lo_k0: __m256i = _mm256_setzero_si256();
261 let mut acc_hi_k0: __m256i = _mm256_setzero_si256();
262 let mut acc_lo_k1: __m256i = _mm256_setzero_si256();
263 let mut acc_hi_k1: __m256i = _mm256_setzero_si256();
264
265 let j0_min: usize = (k0 + 1).saturating_sub(a_size);
266 let j0_max: usize = (k0 + 1).min(b_size);
267
268 if k1 >= bound {
269 let mut a_k0_ptr: *const i64 = a.as_ptr().add(8 * (k0 - j0_min));
270 let mut b_ptr: *const i64 = b.as_ptr().add(j0_min);
271
272 // Contributions to k0 only
273 for _ in 0..j0_max - j0_min {
274 // Broadcast b[j] as i32
275 let br: __m256i = _mm256_set1_epi32(*b_ptr as i32);
276
277 // Load 4×i64 (low half) and 4×i64 (high half)
278 let a_lo_k0: __m256i = _mm256_loadu_si256(a_k0_ptr as *const __m256i);
279 let a_hi_k0: __m256i = _mm256_loadu_si256(a_k0_ptr.add(4) as *const __m256i);
280
281 acc_lo_k0 = _mm256_add_epi64(acc_lo_k0, _mm256_mul_epi32(a_lo_k0, br));
282 acc_hi_k0 = _mm256_add_epi64(acc_hi_k0, _mm256_mul_epi32(a_hi_k0, br));
283
284 a_k0_ptr = a_k0_ptr.sub(8);
285 b_ptr = b_ptr.add(1);
286 }
287 } else {
288 let j1_min: usize = (k1 + 1).saturating_sub(a_size);
289 let j1_max: usize = (k1 + 1).min(b_size);
290
291 let mut a_k0_ptr: *const i64 = a.as_ptr().add(8 * (k0 - j0_min));
292 let mut a_k1_ptr: *const i64 = a.as_ptr().add(8 * (k1 - j1_min));
293 let mut b_ptr: *const i64 = b.as_ptr().add(j0_min);
294
295 // Region 1: k0 only, j ∈ [j0_min, j1_min)
296 for _ in 0..j1_min - j0_min {
297 let br: __m256i = _mm256_set1_epi32(*b_ptr as i32);
298
299 let a_k0_lo: __m256i = _mm256_loadu_si256(a_k0_ptr as *const __m256i);
300 let a_k0_hi: __m256i = _mm256_loadu_si256(a_k0_ptr.add(4) as *const __m256i);
301
302 acc_lo_k0 = _mm256_add_epi64(acc_lo_k0, _mm256_mul_epi32(a_k0_lo, br));
303 acc_hi_k0 = _mm256_add_epi64(acc_hi_k0, _mm256_mul_epi32(a_k0_hi, br));
304
305 a_k0_ptr = a_k0_ptr.sub(8);
306 b_ptr = b_ptr.add(1);
307 }
308
309 // Region 2: overlap, contributions to both k0 and k1, j ∈ [j1_min, j0_max)
310 // Save one load on b: broadcast once and reuse.
311 for _ in 0..j0_max - j1_min {
312 let br: __m256i = _mm256_set1_epi32(*b_ptr as i32);
313
314 let a_lo_k0: __m256i = _mm256_loadu_si256(a_k0_ptr as *const __m256i);
315 let a_hi_k0: __m256i = _mm256_loadu_si256(a_k0_ptr.add(4) as *const __m256i);
316 let a_lo_k1: __m256i = _mm256_loadu_si256(a_k1_ptr as *const __m256i);
317 let a_hi_k1: __m256i = _mm256_loadu_si256(a_k1_ptr.add(4) as *const __m256i);
318
319 // k0
320 acc_lo_k0 = _mm256_add_epi64(acc_lo_k0, _mm256_mul_epi32(a_lo_k0, br));
321 acc_hi_k0 = _mm256_add_epi64(acc_hi_k0, _mm256_mul_epi32(a_hi_k0, br));
322
323 // k1
324 acc_lo_k1 = _mm256_add_epi64(acc_lo_k1, _mm256_mul_epi32(a_lo_k1, br));
325 acc_hi_k1 = _mm256_add_epi64(acc_hi_k1, _mm256_mul_epi32(a_hi_k1, br));
326
327 a_k0_ptr = a_k0_ptr.sub(8);
328 a_k1_ptr = a_k1_ptr.sub(8);
329 b_ptr = b_ptr.add(1);
330 }
331
332 // Region 3: k1 only, j ∈ [j0_max, j1_max)
333 for _ in 0..j1_max - j0_max {
334 let br: __m256i = _mm256_set1_epi32(*b_ptr as i32);
335
336 let a_lo_k1: __m256i = _mm256_loadu_si256(a_k1_ptr as *const __m256i);
337 let a_hi_k1: __m256i = _mm256_loadu_si256(a_k1_ptr.add(4) as *const __m256i);
338
339 acc_lo_k1 = _mm256_add_epi64(acc_lo_k1, _mm256_mul_epi32(a_lo_k1, br));
340 acc_hi_k1 = _mm256_add_epi64(acc_hi_k1, _mm256_mul_epi32(a_hi_k1, br));
341
342 a_k1_ptr = a_k1_ptr.sub(8);
343 b_ptr = b_ptr.add(1);
344 }
345 }
346
347 let dst_ptr: *mut i64 = dst.as_mut_ptr();
348 _mm256_storeu_si256(dst_ptr as *mut __m256i, acc_lo_k0);
349 _mm256_storeu_si256(dst_ptr.add(4) as *mut __m256i, acc_hi_k0);
350 _mm256_storeu_si256(dst_ptr.add(8) as *mut __m256i, acc_lo_k1);
351 _mm256_storeu_si256(dst_ptr.add(12) as *mut __m256i, acc_hi_k1);
352 }
353}
354
355/// # Safety
356/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
357#[target_feature(enable = "avx")]
358pub fn i64_extract_1blk_contiguous_avx(n: usize, offset: usize, rows: usize, blk: usize, dst: &mut [i64], src: &[i64]) {
359 use core::arch::x86_64::{__m256i, _mm256_loadu_si256, _mm256_storeu_si256};
360
361 unsafe {
362 let mut src_ptr: *const __m256i = src.as_ptr().add(offset + (blk << 3)) as *const __m256i; // src + 8*blk
363 let mut dst_ptr: *mut __m256i = dst.as_mut_ptr() as *mut __m256i;
364
365 let step: usize = n >> 2;
366
367 // Each iteration copies 8 i64; advance src by n i64 each row
368 for _ in 0..rows {
369 let v: __m256i = _mm256_loadu_si256(src_ptr);
370 _mm256_storeu_si256(dst_ptr, v);
371 let v: __m256i = _mm256_loadu_si256(src_ptr.add(1));
372 _mm256_storeu_si256(dst_ptr.add(1), v);
373 dst_ptr = dst_ptr.add(2);
374 src_ptr = src_ptr.add(step);
375 }
376 }
377}
378
379/// # Safety
380/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
381#[target_feature(enable = "avx")]
382pub fn i64_save_1blk_contiguous_avx(n: usize, offset: usize, rows: usize, blk: usize, dst: &mut [i64], src: &[i64]) {
383 use core::arch::x86_64::{__m256i, _mm256_loadu_si256, _mm256_storeu_si256};
384
385 unsafe {
386 let mut src_ptr: *const __m256i = src.as_ptr() as *const __m256i;
387 let mut dst_ptr: *mut __m256i = dst.as_mut_ptr().add(offset + (blk << 3)) as *mut __m256i; // dst + 8*blk
388
389 let step: usize = n >> 2;
390
391 // Each iteration copies 8 i64; advance dst by n i64 each row
392 for _ in 0..rows {
393 let v: __m256i = _mm256_loadu_si256(src_ptr);
394 _mm256_storeu_si256(dst_ptr, v);
395 let v: __m256i = _mm256_loadu_si256(src_ptr.add(1));
396 _mm256_storeu_si256(dst_ptr.add(1), v);
397 dst_ptr = dst_ptr.add(step);
398 src_ptr = src_ptr.add(2);
399 }
400 }
401}