strided_perm/

copy.rs

//! Copy/permutation operations on strided views.

use crate::fuse::FusionPlanError;
use crate::hptt::{build_permute_plan, execute_permute_blocked};
use strided_view::{Result, StridedError, StridedView, StridedViewMut};
use threadteam::Exec;

/// Map a permutation planning failure onto the shared view error vocabulary.
///
/// Rank and shape are validated before planning, so a length mismatch cannot
/// reach this point; it still maps to a typed error instead of a panic. An
/// oversized (typically stride-0 broadcast) extent whose fused product cannot
/// be represented is reported as an offset overflow (issue #263).
fn plan_error(err: FusionPlanError) -> StridedError {
    match err {
        FusionPlanError::LengthMismatch { .. } => StridedError::StrideLengthMismatch,
        FusionPlanError::DimensionOverflow => StridedError::OffsetOverflow,
    }
}

fn total_len(dims: &[usize]) -> usize {
    dims.iter().product()
}

/// Check if all strides indicate contiguous column-major or row-major layout.
fn is_both_contiguous(dims: &[usize], dst_strides: &[isize], src_strides: &[isize]) -> bool {
    if dims.is_empty() {
        return true;
    }

    // Check col-major for both
    let mut expected = 1isize;
    let mut col_ok = true;
    for (&d, (&ds, &ss)) in dims.iter().zip(dst_strides.iter().zip(src_strides.iter())) {
        if d <= 1 {
            continue;
        }
        if ds != expected || ss != expected {
            col_ok = false;
            break;
        }
        expected = expected.saturating_mul(d as isize);
    }
    if col_ok {
        return true;
    }

    // Check row-major for both
    let mut expected = 1isize;
    let mut row_ok = true;
    for (&d, (&ds, &ss)) in dims
        .iter()
        .rev()
        .zip(dst_strides.iter().rev().zip(src_strides.iter().rev()))
    {
        if d <= 1 {
            continue;
        }
        if ds != expected || ss != expected {
            row_ok = false;
            break;
        }
        expected = expected.saturating_mul(d as isize);
    }
    row_ok
}

/// The validated raw geometry of one copy: the shared dims, both stride
/// runs and the two origins. `U` is the destination element type: `T` for
/// an initialized destination, `MaybeUninit<T>` for a fresh one.
struct RawCopy<'a, T, U> {
    dims: &'a [usize],
    dst_strides: &'a [isize],
    src_strides: &'a [isize],
    dst: *mut U,
    src: *const T,
}

/// Validates rank and shape and extracts the raw geometry; `None` for a
/// zero-sized copy, which has nothing to write.
///
/// A zero-sized view validates without its other extents or strides forming
/// an in-range product, so neither the contiguous fast path nor the
/// permutation planner may see it.
fn prepare<'a, T, U>(
    dest: &'a mut StridedViewMut<U>,
    src: &'a StridedView<T>,
) -> Result<Option<RawCopy<'a, T, U>>> {
    let dst_dims = dest.dims();
    let src_dims = src.dims();
    if dst_dims.len() != src_dims.len() {
        return Err(StridedError::RankMismatch(dst_dims.len(), src_dims.len()));
    }
    if dst_dims != src_dims {
        return Err(StridedError::ShapeMismatch(
            dst_dims.to_vec(),
            src_dims.to_vec(),
        ));
    }
    if dst_dims.contains(&0) {
        return Ok(None);
    }
    Ok(Some(RawCopy {
        dims: dst_dims,
        dst_strides: dest.strides(),
        src_strides: src.strides(),
        dst: dest.as_mut_ptr(),
        src: src.ptr(),
    }))
}

