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