strided-kernel 0.1.1

//! Dimension fusion logic ported from Strided.jl/src/mapreduce.jl
//!
//! This module implements the core dimension fusion algorithm that merges
//! contiguous dimensions to reduce iteration complexity.

/// Fuse contiguous dimensions across multiple arrays.
///
/// This function fuses subsequent dimensions that are contiguous in memory
/// for all arrays. If `strides[k][i] == dims[i-1] * strides[k][i-1]` for all k,
/// dimensions i-1 and i can be merged.
///
/// # Arguments
/// * `dims` - The shared dimensions of all arrays
/// * `all_strides` - Vector of stride tuples, one per array
///
/// # Returns
/// The fused dimensions (stride values remain unchanged, caller must recompute)
///
/// # Julia equivalent
/// ```julia
/// function _mapreduce_fuse!(f, op, initop, dims, arrays)
///     allstrides = map(strides, arrays)
///     @inbounds for i in length(dims):-1:2
///         merge = true
///         for s in allstrides
///             if s[i] != dims[i - 1] * s[i - 1]
///                 merge = false
///                 break
///             end
///         end
///         if merge
///             dims = setindex(dims, dims[i - 1] * dims[i], i - 1)
///             dims = setindex(dims, 1, i)
///         end
///     end
///     return dims
/// end
/// ```
pub fn fuse_dims(dims: &[usize], all_strides: &[&[isize]]) -> Vec<usize> {
    let n = dims.len();
    if n <= 1 || all_strides.is_empty() {
        return dims.to_vec();
    }

    let mut result = dims.to_vec();

    // Work from the end towards the beginning (Julia: for i in length(dims):-1:2)
    for i in (1..n).rev() {
        let mut can_merge = true;

        // Check all arrays for contiguity. An operand broadcast across both
        // axes has stride 0 for both axes and does not prevent fusing the
        // iteration space.
        for strides in all_strides {
            if strides[i - 1] == 0 && strides[i] == 0 {
                continue;
            }

            // s[i] should equal dims[i-1] * s[i-1] for fusion
            let expected = result[i - 1] as isize * strides[i - 1];
            if strides[i] != expected {
                can_merge = false;
                break;
            }
        }

        if can_merge {
            // Fuse dimensions: merge dimension i into i-1
            result[i - 1] *= result[i];
            result[i] = 1;
        }
    }

    result
}

/// Remove size-1 dimensions from fused dims and all corresponding strides.
///
/// After `fuse_dims()`, many dimensions may be 1 (either originally size-1
/// or merged into a neighbor). These contribute nothing to iteration but
/// increase loop depth. This function strips them out.
///
/// If ALL dimensions are 1 (scalar-like), a single dimension of size 1
/// is preserved so the kernel has something to iterate over.
pub fn compress_dims(dims: &[usize], all_strides: &[Vec<isize>]) -> (Vec<usize>, Vec<Vec<isize>>) {
    let kept: Vec<usize> = (0..dims.len()).filter(|&i| dims[i] != 1).collect();

    if kept.is_empty() {
        // All dims are 1 (or empty). Preserve a single trivial dimension.
        if dims.is_empty() {
            return (vec![], all_strides.to_vec());
        }
        let new_strides = all_strides.iter().map(|s| vec![s[0]]).collect();
        return (vec![1], new_strides);
    }

    let new_dims: Vec<usize> = kept.iter().map(|&i| dims[i]).collect();
    let new_strides: Vec<Vec<isize>> = all_strides
        .iter()
        .map(|s| kept.iter().map(|&i| s[i]).collect())
        .collect();

    (new_dims, new_strides)
}

