strided_perm/
copy.rs

1//! Copy/permutation operations on strided views.
2
3use crate::fuse::FusionPlanError;
4#[cfg(feature = "parallel")]
5use crate::hptt::execute_permute_blocked_par;
6use crate::hptt::{build_permute_plan, execute_permute_blocked};
7use strided_view::{Result, StridedError, StridedView, StridedViewMut};
8
9/// Map a permutation planning failure onto the shared view error vocabulary.
10///
11/// Rank and shape are validated before planning, so a length mismatch cannot
12/// reach this point; it still maps to a typed error instead of a panic. An
13/// oversized (typically stride-0 broadcast) extent whose fused product cannot
14/// be represented is reported as an offset overflow (issue #263).
15fn plan_error(err: FusionPlanError) -> StridedError {
16    match err {
17        FusionPlanError::LengthMismatch { .. } => StridedError::StrideLengthMismatch,
18        FusionPlanError::DimensionOverflow => StridedError::OffsetOverflow,
19    }
20}
21
22fn total_len(dims: &[usize]) -> usize {
23    dims.iter().product()
24}
25
26/// Check if all strides indicate contiguous column-major or row-major layout.
27fn is_both_contiguous(dims: &[usize], dst_strides: &[isize], src_strides: &[isize]) -> bool {
28    if dims.is_empty() {
29        return true;
30    }
31
32    // Check col-major for both
33    let mut expected = 1isize;
34    let mut col_ok = true;
35    for (&d, (&ds, &ss)) in dims.iter().zip(dst_strides.iter().zip(src_strides.iter())) {
36        if d <= 1 {
37            continue;
38        }
39        if ds != expected || ss != expected {
40            col_ok = false;
41            break;
42        }
43        expected = expected.saturating_mul(d as isize);
44    }
45    if col_ok {
46        return true;
47    }
48
49    // Check row-major for both
50    let mut expected = 1isize;
51    let mut row_ok = true;
52    for (&d, (&ds, &ss)) in dims
53        .iter()
54        .rev()
55        .zip(dst_strides.iter().rev().zip(src_strides.iter().rev()))
56    {
57        if d <= 1 {
58            continue;
59        }
60        if ds != expected || ss != expected {
61            row_ok = false;
62            break;
63        }
64        expected = expected.saturating_mul(d as isize);
65    }
66    row_ok
67}
68
69/// Copy elements from source to destination: `dest[i] = src[i]`.
70///
71/// Uses HPTT-inspired blocked permutation with bilateral dimension fusion,
72/// cache-aware blocking, and optimal loop ordering.
73///
74/// This is a simple copy without ElementOp support. For copies with
75/// element operations (conj, transpose, etc.), use `strided_kernel::copy_into`.
76pub fn copy_into<T: Copy>(dest: &mut StridedViewMut<T>, src: &StridedView<T>) -> Result<()> {
77    let dst_dims = dest.dims();
78    let src_dims = src.dims();
79    if dst_dims.len() != src_dims.len() {
80        return Err(StridedError::RankMismatch(dst_dims.len(), src_dims.len()));
81    }
82    if dst_dims != src_dims {
83        return Err(StridedError::ShapeMismatch(
84            dst_dims.to_vec(),
85            src_dims.to_vec(),
86        ));
87    }
88    // A zero-sized view has nothing to copy. It validates without its other
89    // extents or strides forming an in-range product, so neither the
90    // contiguous fast path nor the permutation planner may see it.
91    if dst_dims.contains(&0) {
92        return Ok(());
93    }
94
95    let dst_ptr = dest.as_mut_ptr();
96    let src_ptr = src.ptr();
97    let dst_strides = dest.strides();
98    let src_strides = src.strides();
99
100    // Fast path: both contiguous
101    if is_both_contiguous(dst_dims, dst_strides, src_strides) {
102        let len = total_len(dst_dims);
103        unsafe { std::ptr::copy_nonoverlapping(src_ptr, dst_ptr, len) };
104        return Ok(());
105    }
106
107    // HPTT-inspired blocked permutation
108    let elem_size = std::mem::size_of::<T>();
109    let plan =
110        build_permute_plan(dst_dims, src_strides, dst_strides, elem_size).map_err(plan_error)?;
111    unsafe {
112        execute_permute_blocked(src_ptr, dst_ptr, &plan);
113    }
114    Ok(())
115}
116
117/// Copy elements to a col-major destination.
118///
119/// This now delegates to the same HPTT-inspired blocked permutation as
120/// `copy_into`. The HPTT planner automatically handles the col-major
121/// destination case optimally.
122pub fn copy_into_col_major<T: Copy>(
123    dst: &mut StridedViewMut<T>,
124    src: &StridedView<T>,
125) -> Result<()> {
126    copy_into(dst, src)
127}
128
129/// Conservative sufficient proof that distinct destination elements are
130/// distinct offsets.
131///
132/// Axes with extent > 1 are visited in ascending stride. An axis is separated
133/// when its positive stride lands strictly past the offset span already covered
134/// by the smaller-stride axes, so the two offset sets cannot meet. A negative,
135/// zero, or interleaved stride is left unresolved: the parallel entry then
136/// keeps the serial engine rather than let two workers write one offset. A
137/// layout that fails this test can still be injective; only the parallel fast
138/// path is withheld.
139fn destination_offsets_are_separated(dims: &[usize], strides: &[isize]) -> bool {
140    let mut axes: Vec<(isize, usize)> = dims
141        .iter()
142        .copied()
143        .zip(strides.iter().copied())
144        .filter(|&(dim, _)| dim > 1)
145        .map(|(dim, stride)| (stride, dim))
146        .collect();
147    axes.sort_by_key(|&(stride, _)| stride.unsigned_abs());
148    let mut span: u128 = 0;
149    for (stride, dim) in axes {
150        if stride <= 0 || stride as u128 <= span {
151            return false;
152        }
153        span = span.saturating_add((dim as u128 - 1).saturating_mul(stride as u128));
154    }
155    true
156}
157
158/// Copy elements from source to destination with Rayon parallelism.
159///
160/// Parallelizes the outermost block loop of the HPTT-inspired permutation.
161/// Falls back to single-threaded for small tensors.
162#[cfg(feature = "parallel")]
163pub fn copy_into_par<T: Copy + Send + Sync>(
164    dest: &mut StridedViewMut<T>,
165    src: &StridedView<T>,
166) -> Result<()> {
167    let dst_dims = dest.dims();
168    let src_dims = src.dims();
169    if dst_dims.len() != src_dims.len() {
170        return Err(StridedError::RankMismatch(dst_dims.len(), src_dims.len()));
171    }
172    if dst_dims != src_dims {
173        return Err(StridedError::ShapeMismatch(
174            dst_dims.to_vec(),
175            src_dims.to_vec(),
176        ));
177    }
178    // A zero-sized view has nothing to copy. It validates without its other
179    // extents or strides forming an in-range product, so neither the
180    // contiguous fast path nor the permutation planner may see it.
181    if dst_dims.contains(&0) {
182        return Ok(());
183    }
184
185    let dst_ptr = dest.as_mut_ptr();
186    let src_ptr = src.ptr();
187    let dst_strides = dest.strides();
188    let src_strides = src.strides();
189
190    if !destination_offsets_are_separated(dst_dims, dst_strides) {
191        // Each tile is written by one worker; an unresolved destination layout
192        // could map two tiles onto one offset, which is a data race under
193        // parallel execution. Keep the serial engine for that case.
194        return copy_into(dest, src);
195    }
196
197    // Fast path: both contiguous (memcpy is already bandwidth-limited)
198    if is_both_contiguous(dst_dims, dst_strides, src_strides) {
199        let len = total_len(dst_dims);
200        unsafe { std::ptr::copy_nonoverlapping(src_ptr, dst_ptr, len) };
201        return Ok(());
202    }
203
204    let elem_size = std::mem::size_of::<T>();
205    let plan =
206        build_permute_plan(dst_dims, src_strides, dst_strides, elem_size).map_err(plan_error)?;
207    unsafe {
208        execute_permute_blocked_par(src_ptr, dst_ptr, &plan);
209    }
210    Ok(())
211}
212
213/// Copy elements to a col-major destination with Rayon parallelism.
214#[cfg(feature = "parallel")]
215pub fn copy_into_col_major_par<T: Copy + Send + Sync>(
216    dst: &mut StridedViewMut<T>,
217    src: &StridedView<T>,
218) -> Result<()> {
219    copy_into_par(dst, src)
220}
221
222/// Try to fuse a contiguous dimension group into a single (total_size, innermost_stride).
223///
224/// For the group to be fusable, consecutive dimensions must have contiguous strides.
225/// Returns `None` if strides are not contiguous within the group.
226#[cfg(test)]
227fn try_fuse_group(dims: &[usize], strides: &[isize]) -> Option<(usize, isize)> {
228    match dims.len() {
229        0 => Some((1, 0)),
230        1 => Some((dims[0], strides[0])),
231        _ => {
232            if dims.len() != strides.len() {
233                return None;
234            }
235            for (&d, &s) in dims.iter().zip(strides.iter()) {
236                if d > 1 && s == 0 {
237                    return None;
238                }
239            }
240            let mut base_idx: Option<usize> = None;
241            let mut base_abs = usize::MAX;
242            for (i, (&d, &s)) in dims.iter().zip(strides.iter()).enumerate() {
243                if d <= 1 {
244                    continue;
245                }
246                let abs = s.unsigned_abs();
247                if abs < base_abs {
248                    base_abs = abs;
249                    base_idx = Some(i);
250                }
251            }
252
253            let Some(base) = base_idx else {
254                let stride = *strides
255                    .iter()
256                    .min_by_key(|s| s.unsigned_abs())
257                    .unwrap_or(&0);
258                return Some((dims.iter().product(), stride));
259            };
260
261            let mut used = vec![false; dims.len()];
262            used[base] = true;
263            let mut expected_abs = base_abs.checked_mul(dims[base])?;
264
265            let non_singleton = dims.iter().filter(|&&d| d > 1).count();
266            for _ in 1..non_singleton {
267                let mut next = None;
268                for i in 0..dims.len() {
269                    if used[i] || dims[i] <= 1 {
270                        continue;
271                    }
272                    if strides[i].unsigned_abs() == expected_abs {
273                        next = Some(i);
274                        break;
275                    }
276                }
277                let i = next?;
278                used[i] = true;
279                expected_abs = expected_abs.checked_mul(dims[i])?;
280            }
281
282            Some((dims.iter().product(), strides[base]))
283        }
284    }
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290    use strided_view::StridedArray;
291
292    #[test]
293    fn test_copy_into_contiguous() {
294        let src =
295            StridedArray::<f64>::from_fn_col_major(&[2, 3], |idx| (idx[0] * 10 + idx[1]) as f64);
296        let mut dst = StridedArray::<f64>::col_major(&[2, 3]);
297        copy_into(&mut dst.view_mut(), &src.view()).unwrap();
298        assert_eq!(dst.get(&[0, 0]), 0.0);
299        assert_eq!(dst.get(&[1, 2]), 12.0);
300    }
301
302    #[test]
303    fn test_copy_into_transposed() {
304        // src is row-major [3,2], dst is col-major [3,2]
305        let src = StridedArray::<f64>::from_parts(
306            vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0],
307            &[3, 2],
308            &[2, 1], // row-major
309            0,
310        )
311        .unwrap();
312        let mut dst = StridedArray::<f64>::col_major(&[3, 2]);
313        copy_into(&mut dst.view_mut(), &src.view()).unwrap();
314        assert_eq!(dst.get(&[0, 0]), 0.0);
315        assert_eq!(dst.get(&[0, 1]), 1.0);
316        assert_eq!(dst.get(&[1, 0]), 2.0);
317        assert_eq!(dst.get(&[2, 1]), 5.0);
318    }
319
320    #[test]
321    fn test_copy_into_col_major_basic() {
322        let src =
323            StridedArray::<f64>::from_fn_col_major(&[4, 3], |idx| (idx[0] * 10 + idx[1]) as f64);
324        let mut dst = StridedArray::<f64>::col_major(&[4, 3]);
325        copy_into_col_major(&mut dst.view_mut(), &src.view()).unwrap();
326        assert_eq!(dst.get(&[0, 0]), 0.0);
327        assert_eq!(dst.get(&[3, 2]), 32.0);
328    }
329
330    #[test]
331    fn test_copy_into_col_major_permuted_src() {
332        // src has scattered strides, dst is col-major
333        let data: Vec<f64> = (0..24).map(|i| i as f64).collect();
334        let src = StridedArray::<f64>::from_parts(
335            data,
336            &[2, 3, 4],
337            &[12, 4, 1], // row-major
338            0,
339        )
340        .unwrap();
341        let mut dst = StridedArray::<f64>::col_major(&[2, 3, 4]);
342        copy_into_col_major(&mut dst.view_mut(), &src.view()).unwrap();
343        // Verify element-by-element
344        for i in 0..2 {
345            for j in 0..3 {
346                for k in 0..4 {
347                    assert_eq!(
348                        dst.get(&[i, j, k]),
349                        src.get(&[i, j, k]),
350                        "mismatch at [{},{},{}]",
351                        i,
352                        j,
353                        k
354                    );
355                }
356            }
357        }
358    }
359
360    #[test]
361    fn test_copy_into_grouped_high_rank_binary_permutation() {
362        let dims = vec![2usize; 12];
363        let total: usize = dims.iter().product();
364        let src = StridedArray::<u32>::from_fn_col_major(&dims, |idx| {
365            idx.iter()
366                .enumerate()
367                .map(|(axis, &coord)| (coord as u32) << axis)
368                .sum::<u32>()
369        });
370        let perm = vec![2, 4, 5, 6, 0, 3, 7, 1, 8, 9, 10, 11];
371        let src_view = src.view().permute(&perm).unwrap();
372        let mut dst = StridedArray::<u32>::col_major(&dims);
373
374        copy_into(&mut dst.view_mut(), &src_view).unwrap();
375
376        for flat in 0..total {
377            let mut rem = flat;
378            let mut src_offset = 0usize;
379            for &stride in src_view.strides() {
380                let coord = rem & 1;
381                rem >>= 1;
382                src_offset += coord * stride as usize;
383            }
384            assert_eq!(dst.data()[flat], src.data()[src_offset], "flat={flat}");
385        }
386    }
387
388    #[test]
389    fn test_try_fuse_group_empty() {
390        assert_eq!(try_fuse_group(&[], &[]), Some((1, 0)));
391    }
392
393    #[test]
394    fn test_try_fuse_group_single() {
395        assert_eq!(try_fuse_group(&[5], &[2]), Some((5, 2)));
396    }
397
398    #[test]
399    fn test_try_fuse_group_contiguous_row_major() {
400        assert_eq!(try_fuse_group(&[3, 4], &[4, 1]), Some((12, 1)));
401    }
402
403    #[test]
404    fn test_try_fuse_group_contiguous_col_major() {
405        assert_eq!(try_fuse_group(&[3, 4], &[1, 3]), Some((12, 1)));
406    }
407
408    #[test]
409    fn test_try_fuse_group_non_contiguous() {
410        assert_eq!(try_fuse_group(&[3, 4], &[8, 1]), None);
411    }
412
413    #[test]
414    fn destination_separation_rejects_overlapping_strides() {
415        assert!(destination_offsets_are_separated(&[3, 16384], &[1, 3]));
416        // The first axis' offsets (stride 2) interleave the second axis' (stride
417        // 1), so the parallel entry must keep the serial engine.
418        assert!(!destination_offsets_are_separated(&[3, 16384], &[2, 1]));
419        assert!(!destination_offsets_are_separated(&[4], &[-1]));
420        assert!(destination_offsets_are_separated(&[1, 1], &[0, 1]));
421    }
422
423    #[cfg(feature = "parallel")]
424    #[test]
425    fn parallel_transpose_splits_the_full_tile_grid_for_a_short_row_extent() {
426        // The source's unit-stride extent is the short side of the transpose, so
427        // the whole kernel would be a single row tile over `block`, and the total
428        // is above the threading threshold. The parallel entry must split the
429        // column tiles too; a row-only split would leave the work on one worker.
430        // The result is checked element by element.
431        let rows = 20000usize;
432        let cols = 3usize;
433        let data: Vec<f64> = (0..rows * cols).map(|i| i as f64).collect();
434        let src =
435            StridedArray::<f64>::from_parts(data, &[rows, cols], &[cols as isize, 1], 0).unwrap();
436        let pool = rayon::ThreadPoolBuilder::new()
437            .num_threads(2)
438            .build()
439            .unwrap();
440        let mut dst = StridedArray::<f64>::col_major(&[rows, cols]);
441        pool.install(|| copy_into_par(&mut dst.view_mut(), &src.view()).unwrap());
442        for r in 0..rows {
443            for c in 0..cols {
444                assert_eq!(dst.get(&[r, c]), (r * cols + c) as f64);
445            }
446        }
447    }
448
449    #[test]
450    fn test_copy_shape_mismatch() {
451        let src = StridedArray::<f64>::col_major(&[2, 3]);
452        let mut dst = StridedArray::<f64>::col_major(&[3, 2]);
453        let result = copy_into(&mut dst.view_mut(), &src.view());
454        assert!(result.is_err());
455    }
456}
strided_perm/copy.rs

strided_perm/
copy.rs
