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