strided_basic/copy_plan.rs
1//! Prepared (compile-once, execute-many) copy plans over raw strided layouts.
2//!
3//! [`copy_scale_raw`](crate::copy_scale_raw) and friends rebuild the fused
4//! loop nest on every call; for prepared-replay consumers that issue many
5//! small copies with a fixed layout, that per-call planning dominates
6//! (see issue #139). [`CopyPlan`] splits the work: [`CopyPlan::compile`]
7//! validates the layout pair and builds the fused traversal once,
8//! [`CopyPlan::execute`]/[`CopyPlan::execute_scale`]/[`CopyPlan::execute_conj`]
9//! replay it with no planning and no heap allocation for fused ranks at most
10//! [`RAW_FUSED_RANK_LIMIT`](crate::RAW_FUSED_RANK_LIMIT).
11
12use core::{
13 marker::PhantomData,
14 mem::MaybeUninit,
15 ops::{Add, Mul},
16};
17
18use num_traits::{One, Zero};
19
20use crate::raw_ops::{
21 apply_fused_range, apply_spilled_range, fuse_pair_traversal, fused_total, FusedAxis,
22 FusedPairLayout, PairTraversal,
23};
24use crate::{ElementOpApply, Exec, RawStridedMut, RawStridedRef, Result, StridedError};
25
26// Same pattern as map_view.rs / outer_product.rs: inline storage in both
27// feature configurations. Only `compile` touches these; `execute*` never
28// allocates for them. (A spilled fused layout needs one cursor allocation per
29// replayed range, so an `execute*` can allocate once per worker range above the
30// inline rank bound.)
31type AxisVec<T> = crate::inline::InlineVec<T>;
32
33pub(crate) trait OverwriteWriter<T> {
34 fn dims(&self) -> &[usize];
35 fn strides(&self) -> &[isize];
36 fn offset(&self) -> isize;
37 /// # Safety
38 /// The caller must use the pointer only for the validated allocation and
39 /// logical layout represented by this writer.
40 unsafe fn data_ptr(&mut self) -> *mut T;
41 /// # Safety
42 /// The offset must be an in-bounds logical destination proven by layout.
43 unsafe fn write_at(&mut self, offset: isize, value: T);
44}
45
46pub(crate) trait ReadModifyWrite<T>: OverwriteWriter<T> {
47 /// # Safety
48 /// The offset must be an in-bounds initialized slot covered by the
49 /// traversal's copy and disjointness proof.
50 unsafe fn add_at<F>(&mut self, offset: isize, value: T, combine: F)
51 where
52 F: FnOnce(T, T) -> T;
53}
54
55impl<'a, T> OverwriteWriter<T> for RawStridedMut<'a, T> {
56 fn dims(&self) -> &[usize] {
57 self.dims()
58 }
59 fn strides(&self) -> &[isize] {
60 self.strides()
61 }
62 fn offset(&self) -> isize {
63 self.offset()
64 }
65 unsafe fn data_ptr(&mut self) -> *mut T {
66 self.data_mut().as_mut_ptr()
67 }
68 unsafe fn write_at(&mut self, offset: isize, value: T) {
69 // SAFETY: the prepared layout validates every logical destination.
70 unsafe { self.data_mut().as_mut_ptr().offset(offset).write(value) }
71 }
72}
73
74impl<'a, T> ReadModifyWrite<T> for RawStridedMut<'a, T>
75where
76 T: Add<Output = T>,
77{
78 #[inline(always)]
79 unsafe fn add_at<F>(&mut self, offset: isize, value: T, combine: F)
80 where
81 F: FnOnce(T, T) -> T,
82 {
83 // SAFETY: the copy or initialized caller proves this logical slot.
84 unsafe {
85 let ptr = self.data_mut().as_mut_ptr().offset(offset);
86 ptr.write(combine(ptr.read(), value));
87 }
88 }
89}
90
91impl<'a, T> OverwriteWriter<T> for RawStridedMut<'a, MaybeUninit<T>> {
92 fn dims(&self) -> &[usize] {
93 self.dims()
94 }
95 fn strides(&self) -> &[isize] {
96 self.strides()
97 }
98 fn offset(&self) -> isize {
99 self.offset()
100 }
101 unsafe fn data_ptr(&mut self) -> *mut T {
102 self.data_mut().as_mut_ptr().cast()
103 }
104 unsafe fn write_at(&mut self, offset: isize, value: T) {
105 // SAFETY: the prepared layout validates every logical destination.
106 unsafe {
107 self.data_mut()
108 .as_mut_ptr()
109 .offset(offset)
110 .write(MaybeUninit::new(value))
111 }
112 }
113}
114
115pub(crate) struct InitializedRawDest<'a, T> {
116 ptr: *mut T,
117 extent: usize,
118 dims: &'a [usize],
119 strides: &'a [isize],
120 offset: isize,
121 _marker: PhantomData<&'a mut [MaybeUninit<T>]>,
122}
123
124impl<'a, T> OverwriteWriter<T> for InitializedRawDest<'a, T> {
125 fn dims(&self) -> &[usize] {
126 self.dims
127 }
128 fn strides(&self) -> &[isize] {
129 self.strides
130 }
131 fn offset(&self) -> isize {
132 self.offset
133 }
134 unsafe fn data_ptr(&mut self) -> *mut T {
135 self.ptr
136 }
137 unsafe fn write_at(&mut self, offset: isize, value: T) {
138 debug_assert!(offset >= 0 && (offset as usize) < self.extent);
139 // SAFETY: the copy proof and extent check cover this logical slot.
140 unsafe { self.ptr.offset(offset).write(value) }
141 }
142}
143
144impl<'a, T> ReadModifyWrite<T> for InitializedRawDest<'a, T>
145where
146 T: Add<Output = T>,
147{
148 #[inline(always)]
149 unsafe fn add_at<F>(&mut self, offset: isize, value: T, combine: F)
150 where
151 F: FnOnce(T, T) -> T,
152 {
153 debug_assert!(offset >= 0 && (offset as usize) < self.extent);
154 // SAFETY: the copy proof and extent check cover this logical slot.
155 unsafe {
156 let ptr = self.ptr.offset(offset);
157 ptr.write(combine(ptr.read(), value));
158 }
159 }
160}
161
162/// A compiled copy traversal for one `(dims, dst_strides, src_strides)`
163/// layout pair.
164///
165/// `compile` proves the layout facts once (rank agreement, extent overflow,
166/// destination injectivity) and fuses/orders the loop nest; each `execute*`
167/// call then only re-checks the per-call facts (that the supplied views carry
168/// exactly the compiled layout) before replaying the prepared loops.
169///
170/// Overlapping `src`/`dest` memory is not supported, matching the rest of the
171/// crate.
172///
173/// Ranks above [`RAW_FUSED_RANK_LIMIT`](crate::RAW_FUSED_RANK_LIMIT) are supported by the same fused
174/// traversal, which then stores its axes in one rank-bounded allocation.
175///
176/// # Example
177///
178/// ```rust
179/// use strided_basic::{CopyPlan, Exec, RawStridedMut, RawStridedRef};
180///
181/// let dims = [2usize, 3];
182/// let src_strides = [3isize, 1];
183/// let dst_strides = [1isize, 2]; // transposed destination
184/// let plan = CopyPlan::compile(&dims, &dst_strides, &src_strides).unwrap();
185///
186/// let src = [0.0f64, 1.0, 2.0, 10.0, 11.0, 12.0];
187/// let mut dst = [0.0f64; 6];
188/// let src_ref = RawStridedRef::new(&src, &dims, &src_strides, 0).unwrap();
189/// let mut dst_mut = RawStridedMut::new(&mut dst, &dims, &dst_strides, 0).unwrap();
190/// plan.execute(&mut dst_mut, &src_ref, Exec::serial()).unwrap();
191/// assert_eq!(dst, [0.0, 10.0, 1.0, 11.0, 2.0, 12.0]);
192/// ```
193#[derive(Clone, Debug, Default)]
194pub struct CopyPlan {
195 dims: AxisVec<usize>,
196 dst_strides: AxisVec<isize>,
197 src_strides: AxisVec<isize>,
198 fused: PairTraversal,
199}
200
201impl CopyPlan {
202 /// `None` when the fused traversal is stored in the rank-bounded spill.
203 #[inline]
204 pub(crate) fn fused_layout(&self) -> Option<&FusedPairLayout> {
205 self.fused.inline()
206 }
207
208 pub(crate) fn execute_uninit_then<'a, T, R>(
209 &self,
210 dest: &'a mut RawStridedMut<'a, MaybeUninit<T>>,
211 src: &RawStridedRef<'_, T>,
212 mut exec: Exec<'_>,
213 f: impl for<'b, 'e> FnOnce(InitializedRawDest<'b, T>, Exec<'e>) -> R,
214 ) -> Result<R>
215 where
216 T: Copy + Send + Sync,
217 {
218 self.execute_uninit(dest, src, exec.reborrow())?;
219 let data = dest.data_mut();
220 let receipt = InitializedRawDest {
221 ptr: data.as_mut_ptr().cast(),
222 extent: data.len(),
223 dims: dest.dims(),
224 strides: dest.strides(),
225 offset: dest.offset(),
226 _marker: PhantomData,
227 };
228 Ok(f(receipt, exec))
229 }
230
231 /// Compile a copy plan for the given layout pair.
232 ///
233 /// Performs the layout validation and traversal construction
234 /// (fuse + order) once:
235 ///
236 /// - `dims`, `dst_strides`, and `src_strides` must have equal length
237 /// ([`StridedError::StrideLengthMismatch`]);
238 /// - the total element count must not overflow `usize`
239 /// ([`StridedError::OffsetOverflow`]);
240 /// - the destination layout must be injective, i.e. map distinct logical
241 /// indices to distinct offsets
242 /// ([`StridedError::NonInjectiveOutputLayout`]).
243 pub fn compile(dims: &[usize], dst_strides: &[isize], src_strides: &[isize]) -> Result<Self> {
244 if dims.len() != dst_strides.len() || dims.len() != src_strides.len() {
245 return Err(StridedError::StrideLengthMismatch);
246 }
247 // A zero-sized layout has zero elements even when its other extents
248 // overflow; only a nonempty overflowing count is rejected.
249 crate::kernel::total_len(dims)?;
250 if !crate::layout_check::is_injective_layout(dims, dst_strides) {
251 return Err(StridedError::NonInjectiveOutputLayout);
252 }
253 Ok(Self {
254 dims: dims.into(),
255 dst_strides: dst_strides.into(),
256 src_strides: src_strides.into(),
257 fused: fuse_pair_traversal(dims, dst_strides, src_strides),
258 })
259 }
260
261 /// Check the per-call facts: the supplied views must carry exactly the
262 /// compiled layout. Buffer bounds against that layout were already proven
263 /// by [`RawStridedRef::new`]/[`RawStridedMut::new`] (or asserted by the
264 /// caller of the `new_unchecked` constructors), so layout equality is the
265 /// complete precondition for the pointer-based fused replay below.
266 fn check_call<D, S>(
267 &self,
268 dest: &RawStridedMut<'_, D>,
269 src: &RawStridedRef<'_, S>,
270 ) -> Result<()> {
271 check_layout_pair(&self.dims, &self.dst_strides, &self.src_strides, dest, src)
272 }
273
274 /// `dest = src` into a potentially uninitialized destination.
275 ///
276 /// On success every reachable logical destination element is initialized.
277 /// Non-reachable holes in the backing allocation are not written.
278 pub fn execute_uninit<T>(
279 &self,
280 dest: &mut RawStridedMut<'_, MaybeUninit<T>>,
281 src: &RawStridedRef<'_, T>,
282 exec: Exec<'_>,
283 ) -> Result<()>
284 where
285 T: Copy + Send + Sync,
286 {
287 self.check_call(dest, src)?;
288 replay_fused(
289 dest,
290 src,
291 &self.fused,
292 |dst, value| {
293 dst.write(value);
294 },
295 |value| value,
296 exec,
297 );
298 Ok(())
299 }
300
301 /// `dest = src`. Allocation-free for fused ranks at most
302 /// [`RAW_FUSED_RANK_LIMIT`](crate::RAW_FUSED_RANK_LIMIT).
303 pub fn execute<T>(
304 &self,
305 dest: &mut RawStridedMut<'_, T>,
306 src: &RawStridedRef<'_, T>,
307 exec: Exec<'_>,
308 ) -> Result<()>
309 where
310 T: Copy + Send + Sync,
311 {
312 self.check_call(dest, src)?;
313 replay_fused(
314 dest,
315 src,
316 &self.fused,
317 |dst, value| *dst = value,
318 |value: T| value,
319 exec,
320 );
321 Ok(())
322 }
323
324 /// `dest = scale * src`. Allocation-free for fused ranks at most
325 /// [`RAW_FUSED_RANK_LIMIT`](crate::RAW_FUSED_RANK_LIMIT).
326 pub fn execute_scale<T>(
327 &self,
328 dest: &mut RawStridedMut<'_, T>,
329 src: &RawStridedRef<'_, T>,
330 scale: T,
331 exec: Exec<'_>,
332 ) -> Result<()>
333 where
334 T: Copy + Mul<T, Output = T> + Send + Sync,
335 {
336 self.check_call(dest, src)?;
337 replay_fused(
338 dest,
339 src,
340 &self.fused,
341 |dst, value| *dst = value,
342 |value: T| scale * value,
343 exec,
344 );
345 Ok(())
346 }
347
348 /// `dest = conj(src)`. Allocation-free for fused ranks at most
349 /// [`RAW_FUSED_RANK_LIMIT`](crate::RAW_FUSED_RANK_LIMIT).
350 pub fn execute_conj<T>(
351 &self,
352 dest: &mut RawStridedMut<'_, T>,
353 src: &RawStridedRef<'_, T>,
354 exec: Exec<'_>,
355 ) -> Result<()>
356 where
357 T: Copy + ElementOpApply + Send + Sync,
358 {
359 self.check_call(dest, src)?;
360 replay_fused(
361 dest,
362 src,
363 &self.fused,
364 |dst, value| *dst = value,
365 |value: T| value.conj(),
366 exec,
367 );
368 Ok(())
369 }
370}
371
372/// Verify that a pair of raw views carries exactly a compiled layout.
373///
374/// Shared by [`CopyPlan`] and [`UpdatePlan`], whose `compile` methods prove the
375/// layout facts once. Buffer bounds against that layout were already proven by
376/// [`RawStridedRef::new`]/[`RawStridedMut::new`] (or asserted by the caller of
377/// the `new_unchecked` constructors), so layout equality is the complete
378/// per-call precondition for the pointer-based fused replay.
379fn check_layout_pair<D, S>(
380 dims: &[usize],
381 dst_strides: &[isize],
382 src_strides: &[isize],
383 dest: &RawStridedMut<'_, D>,
384 src: &RawStridedRef<'_, S>,
385) -> Result<()> {
386 if dest.dims() != dims
387 || src.dims() != dims
388 || dest.strides() != dst_strides
389 || src.strides() != src_strides
390 {
391 return Err(StridedError::PlanLayoutMismatch);
392 }
393 Ok(())
394}
395
396/// A compiled read-modify-write traversal for one `(dims, dst_strides,
397/// src_strides)` layout pair.
398///
399/// [`execute`](Self::execute) computes `dest = alpha * f(src) + beta * dest` by
400/// replaying the same fused pair traversal as [`CopyPlan`]: `compile` proves
401/// the layout facts once (rank agreement, extent overflow, destination
402/// injectivity) and each call only re-checks that the supplied views carry the
403/// compiled layout before running monomorphized loops. Execute is
404/// allocation-free for fused ranks at most
405/// [`RAW_FUSED_RANK_LIMIT`](crate::RAW_FUSED_RANK_LIMIT), and the injective
406/// destination makes a large replay parallel-safe exactly as in [`CopyPlan`].
407///
408/// Overlapping `src`/`dest` memory is not supported, matching the rest of the
409/// crate.
410///
411/// # Example
412///
413/// ```rust
414/// use strided_basic::{Exec, RawStridedMut, RawStridedRef, UpdatePlan};
415///
416/// let dims = [2usize, 2];
417/// let dst_strides = [1isize, 2];
418/// let src_strides = [1isize, 2];
419/// let plan = UpdatePlan::compile(&dims, &dst_strides, &src_strides).unwrap();
420///
421/// let src = [1.0f64, 2.0, 3.0, 4.0];
422/// let mut dst = [10.0f64, 20.0, 30.0, 40.0];
423/// let src_ref = RawStridedRef::new(&src, &dims, &src_strides, 0).unwrap();
424/// let mut dst_mut = RawStridedMut::new(&mut dst, &dims, &dst_strides, 0).unwrap();
425/// // dst = 2 * src + 3 * dst
426/// plan.execute(&mut dst_mut, &src_ref, 2.0, 3.0, |x| x, Exec::serial()).unwrap();
427/// assert_eq!(dst, [32.0, 64.0, 96.0, 128.0]);
428/// ```
429#[derive(Clone, Debug)]
430pub struct UpdatePlan {
431 dims: AxisVec<usize>,
432 dst_strides: AxisVec<isize>,
433 src_strides: AxisVec<isize>,
434 fused: PairTraversal,
435}
436
437impl UpdatePlan {
438 /// Compile an update plan for the given layout pair.
439 ///
440 /// Performs the same layout validation as [`CopyPlan::compile`]:
441 ///
442 /// - `dims`, `dst_strides`, and `src_strides` must have equal length
443 /// ([`StridedError::StrideLengthMismatch`]);
444 /// - the total element count must not overflow `usize`
445 /// ([`StridedError::OffsetOverflow`]);
446 /// - the destination layout must be injective
447 /// ([`StridedError::NonInjectiveOutputLayout`]), so two logical indices
448 /// never name one slot and a parallel replay writes each slot once.
449 pub fn compile(dims: &[usize], dst_strides: &[isize], src_strides: &[isize]) -> Result<Self> {
450 if dims.len() != dst_strides.len() || dims.len() != src_strides.len() {
451 return Err(StridedError::StrideLengthMismatch);
452 }
453 crate::kernel::total_len(dims)?;
454 if !crate::layout_check::is_injective_layout(dims, dst_strides) {
455 return Err(StridedError::NonInjectiveOutputLayout);
456 }
457 Ok(Self {
458 dims: dims.into(),
459 dst_strides: dst_strides.into(),
460 src_strides: src_strides.into(),
461 fused: fuse_pair_traversal(dims, dst_strides, src_strides),
462 })
463 }
464
465 /// `dest = alpha * f(src) + beta * dest`.
466 ///
467 /// `f` is applied to the source element before scaling. When `beta` is
468 /// zero (`-0.0` included; `NaN` is not zero) the replay is write-only: the
469 /// old destination element is never read, so a NaN-filled or uninitialized
470 /// destination cannot contaminate the result, matching BLAS `beta = 0`
471 /// semantics. Otherwise the old destination element is read through the
472 /// destination layout. Allocation-free for fused ranks at most
473 /// [`RAW_FUSED_RANK_LIMIT`](crate::RAW_FUSED_RANK_LIMIT).
474 ///
475 /// # Errors
476 ///
477 /// [`StridedError::PlanLayoutMismatch`] when the supplied views do not
478 /// carry the compiled layout, before any write.
479 pub fn execute<T, F>(
480 &self,
481 dest: &mut RawStridedMut<'_, T>,
482 src: &RawStridedRef<'_, T>,
483 alpha: T,
484 beta: T,
485 f: F,
486 exec: Exec<'_>,
487 ) -> Result<()>
488 where
489 T: Copy + Send + Sync + Add<Output = T> + Mul<Output = T> + Zero + One + PartialEq,
490 F: Fn(T) -> T + Send + Sync,
491 {
492 check_layout_pair(&self.dims, &self.dst_strides, &self.src_strides, dest, src)?;
493 if beta.is_zero() {
494 // Write-only: never form a read of the old destination.
495 if alpha.is_one() {
496 replay_fused(dest, src, &self.fused, |dst, value| *dst = value, f, exec);
497 } else {
498 replay_fused(
499 dest,
500 src,
501 &self.fused,
502 |dst, value| *dst = alpha * value,
503 f,
504 exec,
505 );
506 }
507 } else {
508 replay_fused(
509 dest,
510 src,
511 &self.fused,
512 |dst, value| *dst = alpha * value + beta * *dst,
513 f,
514 exec,
515 );
516 }
517 Ok(())
518 }
519
520 /// `dest = alpha * f(src)` into a potentially uninitialized destination.
521 ///
522 /// This is [`execute`](Self::execute) with `beta = 0`: the destination is
523 /// never read. On success every reachable logical destination element is
524 /// initialized; non-reachable holes in the backing allocation are not
525 /// written. Allocation-free for fused ranks at most
526 /// [`RAW_FUSED_RANK_LIMIT`](crate::RAW_FUSED_RANK_LIMIT).
527 ///
528 /// # Errors
529 ///
530 /// [`StridedError::PlanLayoutMismatch`] when the supplied views do not
531 /// carry the compiled layout, before any write.
532 ///
533 /// # Examples
534 ///
535 /// ```rust
536 /// use core::mem::MaybeUninit;
537 /// use strided_basic::{Exec, RawStridedMut, RawStridedRef, UpdatePlan};
538 ///
539 /// let dims = [2usize, 2];
540 /// let strides = [1isize, 2];
541 /// let plan = UpdatePlan::compile(&dims, &strides, &strides).unwrap();
542 /// let src = [1.0f64, 2.0, 3.0, 4.0];
543 /// let mut storage = [MaybeUninit::<f64>::uninit(); 4];
544 /// let mut dst = RawStridedMut::new(&mut storage, &dims, &strides, 0).unwrap();
545 /// let source = RawStridedRef::new(&src, &dims, &strides, 0).unwrap();
546 /// // dst = 2 * src
547 /// plan.execute_uninit(&mut dst, &source, 2.0, |x| x, Exec::serial()).unwrap();
548 /// drop(dst);
549 /// // SAFETY: every reachable destination slot was initialized above.
550 /// let values: Vec<f64> = storage.iter().map(|v| unsafe { v.assume_init() }).collect();
551 /// assert_eq!(values, [2.0, 4.0, 6.0, 8.0]);
552 /// ```
553 pub fn execute_uninit<T, F>(
554 &self,
555 dest: &mut RawStridedMut<'_, MaybeUninit<T>>,
556 src: &RawStridedRef<'_, T>,
557 alpha: T,
558 f: F,
559 exec: Exec<'_>,
560 ) -> Result<()>
561 where
562 T: Copy + Send + Sync + Mul<Output = T> + One + PartialEq,
563 F: Fn(T) -> T + Send + Sync,
564 {
565 check_layout_pair(&self.dims, &self.dst_strides, &self.src_strides, dest, src)?;
566 if alpha.is_one() {
567 replay_fused(
568 dest,
569 src,
570 &self.fused,
571 |dst, value| {
572 dst.write(value);
573 },
574 f,
575 exec,
576 );
577 } else {
578 replay_fused(
579 dest,
580 src,
581 &self.fused,
582 |dst, value| {
583 dst.write(alpha * value);
584 },
585 f,
586 exec,
587 );
588 }
589 Ok(())
590 }
591}
592
593/// Replay a compiled fused layout, splitting it across the budget of `exec`
594/// when the repository threshold allows it.
595///
596/// Parallel replay is sound only because [`CopyPlan::compile`] proved the
597/// destination layout injective: disjoint logical ranges then write disjoint
598/// destination slots. The logical index space `0..total` (column-major over
599/// the fused axes) is split into contiguous worker ranges, which divides the
600/// outer fused axes and, when there are fewer outer positions than workers,
601/// also chunks long inner runs (for example a rank-1 contiguous copy).
602fn replay_fused<D, S, Apply, Op>(
603 dest: &mut RawStridedMut<'_, D>,
604 src: &RawStridedRef<'_, S>,
605 traversal: &PairTraversal,
606 apply: Apply,
607 op: Op,
608 exec: Exec<'_>,
609) where
610 D: Copy + Send + Sync,
611 S: Copy + Send + Sync,
612 Apply: Fn(&mut D, S) + Send + Sync,
613 Op: Fn(S) -> S + Send + Sync,
614{
615 let src_ptr = src.data().as_ptr();
616 let src_base = src.offset();
617 let dst_base = dest.offset();
618 let dst_ptr = dest.data_mut().as_mut_ptr();
619 // SAFETY: `check_call` proved both views carry exactly the compiled
620 // layout, and `RawStridedRef`/`RawStridedMut` guarantee every offset
621 // reachable through that layout lies inside their data. `traversal` is the
622 // fusion of that layout, so every logical index in `0..total` is an
623 // in-bounds slot of both. The compiled destination layout is injective
624 // and the exclusive destination borrow cannot overlap the shared source
625 // borrow.
626 unsafe {
627 match traversal {
628 PairTraversal::Inline(layout) => replay_fused_raw(
629 dst_ptr, dst_base, src_ptr, src_base, layout, apply, op, exec,
630 ),
631 PairTraversal::Spilled(axes) => {
632 replay_spilled_raw(dst_ptr, dst_base, src_ptr, src_base, axes, apply, op, exec)
633 }
634 }
635 }
636}
637
638/// Pointer-level fused replay shared by [`CopyPlan`] and the dynamic slice
639/// plans, which replay a fused window layout at a runtime base offset.
640///
641/// Splits `0..total` across the budget of `exec` when the repository threshold
642/// allows more than one slot, otherwise runs the serial kernel directly.
643///
644/// # Safety
645///
646/// Every logical index of `layout`, mapped through its destination strides
647/// from `dst_base` and its source strides from `src_base`, must be an
648/// in-bounds slot of the destination and source allocations. The
649/// destination strides must be injective over `layout`, the destination
650/// must not overlap the source, and nothing else may access the destination
651/// slots during the call.
652#[allow(clippy::too_many_arguments)]
653pub(crate) unsafe fn replay_fused_raw<D, S, Apply, Op>(
654 dst_ptr: *mut D,
655 dst_base: isize,
656 src_ptr: *const S,
657 src_base: isize,
658 layout: &FusedPairLayout,
659 apply: Apply,
660 op: Op,
661 exec: Exec<'_>,
662) where
663 D: Copy + Send + Sync,
664 S: Copy + Send + Sync,
665 Apply: Fn(&mut D, S) + Send + Sync,
666 Op: Fn(S) -> S + Send + Sync,
667{
668 let total = fused_total(layout);
669 if total == 0 {
670 return;
671 }
672 let exec = crate::threading::for_len(exec, total);
673 if !exec.is_serial() {
674 let dst_ptr = crate::threading::SendPtr(dst_ptr);
675 let src_ptr = crate::threading::SendPtr(src_ptr as *mut S);
676 crate::threading::parallel_for_each(0..total, exec, &|range| {
677 // SAFETY: the caller contract makes every index in bounds; the
678 // worker ranges are disjoint and the destination strides
679 // injective, so no two workers write the same slot, and the
680 // source is only read.
681 unsafe {
682 apply_fused_range(
683 dst_ptr.as_ptr(),
684 dst_base,
685 src_ptr.as_const(),
686 src_base,
687 layout,
688 range.start,
689 range.len(),
690 &apply,
691 &op,
692 );
693 }
694 });
695 return;
696 }
697 // SAFETY: forwarded caller contract over the full range `0..total`.
698 unsafe {
699 apply_fused_range(
700 dst_ptr, dst_base, src_ptr, src_base, layout, 0, total, &apply, &op,
701 );
702 }
703}
704
705/// Pointer-level replay of a spilled axis list.
706///
707/// Mirrors [`replay_fused_raw`] including its worker split; kept separate so
708/// that the inline replay keeps its fixed-rank shape.
709///
710/// # Safety
711///
712/// Every logical index of `axes`, mapped through its destination strides from
713/// `dst_base` and its source strides from `src_base`, must be an in-bounds slot
714/// of the destination and source allocations. The destination strides must be
715/// injective over `axes`, the destination must not overlap the source, and
716/// nothing else may access the destination slots during the call.
717#[allow(clippy::too_many_arguments)]
718pub(crate) unsafe fn replay_spilled_raw<D, S, Apply, Op>(
719 dst_ptr: *mut D,
720 dst_base: isize,
721 src_ptr: *const S,
722 src_base: isize,
723 axes: &[FusedAxis],
724 apply: Apply,
725 op: Op,
726 exec: Exec<'_>,
727) where
728 D: Copy + Send + Sync,
729 S: Copy + Send + Sync,
730 Apply: Fn(&mut D, S) + Send + Sync,
731 Op: Fn(S) -> S + Send + Sync,
732{
733 let total = axes.iter().map(|axis| axis.dim).product::<usize>();
734 if total == 0 {
735 return;
736 }
737 let exec = crate::threading::for_len(exec, total);
738 if !exec.is_serial() {
739 let dst_ptr = crate::threading::SendPtr(dst_ptr);
740 let src_ptr = crate::threading::SendPtr(src_ptr as *mut S);
741 crate::threading::parallel_for_each(0..total, exec, &|range| {
742 // SAFETY: the caller contract makes every index in bounds; the
743 // worker ranges are disjoint and the destination strides
744 // injective, so no two workers write the same slot, and the
745 // source is only read.
746 unsafe {
747 apply_spilled_range(
748 dst_ptr.as_ptr(),
749 dst_base,
750 src_ptr.as_const(),
751 src_base,
752 axes,
753 range.start,
754 range.len(),
755 &apply,
756 &op,
757 );
758 }
759 });
760 return;
761 }
762 // SAFETY: forwarded caller contract over the full range `0..total`.
763 unsafe {
764 apply_spilled_range(
765 dst_ptr, dst_base, src_ptr, src_base, axes, 0, total, &apply, &op,
766 );
767 }
768}
769
770#[cfg(test)]
771#[path = "copy_plan/tests/tests.rs"]
772mod tests;
773
774#[cfg(test)]
775#[path = "copy_plan/tests/parallel_tests.rs"]
776mod parallel_tests;