1use super::bound::{Dim, Idx};
2use crate::internal_prelude::*;
3use core::marker::PhantomData;
4use faer_traits::SimdCapabilities;
5use pulp::Simd;
6pub struct SimdCtx<'N, T: ComplexField, S: Simd> {
7 pub ctx: T::SimdCtx<S>,
8 pub len: Dim<'N>,
9 offset: usize,
10 head_end: usize,
11 body_end: usize,
12 tail_end: usize,
13 head_mask: T::SimdMask<S>,
14 tail_mask: T::SimdMask<S>,
15 head_mem_mask: T::SimdMemMask<S>,
16 tail_mem_mask: T::SimdMemMask<S>,
17}
18impl<'N, T: ComplexField, S: Simd> core::fmt::Debug for SimdCtx<'N, T, S> {
19 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
20 f.debug_struct("SimdCtx")
21 .field("len", &self.len)
22 .field("offset", &self.offset)
23 .field("head_end", &self.head_end)
24 .field("body_end", &self.body_end)
25 .field("tail_end", &self.tail_end)
26 .field("head_mask", &self.head_mask)
27 .field("tail_mask", &self.tail_mask)
28 .finish_non_exhaustive()
29 }
30}
31impl<T: ComplexField, S: Simd> Copy for SimdCtx<'_, T, S> {}
32impl<T: ComplexField, S: Simd> Clone for SimdCtx<'_, T, S> {
33 #[inline]
34 fn clone(&self) -> Self {
35 *self
36 }
37}
38impl<T: ComplexField, S: Simd> core::ops::Deref for SimdCtx<'_, T, S> {
39 type Target = faer_traits::SimdCtx<T, S>;
40
41 #[inline(always)]
42 fn deref(&self) -> &Self::Target {
43 Self::Target::new(&self.ctx)
44 }
45}
46pub trait SimdIndex<'N, T: ComplexField, S: Simd> {
47 fn read(
48 simd: &SimdCtx<'N, T, S>,
49 slice: ColRef<'_, T, Dim<'N>, ContiguousFwd>,
50 index: Self,
51 ) -> T::SimdVec<S>;
52 fn write(
53 simd: &SimdCtx<'N, T, S>,
54 slice: ColMut<'_, T, Dim<'N>, ContiguousFwd>,
55 index: Self,
56 value: T::SimdVec<S>,
57 );
58}
59impl<'N, T: ComplexField, S: Simd> SimdIndex<'N, T, S> for SimdBody<'N, T, S> {
60 #[inline(always)]
61 fn read(
62 simd: &SimdCtx<'N, T, S>,
63 slice: ColRef<'_, T, Dim<'N>, ContiguousFwd>,
64 index: Self,
65 ) -> T::SimdVec<S> {
66 unsafe {
67 simd.load(
68 &*(slice.as_ptr().wrapping_offset(index.start)
69 as *const T::SimdVec<S>),
70 )
71 }
72 }
73
74 #[inline(always)]
75 fn write(
76 simd: &SimdCtx<'N, T, S>,
77 slice: ColMut<'_, T, Dim<'N>, ContiguousFwd>,
78 index: Self,
79 value: T::SimdVec<S>,
80 ) {
81 unsafe {
82 simd.store(
83 &mut *(slice.as_ptr_mut().wrapping_offset(index.start)
84 as *mut T::SimdVec<S>),
85 value,
86 );
87 }
88 }
89}
90impl<'N, T: ComplexField, S: Simd> SimdIndex<'N, T, S> for SimdHead<'N, T, S> {
91 #[inline(always)]
92 fn read(
93 simd: &SimdCtx<'N, T, S>,
94 slice: ColRef<'_, T, Dim<'N>, ContiguousFwd>,
95 index: Self,
96 ) -> T::SimdVec<S> {
97 unsafe {
98 simd.mask_load(
99 simd.head_mem_mask,
100 slice.as_ptr().wrapping_offset(index.start)
101 as *const T::SimdVec<S>,
102 )
103 }
104 }
105
106 #[inline(always)]
107 fn write(
108 simd: &SimdCtx<'N, T, S>,
109 slice: ColMut<'_, T, Dim<'N>, ContiguousFwd>,
110 index: Self,
111 value: T::SimdVec<S>,
112 ) {
113 unsafe {
114 simd.mask_store(
115 simd.head_mem_mask,
116 slice.as_ptr_mut().wrapping_offset(index.start)
117 as *mut T::SimdVec<S>,
118 value,
119 );
120 }
121 }
122}
123impl<'N, T: ComplexField, S: Simd> SimdIndex<'N, T, S> for SimdTail<'N, T, S> {
124 #[inline(always)]
125 fn read(
126 simd: &SimdCtx<'N, T, S>,
127 slice: ColRef<'_, T, Dim<'N>, ContiguousFwd>,
128 index: Self,
129 ) -> T::SimdVec<S> {
130 unsafe {
131 simd.mask_load(
132 simd.tail_mem_mask,
133 slice.as_ptr().wrapping_offset(index.start)
134 as *const T::SimdVec<S>,
135 )
136 }
137 }
138
139 #[inline(always)]
140 fn write(
141 simd: &SimdCtx<'N, T, S>,
142 slice: ColMut<'_, T, Dim<'N>, ContiguousFwd>,
143 index: Self,
144 value: T::SimdVec<S>,
145 ) {
146 unsafe {
147 simd.mask_store(
148 simd.tail_mem_mask,
149 slice.as_ptr_mut().wrapping_offset(index.start)
150 as *mut T::SimdVec<S>,
151 value,
152 );
153 }
154 }
155}
156impl<'N, T: ComplexField, S: Simd> SimdCtx<'N, T, S> {
157 #[inline(always)]
158 pub fn new(simd: T::SimdCtx<S>, len: Dim<'N>) -> Self {
159 core::assert!(
160 const { matches!(T::SIMD_CAPABILITIES, SimdCapabilities::Simd) }
161 );
162 let stride =
163 core::mem::size_of::<T::SimdVec<S>>() / core::mem::size_of::<T>();
164 let iota = T::simd_iota(&simd);
165 let head_start = T::simd_index_splat(&simd, T::Index::truncate(0));
166 let head_end = T::simd_index_splat(&simd, T::Index::truncate(0));
167 let tail_start = T::simd_index_splat(&simd, T::Index::truncate(0));
168 let tail_end =
169 T::simd_index_splat(&simd, T::Index::truncate(*len % stride));
170 Self {
171 ctx: simd,
172 len,
173 offset: 0,
174 head_end: 0,
175 body_end: *len / stride,
176 tail_end: (*len + stride - 1) / stride,
177 head_mask: T::simd_and_mask(
178 &simd,
179 T::simd_index_greater_than_or_equal(&simd, iota, head_start),
180 T::simd_index_less_than(&simd, iota, head_end),
181 ),
182 tail_mask: T::simd_and_mask(
183 &simd,
184 T::simd_index_greater_than_or_equal(&simd, iota, tail_start),
185 T::simd_index_less_than(&simd, iota, tail_end),
186 ),
187 head_mem_mask: T::simd_mem_mask_between(
188 &simd,
189 T::Index::truncate(0),
190 T::Index::truncate(0),
191 ),
192 tail_mem_mask: T::simd_mem_mask_between(
193 &simd,
194 T::Index::truncate(0),
195 T::Index::truncate(*len % stride),
196 ),
197 }
198 }
199
200 #[inline(always)]
201 pub fn new_align(
202 simd: T::SimdCtx<S>,
203 len: Dim<'N>,
204 align_offset: usize,
205 ) -> Self {
206 core::assert!(
207 const { matches!(T::SIMD_CAPABILITIES, SimdCapabilities::Simd) }
208 );
209 let stride =
210 core::mem::size_of::<T::SimdVec<S>>() / core::mem::size_of::<T>();
211 let align_offset = align_offset % stride;
212 let iota = T::simd_iota(&simd);
213 if align_offset == 0 {
214 Self::new(simd, len)
215 } else {
216 let offset = stride - align_offset;
217 let full_len = offset + *len;
218 let head_start =
219 T::simd_index_splat(&simd, T::Index::truncate(offset));
220 let head_end =
221 T::simd_index_splat(&simd, T::Index::truncate(stride));
222 let tail_start = T::simd_index_splat(&simd, T::Index::truncate(0));
223 let tail_end = T::simd_index_splat(
224 &simd,
225 T::Index::truncate(full_len % stride),
226 );
227 if align_offset <= *len {
228 Self {
229 ctx: simd,
230 len,
231 offset,
232 head_end: 1,
233 body_end: full_len / stride,
234 tail_end: (full_len + stride - 1) / stride,
235 head_mask: T::simd_and_mask(
236 &simd,
237 T::simd_index_greater_than_or_equal(
238 &simd, iota, head_start,
239 ),
240 T::simd_index_less_than(&simd, iota, head_end),
241 ),
242 tail_mask: T::simd_and_mask(
243 &simd,
244 T::simd_index_greater_than_or_equal(
245 &simd, iota, tail_start,
246 ),
247 T::simd_index_less_than(&simd, iota, tail_end),
248 ),
249 head_mem_mask: T::simd_mem_mask_between(
250 &simd,
251 T::Index::truncate(offset),
252 T::Index::truncate(stride),
253 ),
254 tail_mem_mask: T::simd_mem_mask_between(
255 &simd,
256 T::Index::truncate(0),
257 T::Index::truncate(full_len % stride),
258 ),
259 }
260 } else {
261 let head_start =
262 T::simd_index_splat(&simd, T::Index::truncate(offset));
263 let head_end = T::simd_index_splat(
264 &simd,
265 T::Index::truncate(full_len % stride),
266 );
267 let tail_start =
268 T::simd_index_splat(&simd, T::Index::truncate(0));
269 let tail_end =
270 T::simd_index_splat(&simd, T::Index::truncate(0));
271 Self {
272 ctx: simd,
273 len,
274 offset,
275 head_end: 1,
276 body_end: 1,
277 tail_end: 1,
278 head_mask: T::simd_and_mask(
279 &simd,
280 T::simd_index_greater_than_or_equal(
281 &simd, iota, head_start,
282 ),
283 T::simd_index_less_than(&simd, iota, head_end),
284 ),
285 tail_mask: T::simd_and_mask(
286 &simd,
287 T::simd_index_greater_than_or_equal(
288 &simd, iota, tail_start,
289 ),
290 T::simd_index_less_than(&simd, iota, tail_end),
291 ),
292 head_mem_mask: T::simd_mem_mask_between(
293 &simd,
294 T::Index::truncate(offset),
295 T::Index::truncate(full_len % stride),
296 ),
297 tail_mem_mask: T::simd_mem_mask_between(
298 &simd,
299 T::Index::truncate(0),
300 T::Index::truncate(0),
301 ),
302 }
303 }
304 }
305 }
306
307 #[inline]
308 pub fn offset(&self) -> usize {
309 self.offset
310 }
311
312 #[inline(always)]
313 pub fn new_force_mask(simd: T::SimdCtx<S>, len: Dim<'N>) -> Self {
314 core::assert!(
315 const { matches!(T::SIMD_CAPABILITIES, SimdCapabilities::Simd) }
316 );
317 crate::assert!(*len != 0);
318 let new_len = *len - 1;
319 let stride =
320 core::mem::size_of::<T::SimdVec<S>>() / core::mem::size_of::<T>();
321 let iota = T::simd_iota(&simd);
322 let head_start = T::simd_index_splat(&simd, T::Index::truncate(0));
323 let head_end = T::simd_index_splat(&simd, T::Index::truncate(0));
324 let tail_start = T::simd_index_splat(&simd, T::Index::truncate(0));
325 let tail_end = T::simd_index_splat(
326 &simd,
327 T::Index::truncate((new_len % stride) + 1),
328 );
329 Self {
330 ctx: simd,
331 len,
332 offset: 0,
333 head_end: 0,
334 body_end: new_len / stride,
335 tail_end: new_len / stride + 1,
336 head_mask: T::simd_and_mask(
337 &simd,
338 T::simd_index_greater_than_or_equal(&simd, iota, head_start),
339 T::simd_index_less_than(&simd, iota, head_end),
340 ),
341 tail_mask: T::simd_and_mask(
342 &simd,
343 T::simd_index_greater_than_or_equal(&simd, iota, tail_start),
344 T::simd_index_less_than(&simd, iota, tail_end),
345 ),
346 head_mem_mask: T::simd_mem_mask_between(
347 &simd,
348 T::Index::truncate(0),
349 T::Index::truncate(0),
350 ),
351 tail_mem_mask: T::simd_mem_mask_between(
352 &simd,
353 T::Index::truncate(0),
354 T::Index::truncate((new_len % stride) + 1),
355 ),
356 }
357 }
358
359 #[inline(always)]
360 pub fn read<I: SimdIndex<'N, T, S>>(
361 &self,
362 slice: ColRef<'_, T, Dim<'N>, ContiguousFwd>,
363 index: I,
364 ) -> T::SimdVec<S> {
365 I::read(self, slice, index)
366 }
367
368 #[inline(always)]
369 pub fn write<I: SimdIndex<'N, T, S>>(
370 &self,
371 slice: ColMut<'_, T, Dim<'N>, ContiguousFwd>,
372 index: I,
373 value: T::SimdVec<S>,
374 ) {
375 I::write(self, slice, index, value)
376 }
377
378 #[inline(always)]
379 pub fn head_mask(&self) -> T::SimdMask<S> {
380 self.head_mask
381 }
382
383 #[inline(always)]
384 pub fn tail_mask(&self) -> T::SimdMask<S> {
385 self.tail_mask
386 }
387
388 #[inline]
389 pub fn indices(
390 &self,
391 ) -> (
392 Option<SimdHead<'N, T, S>>,
393 impl Clone
394 + ExactSizeIterator
395 + DoubleEndedIterator<Item = SimdBody<'N, T, S>>,
396 Option<SimdTail<'N, T, S>>,
397 ) {
398 macro_rules! stride {
399 () => {
400 core::mem::size_of::<T::SimdVec<S>>()
401 / core::mem::size_of::<T>()
402 };
403 }
404 let offset = -(self.offset as isize);
405 (
406 if 0 == self.head_end {
407 None
408 } else {
409 Some(SimdHead {
410 start: offset,
411 mask: PhantomData,
412 })
413 },
414 (self.head_end..self.body_end).map(
415 #[inline(always)]
416 move |i| SimdBody {
417 start: offset + (i * stride!()) as isize,
418 mask: PhantomData,
419 },
420 ),
421 if self.body_end == self.tail_end {
422 None
423 } else {
424 Some(SimdTail {
425 start: offset + (self.body_end * stride!()) as isize,
426 mask: PhantomData,
427 })
428 },
429 )
430 }
431
432 #[inline]
433 pub fn batch_indices<const BATCH: usize>(
434 &self,
435 ) -> (
436 Option<SimdHead<'N, T, S>>,
437 impl Clone
438 + ExactSizeIterator
439 + DoubleEndedIterator<Item = [SimdBody<'N, T, S>; BATCH]>,
440 impl Clone
441 + ExactSizeIterator
442 + DoubleEndedIterator<Item = SimdBody<'N, T, S>>,
443 Option<SimdTail<'N, T, S>>,
444 ) {
445 macro_rules! stride {
446 () => {
447 core::mem::size_of::<T::SimdVec<S>>()
448 / core::mem::size_of::<T>()
449 };
450 }
451 let len = self.body_end - self.head_end;
452 let offset = -(self.offset as isize);
453 (
454 if 0 == self.head_end {
455 None
456 } else {
457 Some(SimdHead {
458 start: offset,
459 mask: PhantomData,
460 })
461 },
462 (self.head_end..self.head_end + len / BATCH * BATCH)
463 .map(move |i| {
464 core::array::from_fn(
465 #[inline(always)]
466 |k| SimdBody {
467 start: offset + ((i + k) * stride!()) as isize,
468 mask: PhantomData,
469 },
470 )
471 })
472 .step_by(BATCH),
473 (self.head_end + len / BATCH * BATCH..self.body_end).map(
474 #[inline(always)]
475 move |i| SimdBody {
476 start: offset + (i * stride!()) as isize,
477 mask: PhantomData,
478 },
479 ),
480 if self.body_end == self.tail_end {
481 None
482 } else {
483 Some(SimdTail {
484 start: offset + (self.body_end * stride!()) as isize,
485 mask: PhantomData,
486 })
487 },
488 )
489 }
490}
491#[repr(transparent)]
492#[derive(Debug)]
493pub struct SimdBody<'N, T: ComplexField, S: Simd> {
494 start: isize,
495 mask: PhantomData<(Idx<'N>, T::SimdMask<S>)>,
496}
497impl<T: ComplexField, S: Simd> SimdBody<'_, T, S> {
498 pub fn offset(&self) -> isize {
499 self.start
500 }
501}
502#[repr(transparent)]
503#[derive(Debug)]
504pub struct SimdHead<'N, T: ComplexField, S: Simd> {
505 start: isize,
506 mask: PhantomData<(Idx<'N>, T::SimdMask<S>)>,
507}
508#[repr(transparent)]
509#[derive(Debug)]
510pub struct SimdTail<'N, T: ComplexField, S: Simd> {
511 start: isize,
512 mask: PhantomData<(Idx<'N>, T::SimdMask<S>)>,
513}
514impl<'N, T: ComplexField, S: Simd> Copy for SimdBody<'N, T, S> {}
515impl<'N, T: ComplexField, S: Simd> Clone for SimdBody<'N, T, S> {
516 #[inline]
517 fn clone(&self) -> Self {
518 *self
519 }
520}
521impl<'N, T: ComplexField, S: Simd> Copy for SimdHead<'N, T, S> {}
522impl<'N, T: ComplexField, S: Simd> Clone for SimdHead<'N, T, S> {
523 #[inline]
524 fn clone(&self) -> Self {
525 *self
526 }
527}
528impl<'N, T: ComplexField, S: Simd> Copy for SimdTail<'N, T, S> {}
529impl<'N, T: ComplexField, S: Simd> Clone for SimdTail<'N, T, S> {
530 #[inline]
531 fn clone(&self) -> Self {
532 *self
533 }
534}