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