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