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<
819 T: Copy + Zero + Add<Output = T> + MaybeSendSync + simd::MaybeSimdOps,
820 Op: ElementOp<T>,
821>(
822 src: &StridedView<T, Op>,
823) -> Result<T> {
824 if total_len(src.dims())? == 0 {
827 return Ok(T::zero());
828 }
829 if Op::IS_IDENTITY {
831 if same_contiguous_layout(src.dims(), &[src.strides()]).is_some() {
832 let len = total_len(src.dims())?;
833 let src_slice = unsafe { std::slice::from_raw_parts(src.ptr(), len) };
834
835 #[cfg(feature = "parallel")]
836 if len > MINTHREADLENGTH {
837 if let Some(result) = parallel_simd_sum(src_slice) {
838 return Ok(result);
839 }
840 }
841
842 if let Some(result) = T::try_simd_sum(src_slice) {
843 return Ok(result);
844 }
845 }
846 }
847 reduce(src, |x| x, |a, b| a + b, T::zero())
848}
849
850pub fn dot<A, B, R, OpA, OpB>(a: &StridedView<A, OpA>, b: &StridedView<B, OpB>) -> Result<R>
875where
876 A: Copy + Mul<B, Output = R> + MaybeSendSync + 'static,
877 B: Copy + MaybeSendSync + 'static,
878 R: Copy + Zero + Add<Output = R> + MaybeSendSync + simd::MaybeSimdOps + 'static,
879 OpA: ElementOp<A>,
880 OpB: ElementOp<B>,
881{
882 ensure_same_shape(a.dims(), b.dims())?;
883
884 let a_ptr = a.ptr();
885 let b_ptr = b.ptr();
886 let a_strides = a.strides();
887 let b_strides = b.strides();
888 let a_dims = a.dims();
889 let len = total_len(a_dims)?;
890
891 if len == 0 {
895 return Ok(R::zero());
896 }
897
898 if same_contiguous_layout(a_dims, &[a_strides, b_strides]).is_some() {
899 if OpA::IS_IDENTITY
901 && OpB::IS_IDENTITY
902 && std::any::TypeId::of::<A>() == std::any::TypeId::of::<R>()
903 && std::any::TypeId::of::<B>() == std::any::TypeId::of::<R>()
904 {
905 let sa = unsafe { std::slice::from_raw_parts(a_ptr as *const R, len) };
906 let sb = unsafe { std::slice::from_raw_parts(b_ptr as *const R, len) };
907 if let Some(result) = R::try_simd_dot(sa, sb) {
908 return Ok(result);
909 }
910 }
911
912 let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
914 let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
915 let mut acc = R::zero();
916 simd::dispatch_if_large(len, || {
917 for i in 0..len {
918 acc = acc + OpA::apply(sa[i]) * OpB::apply(sb[i]);
919 }
920 });
921 return Ok(acc);
922 }
923
924 let strides_list: [&[isize]; 2] = [a_strides, b_strides];
925 let elem_size = std::mem::size_of::<A>()
926 .max(std::mem::size_of::<B>())
927 .max(std::mem::size_of::<R>());
928
929 let total = total_len(a_dims)?;
930 let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
931 build_plan_fused_small(a_dims, &strides_list)
932 } else {
933 build_plan_fused(a_dims, &strides_list, None, elem_size)
934 };
935
936 let mut acc = R::zero();
937 let initial_offsets: crate::inline::OperandVec<isize> =
938 ordered_strides.iter().map(|_| 0isize).collect();
939 for_each_inner_block_preordered(
940 &fused_dims,
941 &plan.block,
942 &ordered_strides,
943 &initial_offsets,
944 |offsets, len, strides| {
945 acc = unsafe {
946 inner_loop_dot::<A, B, R, OpA, OpB>(
947 a_ptr.offset(offsets[0]),
948 strides[0],
949 b_ptr.offset(offsets[1]),
950 strides[1],
951 len,
952 acc,
953 )
954 };
955 Ok(())
956 },
957 )?;
958
959 Ok(acc)
960}
961
962pub fn symmetrize_into<T>(dest: &mut StridedViewMut<T>, src: &StridedView<T>) -> Result<()>
964where
965 T: Copy
966 + Add<Output = T>
967 + Mul<Output = T>
968 + num_traits::FromPrimitive
969 + std::ops::Div<Output = T>
970 + MaybeSendSync,
971{
972 if src.ndim() != 2 {
973 return Err(StridedError::RankMismatch(src.ndim(), 2));
974 }
975 let rows = src.dims()[0];
976 let cols = src.dims()[1];
977 if rows != cols {
978 return Err(StridedError::NonSquare { rows, cols });
979 }
980
981 let src_t = src.permute(&[1, 0])?;
982 let half = T::from_f64(0.5).ok_or(StridedError::ScalarConversion)?;
983
984 zip_map2_into(dest, src, &src_t, |a, b| (a + b) * half)
985}
986
987pub fn symmetrize_conj_into<T>(dest: &mut StridedViewMut<T>, src: &StridedView<T>) -> Result<()>
989where
990 T: Copy
991 + ElementOpApply
992 + Add<Output = T>
993 + Mul<Output = T>
994 + num_traits::FromPrimitive
995 + std::ops::Div<Output = T>
996 + MaybeSendSync,
997{
998 if src.ndim() != 2 {
999 return Err(StridedError::RankMismatch(src.ndim(), 2));
1000 }
1001 let rows = src.dims()[0];
1002 let cols = src.dims()[1];
1003 if rows != cols {
1004 return Err(StridedError::NonSquare { rows, cols });
1005 }
1006
1007 let src_adj = src.adjoint_2d()?;
1009 let half = T::from_f64(0.5).ok_or(StridedError::ScalarConversion)?;
1010
1011 zip_map2_into(dest, src, &src_adj, |a, b| (a + b) * half)
1012}
1013
1014pub fn copy_scale<D, S, A, Op>(
1018 dest: &mut StridedViewMut<D>,
1019 src: &StridedView<S, Op>,
1020 scale: A,
1021) -> Result<()>
1022where
1023 A: Copy + Mul<S, Output = D> + MaybeSync,
1024 D: Copy + MaybeSendSync,
1025 S: Copy + MaybeSendSync,
1026 Op: ElementOp<S>,
1027{
1028 map_into(dest, src, |x| scale * x)
1029}
1030
1031pub fn copy_conj<T: Copy + ElementOpApply + MaybeSendSync>(
1033 dest: &mut StridedViewMut<T>,
1034 src: &StridedView<T>,
1035) -> Result<()> {
1036 let src_conj = src.conj();
1037 copy_into(dest, &src_conj)
1038}
1039
1040#[inline]
1041fn element_transpose_is_identity<T: 'static>() -> bool {
1042 use std::any::TypeId;
1043
1044 macro_rules! matches_type {
1045 ($($ty:ty),* $(,)?) => {{
1046 let id = TypeId::of::<T>();
1047 false $(|| id == TypeId::of::<$ty>())*
1048 }};
1049 }
1050
1051 matches_type!(
1052 f32,
1053 f64,
1054 i8,
1055 i16,
1056 i32,
1057 i64,
1058 i128,
1059 isize,
1060 u8,
1061 u16,
1062 u32,
1063 u64,
1064 u128,
1065 usize,
1066 num_complex::Complex32,
1067 num_complex::Complex64,
1068 )
1069}
1070
1071#[inline]
1072fn element_zero_is_all_bits_zero<T: 'static>() -> bool {
1073 use std::any::TypeId;
1074
1075 macro_rules! matches_type {
1076 ($($ty:ty),* $(,)?) => {{
1077 let id = TypeId::of::<T>();
1078 false $(|| id == TypeId::of::<$ty>())*
1079 }};
1080 }
1081
1082 matches_type!(
1083 f32,
1084 f64,
1085 i8,
1086 i16,
1087 i32,
1088 i64,
1089 i128,
1090 isize,
1091 u8,
1092 u16,
1093 u32,
1094 u64,
1095 u128,
1096 usize,
1097 num_complex::Complex32,
1098 num_complex::Complex64,
1099 )
1100}
1101
1102#[inline]
1103unsafe fn fill_2d<T: Copy + MaybeSendSync>(
1104 dst: *mut T,
1105 dim0: usize,
1106 dim1: usize,
1107 dst_stride0: isize,
1108 dst_stride1: isize,
1109 value: T,
1110) {
1111 #[cfg(feature = "parallel")]
1112 {
1113 let total = dim0.saturating_mul(dim1);
1114 let nthreads = crate::execution_policy::rayon_threads();
1115 if total > MINTHREADLENGTH && nthreads > 1 {
1116 let dst_send = SendPtr(dst);
1117 if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1118 crate::threading::parallel_for_each(0..dim1, nthreads, &|columns| {
1119 for j in columns {
1120 let dst = dst_send.as_ptr();
1121 unsafe {
1122 let base = j as isize * dst_stride1;
1123 for i in 0..dim0 {
1124 *dst.offset(base + i as isize * dst_stride0) = value;
1125 }
1126 }
1127 }
1128 });
1129 } else {
1130 crate::threading::parallel_for_each(0..dim0, nthreads, &|rows| {
1131 for i in rows {
1132 let dst = dst_send.as_ptr();
1133 unsafe {
1134 let base = i as isize * dst_stride0;
1135 for j in 0..dim1 {
1136 *dst.offset(base + j as isize * dst_stride1) = value;
1137 }
1138 }
1139 }
1140 });
1141 }
1142 return;
1143 }
1144 }
1145
1146 if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1147 for j in 0..dim1 {
1148 let base = j as isize * dst_stride1;
1149 for i in 0..dim0 {
1150 *dst.offset(base + i as isize * dst_stride0) = value;
1151 }
1152 }
1153 } else {
1154 for i in 0..dim0 {
1155 let base = i as isize * dst_stride0;
1156 for j in 0..dim1 {
1157 *dst.offset(base + j as isize * dst_stride1) = value;
1158 }
1159 }
1160 }
1161}
1162
1163#[inline]
1164unsafe fn fill_contiguous<T>(dst: *mut T, len: usize, value: T)
1165where
1166 T: Copy + Zero + PartialEq + MaybeSendSync + 'static,
1167{
1168 if element_zero_is_all_bits_zero::<T>() && value == T::zero() {
1169 std::ptr::write_bytes(dst, 0, len);
1170 return;
1171 }
1172
1173 let dst = std::slice::from_raw_parts_mut(dst, len);
1174 dst.fill(value);
1175}
1176
1177#[inline(always)]
1178unsafe fn transpose_scale_4x4_f64(
1179 dst: *mut f64,
1180 dst_stride0: isize,
1181 dst_stride1: isize,
1182 src: *const f64,
1183 src_stride0: isize,
1184 src_stride1: isize,
1185 i: usize,
1186 j: usize,
1187 scale: f64,
1188) {
1189 let src_base = src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1190
1191 let s00 = *src_base;
1192 let s10 = *src_base.offset(src_stride0);
1193 let s20 = *src_base.offset(2 * src_stride0);
1194 let s30 = *src_base.offset(3 * src_stride0);
1195
1196 let src_col1 = src_base.offset(src_stride1);
1197 let s01 = *src_col1;
1198 let s11 = *src_col1.offset(src_stride0);
1199 let s21 = *src_col1.offset(2 * src_stride0);
1200 let s31 = *src_col1.offset(3 * src_stride0);
1201
1202 let src_col2 = src_base.offset(2 * src_stride1);
1203 let s02 = *src_col2;
1204 let s12 = *src_col2.offset(src_stride0);
1205 let s22 = *src_col2.offset(2 * src_stride0);
1206 let s32 = *src_col2.offset(3 * src_stride0);
1207
1208 let src_col3 = src_base.offset(3 * src_stride1);
1209 let s03 = *src_col3;
1210 let s13 = *src_col3.offset(src_stride0);
1211 let s23 = *src_col3.offset(2 * src_stride0);
1212 let s33 = *src_col3.offset(3 * src_stride0);
1213
1214 let dst_row0 = dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1);
1215 *dst_row0 = scale * s00;
1216 *dst_row0.offset(dst_stride0) = scale * s01;
1217 *dst_row0.offset(2 * dst_stride0) = scale * s02;
1218 *dst_row0.offset(3 * dst_stride0) = scale * s03;
1219
1220 let dst_row1 = dst_row0.offset(dst_stride1);
1221 *dst_row1 = scale * s10;
1222 *dst_row1.offset(dst_stride0) = scale * s11;
1223 *dst_row1.offset(2 * dst_stride0) = scale * s12;
1224 *dst_row1.offset(3 * dst_stride0) = scale * s13;
1225
1226 let dst_row2 = dst_row0.offset(2 * dst_stride1);
1227 *dst_row2 = scale * s20;
1228 *dst_row2.offset(dst_stride0) = scale * s21;
1229 *dst_row2.offset(2 * dst_stride0) = scale * s22;
1230 *dst_row2.offset(3 * dst_stride0) = scale * s23;
1231
1232 let dst_row3 = dst_row0.offset(3 * dst_stride1);
1233 *dst_row3 = scale * s30;
1234 *dst_row3.offset(dst_stride0) = scale * s31;
1235 *dst_row3.offset(2 * dst_stride0) = scale * s32;
1236 *dst_row3.offset(3 * dst_stride0) = scale * s33;
1237}
1238
1239#[inline]
1240unsafe fn copy_transpose_scale_2d_f64_tiled_raw(
1241 dst: *mut f64,
1242 dst_stride0: isize,
1243 dst_stride1: isize,
1244 src: *const f64,
1245 src_stride0: isize,
1246 src_stride1: isize,
1247 src_rows: usize,
1248 src_cols: usize,
1249 scale: f64,
1250) {
1251 const TILE: usize = 4;
1252 let row_full = src_rows / TILE * TILE;
1253 let col_full = src_cols / TILE * TILE;
1254
1255 #[cfg(feature = "parallel")]
1256 {
1257 let total = src_rows.saturating_mul(src_cols);
1258 let nthreads = crate::execution_policy::rayon_threads();
1259 if total > MINTHREADLENGTH && nthreads > 1 {
1260 let dst_send = SendPtr(dst);
1261 let src_send = SendPtr(src as *mut f64);
1262 let row_tiles = row_full / TILE;
1263 crate::threading::parallel_for_each(0..row_tiles, nthreads, &|tiles| {
1264 for tile_i in tiles {
1265 let i = tile_i * TILE;
1266 let dst = dst_send.as_ptr();
1267 let src = src_send.as_const();
1268 unsafe {
1269 let mut j = 0;
1270 while j < col_full {
1271 transpose_scale_4x4_f64(
1272 dst,
1273 dst_stride0,
1274 dst_stride1,
1275 src,
1276 src_stride0,
1277 src_stride1,
1278 i,
1279 j,
1280 scale,
1281 );
1282 j += TILE;
1283 }
1284 for j in col_full..src_cols {
1285 for ii in i..i + TILE {
1286 *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1287 scale
1288 * *src.offset(
1289 ii as isize * src_stride0 + j as isize * src_stride1,
1290 );
1291 }
1292 }
1293 }
1294 }
1295 });
1296 for i in row_full..src_rows {
1297 for j in 0..src_cols {
1298 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1299 scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1300 }
1301 }
1302 return;
1303 }
1304 }
1305
1306 let mut i = 0;
1307 while i < row_full {
1308 let mut j = 0;
1309 while j < col_full {
1310 transpose_scale_4x4_f64(
1311 dst,
1312 dst_stride0,
1313 dst_stride1,
1314 src,
1315 src_stride0,
1316 src_stride1,
1317 i,
1318 j,
1319 scale,
1320 );
1321 j += TILE;
1322 }
1323 for j in col_full..src_cols {
1324 for ii in i..i + TILE {
1325 *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1326 scale * *src.offset(ii as isize * src_stride0 + j as isize * src_stride1);
1327 }
1328 }
1329 i += TILE;
1330 }
1331 for i in row_full..src_rows {
1332 for j in 0..src_cols {
1333 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1334 scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1335 }
1336 }
1337}
1338
1339#[inline]
1340#[cfg(test)]
1341unsafe fn try_copy_transpose_scale_2d_f64_tiled(
1342 dest: &mut StridedViewMut<f64>,
1343 src: &StridedView<f64>,
1344 scale: f64,
1345) -> bool {
1346 if src.ndim() != 2 || dest.ndim() != 2 {
1347 return false;
1348 }
1349 let src_dims = src.dims();
1350 if dest.dims() != [src_dims[1], src_dims[0]] {
1351 return false;
1352 }
1353 if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1354 return false;
1355 }
1356
1357 copy_transpose_scale_2d_f64_tiled_raw(
1358 dest.as_mut_ptr(),
1359 dest.strides()[0],
1360 dest.strides()[1],
1361 src.ptr(),
1362 src.strides()[0],
1363 src.strides()[1],
1364 src_dims[0],
1365 src_dims[1],
1366 scale,
1367 );
1368 true
1369}
1370
1371#[inline]
1372unsafe fn try_copy_transpose_scale_2d_f64_tiled_typed<T>(
1373 dest: &mut StridedViewMut<T>,
1374 src: &StridedView<T>,
1375 scale: T,
1376) -> bool
1377where
1378 T: Copy + 'static,
1379{
1380 if std::any::TypeId::of::<T>() != std::any::TypeId::of::<f64>() {
1381 return false;
1382 }
1383
1384 let scale = *(&scale as *const T).cast::<f64>();
1385 if src.ndim() != 2 || dest.ndim() != 2 {
1386 return false;
1387 }
1388 let src_dims = src.dims();
1389 if dest.dims() != [src_dims[1], src_dims[0]] {
1390 return false;
1391 }
1392 if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1393 return false;
1394 }
1395
1396 copy_transpose_scale_2d_f64_tiled_raw(
1397 dest.as_mut_ptr().cast::<f64>(),
1398 dest.strides()[0],
1399 dest.strides()[1],
1400 src.ptr().cast::<f64>(),
1401 src.strides()[0],
1402 src.strides()[1],
1403 src_dims[0],
1404 src_dims[1],
1405 scale,
1406 );
1407 true
1408}
1409
1410#[inline(always)]
1411unsafe fn transpose_scale_4x4_identity<T>(
1412 dst: *mut T,
1413 dst_stride0: isize,
1414 dst_stride1: isize,
1415 src: *const T,
1416 src_stride0: isize,
1417 src_stride1: isize,
1418 i: usize,
1419 j: usize,
1420 scale: T,
1421) where
1422 T: Copy + Mul<Output = T>,
1423{
1424 let src_base = src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1425
1426 let s00 = *src_base;
1427 let s10 = *src_base.offset(src_stride0);
1428 let s20 = *src_base.offset(2 * src_stride0);
1429 let s30 = *src_base.offset(3 * src_stride0);
1430
1431 let src_col1 = src_base.offset(src_stride1);
1432 let s01 = *src_col1;
1433 let s11 = *src_col1.offset(src_stride0);
1434 let s21 = *src_col1.offset(2 * src_stride0);
1435 let s31 = *src_col1.offset(3 * src_stride0);
1436
1437 let src_col2 = src_base.offset(2 * src_stride1);
1438 let s02 = *src_col2;
1439 let s12 = *src_col2.offset(src_stride0);
1440 let s22 = *src_col2.offset(2 * src_stride0);
1441 let s32 = *src_col2.offset(3 * src_stride0);
1442
1443 let src_col3 = src_base.offset(3 * src_stride1);
1444 let s03 = *src_col3;
1445 let s13 = *src_col3.offset(src_stride0);
1446 let s23 = *src_col3.offset(2 * src_stride0);
1447 let s33 = *src_col3.offset(3 * src_stride0);
1448
1449 let dst_row0 = dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1);
1450 *dst_row0 = scale * s00;
1451 *dst_row0.offset(dst_stride0) = scale * s01;
1452 *dst_row0.offset(2 * dst_stride0) = scale * s02;
1453 *dst_row0.offset(3 * dst_stride0) = scale * s03;
1454
1455 let dst_row1 = dst_row0.offset(dst_stride1);
1456 *dst_row1 = scale * s10;
1457 *dst_row1.offset(dst_stride0) = scale * s11;
1458 *dst_row1.offset(2 * dst_stride0) = scale * s12;
1459 *dst_row1.offset(3 * dst_stride0) = scale * s13;
1460
1461 let dst_row2 = dst_row0.offset(2 * dst_stride1);
1462 *dst_row2 = scale * s20;
1463 *dst_row2.offset(dst_stride0) = scale * s21;
1464 *dst_row2.offset(2 * dst_stride0) = scale * s22;
1465 *dst_row2.offset(3 * dst_stride0) = scale * s23;
1466
1467 let dst_row3 = dst_row0.offset(3 * dst_stride1);
1468 *dst_row3 = scale * s30;
1469 *dst_row3.offset(dst_stride0) = scale * s31;
1470 *dst_row3.offset(2 * dst_stride0) = scale * s32;
1471 *dst_row3.offset(3 * dst_stride0) = scale * s33;
1472}
1473
1474#[inline]
1475unsafe fn copy_transpose_scale_2d_identity_tiled_raw<T>(
1476 dst: *mut T,
1477 dst_stride0: isize,
1478 dst_stride1: isize,
1479 src: *const T,
1480 src_stride0: isize,
1481 src_stride1: isize,
1482 src_rows: usize,
1483 src_cols: usize,
1484 scale: T,
1485) where
1486 T: Copy + Mul<Output = T> + MaybeSendSync,
1487{
1488 const TILE: usize = 4;
1489 let row_full = src_rows / TILE * TILE;
1490 let col_full = src_cols / TILE * TILE;
1491
1492 #[cfg(feature = "parallel")]
1493 {
1494 let total = src_rows.saturating_mul(src_cols);
1495 let nthreads = crate::execution_policy::rayon_threads();
1496 if total > MINTHREADLENGTH && nthreads > 1 {
1497 let dst_send = SendPtr(dst);
1498 let src_send = SendPtr(src as *mut T);
1499 let row_tiles = row_full / TILE;
1500 crate::threading::parallel_for_each(0..row_tiles, nthreads, &|tiles| {
1501 for tile_i in tiles {
1502 let i = tile_i * TILE;
1503 let dst = dst_send.as_ptr();
1504 let src = src_send.as_const();
1505 unsafe {
1506 let mut j = 0;
1507 while j < col_full {
1508 transpose_scale_4x4_identity(
1509 dst,
1510 dst_stride0,
1511 dst_stride1,
1512 src,
1513 src_stride0,
1514 src_stride1,
1515 i,
1516 j,
1517 scale,
1518 );
1519 j += TILE;
1520 }
1521 for j in col_full..src_cols {
1522 for ii in i..i + TILE {
1523 *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1524 scale
1525 * *src.offset(
1526 ii as isize * src_stride0 + j as isize * src_stride1,
1527 );
1528 }
1529 }
1530 }
1531 }
1532 });
1533 for i in row_full..src_rows {
1534 for j in 0..src_cols {
1535 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1536 scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1537 }
1538 }
1539 return;
1540 }
1541 }
1542
1543 let mut i = 0;
1544 while i < row_full {
1545 let mut j = 0;
1546 while j < col_full {
1547 transpose_scale_4x4_identity(
1548 dst,
1549 dst_stride0,
1550 dst_stride1,
1551 src,
1552 src_stride0,
1553 src_stride1,
1554 i,
1555 j,
1556 scale,
1557 );
1558 j += TILE;
1559 }
1560 for j in col_full..src_cols {
1561 for ii in i..i + TILE {
1562 *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1563 scale * *src.offset(ii as isize * src_stride0 + j as isize * src_stride1);
1564 }
1565 }
1566 i += TILE;
1567 }
1568 for i in row_full..src_rows {
1569 for j in 0..src_cols {
1570 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1571 scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1572 }
1573 }
1574}
1575
1576#[inline]
1577unsafe fn try_copy_transpose_scale_2d_identity_tiled<T>(
1578 dest: &mut StridedViewMut<T>,
1579 src: &StridedView<T>,
1580 scale: T,
1581) -> bool
1582where
1583 T: Copy + Mul<Output = T> + MaybeSendSync,
1584{
1585 if src.ndim() != 2 || dest.ndim() != 2 {
1586 return false;
1587 }
1588 let src_dims = src.dims();
1589 if dest.dims() != [src_dims[1], src_dims[0]] {
1590 return false;
1591 }
1592 if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1593 return false;
1594 }
1595
1596 copy_transpose_scale_2d_identity_tiled_raw(
1597 dest.as_mut_ptr(),
1598 dest.strides()[0],
1599 dest.strides()[1],
1600 src.ptr(),
1601 src.strides()[0],
1602 src.strides()[1],
1603 src_dims[0],
1604 src_dims[1],
1605 scale,
1606 );
1607 true
1608}
1609
1610#[inline]
1611unsafe fn copy_transpose_scale_2d_loop<T>(
1612 dst: *mut T,
1613 dst_stride0: isize,
1614 dst_stride1: isize,
1615 src: *const T,
1616 src_stride0: isize,
1617 src_stride1: isize,
1618 src_rows: usize,
1619 src_cols: usize,
1620 scale: T,
1621) where
1622 T: Copy + ElementOpApply + Mul<Output = T> + MaybeSendSync,
1623{
1624 #[cfg(feature = "parallel")]
1625 {
1626 let total = src_rows.saturating_mul(src_cols);
1627 let nthreads = crate::execution_policy::rayon_threads();
1628 if total > MINTHREADLENGTH && nthreads > 1 {
1629 let dst_send = SendPtr(dst);
1630 let src_send = SendPtr(src as *mut T);
1631 if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1632 crate::threading::parallel_for_each(0..src_rows, nthreads, &|rows| {
1633 for i in rows {
1634 let dst = dst_send.as_ptr();
1635 let src = src_send.as_const();
1636 unsafe {
1637 for j in 0..src_cols {
1638 let value = (*src
1639 .offset(i as isize * src_stride0 + j as isize * src_stride1))
1640 .transpose();
1641 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1642 scale * value;
1643 }
1644 }
1645 }
1646 });
1647 } else {
1648 crate::threading::parallel_for_each(0..src_cols, nthreads, &|columns| {
1649 for j in columns {
1650 let dst = dst_send.as_ptr();
1651 let src = src_send.as_const();
1652 unsafe {
1653 for i in 0..src_rows {
1654 let value = (*src
1655 .offset(i as isize * src_stride0 + j as isize * src_stride1))
1656 .transpose();
1657 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1658 scale * value;
1659 }
1660 }
1661 }
1662 });
1663 }
1664 return;
1665 }
1666 }
1667
1668 if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1669 for i in 0..src_rows {
1670 for j in 0..src_cols {
1671 let value =
1672 (*src.offset(i as isize * src_stride0 + j as isize * src_stride1)).transpose();
1673 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) = scale * value;
1674 }
1675 }
1676 } else {
1677 for j in 0..src_cols {
1678 for i in 0..src_rows {
1679 let value =
1680 (*src.offset(i as isize * src_stride0 + j as isize * src_stride1)).transpose();
1681 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) = scale * value;
1682 }
1683 }
1684 }
1685}
1686
1687pub fn copy_transpose_scale_into<T>(
1689 dest: &mut StridedViewMut<T>,
1690 src: &StridedView<T>,
1691 scale: T,
1692) -> Result<()>
1693where
1694 T: Copy + ElementOpApply + Mul<Output = T> + Zero + One + PartialEq + MaybeSendSync + 'static,
1695{
1696 if src.ndim() != 2 || dest.ndim() != 2 {
1697 return Err(StridedError::RankMismatch(src.ndim(), 2));
1698 }
1699
1700 let src_dims = src.dims();
1701 let expected_dims = [src_dims[1], src_dims[0]];
1702 ensure_same_shape(dest.dims(), &expected_dims)?;
1703 total_len(src_dims)?;
1706
1707 if scale == T::zero() {
1708 unsafe {
1709 if same_contiguous_layout(dest.dims(), &[dest.strides()]).is_some() {
1710 fill_contiguous(dest.as_mut_ptr(), total_len(dest.dims())?, T::zero());
1711 } else {
1712 fill_2d(
1713 dest.as_mut_ptr(),
1714 dest.dims()[0],
1715 dest.dims()[1],
1716 dest.strides()[0],
1717 dest.strides()[1],
1718 T::zero(),
1719 );
1720 }
1721 }
1722 return Ok(());
1723 }
1724
1725 let transpose_is_identity = element_transpose_is_identity::<T>();
1726
1727 unsafe {
1728 if transpose_is_identity && try_copy_transpose_scale_2d_f64_tiled_typed(dest, src, scale) {
1729 return Ok(());
1730 }
1731 if transpose_is_identity && try_copy_transpose_scale_2d_identity_tiled(dest, src, scale) {
1732 return Ok(());
1733 }
1734 }
1735
1736 if scale == T::one() && transpose_is_identity {
1737 let src_t = src.permute(&[1, 0])?;
1738 #[cfg(feature = "parallel")]
1739 {
1740 return crate::threading::copy_permuted_with_active_policy(dest, &src_t);
1741 }
1742 #[cfg(not(feature = "parallel"))]
1743 return crate::threading::copy_permuted_serial(dest, &src_t);
1744 }
1745
1746 unsafe {
1747 copy_transpose_scale_2d_loop(
1748 dest.as_mut_ptr(),
1749 dest.strides()[0],
1750 dest.strides()[1],
1751 src.ptr(),
1752 src.strides()[0],
1753 src.strides()[1],
1754 src_dims[0],
1755 src_dims[1],
1756 scale,
1757 );
1758 }
1759 Ok(())
1760}
1761
1762#[cfg(test)]
1763#[path = "ops_view/tests/tiled_tests.rs"]
1764mod tiled_tests;