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