1use crate::{ElementOpApply, RawStridedMut, RawStridedRef, Result};
13use core::ops::{Add, Mul};
14
15pub const RAW_FUSED_RANK_LIMIT: usize = 8;
26
27#[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#[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 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#[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 #[inline]
118 fn empty() -> Self {
119 Self::Inline(FusedPairLayout::empty())
120 }
121
122 #[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 #[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
141fn 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 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
204pub(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
231pub(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#[inline]
252pub(crate) fn fused_total(layout: &FusedPairLayout) -> usize {
253 layout.dims[..layout.rank].iter().product()
254}
255
256#[inline]
258fn spilled_total(axes: &[FusedAxis]) -> usize {
259 axes.iter().map(|axis| axis.dim).product()
260}
261
262pub(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 unsafe {
295 apply_fused_range(
296 dst_ptr, dst_base, src_ptr, src_base, layout, 0, total, &apply, &op,
297 );
298 }
299}
300
301#[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 unsafe {
332 apply_spilled_range(
333 dst_ptr, dst_base, src_ptr, src_base, axes, 0, total, &apply, &op,
334 );
335 }
336}
337
338#[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 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 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 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#[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 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 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#[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 let dst_run = unsafe { core::slice::from_raw_parts_mut(dst_ptr.offset(dst_offset), len) };
557 match src_stride {
558 1 => {
559 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 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 let value = unsafe { *src_start.offset(position as isize * src_stride) };
584 apply(dst, op(value));
585 }
586 }
587 }
588 return;
589 }
590 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 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 crate::kernel::total_len(dst)?;
612 Ok(())
613}
614
615pub 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
644pub 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
673pub 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
702pub 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;