Skip to main content

strided_basic/
raw_ops.rs

1//! Allocation-free copy/axpy over borrowed raw strided layouts.
2//!
3//! The map/zip kernels build a traversal plan per call; for small replay
4//! copies that fixed cost dominates. These entry points take
5//! [`RawStridedRef`]/[`RawStridedMut`], fuse the stride pair into a loop nest
6//! (inline storage for fused ranks at most [`RAW_FUSED_RANK_LIMIT`], one
7//! rank-bounded spill above that), and run plain loops - no view plan, no
8//! cache, and no per-element allocation on any call path. Higher *original*
9//! ranks are normalized and fused before that choice is made, so they no
10//! longer leave the fused replay.
11
12use crate::{ElementOpApply, RawStridedMut, RawStridedRef, Result};
13use core::ops::{Add, Mul};
14
15/// Rank up to which a fused loop nest is stored inline.
16///
17/// The stateless [`copy_scale_raw`] family and [`crate::CopyPlan`] fuse any rank:
18/// layouts whose *fused* rank is at most this limit replay from that inline
19/// storage, and higher fused ranks spill to storage bounded by their rank.
20///
21/// Replay is allocation-free for an original rank at most this limit. A higher
22/// *original* rank needs one transient scratch buffer even when the fused layout
23/// still fits inline, because fusion happens before the representation is
24/// chosen; the replay itself is still allocation-free in that case.
25pub const RAW_FUSED_RANK_LIMIT: usize = 8;
26
27/// One fused axis, used while sorting and merging axes and when a layout has to
28/// spill beyond [`RAW_FUSED_RANK_LIMIT`].
29#[derive(Clone, Copy, Debug)]
30pub(crate) struct FusedAxis {
31    pub(crate) dim: usize,
32    pub(crate) dst: isize,
33    pub(crate) src: isize,
34}
35
36const EMPTY_AXIS: FusedAxis = FusedAxis {
37    dim: 1,
38    dst: 0,
39    src: 0,
40};
41
42/// Stack-allocated fused stride pair (axes ordered by destination stride,
43/// adjacent contiguous axes merged). Built once and replayed by both the
44/// per-call raw kernels and the prepared [`crate::CopyPlan`].
45///
46/// This is the representation every call site replays through, so it keeps the
47/// exact shape, size and `Copy`-ness it had before high-rank fusion existed.
48#[derive(Clone, Copy, Debug)]
49pub(crate) struct FusedPairLayout {
50    pub(crate) rank: usize,
51    pub(crate) dims: [usize; RAW_FUSED_RANK_LIMIT],
52    dst_strides: [isize; RAW_FUSED_RANK_LIMIT],
53    src_strides: [isize; RAW_FUSED_RANK_LIMIT],
54}
55
56impl FusedPairLayout {
57    /// An empty rank-zero layout (one element at the base offset).
58    ///
59    /// Used as the neutral value when an inline `TinyVec<[CopyPlan; N]>` is
60    /// default-initialized; it is never executed as a real plan.
61    fn empty() -> Self {
62        Self {
63            rank: 0,
64            dims: [1; RAW_FUSED_RANK_LIMIT],
65            dst_strides: [0; RAW_FUSED_RANK_LIMIT],
66            src_strides: [0; RAW_FUSED_RANK_LIMIT],
67        }
68    }
69
70    fn from_axes(axes: &[FusedAxis]) -> Self {
71        debug_assert!(axes.len() <= RAW_FUSED_RANK_LIMIT);
72        let mut layout = Self {
73            rank: axes.len(),
74            dims: [1; RAW_FUSED_RANK_LIMIT],
75            dst_strides: [0; RAW_FUSED_RANK_LIMIT],
76            src_strides: [0; RAW_FUSED_RANK_LIMIT],
77        };
78        for (slot, axis) in layout.dims[..axes.len()].iter_mut().zip(axes) {
79            *slot = axis.dim;
80        }
81        for (slot, axis) in layout.dst_strides[..axes.len()].iter_mut().zip(axes) {
82            *slot = axis.dst;
83        }
84        for (slot, axis) in layout.src_strides[..axes.len()].iter_mut().zip(axes) {
85            *slot = axis.src;
86        }
87        layout
88    }
89}
90
91/// A fused pair traversal of any rank.
92///
93/// Fused ranks at most [`RAW_FUSED_RANK_LIMIT`] use the inline layout; longer
94/// layouts keep their axes in one `Vec`. That storage is sized by the original
95/// rank (the fusion scratch is reused, and `truncate` keeps its capacity), never
96/// by the tensor's element count.
97#[derive(Clone, Debug)]
98pub(crate) enum PairTraversal {
99    Inline(FusedPairLayout),
100    Spilled(Vec<FusedAxis>),
101}
102
103impl Default for FusedPairLayout {
104    fn default() -> Self {
105        Self::empty()
106    }
107}
108
109impl Default for PairTraversal {
110    fn default() -> Self {
111        Self::empty()
112    }
113}
114
115impl PairTraversal {
116    /// The empty inline layout, used as a neutral default.
117    #[inline]
118    fn empty() -> Self {
119        Self::Inline(FusedPairLayout::empty())
120    }
121
122    /// The inline layout, or `None` when the traversal spilled.
123    #[inline]
124    pub(crate) fn inline(&self) -> Option<&FusedPairLayout> {
125        match self {
126            Self::Inline(layout) => Some(layout),
127            Self::Spilled(_) => None,
128        }
129    }
130
131    /// Number of fused axes.
132    #[cfg(test)]
133    pub(crate) fn rank(&self) -> usize {
134        match self {
135            Self::Inline(layout) => layout.rank,
136            Self::Spilled(axes) => axes.len(),
137        }
138    }
139}
140
141/// Sort `dims`/`dst_strides`/`src_strides` by destination stride, drop size-1
142/// axes, and merge axes that are contiguous in both layouts.
143///
144/// Returns the fused rank written to `scratch[..rank]`. A zero extent collapses
145/// to a single zero-length axis, which every replay treats as an empty range.
146fn fuse_axes(
147    scratch: &mut [FusedAxis],
148    dims: &[usize],
149    dst_strides: &[isize],
150    src_strides: &[isize],
151) -> usize {
152    debug_assert!(scratch.len() >= dims.len());
153    let mut rank = 0usize;
154    for axis in 0..dims.len() {
155        if dims[axis] == 1 {
156            continue;
157        }
158        if dims[axis] == 0 {
159            scratch[0] = FusedAxis {
160                dim: 0,
161                dst: 0,
162                src: 0,
163            };
164            return 1;
165        }
166        let entry = FusedAxis {
167            dim: dims[axis],
168            dst: dst_strides[axis],
169            src: src_strides[axis],
170        };
171        let mut position = rank;
172        while position > 0 && scratch[position - 1].dst > entry.dst {
173            scratch[position] = scratch[position - 1];
174            position -= 1;
175        }
176        scratch[position] = entry;
177        rank += 1;
178    }
179    if rank == 0 {
180        scratch[0] = EMPTY_AXIS;
181        return 1;
182    }
183    let mut fused = 0usize;
184    for axis in 1..rank {
185        // A merged extent that overflows (possible for a stride-0 broadcast
186        // whose element count exceeds usize) leaves the axes unfused.
187        let merged = isize::try_from(scratch[fused].dim)
188            .ok()
189            .filter(|&extent| {
190                scratch[fused].dst.checked_mul(extent) == Some(scratch[axis].dst)
191                    && scratch[fused].src.checked_mul(extent) == Some(scratch[axis].src)
192            })
193            .and_then(|_| scratch[fused].dim.checked_mul(scratch[axis].dim));
194        if let Some(merged) = merged {
195            scratch[fused].dim = merged;
196        } else {
197            fused += 1;
198            scratch[fused] = scratch[axis];
199        }
200    }
201    fused + 1
202}
203
204/// Fuse a raw stride pair into a traversal of any rank.
205///
206/// Inline when the fused rank is at most [`RAW_FUSED_RANK_LIMIT`], otherwise one
207/// rank-bounded spill. An original rank above the limit uses a transient scratch
208/// buffer even when the fused layout still fits inline, so that fusion happens
209/// before the representation is chosen.
210pub(crate) fn fuse_pair_traversal(
211    dims: &[usize],
212    dst_strides: &[isize],
213    src_strides: &[isize],
214) -> PairTraversal {
215    if dims.len() <= RAW_FUSED_RANK_LIMIT {
216        let mut scratch = [EMPTY_AXIS; RAW_FUSED_RANK_LIMIT];
217        let rank = fuse_axes(&mut scratch, dims, dst_strides, src_strides);
218        debug_assert!(rank <= RAW_FUSED_RANK_LIMIT);
219        return PairTraversal::Inline(FusedPairLayout::from_axes(&scratch[..rank]));
220    }
221    let mut scratch = vec![EMPTY_AXIS; dims.len()];
222    let rank = fuse_axes(&mut scratch, dims, dst_strides, src_strides);
223    scratch.truncate(rank);
224    if rank <= RAW_FUSED_RANK_LIMIT {
225        PairTraversal::Inline(FusedPairLayout::from_axes(&scratch))
226    } else {
227        PairTraversal::Spilled(scratch)
228    }
229}
230
231/// Fuse a raw stride pair into the inline representation.
232///
233/// Returns `None` when the fused rank exceeds [`RAW_FUSED_RANK_LIMIT`]. Used by
234/// the window plans, whose non-fused replay is compiled alongside the plan.
235pub(crate) fn fuse_pair_layout(
236    dims: &[usize],
237    dst_strides: &[isize],
238    src_strides: &[isize],
239) -> Option<FusedPairLayout> {
240    match fuse_pair_traversal(dims, dst_strides, src_strides) {
241        PairTraversal::Inline(layout) => Some(layout),
242        PairTraversal::Spilled(_) => None,
243    }
244}
245
246/// Number of logical elements covered by a fused layout.
247///
248/// Every constructor of a [`FusedPairLayout`] starts from dims whose product
249/// was checked by [`crate::kernel::total_len`], and fusion only merges axes
250/// with a checked product, so the product cannot overflow.
251#[inline]
252pub(crate) fn fused_total(layout: &FusedPairLayout) -> usize {
253    layout.dims[..layout.rank].iter().product()
254}
255
256/// Number of logical elements covered by a spilled axis list.
257#[inline]
258fn spilled_total(axes: &[FusedAxis]) -> usize {
259    axes.iter().map(|axis| axis.dim).product()
260}
261
262/// Serial replay of a fused pair layout over borrowed raw views.
263///
264/// The destination layout need not be injective: runs are replayed in order,
265/// so accumulating callers such as [`axpy_raw`] keep their sequential
266/// semantics. Parallel replay is only offered by [`crate::CopyPlan`], whose
267/// compile step proves destination injectivity.
268pub(crate) fn apply_fused_pair<D, S, Apply, Op>(
269    dst: &mut RawStridedMut<'_, D>,
270    src: &RawStridedRef<'_, S>,
271    layout: &FusedPairLayout,
272    apply: Apply,
273    op: Op,
274) where
275    D: Copy,
276    S: Copy,
277    Apply: Fn(&mut D, S),
278    Op: Fn(S) -> S,
279{
280    let total = fused_total(layout);
281    if total == 0 {
282        return;
283    }
284    let src_ptr = src.data().as_ptr();
285    let src_base = src.offset();
286    let dst_base = dst.offset();
287    let dst_ptr = dst.data_mut().as_mut_ptr();
288    // SAFETY: `RawStridedRef`/`RawStridedMut` guarantee (by `new`, or by the
289    // caller of `new_unchecked`) that every offset reachable from their offset
290    // through their dims/strides lies inside their data. Callers pass a layout
291    // fused from exactly those dims/strides, so every logical index in
292    // `0..total` maps to in-bounds source and destination slots. The shared
293    // and exclusive borrows cannot overlap, and this is the only writer.
294    unsafe {
295        apply_fused_range(
296            dst_ptr, dst_base, src_ptr, src_base, layout, 0, total, &apply, &op,
297        );
298    }
299}
300
301/// Serial replay of a spilled fused axis list over borrowed raw views.
302///
303/// Kept out of line and separate from [`apply_fused_pair`] so that the inline
304/// entry points keep the code shape they had before high-rank fusion existed:
305/// fused ranks above [`RAW_FUSED_RANK_LIMIT`] are the rare case, and inlining
306/// this second carry loop into them measurably slowed the inline path (extra
307/// register pressure in a loop that runs once per short run).
308#[inline(never)]
309pub(crate) fn apply_spilled_pair<D, S, Apply, Op>(
310    dst: &mut RawStridedMut<'_, D>,
311    src: &RawStridedRef<'_, S>,
312    axes: &[FusedAxis],
313    apply: Apply,
314    op: Op,
315) where
316    D: Copy,
317    S: Copy,
318    Apply: Fn(&mut D, S),
319    Op: Fn(S) -> S,
320{
321    let total = spilled_total(axes);
322    if total == 0 {
323        return;
324    }
325    let src_ptr = src.data().as_ptr();
326    let src_base = src.offset();
327    let dst_base = dst.offset();
328    let dst_ptr = dst.data_mut().as_mut_ptr();
329    // SAFETY: same contract as `apply_fused_pair`, for the axis list fused from
330    // the same dims/strides.
331    unsafe {
332        apply_spilled_range(
333            dst_ptr, dst_base, src_ptr, src_base, axes, 0, total, &apply, &op,
334        );
335    }
336}
337
338/// Replay logical indices `start..start + len` of a fused pair layout, in
339/// column-major order of the fused axes (axis 0 fastest).
340///
341/// The start coordinate is decoded once; runs along axis 0 then advance the
342/// outer coordinates incrementally, so no per-element range checks or
343/// coordinate rebuilds remain in the hot loop.
344///
345/// # Safety
346///
347/// - `start + len` must not exceed [`fused_total`] of `layout`.
348/// - For every logical index `i` in the range, `dst_base` plus the fused
349///   destination offset of `i` must be an in-bounds, writable slot of the
350///   allocation behind `dst_ptr`, and likewise `src_base` plus the source
351///   offset must be an in-bounds, readable slot behind `src_ptr`.
352/// - No other thread may access the destination slots of this range while it
353///   runs, and the source slots must not be written concurrently. Destination
354///   and source slots must not overlap.
355#[allow(clippy::too_many_arguments)]
356pub(crate) unsafe fn apply_fused_range<D, S, Apply, Op>(
357    dst_ptr: *mut D,
358    dst_base: isize,
359    src_ptr: *const S,
360    src_base: isize,
361    layout: &FusedPairLayout,
362    start: usize,
363    len: usize,
364    apply: &Apply,
365    op: &Op,
366) where
367    D: Copy,
368    S: Copy,
369    Apply: Fn(&mut D, S),
370    Op: Fn(S) -> S,
371{
372    if len == 0 {
373        return;
374    }
375    let rank = layout.rank;
376    let inner_len = layout.dims[0];
377    let inner_dst = layout.dst_strides[0];
378    let inner_src = layout.src_strides[0];
379
380    // Decode the start coordinate once per range. `dst_offset`/`src_offset`
381    // track the offset of the current run start (outer coordinates plus the
382    // inner coordinate `inner_start`); every value is a reachable offset.
383    let mut index = [0usize; RAW_FUSED_RANK_LIMIT];
384    let mut rest = start;
385    let mut dst_outer = dst_base;
386    let mut src_outer = src_base;
387    for axis in 0..rank {
388        let dim = layout.dims[axis];
389        index[axis] = rest % dim;
390        rest /= dim;
391        if axis > 0 {
392            dst_outer += index[axis] as isize * layout.dst_strides[axis];
393            src_outer += index[axis] as isize * layout.src_strides[axis];
394        }
395    }
396    let mut inner_start = index[0];
397    let mut remaining = len;
398    loop {
399        let run = (inner_len - inner_start).min(remaining);
400        let dst_run = dst_outer + inner_start as isize * inner_dst;
401        let src_run = src_outer + inner_start as isize * inner_src;
402        // SAFETY: the run covers logical indices inside the caller's range, so
403        // by this function's contract every slot touched below is in-bounds,
404        // exclusively owned by this call (destination) and unaliased.
405        unsafe {
406            apply_run(
407                dst_ptr, dst_run, inner_dst, src_ptr, src_run, inner_src, run, apply, op,
408            )
409        };
410        remaining -= run;
411        if remaining == 0 {
412            return;
413        }
414        inner_start = 0;
415        // Advance an axis only when another position along it follows, and
416        // rewind it from its last position, so every intermediate base is a
417        // reachable offset of a validated layout (issue #243 follow-up): a
418        // layout ending within one stride of `isize::MAX` must not overflow.
419        // `remaining > 0` guarantees a following run exists, so the loop
420        // always finds an axis to advance before running out of rank.
421        let mut axis = 1;
422        while axis < rank {
423            if index[axis] + 1 < layout.dims[axis] {
424                index[axis] += 1;
425                dst_outer += layout.dst_strides[axis];
426                src_outer += layout.src_strides[axis];
427                break;
428            }
429            let last = (layout.dims[axis] - 1) as isize;
430            dst_outer -= last * layout.dst_strides[axis];
431            src_outer -= last * layout.src_strides[axis];
432            index[axis] = 0;
433            axis += 1;
434        }
435    }
436}
437
438/// Replay logical indices `start..start + len` of a spilled fused axis list, in
439/// column-major order of the axes (axis 0 fastest).
440///
441/// The carry loop is the one [`apply_fused_range`] uses, over interleaved axes
442/// with its carry state in one rank-bounded `Vec`: no allocation scales with
443/// the tensor size and no per-element coordinate decode happens.
444///
445/// # Safety
446///
447/// The [`apply_fused_range`] contract holds for `axes`.
448///
449/// Kept out of line: inlining this second carry loop into the inline replay
450/// measurably slows that one down.
451#[allow(clippy::too_many_arguments)]
452#[inline(never)]
453pub(crate) unsafe fn apply_spilled_range<D, S, Apply, Op>(
454    dst_ptr: *mut D,
455    dst_base: isize,
456    src_ptr: *const S,
457    src_base: isize,
458    axes: &[FusedAxis],
459    start: usize,
460    len: usize,
461    apply: &Apply,
462    op: &Op,
463) where
464    D: Copy,
465    S: Copy,
466    Apply: Fn(&mut D, S),
467    Op: Fn(S) -> S,
468{
469    if len == 0 {
470        return;
471    }
472    let rank = axes.len();
473    let inner_len = axes[0].dim;
474    let inner_dst = axes[0].dst;
475    let inner_src = axes[0].src;
476
477    let mut index = vec![0usize; rank];
478    let mut rest = start;
479    let mut dst_outer = dst_base;
480    let mut src_outer = src_base;
481    for axis in 0..rank {
482        let dim = axes[axis].dim;
483        index[axis] = rest % dim;
484        rest /= dim;
485        if axis > 0 {
486            dst_outer += index[axis] as isize * axes[axis].dst;
487            src_outer += index[axis] as isize * axes[axis].src;
488        }
489    }
490    let mut inner_start = index[0];
491    let mut remaining = len;
492    loop {
493        let run = (inner_len - inner_start).min(remaining);
494        let dst_run = dst_outer + inner_start as isize * inner_dst;
495        let src_run = src_outer + inner_start as isize * inner_src;
496        // SAFETY: the run covers logical indices inside the caller's range, so
497        // by the caller's contract every slot touched below is in-bounds,
498        // exclusively owned by this call (destination) and unaliased.
499        unsafe {
500            apply_run(
501                dst_ptr, dst_run, inner_dst, src_ptr, src_run, inner_src, run, apply, op,
502            )
503        };
504        remaining -= run;
505        if remaining == 0 {
506            return;
507        }
508        inner_start = 0;
509        // Same incremental carry as the inline replay: offsets only ever hold
510        // reachable positions, and `remaining > 0` guarantees an axis advances.
511        let mut axis = 1;
512        while axis < rank {
513            if index[axis] + 1 < axes[axis].dim {
514                index[axis] += 1;
515                dst_outer += axes[axis].dst;
516                src_outer += axes[axis].src;
517                break;
518            }
519            let last = (axes[axis].dim - 1) as isize;
520            dst_outer -= last * axes[axis].dst;
521            src_outer -= last * axes[axis].src;
522            index[axis] = 0;
523            axis += 1;
524        }
525    }
526}
527
528/// One run along the fused inner axis.
529///
530/// # Safety
531///
532/// Every `dst_offset + k * dst_stride` and `src_offset + k * src_stride` for
533/// `k < len` must be an in-bounds slot as described in [`apply_fused_range`],
534/// with exclusive access to the destination slots.
535#[allow(clippy::too_many_arguments)]
536#[inline(always)]
537unsafe fn apply_run<D, S, Apply, Op>(
538    dst_ptr: *mut D,
539    dst_offset: isize,
540    dst_stride: isize,
541    src_ptr: *const S,
542    src_offset: isize,
543    src_stride: isize,
544    len: usize,
545    apply: &Apply,
546    op: &Op,
547) where
548    D: Copy,
549    S: Copy,
550    Apply: Fn(&mut D, S),
551    Op: Fn(S) -> S,
552{
553    if dst_stride == 1 {
554        // SAFETY: a unit destination stride makes the run `len` consecutive
555        // in-bounds slots starting at `dst_offset`, exclusively owned here.
556        let dst_run = unsafe { core::slice::from_raw_parts_mut(dst_ptr.offset(dst_offset), len) };
557        match src_stride {
558            1 => {
559                // SAFETY: `len` consecutive in-bounds source slots.
560                let src_run =
561                    unsafe { core::slice::from_raw_parts(src_ptr.offset(src_offset), len) };
562                for (dst, &value) in dst_run.iter_mut().zip(src_run) {
563                    apply(dst, op(value));
564                }
565            }
566            -1 => {
567                // SAFETY: the run reads `src_offset - (len - 1) ..= src_offset`,
568                // all in-bounds; the lowest one is the run's last source slot.
569                let src_run = unsafe {
570                    core::slice::from_raw_parts(
571                        src_ptr.offset(src_offset - (len as isize - 1)),
572                        len,
573                    )
574                };
575                for (dst, &value) in dst_run.iter_mut().zip(src_run.iter().rev()) {
576                    apply(dst, op(value));
577                }
578            }
579            _ => {
580                let src_start = unsafe { src_ptr.offset(src_offset) };
581                for (position, dst) in dst_run.iter_mut().enumerate() {
582                    // SAFETY: `position < len`, so this is a run source slot.
583                    let value = unsafe { *src_start.offset(position as isize * src_stride) };
584                    apply(dst, op(value));
585                }
586            }
587        }
588        return;
589    }
590    // SAFETY: the run start is an in-bounds slot of each allocation.
591    let dst_start = unsafe { dst_ptr.offset(dst_offset) };
592    let src_start = unsafe { src_ptr.offset(src_offset) };
593    for position in 0..len as isize {
594        // SAFETY: `position < len`, so both are slots of this run.
595        unsafe {
596            let value = *src_start.offset(position * src_stride);
597            apply(&mut *dst_start.offset(position * dst_stride), op(value));
598        }
599    }
600}
601
602fn ensure_same_dims(dst: &[usize], src: &[usize]) -> Result<()> {
603    if dst != src {
604        return Err(crate::StridedError::ShapeMismatch(
605            dst.to_vec(),
606            src.to_vec(),
607        ));
608    }
609    // Reject an element count beyond usize (a huge stride-0 broadcast)
610    // instead of replaying it.
611    crate::kernel::total_len(dst)?;
612    Ok(())
613}
614
615/// `dest = scale * src` over borrowed raw strided layouts.
616pub fn copy_scale_raw<T>(
617    dest: &mut RawStridedMut<'_, T>,
618    src: &RawStridedRef<'_, T>,
619    scale: T,
620) -> Result<()>
621where
622    T: Copy + Mul<T, Output = T> + ElementOpApply + Send + Sync,
623{
624    ensure_same_dims(dest.dims(), src.dims())?;
625    match fuse_pair_traversal(dest.dims(), dest.strides(), src.strides()) {
626        PairTraversal::Inline(layout) => apply_fused_pair(
627            dest,
628            src,
629            &layout,
630            |dst, value| *dst = value,
631            |value: T| scale * value,
632        ),
633        PairTraversal::Spilled(axes) => apply_spilled_pair(
634            dest,
635            src,
636            &axes,
637            |dst, value| *dst = value,
638            |value: T| scale * value,
639        ),
640    }
641    Ok(())
642}
643
644/// `dest = scale * conj(src)` over borrowed raw strided layouts.
645pub fn copy_scale_conj_raw<T>(
646    dest: &mut RawStridedMut<'_, T>,
647    src: &RawStridedRef<'_, T>,
648    scale: T,
649) -> Result<()>
650where
651    T: Copy + Mul<T, Output = T> + ElementOpApply + Send + Sync,
652{
653    ensure_same_dims(dest.dims(), src.dims())?;
654    match fuse_pair_traversal(dest.dims(), dest.strides(), src.strides()) {
655        PairTraversal::Inline(layout) => apply_fused_pair(
656            dest,
657            src,
658            &layout,
659            |dst, value| *dst = value,
660            |value: T| scale * value.conj(),
661        ),
662        PairTraversal::Spilled(axes) => apply_spilled_pair(
663            dest,
664            src,
665            &axes,
666            |dst, value| *dst = value,
667            |value: T| scale * value.conj(),
668        ),
669    }
670    Ok(())
671}
672
673/// `dest = alpha * src + dest` over borrowed raw strided layouts.
674pub fn axpy_raw<T>(
675    dest: &mut RawStridedMut<'_, T>,
676    src: &RawStridedRef<'_, T>,
677    alpha: T,
678) -> Result<()>
679where
680    T: Copy + Add<T, Output = T> + Mul<T, Output = T> + ElementOpApply + Send + Sync,
681{
682    ensure_same_dims(dest.dims(), src.dims())?;
683    match fuse_pair_traversal(dest.dims(), dest.strides(), src.strides()) {
684        PairTraversal::Inline(layout) => apply_fused_pair(
685            dest,
686            src,
687            &layout,
688            |dst, value| *dst = *dst + value,
689            |value: T| alpha * value,
690        ),
691        PairTraversal::Spilled(axes) => apply_spilled_pair(
692            dest,
693            src,
694            &axes,
695            |dst, value| *dst = *dst + value,
696            |value: T| alpha * value,
697        ),
698    }
699    Ok(())
700}
701
702/// `dest = alpha * conj(src) + dest` over borrowed raw strided layouts.
703pub fn axpy_conj_raw<T>(
704    dest: &mut RawStridedMut<'_, T>,
705    src: &RawStridedRef<'_, T>,
706    alpha: T,
707) -> Result<()>
708where
709    T: Copy + Add<T, Output = T> + Mul<T, Output = T> + ElementOpApply + Send + Sync,
710{
711    ensure_same_dims(dest.dims(), src.dims())?;
712    match fuse_pair_traversal(dest.dims(), dest.strides(), src.strides()) {
713        PairTraversal::Inline(layout) => apply_fused_pair(
714            dest,
715            src,
716            &layout,
717            |dst, value| *dst = *dst + value,
718            |value: T| alpha * value.conj(),
719        ),
720        PairTraversal::Spilled(axes) => apply_spilled_pair(
721            dest,
722            src,
723            &axes,
724            |dst, value| *dst = *dst + value,
725            |value: T| alpha * value.conj(),
726        ),
727    }
728    Ok(())
729}
730
731#[cfg(test)]
732#[path = "raw_ops/tests/tests.rs"]
733mod tests;