/// The copy over validated raw geometry, run within the budget of `exec`.
///
/// # Safety
///
/// `raw.dst` must be valid for writes of `T` at every destination offset and
/// `raw.src` valid for reads of initialized `T` at every source offset; the
/// two ranges must not overlap. Every store goes through raw pointers, so the
/// destination may be uninitialized storage.
unsafe fn copy_raw<T: Copy + Send + Sync + 'static, U>(
    raw: RawCopy<'_, T, U>,
    exec: Exec<'_>,
) -> Result<()> {
    let dst_ptr = raw.dst.cast::<T>();
    // The width is the token's decision: the serial time of the copy against the
    // dispatch cost of the one section the engine opens (`for_ranges`).
    let exec = match raw.dims.iter().try_fold(1usize, |n, &d| n.checked_mul(d)) {
        Some(n) => {
            let estimate =
                std::time::Duration::from_secs_f64(n as f64 * crate::copy_ns_per_element() * 1e-9);
            let width = exec.width_for(estimate, 1);
            exec.limit(width)
        }
        // An overflowing count is reported by planning below; run it serially.
        None => exec.limit(std::num::NonZeroUsize::MIN),
    };
    // Fast path: both contiguous. A contiguous layout is injective by
    // construction, so it skips the validation below.
    if is_both_contiguous(raw.dims, raw.dst_strides, raw.src_strides) {
        let len = total_len(raw.dims);
        if exec.is_serial() {
            // SAFETY: the validated contiguous ranges are live and non-overlapping.
            unsafe { std::ptr::copy_nonoverlapping(raw.src, dst_ptr, len) };
        } else {
            let src_addr = raw.src as usize;
            let dst_addr = dst_ptr as usize;
            exec.for_ranges(len, std::num::NonZeroUsize::MIN, |range| {
                // SAFETY: `for_ranges` joins disjoint subranges of the validated
                // contiguous buffers; raw stores also support uninitialized output.
                unsafe {
                    std::ptr::copy_nonoverlapping(
                        (src_addr as *const T).add(range.start),
                        (dst_addr as *mut T).add(range.start),
                        range.len(),
                    );
                }
            });
        }
        return Ok(());
    }

    // HPTT-inspired blocked permutation. Planning runs before the injectivity
    // check so an unrepresentable layout keeps reporting `OffsetOverflow`
    // (issue #263) even when its destination is also non-injective; planning
    // only builds metadata and writes nothing.
    let elem_size = std::mem::size_of::<T>();
    let plan = build_permute_plan(raw.dims, raw.src_strides, raw.dst_strides, elem_size)
        .map_err(plan_error)?;
    // A non-injective destination maps distinct logical elements onto one
    // offset, so the unsafe write loop would produce overlapping writes (and a
    // data race when `exec` fans out). This matches the sibling output-writing
    // APIs (issue #281). The check is global over the destination, so two
    // worker partitions (disjoint logical positions) can never write one
    // offset; no extra per-partition separation proof is needed.
    if !strided_view::is_injective_layout(raw.dims, raw.dst_strides) {
        return Err(StridedError::NonInjectiveOutputLayout);
    }
    execute_permute_blocked(raw.src, dst_ptr, &plan, exec);
    Ok(())
}

/// Copy elements from source to destination: `dest[i] = src[i]`.
///
/// Uses HPTT-inspired blocked permutation with bilateral dimension fusion,
/// cache-aware blocking, and optimal loop ordering. The permutation is split
/// across the budget of `exec`, at the width the token's dispatch cost
/// justifies for the copy ([`copy_ns_per_element`](crate::copy_ns_per_element)
/// times its elements); a small copy stays on the calling thread. A caller that wants a large copy on one thread passes
/// [`Exec::serial`] or narrows the token with [`Exec::limit`].
/// A contiguous copy uses one `memcpy` per range of the narrowed token, or a
/// single `memcpy` on the calling thread when that token is serial.
///
/// This is a simple copy without ElementOp support. For copies with
/// element operations (conj, transpose, etc.), use `strided_kernel::copy_into`.
///
/// # Examples
///
/// ```
/// use threadteam::Exec;
/// use strided_perm::copy_into;
/// use strided_view::StridedArray;
///
/// let src = StridedArray::<f64>::from_fn_col_major(&[2, 3], |idx| (idx[0] * 10 + idx[1]) as f64);
/// let mut dst = StridedArray::<f64>::col_major(&[2, 3]);
/// copy_into(&mut dst.view_mut(), &src.view(), Exec::serial()).unwrap();
/// assert_eq!(dst.get(&[1, 2]), 12.0);
/// ```
pub fn copy_into<T: Copy + Send + Sync + 'static>(
    dest: &mut StridedViewMut<T>,
    src: &StridedView<T>,
    exec: Exec<'_>,
) -> Result<()> {
    match prepare(dest, src)? {
        // SAFETY: both views are validated layouts over live, non-overlapping
        // storage (`StridedViewMut` is exclusive), and `prepare` matched them.
        Some(raw) => unsafe { copy_raw(raw, exec) },
        None => Ok(()),
    }
}

