strided-basic 0.6.3

//! Block size computation ported from Strided.jl
//!
//! This module computes optimal block sizes for cache-efficient iteration.
//! It uses Julia's `totalmemoryregion` for accurate memory region estimation
//! and `_computeblocks` for iterative block size reduction.

use crate::fuse::compute_costs;
use crate::inline::{InlineVec, OperandVec};
use crate::{BLOCK_MEMORY_SIZE, CACHE_LINE_SIZE};
use strided_view::auxiliary::index_order;

/// Compute block sizes for tiled iteration.
///
/// This implementation follows Julia's `_computeblocks` algorithm:
/// 1. Compute byte strides and stride orders
/// 2. If total memory fits in cache, use full dimensions
/// 3. Otherwise, iteratively reduce blocks using cost-weighted halving
///
/// # Arguments
/// * `dims` - The dimensions (in iteration order)
/// * `order` - The iteration order permutation
/// * `strides_list` - Slice of stride arrays, one per array
/// * `elem_size` - Size of each element in bytes
///
/// # Returns
/// Block sizes in iteration order
pub(crate) fn compute_block_sizes(
    dims: &[usize],
    order: &[usize],
    strides_list: &[&[isize]],
    elem_size: usize,
) -> InlineVec<usize> {
    if order.is_empty() {
        return InlineVec::new();
    }

    // Reorder dims to iteration order
    let ordered_dims: InlineVec<usize> = order.iter().map(|&i| dims[i]).collect();

    // Compute byte strides in iteration order
    let byte_strides: OperandVec<InlineVec<isize>> = strides_list
        .iter()
        .map(|strides| {
            order
                .iter()
                .map(|&i| strides[i] * elem_size as isize)
                .collect()
        })
        .collect();

    // Compute stride orders (in iteration order)
    let stride_orders: OperandVec<InlineVec<usize>> =
        byte_strides.iter().map(|bs| index_order(bs)).collect();

    // Reorder strides for cost computation
    let reordered_strides: OperandVec<InlineVec<isize>> = strides_list
        .iter()
        .map(|strides| order.iter().map(|&i| strides[i]).collect())
        .collect();

    let reordered_refs: OperandVec<&[isize]> =
        reordered_strides.iter().map(|s| s.as_slice()).collect();
    let costs = compute_costs(&reordered_refs);

    // Convert byte_strides to slices
    let byte_stride_refs: OperandVec<&[isize]> =
        byte_strides.iter().map(|s| s.as_slice()).collect();
    let stride_order_refs: OperandVec<&[usize]> =
        stride_orders.iter().map(|s| s.as_slice()).collect();

    compute_blocks(
        &ordered_dims,
        &costs,
        &byte_stride_refs,
        &stride_order_refs,
        BLOCK_MEMORY_SIZE,
    )
}

/// Compute block sizes using the Julia algorithm.
///
/// # Julia equivalent
/// ```julia
/// function _computeblocks(dims, costs, bytestrides, strideorders, blocksize)
///     if totalmemoryregion(dims, bytestrides) <= blocksize
///         return dims
///     end
///     # ... reduction logic
/// end
/// ```
fn compute_blocks(
    dims: &[usize],
    costs: &[isize],
    byte_strides: &[&[isize]],
    stride_orders: &[&[usize]],
    block_size: usize,
) -> InlineVec<usize> {
    compute_blocks_at(dims, 0, costs, byte_strides, stride_orders, block_size)
}

