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