/// Copy elements into an uninitialized destination: `dest[i] = src[i]`.
///
/// The same permutation engine as [`copy_into`], selected on `T` rather than
/// on `MaybeUninit<T>`: the engine's native `f32`/`f64` kernels are chosen by
/// element type, so a destination reinterpreted as `MaybeUninit<T>` would
/// otherwise fall to the generic element path (about 2.5x slower on a
/// 4096x4096 `f64` transpose). Every destination element the views describe
/// is written exactly once; the stores go through raw pointers only, so no
/// reference to uninitialized storage is formed.
pub fn copy_into_uninit<T: Copy + Send + Sync + 'static>(
    dest: &mut StridedViewMut<core::mem::MaybeUninit<T>>,
    src: &StridedView<T>,
    exec: Exec<'_>,
) -> Result<()> {
    match prepare(dest, src)? {
        // SAFETY: as `copy_into`; `MaybeUninit<T>` has `T`'s layout and the
        // engine only stores into the destination.
        Some(raw) => unsafe { copy_raw(raw, exec) },
        None => Ok(()),
    }
}

/// Copy elements to a col-major destination.
///
/// This delegates to the same HPTT-inspired blocked permutation as
/// [`copy_into`]. The HPTT planner handles the col-major destination case
/// optimally.
pub fn copy_into_col_major<T: Copy + Send + Sync + 'static>(
    dst: &mut StridedViewMut<T>,
    src: &StridedView<T>,
    exec: Exec<'_>,
) -> Result<()> {
    copy_into(dst, src, exec)
}

