1use crate::kernel::{
10 build_plan_fused, build_plan_fused_small, ensure_same_shape, for_each_inner_block_preordered,
11 sequential_contiguous_layout, total_len, SMALL_TENSOR_THRESHOLD,
12};
13use crate::maybe_sync::{MaybeSendSync, MaybeSync};
14use crate::simd;
15use crate::view::{StridedView, StridedViewMut};
16use crate::{Result, StridedError};
17use core::mem::MaybeUninit;
18use std::ops::Mul;
19use strided_view::{ElementOp, Identity};
20
21#[cfg(feature = "parallel")]
22use crate::fuse::compute_costs;
23#[cfg(feature = "parallel")]
24use crate::threading::{for_each_inner_block_with_offsets, mapreduce_threaded, MINTHREADLENGTH};
25#[cfg(feature = "parallel")]
26use smallvec::SmallVec;
27
28#[cfg(feature = "parallel")]
29type AxisVec<T> = SmallVec<[T; 8]>;
30#[cfg(not(feature = "parallel"))]
31type AxisVec<T> = Vec<T>;
32
33const CONTIGUOUS_RANGE_MIN_LEN: usize = 1 << 15;
34
35const TRANSPOSED_SCALAR_TILE_MAX_INNER_ROW_ASPECT: usize = 8;
44
45#[derive(Clone, Copy, Debug)]
51pub struct ValidatedDestinationLayout(());
52
53#[inline]
54fn validate_destination_layout(
55 dims: &[usize],
56 strides: &[isize],
57) -> Result<ValidatedDestinationLayout> {
58 if crate::layout_check::is_injective_layout(dims, strides) {
59 Ok(ValidatedDestinationLayout(()))
60 } else {
61 Err(StridedError::NonInjectiveOutputLayout)
62 }
63}
64
65#[inline]
66pub fn validate_destination_layout_without_alloc(
80 dims: &[usize],
81 strides: &[isize],
82) -> Result<ValidatedDestinationLayout> {
83 if crate::layout_check::is_injective_layout_without_alloc(dims, strides) {
84 Ok(ValidatedDestinationLayout(()))
85 } else {
86 Err(StridedError::NonInjectiveOutputLayout)
87 }
88}
89
90fn reachable_byte_range(
91 ptr: usize,
92 elem_size: usize,
93 dims: &[usize],
94 strides: &[isize],
95) -> Result<Option<(usize, usize)>> {
96 if elem_size == 0 || dims.contains(&0) {
97 return Ok(None);
98 }
99
100 let mut min_offset = 0isize;
101 let mut max_offset = 0isize;
102 for (&dim, &stride) in dims.iter().zip(strides) {
103 if dim <= 1 {
104 continue;
105 }
106 let extent = isize::try_from(dim - 1).map_err(|_| StridedError::OffsetOverflow)?;
107 let span = stride
108 .checked_mul(extent)
109 .ok_or(StridedError::OffsetOverflow)?;
110 if span < 0 {
111 min_offset = min_offset
112 .checked_add(span)
113 .ok_or(StridedError::OffsetOverflow)?;
114 } else {
115 max_offset = max_offset
116 .checked_add(span)
117 .ok_or(StridedError::OffsetOverflow)?;
118 }
119 }
120
121 let ptr = ptr as i128;
122 let elem_size = elem_size as i128;
123 let start = ptr
124 .checked_add(
125 (min_offset as i128)
126 .checked_mul(elem_size)
127 .ok_or(StridedError::OffsetOverflow)?,
128 )
129 .ok_or(StridedError::OffsetOverflow)?;
130 let end = ptr
131 .checked_add(
132 (max_offset as i128)
133 .checked_mul(elem_size)
134 .and_then(|offset| offset.checked_add(elem_size))
135 .ok_or(StridedError::OffsetOverflow)?,
136 )
137 .ok_or(StridedError::OffsetOverflow)?;
138 if start < 0 || end < 0 || start > usize::MAX as i128 || end > usize::MAX as i128 {
139 return Err(StridedError::OffsetOverflow);
140 }
141 Ok(Some((start as usize, end as usize)))
142}
143
144fn validate_typed_no_overlap<D, A, Op: ElementOp<A>>(
145 dest: &StridedViewMut<MaybeUninit<D>>,
146 input: &StridedView<A, Op>,
147 input_index: usize,
148) -> Result<()> {
149 let dest_range = reachable_byte_range(
150 dest.ptr() as usize,
151 core::mem::size_of::<D>(),
152 dest.dims(),
153 dest.strides(),
154 )?;
155 let input_range = reachable_byte_range(
156 input.ptr() as usize,
157 core::mem::size_of::<A>(),
158 input.dims(),
159 input.strides(),
160 )?;
161 if let (Some((dest_start, dest_end)), Some((input_start, input_end))) =
162 (dest_range, input_range)
163 {
164 if dest_start < input_end && input_start < dest_end {
165 return Err(StridedError::OverlappingInputOutput { input: input_index });
166 }
167 }
168 Ok(())
169}
170
171trait MapKernel<A, D> {
189 fn apply(&self, src: &[A], dst: &mut [D]) -> bool;
191}
192
193struct NoKernel;
194
195impl<A, D> MapKernel<A, D> for NoKernel {
196 #[inline(always)]
197 fn apply(&self, _src: &[A], _dst: &mut [D]) -> bool {
198 false
199 }
200}
201
202struct SliceKernel<K>(K);
203
204impl<A, D, K: Fn(&[A], &mut [D]) -> bool> MapKernel<A, D> for SliceKernel<K> {
205 #[inline(always)]
206 fn apply(&self, src: &[A], dst: &mut [D]) -> bool {
207 (self.0)(src, dst)
208 }
209}
210
211unsafe fn inner_loop_map1<D: Copy, A: Copy, Op: ElementOp<A>, K: MapKernel<A, D>>(
212 dp: *mut D,
213 ds: isize,
214 sp: *const A,
215 ss: isize,
216 len: usize,
217 f: &impl Fn(A) -> D,
218 kernel: &K,
219) {
220 if ds == 1 && ss == 1 {
221 let src = std::slice::from_raw_parts(sp, len);
222 let dst = std::slice::from_raw_parts_mut(dp, len);
223 if kernel.apply(src, dst) {
224 return;
225 }
226 simd::dispatch_if_large(len, || {
227 for (d, s) in dst.iter_mut().zip(src.iter()) {
228 *d = f(Op::apply(*s));
229 }
230 });
231 } else {
232 let mut dp = dp;
233 let mut sp = sp;
234 for _ in 0..len {
235 *dp = f(Op::apply(*sp));
236 dp = dp.wrapping_offset(ds);
237 sp = sp.wrapping_offset(ss);
238 }
239 }
240}
241
242#[inline(always)]
244unsafe fn inner_loop_map2<D: Copy, A: Copy, B: Copy, OpA: ElementOp<A>, OpB: ElementOp<B>>(
245 dp: *mut D,
246 ds: isize,
247 ap: *const A,
248 a_s: isize,
249 bp: *const B,
250 b_s: isize,
251 len: usize,
252 f: &impl Fn(A, B) -> D,
253) {
254 if ds == 1 && a_s == 1 && b_s == 1 {
255 let src_a = std::slice::from_raw_parts(ap, len);
256 let src_b = std::slice::from_raw_parts(bp, len);
257 let dst = std::slice::from_raw_parts_mut(dp, len);
258 simd::dispatch_if_large(len, || {
259 for i in 0..len {
260 dst[i] = f(OpA::apply(src_a[i]), OpB::apply(src_b[i]));
261 }
262 });
263 } else if ds == 1 && a_s == 1 && b_s == 0 {
264 let src_a = std::slice::from_raw_parts(ap, len);
265 let b = OpB::apply(*bp);
266 let dst = std::slice::from_raw_parts_mut(dp, len);
267 simd::dispatch_if_large(len, || {
268 for i in 0..len {
269 dst[i] = f(OpA::apply(src_a[i]), b);
270 }
271 });
272 } else if ds == 1 && a_s == 0 && b_s == 1 {
273 let a = OpA::apply(*ap);
274 let src_b = std::slice::from_raw_parts(bp, len);
275 let dst = std::slice::from_raw_parts_mut(dp, len);
276 simd::dispatch_if_large(len, || {
277 for i in 0..len {
278 dst[i] = f(a, OpB::apply(src_b[i]));
279 }
280 });
281 } else if ds == 1 && a_s == 0 && b_s == 0 {
282 let a = OpA::apply(*ap);
283 let b = OpB::apply(*bp);
284 let dst = std::slice::from_raw_parts_mut(dp, len);
285 simd::dispatch_if_large(len, || {
286 for d in dst.iter_mut() {
287 *d = f(a, b);
288 }
289 });
290 } else if ds == 1 && b_s == 0 {
291 let b = OpB::apply(*bp);
292 let dst = std::slice::from_raw_parts_mut(dp, len);
293 let mut ap = ap;
294 simd::dispatch_if_large(len, || {
295 for d in dst.iter_mut() {
296 *d = f(OpA::apply(*ap), b);
297 ap = ap.wrapping_offset(a_s);
298 }
299 });
300 } else if ds == 1 && a_s == 0 {
301 let a = OpA::apply(*ap);
302 let dst = std::slice::from_raw_parts_mut(dp, len);
303 let mut bp = bp;
304 simd::dispatch_if_large(len, || {
305 for d in dst.iter_mut() {
306 *d = f(a, OpB::apply(*bp));
307 bp = bp.wrapping_offset(b_s);
308 }
309 });
310 } else {
311 let mut dp = dp;
312 let mut ap = ap;
313 let mut bp = bp;
314 for _ in 0..len {
315 *dp = f(OpA::apply(*ap), OpB::apply(*bp));
316 dp = dp.wrapping_offset(ds);
317 ap = ap.wrapping_offset(a_s);
318 bp = bp.wrapping_offset(b_s);
319 }
320 }
321}
322
323trait MulOutput<D: Copy + 'static>: Copy + MaybeSendSync + 'static {
325 type Slot: Copy + MaybeSendSync + 'static;
326
327 unsafe fn write(dst: *mut Self::Slot, value: D);
328
329 unsafe fn try_contiguous<A: 'static, B: 'static>(
330 dst: *mut Self::Slot,
331 len: usize,
332 a: &[A],
333 b: &[B],
334 ) -> bool;
335}
336
337#[derive(Clone, Copy)]
338struct InitializedOutput;
339
340impl<D: Copy + MaybeSendSync + 'static> MulOutput<D> for InitializedOutput {
341 type Slot = D;
342
343 #[inline(always)]
344 unsafe fn write(dst: *mut D, value: D) {
345 unsafe { dst.write(value) };
346 }
347
348 #[inline(always)]
349 unsafe fn try_contiguous<A: 'static, B: 'static>(
350 dst: *mut D,
351 len: usize,
352 a: &[A],
353 b: &[B],
354 ) -> bool {
355 unsafe { simd::try_mul_contiguous_ptr(dst, len, a, b) }
356 }
357}
358
359#[derive(Clone, Copy)]
360struct UninitializedOutput;
361
362impl<D: Copy + MaybeSendSync + 'static> MulOutput<D> for UninitializedOutput {
363 type Slot = MaybeUninit<D>;
364
365 #[inline(always)]
366 unsafe fn write(dst: *mut MaybeUninit<D>, value: D) {
367 unsafe { dst.write(MaybeUninit::new(value)) };
368 }
369
370 #[inline(always)]
371 unsafe fn try_contiguous<A: 'static, B: 'static>(
372 dst: *mut MaybeUninit<D>,
373 len: usize,
374 a: &[A],
375 b: &[B],
376 ) -> bool {
377 unsafe { simd::try_mul_contiguous_ptr(dst.cast::<D>(), len, a, b) }
380 }
381}
382
383#[inline(always)]
384fn multiply_value<A, B, D>(lhs: A, rhs: B) -> D
385where
386 A: Copy + Mul<B, Output = D> + 'static,
387 B: Copy + 'static,
388 D: Copy + 'static,
389{
390 use core::any::TypeId;
391
392 if TypeId::of::<A>() == TypeId::of::<i32>()
393 && TypeId::of::<B>() == TypeId::of::<i32>()
394 && TypeId::of::<D>() == TypeId::of::<i32>()
395 {
396 let lhs = unsafe { *(&lhs as *const A).cast::<i32>() };
398 let rhs = unsafe { *(&rhs as *const B).cast::<i32>() };
399 let value = lhs.wrapping_mul(rhs);
400 return unsafe { core::mem::transmute_copy(&value) };
401 }
402 if TypeId::of::<A>() == TypeId::of::<i64>()
403 && TypeId::of::<B>() == TypeId::of::<i64>()
404 && TypeId::of::<D>() == TypeId::of::<i64>()
405 {
406 let lhs = unsafe { *(&lhs as *const A).cast::<i64>() };
408 let rhs = unsafe { *(&rhs as *const B).cast::<i64>() };
409 let value = lhs.wrapping_mul(rhs);
410 return unsafe { core::mem::transmute_copy(&value) };
411 }
412 lhs * rhs
413}
414
415#[inline(always)]
416unsafe fn inner_loop_mul2<
417 O: MulOutput<D>,
418 D: Copy + 'static,
419 A: Copy + Mul<B, Output = D> + 'static,
420 B: Copy + 'static,
421>(
422 dp: *mut O::Slot,
423 ds: isize,
424 ap: *const A,
425 a_s: isize,
426 bp: *const B,
427 b_s: isize,
428 len: usize,
429) {
430 if ds == 1 && a_s == 1 && b_s == 1 {
431 let src_a = std::slice::from_raw_parts(ap, len);
432 let src_b = std::slice::from_raw_parts(bp, len);
433 if len >= 64 && O::try_contiguous(dp, len, src_a, src_b) {
434 return;
435 }
436 for i in 0..len {
437 O::write(dp.add(i), multiply_value(src_a[i], src_b[i]));
438 }
439 } else if ds == 1 && a_s == 1 && b_s == 0 {
440 let src_a = std::slice::from_raw_parts(ap, len);
441 let b = *bp;
442 for i in 0..len {
443 O::write(dp.add(i), multiply_value(src_a[i], b));
444 }
445 } else if ds == 1 && a_s == 0 && b_s == 1 {
446 let a = *ap;
447 let src_b = std::slice::from_raw_parts(bp, len);
448 for i in 0..len {
449 O::write(dp.add(i), multiply_value(a, src_b[i]));
450 }
451 } else if ds == 1 && a_s == 0 && b_s == 0 {
452 let a = *ap;
453 let b = *bp;
454 for i in 0..len {
455 O::write(dp.add(i), multiply_value(a, b));
456 }
457 } else if ds == 1 && b_s == 0 {
458 let b = *bp;
459 let mut ap = ap;
460 for i in 0..len {
461 O::write(dp.add(i), multiply_value(*ap, b));
462 ap = ap.wrapping_offset(a_s);
463 }
464 } else if ds == 1 && a_s == 0 {
465 let a = *ap;
466 let mut bp = bp;
467 for i in 0..len {
468 O::write(dp.add(i), multiply_value(a, *bp));
469 bp = bp.wrapping_offset(b_s);
470 }
471 } else {
472 let mut dp = dp;
473 let mut ap = ap;
474 let mut bp = bp;
475 for _ in 0..len {
476 O::write(dp, multiply_value(*ap, *bp));
477 dp = dp.wrapping_offset(ds);
478 ap = ap.wrapping_offset(a_s);
479 bp = bp.wrapping_offset(b_s);
480 }
481 }
482}
483
484#[derive(Clone, Debug, Eq, PartialEq)]
485struct ContiguousMulRangePlan {
486 axis_order: AxisVec<usize>,
487 inner_len: usize,
488 inner_axis_count: usize,
489 row_len: usize,
490 outer_axis_start: usize,
491 fast_axis: usize,
492 a_fast_stride: isize,
493 b_fast_stride: isize,
494 a_row_stride: isize,
495 b_row_stride: isize,
496}
497
498impl ContiguousMulRangePlan {
499 fn walks_inputs_without_blocking(&self) -> bool {
509 let streams = |stride: isize| stride == 0 || stride == 1;
510 if streams(self.a_fast_stride) && streams(self.b_fast_stride) {
511 return true;
512 }
513 if transposed_scalar_tile_kind(self).is_some() {
514 return true;
515 }
516 false
517 }
518}
519
520#[derive(Clone, Copy, Debug, Eq, PartialEq)]
521enum TransposedScalarTileKind {
522 RhsScalar,
523 LhsScalar,
524}
525
526fn transposed_scalar_tile_kind(plan: &ContiguousMulRangePlan) -> Option<TransposedScalarTileKind> {
527 if plan.inner_len > TRANSPOSED_SCALAR_TILE_MAX_INNER_ROW_ASPECT.saturating_mul(plan.row_len) {
528 return None;
529 }
530 let row_len = isize::try_from(plan.row_len).ok()?;
531 if plan.a_fast_stride == row_len
532 && plan.a_row_stride == 1
533 && plan.b_fast_stride == 0
534 && plan.b_row_stride == 0
535 {
536 return Some(TransposedScalarTileKind::RhsScalar);
537 }
538
539 if plan.b_fast_stride == row_len
540 && plan.b_row_stride == 1
541 && plan.a_fast_stride == 0
542 && plan.a_row_stride == 0
543 {
544 return Some(TransposedScalarTileKind::LhsScalar);
545 }
546
547 None
548}
549
550fn compact_axis_order(dims: &[usize], strides: &[isize]) -> Option<AxisVec<usize>> {
551 if dims.len() != strides.len() {
552 return None;
553 }
554
555 let mut active = AxisVec::<usize>::new();
556 let mut inactive = AxisVec::<usize>::new();
557 for (axis, (&dim, &stride)) in dims.iter().zip(strides.iter()).enumerate() {
558 if stride < 0 {
559 return None;
560 }
561 if dim > 1 {
562 active.push(axis);
563 } else {
564 inactive.push(axis);
565 }
566 }
567
568 active.sort_by(|&lhs, &rhs| strides[lhs].cmp(&strides[rhs]).then_with(|| lhs.cmp(&rhs)));
569
570 let mut expected = 1isize;
571 for &axis in &active {
572 if strides[axis] != expected {
573 return None;
574 }
575 expected = expected.saturating_mul(dims[axis] as isize);
576 }
577
578 active.extend(inactive);
579 Some(active)
580}
581
582fn can_fuse_contiguous_range_axis(dim: usize, prev_stride: isize, next_stride: isize) -> bool {
583 dim <= 1 || (prev_stride == 0 && next_stride == 0) || next_stride == prev_stride * dim as isize
584}
585
586fn contiguous_mul_range_plan(
587 dims: &[usize],
588 dst_strides: &[isize],
589 a_strides: &[isize],
590 b_strides: &[isize],
591) -> Option<ContiguousMulRangePlan> {
592 let axis_order = compact_axis_order(dims, dst_strides)?;
593 if dims.is_empty() {
594 return Some(ContiguousMulRangePlan {
595 axis_order,
596 inner_len: 1,
597 inner_axis_count: 0,
598 row_len: 1,
599 outer_axis_start: 0,
600 fast_axis: 0,
601 a_fast_stride: 0,
602 b_fast_stride: 0,
603 a_row_stride: 0,
604 b_row_stride: 0,
605 });
606 }
607
608 let first_pos = axis_order
609 .iter()
610 .position(|&axis| dims[axis] > 1)
611 .unwrap_or(0);
612 let first_axis = axis_order[first_pos];
613 let mut inner_len = dims[first_axis].max(1);
614 let mut inner_axis_count = first_pos + 1;
615 let mut prev_axis = first_axis;
616
617 for &axis in axis_order.iter().skip(first_pos + 1) {
618 if can_fuse_contiguous_range_axis(
619 dims[prev_axis],
620 dst_strides[prev_axis],
621 dst_strides[axis],
622 ) && can_fuse_contiguous_range_axis(
623 dims[prev_axis],
624 a_strides[prev_axis],
625 a_strides[axis],
626 ) && can_fuse_contiguous_range_axis(
627 dims[prev_axis],
628 b_strides[prev_axis],
629 b_strides[axis],
630 ) {
631 inner_len = inner_len.checked_mul(dims[axis].max(1))?;
632 inner_axis_count += 1;
633 prev_axis = axis;
634 } else {
635 break;
636 }
637 }
638
639 let row = axis_order
640 .iter()
641 .enumerate()
642 .skip(inner_axis_count)
643 .find(|&(_, &axis)| dims[axis] > 1);
644 let (row_len, outer_axis_start, a_row_stride, b_row_stride) =
645 if let Some((row_pos, &row_axis)) = row {
646 (
647 dims[row_axis],
648 row_pos + 1,
649 a_strides[row_axis],
650 b_strides[row_axis],
651 )
652 } else {
653 (1, axis_order.len(), 0, 0)
654 };
655
656 Some(ContiguousMulRangePlan {
657 axis_order,
658 inner_len,
659 inner_axis_count,
660 row_len,
661 outer_axis_start,
662 fast_axis: first_axis,
663 a_fast_stride: a_strides[first_axis],
664 b_fast_stride: b_strides[first_axis],
665 a_row_stride,
666 b_row_stride,
667 })
668}
669
670struct ContiguousMulOuterCursor<'a> {
671 dims: &'a [usize],
672 a_strides: &'a [isize],
673 b_strides: &'a [isize],
674 axes: AxisVec<usize>,
675 coords: AxisVec<usize>,
676 a_offset: isize,
677 b_offset: isize,
678}
679
680impl<'a> ContiguousMulOuterCursor<'a> {
681 fn new(
682 dims: &'a [usize],
683 a_strides: &'a [isize],
684 b_strides: &'a [isize],
685 plan: &ContiguousMulRangePlan,
686 outer_group: usize,
687 ) -> Self {
688 let axes: AxisVec<usize> = plan
689 .axis_order
690 .iter()
691 .skip(plan.outer_axis_start)
692 .copied()
693 .collect();
694 let mut coords = AxisVec::<usize>::with_capacity(axes.len());
695 let mut rem = outer_group;
696 let mut a_offset = 0isize;
697 let mut b_offset = 0isize;
698
699 for &axis in &axes {
700 let dim = dims[axis].max(1);
701 let coord = rem % dim;
702 rem /= dim;
703 coords.push(coord);
704 a_offset += coord as isize * a_strides[axis];
705 b_offset += coord as isize * b_strides[axis];
706 }
707
708 Self {
709 dims,
710 a_strides,
711 b_strides,
712 axes,
713 coords,
714 a_offset,
715 b_offset,
716 }
717 }
718
719 fn advance(&mut self) {
720 for (i, &axis) in self.axes.iter().enumerate() {
721 let dim = self.dims[axis].max(1);
722 if dim <= 1 {
723 continue;
724 }
725
726 if self.coords[i] + 1 < dim {
729 self.coords[i] += 1;
730 self.a_offset += self.a_strides[axis];
731 self.b_offset += self.b_strides[axis];
732 break;
733 }
734
735 let last = (dim - 1) as isize;
736 self.coords[i] = 0;
737 self.a_offset -= last * self.a_strides[axis];
738 self.b_offset -= last * self.b_strides[axis];
739 }
740 }
741}
742
743#[inline(always)]
744unsafe fn run_contiguous_mul_row_block<
745 O: MulOutput<D>,
746 D: Copy + 'static,
747 A: Copy + Mul<B, Output = D> + 'static,
748 B: Copy + 'static,
749>(
750 dst_ptr: *mut O::Slot,
751 a_ptr: *const A,
752 b_ptr: *const B,
753 plan: &ContiguousMulRangePlan,
754 base_index: usize,
755 total: usize,
756 base_a_offset: isize,
757 base_b_offset: isize,
758) {
759 let inner_len = plan.inner_len.max(1);
760 let row_len = plan.row_len.max(1);
761 let block_len = inner_len.saturating_mul(row_len);
762
763 if total.saturating_sub(base_index) >= block_len {
764 let tiled = match transposed_scalar_tile_kind(plan) {
765 Some(TransposedScalarTileKind::RhsScalar) => {
766 simd::try_mul_transposed_scalar_rhs_2d::<D, A, B>(
767 dst_ptr.add(base_index).cast::<D>(),
768 a_ptr.offset(base_a_offset),
769 b_ptr.offset(base_b_offset),
770 inner_len,
771 row_len,
772 plan.a_fast_stride,
773 plan.a_row_stride,
774 )
775 }
776 Some(TransposedScalarTileKind::LhsScalar) => {
777 simd::try_mul_transposed_scalar_lhs_2d::<D, A, B>(
778 dst_ptr.add(base_index).cast::<D>(),
779 a_ptr.offset(base_a_offset),
780 b_ptr.offset(base_b_offset),
781 inner_len,
782 row_len,
783 plan.b_fast_stride,
784 plan.b_row_stride,
785 )
786 }
787 None => false,
788 };
789 if tiled {
790 return;
791 }
792 }
793
794 let mut index = base_index;
795 let mut a_offset = base_a_offset;
796 let mut b_offset = base_b_offset;
797
798 for _ in 0..row_len {
799 if index >= total {
800 break;
801 }
802 let len = inner_len.min(total - index);
803 inner_loop_mul2::<O, D, A, B>(
804 dst_ptr.add(index),
805 1,
806 a_ptr.offset(a_offset),
807 plan.a_fast_stride,
808 b_ptr.offset(b_offset),
809 plan.b_fast_stride,
810 len,
811 );
812 index += inner_len;
813 a_offset += plan.a_row_stride;
814 b_offset += plan.b_row_stride;
815 }
816}
817
818#[cfg(feature = "parallel")]
819fn strided_offset_for_contiguous_linear_index(
820 dims: &[usize],
821 strides: &[isize],
822 axis_order: &[usize],
823 mut index: usize,
824) -> isize {
825 let mut offset = 0isize;
826 for &axis in axis_order {
827 let dim = dims[axis];
828 if dim == 0 {
829 return 0;
830 }
831 let coord = index % dim;
832 index /= dim;
833 offset += coord as isize * strides[axis];
834 }
835 offset
836}
837
838fn try_contiguous_range_mul<
839 O: MulOutput<D>,
840 D: Copy + MaybeSendSync + 'static,
841 A: Copy + MaybeSendSync + Mul<B, Output = D> + 'static,
842 B: Copy + MaybeSendSync + 'static,
843>(
844 dst_ptr: *mut O::Slot,
845 dims: &[usize],
846 dst_strides: &[isize],
847 a_ptr: *const A,
848 a_strides: &[isize],
849 b_ptr: *const B,
850 b_strides: &[isize],
851) -> bool {
852 let Ok(total) = total_len(dims) else {
855 return false;
856 };
857 if total == 0 {
858 return true;
859 }
860 if total <= CONTIGUOUS_RANGE_MIN_LEN {
861 return false;
862 }
863
864 let Some(plan) = contiguous_mul_range_plan(dims, dst_strides, a_strides, b_strides) else {
865 return false;
866 };
867 if !plan.walks_inputs_without_blocking() {
868 return false;
869 }
870
871 let inner_len = plan.inner_len.max(1);
872 let row_len = plan.row_len.max(1);
873 let block_len = inner_len.saturating_mul(row_len).max(1);
874 let outer_groups = total.div_ceil(block_len);
875
876 #[cfg(feature = "parallel")]
877 {
878 let nthreads = crate::execution_policy::rayon_threads();
879 if nthreads > 1 {
880 use crate::threading::{parallel_for_each, SendPtr};
881
882 let dst = SendPtr(dst_ptr);
883 let a = SendPtr(a_ptr as *mut A);
884 let b = SendPtr(b_ptr as *mut B);
885
886 if outer_groups < nthreads {
887 let chunk_len = total.div_ceil(nthreads);
888 let nchunks = total.div_ceil(chunk_len);
889
890 parallel_for_each(0..nchunks, nthreads, &|chunks| {
891 for chunk in chunks {
892 let start = chunk * chunk_len;
893 let end = (start + chunk_len).min(total);
894 let mut index = start;
895
896 while index < end {
897 let in_inner = index % inner_len;
898 let len = (inner_len - in_inner).min(end - index);
899 let a_offset = strided_offset_for_contiguous_linear_index(
900 dims,
901 a_strides,
902 &plan.axis_order,
903 index,
904 );
905 let b_offset = strided_offset_for_contiguous_linear_index(
906 dims,
907 b_strides,
908 &plan.axis_order,
909 index,
910 );
911
912 unsafe {
913 inner_loop_mul2::<O, D, A, B>(
914 dst.as_ptr().add(index),
915 1,
916 a.as_const().offset(a_offset),
917 plan.a_fast_stride,
918 b.as_const().offset(b_offset),
919 plan.b_fast_stride,
920 len,
921 );
922 }
923 index += len;
924 }
925 }
926 });
927
928 return true;
929 }
930
931 let groups_per_chunk = outer_groups.div_ceil(nthreads);
932 let nchunks = outer_groups.div_ceil(groups_per_chunk);
933
934 parallel_for_each(0..nchunks, nthreads, &|chunks| {
935 for chunk in chunks {
936 let group_start = chunk * groups_per_chunk;
937 let group_end = (group_start + groups_per_chunk).min(outer_groups);
938 let mut cursor = ContiguousMulOuterCursor::new(
939 dims,
940 a_strides,
941 b_strides,
942 &plan,
943 group_start,
944 );
945
946 for group in group_start..group_end {
947 let index = group * block_len;
948 unsafe {
949 run_contiguous_mul_row_block::<O, D, A, B>(
950 dst.as_ptr(),
951 a.as_const(),
952 b.as_const(),
953 &plan,
954 index,
955 total,
956 cursor.a_offset,
957 cursor.b_offset,
958 );
959 }
960 cursor.advance();
961 }
962 }
963 });
964
965 true
966 } else {
967 run_contiguous_range_mul_single_thread::<O, D, A, B>(
968 dst_ptr,
969 dims,
970 a_ptr,
971 a_strides,
972 b_ptr,
973 b_strides,
974 &plan,
975 total,
976 block_len,
977 outer_groups,
978 )
979 }
980 }
981
982 #[cfg(not(feature = "parallel"))]
983 {
984 run_contiguous_range_mul_single_thread::<O, D, A, B>(
985 dst_ptr,
986 dims,
987 a_ptr,
988 a_strides,
989 b_ptr,
990 b_strides,
991 &plan,
992 total,
993 block_len,
994 outer_groups,
995 )
996 }
997}
998
999fn run_contiguous_range_mul_single_thread<
1000 O: MulOutput<D>,
1001 D: Copy + MaybeSendSync + 'static,
1002 A: Copy + MaybeSendSync + Mul<B, Output = D> + 'static,
1003 B: Copy + MaybeSendSync + 'static,
1004>(
1005 dst_ptr: *mut O::Slot,
1006 dims: &[usize],
1007 a_ptr: *const A,
1008 a_strides: &[isize],
1009 b_ptr: *const B,
1010 b_strides: &[isize],
1011 plan: &ContiguousMulRangePlan,
1012 total: usize,
1013 block_len: usize,
1014 outer_groups: usize,
1015) -> bool {
1016 let mut cursor = ContiguousMulOuterCursor::new(dims, a_strides, b_strides, plan, 0);
1017 for group in 0..outer_groups {
1018 let index = group * block_len;
1019 unsafe {
1020 run_contiguous_mul_row_block::<O, D, A, B>(
1021 dst_ptr,
1022 a_ptr,
1023 b_ptr,
1024 plan,
1025 index,
1026 total,
1027 cursor.a_offset,
1028 cursor.b_offset,
1029 );
1030 }
1031 cursor.advance();
1032 }
1033 true
1034}
1035
1036#[inline(always)]
1038unsafe fn inner_loop_map3<
1039 D: Copy,
1040 A: Copy,
1041 B: Copy,
1042 C: Copy,
1043 OpA: ElementOp<A>,
1044 OpB: ElementOp<B>,
1045 OpC: ElementOp<C>,
1046>(
1047 dp: *mut D,
1048 ds: isize,
1049 ap: *const A,
1050 a_s: isize,
1051 bp: *const B,
1052 b_s: isize,
1053 cp: *const C,
1054 c_s: isize,
1055 len: usize,
1056 f: &impl Fn(A, B, C) -> D,
1057) {
1058 if ds == 1 && a_s == 1 && b_s == 1 && c_s == 1 {
1059 let src_a = std::slice::from_raw_parts(ap, len);
1060 let src_b = std::slice::from_raw_parts(bp, len);
1061 let src_c = std::slice::from_raw_parts(cp, len);
1062 let dst = std::slice::from_raw_parts_mut(dp, len);
1063 simd::dispatch_if_large(len, || {
1064 for (((d, &a), &b), &c) in dst.iter_mut().zip(src_a).zip(src_b).zip(src_c) {
1065 *d = f(OpA::apply(a), OpB::apply(b), OpC::apply(c));
1066 }
1067 });
1068 } else {
1069 let mut dp = dp;
1070 let mut ap = ap;
1071 let mut bp = bp;
1072 let mut cp = cp;
1073 for _ in 0..len {
1074 *dp = f(OpA::apply(*ap), OpB::apply(*bp), OpC::apply(*cp));
1075 dp = dp.wrapping_offset(ds);
1076 ap = ap.wrapping_offset(a_s);
1077 bp = bp.wrapping_offset(b_s);
1078 cp = cp.wrapping_offset(c_s);
1079 }
1080 }
1081}
1082
1083#[inline(always)]
1085unsafe fn inner_loop_map4<
1086 D: Copy,
1087 A: Copy,
1088 B: Copy,
1089 C: Copy,
1090 E: Copy,
1091 OpA: ElementOp<A>,
1092 OpB: ElementOp<B>,
1093 OpC: ElementOp<C>,
1094 OpE: ElementOp<E>,
1095>(
1096 dp: *mut D,
1097 ds: isize,
1098 ap: *const A,
1099 a_s: isize,
1100 bp: *const B,
1101 b_s: isize,
1102 cp: *const C,
1103 c_s: isize,
1104 ep: *const E,
1105 e_s: isize,
1106 len: usize,
1107 f: &impl Fn(A, B, C, E) -> D,
1108) {
1109 if ds == 1 && a_s == 1 && b_s == 1 && c_s == 1 && e_s == 1 {
1110 let src_a = std::slice::from_raw_parts(ap, len);
1111 let src_b = std::slice::from_raw_parts(bp, len);
1112 let src_c = std::slice::from_raw_parts(cp, len);
1113 let src_e = std::slice::from_raw_parts(ep, len);
1114 let dst = std::slice::from_raw_parts_mut(dp, len);
1115 simd::dispatch_if_large(len, || {
1116 for i in 0..len {
1117 dst[i] = f(
1118 OpA::apply(src_a[i]),
1119 OpB::apply(src_b[i]),
1120 OpC::apply(src_c[i]),
1121 OpE::apply(src_e[i]),
1122 );
1123 }
1124 });
1125 } else {
1126 let mut dp = dp;
1127 let mut ap = ap;
1128 let mut bp = bp;
1129 let mut cp = cp;
1130 let mut ep = ep;
1131 for _ in 0..len {
1132 *dp = f(
1133 OpA::apply(*ap),
1134 OpB::apply(*bp),
1135 OpC::apply(*cp),
1136 OpE::apply(*ep),
1137 );
1138 dp = dp.wrapping_offset(ds);
1139 ap = ap.wrapping_offset(a_s);
1140 bp = bp.wrapping_offset(b_s);
1141 cp = cp.wrapping_offset(c_s);
1142 ep = ep.wrapping_offset(e_s);
1143 }
1144 }
1145}
1146
1147pub fn map_into<D: Copy + MaybeSendSync, A: Copy + MaybeSendSync, Op: ElementOp<A>>(
1152 dest: &mut StridedViewMut<D>,
1153 src: &StridedView<A, Op>,
1154 f: impl Fn(A) -> D + MaybeSync,
1155) -> Result<()> {
1156 map_parts_into::<D, A, Op, NoKernel>(
1157 dest.as_mut_ptr(),
1158 dest.dims(),
1159 dest.strides(),
1160 src.ptr(),
1161 src.dims(),
1162 src.strides(),
1163 f,
1164 &NoKernel,
1165 )
1166}
1167
1168pub(crate) fn map_into_validated<
1169 D: Copy + MaybeSendSync,
1170 A: Copy + MaybeSendSync,
1171 Op: ElementOp<A>,
1172>(
1173 dest: &mut StridedViewMut<D>,
1174 src: &StridedView<A, Op>,
1175 f: impl Fn(A) -> D + MaybeSync,
1176 validated: ValidatedDestinationLayout,
1177) -> Result<()> {
1178 ensure_same_shape(dest.dims(), src.dims())?;
1179 map_parts_into_validated::<D, A, Op, NoKernel>(
1180 dest.as_mut_ptr(),
1181 dest.dims(),
1182 dest.strides(),
1183 src.ptr(),
1184 src.strides(),
1185 f,
1186 &NoKernel,
1187 validated,
1188 )
1189}
1190
1191pub fn map_into_uninit_with_kernel<T: Copy + MaybeSendSync>(
1232 dest: &mut StridedViewMut<MaybeUninit<T>>,
1233 src: &StridedView<T>,
1234 kernel: impl Fn(&[T], &mut [MaybeUninit<T>]) -> bool + Sync,
1235 scalar: impl Fn(T) -> T + MaybeSync,
1236) -> Result<()> {
1237 map_parts_into::<MaybeUninit<T>, T, Identity, _>(
1238 dest.as_mut_ptr(),
1239 dest.dims(),
1240 dest.strides(),
1241 src.ptr(),
1242 src.dims(),
1243 src.strides(),
1244 |value| MaybeUninit::new(scalar(value)),
1245 &SliceKernel(kernel),
1246 )
1247}
1248
1249#[cfg(test)]
1255pub(crate) fn map_raw_into<D: Copy + MaybeSendSync, A: Copy + MaybeSendSync, Op: ElementOp<A>>(
1256 dest: &mut crate::RawStridedMut<'_, D>,
1257 src: &crate::RawStridedRef<'_, A>,
1258 f: impl Fn(A) -> D + MaybeSync,
1259) -> Result<()> {
1260 map_parts_into::<D, A, Op, NoKernel>(
1261 dest.as_mut_ptr(),
1262 dest.dims(),
1263 dest.strides(),
1264 src.ptr(),
1265 src.dims(),
1266 src.strides(),
1267 f,
1268 &NoKernel,
1269 )
1270}
1271
1272pub(crate) fn map_raw_into_validated<
1273 D: Copy + MaybeSendSync,
1274 A: Copy + MaybeSendSync,
1275 Op: ElementOp<A>,
1276>(
1277 dest: &mut crate::RawStridedMut<'_, D>,
1278 src: &crate::RawStridedRef<'_, A>,
1279 f: impl Fn(A) -> D + MaybeSync,
1280 validated: ValidatedDestinationLayout,
1281) -> Result<()> {
1282 ensure_same_shape(dest.dims(), src.dims())?;
1283 map_parts_into_validated::<D, A, Op, NoKernel>(
1284 dest.as_mut_ptr(),
1285 dest.dims(),
1286 dest.strides(),
1287 src.ptr(),
1288 src.strides(),
1289 f,
1290 &NoKernel,
1291 validated,
1292 )
1293}
1294
1295#[allow(clippy::too_many_arguments)]
1296fn map_parts_into<
1297 D: Copy + MaybeSendSync,
1298 A: Copy + MaybeSendSync,
1299 Op: ElementOp<A>,
1300 K: MapKernel<A, D> + Sync,
1301>(
1302 dst_ptr: *mut D,
1303 dst_dims: &[usize],
1304 dst_strides: &[isize],
1305 src_ptr: *const A,
1306 src_dims: &[usize],
1307 src_strides: &[isize],
1308 f: impl Fn(A) -> D + MaybeSync,
1309 kernel: &K,
1310) -> Result<()> {
1311 ensure_same_shape(dst_dims, src_dims)?;
1312 let validated = validate_destination_layout(dst_dims, dst_strides)?;
1313 map_parts_into_validated::<D, A, Op, K>(
1314 dst_ptr,
1315 dst_dims,
1316 dst_strides,
1317 src_ptr,
1318 src_strides,
1319 f,
1320 kernel,
1321 validated,
1322 )
1323}
1324
1325#[allow(clippy::too_many_arguments)]
1326#[inline(never)]
1333fn map_contiguous<D: Copy, A: Copy, Op: ElementOp<A>>(
1334 dst: &mut [D],
1335 src: &[A],
1336 f: impl Fn(A) -> D,
1337) {
1338 simd::dispatch_if_large(dst.len(), move || {
1339 for i in 0..dst.len() {
1340 dst[i] = f(Op::apply(src[i]));
1341 }
1342 });
1343}
1344
1345fn map_parts_into_validated<
1346 D: Copy + MaybeSendSync,
1347 A: Copy + MaybeSendSync,
1348 Op: ElementOp<A>,
1349 K: MapKernel<A, D> + Sync,
1350>(
1351 dst_ptr: *mut D,
1352 dst_dims: &[usize],
1353 dst_strides: &[isize],
1354 src_ptr: *const A,
1355 src_strides: &[isize],
1356 f: impl Fn(A) -> D + MaybeSync,
1357 kernel: &K,
1358 _validated: ValidatedDestinationLayout,
1359) -> Result<()> {
1360 if sequential_contiguous_layout(dst_dims, &[dst_strides, src_strides])?.is_some() {
1361 let len = total_len(dst_dims)?;
1362 let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
1366 let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
1367 if kernel.apply(src, dst) {
1368 return Ok(());
1369 }
1370 map_contiguous::<D, A, Op>(dst, src, f);
1371 return Ok(());
1372 }
1373
1374 let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
1375 let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<A>());
1376 let total = total_len(dst_dims)?;
1377
1378 let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
1380 build_plan_fused_small(dst_dims, &strides_list)
1381 } else {
1382 build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
1383 };
1384
1385 #[cfg(feature = "parallel")]
1386 {
1387 let total = total_len(&fused_dims)?;
1388 let nthreads = crate::execution_policy::rayon_threads();
1389 if total > MINTHREADLENGTH && nthreads > 1 {
1390 use crate::threading::SendPtr;
1391 let dst_send = SendPtr(dst_ptr);
1392 let src_send = SendPtr(src_ptr as *mut A);
1393
1394 let costs = compute_costs(&ordered_strides);
1395 let initial_offsets = vec![0isize; strides_list.len()];
1396 return mapreduce_threaded(
1397 &fused_dims,
1398 &plan.block,
1399 &ordered_strides,
1400 &initial_offsets,
1401 &costs,
1402 nthreads,
1403 0,
1404 1,
1405 &|dims, blocks, strides_list, offsets| {
1406 for_each_inner_block_with_offsets(
1407 dims,
1408 blocks,
1409 strides_list,
1410 offsets,
1411 |offsets, len, strides| {
1412 let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
1413 let sp = unsafe { src_send.as_const().offset(offsets[1]) };
1414 unsafe {
1415 inner_loop_map1::<D, A, Op, K>(
1416 dp, strides[0], sp, strides[1], len, &f, kernel,
1417 )
1418 };
1419 Ok(())
1420 },
1421 )
1422 },
1423 );
1424 }
1425 }
1426
1427 let initial_offsets = vec![0isize; ordered_strides.len()];
1428 for_each_inner_block_preordered(
1429 &fused_dims,
1430 &plan.block,
1431 &ordered_strides,
1432 &initial_offsets,
1433 |offsets, len, strides| {
1434 let dp = unsafe { dst_ptr.offset(offsets[0]) };
1435 let sp = unsafe { src_ptr.offset(offsets[1]) };
1436 unsafe {
1437 inner_loop_map1::<D, A, Op, K>(dp, strides[0], sp, strides[1], len, &f, kernel)
1438 };
1439 Ok(())
1440 },
1441 )
1442}
1443
1444pub fn zip_map2_into<
1449 D: Copy + MaybeSendSync,
1450 A: Copy + MaybeSendSync,
1451 B: Copy + MaybeSendSync,
1452 OpA: ElementOp<A>,
1453 OpB: ElementOp<B>,
1454>(
1455 dest: &mut StridedViewMut<D>,
1456 a: &StridedView<A, OpA>,
1457 b: &StridedView<B, OpB>,
1458 f: impl Fn(A, B) -> D + MaybeSync,
1459) -> Result<()> {
1460 zip_map2_parts_into::<D, A, B, OpA, OpB>(
1461 dest.as_mut_ptr(),
1462 dest.dims(),
1463 dest.strides(),
1464 a.ptr(),
1465 a.dims(),
1466 a.strides(),
1467 b.ptr(),
1468 b.dims(),
1469 b.strides(),
1470 f,
1471 )
1472}
1473
1474pub(crate) fn zip_map2_into_validated<
1475 D: Copy + MaybeSendSync,
1476 A: Copy + MaybeSendSync,
1477 B: Copy + MaybeSendSync,
1478 OpA: ElementOp<A>,
1479 OpB: ElementOp<B>,
1480>(
1481 dest: &mut StridedViewMut<D>,
1482 a: &StridedView<A, OpA>,
1483 b: &StridedView<B, OpB>,
1484 f: impl Fn(A, B) -> D + MaybeSync,
1485 validated: ValidatedDestinationLayout,
1486) -> Result<()> {
1487 zip_map2_parts_into_validated::<D, A, B, OpA, OpB>(
1488 dest.as_mut_ptr(),
1489 dest.dims(),
1490 dest.strides(),
1491 a.ptr(),
1492 a.strides(),
1493 b.ptr(),
1494 b.strides(),
1495 f,
1496 validated,
1497 )
1498}
1499
1500#[non_exhaustive]
1502#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1503pub enum CompareOp {
1504 Eq,
1505 Lt,
1506 Le,
1507 Gt,
1508 Ge,
1509}
1510
1511pub fn compare_into<T, OpA, OpB>(
1522 dest: &mut StridedViewMut<bool>,
1523 a: &StridedView<T, OpA>,
1524 b: &StridedView<T, OpB>,
1525 op: CompareOp,
1526) -> Result<()>
1527where
1528 T: Copy + MaybeSendSync + PartialOrd,
1529 OpA: ElementOp<T>,
1530 OpB: ElementOp<T>,
1531{
1532 match op {
1533 CompareOp::Eq => zip_map2_into(dest, a, b, |lhs, rhs| lhs == rhs),
1534 CompareOp::Lt => zip_map2_into(dest, a, b, |lhs, rhs| lhs < rhs),
1535 CompareOp::Le => zip_map2_into(dest, a, b, |lhs, rhs| lhs <= rhs),
1536 CompareOp::Gt => zip_map2_into(dest, a, b, |lhs, rhs| lhs > rhs),
1537 CompareOp::Ge => zip_map2_into(dest, a, b, |lhs, rhs| lhs >= rhs),
1538 }
1539}
1540
1541pub fn compare_into_uninit<T, OpA, OpB>(
1560 dest: &mut StridedViewMut<MaybeUninit<bool>>,
1561 a: &StridedView<T, OpA>,
1562 b: &StridedView<T, OpB>,
1563 op: CompareOp,
1564) -> Result<()>
1565where
1566 T: Copy + MaybeSendSync + PartialOrd,
1567 OpA: ElementOp<T>,
1568 OpB: ElementOp<T>,
1569{
1570 ensure_same_shape(dest.dims(), a.dims())?;
1571 ensure_same_shape(dest.dims(), b.dims())?;
1572 let validated = validate_destination_layout(dest.dims(), dest.strides())?;
1573 validate_typed_no_overlap(dest, a, 0)?;
1574 validate_typed_no_overlap(dest, b, 1)?;
1575 match op {
1576 CompareOp::Eq => zip_map2_into_validated(
1577 dest,
1578 a,
1579 b,
1580 |lhs, rhs| MaybeUninit::new(lhs == rhs),
1581 validated,
1582 ),
1583 CompareOp::Lt => zip_map2_into_validated(
1584 dest,
1585 a,
1586 b,
1587 |lhs, rhs| MaybeUninit::new(lhs < rhs),
1588 validated,
1589 ),
1590 CompareOp::Le => zip_map2_into_validated(
1591 dest,
1592 a,
1593 b,
1594 |lhs, rhs| MaybeUninit::new(lhs <= rhs),
1595 validated,
1596 ),
1597 CompareOp::Gt => zip_map2_into_validated(
1598 dest,
1599 a,
1600 b,
1601 |lhs, rhs| MaybeUninit::new(lhs > rhs),
1602 validated,
1603 ),
1604 CompareOp::Ge => zip_map2_into_validated(
1605 dest,
1606 a,
1607 b,
1608 |lhs, rhs| MaybeUninit::new(lhs >= rhs),
1609 validated,
1610 ),
1611 }
1612}
1613
1614pub(crate) fn zip_map2_raw_into_validated<
1615 D: Copy + MaybeSendSync,
1616 A: Copy + MaybeSendSync,
1617 B: Copy + MaybeSendSync,
1618 OpA: ElementOp<A>,
1619 OpB: ElementOp<B>,
1620>(
1621 dest: &mut crate::RawStridedMut<'_, D>,
1622 a: &crate::RawStridedRef<'_, A>,
1623 b: &crate::RawStridedRef<'_, B>,
1624 f: impl Fn(A, B) -> D + MaybeSync,
1625 validated: ValidatedDestinationLayout,
1626) -> Result<()> {
1627 ensure_same_shape(dest.dims(), a.dims())?;
1628 ensure_same_shape(dest.dims(), b.dims())?;
1629 zip_map2_parts_into_validated::<D, A, B, OpA, OpB>(
1630 dest.as_mut_ptr(),
1631 dest.dims(),
1632 dest.strides(),
1633 a.ptr(),
1634 a.strides(),
1635 b.ptr(),
1636 b.strides(),
1637 f,
1638 validated,
1639 )
1640}
1641
1642#[allow(clippy::too_many_arguments)]
1643fn zip_map2_parts_into<
1644 D: Copy + MaybeSendSync,
1645 A: Copy + MaybeSendSync,
1646 B: Copy + MaybeSendSync,
1647 OpA: ElementOp<A>,
1648 OpB: ElementOp<B>,
1649>(
1650 dst_ptr: *mut D,
1651 dst_dims: &[usize],
1652 dst_strides: &[isize],
1653 a_ptr: *const A,
1654 a_dims: &[usize],
1655 a_strides: &[isize],
1656 b_ptr: *const B,
1657 b_dims: &[usize],
1658 b_strides: &[isize],
1659 f: impl Fn(A, B) -> D + MaybeSync,
1660) -> Result<()> {
1661 ensure_same_shape(dst_dims, a_dims)?;
1662 ensure_same_shape(dst_dims, b_dims)?;
1663 let validated = validate_destination_layout(dst_dims, dst_strides)?;
1664 zip_map2_parts_into_validated::<D, A, B, OpA, OpB>(
1665 dst_ptr,
1666 dst_dims,
1667 dst_strides,
1668 a_ptr,
1669 a_strides,
1670 b_ptr,
1671 b_strides,
1672 f,
1673 validated,
1674 )
1675}
1676
1677#[allow(clippy::too_many_arguments)]
1678fn zip_map2_parts_into_validated<
1679 D: Copy + MaybeSendSync,
1680 A: Copy + MaybeSendSync,
1681 B: Copy + MaybeSendSync,
1682 OpA: ElementOp<A>,
1683 OpB: ElementOp<B>,
1684>(
1685 dst_ptr: *mut D,
1686 dst_dims: &[usize],
1687 dst_strides: &[isize],
1688 a_ptr: *const A,
1689 a_strides: &[isize],
1690 b_ptr: *const B,
1691 b_strides: &[isize],
1692 f: impl Fn(A, B) -> D + MaybeSync,
1693 _validated: ValidatedDestinationLayout,
1694) -> Result<()> {
1695 if sequential_contiguous_layout(dst_dims, &[dst_strides, a_strides, b_strides])?.is_some() {
1696 let len = total_len(dst_dims)?;
1697 let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
1698 let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
1699 let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
1700 simd::dispatch_if_large(len, || {
1701 for i in 0..len {
1702 dst[i] = f(OpA::apply(sa[i]), OpB::apply(sb[i]));
1703 }
1704 });
1705 return Ok(());
1706 }
1707
1708 let strides_list: [&[isize]; 3] = [dst_strides, a_strides, b_strides];
1709 let elem_size = std::mem::size_of::<D>()
1710 .max(std::mem::size_of::<A>())
1711 .max(std::mem::size_of::<B>());
1712 let total = total_len(dst_dims)?;
1713
1714 let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
1716 build_plan_fused_small(dst_dims, &strides_list)
1717 } else {
1718 build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
1719 };
1720
1721 #[cfg(feature = "parallel")]
1722 {
1723 let total = total_len(&fused_dims)?;
1724 let nthreads = crate::execution_policy::rayon_threads();
1725 if total > MINTHREADLENGTH && nthreads > 1 {
1726 use crate::threading::SendPtr;
1727 let dst_send = SendPtr(dst_ptr);
1728 let a_send = SendPtr(a_ptr as *mut A);
1729 let b_send = SendPtr(b_ptr as *mut B);
1730
1731 let costs = compute_costs(&ordered_strides);
1732 let initial_offsets = vec![0isize; strides_list.len()];
1733 return mapreduce_threaded(
1734 &fused_dims,
1735 &plan.block,
1736 &ordered_strides,
1737 &initial_offsets,
1738 &costs,
1739 nthreads,
1740 0,
1741 1,
1742 &|dims, blocks, strides_list, offsets| {
1743 for_each_inner_block_with_offsets(
1744 dims,
1745 blocks,
1746 strides_list,
1747 offsets,
1748 |offsets, len, strides| {
1749 let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
1750 let ap = unsafe { a_send.as_const().offset(offsets[1]) };
1751 let bp = unsafe { b_send.as_const().offset(offsets[2]) };
1752 unsafe {
1753 inner_loop_map2::<D, A, B, OpA, OpB>(
1754 dp, strides[0], ap, strides[1], bp, strides[2], len, &f,
1755 )
1756 };
1757 Ok(())
1758 },
1759 )
1760 },
1761 );
1762 }
1763 }
1764
1765 let initial_offsets = vec![0isize; ordered_strides.len()];
1766 for_each_inner_block_preordered(
1767 &fused_dims,
1768 &plan.block,
1769 &ordered_strides,
1770 &initial_offsets,
1771 |offsets, len, strides| {
1772 let dp = unsafe { dst_ptr.offset(offsets[0]) };
1773 let ap = unsafe { a_ptr.offset(offsets[1]) };
1774 let bp = unsafe { b_ptr.offset(offsets[2]) };
1775 unsafe {
1776 inner_loop_map2::<D, A, B, OpA, OpB>(
1777 dp, strides[0], ap, strides[1], bp, strides[2], len, &f,
1778 )
1779 };
1780 Ok(())
1781 },
1782 )
1783}
1784
1785fn mul_identity_into_raw<
1786 O: MulOutput<D>,
1787 D: Copy + MaybeSendSync + 'static,
1788 A: Copy + MaybeSendSync + Mul<B, Output = D> + 'static,
1789 B: Copy + MaybeSendSync + 'static,
1790>(
1791 dst_ptr: *mut O::Slot,
1792 dst_dims: &[usize],
1793 dst_strides: &[isize],
1794 a_ptr: *const A,
1795 a_strides: &[isize],
1796 b_ptr: *const B,
1797 b_strides: &[isize],
1798 _validated: ValidatedDestinationLayout,
1799) -> Result<()> {
1800 debug_assert_eq!(dst_dims.len(), a_strides.len());
1801 debug_assert_eq!(dst_dims.len(), b_strides.len());
1802
1803 if sequential_contiguous_layout(dst_dims, &[dst_strides, a_strides, b_strides])?.is_some() {
1804 let len = total_len(dst_dims)?;
1805 let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
1806 let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
1807 if unsafe { O::try_contiguous(dst_ptr, len, sa, sb) } {
1808 return Ok(());
1809 }
1810 for i in 0..len {
1811 unsafe { O::write(dst_ptr.add(i), multiply_value(sa[i], sb[i])) };
1812 }
1813 return Ok(());
1814 }
1815
1816 let strides_list: [&[isize]; 3] = [dst_strides, a_strides, b_strides];
1817 let elem_size = std::mem::size_of::<D>()
1818 .max(std::mem::size_of::<A>())
1819 .max(std::mem::size_of::<B>());
1820 let total = total_len(dst_dims)?;
1821
1822 if try_contiguous_range_mul::<O, D, A, B>(
1823 dst_ptr,
1824 dst_dims,
1825 dst_strides,
1826 a_ptr,
1827 a_strides,
1828 b_ptr,
1829 b_strides,
1830 ) {
1831 return Ok(());
1832 }
1833
1834 let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
1835 build_plan_fused_small(dst_dims, &strides_list)
1836 } else {
1837 build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
1838 };
1839
1840 #[cfg(feature = "parallel")]
1841 {
1842 let total = total_len(&fused_dims)?;
1843 let nthreads = crate::execution_policy::rayon_threads();
1844 if total > MINTHREADLENGTH && nthreads > 1 {
1845 use crate::threading::SendPtr;
1846 let dst_send = SendPtr(dst_ptr);
1847 let a_send = SendPtr(a_ptr as *mut A);
1848 let b_send = SendPtr(b_ptr as *mut B);
1849
1850 let costs = compute_costs(&ordered_strides);
1851 let initial_offsets = vec![0isize; strides_list.len()];
1852 return mapreduce_threaded(
1853 &fused_dims,
1854 &plan.block,
1855 &ordered_strides,
1856 &initial_offsets,
1857 &costs,
1858 nthreads,
1859 0,
1860 1,
1861 &|dims, blocks, strides_list, offsets| {
1862 for_each_inner_block_with_offsets(
1863 dims,
1864 blocks,
1865 strides_list,
1866 offsets,
1867 |offsets, len, strides| {
1868 let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
1869 let ap = unsafe { a_send.as_const().offset(offsets[1]) };
1870 let bp = unsafe { b_send.as_const().offset(offsets[2]) };
1871 unsafe {
1872 inner_loop_mul2::<O, D, A, B>(
1873 dp, strides[0], ap, strides[1], bp, strides[2], len,
1874 )
1875 };
1876 Ok(())
1877 },
1878 )
1879 },
1880 );
1881 }
1882 }
1883
1884 let initial_offsets = vec![0isize; ordered_strides.len()];
1885 for_each_inner_block_preordered(
1886 &fused_dims,
1887 &plan.block,
1888 &ordered_strides,
1889 &initial_offsets,
1890 |offsets, len, strides| {
1891 let dp = unsafe { dst_ptr.offset(offsets[0]) };
1892 let ap = unsafe { a_ptr.offset(offsets[1]) };
1893 let bp = unsafe { b_ptr.offset(offsets[2]) };
1894 unsafe {
1895 inner_loop_mul2::<O, D, A, B>(dp, strides[0], ap, strides[1], bp, strides[2], len)
1896 };
1897 Ok(())
1898 },
1899 )
1900}
1901
1902pub fn mul_into<
1907 D: Copy + MaybeSendSync + 'static,
1908 A: Copy + Mul<B, Output = D> + MaybeSendSync + 'static,
1909 B: Copy + MaybeSendSync + 'static,
1910 OpA: ElementOp<A>,
1911 OpB: ElementOp<B>,
1912>(
1913 dest: &mut StridedViewMut<D>,
1914 a: &StridedView<A, OpA>,
1915 b: &StridedView<B, OpB>,
1916) -> Result<()> {
1917 ensure_same_shape(dest.dims(), a.dims())?;
1918 ensure_same_shape(dest.dims(), b.dims())?;
1919 let validated = validate_destination_layout(dest.dims(), dest.strides())?;
1920
1921 if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
1922 return mul_identity_into_raw::<InitializedOutput, D, A, B>(
1923 dest.as_mut_ptr(),
1924 dest.dims(),
1925 dest.strides(),
1926 a.ptr(),
1927 a.strides(),
1928 b.ptr(),
1929 b.strides(),
1930 validated,
1931 );
1932 }
1933
1934 zip_map2_into_validated(dest, a, b, multiply_value, validated)
1935}
1936
1937pub fn mul_into_uninit<
1956 D: Copy + MaybeSendSync + 'static,
1957 A: Copy + Mul<B, Output = D> + MaybeSendSync + 'static,
1958 B: Copy + MaybeSendSync + 'static,
1959 OpA: ElementOp<A>,
1960 OpB: ElementOp<B>,
1961>(
1962 dest: &mut StridedViewMut<MaybeUninit<D>>,
1963 a: &StridedView<A, OpA>,
1964 b: &StridedView<B, OpB>,
1965) -> Result<()> {
1966 ensure_same_shape(dest.dims(), a.dims())?;
1967 ensure_same_shape(dest.dims(), b.dims())?;
1968 let validated = validate_destination_layout(dest.dims(), dest.strides())?;
1969 validate_typed_no_overlap(dest, a, 0)?;
1970 validate_typed_no_overlap(dest, b, 1)?;
1971
1972 if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
1973 return mul_identity_into_raw::<UninitializedOutput, D, A, B>(
1974 dest.as_mut_ptr(),
1975 dest.dims(),
1976 dest.strides(),
1977 a.ptr(),
1978 a.strides(),
1979 b.ptr(),
1980 b.strides(),
1981 validated,
1982 );
1983 }
1984
1985 zip_map2_into_validated(
1986 dest,
1987 a,
1988 b,
1989 |lhs, rhs| MaybeUninit::new(multiply_value(lhs, rhs)),
1990 validated,
1991 )
1992}
1993
1994fn broadcast_strides_for_axes(
1995 source_dims: &[usize],
1996 source_strides: &[isize],
1997 target_dims: &[usize],
1998 axes: &[usize],
1999) -> Result<AxisVec<isize>> {
2000 if source_dims.len() != axes.len() {
2001 return Err(StridedError::RankMismatch(source_dims.len(), axes.len()));
2002 }
2003 debug_assert_eq!(source_dims.len(), source_strides.len());
2004
2005 let mut seen = AxisVec::<bool>::new();
2006 seen.resize(target_dims.len(), false);
2007 let mut strides = AxisVec::<isize>::new();
2008 strides.resize(target_dims.len(), 0);
2009 for (src_axis, &dst_axis) in axes.iter().enumerate() {
2010 if dst_axis >= target_dims.len() {
2011 return Err(StridedError::InvalidAxis {
2012 axis: dst_axis,
2013 rank: target_dims.len(),
2014 });
2015 }
2016 if seen[dst_axis] {
2017 return Err(StridedError::InvalidAxis {
2018 axis: dst_axis,
2019 rank: target_dims.len(),
2020 });
2021 }
2022 seen[dst_axis] = true;
2023
2024 let source_dim = source_dims[src_axis];
2025 let target_dim = target_dims[dst_axis];
2026 if source_dim != target_dim && source_dim != 1 {
2027 return Err(StridedError::ShapeMismatch(
2028 source_dims.to_vec(),
2029 target_dims.to_vec(),
2030 ));
2031 }
2032 if source_dim == target_dim {
2033 strides[dst_axis] = source_strides[src_axis];
2034 }
2035 }
2036
2037 Ok(strides)
2038}
2039
2040fn broadcast_view_with_strides<'a, T, Op: ElementOp<T>>(
2045 view: &StridedView<'a, T, Op>,
2046 target_dims: &[usize],
2047 strides: &[isize],
2048) -> StridedView<'a, T, Op> {
2049 unsafe {
2050 StridedView::new_owned_unchecked(
2051 view.data(),
2052 target_dims.to_vec(),
2053 strides.to_vec(),
2054 view.offset(),
2055 )
2056 }
2057}
2058
2059pub fn broadcast_mul_into<
2064 D: Copy + MaybeSendSync + 'static,
2065 A: Copy + Mul<B, Output = D> + MaybeSendSync + 'static,
2066 B: Copy + MaybeSendSync + 'static,
2067 OpA: ElementOp<A>,
2068 OpB: ElementOp<B>,
2069>(
2070 dest: &mut StridedViewMut<D>,
2071 a: &StridedView<A, OpA>,
2072 a_axes: &[usize],
2073 b: &StridedView<B, OpB>,
2074 b_axes: &[usize],
2075) -> Result<()> {
2076 let validated = validate_destination_layout(dest.dims(), dest.strides())?;
2077 let a_strides = broadcast_strides_for_axes(a.dims(), a.strides(), dest.dims(), a_axes)?;
2078 let b_strides = broadcast_strides_for_axes(b.dims(), b.strides(), dest.dims(), b_axes)?;
2079
2080 if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
2081 return mul_identity_into_raw::<InitializedOutput, D, A, B>(
2082 dest.as_mut_ptr(),
2083 dest.dims(),
2084 dest.strides(),
2085 a.ptr(),
2086 &a_strides,
2087 b.ptr(),
2088 &b_strides,
2089 validated,
2090 );
2091 }
2092
2093 let a = broadcast_view_with_strides(a, dest.dims(), &a_strides);
2094 let b = broadcast_view_with_strides(b, dest.dims(), &b_strides);
2095 zip_map2_into_validated(dest, &a, &b, multiply_value, validated)
2096}
2097
2098pub fn broadcast_mul_into_uninit<
2117 D: Copy + MaybeSendSync + 'static,
2118 A: Copy + Mul<B, Output = D> + MaybeSendSync + 'static,
2119 B: Copy + MaybeSendSync + 'static,
2120 OpA: ElementOp<A>,
2121 OpB: ElementOp<B>,
2122>(
2123 dest: &mut StridedViewMut<MaybeUninit<D>>,
2124 a: &StridedView<A, OpA>,
2125 a_axes: &[usize],
2126 b: &StridedView<B, OpB>,
2127 b_axes: &[usize],
2128) -> Result<()> {
2129 let validated = validate_destination_layout(dest.dims(), dest.strides())?;
2130 let a_strides = broadcast_strides_for_axes(a.dims(), a.strides(), dest.dims(), a_axes)?;
2131 let b_strides = broadcast_strides_for_axes(b.dims(), b.strides(), dest.dims(), b_axes)?;
2132 validate_typed_no_overlap(dest, a, 0)?;
2133 validate_typed_no_overlap(dest, b, 1)?;
2134
2135 if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
2136 return mul_identity_into_raw::<UninitializedOutput, D, A, B>(
2137 dest.as_mut_ptr(),
2138 dest.dims(),
2139 dest.strides(),
2140 a.ptr(),
2141 &a_strides,
2142 b.ptr(),
2143 &b_strides,
2144 validated,
2145 );
2146 }
2147
2148 let a = broadcast_view_with_strides(a, dest.dims(), &a_strides);
2149 let b = broadcast_view_with_strides(b, dest.dims(), &b_strides);
2150 zip_map2_into_validated(
2151 dest,
2152 &a,
2153 &b,
2154 |lhs, rhs| MaybeUninit::new(multiply_value(lhs, rhs)),
2155 validated,
2156 )
2157}
2158
2159pub fn zip_map3_into<
2161 D: Copy + MaybeSendSync,
2162 A: Copy + MaybeSendSync,
2163 B: Copy + MaybeSendSync,
2164 C: Copy + MaybeSendSync,
2165 OpA: ElementOp<A>,
2166 OpB: ElementOp<B>,
2167 OpC: ElementOp<C>,
2168>(
2169 dest: &mut StridedViewMut<D>,
2170 a: &StridedView<A, OpA>,
2171 b: &StridedView<B, OpB>,
2172 c: &StridedView<C, OpC>,
2173 f: impl Fn(A, B, C) -> D + MaybeSync,
2174) -> Result<()> {
2175 ensure_same_shape(dest.dims(), a.dims())?;
2176 ensure_same_shape(dest.dims(), b.dims())?;
2177 ensure_same_shape(dest.dims(), c.dims())?;
2178 let validated = validate_destination_layout(dest.dims(), dest.strides())?;
2179 zip_map3_into_validated(dest, a, b, c, f, validated)
2180}
2181
2182pub(crate) fn zip_map3_into_validated<
2183 D: Copy + MaybeSendSync,
2184 A: Copy + MaybeSendSync,
2185 B: Copy + MaybeSendSync,
2186 C: Copy + MaybeSendSync,
2187 OpA: ElementOp<A>,
2188 OpB: ElementOp<B>,
2189 OpC: ElementOp<C>,
2190>(
2191 dest: &mut StridedViewMut<D>,
2192 a: &StridedView<A, OpA>,
2193 b: &StridedView<B, OpB>,
2194 c: &StridedView<C, OpC>,
2195 f: impl Fn(A, B, C) -> D + MaybeSync,
2196 _validated: ValidatedDestinationLayout,
2197) -> Result<()> {
2198 ensure_same_shape(dest.dims(), a.dims())?;
2199 ensure_same_shape(dest.dims(), b.dims())?;
2200 ensure_same_shape(dest.dims(), c.dims())?;
2201 let dst_ptr = dest.as_mut_ptr();
2202 let a_ptr = a.ptr();
2203 let b_ptr = b.ptr();
2204 let c_ptr = c.ptr();
2205
2206 let dst_dims = dest.dims();
2207 let dst_strides = dest.strides();
2208
2209 if sequential_contiguous_layout(
2210 dst_dims,
2211 &[dst_strides, a.strides(), b.strides(), c.strides()],
2212 )?
2213 .is_some()
2214 {
2215 let len = total_len(dst_dims)?;
2216 let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
2217 let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
2218 let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
2219 let sc = unsafe { std::slice::from_raw_parts(c_ptr, len) };
2220 simd::dispatch_if_large(len, || {
2221 for (((d, &a), &b), &c) in dst.iter_mut().zip(sa).zip(sb).zip(sc) {
2222 *d = f(OpA::apply(a), OpB::apply(b), OpC::apply(c));
2223 }
2224 });
2225 return Ok(());
2226 }
2227
2228 let strides_list: [&[isize]; 4] = [dst_strides, a.strides(), b.strides(), c.strides()];
2229 let elem_size = std::mem::size_of::<D>()
2230 .max(std::mem::size_of::<A>())
2231 .max(std::mem::size_of::<B>())
2232 .max(std::mem::size_of::<C>());
2233 let total = total_len(dst_dims)?;
2234
2235 let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
2237 build_plan_fused_small(dst_dims, &strides_list)
2238 } else {
2239 build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
2240 };
2241
2242 #[cfg(feature = "parallel")]
2243 {
2244 let total = total_len(&fused_dims)?;
2245 let nthreads = crate::execution_policy::rayon_threads();
2246 if total > MINTHREADLENGTH && nthreads > 1 {
2247 use crate::threading::SendPtr;
2248 let dst_send = SendPtr(dst_ptr);
2249 let a_send = SendPtr(a_ptr as *mut A);
2250 let b_send = SendPtr(b_ptr as *mut B);
2251 let c_send = SendPtr(c_ptr as *mut C);
2252
2253 let costs = compute_costs(&ordered_strides);
2254 let initial_offsets = vec![0isize; strides_list.len()];
2255 return mapreduce_threaded(
2256 &fused_dims,
2257 &plan.block,
2258 &ordered_strides,
2259 &initial_offsets,
2260 &costs,
2261 nthreads,
2262 0,
2263 1,
2264 &|dims, blocks, strides_list, offsets| {
2265 for_each_inner_block_with_offsets(
2266 dims,
2267 blocks,
2268 strides_list,
2269 offsets,
2270 |offsets, len, strides| {
2271 let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
2272 let ap = unsafe { a_send.as_const().offset(offsets[1]) };
2273 let bp = unsafe { b_send.as_const().offset(offsets[2]) };
2274 let cp = unsafe { c_send.as_const().offset(offsets[3]) };
2275 unsafe {
2276 inner_loop_map3::<D, A, B, C, OpA, OpB, OpC>(
2277 dp, strides[0], ap, strides[1], bp, strides[2], cp, strides[3],
2278 len, &f,
2279 )
2280 };
2281 Ok(())
2282 },
2283 )
2284 },
2285 );
2286 }
2287 }
2288
2289 let initial_offsets = vec![0isize; ordered_strides.len()];
2290 for_each_inner_block_preordered(
2291 &fused_dims,
2292 &plan.block,
2293 &ordered_strides,
2294 &initial_offsets,
2295 |offsets, len, strides| {
2296 let dp = unsafe { dst_ptr.offset(offsets[0]) };
2297 let ap = unsafe { a_ptr.offset(offsets[1]) };
2298 let bp = unsafe { b_ptr.offset(offsets[2]) };
2299 let cp = unsafe { c_ptr.offset(offsets[3]) };
2300 unsafe {
2301 inner_loop_map3::<D, A, B, C, OpA, OpB, OpC>(
2302 dp, strides[0], ap, strides[1], bp, strides[2], cp, strides[3], len, &f,
2303 )
2304 };
2305 Ok(())
2306 },
2307 )
2308}
2309
2310pub fn zip_map4_into<
2312 D: Copy + MaybeSendSync,
2313 A: Copy + MaybeSendSync,
2314 B: Copy + MaybeSendSync,
2315 C: Copy + MaybeSendSync,
2316 E: Copy + MaybeSendSync,
2317 OpA: ElementOp<A>,
2318 OpB: ElementOp<B>,
2319 OpC: ElementOp<C>,
2320 OpE: ElementOp<E>,
2321>(
2322 dest: &mut StridedViewMut<D>,
2323 a: &StridedView<A, OpA>,
2324 b: &StridedView<B, OpB>,
2325 c: &StridedView<C, OpC>,
2326 e: &StridedView<E, OpE>,
2327 f: impl Fn(A, B, C, E) -> D + MaybeSync,
2328) -> Result<()> {
2329 ensure_same_shape(dest.dims(), a.dims())?;
2330 ensure_same_shape(dest.dims(), b.dims())?;
2331 ensure_same_shape(dest.dims(), c.dims())?;
2332 ensure_same_shape(dest.dims(), e.dims())?;
2333 let validated = validate_destination_layout(dest.dims(), dest.strides())?;
2334 zip_map4_into_validated(dest, a, b, c, e, f, validated)
2335}
2336
2337pub(crate) fn zip_map4_into_validated<
2338 D: Copy + MaybeSendSync,
2339 A: Copy + MaybeSendSync,
2340 B: Copy + MaybeSendSync,
2341 C: Copy + MaybeSendSync,
2342 E: Copy + MaybeSendSync,
2343 OpA: ElementOp<A>,
2344 OpB: ElementOp<B>,
2345 OpC: ElementOp<C>,
2346 OpE: ElementOp<E>,
2347>(
2348 dest: &mut StridedViewMut<D>,
2349 a: &StridedView<A, OpA>,
2350 b: &StridedView<B, OpB>,
2351 c: &StridedView<C, OpC>,
2352 e: &StridedView<E, OpE>,
2353 f: impl Fn(A, B, C, E) -> D + MaybeSync,
2354 _validated: ValidatedDestinationLayout,
2355) -> Result<()> {
2356 ensure_same_shape(dest.dims(), a.dims())?;
2357 ensure_same_shape(dest.dims(), b.dims())?;
2358 ensure_same_shape(dest.dims(), c.dims())?;
2359 ensure_same_shape(dest.dims(), e.dims())?;
2360 let dst_ptr = dest.as_mut_ptr();
2361 let a_ptr = a.ptr();
2362 let b_ptr = b.ptr();
2363 let c_ptr = c.ptr();
2364 let e_ptr = e.ptr();
2365
2366 let dst_dims = dest.dims();
2367 let dst_strides = dest.strides();
2368
2369 if sequential_contiguous_layout(
2370 dst_dims,
2371 &[
2372 dst_strides,
2373 a.strides(),
2374 b.strides(),
2375 c.strides(),
2376 e.strides(),
2377 ],
2378 )?
2379 .is_some()
2380 {
2381 let len = total_len(dst_dims)?;
2382 let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
2383 let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
2384 let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
2385 let sc = unsafe { std::slice::from_raw_parts(c_ptr, len) };
2386 let se = unsafe { std::slice::from_raw_parts(e_ptr, len) };
2387 simd::dispatch_if_large(len, || {
2388 for i in 0..len {
2389 dst[i] = f(
2390 OpA::apply(sa[i]),
2391 OpB::apply(sb[i]),
2392 OpC::apply(sc[i]),
2393 OpE::apply(se[i]),
2394 );
2395 }
2396 });
2397 return Ok(());
2398 }
2399
2400 let strides_list: [&[isize]; 5] = [
2401 dst_strides,
2402 a.strides(),
2403 b.strides(),
2404 c.strides(),
2405 e.strides(),
2406 ];
2407 let elem_size = std::mem::size_of::<D>()
2408 .max(std::mem::size_of::<A>())
2409 .max(std::mem::size_of::<B>())
2410 .max(std::mem::size_of::<C>())
2411 .max(std::mem::size_of::<E>());
2412 let total = total_len(dst_dims)?;
2413
2414 let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
2416 build_plan_fused_small(dst_dims, &strides_list)
2417 } else {
2418 build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
2419 };
2420
2421 #[cfg(feature = "parallel")]
2422 {
2423 let total = total_len(&fused_dims)?;
2424 let nthreads = crate::execution_policy::rayon_threads();
2425 if total > MINTHREADLENGTH && nthreads > 1 {
2426 use crate::threading::SendPtr;
2427 let dst_send = SendPtr(dst_ptr);
2428 let a_send = SendPtr(a_ptr as *mut A);
2429 let b_send = SendPtr(b_ptr as *mut B);
2430 let c_send = SendPtr(c_ptr as *mut C);
2431 let e_send = SendPtr(e_ptr as *mut E);
2432
2433 let costs = compute_costs(&ordered_strides);
2434 let initial_offsets = vec![0isize; strides_list.len()];
2435 return mapreduce_threaded(
2436 &fused_dims,
2437 &plan.block,
2438 &ordered_strides,
2439 &initial_offsets,
2440 &costs,
2441 nthreads,
2442 0,
2443 1,
2444 &|dims, blocks, strides_list, offsets| {
2445 for_each_inner_block_with_offsets(
2446 dims,
2447 blocks,
2448 strides_list,
2449 offsets,
2450 |offsets, len, strides| {
2451 let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
2452 let ap = unsafe { a_send.as_const().offset(offsets[1]) };
2453 let bp = unsafe { b_send.as_const().offset(offsets[2]) };
2454 let cp = unsafe { c_send.as_const().offset(offsets[3]) };
2455 let ep = unsafe { e_send.as_const().offset(offsets[4]) };
2456 unsafe {
2457 inner_loop_map4::<D, A, B, C, E, OpA, OpB, OpC, OpE>(
2458 dp, strides[0], ap, strides[1], bp, strides[2], cp, strides[3],
2459 ep, strides[4], len, &f,
2460 )
2461 };
2462 Ok(())
2463 },
2464 )
2465 },
2466 );
2467 }
2468 }
2469
2470 let initial_offsets = vec![0isize; ordered_strides.len()];
2471 for_each_inner_block_preordered(
2472 &fused_dims,
2473 &plan.block,
2474 &ordered_strides,
2475 &initial_offsets,
2476 |offsets, len, strides| {
2477 let dp = unsafe { dst_ptr.offset(offsets[0]) };
2478 let ap = unsafe { a_ptr.offset(offsets[1]) };
2479 let bp = unsafe { b_ptr.offset(offsets[2]) };
2480 let cp = unsafe { c_ptr.offset(offsets[3]) };
2481 let ep = unsafe { e_ptr.offset(offsets[4]) };
2482 unsafe {
2483 inner_loop_map4::<D, A, B, C, E, OpA, OpB, OpC, OpE>(
2484 dp, strides[0], ap, strides[1], bp, strides[2], cp, strides[3], ep, strides[4],
2485 len, &f,
2486 )
2487 };
2488 Ok(())
2489 },
2490 )
2491}
2492
2493#[cfg(test)]
2494#[path = "map_view/tests/scalar_branch_tests.rs"]
2495mod scalar_branch_tests;
2496
2497#[cfg(all(test, feature = "parallel"))]
2498#[path = "map_view/tests/tests.rs"]
2499mod tests;