Skip to main content

strided_basic/
ops_view.rs

1//! High-level operations on dynamic-rank strided views.
2
3use crate::kernel::{
4    build_plan_fused, build_plan_fused_small, ensure_same_shape, for_each_inner_block_preordered,
5    same_contiguous_layout, sequential_contiguous_layout, total_len, SMALL_TENSOR_THRESHOLD,
6};
7use crate::map_view::{copy_map_into, map_into, zip_map2_into};
8use crate::reduce_view::reduce;
9use crate::simd;
10use crate::view::{StridedView, StridedViewMut};
11use crate::{Exec, Result, StridedError};
12use num_traits::{One, Zero};
13use std::mem::MaybeUninit;
14use std::ops::{Add, Mul};
15use strided_view::{ElementOp, ElementOpApply};
16
17use crate::cost::KernelClass;
18use crate::fuse::compute_costs;
19use crate::threading::{
20    for_each_inner_block_with_offsets, is_parallel, mapreduce_threaded, SendPtr,
21};
22// ============================================================================
23// Stride-specialized inner loop helpers for ops_view
24//
25// When all inner strides are 1 (contiguous in the innermost dimension),
26// we use slice-based iteration so LLVM can auto-vectorize effectively.
27// This mirrors the inner_loop_map* helpers in map_view.rs.
28// ============================================================================
29
30// INVARIANT: validated layouts keep all dereferences in bounds. Wrapping
31// advances permit an unused final cursor outside a reversed/gapped allocation.
32/// Inner loop for add: `dst[i] += Op::apply(src[i])`.
33#[inline(always)]
34unsafe fn inner_loop_add<D: Copy + Add<S, Output = D>, S: Copy, Op: ElementOp<S>>(
35    dp: *mut D,
36    ds: isize,
37    sp: *const S,
38    ss: isize,
39    len: usize,
40) {
41    if ds == 1 && ss == 1 {
42        let dst = std::slice::from_raw_parts_mut(dp, len);
43        let src = std::slice::from_raw_parts(sp, len);
44        simd::dispatch_if_large(len, || {
45            for i in 0..len {
46                dst[i] = dst[i] + Op::apply(src[i]);
47            }
48        });
49    } else {
50        let mut dp = dp;
51        let mut sp = sp;
52        for _ in 0..len {
53            *dp = *dp + Op::apply(*sp);
54            dp = dp.wrapping_offset(ds);
55            sp = sp.wrapping_offset(ss);
56        }
57    }
58}
59
60/// Inner loop for mul: `dst[i] *= Op::apply(src[i])`.
61#[inline(always)]
62unsafe fn inner_loop_mul<D: Copy + Mul<S, Output = D>, S: Copy, Op: ElementOp<S>>(
63    dp: *mut D,
64    ds: isize,
65    sp: *const S,
66    ss: isize,
67    len: usize,
68) {
69    if ds == 1 && ss == 1 {
70        let dst = std::slice::from_raw_parts_mut(dp, len);
71        let src = std::slice::from_raw_parts(sp, len);
72        simd::dispatch_if_large(len, || {
73            for i in 0..len {
74                dst[i] = dst[i] * Op::apply(src[i]);
75            }
76        });
77    } else {
78        let mut dp = dp;
79        let mut sp = sp;
80        for _ in 0..len {
81            *dp = *dp * Op::apply(*sp);
82            dp = dp.wrapping_offset(ds);
83            sp = sp.wrapping_offset(ss);
84        }
85    }
86}
87
88/// Inner loop for axpy: `dst[i] = alpha * Op::apply(src[i]) + dst[i]`.
89#[inline(always)]
90unsafe fn inner_loop_axpy<
91    D: Copy + Add<D, Output = D>,
92    S: Copy,
93    A: Copy + Mul<S, Output = D>,
94    Op: ElementOp<S>,
95>(
96    dp: *mut D,
97    ds: isize,
98    sp: *const S,
99    ss: isize,
100    len: usize,
101    alpha: A,
102) {
103    if ds == 1 && ss == 1 {
104        let dst = std::slice::from_raw_parts_mut(dp, len);
105        let src = std::slice::from_raw_parts(sp, len);
106        simd::dispatch_if_large(len, || {
107            for i in 0..len {
108                dst[i] = alpha * Op::apply(src[i]) + dst[i];
109            }
110        });
111    } else {
112        let mut dp = dp;
113        let mut sp = sp;
114        for _ in 0..len {
115            *dp = alpha * Op::apply(*sp) + *dp;
116            dp = dp.wrapping_offset(ds);
117            sp = sp.wrapping_offset(ss);
118        }
119    }
120}
121
122/// Inner loop for fma: `dst[i] += OpA::apply(a[i]) * OpB::apply(b[i])`.
123#[inline(always)]
124unsafe fn inner_loop_fma<
125    D: Copy + Add<D, Output = D>,
126    A: Copy + Mul<B, Output = D>,
127    B: Copy,
128    OpA: ElementOp<A>,
129    OpB: ElementOp<B>,
130>(
131    dp: *mut D,
132    ds: isize,
133    ap: *const A,
134    a_s: isize,
135    bp: *const B,
136    b_s: isize,
137    len: usize,
138) {
139    if ds == 1 && a_s == 1 && b_s == 1 {
140        let dst = std::slice::from_raw_parts_mut(dp, len);
141        let sa = std::slice::from_raw_parts(ap, len);
142        let sb = std::slice::from_raw_parts(bp, len);
143        simd::dispatch_if_large(len, || {
144            for i in 0..len {
145                dst[i] = dst[i] + OpA::apply(sa[i]) * OpB::apply(sb[i]);
146            }
147        });
148    } else {
149        let mut dp = dp;
150        let mut ap = ap;
151        let mut bp = bp;
152        for _ in 0..len {
153            *dp = *dp + OpA::apply(*ap) * OpB::apply(*bp);
154            dp = dp.wrapping_offset(ds);
155            ap = ap.wrapping_offset(a_s);
156            bp = bp.wrapping_offset(b_s);
157        }
158    }
159}
160
161/// Inner loop for dot: `acc += OpA::apply(a[i]) * OpB::apply(b[i])`.
162#[inline(always)]
163unsafe fn inner_loop_dot<
164    A: Copy + Mul<B, Output = R>,
165    B: Copy,
166    R: Copy + Add<R, Output = R>,
167    OpA: ElementOp<A>,
168    OpB: ElementOp<B>,
169>(
170    ap: *const A,
171    a_s: isize,
172    bp: *const B,
173    b_s: isize,
174    len: usize,
175    mut acc: R,
176) -> R {
177    if a_s == 1 && b_s == 1 {
178        let sa = std::slice::from_raw_parts(ap, len);
179        let sb = std::slice::from_raw_parts(bp, len);
180        simd::dispatch_if_large(len, || {
181            for i in 0..len {
182                acc = acc + OpA::apply(sa[i]) * OpB::apply(sb[i]);
183            }
184        });
185    } else {
186        let mut ap = ap;
187        let mut bp = bp;
188        for _ in 0..len {
189            acc = acc + OpA::apply(*ap) * OpB::apply(*bp);
190            ap = ap.wrapping_offset(a_s);
191            bp = bp.wrapping_offset(b_s);
192        }
193    }
194    acc
195}
196
197/// Copy elements from source to destination: `dest[i] = src[i]`.
198pub fn copy_into<T: Copy + Send + Sync, Op: ElementOp<T>>(
199    dest: &mut StridedViewMut<T>,
200    src: &StridedView<T, Op>,
201    exec: Exec<'_>,
202) -> Result<()> {
203    ensure_same_shape(dest.dims(), src.dims())?;
204
205    let dst_ptr = dest.as_mut_ptr();
206    let src_ptr = src.ptr();
207    let dst_dims = dest.dims();
208    let dst_strides = dest.strides();
209    let src_strides = src.strides();
210
211    if sequential_contiguous_layout(
212        KernelClass::Copy,
213        dst_dims,
214        &[dst_strides, src_strides],
215        &exec,
216    )?
217    .is_some()
218    {
219        let len = total_len(dst_dims)?;
220        if Op::IS_IDENTITY {
221            debug_assert!(
222                {
223                    let nbytes = len
224                        .checked_mul(std::mem::size_of::<T>())
225                        .expect("copy size must not overflow");
226                    let dst_start = dst_ptr as usize;
227                    let src_start = src_ptr as usize;
228                    let dst_end = dst_start.saturating_add(nbytes);
229                    let src_end = src_start.saturating_add(nbytes);
230                    dst_end <= src_start || src_end <= dst_start
231                },
232                "overlapping src/dest is not supported"
233            );
234            unsafe { std::ptr::copy_nonoverlapping(src_ptr, dst_ptr, len) };
235        } else {
236            let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
237            let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
238            simd::dispatch_if_large(len, || {
239                for i in 0..len {
240                    dst[i] = Op::apply(src[i]);
241                }
242            });
243        }
244        return Ok(());
245    }
246
247    copy_map_into(dest, src, |x| x, exec)
248}
249
250/// Copy initialized values into a potentially uninitialized destination.
251///
252/// On success every logical destination element is initialized; holes are
253/// untouched. Accepts identity views; use `map_into` for element operations
254/// such as conjugation. Equal contiguous f32/f64 layouts and every other type
255/// keep the element map; a non-contiguous f32/f64 copy uses the permutation
256/// engine, split within the budget of `exec`.
257///
258/// # Errors
259/// Returns a shape/rank error for mismatched views, `NonInjectiveOutputLayout`
260/// for overlapping destination elements, or `OffsetOverflow` for size overflow.
261///
262/// # Examples
263/// ```
264/// use std::mem::MaybeUninit;
265/// use strided_basic::{copy_into_uninit, Exec, StridedView, StridedViewMut};
266/// let src = [1.0_f64, 2.0];
267/// let mut dst = [MaybeUninit::uninit(); 2];
268/// copy_into_uninit(
269///     &mut StridedViewMut::new(&mut dst, &[2], &[1], 0).unwrap(),
270///     &StridedView::new(&src, &[2], &[1], 0).unwrap(),
271///     Exec::serial(),
272/// ).unwrap();
273/// // SAFETY: the successful copy initialized both destination elements.
274/// assert_eq!(unsafe { dst[1].assume_init() }, 2.0);
275/// ```
276pub fn copy_into_uninit<T: Copy + Send + Sync + 'static>(
277    dest: &mut StridedViewMut<MaybeUninit<T>>,
278    src: &StridedView<T>,
279    exec: Exec<'_>,
280) -> Result<()> {
281    ensure_same_shape(dest.dims(), src.dims())?;
282    if !crate::layout_check::is_injective_layout(dest.dims(), dest.strides()) {
283        return Err(StridedError::NonInjectiveOutputLayout);
284    }
285    if src.is_empty() {
286        return Ok(());
287    }
288    let _len = src
289        .dims()
290        .iter()
291        .try_fold(1usize, |n, &dim| n.checked_mul(dim))
292        .ok_or(StridedError::OffsetOverflow)?;
293    // The permutation engine's 4/8-byte paths reinterpret storage as native
294    // floats. Restrict that path to those exact types: arbitrary Copy types can
295    // contain uninitialized padding or have weaker alignment (e.g. Complex32).
296    let native_float = std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>()
297        || std::any::TypeId::of::<T>() == std::any::TypeId::of::<f64>();
298    // Equal contiguous layouts go to the element map, whose contiguous path
299    // parallelizes a large copy; the permutation engine has only a serial
300    // contiguous shortcut. Non-contiguous native-float copies use the
301    // permutation engine, serial or parallel.
302    if !native_float
303        || same_contiguous_layout(dest.dims(), &[dest.strides(), src.strides()]).is_some()
304    {
305        return copy_map_into(dest, src, MaybeUninit::new, exec);
306    }
307    // INVARIANT: the permutation engine selects its float paths by the
308    // source element type and only stores into the destination, so an
309    // uninitialized destination with an initialized source is valid.
310    crate::threading::copy_permuted_uninit(dest, src, exec)
311}
312
313/// Copy elements from `src` to `dst`, optimized for col-major destination.
314///
315/// Delegates to the centralized permutation-copy dispatch for the actual work.
316pub fn copy_into_col_major<T: Copy + Send + Sync + 'static>(
317    dst: &mut StridedViewMut<T>,
318    src: &StridedView<T>,
319) -> Result<()> {
320    strided_perm::copy_into_col_major(dst, src, Exec::serial())
321}
322
323/// Element-wise addition: `dest[i] += src[i]`.
324///
325/// Source may have a different element type from destination.
326pub fn add<D: Copy + Add<S, Output = D> + Send + Sync, S: Copy + Send + Sync, Op: ElementOp<S>>(
327    dest: &mut StridedViewMut<D>,
328    src: &StridedView<S, Op>,
329    exec: Exec<'_>,
330) -> Result<()> {
331    ensure_same_shape(dest.dims(), src.dims())?;
332
333    let dst_ptr = dest.as_mut_ptr();
334    let src_ptr = src.ptr();
335    let dst_dims = dest.dims();
336    let dst_strides = dest.strides();
337    let src_strides = src.strides();
338
339    if sequential_contiguous_layout(
340        KernelClass::Map,
341        dst_dims,
342        &[dst_strides, src_strides],
343        &exec,
344    )?
345    .is_some()
346    {
347        let len = total_len(dst_dims)?;
348        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
349        let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
350        simd::dispatch_if_large(len, || {
351            for i in 0..len {
352                dst[i] = dst[i] + Op::apply(src[i]);
353            }
354        });
355        return Ok(());
356    }
357
358    let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
359    let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<S>());
360
361    let total = total_len(dst_dims)?;
362    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
363        build_plan_fused_small(dst_dims, &strides_list)
364    } else {
365        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
366    };
367
368    let total = total_len(&fused_dims)?;
369    if is_parallel(KernelClass::Map, total, &exec) {
370        let dst_send = SendPtr(dst_ptr);
371        let src_send = SendPtr(src_ptr as *mut S);
372
373        let costs = compute_costs(&ordered_strides);
374        let initial_offsets: crate::inline::OperandVec<isize> =
375            strides_list.iter().map(|_| 0isize).collect();
376        return mapreduce_threaded(
377            &fused_dims,
378            &plan.block,
379            &ordered_strides,
380            &initial_offsets,
381            &costs,
382            KernelClass::Map,
383            exec,
384            &|dims, blocks, strides_list, offsets| {
385                for_each_inner_block_with_offsets(
386                    dims,
387                    blocks,
388                    strides_list,
389                    offsets,
390                    |offsets, len, strides| {
391                        unsafe {
392                            inner_loop_add::<D, S, Op>(
393                                dst_send.as_ptr().offset(offsets[0]),
394                                strides[0],
395                                src_send.as_const().offset(offsets[1]),
396                                strides[1],
397                                len,
398                            )
399                        };
400                        Ok(())
401                    },
402                )
403            },
404        );
405    }
406
407    let initial_offsets: crate::inline::OperandVec<isize> =
408        ordered_strides.iter().map(|_| 0isize).collect();
409    for_each_inner_block_preordered(
410        &fused_dims,
411        &plan.block,
412        &ordered_strides,
413        &initial_offsets,
414        |offsets, len, strides| {
415            unsafe {
416                inner_loop_add::<D, S, Op>(
417                    dst_ptr.offset(offsets[0]),
418                    strides[0],
419                    src_ptr.offset(offsets[1]),
420                    strides[1],
421                    len,
422                )
423            };
424            Ok(())
425        },
426    )
427}
428
429/// Element-wise multiplication: `dest[i] *= src[i]`.
430///
431/// Source may have a different element type from destination.
432pub fn mul<D: Copy + Mul<S, Output = D> + Send + Sync, S: Copy + Send + Sync, Op: ElementOp<S>>(
433    dest: &mut StridedViewMut<D>,
434    src: &StridedView<S, Op>,
435    exec: Exec<'_>,
436) -> Result<()> {
437    ensure_same_shape(dest.dims(), src.dims())?;
438
439    let dst_ptr = dest.as_mut_ptr();
440    let src_ptr = src.ptr();
441    let dst_dims = dest.dims();
442    let dst_strides = dest.strides();
443    let src_strides = src.strides();
444
445    if sequential_contiguous_layout(
446        KernelClass::Map,
447        dst_dims,
448        &[dst_strides, src_strides],
449        &exec,
450    )?
451    .is_some()
452    {
453        let len = total_len(dst_dims)?;
454        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
455        let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
456        simd::dispatch_if_large(len, || {
457            for i in 0..len {
458                dst[i] = dst[i] * Op::apply(src[i]);
459            }
460        });
461        return Ok(());
462    }
463
464    let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
465    let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<S>());
466
467    let total = total_len(dst_dims)?;
468    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
469        build_plan_fused_small(dst_dims, &strides_list)
470    } else {
471        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
472    };
473
474    let total = total_len(&fused_dims)?;
475    if is_parallel(KernelClass::Map, total, &exec) {
476        let dst_send = SendPtr(dst_ptr);
477        let src_send = SendPtr(src_ptr as *mut S);
478
479        let costs = compute_costs(&ordered_strides);
480        let initial_offsets: crate::inline::OperandVec<isize> =
481            strides_list.iter().map(|_| 0isize).collect();
482        return mapreduce_threaded(
483            &fused_dims,
484            &plan.block,
485            &ordered_strides,
486            &initial_offsets,
487            &costs,
488            KernelClass::Map,
489            exec,
490            &|dims, blocks, strides_list, offsets| {
491                for_each_inner_block_with_offsets(
492                    dims,
493                    blocks,
494                    strides_list,
495                    offsets,
496                    |offsets, len, strides| {
497                        unsafe {
498                            inner_loop_mul::<D, S, Op>(
499                                dst_send.as_ptr().offset(offsets[0]),
500                                strides[0],
501                                src_send.as_const().offset(offsets[1]),
502                                strides[1],
503                                len,
504                            )
505                        };
506                        Ok(())
507                    },
508                )
509            },
510        );
511    }
512
513    let initial_offsets: crate::inline::OperandVec<isize> =
514        ordered_strides.iter().map(|_| 0isize).collect();
515    for_each_inner_block_preordered(
516        &fused_dims,
517        &plan.block,
518        &ordered_strides,
519        &initial_offsets,
520        |offsets, len, strides| {
521            unsafe {
522                inner_loop_mul::<D, S, Op>(
523                    dst_ptr.offset(offsets[0]),
524                    strides[0],
525                    src_ptr.offset(offsets[1]),
526                    strides[1],
527                    len,
528                )
529            };
530            Ok(())
531        },
532    )
533}
534
535/// AXPY: `dest[i] = alpha * src[i] + dest[i]`.
536///
537/// Alpha, source, and destination may have different element types.
538pub fn axpy<D, S, A, Op>(
539    dest: &mut StridedViewMut<D>,
540    src: &StridedView<S, Op>,
541    alpha: A,
542    exec: Exec<'_>,
543) -> Result<()>
544where
545    A: Copy + Mul<S, Output = D> + Sync,
546    D: Copy + Add<D, Output = D> + Send + Sync,
547    S: Copy + Send + Sync,
548    Op: ElementOp<S>,
549{
550    ensure_same_shape(dest.dims(), src.dims())?;
551
552    let dst_ptr = dest.as_mut_ptr();
553    let src_ptr = src.ptr();
554    let dst_dims = dest.dims();
555    let dst_strides = dest.strides();
556    let src_strides = src.strides();
557
558    if sequential_contiguous_layout(
559        KernelClass::Map,
560        dst_dims,
561        &[dst_strides, src_strides],
562        &exec,
563    )?
564    .is_some()
565    {
566        let len = total_len(dst_dims)?;
567        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
568        let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
569        simd::dispatch_if_large(len, || {
570            for i in 0..len {
571                dst[i] = alpha * Op::apply(src[i]) + dst[i];
572            }
573        });
574        return Ok(());
575    }
576
577    let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
578    let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<S>());
579
580    let total = total_len(dst_dims)?;
581    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
582        build_plan_fused_small(dst_dims, &strides_list)
583    } else {
584        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
585    };
586
587    let total = total_len(&fused_dims)?;
588    if is_parallel(KernelClass::Map, total, &exec) {
589        let dst_send = SendPtr(dst_ptr);
590        let src_send = SendPtr(src_ptr as *mut S);
591
592        let costs = compute_costs(&ordered_strides);
593        let initial_offsets: crate::inline::OperandVec<isize> =
594            strides_list.iter().map(|_| 0isize).collect();
595        return mapreduce_threaded(
596            &fused_dims,
597            &plan.block,
598            &ordered_strides,
599            &initial_offsets,
600            &costs,
601            KernelClass::Map,
602            exec,
603            &|dims, blocks, strides_list, offsets| {
604                for_each_inner_block_with_offsets(
605                    dims,
606                    blocks,
607                    strides_list,
608                    offsets,
609                    |offsets, len, strides| {
610                        unsafe {
611                            inner_loop_axpy::<D, S, A, Op>(
612                                dst_send.as_ptr().offset(offsets[0]),
613                                strides[0],
614                                src_send.as_const().offset(offsets[1]),
615                                strides[1],
616                                len,
617                                alpha,
618                            )
619                        };
620                        Ok(())
621                    },
622                )
623            },
624        );
625    }
626
627    let initial_offsets: crate::inline::OperandVec<isize> =
628        ordered_strides.iter().map(|_| 0isize).collect();
629    for_each_inner_block_preordered(
630        &fused_dims,
631        &plan.block,
632        &ordered_strides,
633        &initial_offsets,
634        |offsets, len, strides| {
635            unsafe {
636                inner_loop_axpy::<D, S, A, Op>(
637                    dst_ptr.offset(offsets[0]),
638                    strides[0],
639                    src_ptr.offset(offsets[1]),
640                    strides[1],
641                    len,
642                    alpha,
643                )
644            };
645            Ok(())
646        },
647    )
648}
649
650/// Fused multiply-add: `dest[i] += OpA::apply(a[i]) * OpB::apply(b[i])`.
651///
652/// Operands may have different element types. Element operations are applied lazily.
653pub fn fma<D, A, B, OpA, OpB>(
654    dest: &mut StridedViewMut<D>,
655    a: &StridedView<A, OpA>,
656    b: &StridedView<B, OpB>,
657    exec: Exec<'_>,
658) -> Result<()>
659where
660    A: Copy + Mul<B, Output = D> + Send + Sync,
661    B: Copy + Send + Sync,
662    D: Copy + Add<D, Output = D> + Send + Sync,
663    OpA: ElementOp<A>,
664    OpB: ElementOp<B>,
665{
666    ensure_same_shape(dest.dims(), a.dims())?;
667    ensure_same_shape(dest.dims(), b.dims())?;
668
669    let dst_ptr = dest.as_mut_ptr();
670    let a_ptr = a.ptr();
671    let b_ptr = b.ptr();
672    let dst_dims = dest.dims();
673    let dst_strides = dest.strides();
674    let a_strides = a.strides();
675    let b_strides = b.strides();
676
677    if sequential_contiguous_layout(
678        KernelClass::Map,
679        dst_dims,
680        &[dst_strides, a_strides, b_strides],
681        &exec,
682    )?
683    .is_some()
684    {
685        let len = total_len(dst_dims)?;
686        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
687        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
688        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
689        simd::dispatch_if_large(len, || {
690            for i in 0..len {
691                dst[i] = dst[i] + OpA::apply(sa[i]) * OpB::apply(sb[i]);
692            }
693        });
694        return Ok(());
695    }
696
697    let strides_list: [&[isize]; 3] = [dst_strides, a_strides, b_strides];
698    let elem_size = std::mem::size_of::<D>()
699        .max(std::mem::size_of::<A>())
700        .max(std::mem::size_of::<B>());
701
702    let total = total_len(dst_dims)?;
703    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
704        build_plan_fused_small(dst_dims, &strides_list)
705    } else {
706        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
707    };
708
709    let total = total_len(&fused_dims)?;
710    if is_parallel(KernelClass::Map, total, &exec) {
711        let dst_send = SendPtr(dst_ptr);
712        let a_send = SendPtr(a_ptr as *mut A);
713        let b_send = SendPtr(b_ptr as *mut B);
714
715        let costs = compute_costs(&ordered_strides);
716        let initial_offsets: crate::inline::OperandVec<isize> =
717            strides_list.iter().map(|_| 0isize).collect();
718        return mapreduce_threaded(
719            &fused_dims,
720            &plan.block,
721            &ordered_strides,
722            &initial_offsets,
723            &costs,
724            KernelClass::Map,
725            exec,
726            &|dims, blocks, strides_list, offsets| {
727                for_each_inner_block_with_offsets(
728                    dims,
729                    blocks,
730                    strides_list,
731                    offsets,
732                    |offsets, len, strides| {
733                        unsafe {
734                            inner_loop_fma::<D, A, B, OpA, OpB>(
735                                dst_send.as_ptr().offset(offsets[0]),
736                                strides[0],
737                                a_send.as_const().offset(offsets[1]),
738                                strides[1],
739                                b_send.as_const().offset(offsets[2]),
740                                strides[2],
741                                len,
742                            )
743                        };
744                        Ok(())
745                    },
746                )
747            },
748        );
749    }
750
751    let initial_offsets: crate::inline::OperandVec<isize> =
752        ordered_strides.iter().map(|_| 0isize).collect();
753    for_each_inner_block_preordered(
754        &fused_dims,
755        &plan.block,
756        &ordered_strides,
757        &initial_offsets,
758        |offsets, len, strides| {
759            unsafe {
760                inner_loop_fma::<D, A, B, OpA, OpB>(
761                    dst_ptr.offset(offsets[0]),
762                    strides[0],
763                    a_ptr.offset(offsets[1]),
764                    strides[1],
765                    b_ptr.offset(offsets[2]),
766                    strides[2],
767                    len,
768                )
769            };
770            Ok(())
771        },
772    )
773}
774
775/// Sum all elements: `sum(src)`.
776///
777/// # Evaluation order
778///
779/// The result is independent of the pool when the call runs with a serial
780/// token ([`Exec::serial`]) or a budget of one. With a larger budget a
781/// contiguous identity operand above the threading threshold is summed per
782/// task and the partial sums are combined in a tree, so the result is
783/// repeatable for a fixed budget but can differ between budgets.
784///
785/// Neither path is a strict left to right sum: a contiguous run is folded with
786/// independent SIMD or scalar lanes (see [`reduce`] for the lane scheme of the
787/// generic path). The last bits can therefore differ from a naive loop, and no
788/// bitwise equality across machines or target features is guaranteed.
789///
790/// For a deterministic accumulation of `f32` data in `f64`, use
791/// [`zip_reduce`](crate::zip_reduce) with an `f64` map and a `0.0f64` init.
792///
793/// # Allocation
794///
795/// At rank ≤ [`INLINE_RANK`](strided_view::INLINE_RANK) the view metadata and
796/// the traversal plan are stored inline, so a call performs no heap allocation.
797pub fn sum<
798    T: Copy + Zero + Add<Output = T> + Send + Sync + simd::MaybeSimdOps,
799    Op: ElementOp<T>,
800>(
801    src: &StridedView<T, Op>,
802    exec: Exec<'_>,
803) -> Result<T> {
804    // A zero-element sum is the additive identity. Return before any planning
805    // so an empty view with an enormous stride never reaches block sizing.
806    if total_len(src.dims())? == 0 {
807        return Ok(T::zero());
808    }
809    // SIMD fast path: contiguous Identity view with SIMD support
810    if Op::IS_IDENTITY {
811        if same_contiguous_layout(src.dims(), &[src.strides()]).is_some() {
812            let len = total_len(src.dims())?;
813            let src_slice = unsafe { std::slice::from_raw_parts(src.ptr(), len) };
814
815            // The partial sums of the halves are combined in a tree; a type
816            // without SIMD support takes the generic reduction below.
817            if T::try_simd_sum(&[]).is_some() && is_parallel(KernelClass::Reduce, len, &exec) {
818                return Ok(crate::threading::parallel_map_reduce(
819                    0..len,
820                    crate::cost::for_len_tree(exec, KernelClass::Reduce, len),
821                    &|range| T::try_simd_sum(&src_slice[range]).unwrap(),
822                    &|left, right| left + right,
823                ));
824            }
825
826            if let Some(result) = T::try_simd_sum(src_slice) {
827                return Ok(result);
828            }
829        }
830    }
831    reduce(src, |x| x, |a, b| a + b, T::zero(), exec)
832}
833
834/// Dot product: `sum(OpA::apply(a[i]) * OpB::apply(b[i]))`.
835///
836/// Operands may have different element types. Result type `R` must be `A * B`.
837/// SIMD fast path fires only when `A == B == R` (same type) and both Identity ops.
838///
839/// A widening dot — `f32` operands with an `f64` accumulator — is not
840/// expressible here because `R` is pinned to `A * B`; use
841/// [`zip_reduce`](crate::zip_reduce) with a
842/// `|x, y| x as f64 * y as f64` map instead.
843///
844/// # Evaluation order
845///
846/// This entry takes no parallel branch and no token, so its result does not
847/// depend on any pool: it is computed on the calling thread. The homogeneous contiguous fast path folds with the SIMD
848/// lane scheme of the element type; every other layout folds a single
849/// accumulator along the traversal order of the fused plan, which is left to
850/// right over each unit-stride run. Neither is guaranteed to be bitwise equal
851/// across machines or target features.
852///
853/// # Allocation
854///
855/// At rank ≤ [`INLINE_RANK`](strided_view::INLINE_RANK) the view metadata and
856/// the traversal plan are stored inline, so a call performs no heap allocation.
857pub fn dot<A, B, R, OpA, OpB>(a: &StridedView<A, OpA>, b: &StridedView<B, OpB>) -> Result<R>
858where
859    A: Copy + Mul<B, Output = R> + Send + Sync + 'static,
860    B: Copy + Send + Sync + 'static,
861    R: Copy + Zero + Add<Output = R> + Send + Sync + simd::MaybeSimdOps + 'static,
862    OpA: ElementOp<A>,
863    OpB: ElementOp<B>,
864{
865    ensure_same_shape(a.dims(), b.dims())?;
866
867    let a_ptr = a.ptr();
868    let b_ptr = b.ptr();
869    let a_strides = a.strides();
870    let b_strides = b.strides();
871    let a_dims = a.dims();
872    let len = total_len(a_dims)?;
873
874    // A zero-element dot product is the additive identity. Return before any
875    // planning so an empty view with an enormous stride never reaches block
876    // sizing.
877    if len == 0 {
878        return Ok(R::zero());
879    }
880
881    if same_contiguous_layout(a_dims, &[a_strides, b_strides]).is_some() {
882        // SIMD fast path: both contiguous, both Identity ops, same type
883        if OpA::IS_IDENTITY
884            && OpB::IS_IDENTITY
885            && std::any::TypeId::of::<A>() == std::any::TypeId::of::<R>()
886            && std::any::TypeId::of::<B>() == std::any::TypeId::of::<R>()
887        {
888            let sa = unsafe { std::slice::from_raw_parts(a_ptr as *const R, len) };
889            let sb = unsafe { std::slice::from_raw_parts(b_ptr as *const R, len) };
890            if let Some(result) = R::try_simd_dot(sa, sb) {
891                return Ok(result);
892            }
893        }
894
895        // Generic contiguous fast path
896        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
897        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
898        let mut acc = R::zero();
899        simd::dispatch_if_large(len, || {
900            for i in 0..len {
901                acc = acc + OpA::apply(sa[i]) * OpB::apply(sb[i]);
902            }
903        });
904        return Ok(acc);
905    }
906
907    let strides_list: [&[isize]; 2] = [a_strides, b_strides];
908    let elem_size = std::mem::size_of::<A>()
909        .max(std::mem::size_of::<B>())
910        .max(std::mem::size_of::<R>());
911
912    let total = total_len(a_dims)?;
913    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
914        build_plan_fused_small(a_dims, &strides_list)
915    } else {
916        build_plan_fused(a_dims, &strides_list, None, elem_size)
917    };
918
919    let mut acc = R::zero();
920    let initial_offsets: crate::inline::OperandVec<isize> =
921        ordered_strides.iter().map(|_| 0isize).collect();
922    for_each_inner_block_preordered(
923        &fused_dims,
924        &plan.block,
925        &ordered_strides,
926        &initial_offsets,
927        |offsets, len, strides| {
928            acc = unsafe {
929                inner_loop_dot::<A, B, R, OpA, OpB>(
930                    a_ptr.offset(offsets[0]),
931                    strides[0],
932                    b_ptr.offset(offsets[1]),
933                    strides[1],
934                    len,
935                    acc,
936                )
937            };
938            Ok(())
939        },
940    )?;
941
942    Ok(acc)
943}
944
945/// Symmetrize a square matrix: `dest = (src + src^T) / 2`.
946pub fn symmetrize_into<T>(
947    dest: &mut StridedViewMut<T>,
948    src: &StridedView<T>,
949    exec: Exec<'_>,
950) -> Result<()>
951where
952    T: Copy
953        + Add<Output = T>
954        + Mul<Output = T>
955        + num_traits::FromPrimitive
956        + std::ops::Div<Output = T>
957        + Send
958        + Sync,
959{
960    if src.ndim() != 2 {
961        return Err(StridedError::RankMismatch(src.ndim(), 2));
962    }
963    let rows = src.dims()[0];
964    let cols = src.dims()[1];
965    if rows != cols {
966        return Err(StridedError::NonSquare { rows, cols });
967    }
968
969    let src_t = src.permute(&[1, 0])?;
970    let half = T::from_f64(0.5).ok_or(StridedError::ScalarConversion)?;
971
972    zip_map2_into(dest, src, &src_t, |a, b| (a + b) * half, exec)
973}
974
975/// Conjugate-symmetrize a square matrix: `dest = (src + conj(src^T)) / 2`.
976pub fn symmetrize_conj_into<T>(
977    dest: &mut StridedViewMut<T>,
978    src: &StridedView<T>,
979    exec: Exec<'_>,
980) -> Result<()>
981where
982    T: Copy
983        + ElementOpApply
984        + Add<Output = T>
985        + Mul<Output = T>
986        + num_traits::FromPrimitive
987        + std::ops::Div<Output = T>
988        + Send
989        + Sync,
990{
991    if src.ndim() != 2 {
992        return Err(StridedError::RankMismatch(src.ndim(), 2));
993    }
994    let rows = src.dims()[0];
995    let cols = src.dims()[1];
996    if rows != cols {
997        return Err(StridedError::NonSquare { rows, cols });
998    }
999
1000    // adjoint = conj + transpose
1001    let src_adj = src.adjoint_2d()?;
1002    let half = T::from_f64(0.5).ok_or(StridedError::ScalarConversion)?;
1003
1004    zip_map2_into(dest, src, &src_adj, |a, b| (a + b) * half, exec)
1005}
1006
1007/// Copy with scaling: `dest[i] = scale * src[i]`.
1008///
1009/// Scale, source, and destination may have different element types.
1010pub fn copy_scale<D, S, A, Op>(
1011    dest: &mut StridedViewMut<D>,
1012    src: &StridedView<S, Op>,
1013    scale: A,
1014    exec: Exec<'_>,
1015) -> Result<()>
1016where
1017    A: Copy + Mul<S, Output = D> + Sync,
1018    D: Copy + Send + Sync,
1019    S: Copy + Send + Sync,
1020    Op: ElementOp<S>,
1021{
1022    map_into(dest, src, |x| scale * x, exec)
1023}
1024
1025/// Copy with complex conjugation: `dest[i] = conj(src[i])`.
1026pub fn copy_conj<T: Copy + ElementOpApply + Send + Sync>(
1027    dest: &mut StridedViewMut<T>,
1028    src: &StridedView<T>,
1029    exec: Exec<'_>,
1030) -> Result<()> {
1031    let src_conj = src.conj();
1032    copy_into(dest, &src_conj, exec)
1033}
1034
1035#[inline]
1036fn element_transpose_is_identity<T: 'static>() -> bool {
1037    use std::any::TypeId;
1038
1039    macro_rules! matches_type {
1040        ($($ty:ty),* $(,)?) => {{
1041            let id = TypeId::of::<T>();
1042            false $(|| id == TypeId::of::<$ty>())*
1043        }};
1044    }
1045
1046    matches_type!(
1047        f32,
1048        f64,
1049        i8,
1050        i16,
1051        i32,
1052        i64,
1053        i128,
1054        isize,
1055        u8,
1056        u16,
1057        u32,
1058        u64,
1059        u128,
1060        usize,
1061        num_complex::Complex32,
1062        num_complex::Complex64,
1063    )
1064}
1065
1066#[inline]
1067fn element_zero_is_all_bits_zero<T: 'static>() -> bool {
1068    use std::any::TypeId;
1069
1070    macro_rules! matches_type {
1071        ($($ty:ty),* $(,)?) => {{
1072            let id = TypeId::of::<T>();
1073            false $(|| id == TypeId::of::<$ty>())*
1074        }};
1075    }
1076
1077    matches_type!(
1078        f32,
1079        f64,
1080        i8,
1081        i16,
1082        i32,
1083        i64,
1084        i128,
1085        isize,
1086        u8,
1087        u16,
1088        u32,
1089        u64,
1090        u128,
1091        usize,
1092        num_complex::Complex32,
1093        num_complex::Complex64,
1094    )
1095}
1096
1097#[inline]
1098unsafe fn fill_2d<T: Copy + Send + Sync>(
1099    dst: *mut T,
1100    dim0: usize,
1101    dim1: usize,
1102    dst_stride0: isize,
1103    dst_stride1: isize,
1104    value: T,
1105    exec: Exec<'_>,
1106) {
1107    let total = dim0.saturating_mul(dim1);
1108    let exec = crate::cost::for_len(exec, KernelClass::Copy, total);
1109    if !exec.is_serial() {
1110        let dst_send = SendPtr(dst);
1111        if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1112            crate::threading::parallel_for_each(0..dim1, exec, &|columns| {
1113                for j in columns {
1114                    let dst = dst_send.as_ptr();
1115                    unsafe {
1116                        let base = j as isize * dst_stride1;
1117                        for i in 0..dim0 {
1118                            *dst.offset(base + i as isize * dst_stride0) = value;
1119                        }
1120                    }
1121                }
1122            });
1123        } else {
1124            crate::threading::parallel_for_each(0..dim0, exec, &|rows| {
1125                for i in rows {
1126                    let dst = dst_send.as_ptr();
1127                    unsafe {
1128                        let base = i as isize * dst_stride0;
1129                        for j in 0..dim1 {
1130                            *dst.offset(base + j as isize * dst_stride1) = value;
1131                        }
1132                    }
1133                }
1134            });
1135        }
1136        return;
1137    }
1138
1139    if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1140        for j in 0..dim1 {
1141            let base = j as isize * dst_stride1;
1142            for i in 0..dim0 {
1143                *dst.offset(base + i as isize * dst_stride0) = value;
1144            }
1145        }
1146    } else {
1147        for i in 0..dim0 {
1148            let base = i as isize * dst_stride0;
1149            for j in 0..dim1 {
1150                *dst.offset(base + j as isize * dst_stride1) = value;
1151            }
1152        }
1153    }
1154}
1155
1156#[inline]
1157unsafe fn fill_contiguous<T>(dst: *mut T, len: usize, value: T)
1158where
1159    T: Copy + Zero + PartialEq + Send + Sync + 'static,
1160{
1161    if element_zero_is_all_bits_zero::<T>() && value == T::zero() {
1162        std::ptr::write_bytes(dst, 0, len);
1163        return;
1164    }
1165
1166    let dst = std::slice::from_raw_parts_mut(dst, len);
1167    dst.fill(value);
1168}
1169
1170#[inline(always)]
1171unsafe fn transpose_scale_4x4_f64(
1172    dst: *mut f64,
1173    dst_stride0: isize,
1174    dst_stride1: isize,
1175    src: *const f64,
1176    src_stride0: isize,
1177    src_stride1: isize,
1178    i: usize,
1179    j: usize,
1180    scale: f64,
1181) {
1182    let src_base = src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1183
1184    let s00 = *src_base;
1185    let s10 = *src_base.offset(src_stride0);
1186    let s20 = *src_base.offset(2 * src_stride0);
1187    let s30 = *src_base.offset(3 * src_stride0);
1188
1189    let src_col1 = src_base.offset(src_stride1);
1190    let s01 = *src_col1;
1191    let s11 = *src_col1.offset(src_stride0);
1192    let s21 = *src_col1.offset(2 * src_stride0);
1193    let s31 = *src_col1.offset(3 * src_stride0);
1194
1195    let src_col2 = src_base.offset(2 * src_stride1);
1196    let s02 = *src_col2;
1197    let s12 = *src_col2.offset(src_stride0);
1198    let s22 = *src_col2.offset(2 * src_stride0);
1199    let s32 = *src_col2.offset(3 * src_stride0);
1200
1201    let src_col3 = src_base.offset(3 * src_stride1);
1202    let s03 = *src_col3;
1203    let s13 = *src_col3.offset(src_stride0);
1204    let s23 = *src_col3.offset(2 * src_stride0);
1205    let s33 = *src_col3.offset(3 * src_stride0);
1206
1207    let dst_row0 = dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1);
1208    *dst_row0 = scale * s00;
1209    *dst_row0.offset(dst_stride0) = scale * s01;
1210    *dst_row0.offset(2 * dst_stride0) = scale * s02;
1211    *dst_row0.offset(3 * dst_stride0) = scale * s03;
1212
1213    let dst_row1 = dst_row0.offset(dst_stride1);
1214    *dst_row1 = scale * s10;
1215    *dst_row1.offset(dst_stride0) = scale * s11;
1216    *dst_row1.offset(2 * dst_stride0) = scale * s12;
1217    *dst_row1.offset(3 * dst_stride0) = scale * s13;
1218
1219    let dst_row2 = dst_row0.offset(2 * dst_stride1);
1220    *dst_row2 = scale * s20;
1221    *dst_row2.offset(dst_stride0) = scale * s21;
1222    *dst_row2.offset(2 * dst_stride0) = scale * s22;
1223    *dst_row2.offset(3 * dst_stride0) = scale * s23;
1224
1225    let dst_row3 = dst_row0.offset(3 * dst_stride1);
1226    *dst_row3 = scale * s30;
1227    *dst_row3.offset(dst_stride0) = scale * s31;
1228    *dst_row3.offset(2 * dst_stride0) = scale * s32;
1229    *dst_row3.offset(3 * dst_stride0) = scale * s33;
1230}
1231
1232#[inline]
1233#[allow(clippy::too_many_arguments)]
1234unsafe fn copy_transpose_scale_2d_f64_tiled_raw(
1235    dst: *mut f64,
1236    dst_stride0: isize,
1237    dst_stride1: isize,
1238    src: *const f64,
1239    src_stride0: isize,
1240    src_stride1: isize,
1241    src_rows: usize,
1242    src_cols: usize,
1243    scale: f64,
1244    exec: Exec<'_>,
1245) {
1246    const TILE: usize = 4;
1247    let row_full = src_rows / TILE * TILE;
1248    let col_full = src_cols / TILE * TILE;
1249
1250    let total = src_rows.saturating_mul(src_cols);
1251    let exec = crate::cost::for_len(exec, KernelClass::Map, total);
1252    if !exec.is_serial() {
1253        let dst_send = SendPtr(dst);
1254        let src_send = SendPtr(src as *mut f64);
1255        let row_tiles = row_full / TILE;
1256        crate::threading::parallel_for_each(0..row_tiles, exec, &|tiles| {
1257            for tile_i in tiles {
1258                let i = tile_i * TILE;
1259                let dst = dst_send.as_ptr();
1260                let src = src_send.as_const();
1261                unsafe {
1262                    let mut j = 0;
1263                    while j < col_full {
1264                        transpose_scale_4x4_f64(
1265                            dst,
1266                            dst_stride0,
1267                            dst_stride1,
1268                            src,
1269                            src_stride0,
1270                            src_stride1,
1271                            i,
1272                            j,
1273                            scale,
1274                        );
1275                        j += TILE;
1276                    }
1277                    for j in col_full..src_cols {
1278                        for ii in i..i + TILE {
1279                            *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1280                                scale
1281                                    * *src.offset(
1282                                        ii as isize * src_stride0 + j as isize * src_stride1,
1283                                    );
1284                        }
1285                    }
1286                }
1287            }
1288        });
1289        for i in row_full..src_rows {
1290            for j in 0..src_cols {
1291                *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1292                    scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1293            }
1294        }
1295        return;
1296    }
1297
1298    let mut i = 0;
1299    while i < row_full {
1300        let mut j = 0;
1301        while j < col_full {
1302            transpose_scale_4x4_f64(
1303                dst,
1304                dst_stride0,
1305                dst_stride1,
1306                src,
1307                src_stride0,
1308                src_stride1,
1309                i,
1310                j,
1311                scale,
1312            );
1313            j += TILE;
1314        }
1315        for j in col_full..src_cols {
1316            for ii in i..i + TILE {
1317                *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1318                    scale * *src.offset(ii as isize * src_stride0 + j as isize * src_stride1);
1319            }
1320        }
1321        i += TILE;
1322    }
1323    for i in row_full..src_rows {
1324        for j in 0..src_cols {
1325            *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1326                scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1327        }
1328    }
1329}
1330
1331#[inline]
1332#[cfg(test)]
1333unsafe fn try_copy_transpose_scale_2d_f64_tiled(
1334    dest: &mut StridedViewMut<f64>,
1335    src: &StridedView<f64>,
1336    scale: f64,
1337    exec: Exec<'_>,
1338) -> bool {
1339    if src.ndim() != 2 || dest.ndim() != 2 {
1340        return false;
1341    }
1342    let src_dims = src.dims();
1343    if dest.dims() != [src_dims[1], src_dims[0]] {
1344        return false;
1345    }
1346    if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1347        return false;
1348    }
1349
1350    copy_transpose_scale_2d_f64_tiled_raw(
1351        dest.as_mut_ptr(),
1352        dest.strides()[0],
1353        dest.strides()[1],
1354        src.ptr(),
1355        src.strides()[0],
1356        src.strides()[1],
1357        src_dims[0],
1358        src_dims[1],
1359        scale,
1360        exec,
1361    );
1362    true
1363}
1364
1365#[inline]
1366unsafe fn try_copy_transpose_scale_2d_f64_tiled_typed<T>(
1367    dest: &mut StridedViewMut<T>,
1368    src: &StridedView<T>,
1369    scale: T,
1370    exec: Exec<'_>,
1371) -> bool
1372where
1373    T: Copy + 'static,
1374{
1375    if std::any::TypeId::of::<T>() != std::any::TypeId::of::<f64>() {
1376        return false;
1377    }
1378
1379    let scale = *(&scale as *const T).cast::<f64>();
1380    if src.ndim() != 2 || dest.ndim() != 2 {
1381        return false;
1382    }
1383    let src_dims = src.dims();
1384    if dest.dims() != [src_dims[1], src_dims[0]] {
1385        return false;
1386    }
1387    if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1388        return false;
1389    }
1390
1391    copy_transpose_scale_2d_f64_tiled_raw(
1392        dest.as_mut_ptr().cast::<f64>(),
1393        dest.strides()[0],
1394        dest.strides()[1],
1395        src.ptr().cast::<f64>(),
1396        src.strides()[0],
1397        src.strides()[1],
1398        src_dims[0],
1399        src_dims[1],
1400        scale,
1401        exec,
1402    );
1403    true
1404}
1405
1406#[inline(always)]
1407unsafe fn transpose_scale_4x4_identity<T>(
1408    dst: *mut T,
1409    dst_stride0: isize,
1410    dst_stride1: isize,
1411    src: *const T,
1412    src_stride0: isize,
1413    src_stride1: isize,
1414    i: usize,
1415    j: usize,
1416    scale: T,
1417) where
1418    T: Copy + Mul<Output = T>,
1419{
1420    let src_base = src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1421
1422    let s00 = *src_base;
1423    let s10 = *src_base.offset(src_stride0);
1424    let s20 = *src_base.offset(2 * src_stride0);
1425    let s30 = *src_base.offset(3 * src_stride0);
1426
1427    let src_col1 = src_base.offset(src_stride1);
1428    let s01 = *src_col1;
1429    let s11 = *src_col1.offset(src_stride0);
1430    let s21 = *src_col1.offset(2 * src_stride0);
1431    let s31 = *src_col1.offset(3 * src_stride0);
1432
1433    let src_col2 = src_base.offset(2 * src_stride1);
1434    let s02 = *src_col2;
1435    let s12 = *src_col2.offset(src_stride0);
1436    let s22 = *src_col2.offset(2 * src_stride0);
1437    let s32 = *src_col2.offset(3 * src_stride0);
1438
1439    let src_col3 = src_base.offset(3 * src_stride1);
1440    let s03 = *src_col3;
1441    let s13 = *src_col3.offset(src_stride0);
1442    let s23 = *src_col3.offset(2 * src_stride0);
1443    let s33 = *src_col3.offset(3 * src_stride0);
1444
1445    let dst_row0 = dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1);
1446    *dst_row0 = scale * s00;
1447    *dst_row0.offset(dst_stride0) = scale * s01;
1448    *dst_row0.offset(2 * dst_stride0) = scale * s02;
1449    *dst_row0.offset(3 * dst_stride0) = scale * s03;
1450
1451    let dst_row1 = dst_row0.offset(dst_stride1);
1452    *dst_row1 = scale * s10;
1453    *dst_row1.offset(dst_stride0) = scale * s11;
1454    *dst_row1.offset(2 * dst_stride0) = scale * s12;
1455    *dst_row1.offset(3 * dst_stride0) = scale * s13;
1456
1457    let dst_row2 = dst_row0.offset(2 * dst_stride1);
1458    *dst_row2 = scale * s20;
1459    *dst_row2.offset(dst_stride0) = scale * s21;
1460    *dst_row2.offset(2 * dst_stride0) = scale * s22;
1461    *dst_row2.offset(3 * dst_stride0) = scale * s23;
1462
1463    let dst_row3 = dst_row0.offset(3 * dst_stride1);
1464    *dst_row3 = scale * s30;
1465    *dst_row3.offset(dst_stride0) = scale * s31;
1466    *dst_row3.offset(2 * dst_stride0) = scale * s32;
1467    *dst_row3.offset(3 * dst_stride0) = scale * s33;
1468}
1469
1470#[inline]
1471#[allow(clippy::too_many_arguments)]
1472unsafe fn copy_transpose_scale_2d_identity_tiled_raw<T>(
1473    dst: *mut T,
1474    dst_stride0: isize,
1475    dst_stride1: isize,
1476    src: *const T,
1477    src_stride0: isize,
1478    src_stride1: isize,
1479    src_rows: usize,
1480    src_cols: usize,
1481    scale: T,
1482    exec: Exec<'_>,
1483) where
1484    T: Copy + Mul<Output = T> + Send + Sync,
1485{
1486    const TILE: usize = 4;
1487    let row_full = src_rows / TILE * TILE;
1488    let col_full = src_cols / TILE * TILE;
1489
1490    let total = src_rows.saturating_mul(src_cols);
1491    let exec = crate::cost::for_len(exec, KernelClass::Map, total);
1492    if !exec.is_serial() {
1493        let dst_send = SendPtr(dst);
1494        let src_send = SendPtr(src as *mut T);
1495        let row_tiles = row_full / TILE;
1496        crate::threading::parallel_for_each(0..row_tiles, exec, &|tiles| {
1497            for tile_i in tiles {
1498                let i = tile_i * TILE;
1499                let dst = dst_send.as_ptr();
1500                let src = src_send.as_const();
1501                unsafe {
1502                    let mut j = 0;
1503                    while j < col_full {
1504                        transpose_scale_4x4_identity(
1505                            dst,
1506                            dst_stride0,
1507                            dst_stride1,
1508                            src,
1509                            src_stride0,
1510                            src_stride1,
1511                            i,
1512                            j,
1513                            scale,
1514                        );
1515                        j += TILE;
1516                    }
1517                    for j in col_full..src_cols {
1518                        for ii in i..i + TILE {
1519                            *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1520                                scale
1521                                    * *src.offset(
1522                                        ii as isize * src_stride0 + j as isize * src_stride1,
1523                                    );
1524                        }
1525                    }
1526                }
1527            }
1528        });
1529        for i in row_full..src_rows {
1530            for j in 0..src_cols {
1531                *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1532                    scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1533            }
1534        }
1535        return;
1536    }
1537
1538    let mut i = 0;
1539    while i < row_full {
1540        let mut j = 0;
1541        while j < col_full {
1542            transpose_scale_4x4_identity(
1543                dst,
1544                dst_stride0,
1545                dst_stride1,
1546                src,
1547                src_stride0,
1548                src_stride1,
1549                i,
1550                j,
1551                scale,
1552            );
1553            j += TILE;
1554        }
1555        for j in col_full..src_cols {
1556            for ii in i..i + TILE {
1557                *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1558                    scale * *src.offset(ii as isize * src_stride0 + j as isize * src_stride1);
1559            }
1560        }
1561        i += TILE;
1562    }
1563    for i in row_full..src_rows {
1564        for j in 0..src_cols {
1565            *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1566                scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1567        }
1568    }
1569}
1570
1571#[inline]
1572unsafe fn try_copy_transpose_scale_2d_identity_tiled<T>(
1573    dest: &mut StridedViewMut<T>,
1574    src: &StridedView<T>,
1575    scale: T,
1576    exec: Exec<'_>,
1577) -> bool
1578where
1579    T: Copy + Mul<Output = T> + Send + Sync,
1580{
1581    if src.ndim() != 2 || dest.ndim() != 2 {
1582        return false;
1583    }
1584    let src_dims = src.dims();
1585    if dest.dims() != [src_dims[1], src_dims[0]] {
1586        return false;
1587    }
1588    if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1589        return false;
1590    }
1591
1592    copy_transpose_scale_2d_identity_tiled_raw(
1593        dest.as_mut_ptr(),
1594        dest.strides()[0],
1595        dest.strides()[1],
1596        src.ptr(),
1597        src.strides()[0],
1598        src.strides()[1],
1599        src_dims[0],
1600        src_dims[1],
1601        scale,
1602        exec,
1603    );
1604    true
1605}
1606
1607#[inline]
1608#[allow(clippy::too_many_arguments)]
1609unsafe fn copy_transpose_scale_2d_loop<T>(
1610    dst: *mut T,
1611    dst_stride0: isize,
1612    dst_stride1: isize,
1613    src: *const T,
1614    src_stride0: isize,
1615    src_stride1: isize,
1616    src_rows: usize,
1617    src_cols: usize,
1618    scale: T,
1619    exec: Exec<'_>,
1620) where
1621    T: Copy + ElementOpApply + Mul<Output = T> + Send + Sync,
1622{
1623    let total = src_rows.saturating_mul(src_cols);
1624    let exec = crate::cost::for_len(exec, KernelClass::Map, total);
1625    if !exec.is_serial() {
1626        let dst_send = SendPtr(dst);
1627        let src_send = SendPtr(src as *mut T);
1628        if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1629            crate::threading::parallel_for_each(0..src_rows, exec, &|rows| {
1630                for i in rows {
1631                    let dst = dst_send.as_ptr();
1632                    let src = src_send.as_const();
1633                    unsafe {
1634                        for j in 0..src_cols {
1635                            let value = (*src
1636                                .offset(i as isize * src_stride0 + j as isize * src_stride1))
1637                            .transpose();
1638                            *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1639                                scale * value;
1640                        }
1641                    }
1642                }
1643            });
1644        } else {
1645            crate::threading::parallel_for_each(0..src_cols, exec, &|columns| {
1646                for j in columns {
1647                    let dst = dst_send.as_ptr();
1648                    let src = src_send.as_const();
1649                    unsafe {
1650                        for i in 0..src_rows {
1651                            let value = (*src
1652                                .offset(i as isize * src_stride0 + j as isize * src_stride1))
1653                            .transpose();
1654                            *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1655                                scale * value;
1656                        }
1657                    }
1658                }
1659            });
1660        }
1661        return;
1662    }
1663
1664    if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1665        for i in 0..src_rows {
1666            for j in 0..src_cols {
1667                let value =
1668                    (*src.offset(i as isize * src_stride0 + j as isize * src_stride1)).transpose();
1669                *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) = scale * value;
1670            }
1671        }
1672    } else {
1673        for j in 0..src_cols {
1674            for i in 0..src_rows {
1675                let value =
1676                    (*src.offset(i as isize * src_stride0 + j as isize * src_stride1)).transpose();
1677                *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) = scale * value;
1678            }
1679        }
1680    }
1681}
1682
1683/// Copy with transpose and scaling: `dest[j,i] = scale * src[i,j]`.
1684pub fn copy_transpose_scale_into<T>(
1685    dest: &mut StridedViewMut<T>,
1686    src: &StridedView<T>,
1687    scale: T,
1688    exec: Exec<'_>,
1689) -> Result<()>
1690where
1691    T: Copy + ElementOpApply + Mul<Output = T> + Zero + One + PartialEq + Send + Sync + 'static,
1692{
1693    let mut exec = exec;
1694    if src.ndim() != 2 || dest.ndim() != 2 {
1695        return Err(StridedError::RankMismatch(src.ndim(), 2));
1696    }
1697
1698    let src_dims = src.dims();
1699    let expected_dims = [src_dims[1], src_dims[0]];
1700    ensure_same_shape(dest.dims(), &expected_dims)?;
1701    // Reject an element count beyond usize (a huge stride-0 broadcast) up
1702    // front instead of looping over it.
1703    total_len(src_dims)?;
1704
1705    if scale == T::zero() {
1706        unsafe {
1707            if same_contiguous_layout(dest.dims(), &[dest.strides()]).is_some() {
1708                fill_contiguous(dest.as_mut_ptr(), total_len(dest.dims())?, T::zero());
1709            } else {
1710                fill_2d(
1711                    dest.as_mut_ptr(),
1712                    dest.dims()[0],
1713                    dest.dims()[1],
1714                    dest.strides()[0],
1715                    dest.strides()[1],
1716                    T::zero(),
1717                    exec,
1718                );
1719            }
1720        }
1721        return Ok(());
1722    }
1723
1724    let transpose_is_identity = element_transpose_is_identity::<T>();
1725
1726    unsafe {
1727        if transpose_is_identity
1728            && try_copy_transpose_scale_2d_f64_tiled_typed(dest, src, scale, exec.reborrow())
1729        {
1730            return Ok(());
1731        }
1732        if transpose_is_identity
1733            && try_copy_transpose_scale_2d_identity_tiled(dest, src, scale, exec.reborrow())
1734        {
1735            return Ok(());
1736        }
1737    }
1738
1739    if scale == T::one() && transpose_is_identity {
1740        let src_t = src.permute(&[1, 0])?;
1741        return crate::threading::copy_permuted(dest, &src_t, exec);
1742    }
1743
1744    unsafe {
1745        copy_transpose_scale_2d_loop(
1746            dest.as_mut_ptr(),
1747            dest.strides()[0],
1748            dest.strides()[1],
1749            src.ptr(),
1750            src.strides()[0],
1751            src.strides()[1],
1752            src_dims[0],
1753            src_dims[1],
1754            scale,
1755            exec,
1756        );
1757    }
1758    Ok(())
1759}
1760
1761#[cfg(test)]
1762#[path = "ops_view/tests/tiled_tests.rs"]
1763mod tiled_tests;