cubek-pool 0.3.0

mod backward;
mod forward;
pub mod geometry;

use crate::{
    definition::{
        MaxPoolOptions, PoolBackward, PoolBackwardProblem, PoolForward, PoolForwardProblem,
        PoolMode, PoolProblem,
    },
    eval::cpu_reference::{
        backward::{run_adaptive_avg_pool_backward, run_avg_pool_backward, run_max_pool_backward},
        forward::{run_adaptive_avg_pool, run_avg_pool, run_max_pool, run_max_pool_with_indices},
        geometry::PoolGeometry,
    },
};
use cubecl::{
    client::Client,
    prelude::*,
    std::tensor::TensorHandle,
    zspace::{Shape, Strides},
};
use cubek_test_utils::{HostData, HostDataVec, Progress, TestInput};

pub(crate) fn f32_elem_type() -> ElemType {
    f32::elem_type_native()
}

pub(crate) fn i32_elem_type() -> ElemType {
    i32::elem_type_native()
}

pub(crate) fn make_random_f32_host(
    client: &Client,
    shape: Vec<usize>,
    seed: u64,
) -> (TensorHandle, HostData) {
    TestInput::builder(client.clone(), shape)
        .uniform(seed, -1., 1.)
        .generate_with_f32_host_data()
}

pub(crate) fn make_zero_handle(
    client: &Client,
    shape: Vec<usize>,
    dtype: ElemType,
) -> TensorHandle {
    TestInput::builder(client.clone(), shape)
        .dtype(dtype)
        .zeros()
        .generate_without_host_data()
}

pub fn strategy_result(
    client: Client,
    problem: PoolProblem,
    seed: u64,
) -> Result<HostData, String> {
    match problem {
        PoolProblem::Forward(PoolForward::D2(prob)) => forward::strategy_result(client, prob, seed),
        PoolProblem::Forward(_) => Err("cpu reference only supports 2d pool forward".to_string()),
        PoolProblem::Backward(PoolBackward::D2(prob)) => {
            backward::strategy_result(client, prob, seed)
        }
        PoolProblem::Backward(_) => Err("cpu reference only supports 2d pool backward".to_string()),
    }
}

pub fn cpu_reference_result(
    client: Client,
    problem: PoolProblem,
    seed: u64,
    progress: Option<&Progress>,
) -> Result<HostData, String> {
    match problem {
        PoolProblem::Forward(PoolForward::D2(prob)) => {
            forward::cpu_reference_result(client, prob, seed, progress)
        }
        PoolProblem::Forward(_) => Err("cpu reference only supports 2d pool forward".to_string()),
        PoolProblem::Backward(PoolBackward::D2(prob)) => {
            backward::cpu_reference_result(client, prob, seed, progress)
        }
        PoolProblem::Backward(_) => Err("cpu reference only supports 2d pool backward".to_string()),
    }
}

pub fn cpu_reference_pool<const N: usize>(
    input: &HostData,
    problem: PoolForwardProblem<N>,
) -> HostData {
    let output_shape_struct = problem.output_shape(&problem.input_shape);
    let out_dims = output_shape_struct.to_vec();
    let in_dims = problem.input_shape.to_vec();

    let out_strides = row_major_strides_vec(&out_dims);

    let output_data = match &problem.mode {
        PoolMode::Max(opts) => run_max_pool(input, opts, &out_dims, &in_dims, &out_strides),
        PoolMode::Avg(opts) => run_avg_pool(input, opts, &out_dims, &in_dims, &out_strides),
        PoolMode::AdaptiveAvg(opts) => {
            run_adaptive_avg_pool(input, opts, &out_dims, &in_dims, &out_strides)
        }
    };

    HostData {
        data: HostDataVec::F32(output_data),
        shape: output_shape_struct,
        strides: Strides::new(&out_strides),
    }
}

pub fn cpu_reference_pool_backward<const N: usize>(
    grad_output: &HostData,
    indices: &HostData,
    problem: PoolBackwardProblem<N>,
) -> HostData {
    let out_dims = grad_output.shape.to_vec();
    let input_shape = problem.input_shape();
    let in_dims = input_shape.to_vec();
    let in_strides = row_major_strides_vec(&in_dims);

    let output_grad = match &problem.mode {
        PoolMode::Max(opts) => {
            if problem.with_indices {
                run_max_pool_backward(grad_output, indices, opts, &in_dims, &out_dims, &in_strides)
            } else {
                unimplemented!("Max pool backward without indices is not implemented yet")
            }
        }
        PoolMode::Avg(_opts) => {
            run_avg_pool_backward(grad_output, _opts, &in_dims, &out_dims, &in_strides)
        }
        PoolMode::AdaptiveAvg(opts) => {
            assert_eq!(
                &out_dims[1..N + 1],
                opts.output_size.as_slice(),
                "adaptive output-gradient shape must match options"
            );
            run_adaptive_avg_pool_backward::<N>(grad_output, &in_dims, &out_dims, &in_strides)
        }
    };

    HostData {
        data: HostDataVec::F32(output_grad),
        shape: input_shape,
        strides: Strides::new(&row_major_strides_vec(&in_dims)),
    }
}

pub fn cpu_reference_max_pool_indices<const N: usize>(
    input: &HostData,
    opts: &MaxPoolOptions<N>,
    out_dims: &[usize],
) -> HostData {
    let in_dims = input.shape.to_vec();
    let out_strides = row_major_strides_vec(out_dims);
    let (_output, indices) =
        run_max_pool_with_indices(input, opts, out_dims, &in_dims, &out_strides);

    HostData {
        data: HostDataVec::I32(indices),
        shape: Shape::from(out_dims.to_vec()),
        strides: Strides::new(&out_strides),
    }
}

pub(crate) fn decode_index(mut index: usize, shape: &[usize], strides: &[usize]) -> Vec<usize> {
    let mut coords = vec![0; shape.len()];
    for i in 0..shape.len() {
        coords[i] = index / strides[i];
        index %= strides[i];
    }
    coords
}

pub(crate) fn decode_index_simple(index: usize, shape: &[usize]) -> Vec<usize> {
    let strides = row_major_strides_vec(shape);
    decode_index(index, shape, &strides)
}

pub(crate) fn row_major_strides_vec(shape: &[usize]) -> Vec<usize> {
    let mut strides = vec![1; shape.len()];
    for i in (0..shape.len() - 1).rev() {
        strides[i] = strides[i + 1] * shape[i + 1];
    }
    strides
}