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