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