cubek-pool 0.3.0

CubeK: Pool
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

mod adaptive_avg_pool2d;
mod adaptive_avg_pool3d;
mod avg_pool2d;
mod max_pool2d;

use super::{
    build_output_tensor, indices_elem_type, output_host_f32, output_host_i32, validate_indices,
    validate_test,
};
use cubecl::{client::Client, zspace::Shape};
use cubek_pool::{
    definition::{PoolForwardProblem, PoolMode},
    eval::cpu_reference::{
        cpu_reference_max_pool_indices, cpu_reference_pool, geometry::PoolGeometry,
    },
    pool2d, pool2d_with_indices,
};
use cubek_test_utils::TestInput;

pub fn make_problem(
    input_shape: Shape,
    with_indices: bool,
    mode: impl Into<PoolMode<2>>,
) -> PoolForwardProblem<2> {
    PoolForwardProblem {
        input_shape,
        with_indices,
        mode: mode.into(),
    }
}

pub fn run_pool_test(
    client: Client,
    seed: u64,
    input_min: f32,
    input_max: f32,
    problem: PoolForwardProblem<2>,
    tolerance: f32,
) {
    let (input, input_data) = TestInput::builder(client.clone(), problem.input_shape.to_vec())
        .uniform(seed, input_min, input_max)
        .generate_with_f32_host_data();

    let output_shape = problem.output_shape(&problem.input_shape).to_vec();
    let output = build_output_tensor(&client, output_shape.clone(), input.dtype);

    let indices = if problem.with_indices {
        Some(build_output_tensor(
            &client,
            output_shape.to_vec(),
            indices_elem_type(),
        ))
    } else {
        None
    };

    let result = if problem.with_indices {
        pool2d_with_indices(
            &client,
            input.clone().binding(),
            output.clone().binding(),
            indices.clone().unwrap().binding(),
            problem.mode.clone(),
            input.dtype,
        )
    } else {
        pool2d(
            &client,
            input.clone().binding(),
            output.clone().binding(),
            problem.mode.clone(),
            input.dtype,
        )
    };

    let output_host = output_host_f32(&client, output);
    let indices_host = if problem.with_indices {
        Some(output_host_i32(
            &client,
            indices.expect("indices tensor missing"),
        ))
    } else {
        None
    };

    if problem.with_indices
        && let PoolMode::Max(opts) = &problem.mode
    {
        let indices_reference = cpu_reference_max_pool_indices(&input_data, opts, &output_shape);
        validate_indices(
            indices_host
                .as_ref()
                .expect("indices host data missing")
                .clone(),
            indices_reference,
        );
    }

    let reference = cpu_reference_pool(&input_data, problem);

    validate_test(result, output_host, reference, tolerance);
}