/// Compute the "importance" of each dimension for loop ordering.
///
/// This encodes stride order information into importance scores that determine
/// the optimal iteration order. Zero-stride broadcast axes do not contribute to
/// importance, and the output array gets strong weight so contiguous stores
/// remain in the inner loop.
///
/// # Arguments
/// * `dims` - The dimensions
/// * `all_strides` - Vector of stride tuples
/// * `index_orders` - Pre-computed index orders for each stride tuple
///
/// # Returns
/// Importance scores for each dimension
pub fn compute_importance(
    dims: &[usize],
    all_strides: &[&[isize]],
    index_orders: &[Vec<usize>],
) -> Vec<u64> {
    let n = dims.len();
    let m = all_strides.len();

    if n == 0 || m == 0 {
        return vec![];
    }

    // g = ceil(log2(M + 2)) = number of bits needed to encode array count
    let g = (64 - (m as u64 + 1).leading_zeros()) as u64;

    let mut importance = vec![0u64; n];

    let output_weight = 1u64 << (g + 1);

    // First array (output) gets a strong weight. For elementwise kernels the
    // store stream determines whether the inner loop can be emitted as a
    // contiguous vector loop, so a broadcasted input's local stride should not
    // break a contiguous output group.
    for i in 0..n {
        if all_strides[0][i] != 0 {
            let shift = g * (n - index_orders[0][i]) as u64;
            importance[i] = output_weight * (1u64 << shift);
        }
    }

    // Add contributions from remaining arrays
    #[allow(clippy::needless_range_loop)]
    for k in 1..m {
        for i in 0..n {
            if all_strides[k][i] != 0 {
                let shift = g * (n - index_orders[k][i]) as u64;
                importance[i] += 1u64 << shift;
            }
        }
    }

    // Zero importance for size-1 dimensions (put them at the back)
    for i in 0..n {
        if dims[i] <= 1 {
            importance[i] = 0;
        }
    }

    importance
}

/// Get the permutation that sorts by importance (descending).
///
/// Returns indices that would sort the importance array in descending order.
pub fn sort_by_importance(importance: &[u64]) -> Vec<usize> {
    let mut indices: Vec<usize> = (0..importance.len()).collect();
    indices.sort_by(|&a, &b| importance[b].cmp(&importance[a]));
    indices
}