/// Recursive core of [`compute_blocks`] with `offset` tracking how many
/// leading axes have already been fixed. Indexing each operand's stride slice
/// from `offset` keeps the recursion allocation-free at the inline bound
/// (#306 P1).
fn compute_blocks_at(
    dims: &[usize],
    offset: usize,
    costs: &[isize],
    byte_strides: &[&[isize]],
    stride_orders: &[&[usize]],
    block_size: usize,
) -> InlineVec<usize> {
    let n = dims.len();
    if n == 0 {
        return InlineVec::new();
    }

    // If everything fits in cache, use full dims
    if total_memory_region_at(dims, byte_strides, offset) <= block_size {
        return dims.iter().copied().collect();
    }

    // Check if first dimension is smallest stride for all arrays
    let min_order = stride_orders
        .iter()
        .filter_map(|orders| orders[offset..].iter().min().copied())
        .min()
        .unwrap_or(1);

    if stride_orders
        .iter()
        .all(|orders| !orders[offset..].is_empty() && orders[offset] == min_order)
    {
        // First dimension has smallest stride in all arrays
        // Keep first dimension, recurse on rest
        let tail_blocks = compute_blocks_at(
            &dims[1..],
            offset + 1,
            &costs[1..],
            byte_strides,
            stride_orders,
            block_size,
        );

        let mut result: InlineVec<usize> = InlineVec::new();
        result.push(dims[0]);
        result.extend(tail_blocks);
        return result;
    }

    // Check if minimum stride is larger than block size
    let min_stride = byte_strides
        .iter()
        .filter_map(|s| s[offset..].iter().map(|x| x.unsigned_abs()).min())
        .min()
        .unwrap_or(0);

    if min_stride > block_size {
        return (0..n).map(|_| 1).collect();
    }

    // Iteratively reduce blocks
    let mut blocks: InlineVec<usize> = dims.iter().copied().collect();

    // Phase 1: Halve until within 2x of target
    while total_memory_region_at(&blocks, byte_strides, offset) >= 2 * block_size {
        let i = last_argmax_weighted(&blocks, costs);
        if i.is_none() || blocks[i.unwrap()] <= 1 {
            break;
        }
        let i = i.unwrap();
        blocks[i] = blocks[i].div_ceil(2);
    }

    // Phase 2: Decrement until within target
    while total_memory_region_at(&blocks, byte_strides, offset) > block_size {
        let i = last_argmax_weighted(&blocks, costs);
        if i.is_none() || blocks[i.unwrap()] <= 1 {
            break;
        }
        let i = i.unwrap();
        blocks[i] -= 1;
    }

    blocks
}

/// Compute total memory region for cache considerations.
///
/// This estimates the memory footprint considering cache line effects.
/// Strides smaller than cache line contribute to contiguous region,
/// larger strides multiply the number of cache line blocks.
///
/// # Julia equivalent
/// ```julia
/// function totalmemoryregion(dims, bytestrides)
///     memoryregion = 0
///     for i in 1:length(bytestrides)
///         strides = bytestrides[i]
///         numcontigeouscachelines = 0
///         numcachelineblocks = 1
///         for (d, s) in zip(dims, strides)
///             if s < _cachelinelength
///                 numcontigeouscachelines += (d - 1) * s
///             else
///                 numcachelineblocks *= d
///             end
///         end
///         numcontigeouscachelines = div(numcontigeouscachelines, _cachelinelength) + 1
///         memoryregion += _cachelinelength * numcontigeouscachelines * numcachelineblocks
///     end
///     return memoryregion
/// end
/// ```
#[cfg(test)]
fn total_memory_region(dims: &[usize], byte_strides: &[&[isize]]) -> usize {
    total_memory_region_at(dims, byte_strides, 0)
}

/// As [`total_memory_region`], but with the leading `offset` axes fixed. The
/// recursion in `compute_blocks_at` uses it to avoid materializing tail slices.
fn total_memory_region_at(dims: &[usize], byte_strides: &[&[isize]], offset: usize) -> usize {
    let cache_line = CACHE_LINE_SIZE;
    let mut memory_region = 0usize;

    for strides in byte_strides {
        let strides = &strides[offset..];
        let mut num_contiguous_cache_lines = 0isize;
        let mut num_cache_line_blocks = 1usize;

        for (&d, &s) in dims.iter().zip(strides.iter()) {
            let s_abs = s.unsigned_abs();
            if s_abs < cache_line {
                // Small stride: contributes to contiguous region
                num_contiguous_cache_lines += (d.saturating_sub(1) as isize) * (s_abs as isize);
            } else {
                // Large stride: multiplies cache line blocks
                num_cache_line_blocks *= d;
            }
        }

        // Convert to cache lines
        let contiguous_lines = (num_contiguous_cache_lines as usize / cache_line) + 1;

        memory_region += cache_line * contiguous_lines * num_cache_line_blocks;
    }

    memory_region
}

/// Find the last index with maximum (value - 1) * cost.
///
/// Julia: `_lastargmax((blocks .- 1) .* costs)`
fn last_argmax_weighted(blocks: &[usize], costs: &[isize]) -> Option<usize> {
    if blocks.is_empty() {
        return None;
    }

    let mut max_score = 0isize;
    let mut max_idx = None;

    for (i, (&b, &c)) in blocks.iter().zip(costs.iter()).enumerate() {
        if b <= 1 {
            continue;
        }
        let score = (b as isize - 1) * c;
        if score >= max_score {
            max_score = score;
            max_idx = Some(i);
        }
    }

    max_idx
}

#[cfg(test)]
#[path = "block/tests/tests.rs"]
mod tests;