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)`.
794pub fn sum<
795    T: Copy + Zero + Add<Output = T> + MaybeSendSync + simd::MaybeSimdOps,
796    Op: ElementOp<T>,
797>(
798    src: &StridedView<T, Op>,
799) -> Result<T> {
800    // SIMD fast path: contiguous Identity view with SIMD support
801    if Op::IS_IDENTITY {
802        if same_contiguous_layout(src.dims(), &[src.strides()]).is_some() {
803            let len = total_len(src.dims())?;
804            let src_slice = unsafe { std::slice::from_raw_parts(src.ptr(), len) };
805
806            #[cfg(feature = "parallel")]
807            if len > MINTHREADLENGTH {
808                if let Some(result) = parallel_simd_sum(src_slice) {
809                    return Ok(result);
810                }
811            }
812
813            if let Some(result) = T::try_simd_sum(src_slice) {
814                return Ok(result);
815            }
816        }
817    }
818    reduce(src, |x| x, |a, b| a + b, T::zero())
819}
820
821/// Dot product: `sum(OpA::apply(a[i]) * OpB::apply(b[i]))`.
822///
823/// Operands may have different element types. Result type `R` must be `A * B`.
824/// SIMD fast path fires only when `A == B == R` (same type) and both Identity ops.
825pub fn dot<A, B, R, OpA, OpB>(a: &StridedView<A, OpA>, b: &StridedView<B, OpB>) -> Result<R>
826where
827    A: Copy + Mul<B, Output = R> + MaybeSendSync + 'static,
828    B: Copy + MaybeSendSync + 'static,
829    R: Copy + Zero + Add<Output = R> + MaybeSendSync + simd::MaybeSimdOps + 'static,
830    OpA: ElementOp<A>,
831    OpB: ElementOp<B>,
832{
833    ensure_same_shape(a.dims(), b.dims())?;
834
835    let a_ptr = a.ptr();
836    let b_ptr = b.ptr();
837    let a_strides = a.strides();
838    let b_strides = b.strides();
839    let a_dims = a.dims();
840    let len = total_len(a_dims)?;
841
842    if same_contiguous_layout(a_dims, &[a_strides, b_strides]).is_some() {
843        // SIMD fast path: both contiguous, both Identity ops, same type
844        if OpA::IS_IDENTITY
845            && OpB::IS_IDENTITY
846            && std::any::TypeId::of::<A>() == std::any::TypeId::of::<R>()
847            && std::any::TypeId::of::<B>() == std::any::TypeId::of::<R>()
848        {
849            let sa = unsafe { std::slice::from_raw_parts(a_ptr as *const R, len) };
850            let sb = unsafe { std::slice::from_raw_parts(b_ptr as *const R, len) };
851            if let Some(result) = R::try_simd_dot(sa, sb) {
852                return Ok(result);
853            }
854        }
855
856        // Generic contiguous fast path
857        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
858        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
859        let mut acc = R::zero();
860        simd::dispatch_if_large(len, || {
861            for i in 0..len {
862                acc = acc + OpA::apply(sa[i]) * OpB::apply(sb[i]);
863            }
864        });
865        return Ok(acc);
866    }
867
868    let strides_list: [&[isize]; 2] = [a_strides, b_strides];
869    let elem_size = std::mem::size_of::<A>()
870        .max(std::mem::size_of::<B>())
871        .max(std::mem::size_of::<R>());
872
873    let total = total_len(a_dims)?;
874    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
875        build_plan_fused_small(a_dims, &strides_list)
876    } else {
877        build_plan_fused(a_dims, &strides_list, None, elem_size)
878    };
879
880    let mut acc = R::zero();
881    let initial_offsets: crate::inline::OperandVec<isize> =
882        ordered_strides.iter().map(|_| 0isize).collect();
883    for_each_inner_block_preordered(
884        &fused_dims,
885        &plan.block,
886        &ordered_strides,
887        &initial_offsets,
888        |offsets, len, strides| {
889            acc = unsafe {
890                inner_loop_dot::<A, B, R, OpA, OpB>(
891                    a_ptr.offset(offsets[0]),
892                    strides[0],
893                    b_ptr.offset(offsets[1]),
894                    strides[1],
895                    len,
896                    acc,
897                )
898            };
899            Ok(())
900        },
901    )?;
902
903    Ok(acc)
904}
905
906/// Symmetrize a square matrix: `dest = (src + src^T) / 2`.
907pub fn symmetrize_into<T>(dest: &mut StridedViewMut<T>, src: &StridedView<T>) -> Result<()>
908where
909    T: Copy
910        + Add<Output = T>
911        + Mul<Output = T>
912        + num_traits::FromPrimitive
913        + std::ops::Div<Output = T>
914        + MaybeSendSync,
915{
916    if src.ndim() != 2 {
917        return Err(StridedError::RankMismatch(src.ndim(), 2));
918    }
919    let rows = src.dims()[0];
920    let cols = src.dims()[1];
921    if rows != cols {
922        return Err(StridedError::NonSquare { rows, cols });
923    }
924
925    let src_t = src.permute(&[1, 0])?;
926    let half = T::from_f64(0.5).ok_or(StridedError::ScalarConversion)?;
927
928    zip_map2_into(dest, src, &src_t, |a, b| (a + b) * half)
929}
930
931/// Conjugate-symmetrize a square matrix: `dest = (src + conj(src^T)) / 2`.
932pub fn symmetrize_conj_into<T>(dest: &mut StridedViewMut<T>, src: &StridedView<T>) -> Result<()>
933where
934    T: Copy
935        + ElementOpApply
936        + Add<Output = T>
937        + Mul<Output = T>
938        + num_traits::FromPrimitive
939        + std::ops::Div<Output = T>
940        + MaybeSendSync,
941{
942    if src.ndim() != 2 {
943        return Err(StridedError::RankMismatch(src.ndim(), 2));
944    }
945    let rows = src.dims()[0];
946    let cols = src.dims()[1];
947    if rows != cols {
948        return Err(StridedError::NonSquare { rows, cols });
949    }
950
951    // adjoint = conj + transpose
952    let src_adj = src.adjoint_2d()?;
953    let half = T::from_f64(0.5).ok_or(StridedError::ScalarConversion)?;
954
955    zip_map2_into(dest, src, &src_adj, |a, b| (a + b) * half)
956}
957
958/// Copy with scaling: `dest[i] = scale * src[i]`.
959///
960/// Scale, source, and destination may have different element types.
961pub fn copy_scale<D, S, A, Op>(
962    dest: &mut StridedViewMut<D>,
963    src: &StridedView<S, Op>,
964    scale: A,
965) -> Result<()>
966where
967    A: Copy + Mul<S, Output = D> + MaybeSync,
968    D: Copy + MaybeSendSync,
969    S: Copy + MaybeSendSync,
970    Op: ElementOp<S>,
971{
972    map_into(dest, src, |x| scale * x)
973}
974
975/// Copy with complex conjugation: `dest[i] = conj(src[i])`.
976pub fn copy_conj<T: Copy + ElementOpApply + MaybeSendSync>(
977    dest: &mut StridedViewMut<T>,
978    src: &StridedView<T>,
979) -> Result<()> {
980    let src_conj = src.conj();
981    copy_into(dest, &src_conj)
982}
983
984#[inline]
985fn element_transpose_is_identity<T: 'static>() -> bool {
986    use std::any::TypeId;
987
988    macro_rules! matches_type {
989        ($($ty:ty),* $(,)?) => {{
990            let id = TypeId::of::<T>();
991            false $(|| id == TypeId::of::<$ty>())*
992        }};
993    }
994
995    matches_type!(
996        f32,
997        f64,
998        i8,
999        i16,
1000        i32,
1001        i64,
1002        i128,
1003        isize,
1004        u8,
1005        u16,
1006        u32,
1007        u64,
1008        u128,
1009        usize,
1010        num_complex::Complex32,
1011        num_complex::Complex64,
1012    )
1013}
1014
1015#[inline]
1016fn element_zero_is_all_bits_zero<T: 'static>() -> bool {
1017    use std::any::TypeId;
1018
1019    macro_rules! matches_type {
1020        ($($ty:ty),* $(,)?) => {{
1021            let id = TypeId::of::<T>();
1022            false $(|| id == TypeId::of::<$ty>())*
1023        }};
1024    }
1025
1026    matches_type!(
1027        f32,
1028        f64,
1029        i8,
1030        i16,
1031        i32,
1032        i64,
1033        i128,
1034        isize,
1035        u8,
1036        u16,
1037        u32,
1038        u64,
1039        u128,
1040        usize,
1041        num_complex::Complex32,
1042        num_complex::Complex64,
1043    )
1044}
1045
1046#[inline]
1047unsafe fn fill_2d<T: Copy + MaybeSendSync>(
1048    dst: *mut T,
1049    dim0: usize,
1050    dim1: usize,
1051    dst_stride0: isize,
1052    dst_stride1: isize,
1053    value: T,
1054) {
1055    #[cfg(feature = "parallel")]
1056    {
1057        let total = dim0.saturating_mul(dim1);
1058        let nthreads = crate::execution_policy::rayon_threads();
1059        if total > MINTHREADLENGTH && nthreads > 1 {
1060            let dst_send = SendPtr(dst);
1061            if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1062                crate::threading::parallel_for_each(0..dim1, nthreads, &|columns| {
1063                    for j in columns {
1064                        let dst = dst_send.as_ptr();
1065                        unsafe {
1066                            let base = j as isize * dst_stride1;
1067                            for i in 0..dim0 {
1068                                *dst.offset(base + i as isize * dst_stride0) = value;
1069                            }
1070                        }
1071                    }
1072                });
1073            } else {
1074                crate::threading::parallel_for_each(0..dim0, nthreads, &|rows| {
1075                    for i in rows {
1076                        let dst = dst_send.as_ptr();
1077                        unsafe {
1078                            let base = i as isize * dst_stride0;
1079                            for j in 0..dim1 {
1080                                *dst.offset(base + j as isize * dst_stride1) = value;
1081                            }
1082                        }
1083                    }
1084                });
1085            }
1086            return;
1087        }
1088    }
1089
1090    if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1091        for j in 0..dim1 {
1092            let base = j as isize * dst_stride1;
1093            for i in 0..dim0 {
1094                *dst.offset(base + i as isize * dst_stride0) = value;
1095            }
1096        }
1097    } else {
1098        for i in 0..dim0 {
1099            let base = i as isize * dst_stride0;
1100            for j in 0..dim1 {
1101                *dst.offset(base + j as isize * dst_stride1) = value;
1102            }
1103        }
1104    }
1105}
1106
1107#[inline]
1108unsafe fn fill_contiguous<T>(dst: *mut T, len: usize, value: T)
1109where
1110    T: Copy + Zero + PartialEq + MaybeSendSync + 'static,
1111{
1112    if element_zero_is_all_bits_zero::<T>() && value == T::zero() {
1113        std::ptr::write_bytes(dst, 0, len);
1114        return;
1115    }
1116
1117    let dst = std::slice::from_raw_parts_mut(dst, len);
1118    dst.fill(value);
1119}
1120
1121#[inline(always)]
1122unsafe fn transpose_scale_4x4_f64(
1123    dst: *mut f64,
1124    dst_stride0: isize,
1125    dst_stride1: isize,
1126    src: *const f64,
1127    src_stride0: isize,
1128    src_stride1: isize,
1129    i: usize,
1130    j: usize,
1131    scale: f64,
1132) {
1133    let src_base = src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1134
1135    let s00 = *src_base;
1136    let s10 = *src_base.offset(src_stride0);
1137    let s20 = *src_base.offset(2 * src_stride0);
1138    let s30 = *src_base.offset(3 * src_stride0);
1139
1140    let src_col1 = src_base.offset(src_stride1);
1141    let s01 = *src_col1;
1142    let s11 = *src_col1.offset(src_stride0);
1143    let s21 = *src_col1.offset(2 * src_stride0);
1144    let s31 = *src_col1.offset(3 * src_stride0);
1145
1146    let src_col2 = src_base.offset(2 * src_stride1);
1147    let s02 = *src_col2;
1148    let s12 = *src_col2.offset(src_stride0);
1149    let s22 = *src_col2.offset(2 * src_stride0);
1150    let s32 = *src_col2.offset(3 * src_stride0);
1151
1152    let src_col3 = src_base.offset(3 * src_stride1);
1153    let s03 = *src_col3;
1154    let s13 = *src_col3.offset(src_stride0);
1155    let s23 = *src_col3.offset(2 * src_stride0);
1156    let s33 = *src_col3.offset(3 * src_stride0);
1157
1158    let dst_row0 = dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1);
1159    *dst_row0 = scale * s00;
1160    *dst_row0.offset(dst_stride0) = scale * s01;
1161    *dst_row0.offset(2 * dst_stride0) = scale * s02;
1162    *dst_row0.offset(3 * dst_stride0) = scale * s03;
1163
1164    let dst_row1 = dst_row0.offset(dst_stride1);
1165    *dst_row1 = scale * s10;
1166    *dst_row1.offset(dst_stride0) = scale * s11;
1167    *dst_row1.offset(2 * dst_stride0) = scale * s12;
1168    *dst_row1.offset(3 * dst_stride0) = scale * s13;
1169
1170    let dst_row2 = dst_row0.offset(2 * dst_stride1);
1171    *dst_row2 = scale * s20;
1172    *dst_row2.offset(dst_stride0) = scale * s21;
1173    *dst_row2.offset(2 * dst_stride0) = scale * s22;
1174    *dst_row2.offset(3 * dst_stride0) = scale * s23;
1175
1176    let dst_row3 = dst_row0.offset(3 * dst_stride1);
1177    *dst_row3 = scale * s30;
1178    *dst_row3.offset(dst_stride0) = scale * s31;
1179    *dst_row3.offset(2 * dst_stride0) = scale * s32;
1180    *dst_row3.offset(3 * dst_stride0) = scale * s33;
1181}
1182
1183#[inline]
1184unsafe fn copy_transpose_scale_2d_f64_tiled_raw(
1185    dst: *mut f64,
1186    dst_stride0: isize,
1187    dst_stride1: isize,
1188    src: *const f64,
1189    src_stride0: isize,
1190    src_stride1: isize,
1191    src_rows: usize,
1192    src_cols: usize,
1193    scale: f64,
1194) {
1195    const TILE: usize = 4;
1196    let row_full = src_rows / TILE * TILE;
1197    let col_full = src_cols / TILE * TILE;
1198
1199    #[cfg(feature = "parallel")]
1200    {
1201        let total = src_rows.saturating_mul(src_cols);
1202        let nthreads = crate::execution_policy::rayon_threads();
1203        if total > MINTHREADLENGTH && nthreads > 1 {
1204            let dst_send = SendPtr(dst);
1205            let src_send = SendPtr(src as *mut f64);
1206            let row_tiles = row_full / TILE;
1207            crate::threading::parallel_for_each(0..row_tiles, nthreads, &|tiles| {
1208                for tile_i in tiles {
1209                    let i = tile_i * TILE;
1210                    let dst = dst_send.as_ptr();
1211                    let src = src_send.as_const();
1212                    unsafe {
1213                        let mut j = 0;
1214                        while j < col_full {
1215                            transpose_scale_4x4_f64(
1216                                dst,
1217                                dst_stride0,
1218                                dst_stride1,
1219                                src,
1220                                src_stride0,
1221                                src_stride1,
1222                                i,
1223                                j,
1224                                scale,
1225                            );
1226                            j += TILE;
1227                        }
1228                        for j in col_full..src_cols {
1229                            for ii in i..i + TILE {
1230                                *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1231                                    scale
1232                                        * *src.offset(
1233                                            ii as isize * src_stride0 + j as isize * src_stride1,
1234                                        );
1235                            }
1236                        }
1237                    }
1238                }
1239            });
1240            for i in row_full..src_rows {
1241                for j in 0..src_cols {
1242                    *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1243                        scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1244                }
1245            }
1246            return;
1247        }
1248    }
1249
1250    let mut i = 0;
1251    while i < row_full {
1252        let mut j = 0;
1253        while j < col_full {
1254            transpose_scale_4x4_f64(
1255                dst,
1256                dst_stride0,
1257                dst_stride1,
1258                src,
1259                src_stride0,
1260                src_stride1,
1261                i,
1262                j,
1263                scale,
1264            );
1265            j += TILE;
1266        }
1267        for j in col_full..src_cols {
1268            for ii in i..i + TILE {
1269                *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1270                    scale * *src.offset(ii as isize * src_stride0 + j as isize * src_stride1);
1271            }
1272        }
1273        i += TILE;
1274    }
1275    for i in row_full..src_rows {
1276        for j in 0..src_cols {
1277            *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1278                scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1279        }
1280    }
1281}
1282
1283#[inline]
1284#[cfg(test)]
1285unsafe fn try_copy_transpose_scale_2d_f64_tiled(
1286    dest: &mut StridedViewMut<f64>,
1287    src: &StridedView<f64>,
1288    scale: f64,
1289) -> bool {
1290    if src.ndim() != 2 || dest.ndim() != 2 {
1291        return false;
1292    }
1293    let src_dims = src.dims();
1294    if dest.dims() != [src_dims[1], src_dims[0]] {
1295        return false;
1296    }
1297    if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1298        return false;
1299    }
1300
1301    copy_transpose_scale_2d_f64_tiled_raw(
1302        dest.as_mut_ptr(),
1303        dest.strides()[0],
1304        dest.strides()[1],
1305        src.ptr(),
1306        src.strides()[0],
1307        src.strides()[1],
1308        src_dims[0],
1309        src_dims[1],
1310        scale,
1311    );
1312    true
1313}
1314
1315#[inline]
1316unsafe fn try_copy_transpose_scale_2d_f64_tiled_typed<T>(
1317    dest: &mut StridedViewMut<T>,
1318    src: &StridedView<T>,
1319    scale: T,
1320) -> bool
1321where
1322    T: Copy + 'static,
1323{
1324    if std::any::TypeId::of::<T>() != std::any::TypeId::of::<f64>() {
1325        return false;
1326    }
1327
1328    let scale = *(&scale as *const T).cast::<f64>();
1329    if src.ndim() != 2 || dest.ndim() != 2 {
1330        return false;
1331    }
1332    let src_dims = src.dims();
1333    if dest.dims() != [src_dims[1], src_dims[0]] {
1334        return false;
1335    }
1336    if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1337        return false;
1338    }
1339
1340    copy_transpose_scale_2d_f64_tiled_raw(
1341        dest.as_mut_ptr().cast::<f64>(),
1342        dest.strides()[0],
1343        dest.strides()[1],
1344        src.ptr().cast::<f64>(),
1345        src.strides()[0],
1346        src.strides()[1],
1347        src_dims[0],
1348        src_dims[1],
1349        scale,
1350    );
1351    true
1352}
1353
1354#[inline(always)]
1355unsafe fn transpose_scale_4x4_identity<T>(
1356    dst: *mut T,
1357    dst_stride0: isize,
1358    dst_stride1: isize,
1359    src: *const T,
1360    src_stride0: isize,
1361    src_stride1: isize,
1362    i: usize,
1363    j: usize,
1364    scale: T,
1365) where
1366    T: Copy + Mul<Output = T>,
1367{
1368    let src_base = src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1369
1370    let s00 = *src_base;
1371    let s10 = *src_base.offset(src_stride0);
1372    let s20 = *src_base.offset(2 * src_stride0);
1373    let s30 = *src_base.offset(3 * src_stride0);
1374
1375    let src_col1 = src_base.offset(src_stride1);
1376    let s01 = *src_col1;
1377    let s11 = *src_col1.offset(src_stride0);
1378    let s21 = *src_col1.offset(2 * src_stride0);
1379    let s31 = *src_col1.offset(3 * src_stride0);
1380
1381    let src_col2 = src_base.offset(2 * src_stride1);
1382    let s02 = *src_col2;
1383    let s12 = *src_col2.offset(src_stride0);
1384    let s22 = *src_col2.offset(2 * src_stride0);
1385    let s32 = *src_col2.offset(3 * src_stride0);
1386
1387    let src_col3 = src_base.offset(3 * src_stride1);
1388    let s03 = *src_col3;
1389    let s13 = *src_col3.offset(src_stride0);
1390    let s23 = *src_col3.offset(2 * src_stride0);
1391    let s33 = *src_col3.offset(3 * src_stride0);
1392
1393    let dst_row0 = dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1);
1394    *dst_row0 = scale * s00;
1395    *dst_row0.offset(dst_stride0) = scale * s01;
1396    *dst_row0.offset(2 * dst_stride0) = scale * s02;
1397    *dst_row0.offset(3 * dst_stride0) = scale * s03;
1398
1399    let dst_row1 = dst_row0.offset(dst_stride1);
1400    *dst_row1 = scale * s10;
1401    *dst_row1.offset(dst_stride0) = scale * s11;
1402    *dst_row1.offset(2 * dst_stride0) = scale * s12;
1403    *dst_row1.offset(3 * dst_stride0) = scale * s13;
1404
1405    let dst_row2 = dst_row0.offset(2 * dst_stride1);
1406    *dst_row2 = scale * s20;
1407    *dst_row2.offset(dst_stride0) = scale * s21;
1408    *dst_row2.offset(2 * dst_stride0) = scale * s22;
1409    *dst_row2.offset(3 * dst_stride0) = scale * s23;
1410
1411    let dst_row3 = dst_row0.offset(3 * dst_stride1);
1412    *dst_row3 = scale * s30;
1413    *dst_row3.offset(dst_stride0) = scale * s31;
1414    *dst_row3.offset(2 * dst_stride0) = scale * s32;
1415    *dst_row3.offset(3 * dst_stride0) = scale * s33;
1416}
1417
1418#[inline]
1419unsafe fn copy_transpose_scale_2d_identity_tiled_raw<T>(
1420    dst: *mut T,
1421    dst_stride0: isize,
1422    dst_stride1: isize,
1423    src: *const T,
1424    src_stride0: isize,
1425    src_stride1: isize,
1426    src_rows: usize,
1427    src_cols: usize,
1428    scale: T,
1429) where
1430    T: Copy + Mul<Output = T> + MaybeSendSync,
1431{
1432    const TILE: usize = 4;
1433    let row_full = src_rows / TILE * TILE;
1434    let col_full = src_cols / TILE * TILE;
1435
1436    #[cfg(feature = "parallel")]
1437    {
1438        let total = src_rows.saturating_mul(src_cols);
1439        let nthreads = crate::execution_policy::rayon_threads();
1440        if total > MINTHREADLENGTH && nthreads > 1 {
1441            let dst_send = SendPtr(dst);
1442            let src_send = SendPtr(src as *mut T);
1443            let row_tiles = row_full / TILE;
1444            crate::threading::parallel_for_each(0..row_tiles, nthreads, &|tiles| {
1445                for tile_i in tiles {
1446                    let i = tile_i * TILE;
1447                    let dst = dst_send.as_ptr();
1448                    let src = src_send.as_const();
1449                    unsafe {
1450                        let mut j = 0;
1451                        while j < col_full {
1452                            transpose_scale_4x4_identity(
1453                                dst,
1454                                dst_stride0,
1455                                dst_stride1,
1456                                src,
1457                                src_stride0,
1458                                src_stride1,
1459                                i,
1460                                j,
1461                                scale,
1462                            );
1463                            j += TILE;
1464                        }
1465                        for j in col_full..src_cols {
1466                            for ii in i..i + TILE {
1467                                *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1468                                    scale
1469                                        * *src.offset(
1470                                            ii as isize * src_stride0 + j as isize * src_stride1,
1471                                        );
1472                            }
1473                        }
1474                    }
1475                }
1476            });
1477            for i in row_full..src_rows {
1478                for j in 0..src_cols {
1479                    *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1480                        scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1481                }
1482            }
1483            return;
1484        }
1485    }
1486
1487    let mut i = 0;
1488    while i < row_full {
1489        let mut j = 0;
1490        while j < col_full {
1491            transpose_scale_4x4_identity(
1492                dst,
1493                dst_stride0,
1494                dst_stride1,
1495                src,
1496                src_stride0,
1497                src_stride1,
1498                i,
1499                j,
1500                scale,
1501            );
1502            j += TILE;
1503        }
1504        for j in col_full..src_cols {
1505            for ii in i..i + TILE {
1506                *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1507                    scale * *src.offset(ii as isize * src_stride0 + j as isize * src_stride1);
1508            }
1509        }
1510        i += TILE;
1511    }
1512    for i in row_full..src_rows {
1513        for j in 0..src_cols {
1514            *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1515                scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1516        }
1517    }
1518}
1519
1520#[inline]
1521unsafe fn try_copy_transpose_scale_2d_identity_tiled<T>(
1522    dest: &mut StridedViewMut<T>,
1523    src: &StridedView<T>,
1524    scale: T,
1525) -> bool
1526where
1527    T: Copy + Mul<Output = T> + MaybeSendSync,
1528{
1529    if src.ndim() != 2 || dest.ndim() != 2 {
1530        return false;
1531    }
1532    let src_dims = src.dims();
1533    if dest.dims() != [src_dims[1], src_dims[0]] {
1534        return false;
1535    }
1536    if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1537        return false;
1538    }
1539
1540    copy_transpose_scale_2d_identity_tiled_raw(
1541        dest.as_mut_ptr(),
1542        dest.strides()[0],
1543        dest.strides()[1],
1544        src.ptr(),
1545        src.strides()[0],
1546        src.strides()[1],
1547        src_dims[0],
1548        src_dims[1],
1549        scale,
1550    );
1551    true
1552}
1553
1554#[inline]
1555unsafe fn copy_transpose_scale_2d_loop<T>(
1556    dst: *mut T,
1557    dst_stride0: isize,
1558    dst_stride1: isize,
1559    src: *const T,
1560    src_stride0: isize,
1561    src_stride1: isize,
1562    src_rows: usize,
1563    src_cols: usize,
1564    scale: T,
1565) where
1566    T: Copy + ElementOpApply + Mul<Output = T> + MaybeSendSync,
1567{
1568    #[cfg(feature = "parallel")]
1569    {
1570        let total = src_rows.saturating_mul(src_cols);
1571        let nthreads = crate::execution_policy::rayon_threads();
1572        if total > MINTHREADLENGTH && nthreads > 1 {
1573            let dst_send = SendPtr(dst);
1574            let src_send = SendPtr(src as *mut T);
1575            if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1576                crate::threading::parallel_for_each(0..src_rows, nthreads, &|rows| {
1577                    for i in rows {
1578                        let dst = dst_send.as_ptr();
1579                        let src = src_send.as_const();
1580                        unsafe {
1581                            for j in 0..src_cols {
1582                                let value = (*src
1583                                    .offset(i as isize * src_stride0 + j as isize * src_stride1))
1584                                .transpose();
1585                                *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1586                                    scale * value;
1587                            }
1588                        }
1589                    }
1590                });
1591            } else {
1592                crate::threading::parallel_for_each(0..src_cols, nthreads, &|columns| {
1593                    for j in columns {
1594                        let dst = dst_send.as_ptr();
1595                        let src = src_send.as_const();
1596                        unsafe {
1597                            for i in 0..src_rows {
1598                                let value = (*src
1599                                    .offset(i as isize * src_stride0 + j as isize * src_stride1))
1600                                .transpose();
1601                                *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1602                                    scale * value;
1603                            }
1604                        }
1605                    }
1606                });
1607            }
1608            return;
1609        }
1610    }
1611
1612    if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1613        for i in 0..src_rows {
1614            for j in 0..src_cols {
1615                let value =
1616                    (*src.offset(i as isize * src_stride0 + j as isize * src_stride1)).transpose();
1617                *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) = scale * value;
1618            }
1619        }
1620    } else {
1621        for j in 0..src_cols {
1622            for i in 0..src_rows {
1623                let value =
1624                    (*src.offset(i as isize * src_stride0 + j as isize * src_stride1)).transpose();
1625                *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) = scale * value;
1626            }
1627        }
1628    }
1629}
1630
1631/// Copy with transpose and scaling: `dest[j,i] = scale * src[i,j]`.
1632pub fn copy_transpose_scale_into<T>(
1633    dest: &mut StridedViewMut<T>,
1634    src: &StridedView<T>,
1635    scale: T,
1636) -> Result<()>
1637where
1638    T: Copy + ElementOpApply + Mul<Output = T> + Zero + One + PartialEq + MaybeSendSync + 'static,
1639{
1640    if src.ndim() != 2 || dest.ndim() != 2 {
1641        return Err(StridedError::RankMismatch(src.ndim(), 2));
1642    }
1643
1644    let src_dims = src.dims();
1645    let expected_dims = [src_dims[1], src_dims[0]];
1646    ensure_same_shape(dest.dims(), &expected_dims)?;
1647    // Reject an element count beyond usize (a huge stride-0 broadcast) up
1648    // front instead of looping over it.
1649    total_len(src_dims)?;
1650
1651    if scale == T::zero() {
1652        unsafe {
1653            if same_contiguous_layout(dest.dims(), &[dest.strides()]).is_some() {
1654                fill_contiguous(dest.as_mut_ptr(), total_len(dest.dims())?, T::zero());
1655            } else {
1656                fill_2d(
1657                    dest.as_mut_ptr(),
1658                    dest.dims()[0],
1659                    dest.dims()[1],
1660                    dest.strides()[0],
1661                    dest.strides()[1],
1662                    T::zero(),
1663                );
1664            }
1665        }
1666        return Ok(());
1667    }
1668
1669    let transpose_is_identity = element_transpose_is_identity::<T>();
1670
1671    unsafe {
1672        if transpose_is_identity && try_copy_transpose_scale_2d_f64_tiled_typed(dest, src, scale) {
1673            return Ok(());
1674        }
1675        if transpose_is_identity && try_copy_transpose_scale_2d_identity_tiled(dest, src, scale) {
1676            return Ok(());
1677        }
1678    }
1679
1680    if scale == T::one() && transpose_is_identity {
1681        let src_t = src.permute(&[1, 0])?;
1682        #[cfg(feature = "parallel")]
1683        {
1684            return crate::threading::copy_permuted_with_active_policy(dest, &src_t);
1685        }
1686        #[cfg(not(feature = "parallel"))]
1687        return crate::threading::copy_permuted_serial(dest, &src_t);
1688    }
1689
1690    unsafe {
1691        copy_transpose_scale_2d_loop(
1692            dest.as_mut_ptr(),
1693            dest.strides()[0],
1694            dest.strides()[1],
1695            src.ptr(),
1696            src.strides()[0],
1697            src.strides()[1],
1698            src_dims[0],
1699            src_dims[1],
1700            scale,
1701        );
1702    }
1703    Ok(())
1704}
1705
1706#[cfg(test)]
1707#[path = "ops_view/tests/tiled_tests.rs"]
1708mod tiled_tests;