/// Compute the minimum stride cost for each dimension.
///
/// Julia: `costs = map(a -> ifelse(iszero(a), 1, a << 1), map(min, strides...))`
pub fn compute_costs<S: AsRef<[isize]>>(all_strides: &[S]) -> Vec<isize> {
    if all_strides.is_empty() {
        return vec![];
    }

    let n = all_strides[0].as_ref().len();
    let mut costs = vec![isize::MAX; n];

    for strides in all_strides {
        let strides = strides.as_ref();
        for i in 0..n {
            costs[i] = costs[i].min(strides[i].abs());
        }
    }

    // Transform: zero -> 1, nonzero -> 2*abs
    for cost in &mut costs {
        if *cost == 0 {
            *cost = 1;
        } else {
            *cost *= 2;
        }
    }

    costs
}

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

    #[test]
    fn test_fuse_dims_contiguous() {
        // Two contiguous dimensions: [3, 4] with strides [1, 3] -> fused to [12, 1]
        let dims = [3, 4];
        let strides1 = [1isize, 3];
        let strides2 = [1isize, 3];
        let all_strides: Vec<&[isize]> = vec![&strides1, &strides2];

        let fused = fuse_dims(&dims, &all_strides);
        assert_eq!(fused, vec![12, 1]);
    }

    #[test]
    fn test_fuse_dims_allows_broadcast_operand_across_fused_axes() {
        let dims = [16usize, 16, 64, 64];
        let out = [1isize, 16, 256, 16_384];
        let lhs = [1isize, 16, 0, 256];
        let rhs = [0isize, 0, 1, 64];
        let all_strides: Vec<&[isize]> = vec![&out, &lhs, &rhs];

        let fused = fuse_dims(&dims, &all_strides);

        assert_eq!(fused, vec![256, 1, 64, 64]);
    }

    #[test]
    fn test_fuse_dims_non_contiguous() {
        // Non-contiguous: strides don't match
        let dims = [3, 4];
        let strides1 = [1isize, 10]; // Not contiguous (should be 3)
        let all_strides: Vec<&[isize]> = vec![&strides1];

        let fused = fuse_dims(&dims, &all_strides);
        assert_eq!(fused, vec![3, 4]); // No fusion
    }

    #[test]
    fn test_fuse_dims_partial() {
        // 3D: first two fuse, third doesn't
        let dims = [2, 3, 4];
        let strides = [1isize, 2, 100]; // dims[0]*strides[0]=2=strides[1], but 6≠100
        let all_strides: Vec<&[isize]> = vec![&strides];

        let fused = fuse_dims(&dims, &all_strides);
        assert_eq!(fused, vec![6, 1, 4]); // Fused first two
    }

    #[test]
    fn test_fuse_dims_multiple_arrays() {
        // Only fuse if ALL arrays are contiguous
        let dims = [3, 4];
        let strides1 = [1isize, 3]; // Contiguous
        let strides2 = [1isize, 10]; // Not contiguous
        let all_strides: Vec<&[isize]> = vec![&strides1, &strides2];

        let fused = fuse_dims(&dims, &all_strides);
        assert_eq!(fused, vec![3, 4]); // No fusion because strides2 isn't contiguous
    }

    #[test]
    fn test_compute_importance_2_arrays() {
        // Example with 2 arrays, dims [4, 5]
        let dims = [4usize, 5];
        let strides1 = [1isize, 4]; // Column-major output
        let strides2 = [5isize, 1]; // Row-major input
        let all_strides: Vec<&[isize]> = vec![&strides1, &strides2];

        let order1 = index_order(&strides1);
        let order2 = index_order(&strides2);
        let index_orders = vec![order1, order2];

        let importance = compute_importance(&dims, &all_strides, &index_orders);

        // With output weighted 2x, dimension 0 should have higher importance
        // since it has smaller stride in the output array
        assert!(importance[0] > importance[1]);
    }

    #[test]
    fn test_sort_by_importance() {
        let importance = vec![100u64, 50, 200, 10];
        let perm = sort_by_importance(&importance);
        assert_eq!(perm, vec![2, 0, 1, 3]); // Indices sorted by descending importance
    }

    #[test]
    fn test_compute_costs() {
        let strides1 = [1isize, 4, 0];
        let strides2 = [2isize, 1, 0];
        let all_strides: Vec<&[isize]> = vec![&strides1, &strides2];

        let costs = compute_costs(&all_strides);
        // min strides: [1, 1, 0], transformed: [2, 2, 1]
        assert_eq!(costs, vec![2, 2, 1]);
    }

    #[test]
    fn test_compute_importance_with_zero_stride() {
        // Zero stride (broadcast) still gets index_order = 1 for compatibility,
        // but it does not contribute to ordering importance.
        let dims = [4usize, 5];
        let strides1 = [0isize, 1]; // First dim is broadcast
        let all_strides: Vec<&[isize]> = vec![&strides1];

        let order1 = index_order(&strides1);
        // For stride 0: order = 1 (zero strides always get 1)
        // For stride 1: order = 1 (no non-zero stride < 1)
        assert_eq!(order1, vec![1, 1]);

        let index_orders = vec![order1];
        let importance = compute_importance(&dims, &all_strides, &index_orders);

        assert_eq!(importance[0], 0);
        assert!(importance[1] > 0);
    }

    #[test]
    fn test_compute_importance_size_one_dim() {
        // Size-1 dimensions get zero importance
        let dims = [4usize, 1, 5];
        let strides1 = [1isize, 4, 4];
        let all_strides: Vec<&[isize]> = vec![&strides1];

        let order1 = index_order(&strides1);
        let index_orders = vec![order1];
        let importance = compute_importance(&dims, &all_strides, &index_orders);

        // Dimension 1 has size 1 -> importance = 0
        assert_eq!(importance[1], 0);
        // Other dimensions should have non-zero importance
        assert!(importance[0] > 0);
        assert!(importance[2] > 0);
    }

    #[test]
    fn test_compute_importance_output_weight() {
        // Output (first array) is strongly weighted
        // With same strides, dimension with smaller stride in output wins
        let dims = [4usize, 5];
        let out_strides = [1isize, 4]; // Column-major output
        let in_strides = [5isize, 1]; // Row-major input
        let all_strides: Vec<&[isize]> = vec![&out_strides, &in_strides];

        let order_out = index_order(&out_strides); // [1, 2]
        let order_in = index_order(&in_strides); // [2, 1]
        let index_orders = vec![order_out, order_in];

        let importance = compute_importance(&dims, &all_strides, &index_orders);

        // Output weighting makes dimension 0 (smaller stride in output) win.
        assert!(importance[0] > importance[1]);
    }

    #[test]
    fn test_compute_costs_owned_vecs() {
        // Ported from threading.rs: verify compute_costs works with Vec<Vec<isize>>
        let strides_list: Vec<Vec<isize>> = vec![vec![1, 0, 3], vec![2, 0, 4]];
        let costs = compute_costs(&strides_list);
        assert_eq!(costs, vec![2, 1, 6]);
    }

    #[test]
    fn test_compute_costs_with_zero() {
        // Zero strides become cost 1, non-zero become 2*abs
        let strides1 = [0isize, 2, -3];
        let strides2 = [1isize, 0, 2];
        let all_strides: Vec<&[isize]> = vec![&strides1, &strides2];

        let costs = compute_costs(&all_strides);
        // min abs: [0, 0, 2]
        // transform: [1, 1, 4]
        assert_eq!(costs, vec![1, 1, 4]);
    }

    // ---- compress_dims tests ----

    #[test]
    fn test_compress_dims_removes_fused() {
        let dims = vec![12usize, 1];
        let strides = vec![vec![1isize, 3]];
        let (cd, cs) = compress_dims(&dims, &strides);
        assert_eq!(cd, vec![12]);
        assert_eq!(cs, vec![vec![1]]);
    }

    #[test]
    fn test_compress_dims_removes_multiple() {
        let dims = vec![6usize, 1, 4];
        let strides = vec![vec![1isize, 2, 100]];
        let (cd, cs) = compress_dims(&dims, &strides);
        assert_eq!(cd, vec![6, 4]);
        assert_eq!(cs, vec![vec![1, 100]]);
    }

    #[test]
    fn test_compress_dims_no_removal() {
        let dims = vec![3usize, 4];
        let strides = vec![vec![1isize, 3]];
        let (cd, cs) = compress_dims(&dims, &strides);
        assert_eq!(cd, vec![3, 4]);
        assert_eq!(cs, vec![vec![1, 3]]);
    }

    #[test]
    fn test_compress_dims_all_ones() {
        let dims = vec![1usize, 1, 1];
        let strides = vec![vec![1isize, 1, 1]];
        let (cd, cs) = compress_dims(&dims, &strides);
        assert_eq!(cd, vec![1]);
        assert_eq!(cs, vec![vec![1]]);
    }

    #[test]
    fn test_compress_dims_multi_arrays() {
        let dims = vec![6usize, 1, 4];
        let strides = vec![vec![1isize, 6, 6], vec![4isize, 24, 1]];
        let (cd, cs) = compress_dims(&dims, &strides);
        assert_eq!(cd, vec![6, 4]);
        assert_eq!(cs, vec![vec![1, 6], vec![4, 1]]);
    }

    #[test]
    fn test_compress_dims_single_dim() {
        let dims = vec![5usize];
        let strides = vec![vec![1isize]];
        let (cd, cs) = compress_dims(&dims, &strides);
        assert_eq!(cd, vec![5]);
        assert_eq!(cs, vec![vec![1]]);
    }

    #[test]
    fn test_compress_dims_single_dim_one() {
        let dims = vec![1usize];
        let strides = vec![vec![1isize]];
        let (cd, cs) = compress_dims(&dims, &strides);
        assert_eq!(cd, vec![1]);
        assert_eq!(cs, vec![vec![1]]);
    }

    #[test]
    fn test_compress_dims_empty() {
        let dims: Vec<usize> = vec![];
        let strides: Vec<Vec<isize>> = vec![];
        let (cd, cs) = compress_dims(&dims, &strides);
        assert!(cd.is_empty());
        assert!(cs.is_empty());
    }
}