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::{copy_map_into, map_into, zip_map2_into};
8use crate::reduce_view::reduce;
9use crate::simd;
10use crate::view::{StridedView, StridedViewMut};
11use crate::{Exec, Result, StridedError};
12use num_traits::{One, Zero};
13use std::mem::MaybeUninit;
14use std::ops::{Add, Mul};
15use strided_view::{ElementOp, ElementOpApply};
16
17use crate::cost::KernelClass;
18use crate::fuse::compute_costs;
19use crate::threading::{
20 for_each_inner_block_with_offsets, is_parallel, mapreduce_threaded, SendPtr,
21};
22#[inline(always)]
34unsafe fn inner_loop_add<D: Copy + Add<S, Output = D>, S: Copy, Op: ElementOp<S>>(
35 dp: *mut D,
36 ds: isize,
37 sp: *const S,
38 ss: isize,
39 len: usize,
40) {
41 if ds == 1 && ss == 1 {
42 let dst = std::slice::from_raw_parts_mut(dp, len);
43 let src = std::slice::from_raw_parts(sp, len);
44 simd::dispatch_if_large(len, || {
45 for i in 0..len {
46 dst[i] = dst[i] + Op::apply(src[i]);
47 }
48 });
49 } else {
50 let mut dp = dp;
51 let mut sp = sp;
52 for _ in 0..len {
53 *dp = *dp + Op::apply(*sp);
54 dp = dp.wrapping_offset(ds);
55 sp = sp.wrapping_offset(ss);
56 }
57 }
58}
59
60#[inline(always)]
62unsafe fn inner_loop_mul<D: Copy + Mul<S, Output = D>, S: Copy, Op: ElementOp<S>>(
63 dp: *mut D,
64 ds: isize,
65 sp: *const S,
66 ss: isize,
67 len: usize,
68) {
69 if ds == 1 && ss == 1 {
70 let dst = std::slice::from_raw_parts_mut(dp, len);
71 let src = std::slice::from_raw_parts(sp, len);
72 simd::dispatch_if_large(len, || {
73 for i in 0..len {
74 dst[i] = dst[i] * Op::apply(src[i]);
75 }
76 });
77 } else {
78 let mut dp = dp;
79 let mut sp = sp;
80 for _ in 0..len {
81 *dp = *dp * Op::apply(*sp);
82 dp = dp.wrapping_offset(ds);
83 sp = sp.wrapping_offset(ss);
84 }
85 }
86}
87
88#[inline(always)]
90unsafe fn inner_loop_axpy<
91 D: Copy + Add<D, Output = D>,
92 S: Copy,
93 A: Copy + Mul<S, Output = D>,
94 Op: ElementOp<S>,
95>(
96 dp: *mut D,
97 ds: isize,
98 sp: *const S,
99 ss: isize,
100 len: usize,
101 alpha: A,
102) {
103 if ds == 1 && ss == 1 {
104 let dst = std::slice::from_raw_parts_mut(dp, len);
105 let src = std::slice::from_raw_parts(sp, len);
106 simd::dispatch_if_large(len, || {
107 for i in 0..len {
108 dst[i] = alpha * Op::apply(src[i]) + dst[i];
109 }
110 });
111 } else {
112 let mut dp = dp;
113 let mut sp = sp;
114 for _ in 0..len {
115 *dp = alpha * Op::apply(*sp) + *dp;
116 dp = dp.wrapping_offset(ds);
117 sp = sp.wrapping_offset(ss);
118 }
119 }
120}
121
122#[inline(always)]
124unsafe fn inner_loop_fma<
125 D: Copy + Add<D, Output = D>,
126 A: Copy + Mul<B, Output = D>,
127 B: Copy,
128 OpA: ElementOp<A>,
129 OpB: ElementOp<B>,
130>(
131 dp: *mut D,
132 ds: isize,
133 ap: *const A,
134 a_s: isize,
135 bp: *const B,
136 b_s: isize,
137 len: usize,
138) {
139 if ds == 1 && a_s == 1 && b_s == 1 {
140 let dst = std::slice::from_raw_parts_mut(dp, len);
141 let sa = std::slice::from_raw_parts(ap, len);
142 let sb = std::slice::from_raw_parts(bp, len);
143 simd::dispatch_if_large(len, || {
144 for i in 0..len {
145 dst[i] = dst[i] + OpA::apply(sa[i]) * OpB::apply(sb[i]);
146 }
147 });
148 } else {
149 let mut dp = dp;
150 let mut ap = ap;
151 let mut bp = bp;
152 for _ in 0..len {
153 *dp = *dp + OpA::apply(*ap) * OpB::apply(*bp);
154 dp = dp.wrapping_offset(ds);
155 ap = ap.wrapping_offset(a_s);
156 bp = bp.wrapping_offset(b_s);
157 }
158 }
159}
160
161#[inline(always)]
163unsafe fn inner_loop_dot<
164 A: Copy + Mul<B, Output = R>,
165 B: Copy,
166 R: Copy + Add<R, Output = R>,
167 OpA: ElementOp<A>,
168 OpB: ElementOp<B>,
169>(
170 ap: *const A,
171 a_s: isize,
172 bp: *const B,
173 b_s: isize,
174 len: usize,
175 mut acc: R,
176) -> R {
177 if a_s == 1 && b_s == 1 {
178 let sa = std::slice::from_raw_parts(ap, len);
179 let sb = std::slice::from_raw_parts(bp, len);
180 simd::dispatch_if_large(len, || {
181 for i in 0..len {
182 acc = acc + OpA::apply(sa[i]) * OpB::apply(sb[i]);
183 }
184 });
185 } else {
186 let mut ap = ap;
187 let mut bp = bp;
188 for _ in 0..len {
189 acc = acc + OpA::apply(*ap) * OpB::apply(*bp);
190 ap = ap.wrapping_offset(a_s);
191 bp = bp.wrapping_offset(b_s);
192 }
193 }
194 acc
195}
196
197pub fn copy_into<T: Copy + Send + Sync, Op: ElementOp<T>>(
199 dest: &mut StridedViewMut<T>,
200 src: &StridedView<T, Op>,
201 exec: Exec<'_>,
202) -> Result<()> {
203 ensure_same_shape(dest.dims(), src.dims())?;
204
205 let dst_ptr = dest.as_mut_ptr();
206 let src_ptr = src.ptr();
207 let dst_dims = dest.dims();
208 let dst_strides = dest.strides();
209 let src_strides = src.strides();
210
211 if sequential_contiguous_layout(
212 KernelClass::Copy,
213 dst_dims,
214 &[dst_strides, src_strides],
215 &exec,
216 )?
217 .is_some()
218 {
219 let len = total_len(dst_dims)?;
220 if Op::IS_IDENTITY {
221 debug_assert!(
222 {
223 let nbytes = len
224 .checked_mul(std::mem::size_of::<T>())
225 .expect("copy size must not overflow");
226 let dst_start = dst_ptr as usize;
227 let src_start = src_ptr as usize;
228 let dst_end = dst_start.saturating_add(nbytes);
229 let src_end = src_start.saturating_add(nbytes);
230 dst_end <= src_start || src_end <= dst_start
231 },
232 "overlapping src/dest is not supported"
233 );
234 unsafe { std::ptr::copy_nonoverlapping(src_ptr, dst_ptr, len) };
235 } else {
236 let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
237 let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
238 simd::dispatch_if_large(len, || {
239 for i in 0..len {
240 dst[i] = Op::apply(src[i]);
241 }
242 });
243 }
244 return Ok(());
245 }
246
247 copy_map_into(dest, src, |x| x, exec)
248}
249
250pub fn copy_into_uninit<T: Copy + Send + Sync + 'static>(
277 dest: &mut StridedViewMut<MaybeUninit<T>>,
278 src: &StridedView<T>,
279 exec: Exec<'_>,
280) -> Result<()> {
281 ensure_same_shape(dest.dims(), src.dims())?;
282 if !crate::layout_check::is_injective_layout(dest.dims(), dest.strides()) {
283 return Err(StridedError::NonInjectiveOutputLayout);
284 }
285 if src.is_empty() {
286 return Ok(());
287 }
288 let _len = src
289 .dims()
290 .iter()
291 .try_fold(1usize, |n, &dim| n.checked_mul(dim))
292 .ok_or(StridedError::OffsetOverflow)?;
293 let native_float = std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>()
297 || std::any::TypeId::of::<T>() == std::any::TypeId::of::<f64>();
298 if !native_float
303 || same_contiguous_layout(dest.dims(), &[dest.strides(), src.strides()]).is_some()
304 {
305 return copy_map_into(dest, src, MaybeUninit::new, exec);
306 }
307 crate::threading::copy_permuted_uninit(dest, src, exec)
311}
312
313pub fn copy_into_col_major<T: Copy + Send + Sync + 'static>(
317 dst: &mut StridedViewMut<T>,
318 src: &StridedView<T>,
319) -> Result<()> {
320 strided_perm::copy_into_col_major(dst, src, Exec::serial())
321}
322
323pub fn add<D: Copy + Add<S, Output = D> + Send + Sync, S: Copy + Send + Sync, Op: ElementOp<S>>(
327 dest: &mut StridedViewMut<D>,
328 src: &StridedView<S, Op>,
329 exec: Exec<'_>,
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(
340 KernelClass::Map,
341 dst_dims,
342 &[dst_strides, src_strides],
343 &exec,
344 )?
345 .is_some()
346 {
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 let total = total_len(&fused_dims)?;
369 if is_parallel(KernelClass::Map, total, &exec) {
370 let dst_send = SendPtr(dst_ptr);
371 let src_send = SendPtr(src_ptr as *mut S);
372
373 let costs = compute_costs(&ordered_strides);
374 let initial_offsets: crate::inline::OperandVec<isize> =
375 strides_list.iter().map(|_| 0isize).collect();
376 return mapreduce_threaded(
377 &fused_dims,
378 &plan.block,
379 &ordered_strides,
380 &initial_offsets,
381 &costs,
382 KernelClass::Map,
383 exec,
384 &|dims, blocks, strides_list, offsets| {
385 for_each_inner_block_with_offsets(
386 dims,
387 blocks,
388 strides_list,
389 offsets,
390 |offsets, len, strides| {
391 unsafe {
392 inner_loop_add::<D, S, Op>(
393 dst_send.as_ptr().offset(offsets[0]),
394 strides[0],
395 src_send.as_const().offset(offsets[1]),
396 strides[1],
397 len,
398 )
399 };
400 Ok(())
401 },
402 )
403 },
404 );
405 }
406
407 let initial_offsets: crate::inline::OperandVec<isize> =
408 ordered_strides.iter().map(|_| 0isize).collect();
409 for_each_inner_block_preordered(
410 &fused_dims,
411 &plan.block,
412 &ordered_strides,
413 &initial_offsets,
414 |offsets, len, strides| {
415 unsafe {
416 inner_loop_add::<D, S, Op>(
417 dst_ptr.offset(offsets[0]),
418 strides[0],
419 src_ptr.offset(offsets[1]),
420 strides[1],
421 len,
422 )
423 };
424 Ok(())
425 },
426 )
427}
428
429pub fn mul<D: Copy + Mul<S, Output = D> + Send + Sync, S: Copy + Send + Sync, Op: ElementOp<S>>(
433 dest: &mut StridedViewMut<D>,
434 src: &StridedView<S, Op>,
435 exec: Exec<'_>,
436) -> Result<()> {
437 ensure_same_shape(dest.dims(), src.dims())?;
438
439 let dst_ptr = dest.as_mut_ptr();
440 let src_ptr = src.ptr();
441 let dst_dims = dest.dims();
442 let dst_strides = dest.strides();
443 let src_strides = src.strides();
444
445 if sequential_contiguous_layout(
446 KernelClass::Map,
447 dst_dims,
448 &[dst_strides, src_strides],
449 &exec,
450 )?
451 .is_some()
452 {
453 let len = total_len(dst_dims)?;
454 let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
455 let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
456 simd::dispatch_if_large(len, || {
457 for i in 0..len {
458 dst[i] = dst[i] * Op::apply(src[i]);
459 }
460 });
461 return Ok(());
462 }
463
464 let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
465 let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<S>());
466
467 let total = total_len(dst_dims)?;
468 let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
469 build_plan_fused_small(dst_dims, &strides_list)
470 } else {
471 build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
472 };
473
474 let total = total_len(&fused_dims)?;
475 if is_parallel(KernelClass::Map, total, &exec) {
476 let dst_send = SendPtr(dst_ptr);
477 let src_send = SendPtr(src_ptr as *mut S);
478
479 let costs = compute_costs(&ordered_strides);
480 let initial_offsets: crate::inline::OperandVec<isize> =
481 strides_list.iter().map(|_| 0isize).collect();
482 return mapreduce_threaded(
483 &fused_dims,
484 &plan.block,
485 &ordered_strides,
486 &initial_offsets,
487 &costs,
488 KernelClass::Map,
489 exec,
490 &|dims, blocks, strides_list, offsets| {
491 for_each_inner_block_with_offsets(
492 dims,
493 blocks,
494 strides_list,
495 offsets,
496 |offsets, len, strides| {
497 unsafe {
498 inner_loop_mul::<D, S, Op>(
499 dst_send.as_ptr().offset(offsets[0]),
500 strides[0],
501 src_send.as_const().offset(offsets[1]),
502 strides[1],
503 len,
504 )
505 };
506 Ok(())
507 },
508 )
509 },
510 );
511 }
512
513 let initial_offsets: crate::inline::OperandVec<isize> =
514 ordered_strides.iter().map(|_| 0isize).collect();
515 for_each_inner_block_preordered(
516 &fused_dims,
517 &plan.block,
518 &ordered_strides,
519 &initial_offsets,
520 |offsets, len, strides| {
521 unsafe {
522 inner_loop_mul::<D, S, Op>(
523 dst_ptr.offset(offsets[0]),
524 strides[0],
525 src_ptr.offset(offsets[1]),
526 strides[1],
527 len,
528 )
529 };
530 Ok(())
531 },
532 )
533}
534
535pub fn axpy<D, S, A, Op>(
539 dest: &mut StridedViewMut<D>,
540 src: &StridedView<S, Op>,
541 alpha: A,
542 exec: Exec<'_>,
543) -> Result<()>
544where
545 A: Copy + Mul<S, Output = D> + Sync,
546 D: Copy + Add<D, Output = D> + Send + Sync,
547 S: Copy + Send + Sync,
548 Op: ElementOp<S>,
549{
550 ensure_same_shape(dest.dims(), src.dims())?;
551
552 let dst_ptr = dest.as_mut_ptr();
553 let src_ptr = src.ptr();
554 let dst_dims = dest.dims();
555 let dst_strides = dest.strides();
556 let src_strides = src.strides();
557
558 if sequential_contiguous_layout(
559 KernelClass::Map,
560 dst_dims,
561 &[dst_strides, src_strides],
562 &exec,
563 )?
564 .is_some()
565 {
566 let len = total_len(dst_dims)?;
567 let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
568 let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
569 simd::dispatch_if_large(len, || {
570 for i in 0..len {
571 dst[i] = alpha * Op::apply(src[i]) + dst[i];
572 }
573 });
574 return Ok(());
575 }
576
577 let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
578 let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<S>());
579
580 let total = total_len(dst_dims)?;
581 let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
582 build_plan_fused_small(dst_dims, &strides_list)
583 } else {
584 build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
585 };
586
587 let total = total_len(&fused_dims)?;
588 if is_parallel(KernelClass::Map, total, &exec) {
589 let dst_send = SendPtr(dst_ptr);
590 let src_send = SendPtr(src_ptr as *mut S);
591
592 let costs = compute_costs(&ordered_strides);
593 let initial_offsets: crate::inline::OperandVec<isize> =
594 strides_list.iter().map(|_| 0isize).collect();
595 return mapreduce_threaded(
596 &fused_dims,
597 &plan.block,
598 &ordered_strides,
599 &initial_offsets,
600 &costs,
601 KernelClass::Map,
602 exec,
603 &|dims, blocks, strides_list, offsets| {
604 for_each_inner_block_with_offsets(
605 dims,
606 blocks,
607 strides_list,
608 offsets,
609 |offsets, len, strides| {
610 unsafe {
611 inner_loop_axpy::<D, S, A, Op>(
612 dst_send.as_ptr().offset(offsets[0]),
613 strides[0],
614 src_send.as_const().offset(offsets[1]),
615 strides[1],
616 len,
617 alpha,
618 )
619 };
620 Ok(())
621 },
622 )
623 },
624 );
625 }
626
627 let initial_offsets: crate::inline::OperandVec<isize> =
628 ordered_strides.iter().map(|_| 0isize).collect();
629 for_each_inner_block_preordered(
630 &fused_dims,
631 &plan.block,
632 &ordered_strides,
633 &initial_offsets,
634 |offsets, len, strides| {
635 unsafe {
636 inner_loop_axpy::<D, S, A, Op>(
637 dst_ptr.offset(offsets[0]),
638 strides[0],
639 src_ptr.offset(offsets[1]),
640 strides[1],
641 len,
642 alpha,
643 )
644 };
645 Ok(())
646 },
647 )
648}
649
650pub fn fma<D, A, B, OpA, OpB>(
654 dest: &mut StridedViewMut<D>,
655 a: &StridedView<A, OpA>,
656 b: &StridedView<B, OpB>,
657 exec: Exec<'_>,
658) -> Result<()>
659where
660 A: Copy + Mul<B, Output = D> + Send + Sync,
661 B: Copy + Send + Sync,
662 D: Copy + Add<D, Output = D> + Send + Sync,
663 OpA: ElementOp<A>,
664 OpB: ElementOp<B>,
665{
666 ensure_same_shape(dest.dims(), a.dims())?;
667 ensure_same_shape(dest.dims(), b.dims())?;
668
669 let dst_ptr = dest.as_mut_ptr();
670 let a_ptr = a.ptr();
671 let b_ptr = b.ptr();
672 let dst_dims = dest.dims();
673 let dst_strides = dest.strides();
674 let a_strides = a.strides();
675 let b_strides = b.strides();
676
677 if sequential_contiguous_layout(
678 KernelClass::Map,
679 dst_dims,
680 &[dst_strides, a_strides, b_strides],
681 &exec,
682 )?
683 .is_some()
684 {
685 let len = total_len(dst_dims)?;
686 let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
687 let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
688 let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
689 simd::dispatch_if_large(len, || {
690 for i in 0..len {
691 dst[i] = dst[i] + OpA::apply(sa[i]) * OpB::apply(sb[i]);
692 }
693 });
694 return Ok(());
695 }
696
697 let strides_list: [&[isize]; 3] = [dst_strides, a_strides, b_strides];
698 let elem_size = std::mem::size_of::<D>()
699 .max(std::mem::size_of::<A>())
700 .max(std::mem::size_of::<B>());
701
702 let total = total_len(dst_dims)?;
703 let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
704 build_plan_fused_small(dst_dims, &strides_list)
705 } else {
706 build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
707 };
708
709 let total = total_len(&fused_dims)?;
710 if is_parallel(KernelClass::Map, total, &exec) {
711 let dst_send = SendPtr(dst_ptr);
712 let a_send = SendPtr(a_ptr as *mut A);
713 let b_send = SendPtr(b_ptr as *mut B);
714
715 let costs = compute_costs(&ordered_strides);
716 let initial_offsets: crate::inline::OperandVec<isize> =
717 strides_list.iter().map(|_| 0isize).collect();
718 return mapreduce_threaded(
719 &fused_dims,
720 &plan.block,
721 &ordered_strides,
722 &initial_offsets,
723 &costs,
724 KernelClass::Map,
725 exec,
726 &|dims, blocks, strides_list, offsets| {
727 for_each_inner_block_with_offsets(
728 dims,
729 blocks,
730 strides_list,
731 offsets,
732 |offsets, len, strides| {
733 unsafe {
734 inner_loop_fma::<D, A, B, OpA, OpB>(
735 dst_send.as_ptr().offset(offsets[0]),
736 strides[0],
737 a_send.as_const().offset(offsets[1]),
738 strides[1],
739 b_send.as_const().offset(offsets[2]),
740 strides[2],
741 len,
742 )
743 };
744 Ok(())
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
775pub fn sum<
798 T: Copy + Zero + Add<Output = T> + Send + Sync + simd::MaybeSimdOps,
799 Op: ElementOp<T>,
800>(
801 src: &StridedView<T, Op>,
802 exec: Exec<'_>,
803) -> Result<T> {
804 if total_len(src.dims())? == 0 {
807 return Ok(T::zero());
808 }
809 if Op::IS_IDENTITY {
811 if same_contiguous_layout(src.dims(), &[src.strides()]).is_some() {
812 let len = total_len(src.dims())?;
813 let src_slice = unsafe { std::slice::from_raw_parts(src.ptr(), len) };
814
815 if T::try_simd_sum(&[]).is_some() && is_parallel(KernelClass::Reduce, len, &exec) {
818 return Ok(crate::threading::parallel_map_reduce(
819 0..len,
820 crate::cost::for_len_tree(exec, KernelClass::Reduce, len),
821 &|range| T::try_simd_sum(&src_slice[range]).unwrap(),
822 &|left, right| left + right,
823 ));
824 }
825
826 if let Some(result) = T::try_simd_sum(src_slice) {
827 return Ok(result);
828 }
829 }
830 }
831 reduce(src, |x| x, |a, b| a + b, T::zero(), exec)
832}
833
834pub fn dot<A, B, R, OpA, OpB>(a: &StridedView<A, OpA>, b: &StridedView<B, OpB>) -> Result<R>
858where
859 A: Copy + Mul<B, Output = R> + Send + Sync + 'static,
860 B: Copy + Send + Sync + 'static,
861 R: Copy + Zero + Add<Output = R> + Send + Sync + simd::MaybeSimdOps + 'static,
862 OpA: ElementOp<A>,
863 OpB: ElementOp<B>,
864{
865 ensure_same_shape(a.dims(), b.dims())?;
866
867 let a_ptr = a.ptr();
868 let b_ptr = b.ptr();
869 let a_strides = a.strides();
870 let b_strides = b.strides();
871 let a_dims = a.dims();
872 let len = total_len(a_dims)?;
873
874 if len == 0 {
878 return Ok(R::zero());
879 }
880
881 if same_contiguous_layout(a_dims, &[a_strides, b_strides]).is_some() {
882 if OpA::IS_IDENTITY
884 && OpB::IS_IDENTITY
885 && std::any::TypeId::of::<A>() == std::any::TypeId::of::<R>()
886 && std::any::TypeId::of::<B>() == std::any::TypeId::of::<R>()
887 {
888 let sa = unsafe { std::slice::from_raw_parts(a_ptr as *const R, len) };
889 let sb = unsafe { std::slice::from_raw_parts(b_ptr as *const R, len) };
890 if let Some(result) = R::try_simd_dot(sa, sb) {
891 return Ok(result);
892 }
893 }
894
895 let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
897 let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
898 let mut acc = R::zero();
899 simd::dispatch_if_large(len, || {
900 for i in 0..len {
901 acc = acc + OpA::apply(sa[i]) * OpB::apply(sb[i]);
902 }
903 });
904 return Ok(acc);
905 }
906
907 let strides_list: [&[isize]; 2] = [a_strides, b_strides];
908 let elem_size = std::mem::size_of::<A>()
909 .max(std::mem::size_of::<B>())
910 .max(std::mem::size_of::<R>());
911
912 let total = total_len(a_dims)?;
913 let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
914 build_plan_fused_small(a_dims, &strides_list)
915 } else {
916 build_plan_fused(a_dims, &strides_list, None, elem_size)
917 };
918
919 let mut acc = R::zero();
920 let initial_offsets: crate::inline::OperandVec<isize> =
921 ordered_strides.iter().map(|_| 0isize).collect();
922 for_each_inner_block_preordered(
923 &fused_dims,
924 &plan.block,
925 &ordered_strides,
926 &initial_offsets,
927 |offsets, len, strides| {
928 acc = unsafe {
929 inner_loop_dot::<A, B, R, OpA, OpB>(
930 a_ptr.offset(offsets[0]),
931 strides[0],
932 b_ptr.offset(offsets[1]),
933 strides[1],
934 len,
935 acc,
936 )
937 };
938 Ok(())
939 },
940 )?;
941
942 Ok(acc)
943}
944
945pub fn symmetrize_into<T>(
947 dest: &mut StridedViewMut<T>,
948 src: &StridedView<T>,
949 exec: Exec<'_>,
950) -> Result<()>
951where
952 T: Copy
953 + Add<Output = T>
954 + Mul<Output = T>
955 + num_traits::FromPrimitive
956 + std::ops::Div<Output = T>
957 + Send
958 + Sync,
959{
960 if src.ndim() != 2 {
961 return Err(StridedError::RankMismatch(src.ndim(), 2));
962 }
963 let rows = src.dims()[0];
964 let cols = src.dims()[1];
965 if rows != cols {
966 return Err(StridedError::NonSquare { rows, cols });
967 }
968
969 let src_t = src.permute(&[1, 0])?;
970 let half = T::from_f64(0.5).ok_or(StridedError::ScalarConversion)?;
971
972 zip_map2_into(dest, src, &src_t, |a, b| (a + b) * half, exec)
973}
974
975pub fn symmetrize_conj_into<T>(
977 dest: &mut StridedViewMut<T>,
978 src: &StridedView<T>,
979 exec: Exec<'_>,
980) -> Result<()>
981where
982 T: Copy
983 + ElementOpApply
984 + Add<Output = T>
985 + Mul<Output = T>
986 + num_traits::FromPrimitive
987 + std::ops::Div<Output = T>
988 + Send
989 + Sync,
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, exec)
1005}
1006
1007pub fn copy_scale<D, S, A, Op>(
1011 dest: &mut StridedViewMut<D>,
1012 src: &StridedView<S, Op>,
1013 scale: A,
1014 exec: Exec<'_>,
1015) -> Result<()>
1016where
1017 A: Copy + Mul<S, Output = D> + Sync,
1018 D: Copy + Send + Sync,
1019 S: Copy + Send + Sync,
1020 Op: ElementOp<S>,
1021{
1022 map_into(dest, src, |x| scale * x, exec)
1023}
1024
1025pub fn copy_conj<T: Copy + ElementOpApply + Send + Sync>(
1027 dest: &mut StridedViewMut<T>,
1028 src: &StridedView<T>,
1029 exec: Exec<'_>,
1030) -> Result<()> {
1031 let src_conj = src.conj();
1032 copy_into(dest, &src_conj, exec)
1033}
1034
1035#[inline]
1036fn element_transpose_is_identity<T: 'static>() -> bool {
1037 use std::any::TypeId;
1038
1039 macro_rules! matches_type {
1040 ($($ty:ty),* $(,)?) => {{
1041 let id = TypeId::of::<T>();
1042 false $(|| id == TypeId::of::<$ty>())*
1043 }};
1044 }
1045
1046 matches_type!(
1047 f32,
1048 f64,
1049 i8,
1050 i16,
1051 i32,
1052 i64,
1053 i128,
1054 isize,
1055 u8,
1056 u16,
1057 u32,
1058 u64,
1059 u128,
1060 usize,
1061 num_complex::Complex32,
1062 num_complex::Complex64,
1063 )
1064}
1065
1066#[inline]
1067fn element_zero_is_all_bits_zero<T: 'static>() -> bool {
1068 use std::any::TypeId;
1069
1070 macro_rules! matches_type {
1071 ($($ty:ty),* $(,)?) => {{
1072 let id = TypeId::of::<T>();
1073 false $(|| id == TypeId::of::<$ty>())*
1074 }};
1075 }
1076
1077 matches_type!(
1078 f32,
1079 f64,
1080 i8,
1081 i16,
1082 i32,
1083 i64,
1084 i128,
1085 isize,
1086 u8,
1087 u16,
1088 u32,
1089 u64,
1090 u128,
1091 usize,
1092 num_complex::Complex32,
1093 num_complex::Complex64,
1094 )
1095}
1096
1097#[inline]
1098unsafe fn fill_2d<T: Copy + Send + Sync>(
1099 dst: *mut T,
1100 dim0: usize,
1101 dim1: usize,
1102 dst_stride0: isize,
1103 dst_stride1: isize,
1104 value: T,
1105 exec: Exec<'_>,
1106) {
1107 let total = dim0.saturating_mul(dim1);
1108 let exec = crate::cost::for_len(exec, KernelClass::Copy, total);
1109 if !exec.is_serial() {
1110 let dst_send = SendPtr(dst);
1111 if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1112 crate::threading::parallel_for_each(0..dim1, exec, &|columns| {
1113 for j in columns {
1114 let dst = dst_send.as_ptr();
1115 unsafe {
1116 let base = j as isize * dst_stride1;
1117 for i in 0..dim0 {
1118 *dst.offset(base + i as isize * dst_stride0) = value;
1119 }
1120 }
1121 }
1122 });
1123 } else {
1124 crate::threading::parallel_for_each(0..dim0, exec, &|rows| {
1125 for i in rows {
1126 let dst = dst_send.as_ptr();
1127 unsafe {
1128 let base = i as isize * dst_stride0;
1129 for j in 0..dim1 {
1130 *dst.offset(base + j as isize * dst_stride1) = value;
1131 }
1132 }
1133 }
1134 });
1135 }
1136 return;
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 + Send + Sync + '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]
1233#[allow(clippy::too_many_arguments)]
1234unsafe fn copy_transpose_scale_2d_f64_tiled_raw(
1235 dst: *mut f64,
1236 dst_stride0: isize,
1237 dst_stride1: isize,
1238 src: *const f64,
1239 src_stride0: isize,
1240 src_stride1: isize,
1241 src_rows: usize,
1242 src_cols: usize,
1243 scale: f64,
1244 exec: Exec<'_>,
1245) {
1246 const TILE: usize = 4;
1247 let row_full = src_rows / TILE * TILE;
1248 let col_full = src_cols / TILE * TILE;
1249
1250 let total = src_rows.saturating_mul(src_cols);
1251 let exec = crate::cost::for_len(exec, KernelClass::Map, total);
1252 if !exec.is_serial() {
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, exec, &|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 let mut i = 0;
1299 while i < row_full {
1300 let mut j = 0;
1301 while j < col_full {
1302 transpose_scale_4x4_f64(
1303 dst,
1304 dst_stride0,
1305 dst_stride1,
1306 src,
1307 src_stride0,
1308 src_stride1,
1309 i,
1310 j,
1311 scale,
1312 );
1313 j += TILE;
1314 }
1315 for j in col_full..src_cols {
1316 for ii in i..i + TILE {
1317 *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1318 scale * *src.offset(ii as isize * src_stride0 + j as isize * src_stride1);
1319 }
1320 }
1321 i += TILE;
1322 }
1323 for i in row_full..src_rows {
1324 for j in 0..src_cols {
1325 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1326 scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1327 }
1328 }
1329}
1330
1331#[inline]
1332#[cfg(test)]
1333unsafe fn try_copy_transpose_scale_2d_f64_tiled(
1334 dest: &mut StridedViewMut<f64>,
1335 src: &StridedView<f64>,
1336 scale: f64,
1337 exec: Exec<'_>,
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 exec,
1361 );
1362 true
1363}
1364
1365#[inline]
1366unsafe fn try_copy_transpose_scale_2d_f64_tiled_typed<T>(
1367 dest: &mut StridedViewMut<T>,
1368 src: &StridedView<T>,
1369 scale: T,
1370 exec: Exec<'_>,
1371) -> bool
1372where
1373 T: Copy + 'static,
1374{
1375 if std::any::TypeId::of::<T>() != std::any::TypeId::of::<f64>() {
1376 return false;
1377 }
1378
1379 let scale = *(&scale as *const T).cast::<f64>();
1380 if src.ndim() != 2 || dest.ndim() != 2 {
1381 return false;
1382 }
1383 let src_dims = src.dims();
1384 if dest.dims() != [src_dims[1], src_dims[0]] {
1385 return false;
1386 }
1387 if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1388 return false;
1389 }
1390
1391 copy_transpose_scale_2d_f64_tiled_raw(
1392 dest.as_mut_ptr().cast::<f64>(),
1393 dest.strides()[0],
1394 dest.strides()[1],
1395 src.ptr().cast::<f64>(),
1396 src.strides()[0],
1397 src.strides()[1],
1398 src_dims[0],
1399 src_dims[1],
1400 scale,
1401 exec,
1402 );
1403 true
1404}
1405
1406#[inline(always)]
1407unsafe fn transpose_scale_4x4_identity<T>(
1408 dst: *mut T,
1409 dst_stride0: isize,
1410 dst_stride1: isize,
1411 src: *const T,
1412 src_stride0: isize,
1413 src_stride1: isize,
1414 i: usize,
1415 j: usize,
1416 scale: T,
1417) where
1418 T: Copy + Mul<Output = T>,
1419{
1420 let src_base = src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1421
1422 let s00 = *src_base;
1423 let s10 = *src_base.offset(src_stride0);
1424 let s20 = *src_base.offset(2 * src_stride0);
1425 let s30 = *src_base.offset(3 * src_stride0);
1426
1427 let src_col1 = src_base.offset(src_stride1);
1428 let s01 = *src_col1;
1429 let s11 = *src_col1.offset(src_stride0);
1430 let s21 = *src_col1.offset(2 * src_stride0);
1431 let s31 = *src_col1.offset(3 * src_stride0);
1432
1433 let src_col2 = src_base.offset(2 * src_stride1);
1434 let s02 = *src_col2;
1435 let s12 = *src_col2.offset(src_stride0);
1436 let s22 = *src_col2.offset(2 * src_stride0);
1437 let s32 = *src_col2.offset(3 * src_stride0);
1438
1439 let src_col3 = src_base.offset(3 * src_stride1);
1440 let s03 = *src_col3;
1441 let s13 = *src_col3.offset(src_stride0);
1442 let s23 = *src_col3.offset(2 * src_stride0);
1443 let s33 = *src_col3.offset(3 * src_stride0);
1444
1445 let dst_row0 = dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1);
1446 *dst_row0 = scale * s00;
1447 *dst_row0.offset(dst_stride0) = scale * s01;
1448 *dst_row0.offset(2 * dst_stride0) = scale * s02;
1449 *dst_row0.offset(3 * dst_stride0) = scale * s03;
1450
1451 let dst_row1 = dst_row0.offset(dst_stride1);
1452 *dst_row1 = scale * s10;
1453 *dst_row1.offset(dst_stride0) = scale * s11;
1454 *dst_row1.offset(2 * dst_stride0) = scale * s12;
1455 *dst_row1.offset(3 * dst_stride0) = scale * s13;
1456
1457 let dst_row2 = dst_row0.offset(2 * dst_stride1);
1458 *dst_row2 = scale * s20;
1459 *dst_row2.offset(dst_stride0) = scale * s21;
1460 *dst_row2.offset(2 * dst_stride0) = scale * s22;
1461 *dst_row2.offset(3 * dst_stride0) = scale * s23;
1462
1463 let dst_row3 = dst_row0.offset(3 * dst_stride1);
1464 *dst_row3 = scale * s30;
1465 *dst_row3.offset(dst_stride0) = scale * s31;
1466 *dst_row3.offset(2 * dst_stride0) = scale * s32;
1467 *dst_row3.offset(3 * dst_stride0) = scale * s33;
1468}
1469
1470#[inline]
1471#[allow(clippy::too_many_arguments)]
1472unsafe fn copy_transpose_scale_2d_identity_tiled_raw<T>(
1473 dst: *mut T,
1474 dst_stride0: isize,
1475 dst_stride1: isize,
1476 src: *const T,
1477 src_stride0: isize,
1478 src_stride1: isize,
1479 src_rows: usize,
1480 src_cols: usize,
1481 scale: T,
1482 exec: Exec<'_>,
1483) where
1484 T: Copy + Mul<Output = T> + Send + Sync,
1485{
1486 const TILE: usize = 4;
1487 let row_full = src_rows / TILE * TILE;
1488 let col_full = src_cols / TILE * TILE;
1489
1490 let total = src_rows.saturating_mul(src_cols);
1491 let exec = crate::cost::for_len(exec, KernelClass::Map, total);
1492 if !exec.is_serial() {
1493 let dst_send = SendPtr(dst);
1494 let src_send = SendPtr(src as *mut T);
1495 let row_tiles = row_full / TILE;
1496 crate::threading::parallel_for_each(0..row_tiles, exec, &|tiles| {
1497 for tile_i in tiles {
1498 let i = tile_i * TILE;
1499 let dst = dst_send.as_ptr();
1500 let src = src_send.as_const();
1501 unsafe {
1502 let mut j = 0;
1503 while j < col_full {
1504 transpose_scale_4x4_identity(
1505 dst,
1506 dst_stride0,
1507 dst_stride1,
1508 src,
1509 src_stride0,
1510 src_stride1,
1511 i,
1512 j,
1513 scale,
1514 );
1515 j += TILE;
1516 }
1517 for j in col_full..src_cols {
1518 for ii in i..i + TILE {
1519 *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1520 scale
1521 * *src.offset(
1522 ii as isize * src_stride0 + j as isize * src_stride1,
1523 );
1524 }
1525 }
1526 }
1527 }
1528 });
1529 for i in row_full..src_rows {
1530 for j in 0..src_cols {
1531 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1532 scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1533 }
1534 }
1535 return;
1536 }
1537
1538 let mut i = 0;
1539 while i < row_full {
1540 let mut j = 0;
1541 while j < col_full {
1542 transpose_scale_4x4_identity(
1543 dst,
1544 dst_stride0,
1545 dst_stride1,
1546 src,
1547 src_stride0,
1548 src_stride1,
1549 i,
1550 j,
1551 scale,
1552 );
1553 j += TILE;
1554 }
1555 for j in col_full..src_cols {
1556 for ii in i..i + TILE {
1557 *dst.offset(j as isize * dst_stride0 + ii as isize * dst_stride1) =
1558 scale * *src.offset(ii as isize * src_stride0 + j as isize * src_stride1);
1559 }
1560 }
1561 i += TILE;
1562 }
1563 for i in row_full..src_rows {
1564 for j in 0..src_cols {
1565 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1566 scale * *src.offset(i as isize * src_stride0 + j as isize * src_stride1);
1567 }
1568 }
1569}
1570
1571#[inline]
1572unsafe fn try_copy_transpose_scale_2d_identity_tiled<T>(
1573 dest: &mut StridedViewMut<T>,
1574 src: &StridedView<T>,
1575 scale: T,
1576 exec: Exec<'_>,
1577) -> bool
1578where
1579 T: Copy + Mul<Output = T> + Send + Sync,
1580{
1581 if src.ndim() != 2 || dest.ndim() != 2 {
1582 return false;
1583 }
1584 let src_dims = src.dims();
1585 if dest.dims() != [src_dims[1], src_dims[0]] {
1586 return false;
1587 }
1588 if src.strides()[0] != 1 || dest.strides()[0] != 1 {
1589 return false;
1590 }
1591
1592 copy_transpose_scale_2d_identity_tiled_raw(
1593 dest.as_mut_ptr(),
1594 dest.strides()[0],
1595 dest.strides()[1],
1596 src.ptr(),
1597 src.strides()[0],
1598 src.strides()[1],
1599 src_dims[0],
1600 src_dims[1],
1601 scale,
1602 exec,
1603 );
1604 true
1605}
1606
1607#[inline]
1608#[allow(clippy::too_many_arguments)]
1609unsafe fn copy_transpose_scale_2d_loop<T>(
1610 dst: *mut T,
1611 dst_stride0: isize,
1612 dst_stride1: isize,
1613 src: *const T,
1614 src_stride0: isize,
1615 src_stride1: isize,
1616 src_rows: usize,
1617 src_cols: usize,
1618 scale: T,
1619 exec: Exec<'_>,
1620) where
1621 T: Copy + ElementOpApply + Mul<Output = T> + Send + Sync,
1622{
1623 let total = src_rows.saturating_mul(src_cols);
1624 let exec = crate::cost::for_len(exec, KernelClass::Map, total);
1625 if !exec.is_serial() {
1626 let dst_send = SendPtr(dst);
1627 let src_send = SendPtr(src as *mut T);
1628 if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1629 crate::threading::parallel_for_each(0..src_rows, exec, &|rows| {
1630 for i in rows {
1631 let dst = dst_send.as_ptr();
1632 let src = src_send.as_const();
1633 unsafe {
1634 for j in 0..src_cols {
1635 let value = (*src
1636 .offset(i as isize * src_stride0 + j as isize * src_stride1))
1637 .transpose();
1638 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1639 scale * value;
1640 }
1641 }
1642 }
1643 });
1644 } else {
1645 crate::threading::parallel_for_each(0..src_cols, exec, &|columns| {
1646 for j in columns {
1647 let dst = dst_send.as_ptr();
1648 let src = src_send.as_const();
1649 unsafe {
1650 for i in 0..src_rows {
1651 let value = (*src
1652 .offset(i as isize * src_stride0 + j as isize * src_stride1))
1653 .transpose();
1654 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) =
1655 scale * value;
1656 }
1657 }
1658 }
1659 });
1660 }
1661 return;
1662 }
1663
1664 if dst_stride0.unsigned_abs() <= dst_stride1.unsigned_abs() {
1665 for i in 0..src_rows {
1666 for j in 0..src_cols {
1667 let value =
1668 (*src.offset(i as isize * src_stride0 + j as isize * src_stride1)).transpose();
1669 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) = scale * value;
1670 }
1671 }
1672 } else {
1673 for j in 0..src_cols {
1674 for i in 0..src_rows {
1675 let value =
1676 (*src.offset(i as isize * src_stride0 + j as isize * src_stride1)).transpose();
1677 *dst.offset(j as isize * dst_stride0 + i as isize * dst_stride1) = scale * value;
1678 }
1679 }
1680 }
1681}
1682
1683pub fn copy_transpose_scale_into<T>(
1685 dest: &mut StridedViewMut<T>,
1686 src: &StridedView<T>,
1687 scale: T,
1688 exec: Exec<'_>,
1689) -> Result<()>
1690where
1691 T: Copy + ElementOpApply + Mul<Output = T> + Zero + One + PartialEq + Send + Sync + 'static,
1692{
1693 let mut exec = exec;
1694 if src.ndim() != 2 || dest.ndim() != 2 {
1695 return Err(StridedError::RankMismatch(src.ndim(), 2));
1696 }
1697
1698 let src_dims = src.dims();
1699 let expected_dims = [src_dims[1], src_dims[0]];
1700 ensure_same_shape(dest.dims(), &expected_dims)?;
1701 total_len(src_dims)?;
1704
1705 if scale == T::zero() {
1706 unsafe {
1707 if same_contiguous_layout(dest.dims(), &[dest.strides()]).is_some() {
1708 fill_contiguous(dest.as_mut_ptr(), total_len(dest.dims())?, T::zero());
1709 } else {
1710 fill_2d(
1711 dest.as_mut_ptr(),
1712 dest.dims()[0],
1713 dest.dims()[1],
1714 dest.strides()[0],
1715 dest.strides()[1],
1716 T::zero(),
1717 exec,
1718 );
1719 }
1720 }
1721 return Ok(());
1722 }
1723
1724 let transpose_is_identity = element_transpose_is_identity::<T>();
1725
1726 unsafe {
1727 if transpose_is_identity
1728 && try_copy_transpose_scale_2d_f64_tiled_typed(dest, src, scale, exec.reborrow())
1729 {
1730 return Ok(());
1731 }
1732 if transpose_is_identity
1733 && try_copy_transpose_scale_2d_identity_tiled(dest, src, scale, exec.reborrow())
1734 {
1735 return Ok(());
1736 }
1737 }
1738
1739 if scale == T::one() && transpose_is_identity {
1740 let src_t = src.permute(&[1, 0])?;
1741 return crate::threading::copy_permuted(dest, &src_t, exec);
1742 }
1743
1744 unsafe {
1745 copy_transpose_scale_2d_loop(
1746 dest.as_mut_ptr(),
1747 dest.strides()[0],
1748 dest.strides()[1],
1749 src.ptr(),
1750 src.strides()[0],
1751 src.strides()[1],
1752 src_dims[0],
1753 src_dims[1],
1754 scale,
1755 exec,
1756 );
1757 }
1758 Ok(())
1759}
1760
1761#[cfg(test)]
1762#[path = "ops_view/tests/tiled_tests.rs"]
1763mod tiled_tests;