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