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::{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    #[cfg(feature = "parallel")]
300    {
301        // INVARIANT: the permutation engine selects its float paths by the
302        // source element type and only stores into the destination, so an
303        // uninitialized destination with an initialized source is valid.
304        return crate::threading::copy_permuted_uninit_with_active_policy(dest, src);
305    }
306    #[cfg(not(feature = "parallel"))]
307    crate::threading::copy_permuted_uninit_serial(dest, src)
308}
309
310/// Copy elements from `src` to `dst`, optimized for col-major destination.
311///
312/// Delegates to the centralized permutation-copy dispatch for the actual work.
313pub fn copy_into_col_major<T: Copy + MaybeSendSync + 'static>(
314    dst: &mut StridedViewMut<T>,
315    src: &StridedView<T>,
316) -> Result<()> {
317    crate::threading::copy_into_col_major(dst, src)
318}
319
320/// Element-wise addition: `dest[i] += src[i]`.
321///
322/// Source may have a different element type from destination.
323pub fn add<
324    D: Copy + Add<S, Output = D> + MaybeSendSync,
325    S: Copy + MaybeSendSync,
326    Op: ElementOp<S>,
327>(
328    dest: &mut StridedViewMut<D>,
329    src: &StridedView<S, Op>,
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(dst_dims, &[dst_strides, src_strides])?.is_some() {
340        let len = total_len(dst_dims)?;
341        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
342        let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
343        simd::dispatch_if_large(len, || {
344            for i in 0..len {
345                dst[i] = dst[i] + Op::apply(src[i]);
346            }
347        });
348        return Ok(());
349    }
350
351    let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
352    let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<S>());
353
354    let total = total_len(dst_dims)?;
355    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
356        build_plan_fused_small(dst_dims, &strides_list)
357    } else {
358        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
359    };
360
361    #[cfg(feature = "parallel")]
362    {
363        let total = total_len(&fused_dims)?;
364        let nthreads = crate::execution_policy::rayon_threads();
365        if total > MINTHREADLENGTH && nthreads > 1 {
366            let dst_send = SendPtr(dst_ptr);
367            let src_send = SendPtr(src_ptr as *mut S);
368
369            let costs = compute_costs(&ordered_strides);
370            let initial_offsets: crate::inline::OperandVec<isize> =
371                strides_list.iter().map(|_| 0isize).collect();
372            return mapreduce_threaded(
373                &fused_dims,
374                &plan.block,
375                &ordered_strides,
376                &initial_offsets,
377                &costs,
378                nthreads,
379                0,
380                1,
381                &|dims, blocks, strides_list, offsets| {
382                    for_each_inner_block_with_offsets(
383                        dims,
384                        blocks,
385                        strides_list,
386                        offsets,
387                        |offsets, len, strides| {
388                            unsafe {
389                                inner_loop_add::<D, S, Op>(
390                                    dst_send.as_ptr().offset(offsets[0]),
391                                    strides[0],
392                                    src_send.as_const().offset(offsets[1]),
393                                    strides[1],
394                                    len,
395                                )
396                            };
397                            Ok(())
398                        },
399                    )
400                },
401            );
402        }
403    }
404
405    let initial_offsets: crate::inline::OperandVec<isize> =
406        ordered_strides.iter().map(|_| 0isize).collect();
407    for_each_inner_block_preordered(
408        &fused_dims,
409        &plan.block,
410        &ordered_strides,
411        &initial_offsets,
412        |offsets, len, strides| {
413            unsafe {
414                inner_loop_add::<D, S, Op>(
415                    dst_ptr.offset(offsets[0]),
416                    strides[0],
417                    src_ptr.offset(offsets[1]),
418                    strides[1],
419                    len,
420                )
421            };
422            Ok(())
423        },
424    )
425}
426
427/// Element-wise multiplication: `dest[i] *= src[i]`.
428///
429/// Source may have a different element type from destination.
430pub fn mul<
431    D: Copy + Mul<S, Output = D> + MaybeSendSync,
432    S: Copy + MaybeSendSync,
433    Op: ElementOp<S>,
434>(
435    dest: &mut StridedViewMut<D>,
436    src: &StridedView<S, Op>,
437) -> Result<()> {
438    ensure_same_shape(dest.dims(), src.dims())?;
439
440    let dst_ptr = dest.as_mut_ptr();
441    let src_ptr = src.ptr();
442    let dst_dims = dest.dims();
443    let dst_strides = dest.strides();
444    let src_strides = src.strides();
445
446    if sequential_contiguous_layout(dst_dims, &[dst_strides, src_strides])?.is_some() {
447        let len = total_len(dst_dims)?;
448        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
449        let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
450        simd::dispatch_if_large(len, || {
451            for i in 0..len {
452                dst[i] = dst[i] * Op::apply(src[i]);
453            }
454        });
455        return Ok(());
456    }
457
458    let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
459    let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<S>());
460
461    let total = total_len(dst_dims)?;
462    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
463        build_plan_fused_small(dst_dims, &strides_list)
464    } else {
465        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
466    };
467
468    #[cfg(feature = "parallel")]
469    {
470        let total = total_len(&fused_dims)?;
471        let nthreads = crate::execution_policy::rayon_threads();
472        if total > MINTHREADLENGTH && nthreads > 1 {
473            let dst_send = SendPtr(dst_ptr);
474            let src_send = SendPtr(src_ptr as *mut S);
475
476            let costs = compute_costs(&ordered_strides);
477            let initial_offsets: crate::inline::OperandVec<isize> =
478                strides_list.iter().map(|_| 0isize).collect();
479            return mapreduce_threaded(
480                &fused_dims,
481                &plan.block,
482                &ordered_strides,
483                &initial_offsets,
484                &costs,
485                nthreads,
486                0,
487                1,
488                &|dims, blocks, strides_list, offsets| {
489                    for_each_inner_block_with_offsets(
490                        dims,
491                        blocks,
492                        strides_list,
493                        offsets,
494                        |offsets, len, strides| {
495                            unsafe {
496                                inner_loop_mul::<D, S, Op>(
497                                    dst_send.as_ptr().offset(offsets[0]),
498                                    strides[0],
499                                    src_send.as_const().offset(offsets[1]),
500                                    strides[1],
501                                    len,
502                                )
503                            };
504                            Ok(())
505                        },
506                    )
507                },
508            );
509        }
510    }
511
512    let initial_offsets: crate::inline::OperandVec<isize> =
513        ordered_strides.iter().map(|_| 0isize).collect();
514    for_each_inner_block_preordered(
515        &fused_dims,
516        &plan.block,
517        &ordered_strides,
518        &initial_offsets,
519        |offsets, len, strides| {
520            unsafe {
521                inner_loop_mul::<D, S, Op>(
522                    dst_ptr.offset(offsets[0]),
523                    strides[0],
524                    src_ptr.offset(offsets[1]),
525                    strides[1],
526                    len,
527                )
528            };
529            Ok(())
530        },
531    )
532}
533
534/// AXPY: `dest[i] = alpha * src[i] + dest[i]`.
535///
536/// Alpha, source, and destination may have different element types.
537pub fn axpy<D, S, A, Op>(
538    dest: &mut StridedViewMut<D>,
539    src: &StridedView<S, Op>,
540    alpha: A,
541) -> Result<()>
542where
543    A: Copy + Mul<S, Output = D> + MaybeSync,
544    D: Copy + Add<D, Output = D> + MaybeSendSync,
545    S: Copy + MaybeSendSync,
546    Op: ElementOp<S>,
547{
548    ensure_same_shape(dest.dims(), src.dims())?;
549
550    let dst_ptr = dest.as_mut_ptr();
551    let src_ptr = src.ptr();
552    let dst_dims = dest.dims();
553    let dst_strides = dest.strides();
554    let src_strides = src.strides();
555
556    if sequential_contiguous_layout(dst_dims, &[dst_strides, src_strides])?.is_some() {
557        let len = total_len(dst_dims)?;
558        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
559        let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
560        simd::dispatch_if_large(len, || {
561            for i in 0..len {
562                dst[i] = alpha * Op::apply(src[i]) + dst[i];
563            }
564        });
565        return Ok(());
566    }
567
568    let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
569    let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<S>());
570
571    let total = total_len(dst_dims)?;
572    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
573        build_plan_fused_small(dst_dims, &strides_list)
574    } else {
575        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
576    };
577
578    #[cfg(feature = "parallel")]
579    {
580        let total = total_len(&fused_dims)?;
581        let nthreads = crate::execution_policy::rayon_threads();
582        if total > MINTHREADLENGTH && nthreads > 1 {
583            let dst_send = SendPtr(dst_ptr);
584            let src_send = SendPtr(src_ptr as *mut S);
585
586            let costs = compute_costs(&ordered_strides);
587            let initial_offsets: crate::inline::OperandVec<isize> =
588                strides_list.iter().map(|_| 0isize).collect();
589            return mapreduce_threaded(
590                &fused_dims,
591                &plan.block,
592                &ordered_strides,
593                &initial_offsets,
594                &costs,
595                nthreads,
596                0,
597                1,
598                &|dims, blocks, strides_list, offsets| {
599                    for_each_inner_block_with_offsets(
600                        dims,
601                        blocks,
602                        strides_list,
603                        offsets,
604                        |offsets, len, strides| {
605                            unsafe {
606                                inner_loop_axpy::<D, S, A, Op>(
607                                    dst_send.as_ptr().offset(offsets[0]),
608                                    strides[0],
609                                    src_send.as_const().offset(offsets[1]),
610                                    strides[1],
611                                    len,
612                                    alpha,
613                                )
614                            };
615                            Ok(())
616                        },
617                    )
618                },
619            );
620        }
621    }
622
623    let initial_offsets: crate::inline::OperandVec<isize> =
624        ordered_strides.iter().map(|_| 0isize).collect();
625    for_each_inner_block_preordered(
626        &fused_dims,
627        &plan.block,
628        &ordered_strides,
629        &initial_offsets,
630        |offsets, len, strides| {
631            unsafe {
632                inner_loop_axpy::<D, S, A, Op>(
633                    dst_ptr.offset(offsets[0]),
634                    strides[0],
635                    src_ptr.offset(offsets[1]),
636                    strides[1],
637                    len,
638                    alpha,
639                )
640            };
641            Ok(())
642        },
643    )
644}
645
646/// Fused multiply-add: `dest[i] += OpA::apply(a[i]) * OpB::apply(b[i])`.
647///
648/// Operands may have different element types. Element operations are applied lazily.
649pub fn fma<D, A, B, OpA, OpB>(
650    dest: &mut StridedViewMut<D>,
651    a: &StridedView<A, OpA>,
652    b: &StridedView<B, OpB>,
653) -> Result<()>
654where
655    A: Copy + Mul<B, Output = D> + MaybeSendSync,
656    B: Copy + MaybeSendSync,
657    D: Copy + Add<D, Output = D> + MaybeSendSync,
658    OpA: ElementOp<A>,
659    OpB: ElementOp<B>,
660{
661    ensure_same_shape(dest.dims(), a.dims())?;
662    ensure_same_shape(dest.dims(), b.dims())?;
663
664    let dst_ptr = dest.as_mut_ptr();
665    let a_ptr = a.ptr();
666    let b_ptr = b.ptr();
667    let dst_dims = dest.dims();
668    let dst_strides = dest.strides();
669    let a_strides = a.strides();
670    let b_strides = b.strides();
671
672    if sequential_contiguous_layout(dst_dims, &[dst_strides, a_strides, b_strides])?.is_some() {
673        let len = total_len(dst_dims)?;
674        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
675        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
676        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
677        simd::dispatch_if_large(len, || {
678            for i in 0..len {
679                dst[i] = dst[i] + OpA::apply(sa[i]) * OpB::apply(sb[i]);
680            }
681        });
682        return Ok(());
683    }
684
685    let strides_list: [&[isize]; 3] = [dst_strides, a_strides, b_strides];
686    let elem_size = std::mem::size_of::<D>()
687        .max(std::mem::size_of::<A>())
688        .max(std::mem::size_of::<B>());
689
690    let total = total_len(dst_dims)?;
691    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
692        build_plan_fused_small(dst_dims, &strides_list)
693    } else {
694        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
695    };
696
697    #[cfg(feature = "parallel")]
698    {
699        let total = total_len(&fused_dims)?;
700        let nthreads = crate::execution_policy::rayon_threads();
701        if total > MINTHREADLENGTH && nthreads > 1 {
702            let dst_send = SendPtr(dst_ptr);
703            let a_send = SendPtr(a_ptr as *mut A);
704            let b_send = SendPtr(b_ptr as *mut B);
705
706            let costs = compute_costs(&ordered_strides);
707            let initial_offsets: crate::inline::OperandVec<isize> =
708                strides_list.iter().map(|_| 0isize).collect();
709            return mapreduce_threaded(
710                &fused_dims,
711                &plan.block,
712                &ordered_strides,
713                &initial_offsets,
714                &costs,
715                nthreads,
716                0,
717                1,
718                &|dims, blocks, strides_list, offsets| {
719                    for_each_inner_block_with_offsets(
720                        dims,
721                        blocks,
722                        strides_list,
723                        offsets,
724                        |offsets, len, strides| {
725                            unsafe {
726                                inner_loop_fma::<D, A, B, OpA, OpB>(
727                                    dst_send.as_ptr().offset(offsets[0]),
728                                    strides[0],
729                                    a_send.as_const().offset(offsets[1]),
730                                    strides[1],
731                                    b_send.as_const().offset(offsets[2]),
732                                    strides[2],
733                                    len,
734                                )
735                            };
736                            Ok(())
737                        },
738                    )
739                },
740            );
741        }
742    }
743
744    let initial_offsets: crate::inline::OperandVec<isize> =
745        ordered_strides.iter().map(|_| 0isize).collect();
746    for_each_inner_block_preordered(
747        &fused_dims,
748        &plan.block,
749        &ordered_strides,
750        &initial_offsets,
751        |offsets, len, strides| {
752            unsafe {
753                inner_loop_fma::<D, A, B, OpA, OpB>(
754                    dst_ptr.offset(offsets[0]),
755                    strides[0],
756                    a_ptr.offset(offsets[1]),
757                    strides[1],
758                    b_ptr.offset(offsets[2]),
759                    strides[2],
760                    len,
761                )
762            };
763            Ok(())
764        },
765    )
766}
767
768#[cfg(feature = "parallel")]
769fn parallel_simd_sum<T: Copy + Zero + Add<Output = T> + simd::MaybeSimdOps + Send + Sync>(
770    src: &[T],
771) -> Option<T> {
772    // Check that T has SIMD support
773    if T::try_simd_sum(&[]).is_none() {
774        return None;
775    }
776    let nthreads = crate::execution_policy::rayon_threads();
777    let result = crate::threading::parallel_map_reduce(
778        0..src.len(),
779        nthreads,
780        &|range| T::try_simd_sum(&src[range]).unwrap(),
781        &|left, right| left + right,
782    );
783    Some(result)
784}
785
786/// Sum all elements: `sum(src)`.
787///
788/// # Evaluation order
789///
790/// The result is independent of the ambient pool width when the call runs under
791/// the sequential policy — [`ExecContext::serial`](crate::ExecContext::serial)
792/// or [`with_execution_policy`](crate::with_execution_policy) with
793/// [`ExecutionPolicy::Sequential`](crate::ExecutionPolicy::Sequential). Under a
794/// parallel policy a contiguous identity operand above the threading threshold
795/// is summed per worker and the partial sums are combined in a tree, so the
796/// result is repeatable for a fixed thread count but can differ between thread
797/// counts.
798///
799/// Neither path is a strict left to right sum: a contiguous run is folded with
800/// independent SIMD or scalar lanes (see [`reduce`] for the lane scheme of the
801/// generic path). The last bits can therefore differ from a naive loop, and no
802/// bitwise equality across machines or target features is guaranteed.
803///
804/// For a deterministic accumulation of `f32` data in `f64`, use
805/// [`zip_reduce`](crate::zip_reduce) with an `f64` map and a `0.0f64` init.
806///
807/// # Allocation
808///
809/// At rank ≤ [`INLINE_RANK`](strided_view::INLINE_RANK) the view metadata and
810/// the traversal plan are stored inline, so a call performs no heap allocation.
811pub fn sum<
812    T: Copy + Zero + Add<Output = T> + MaybeSendSync + simd::MaybeSimdOps,
813    Op: ElementOp<T>,
814>(
815    src: &StridedView<T, Op>,
816) -> Result<T> {
817    // A zero-element sum is the additive identity. Return before any planning
818    // so an empty view with an enormous stride never reaches block sizing.
819    if total_len(src.dims())? == 0 {
820        return Ok(T::zero());
821    }
822    // SIMD fast path: contiguous Identity view with SIMD support
823    if Op::IS_IDENTITY {
824        if same_contiguous_layout(src.dims(), &[src.strides()]).is_some() {
825            let len = total_len(src.dims())?;
826            let src_slice = unsafe { std::slice::from_raw_parts(src.ptr(), len) };
827
828            #[cfg(feature = "parallel")]
829            if len > MINTHREADLENGTH {
830                if let Some(result) = parallel_simd_sum(src_slice) {
831                    return Ok(result);
832                }
833            }
834
835            if let Some(result) = T::try_simd_sum(src_slice) {
836                return Ok(result);
837            }
838        }
839    }
840    reduce(src, |x| x, |a, b| a + b, T::zero())
841}
842
843/// Dot product: `sum(OpA::apply(a[i]) * OpB::apply(b[i]))`.
844///
845/// Operands may have different element types. Result type `R` must be `A * B`.
846/// SIMD fast path fires only when `A == B == R` (same type) and both Identity ops.
847///
848/// A widening dot — `f32` operands with an `f64` accumulator — is not
849/// expressible here because `R` is pinned to `A * B`; use
850/// [`zip_reduce`](crate::zip_reduce) with a
851/// `|x, y| x as f64 * y as f64` map instead.
852///
853/// # Evaluation order
854///
855/// This entry takes no parallel branch, so its result does not depend on the
856/// ambient pool width at all: it is computed on the calling thread whatever the
857/// execution policy is. The homogeneous contiguous fast path folds with the SIMD
858/// lane scheme of the element type; every other layout folds a single
859/// accumulator along the traversal order of the fused plan, which is left to
860/// right over each unit-stride run. Neither is guaranteed to be bitwise equal
861/// across machines or target features.
862///
863/// # Allocation
864///
865/// At rank ≤ [`INLINE_RANK`](strided_view::INLINE_RANK) the view metadata and
866/// the traversal plan are stored inline, so a call performs no heap allocation.
867pub fn dot<A, B, R, OpA, OpB>(a: &StridedView<A, OpA>, b: &StridedView<B, OpB>) -> Result<R>
868where
869    A: Copy + Mul<B, Output = R> + MaybeSendSync + 'static,
870    B: Copy + MaybeSendSync + 'static,
871    R: Copy + Zero + Add<Output = R> + MaybeSendSync + simd::MaybeSimdOps + 'static,
872    OpA: ElementOp<A>,
873    OpB: ElementOp<B>,
874{
875    ensure_same_shape(a.dims(), b.dims())?;
876
877    let a_ptr = a.ptr();
878    let b_ptr = b.ptr();
879    let a_strides = a.strides();
880    let b_strides = b.strides();
881    let a_dims = a.dims();
882    let len = total_len(a_dims)?;
883
884    // A zero-element dot product is the additive identity. Return before any
885    // planning so an empty view with an enormous stride never reaches block
886    // sizing.
887    if len == 0 {
888        return Ok(R::zero());
889    }
890
891    if same_contiguous_layout(a_dims, &[a_strides, b_strides]).is_some() {
892        // SIMD fast path: both contiguous, both Identity ops, same type
893        if OpA::IS_IDENTITY
894            && OpB::IS_IDENTITY
895            && std::any::TypeId::of::<A>() == std::any::TypeId::of::<R>()
896            && std::any::TypeId::of::<B>() == std::any::TypeId::of::<R>()
897        {
898            let sa = unsafe { std::slice::from_raw_parts(a_ptr as *const R, len) };
899            let sb = unsafe { std::slice::from_raw_parts(b_ptr as *const R, len) };
900            if let Some(result) = R::try_simd_dot(sa, sb) {
901                return Ok(result);
902            }
903        }
904
905        // Generic contiguous fast path
906        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
907        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
908        let mut acc = R::zero();
909        simd::dispatch_if_large(len, || {
910            for i in 0..len {
911                acc = acc + OpA::apply(sa[i]) * OpB::apply(sb[i]);
912            }
913        });
914        return Ok(acc);
915    }
916
917    let strides_list: [&[isize]; 2] = [a_strides, b_strides];
918    let elem_size = std::mem::size_of::<A>()
919        .max(std::mem::size_of::<B>())
920        .max(std::mem::size_of::<R>());
921
922    let total = total_len(a_dims)?;
923    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
924        build_plan_fused_small(a_dims, &strides_list)
925    } else {
926        build_plan_fused(a_dims, &strides_list, None, elem_size)
927    };
928
929    let mut acc = R::zero();
930    let initial_offsets: crate::inline::OperandVec<isize> =
931        ordered_strides.iter().map(|_| 0isize).collect();
932    for_each_inner_block_preordered(
933        &fused_dims,
934        &plan.block,
935        &ordered_strides,
936        &initial_offsets,
937        |offsets, len, strides| {
938            acc = unsafe {
939                inner_loop_dot::<A, B, R, OpA, OpB>(
940                    a_ptr.offset(offsets[0]),
941                    strides[0],
942                    b_ptr.offset(offsets[1]),
943                    strides[1],
944                    len,
945                    acc,
946                )
947            };
948            Ok(())
949        },
950    )?;
951
952    Ok(acc)
953}
954
955/// Symmetrize a square matrix: `dest = (src + src^T) / 2`.
956pub fn symmetrize_into<T>(dest: &mut StridedViewMut<T>, src: &StridedView<T>) -> Result<()>
957where
958    T: Copy
959        + Add<Output = T>
960        + Mul<Output = T>
961        + num_traits::FromPrimitive
962        + std::ops::Div<Output = T>
963        + MaybeSendSync,
964{
965    if src.ndim() != 2 {
966        return Err(StridedError::RankMismatch(src.ndim(), 2));
967    }
968    let rows = src.dims()[0];
969    let cols = src.dims()[1];
970    if rows != cols {
971        return Err(StridedError::NonSquare { rows, cols });
972    }
973
974    let src_t = src.permute(&[1, 0])?;
975    let half = T::from_f64(0.5).ok_or(StridedError::ScalarConversion)?;
976
977    zip_map2_into(dest, src, &src_t, |a, b| (a + b) * half)
978}
979
980/// Conjugate-symmetrize a square matrix: `dest = (src + conj(src^T)) / 2`.
981pub fn symmetrize_conj_into<T>(dest: &mut StridedViewMut<T>, src: &StridedView<T>) -> Result<()>
982where
983    T: Copy
984        + ElementOpApply
985        + Add<Output = T>
986        + Mul<Output = T>
987        + num_traits::FromPrimitive
988        + std::ops::Div<Output = T>
989        + MaybeSendSync,
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)
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) -> Result<()>
1015where
1016    A: Copy + Mul<S, Output = D> + MaybeSync,
1017    D: Copy + MaybeSendSync,
1018    S: Copy + MaybeSendSync,
1019    Op: ElementOp<S>,
1020{
1021    map_into(dest, src, |x| scale * x)
1022}
1023
1024/// Copy with complex conjugation: `dest[i] = conj(src[i])`.
1025pub fn copy_conj<T: Copy + ElementOpApply + MaybeSendSync>(
1026    dest: &mut StridedViewMut<T>,
1027    src: &StridedView<T>,
1028) -> Result<()> {
1029    let src_conj = src.conj();
1030    copy_into(dest, &src_conj)
1031}
1032
1033#[inline]
1034fn element_transpose_is_identity<T: 'static>() -> bool {
1035    use std::any::TypeId;
1036
1037    macro_rules! matches_type {
1038        ($($ty:ty),* $(,)?) => {{
1039            let id = TypeId::of::<T>();
1040            false $(|| id == TypeId::of::<$ty>())*
1041        }};
1042    }
1043
1044    matches_type!(
1045        f32,
1046        f64,
1047        i8,
1048        i16,
1049        i32,
1050        i64,
1051        i128,
1052        isize,
1053        u8,
1054        u16,
1055        u32,
1056        u64,
1057        u128,
1058        usize,
1059        num_complex::Complex32,
1060        num_complex::Complex64,
1061    )
1062}
1063
1064#[inline]
1065fn element_zero_is_all_bits_zero<T: 'static>() -> bool {
1066    use std::any::TypeId;
1067
1068    macro_rules! matches_type {
1069        ($($ty:ty),* $(,)?) => {{
1070            let id = TypeId::of::<T>();
1071            false $(|| id == TypeId::of::<$ty>())*
1072        }};
1073    }
1074
1075    matches_type!(
1076        f32,
1077        f64,
1078        i8,
1079        i16,
1080        i32,
1081        i64,
1082        i128,
1083        isize,
1084        u8,
1085        u16,
1086        u32,
1087        u64,
1088        u128,
1089        usize,
1090        num_complex::Complex32,
1091        num_complex::Complex64,
1092    )
1093}
1094
1095#[inline]
1096unsafe fn fill_2d<T: Copy + MaybeSendSync>(
1097    dst: *mut T,
1098    dim0: usize,
1099    dim1: usize,
1100    dst_stride0: isize,
1101    dst_stride1: isize,
1102    value: T,
1103) {
1104    #[cfg(feature = "parallel")]
1105    {
1106        let total = dim0.saturating_mul(dim1);
1107        let nthreads = crate::execution_policy::rayon_threads();
1108        if total > MINTHREADLENGTH && nthreads > 1 {
1109            let dst_send = SendPtr(dst);
1110            if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1111                crate::threading::parallel_for_each(0..dim1, nthreads, &|columns| {
1112                    for j in columns {
1113                        let dst = dst_send.as_ptr();
1114                        unsafe {
1115                            let base = j as isize * dst_stride1;
1116                            for i in 0..dim0 {
1117                                *dst.offset(base + i as isize * dst_stride0) = value;
1118                            }
1119                        }
1120                    }
1121                });
1122            } else {
1123                crate::threading::parallel_for_each(0..dim0, nthreads, &|rows| {
1124                    for i in rows {
1125                        let dst = dst_send.as_ptr();
1126                        unsafe {
1127                            let base = i as isize * dst_stride0;
1128                            for j in 0..dim1 {
1129                                *dst.offset(base + j as isize * dst_stride1) = value;
1130                            }
1131                        }
1132                    }
1133                });
1134            }
1135            return;
1136        }
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 + MaybeSendSync + '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]
1233unsafe fn copy_transpose_scale_2d_f64_tiled_raw(
1234    dst: *mut f64,
1235    dst_stride0: isize,
1236    dst_stride1: isize,
1237    src: *const f64,
1238    src_stride0: isize,
1239    src_stride1: isize,
1240    src_rows: usize,
1241    src_cols: usize,
1242    scale: f64,
1243) {
1244    const TILE: usize = 4;
1245    let row_full = src_rows / TILE * TILE;
1246    let col_full = src_cols / TILE * TILE;
1247
1248    #[cfg(feature = "parallel")]
1249    {
1250        let total = src_rows.saturating_mul(src_cols);
1251        let nthreads = crate::execution_policy::rayon_threads();
1252        if total > MINTHREADLENGTH && nthreads > 1 {
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, nthreads, &|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
1299    let mut i = 0;
1300    while i < row_full {
1301        let mut j = 0;
1302        while j < col_full {
1303            transpose_scale_4x4_f64(
1304                dst,
1305                dst_stride0,
1306                dst_stride1,
1307                src,
1308                src_stride0,
1309                src_stride1,
1310                i,
1311                j,
1312                scale,
1313            );
1314            j += TILE;
1315        }
1316        for j in col_full..src_cols {
1317            for ii in i..i + TILE {
1318                *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1319                    scale * *src.offset(ii as isize * src_stride0 + j as isize * src_stride1);
1320            }
1321        }
1322        i += TILE;
1323    }
1324    for i in row_full..src_rows {
1325        for j in 0..src_cols {
1326            *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1327                scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1328        }
1329    }
1330}
1331
1332#[inline]
1333#[cfg(test)]
1334unsafe fn try_copy_transpose_scale_2d_f64_tiled(
1335    dest: &mut StridedViewMut<f64>,
1336    src: &StridedView<f64>,
1337    scale: f64,
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    );
1361    true
1362}
1363
1364#[inline]
1365unsafe fn try_copy_transpose_scale_2d_f64_tiled_typed<T>(
1366    dest: &mut StridedViewMut<T>,
1367    src: &StridedView<T>,
1368    scale: T,
1369) -> bool
1370where
1371    T: Copy + 'static,
1372{
1373    if std::any::TypeId::of::<T>() != std::any::TypeId::of::<f64>() {
1374        return false;
1375    }
1376
1377    let scale = *(&scale as *const T).cast::<f64>();
1378    if src.ndim() != 2 || dest.ndim() != 2 {
1379        return false;
1380    }
1381    let src_dims = src.dims();
1382    if dest.dims() != [src_dims[1], src_dims[0]] {
1383        return false;
1384    }
1385    if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1386        return false;
1387    }
1388
1389    copy_transpose_scale_2d_f64_tiled_raw(
1390        dest.as_mut_ptr().cast::<f64>(),
1391        dest.strides()[0],
1392        dest.strides()[1],
1393        src.ptr().cast::<f64>(),
1394        src.strides()[0],
1395        src.strides()[1],
1396        src_dims[0],
1397        src_dims[1],
1398        scale,
1399    );
1400    true
1401}
1402
1403#[inline(always)]
1404unsafe fn transpose_scale_4x4_identity<T>(
1405    dst: *mut T,
1406    dst_stride0: isize,
1407    dst_stride1: isize,
1408    src: *const T,
1409    src_stride0: isize,
1410    src_stride1: isize,
1411    i: usize,
1412    j: usize,
1413    scale: T,
1414) where
1415    T: Copy + Mul<Output = T>,
1416{
1417    let src_base = src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1418
1419    let s00 = *src_base;
1420    let s10 = *src_base.offset(src_stride0);
1421    let s20 = *src_base.offset(2 * src_stride0);
1422    let s30 = *src_base.offset(3 * src_stride0);
1423
1424    let src_col1 = src_base.offset(src_stride1);
1425    let s01 = *src_col1;
1426    let s11 = *src_col1.offset(src_stride0);
1427    let s21 = *src_col1.offset(2 * src_stride0);
1428    let s31 = *src_col1.offset(3 * src_stride0);
1429
1430    let src_col2 = src_base.offset(2 * src_stride1);
1431    let s02 = *src_col2;
1432    let s12 = *src_col2.offset(src_stride0);
1433    let s22 = *src_col2.offset(2 * src_stride0);
1434    let s32 = *src_col2.offset(3 * src_stride0);
1435
1436    let src_col3 = src_base.offset(3 * src_stride1);
1437    let s03 = *src_col3;
1438    let s13 = *src_col3.offset(src_stride0);
1439    let s23 = *src_col3.offset(2 * src_stride0);
1440    let s33 = *src_col3.offset(3 * src_stride0);
1441
1442    let dst_row0 = dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1);
1443    *dst_row0 = scale * s00;
1444    *dst_row0.offset(dst_stride0) = scale * s01;
1445    *dst_row0.offset(2 * dst_stride0) = scale * s02;
1446    *dst_row0.offset(3 * dst_stride0) = scale * s03;
1447
1448    let dst_row1 = dst_row0.offset(dst_stride1);
1449    *dst_row1 = scale * s10;
1450    *dst_row1.offset(dst_stride0) = scale * s11;
1451    *dst_row1.offset(2 * dst_stride0) = scale * s12;
1452    *dst_row1.offset(3 * dst_stride0) = scale * s13;
1453
1454    let dst_row2 = dst_row0.offset(2 * dst_stride1);
1455    *dst_row2 = scale * s20;
1456    *dst_row2.offset(dst_stride0) = scale * s21;
1457    *dst_row2.offset(2 * dst_stride0) = scale * s22;
1458    *dst_row2.offset(3 * dst_stride0) = scale * s23;
1459
1460    let dst_row3 = dst_row0.offset(3 * dst_stride1);
1461    *dst_row3 = scale * s30;
1462    *dst_row3.offset(dst_stride0) = scale * s31;
1463    *dst_row3.offset(2 * dst_stride0) = scale * s32;
1464    *dst_row3.offset(3 * dst_stride0) = scale * s33;
1465}
1466
1467#[inline]
1468unsafe fn copy_transpose_scale_2d_identity_tiled_raw<T>(
1469    dst: *mut T,
1470    dst_stride0: isize,
1471    dst_stride1: isize,
1472    src: *const T,
1473    src_stride0: isize,
1474    src_stride1: isize,
1475    src_rows: usize,
1476    src_cols: usize,
1477    scale: T,
1478) where
1479    T: Copy + Mul<Output = T> + MaybeSendSync,
1480{
1481    const TILE: usize = 4;
1482    let row_full = src_rows / TILE * TILE;
1483    let col_full = src_cols / TILE * TILE;
1484
1485    #[cfg(feature = "parallel")]
1486    {
1487        let total = src_rows.saturating_mul(src_cols);
1488        let nthreads = crate::execution_policy::rayon_threads();
1489        if total > MINTHREADLENGTH && nthreads > 1 {
1490            let dst_send = SendPtr(dst);
1491            let src_send = SendPtr(src as *mut T);
1492            let row_tiles = row_full / TILE;
1493            crate::threading::parallel_for_each(0..row_tiles, nthreads, &|tiles| {
1494                for tile_i in tiles {
1495                    let i = tile_i * TILE;
1496                    let dst = dst_send.as_ptr();
1497                    let src = src_send.as_const();
1498                    unsafe {
1499                        let mut j = 0;
1500                        while j < col_full {
1501                            transpose_scale_4x4_identity(
1502                                dst,
1503                                dst_stride0,
1504                                dst_stride1,
1505                                src,
1506                                src_stride0,
1507                                src_stride1,
1508                                i,
1509                                j,
1510                                scale,
1511                            );
1512                            j += TILE;
1513                        }
1514                        for j in col_full..src_cols {
1515                            for ii in i..i + TILE {
1516                                *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1517                                    scale
1518                                        * *src.offset(
1519                                            ii as isize * src_stride0 + j as isize * src_stride1,
1520                                        );
1521                            }
1522                        }
1523                    }
1524                }
1525            });
1526            for i in row_full..src_rows {
1527                for j in 0..src_cols {
1528                    *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1529                        scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1530                }
1531            }
1532            return;
1533        }
1534    }
1535
1536    let mut i = 0;
1537    while i < row_full {
1538        let mut j = 0;
1539        while j < col_full {
1540            transpose_scale_4x4_identity(
1541                dst,
1542                dst_stride0,
1543                dst_stride1,
1544                src,
1545                src_stride0,
1546                src_stride1,
1547                i,
1548                j,
1549                scale,
1550            );
1551            j += TILE;
1552        }
1553        for j in col_full..src_cols {
1554            for ii in i..i + TILE {
1555                *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1556                    scale * *src.offset(ii as isize * src_stride0 + j as isize * src_stride1);
1557            }
1558        }
1559        i += TILE;
1560    }
1561    for i in row_full..src_rows {
1562        for j in 0..src_cols {
1563            *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1564                scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1565        }
1566    }
1567}
1568
1569#[inline]
1570unsafe fn try_copy_transpose_scale_2d_identity_tiled<T>(
1571    dest: &mut StridedViewMut<T>,
1572    src: &StridedView<T>,
1573    scale: T,
1574) -> bool
1575where
1576    T: Copy + Mul<Output = T> + MaybeSendSync,
1577{
1578    if src.ndim() != 2 || dest.ndim() != 2 {
1579        return false;
1580    }
1581    let src_dims = src.dims();
1582    if dest.dims() != [src_dims[1], src_dims[0]] {
1583        return false;
1584    }
1585    if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1586        return false;
1587    }
1588
1589    copy_transpose_scale_2d_identity_tiled_raw(
1590        dest.as_mut_ptr(),
1591        dest.strides()[0],
1592        dest.strides()[1],
1593        src.ptr(),
1594        src.strides()[0],
1595        src.strides()[1],
1596        src_dims[0],
1597        src_dims[1],
1598        scale,
1599    );
1600    true
1601}
1602
1603#[inline]
1604unsafe fn copy_transpose_scale_2d_loop<T>(
1605    dst: *mut T,
1606    dst_stride0: isize,
1607    dst_stride1: isize,
1608    src: *const T,
1609    src_stride0: isize,
1610    src_stride1: isize,
1611    src_rows: usize,
1612    src_cols: usize,
1613    scale: T,
1614) where
1615    T: Copy + ElementOpApply + Mul<Output = T> + MaybeSendSync,
1616{
1617    #[cfg(feature = "parallel")]
1618    {
1619        let total = src_rows.saturating_mul(src_cols);
1620        let nthreads = crate::execution_policy::rayon_threads();
1621        if total > MINTHREADLENGTH && nthreads > 1 {
1622            let dst_send = SendPtr(dst);
1623            let src_send = SendPtr(src as *mut T);
1624            if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1625                crate::threading::parallel_for_each(0..src_rows, nthreads, &|rows| {
1626                    for i in rows {
1627                        let dst = dst_send.as_ptr();
1628                        let src = src_send.as_const();
1629                        unsafe {
1630                            for j in 0..src_cols {
1631                                let value = (*src
1632                                    .offset(i as isize * src_stride0 + j as isize * src_stride1))
1633                                .transpose();
1634                                *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1635                                    scale * value;
1636                            }
1637                        }
1638                    }
1639                });
1640            } else {
1641                crate::threading::parallel_for_each(0..src_cols, nthreads, &|columns| {
1642                    for j in columns {
1643                        let dst = dst_send.as_ptr();
1644                        let src = src_send.as_const();
1645                        unsafe {
1646                            for i in 0..src_rows {
1647                                let value = (*src
1648                                    .offset(i as isize * src_stride0 + j as isize * src_stride1))
1649                                .transpose();
1650                                *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1651                                    scale * value;
1652                            }
1653                        }
1654                    }
1655                });
1656            }
1657            return;
1658        }
1659    }
1660
1661    if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1662        for i in 0..src_rows {
1663            for j in 0..src_cols {
1664                let value =
1665                    (*src.offset(i as isize * src_stride0 + j as isize * src_stride1)).transpose();
1666                *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) = scale * value;
1667            }
1668        }
1669    } else {
1670        for j in 0..src_cols {
1671            for i in 0..src_rows {
1672                let value =
1673                    (*src.offset(i as isize * src_stride0 + j as isize * src_stride1)).transpose();
1674                *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) = scale * value;
1675            }
1676        }
1677    }
1678}
1679
1680/// Copy with transpose and scaling: `dest[j,i] = scale * src[i,j]`.
1681pub fn copy_transpose_scale_into<T>(
1682    dest: &mut StridedViewMut<T>,
1683    src: &StridedView<T>,
1684    scale: T,
1685) -> Result<()>
1686where
1687    T: Copy + ElementOpApply + Mul<Output = T> + Zero + One + PartialEq + MaybeSendSync + 'static,
1688{
1689    if src.ndim() != 2 || dest.ndim() != 2 {
1690        return Err(StridedError::RankMismatch(src.ndim(), 2));
1691    }
1692
1693    let src_dims = src.dims();
1694    let expected_dims = [src_dims[1], src_dims[0]];
1695    ensure_same_shape(dest.dims(), &expected_dims)?;
1696    // Reject an element count beyond usize (a huge stride-0 broadcast) up
1697    // front instead of looping over it.
1698    total_len(src_dims)?;
1699
1700    if scale == T::zero() {
1701        unsafe {
1702            if same_contiguous_layout(dest.dims(), &[dest.strides()]).is_some() {
1703                fill_contiguous(dest.as_mut_ptr(), total_len(dest.dims())?, T::zero());
1704            } else {
1705                fill_2d(
1706                    dest.as_mut_ptr(),
1707                    dest.dims()[0],
1708                    dest.dims()[1],
1709                    dest.strides()[0],
1710                    dest.strides()[1],
1711                    T::zero(),
1712                );
1713            }
1714        }
1715        return Ok(());
1716    }
1717
1718    let transpose_is_identity = element_transpose_is_identity::<T>();
1719
1720    unsafe {
1721        if transpose_is_identity && try_copy_transpose_scale_2d_f64_tiled_typed(dest, src, scale) {
1722            return Ok(());
1723        }
1724        if transpose_is_identity && try_copy_transpose_scale_2d_identity_tiled(dest, src, scale) {
1725            return Ok(());
1726        }
1727    }
1728
1729    if scale == T::one() && transpose_is_identity {
1730        let src_t = src.permute(&[1, 0])?;
1731        #[cfg(feature = "parallel")]
1732        {
1733            return crate::threading::copy_permuted_with_active_policy(dest, &src_t);
1734        }
1735        #[cfg(not(feature = "parallel"))]
1736        return crate::threading::copy_permuted_serial(dest, &src_t);
1737    }
1738
1739    unsafe {
1740        copy_transpose_scale_2d_loop(
1741            dest.as_mut_ptr(),
1742            dest.strides()[0],
1743            dest.strides()[1],
1744            src.ptr(),
1745            src.strides()[0],
1746            src.strides()[1],
1747            src_dims[0],
1748            src_dims[1],
1749            scale,
1750        );
1751    }
1752    Ok(())
1753}
1754
1755#[cfg(test)]
1756#[path = "ops_view/tests/tiled_tests.rs"]
1757mod tiled_tests;