#[cfg(test)]
fn try_fuse_group(dims: &[usize], strides: &[isize]) -> Option<(usize, isize)> {
    match dims.len() {
        0 => Some((1, 0)),
        1 => Some((dims[0], strides[0])),
        _ => {
            if dims.len() != strides.len() {
                return None;
            }
            for (&d, &s) in dims.iter().zip(strides.iter()) {
                if d > 1 && s == 0 {
                    return None;
                }
            }
            let mut base_idx: Option<usize> = None;
            let mut base_abs = usize::MAX;
            for (i, (&d, &s)) in dims.iter().zip(strides.iter()).enumerate() {
                if d <= 1 {
                    continue;
                }
                let abs = s.unsigned_abs();
                if abs < base_abs {
                    base_abs = abs;
                    base_idx = Some(i);
                }
            }

            let Some(base) = base_idx else {
                let stride = *strides
                    .iter()
                    .min_by_key(|s| s.unsigned_abs())
                    .unwrap_or(&0);
                return Some((dims.iter().product(), stride));
            };

            let mut used = vec![false; dims.len()];
            used[base] = true;
            let mut expected_abs = base_abs.checked_mul(dims[base])?;

            let non_singleton = dims.iter().filter(|&&d| d > 1).count();
            for _ in 1..non_singleton {
                let mut next = None;
                for i in 0..dims.len() {
                    if used[i] || dims[i] <= 1 {
                        continue;
                    }
                    if strides[i].unsigned_abs() == expected_abs {
                        next = Some(i);
                        break;
                    }
                }
                let i = next?;
                used[i] = true;
                expected_abs = expected_abs.checked_mul(dims[i])?;
            }

            Some((dims.iter().product(), strides[base]))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use strided_view::StridedArray;

    #[test]
    fn test_copy_into_contiguous() {
        let src =
            StridedArray::<f64>::from_fn_col_major(&[2, 3], |idx| (idx[0] * 10 + idx[1]) as f64);
        let mut dst = StridedArray::<f64>::col_major(&[2, 3]);
        copy_into(&mut dst.view_mut(), &src.view(), Exec::serial()).unwrap();
        assert_eq!(dst.get(&[0, 0]), 0.0);
        assert_eq!(dst.get(&[1, 2]), 12.0);
    }

    #[test]
    fn test_copy_into_transposed() {
        // src is row-major [3,2], dst is col-major [3,2]
        let src = StridedArray::<f64>::from_parts(
            vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0],
            &[3, 2],
            &[2, 1], // row-major
            0,
        )
        .unwrap();
        let mut dst = StridedArray::<f64>::col_major(&[3, 2]);
        copy_into(&mut dst.view_mut(), &src.view(), Exec::serial()).unwrap();
        assert_eq!(dst.get(&[0, 0]), 0.0);
        assert_eq!(dst.get(&[0, 1]), 1.0);
        assert_eq!(dst.get(&[1, 0]), 2.0);
        assert_eq!(dst.get(&[2, 1]), 5.0);
    }

    #[test]
    fn test_copy_into_col_major_basic() {
        let src =
            StridedArray::<f64>::from_fn_col_major(&[4, 3], |idx| (idx[0] * 10 + idx[1]) as f64);
        let mut dst = StridedArray::<f64>::col_major(&[4, 3]);
        copy_into_col_major(&mut dst.view_mut(), &src.view(), Exec::serial()).unwrap();
        assert_eq!(dst.get(&[0, 0]), 0.0);
        assert_eq!(dst.get(&[3, 2]), 32.0);
    }

    #[test]
    fn test_copy_into_col_major_permuted_src() {
        // src has scattered strides, dst is col-major
        let data: Vec<f64> = (0..24).map(|i| i as f64).collect();
        let src = StridedArray::<f64>::from_parts(
            data,
            &[2, 3, 4],
            &[12, 4, 1], // row-major
            0,
        )
        .unwrap();
        let mut dst = StridedArray::<f64>::col_major(&[2, 3, 4]);
        copy_into_col_major(&mut dst.view_mut(), &src.view(), Exec::serial()).unwrap();
        // Verify element-by-element
        for i in 0..2 {
            for j in 0..3 {
                for k in 0..4 {
                    assert_eq!(
                        dst.get(&[i, j, k]),
                        src.get(&[i, j, k]),
                        "mismatch at [{},{},{}]",
                        i,
                        j,
                        k
                    );
                }
            }
        }
    }

    #[test]
    fn test_copy_into_grouped_high_rank_binary_permutation() {
        let dims = vec![2usize; 12];
        let total: usize = dims.iter().product();
        let src = StridedArray::<u32>::from_fn_col_major(&dims, |idx| {
            idx.iter()
                .enumerate()
                .map(|(axis, &coord)| (coord as u32) << axis)
                .sum::<u32>()
        });
        let perm = vec![2, 4, 5, 6, 0, 3, 7, 1, 8, 9, 10, 11];
        let src_view = src.view().permute(&perm).unwrap();
        let mut dst = StridedArray::<u32>::col_major(&dims);

        copy_into(&mut dst.view_mut(), &src_view, Exec::serial()).unwrap();

        for flat in 0..total {
            let mut rem = flat;
            let mut src_offset = 0usize;
            for &stride in src_view.strides() {
                let coord = rem & 1;
                rem >>= 1;
                src_offset += coord * stride as usize;
            }
            assert_eq!(dst.data()[flat], src.data()[src_offset], "flat={flat}");
        }
    }

    #[test]
    fn test_try_fuse_group_empty() {
        assert_eq!(try_fuse_group(&[], &[]), Some((1, 0)));
    }

    #[test]
    fn test_try_fuse_group_single() {
        assert_eq!(try_fuse_group(&[5], &[2]), Some((5, 2)));
    }

    #[test]
    fn test_try_fuse_group_contiguous_row_major() {
        assert_eq!(try_fuse_group(&[3, 4], &[4, 1]), Some((12, 1)));
    }

    #[test]
    fn test_try_fuse_group_contiguous_col_major() {
        assert_eq!(try_fuse_group(&[3, 4], &[1, 3]), Some((12, 1)));
    }

    #[test]
    fn test_try_fuse_group_non_contiguous() {
        assert_eq!(try_fuse_group(&[3, 4], &[8, 1]), None);
    }

    #[test]
    fn parallel_transpose_splits_the_full_tile_grid_for_a_short_row_extent() {
        // The source's unit-stride extent is the short side of the transpose, so
        // the whole kernel would be a single row tile over `block`, and the budget is
        // two. The parallel entry must split the
        // column tiles too; a row-only split would leave the work on one worker.
        // The result is checked element by element.
        let rows = 20000usize;
        let cols = 3usize;
        let data: Vec<f64> = (0..rows * cols).map(|i| i as f64).collect();
        let src =
            StridedArray::<f64>::from_parts(data, &[rows, cols], &[cols as isize, 1], 0).unwrap();
        let pool = threadteam::Pool::builder().threads(2).build().unwrap();
        let mut dst = StridedArray::<f64>::col_major(&[rows, cols]);
        let exec = Exec::new(&pool, 2).unwrap();
        copy_into(&mut dst.view_mut(), &src.view(), exec).unwrap();
        for r in 0..rows {
            for c in 0..cols {
                assert_eq!(dst.get(&[r, c]), (r * cols + c) as f64);
            }
        }
    }

    #[test]
    fn test_copy_shape_mismatch() {
        let src = StridedArray::<f64>::col_major(&[2, 3]);
        let mut dst = StridedArray::<f64>::col_major(&[3, 2]);
        let result = copy_into(&mut dst.view_mut(), &src.view(), Exec::serial());
        assert!(result.is_err());
    }
}