1use crate::kernel::{
4 build_plan_fused, ensure_same_shape, for_each_inner_block_preordered, same_contiguous_layout,
5 sequential_contiguous_layout, total_len,
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>(
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 (fused_dims, ordered_strides, plan) =
362 build_plan_fused(dst_dims, &strides_list, Some(0), elem_size);
363
364 #[cfg(feature = "parallel")]
365 {
366 let total = total_len(&fused_dims)?;
367 let nthreads = crate::execution_policy::rayon_threads();
368 if total > MINTHREADLENGTH && nthreads > 1 {
369 let dst_send = SendPtr(dst_ptr);
370 let src_send = SendPtr(src_ptr as *mut S);
371
372 let costs = compute_costs(&ordered_strides);
373 let initial_offsets = vec![0isize; strides_list.len()];
374 return mapreduce_threaded(
375 &fused_dims,
376 &plan.block,
377 &ordered_strides,
378 &initial_offsets,
379 &costs,
380 nthreads,
381 0,
382 1,
383 &|dims, blocks, strides_list, offsets| {
384 for_each_inner_block_with_offsets(
385 dims,
386 blocks,
387 strides_list,
388 offsets,
389 |offsets, len, strides| {
390 unsafe {
391 inner_loop_add::<D, S, Op>(
392 dst_send.as_ptr().offset(offsets[0]),
393 strides[0],
394 src_send.as_const().offset(offsets[1]),
395 strides[1],
396 len,
397 )
398 };
399 Ok(())
400 },
401 )
402 },
403 );
404 }
405 }
406
407 let initial_offsets = vec![0isize; ordered_strides.len()];
408 for_each_inner_block_preordered(
409 &fused_dims,
410 &plan.block,
411 &ordered_strides,
412 &initial_offsets,
413 |offsets, len, strides| {
414 unsafe {
415 inner_loop_add::<D, S, Op>(
416 dst_ptr.offset(offsets[0]),
417 strides[0],
418 src_ptr.offset(offsets[1]),
419 strides[1],
420 len,
421 )
422 };
423 Ok(())
424 },
425 )
426}
427
428pub fn mul<
432 D: Copy + Mul<S, Output = D> + MaybeSendSync,
433 S: Copy + MaybeSendSync,
434 Op: ElementOp<S>,
435>(
436 dest: &mut StridedViewMut<D>,
437 src: &StridedView<S, Op>,
438) -> Result<()> {
439 ensure_same_shape(dest.dims(), src.dims())?;
440
441 let dst_ptr = dest.as_mut_ptr();
442 let src_ptr = src.ptr();
443 let dst_dims = dest.dims();
444 let dst_strides = dest.strides();
445 let src_strides = src.strides();
446
447 if sequential_contiguous_layout(dst_dims, &[dst_strides, src_strides])?.is_some() {
448 let len = total_len(dst_dims)?;
449 let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
450 let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
451 simd::dispatch_if_large(len, || {
452 for i in 0..len {
453 dst[i] = dst[i] * Op::apply(src[i]);
454 }
455 });
456 return Ok(());
457 }
458
459 let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
460 let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<S>());
461
462 let (fused_dims, ordered_strides, plan) =
463 build_plan_fused(dst_dims, &strides_list, Some(0), elem_size);
464
465 #[cfg(feature = "parallel")]
466 {
467 let total = total_len(&fused_dims)?;
468 let nthreads = crate::execution_policy::rayon_threads();
469 if total > MINTHREADLENGTH && nthreads > 1 {
470 let dst_send = SendPtr(dst_ptr);
471 let src_send = SendPtr(src_ptr as *mut S);
472
473 let costs = compute_costs(&ordered_strides);
474 let initial_offsets = vec![0isize; strides_list.len()];
475 return mapreduce_threaded(
476 &fused_dims,
477 &plan.block,
478 &ordered_strides,
479 &initial_offsets,
480 &costs,
481 nthreads,
482 0,
483 1,
484 &|dims, blocks, strides_list, offsets| {
485 for_each_inner_block_with_offsets(
486 dims,
487 blocks,
488 strides_list,
489 offsets,
490 |offsets, len, strides| {
491 unsafe {
492 inner_loop_mul::<D, S, Op>(
493 dst_send.as_ptr().offset(offsets[0]),
494 strides[0],
495 src_send.as_const().offset(offsets[1]),
496 strides[1],
497 len,
498 )
499 };
500 Ok(())
501 },
502 )
503 },
504 );
505 }
506 }
507
508 let initial_offsets = vec![0isize; ordered_strides.len()];
509 for_each_inner_block_preordered(
510 &fused_dims,
511 &plan.block,
512 &ordered_strides,
513 &initial_offsets,
514 |offsets, len, strides| {
515 unsafe {
516 inner_loop_mul::<D, S, Op>(
517 dst_ptr.offset(offsets[0]),
518 strides[0],
519 src_ptr.offset(offsets[1]),
520 strides[1],
521 len,
522 )
523 };
524 Ok(())
525 },
526 )
527}
528
529pub fn axpy<D, S, A, Op>(
533 dest: &mut StridedViewMut<D>,
534 src: &StridedView<S, Op>,
535 alpha: A,
536) -> Result<()>
537where
538 A: Copy + Mul<S, Output = D> + MaybeSync,
539 D: Copy + Add<D, Output = D> + MaybeSendSync,
540 S: Copy + MaybeSendSync,
541 Op: ElementOp<S>,
542{
543 ensure_same_shape(dest.dims(), src.dims())?;
544
545 let dst_ptr = dest.as_mut_ptr();
546 let src_ptr = src.ptr();
547 let dst_dims = dest.dims();
548 let dst_strides = dest.strides();
549 let src_strides = src.strides();
550
551 if sequential_contiguous_layout(dst_dims, &[dst_strides, src_strides])?.is_some() {
552 let len = total_len(dst_dims)?;
553 let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
554 let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
555 simd::dispatch_if_large(len, || {
556 for i in 0..len {
557 dst[i] = alpha * Op::apply(src[i]) + dst[i];
558 }
559 });
560 return Ok(());
561 }
562
563 let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
564 let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<S>());
565
566 let (fused_dims, ordered_strides, plan) =
567 build_plan_fused(dst_dims, &strides_list, Some(0), elem_size);
568
569 #[cfg(feature = "parallel")]
570 {
571 let total = total_len(&fused_dims)?;
572 let nthreads = crate::execution_policy::rayon_threads();
573 if total > MINTHREADLENGTH && nthreads > 1 {
574 let dst_send = SendPtr(dst_ptr);
575 let src_send = SendPtr(src_ptr as *mut S);
576
577 let costs = compute_costs(&ordered_strides);
578 let initial_offsets = vec![0isize; strides_list.len()];
579 return mapreduce_threaded(
580 &fused_dims,
581 &plan.block,
582 &ordered_strides,
583 &initial_offsets,
584 &costs,
585 nthreads,
586 0,
587 1,
588 &|dims, blocks, strides_list, offsets| {
589 for_each_inner_block_with_offsets(
590 dims,
591 blocks,
592 strides_list,
593 offsets,
594 |offsets, len, strides| {
595 unsafe {
596 inner_loop_axpy::<D, S, A, Op>(
597 dst_send.as_ptr().offset(offsets[0]),
598 strides[0],
599 src_send.as_const().offset(offsets[1]),
600 strides[1],
601 len,
602 alpha,
603 )
604 };
605 Ok(())
606 },
607 )
608 },
609 );
610 }
611 }
612
613 let initial_offsets = vec![0isize; ordered_strides.len()];
614 for_each_inner_block_preordered(
615 &fused_dims,
616 &plan.block,
617 &ordered_strides,
618 &initial_offsets,
619 |offsets, len, strides| {
620 unsafe {
621 inner_loop_axpy::<D, S, A, Op>(
622 dst_ptr.offset(offsets[0]),
623 strides[0],
624 src_ptr.offset(offsets[1]),
625 strides[1],
626 len,
627 alpha,
628 )
629 };
630 Ok(())
631 },
632 )
633}
634
635pub fn fma<D, A, B, OpA, OpB>(
639 dest: &mut StridedViewMut<D>,
640 a: &StridedView<A, OpA>,
641 b: &StridedView<B, OpB>,
642) -> Result<()>
643where
644 A: Copy + Mul<B, Output = D> + MaybeSendSync,
645 B: Copy + MaybeSendSync,
646 D: Copy + Add<D, Output = D> + MaybeSendSync,
647 OpA: ElementOp<A>,
648 OpB: ElementOp<B>,
649{
650 ensure_same_shape(dest.dims(), a.dims())?;
651 ensure_same_shape(dest.dims(), b.dims())?;
652
653 let dst_ptr = dest.as_mut_ptr();
654 let a_ptr = a.ptr();
655 let b_ptr = b.ptr();
656 let dst_dims = dest.dims();
657 let dst_strides = dest.strides();
658 let a_strides = a.strides();
659 let b_strides = b.strides();
660
661 if sequential_contiguous_layout(dst_dims, &[dst_strides, a_strides, b_strides])?.is_some() {
662 let len = total_len(dst_dims)?;
663 let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
664 let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
665 let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
666 simd::dispatch_if_large(len, || {
667 for i in 0..len {
668 dst[i] = dst[i] + OpA::apply(sa[i]) * OpB::apply(sb[i]);
669 }
670 });
671 return Ok(());
672 }
673
674 let strides_list: [&[isize]; 3] = [dst_strides, a_strides, b_strides];
675 let elem_size = std::mem::size_of::<D>()
676 .max(std::mem::size_of::<A>())
677 .max(std::mem::size_of::<B>());
678
679 let (fused_dims, ordered_strides, plan) =
680 build_plan_fused(dst_dims, &strides_list, Some(0), elem_size);
681
682 #[cfg(feature = "parallel")]
683 {
684 let total = total_len(&fused_dims)?;
685 let nthreads = crate::execution_policy::rayon_threads();
686 if total > MINTHREADLENGTH && nthreads > 1 {
687 let dst_send = SendPtr(dst_ptr);
688 let a_send = SendPtr(a_ptr as *mut A);
689 let b_send = SendPtr(b_ptr as *mut B);
690
691 let costs = compute_costs(&ordered_strides);
692 let initial_offsets = vec![0isize; strides_list.len()];
693 return mapreduce_threaded(
694 &fused_dims,
695 &plan.block,
696 &ordered_strides,
697 &initial_offsets,
698 &costs,
699 nthreads,
700 0,
701 1,
702 &|dims, blocks, strides_list, offsets| {
703 for_each_inner_block_with_offsets(
704 dims,
705 blocks,
706 strides_list,
707 offsets,
708 |offsets, len, strides| {
709 unsafe {
710 inner_loop_fma::<D, A, B, OpA, OpB>(
711 dst_send.as_ptr().offset(offsets[0]),
712 strides[0],
713 a_send.as_const().offset(offsets[1]),
714 strides[1],
715 b_send.as_const().offset(offsets[2]),
716 strides[2],
717 len,
718 )
719 };
720 Ok(())
721 },
722 )
723 },
724 );
725 }
726 }
727
728 let initial_offsets = vec![0isize; ordered_strides.len()];
729 for_each_inner_block_preordered(
730 &fused_dims,
731 &plan.block,
732 &ordered_strides,
733 &initial_offsets,
734 |offsets, len, strides| {
735 unsafe {
736 inner_loop_fma::<D, A, B, OpA, OpB>(
737 dst_ptr.offset(offsets[0]),
738 strides[0],
739 a_ptr.offset(offsets[1]),
740 strides[1],
741 b_ptr.offset(offsets[2]),
742 strides[2],
743 len,
744 )
745 };
746 Ok(())
747 },
748 )
749}
750
751#[cfg(feature = "parallel")]
752fn parallel_simd_sum<T: Copy + Zero + Add<Output = T> + simd::MaybeSimdOps + Send + Sync>(
753 src: &[T],
754) -> Option<T> {
755 if T::try_simd_sum(&[]).is_none() {
757 return None;
758 }
759 let nthreads = crate::execution_policy::rayon_threads();
760 let result = crate::threading::parallel_map_reduce(
761 0..src.len(),
762 nthreads,
763 &|range| T::try_simd_sum(&src[range]).unwrap(),
764 &|left, right| left + right,
765 );
766 Some(result)
767}
768
769pub fn sum<
771 T: Copy + Zero + Add<Output = T> + MaybeSendSync + simd::MaybeSimdOps,
772 Op: ElementOp<T>,
773>(
774 src: &StridedView<T, Op>,
775) -> Result<T> {
776 if Op::IS_IDENTITY {
778 if same_contiguous_layout(src.dims(), &[src.strides()]).is_some() {
779 let len = total_len(src.dims())?;
780 let src_slice = unsafe { std::slice::from_raw_parts(src.ptr(), len) };
781
782 #[cfg(feature = "parallel")]
783 if len > MINTHREADLENGTH {
784 if let Some(result) = parallel_simd_sum(src_slice) {
785 return Ok(result);
786 }
787 }
788
789 if let Some(result) = T::try_simd_sum(src_slice) {
790 return Ok(result);
791 }
792 }
793 }
794 reduce(src, |x| x, |a, b| a + b, T::zero())
795}
796
797pub fn dot<A, B, R, OpA, OpB>(a: &StridedView<A, OpA>, b: &StridedView<B, OpB>) -> Result<R>
802where
803 A: Copy + Mul<B, Output = R> + MaybeSendSync + 'static,
804 B: Copy + MaybeSendSync + 'static,
805 R: Copy + Zero + Add<Output = R> + MaybeSendSync + simd::MaybeSimdOps + 'static,
806 OpA: ElementOp<A>,
807 OpB: ElementOp<B>,
808{
809 ensure_same_shape(a.dims(), b.dims())?;
810
811 let a_ptr = a.ptr();
812 let b_ptr = b.ptr();
813 let a_strides = a.strides();
814 let b_strides = b.strides();
815 let a_dims = a.dims();
816 let len = total_len(a_dims)?;
817
818 if same_contiguous_layout(a_dims, &[a_strides, b_strides]).is_some() {
819 if OpA::IS_IDENTITY
821 && OpB::IS_IDENTITY
822 && std::any::TypeId::of::<A>() == std::any::TypeId::of::<R>()
823 && std::any::TypeId::of::<B>() == std::any::TypeId::of::<R>()
824 {
825 let sa = unsafe { std::slice::from_raw_parts(a_ptr as *const R, len) };
826 let sb = unsafe { std::slice::from_raw_parts(b_ptr as *const R, len) };
827 if let Some(result) = R::try_simd_dot(sa, sb) {
828 return Ok(result);
829 }
830 }
831
832 let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
834 let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
835 let mut acc = R::zero();
836 simd::dispatch_if_large(len, || {
837 for i in 0..len {
838 acc = acc + OpA::apply(sa[i]) * OpB::apply(sb[i]);
839 }
840 });
841 return Ok(acc);
842 }
843
844 let strides_list: [&[isize]; 2] = [a_strides, b_strides];
845 let elem_size = std::mem::size_of::<A>()
846 .max(std::mem::size_of::<B>())
847 .max(std::mem::size_of::<R>());
848
849 let (fused_dims, ordered_strides, plan) =
850 build_plan_fused(a_dims, &strides_list, None, elem_size);
851
852 let mut acc = R::zero();
853 let initial_offsets = vec![0isize; ordered_strides.len()];
854 for_each_inner_block_preordered(
855 &fused_dims,
856 &plan.block,
857 &ordered_strides,
858 &initial_offsets,
859 |offsets, len, strides| {
860 acc = unsafe {
861 inner_loop_dot::<A, B, R, OpA, OpB>(
862 a_ptr.offset(offsets[0]),
863 strides[0],
864 b_ptr.offset(offsets[1]),
865 strides[1],
866 len,
867 acc,
868 )
869 };
870 Ok(())
871 },
872 )?;
873
874 Ok(acc)
875}
876
877pub fn symmetrize_into<T>(dest: &mut StridedViewMut<T>, src: &StridedView<T>) -> Result<()>
879where
880 T: Copy
881 + Add<Output = T>
882 + Mul<Output = T>
883 + num_traits::FromPrimitive
884 + std::ops::Div<Output = T>
885 + MaybeSendSync,
886{
887 if src.ndim() != 2 {
888 return Err(StridedError::RankMismatch(src.ndim(), 2));
889 }
890 let rows = src.dims()[0];
891 let cols = src.dims()[1];
892 if rows != cols {
893 return Err(StridedError::NonSquare { rows, cols });
894 }
895
896 let src_t = src.permute(&[1, 0])?;
897 let half = T::from_f64(0.5).ok_or(StridedError::ScalarConversion)?;
898
899 zip_map2_into(dest, src, &src_t, |a, b| (a + b) * half)
900}
901
902pub fn symmetrize_conj_into<T>(dest: &mut StridedViewMut<T>, src: &StridedView<T>) -> Result<()>
904where
905 T: Copy
906 + ElementOpApply
907 + Add<Output = T>
908 + Mul<Output = T>
909 + num_traits::FromPrimitive
910 + std::ops::Div<Output = T>
911 + MaybeSendSync,
912{
913 if src.ndim() != 2 {
914 return Err(StridedError::RankMismatch(src.ndim(), 2));
915 }
916 let rows = src.dims()[0];
917 let cols = src.dims()[1];
918 if rows != cols {
919 return Err(StridedError::NonSquare { rows, cols });
920 }
921
922 let src_adj = src.adjoint_2d()?;
924 let half = T::from_f64(0.5).ok_or(StridedError::ScalarConversion)?;
925
926 zip_map2_into(dest, src, &src_adj, |a, b| (a + b) * half)
927}
928
929pub fn copy_scale<D, S, A, Op>(
933 dest: &mut StridedViewMut<D>,
934 src: &StridedView<S, Op>,
935 scale: A,
936) -> Result<()>
937where
938 A: Copy + Mul<S, Output = D> + MaybeSync,
939 D: Copy + MaybeSendSync,
940 S: Copy + MaybeSendSync,
941 Op: ElementOp<S>,
942{
943 map_into(dest, src, |x| scale * x)
944}
945
946pub fn copy_conj<T: Copy + ElementOpApply + MaybeSendSync>(
948 dest: &mut StridedViewMut<T>,
949 src: &StridedView<T>,
950) -> Result<()> {
951 let src_conj = src.conj();
952 copy_into(dest, &src_conj)
953}
954
955#[inline]
956fn element_transpose_is_identity<T: 'static>() -> bool {
957 use std::any::TypeId;
958
959 macro_rules! matches_type {
960 ($($ty:ty),* $(,)?) => {{
961 let id = TypeId::of::<T>();
962 false $(|| id == TypeId::of::<$ty>())*
963 }};
964 }
965
966 matches_type!(
967 f32,
968 f64,
969 i8,
970 i16,
971 i32,
972 i64,
973 i128,
974 isize,
975 u8,
976 u16,
977 u32,
978 u64,
979 u128,
980 usize,
981 num_complex::Complex32,
982 num_complex::Complex64,
983 )
984}
985
986#[inline]
987fn element_zero_is_all_bits_zero<T: 'static>() -> bool {
988 use std::any::TypeId;
989
990 macro_rules! matches_type {
991 ($($ty:ty),* $(,)?) => {{
992 let id = TypeId::of::<T>();
993 false $(|| id == TypeId::of::<$ty>())*
994 }};
995 }
996
997 matches_type!(
998 f32,
999 f64,
1000 i8,
1001 i16,
1002 i32,
1003 i64,
1004 i128,
1005 isize,
1006 u8,
1007 u16,
1008 u32,
1009 u64,
1010 u128,
1011 usize,
1012 num_complex::Complex32,
1013 num_complex::Complex64,
1014 )
1015}
1016
1017#[inline]
1018unsafe fn fill_2d<T: Copy + MaybeSendSync>(
1019 dst: *mut T,
1020 dim0: usize,
1021 dim1: usize,
1022 dst_stride0: isize,
1023 dst_stride1: isize,
1024 value: T,
1025) {
1026 #[cfg(feature = "parallel")]
1027 {
1028 let total = dim0.saturating_mul(dim1);
1029 let nthreads = crate::execution_policy::rayon_threads();
1030 if total > MINTHREADLENGTH && nthreads > 1 {
1031 let dst_send = SendPtr(dst);
1032 if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1033 crate::threading::parallel_for_each(0..dim1, nthreads, &|columns| {
1034 for j in columns {
1035 let dst = dst_send.as_ptr();
1036 unsafe {
1037 let base = j as isize * dst_stride1;
1038 for i in 0..dim0 {
1039 *dst.offset(base + i as isize * dst_stride0) = value;
1040 }
1041 }
1042 }
1043 });
1044 } else {
1045 crate::threading::parallel_for_each(0..dim0, nthreads, &|rows| {
1046 for i in rows {
1047 let dst = dst_send.as_ptr();
1048 unsafe {
1049 let base = i as isize * dst_stride0;
1050 for j in 0..dim1 {
1051 *dst.offset(base + j as isize * dst_stride1) = value;
1052 }
1053 }
1054 }
1055 });
1056 }
1057 return;
1058 }
1059 }
1060
1061 if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1062 for j in 0..dim1 {
1063 let base = j as isize * dst_stride1;
1064 for i in 0..dim0 {
1065 *dst.offset(base + i as isize * dst_stride0) = value;
1066 }
1067 }
1068 } else {
1069 for i in 0..dim0 {
1070 let base = i as isize * dst_stride0;
1071 for j in 0..dim1 {
1072 *dst.offset(base + j as isize * dst_stride1) = value;
1073 }
1074 }
1075 }
1076}
1077
1078#[inline]
1079unsafe fn fill_contiguous<T>(dst: *mut T, len: usize, value: T)
1080where
1081 T: Copy + Zero + PartialEq + MaybeSendSync + 'static,
1082{
1083 if element_zero_is_all_bits_zero::<T>() && value == T::zero() {
1084 std::ptr::write_bytes(dst, 0, len);
1085 return;
1086 }
1087
1088 let dst = std::slice::from_raw_parts_mut(dst, len);
1089 dst.fill(value);
1090}
1091
1092#[inline(always)]
1093unsafe fn transpose_scale_4x4_f64(
1094 dst: *mut f64,
1095 dst_stride0: isize,
1096 dst_stride1: isize,
1097 src: *const f64,
1098 src_stride0: isize,
1099 src_stride1: isize,
1100 i: usize,
1101 j: usize,
1102 scale: f64,
1103) {
1104 let src_base = src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1105
1106 let s00 = *src_base;
1107 let s10 = *src_base.offset(src_stride0);
1108 let s20 = *src_base.offset(2 * src_stride0);
1109 let s30 = *src_base.offset(3 * src_stride0);
1110
1111 let src_col1 = src_base.offset(src_stride1);
1112 let s01 = *src_col1;
1113 let s11 = *src_col1.offset(src_stride0);
1114 let s21 = *src_col1.offset(2 * src_stride0);
1115 let s31 = *src_col1.offset(3 * src_stride0);
1116
1117 let src_col2 = src_base.offset(2 * src_stride1);
1118 let s02 = *src_col2;
1119 let s12 = *src_col2.offset(src_stride0);
1120 let s22 = *src_col2.offset(2 * src_stride0);
1121 let s32 = *src_col2.offset(3 * src_stride0);
1122
1123 let src_col3 = src_base.offset(3 * src_stride1);
1124 let s03 = *src_col3;
1125 let s13 = *src_col3.offset(src_stride0);
1126 let s23 = *src_col3.offset(2 * src_stride0);
1127 let s33 = *src_col3.offset(3 * src_stride0);
1128
1129 let dst_row0 = dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1);
1130 *dst_row0 = scale * s00;
1131 *dst_row0.offset(dst_stride0) = scale * s01;
1132 *dst_row0.offset(2 * dst_stride0) = scale * s02;
1133 *dst_row0.offset(3 * dst_stride0) = scale * s03;
1134
1135 let dst_row1 = dst_row0.offset(dst_stride1);
1136 *dst_row1 = scale * s10;
1137 *dst_row1.offset(dst_stride0) = scale * s11;
1138 *dst_row1.offset(2 * dst_stride0) = scale * s12;
1139 *dst_row1.offset(3 * dst_stride0) = scale * s13;
1140
1141 let dst_row2 = dst_row0.offset(2 * dst_stride1);
1142 *dst_row2 = scale * s20;
1143 *dst_row2.offset(dst_stride0) = scale * s21;
1144 *dst_row2.offset(2 * dst_stride0) = scale * s22;
1145 *dst_row2.offset(3 * dst_stride0) = scale * s23;
1146
1147 let dst_row3 = dst_row0.offset(3 * dst_stride1);
1148 *dst_row3 = scale * s30;
1149 *dst_row3.offset(dst_stride0) = scale * s31;
1150 *dst_row3.offset(2 * dst_stride0) = scale * s32;
1151 *dst_row3.offset(3 * dst_stride0) = scale * s33;
1152}
1153
1154#[inline]
1155unsafe fn copy_transpose_scale_2d_f64_tiled_raw(
1156 dst: *mut f64,
1157 dst_stride0: isize,
1158 dst_stride1: isize,
1159 src: *const f64,
1160 src_stride0: isize,
1161 src_stride1: isize,
1162 src_rows: usize,
1163 src_cols: usize,
1164 scale: f64,
1165) {
1166 const TILE: usize = 4;
1167 let row_full = src_rows / TILE * TILE;
1168 let col_full = src_cols / TILE * TILE;
1169
1170 #[cfg(feature = "parallel")]
1171 {
1172 let total = src_rows.saturating_mul(src_cols);
1173 let nthreads = crate::execution_policy::rayon_threads();
1174 if total > MINTHREADLENGTH && nthreads > 1 {
1175 let dst_send = SendPtr(dst);
1176 let src_send = SendPtr(src as *mut f64);
1177 let row_tiles = row_full / TILE;
1178 crate::threading::parallel_for_each(0..row_tiles, nthreads, &|tiles| {
1179 for tile_i in tiles {
1180 let i = tile_i * TILE;
1181 let dst = dst_send.as_ptr();
1182 let src = src_send.as_const();
1183 unsafe {
1184 let mut j = 0;
1185 while j < col_full {
1186 transpose_scale_4x4_f64(
1187 dst,
1188 dst_stride0,
1189 dst_stride1,
1190 src,
1191 src_stride0,
1192 src_stride1,
1193 i,
1194 j,
1195 scale,
1196 );
1197 j += TILE;
1198 }
1199 for j in col_full..src_cols {
1200 for ii in i..i + TILE {
1201 *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1202 scale
1203 * *src.offset(
1204 ii as isize * src_stride0 + j as isize * src_stride1,
1205 );
1206 }
1207 }
1208 }
1209 }
1210 });
1211 for i in row_full..src_rows {
1212 for j in 0..src_cols {
1213 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1214 scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1215 }
1216 }
1217 return;
1218 }
1219 }
1220
1221 let mut i = 0;
1222 while i < row_full {
1223 let mut j = 0;
1224 while j < col_full {
1225 transpose_scale_4x4_f64(
1226 dst,
1227 dst_stride0,
1228 dst_stride1,
1229 src,
1230 src_stride0,
1231 src_stride1,
1232 i,
1233 j,
1234 scale,
1235 );
1236 j += TILE;
1237 }
1238 for j in col_full..src_cols {
1239 for ii in i..i + TILE {
1240 *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1241 scale * *src.offset(ii as isize * src_stride0 + j as isize * src_stride1);
1242 }
1243 }
1244 i += TILE;
1245 }
1246 for i in row_full..src_rows {
1247 for j in 0..src_cols {
1248 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1249 scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1250 }
1251 }
1252}
1253
1254#[inline]
1255#[cfg(test)]
1256unsafe fn try_copy_transpose_scale_2d_f64_tiled(
1257 dest: &mut StridedViewMut<f64>,
1258 src: &StridedView<f64>,
1259 scale: f64,
1260) -> bool {
1261 if src.ndim() != 2 || dest.ndim() != 2 {
1262 return false;
1263 }
1264 let src_dims = src.dims();
1265 if dest.dims() != [src_dims[1], src_dims[0]] {
1266 return false;
1267 }
1268 if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1269 return false;
1270 }
1271
1272 copy_transpose_scale_2d_f64_tiled_raw(
1273 dest.as_mut_ptr(),
1274 dest.strides()[0],
1275 dest.strides()[1],
1276 src.ptr(),
1277 src.strides()[0],
1278 src.strides()[1],
1279 src_dims[0],
1280 src_dims[1],
1281 scale,
1282 );
1283 true
1284}
1285
1286#[inline]
1287unsafe fn try_copy_transpose_scale_2d_f64_tiled_typed<T>(
1288 dest: &mut StridedViewMut<T>,
1289 src: &StridedView<T>,
1290 scale: T,
1291) -> bool
1292where
1293 T: Copy + 'static,
1294{
1295 if std::any::TypeId::of::<T>() != std::any::TypeId::of::<f64>() {
1296 return false;
1297 }
1298
1299 let scale = *(&scale as *const T).cast::<f64>();
1300 if src.ndim() != 2 || dest.ndim() != 2 {
1301 return false;
1302 }
1303 let src_dims = src.dims();
1304 if dest.dims() != [src_dims[1], src_dims[0]] {
1305 return false;
1306 }
1307 if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1308 return false;
1309 }
1310
1311 copy_transpose_scale_2d_f64_tiled_raw(
1312 dest.as_mut_ptr().cast::<f64>(),
1313 dest.strides()[0],
1314 dest.strides()[1],
1315 src.ptr().cast::<f64>(),
1316 src.strides()[0],
1317 src.strides()[1],
1318 src_dims[0],
1319 src_dims[1],
1320 scale,
1321 );
1322 true
1323}
1324
1325#[inline(always)]
1326unsafe fn transpose_scale_4x4_identity<T>(
1327 dst: *mut T,
1328 dst_stride0: isize,
1329 dst_stride1: isize,
1330 src: *const T,
1331 src_stride0: isize,
1332 src_stride1: isize,
1333 i: usize,
1334 j: usize,
1335 scale: T,
1336) where
1337 T: Copy + Mul<Output = T>,
1338{
1339 let src_base = src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1340
1341 let s00 = *src_base;
1342 let s10 = *src_base.offset(src_stride0);
1343 let s20 = *src_base.offset(2 * src_stride0);
1344 let s30 = *src_base.offset(3 * src_stride0);
1345
1346 let src_col1 = src_base.offset(src_stride1);
1347 let s01 = *src_col1;
1348 let s11 = *src_col1.offset(src_stride0);
1349 let s21 = *src_col1.offset(2 * src_stride0);
1350 let s31 = *src_col1.offset(3 * src_stride0);
1351
1352 let src_col2 = src_base.offset(2 * src_stride1);
1353 let s02 = *src_col2;
1354 let s12 = *src_col2.offset(src_stride0);
1355 let s22 = *src_col2.offset(2 * src_stride0);
1356 let s32 = *src_col2.offset(3 * src_stride0);
1357
1358 let src_col3 = src_base.offset(3 * src_stride1);
1359 let s03 = *src_col3;
1360 let s13 = *src_col3.offset(src_stride0);
1361 let s23 = *src_col3.offset(2 * src_stride0);
1362 let s33 = *src_col3.offset(3 * src_stride0);
1363
1364 let dst_row0 = dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1);
1365 *dst_row0 = scale * s00;
1366 *dst_row0.offset(dst_stride0) = scale * s01;
1367 *dst_row0.offset(2 * dst_stride0) = scale * s02;
1368 *dst_row0.offset(3 * dst_stride0) = scale * s03;
1369
1370 let dst_row1 = dst_row0.offset(dst_stride1);
1371 *dst_row1 = scale * s10;
1372 *dst_row1.offset(dst_stride0) = scale * s11;
1373 *dst_row1.offset(2 * dst_stride0) = scale * s12;
1374 *dst_row1.offset(3 * dst_stride0) = scale * s13;
1375
1376 let dst_row2 = dst_row0.offset(2 * dst_stride1);
1377 *dst_row2 = scale * s20;
1378 *dst_row2.offset(dst_stride0) = scale * s21;
1379 *dst_row2.offset(2 * dst_stride0) = scale * s22;
1380 *dst_row2.offset(3 * dst_stride0) = scale * s23;
1381
1382 let dst_row3 = dst_row0.offset(3 * dst_stride1);
1383 *dst_row3 = scale * s30;
1384 *dst_row3.offset(dst_stride0) = scale * s31;
1385 *dst_row3.offset(2 * dst_stride0) = scale * s32;
1386 *dst_row3.offset(3 * dst_stride0) = scale * s33;
1387}
1388
1389#[inline]
1390unsafe fn copy_transpose_scale_2d_identity_tiled_raw<T>(
1391 dst: *mut T,
1392 dst_stride0: isize,
1393 dst_stride1: isize,
1394 src: *const T,
1395 src_stride0: isize,
1396 src_stride1: isize,
1397 src_rows: usize,
1398 src_cols: usize,
1399 scale: T,
1400) where
1401 T: Copy + Mul<Output = T> + MaybeSendSync,
1402{
1403 const TILE: usize = 4;
1404 let row_full = src_rows / TILE * TILE;
1405 let col_full = src_cols / TILE * TILE;
1406
1407 #[cfg(feature = "parallel")]
1408 {
1409 let total = src_rows.saturating_mul(src_cols);
1410 let nthreads = crate::execution_policy::rayon_threads();
1411 if total > MINTHREADLENGTH && nthreads > 1 {
1412 let dst_send = SendPtr(dst);
1413 let src_send = SendPtr(src as *mut T);
1414 let row_tiles = row_full / TILE;
1415 crate::threading::parallel_for_each(0..row_tiles, nthreads, &|tiles| {
1416 for tile_i in tiles {
1417 let i = tile_i * TILE;
1418 let dst = dst_send.as_ptr();
1419 let src = src_send.as_const();
1420 unsafe {
1421 let mut j = 0;
1422 while j < col_full {
1423 transpose_scale_4x4_identity(
1424 dst,
1425 dst_stride0,
1426 dst_stride1,
1427 src,
1428 src_stride0,
1429 src_stride1,
1430 i,
1431 j,
1432 scale,
1433 );
1434 j += TILE;
1435 }
1436 for j in col_full..src_cols {
1437 for ii in i..i + TILE {
1438 *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1439 scale
1440 * *src.offset(
1441 ii as isize * src_stride0 + j as isize * src_stride1,
1442 );
1443 }
1444 }
1445 }
1446 }
1447 });
1448 for i in row_full..src_rows {
1449 for j in 0..src_cols {
1450 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1451 scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1452 }
1453 }
1454 return;
1455 }
1456 }
1457
1458 let mut i = 0;
1459 while i < row_full {
1460 let mut j = 0;
1461 while j < col_full {
1462 transpose_scale_4x4_identity(
1463 dst,
1464 dst_stride0,
1465 dst_stride1,
1466 src,
1467 src_stride0,
1468 src_stride1,
1469 i,
1470 j,
1471 scale,
1472 );
1473 j += TILE;
1474 }
1475 for j in col_full..src_cols {
1476 for ii in i..i + TILE {
1477 *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1478 scale * *src.offset(ii as isize * src_stride0 + j as isize * src_stride1);
1479 }
1480 }
1481 i += TILE;
1482 }
1483 for i in row_full..src_rows {
1484 for j in 0..src_cols {
1485 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1486 scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1487 }
1488 }
1489}
1490
1491#[inline]
1492unsafe fn try_copy_transpose_scale_2d_identity_tiled<T>(
1493 dest: &mut StridedViewMut<T>,
1494 src: &StridedView<T>,
1495 scale: T,
1496) -> bool
1497where
1498 T: Copy + Mul<Output = T> + MaybeSendSync,
1499{
1500 if src.ndim() != 2 || dest.ndim() != 2 {
1501 return false;
1502 }
1503 let src_dims = src.dims();
1504 if dest.dims() != [src_dims[1], src_dims[0]] {
1505 return false;
1506 }
1507 if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1508 return false;
1509 }
1510
1511 copy_transpose_scale_2d_identity_tiled_raw(
1512 dest.as_mut_ptr(),
1513 dest.strides()[0],
1514 dest.strides()[1],
1515 src.ptr(),
1516 src.strides()[0],
1517 src.strides()[1],
1518 src_dims[0],
1519 src_dims[1],
1520 scale,
1521 );
1522 true
1523}
1524
1525#[inline]
1526unsafe fn copy_transpose_scale_2d_loop<T>(
1527 dst: *mut T,
1528 dst_stride0: isize,
1529 dst_stride1: isize,
1530 src: *const T,
1531 src_stride0: isize,
1532 src_stride1: isize,
1533 src_rows: usize,
1534 src_cols: usize,
1535 scale: T,
1536) where
1537 T: Copy + ElementOpApply + Mul<Output = T> + MaybeSendSync,
1538{
1539 #[cfg(feature = "parallel")]
1540 {
1541 let total = src_rows.saturating_mul(src_cols);
1542 let nthreads = crate::execution_policy::rayon_threads();
1543 if total > MINTHREADLENGTH && nthreads > 1 {
1544 let dst_send = SendPtr(dst);
1545 let src_send = SendPtr(src as *mut T);
1546 if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1547 crate::threading::parallel_for_each(0..src_rows, nthreads, &|rows| {
1548 for i in rows {
1549 let dst = dst_send.as_ptr();
1550 let src = src_send.as_const();
1551 unsafe {
1552 for j in 0..src_cols {
1553 let value = (*src
1554 .offset(i as isize * src_stride0 + j as isize * src_stride1))
1555 .transpose();
1556 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1557 scale * value;
1558 }
1559 }
1560 }
1561 });
1562 } else {
1563 crate::threading::parallel_for_each(0..src_cols, nthreads, &|columns| {
1564 for j in columns {
1565 let dst = dst_send.as_ptr();
1566 let src = src_send.as_const();
1567 unsafe {
1568 for i in 0..src_rows {
1569 let value = (*src
1570 .offset(i as isize * src_stride0 + j as isize * src_stride1))
1571 .transpose();
1572 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1573 scale * value;
1574 }
1575 }
1576 }
1577 });
1578 }
1579 return;
1580 }
1581 }
1582
1583 if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1584 for i in 0..src_rows {
1585 for j in 0..src_cols {
1586 let value =
1587 (*src.offset(i as isize * src_stride0 + j as isize * src_stride1)).transpose();
1588 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) = scale * value;
1589 }
1590 }
1591 } else {
1592 for j in 0..src_cols {
1593 for i in 0..src_rows {
1594 let value =
1595 (*src.offset(i as isize * src_stride0 + j as isize * src_stride1)).transpose();
1596 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) = scale * value;
1597 }
1598 }
1599 }
1600}
1601
1602pub fn copy_transpose_scale_into<T>(
1604 dest: &mut StridedViewMut<T>,
1605 src: &StridedView<T>,
1606 scale: T,
1607) -> Result<()>
1608where
1609 T: Copy + ElementOpApply + Mul<Output = T> + Zero + One + PartialEq + MaybeSendSync + 'static,
1610{
1611 if src.ndim() != 2 || dest.ndim() != 2 {
1612 return Err(StridedError::RankMismatch(src.ndim(), 2));
1613 }
1614
1615 let src_dims = src.dims();
1616 let expected_dims = [src_dims[1], src_dims[0]];
1617 ensure_same_shape(dest.dims(), &expected_dims)?;
1618 total_len(src_dims)?;
1621
1622 if scale == T::zero() {
1623 unsafe {
1624 if same_contiguous_layout(dest.dims(), &[dest.strides()]).is_some() {
1625 fill_contiguous(dest.as_mut_ptr(), total_len(dest.dims())?, T::zero());
1626 } else {
1627 fill_2d(
1628 dest.as_mut_ptr(),
1629 dest.dims()[0],
1630 dest.dims()[1],
1631 dest.strides()[0],
1632 dest.strides()[1],
1633 T::zero(),
1634 );
1635 }
1636 }
1637 return Ok(());
1638 }
1639
1640 let transpose_is_identity = element_transpose_is_identity::<T>();
1641
1642 unsafe {
1643 if transpose_is_identity && try_copy_transpose_scale_2d_f64_tiled_typed(dest, src, scale) {
1644 return Ok(());
1645 }
1646 if transpose_is_identity && try_copy_transpose_scale_2d_identity_tiled(dest, src, scale) {
1647 return Ok(());
1648 }
1649 }
1650
1651 if scale == T::one() && transpose_is_identity {
1652 let src_t = src.permute(&[1, 0])?;
1653 #[cfg(feature = "parallel")]
1654 {
1655 return crate::threading::copy_permuted_with_active_policy(dest, &src_t);
1656 }
1657 #[cfg(not(feature = "parallel"))]
1658 return crate::threading::copy_permuted_serial(dest, &src_t);
1659 }
1660
1661 unsafe {
1662 copy_transpose_scale_2d_loop(
1663 dest.as_mut_ptr(),
1664 dest.strides()[0],
1665 dest.strides()[1],
1666 src.ptr(),
1667 src.strides()[0],
1668 src.strides()[1],
1669 src_dims[0],
1670 src_dims[1],
1671 scale,
1672 );
1673 }
1674 Ok(())
1675}
1676
1677#[cfg(test)]
1678#[path = "ops_view/tests/tiled_tests.rs"]
1679mod tiled_tests;