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

use crate::{
    definition::AdaptiveAvgPoolOptions,
    eval::cpu_reference::{decode_index, decode_index_simple},
};
use cubek_test_utils::HostData;

pub fn run_adaptive_avg_pool<const N: usize>(
    input: &HostData,
    _opts: &AdaptiveAvgPoolOptions<N>,
    out_dims: &[usize],
    in_dims: &[usize],
    out_strides: &[usize],
) -> Vec<f32> {
    let total: usize = out_dims.iter().product();
    let mut out = vec![0.0; total];

    for (i, out_val) in out.iter_mut().enumerate().take(total) {
        let coords = decode_index(i, out_dims, out_strides);
        let (batch, channel, spatial_out) = (coords[0], coords[N + 1], &coords[1..N + 1]);

        let mut starts = [0; N];
        let mut ends = [0; N];
        for d in 0..N {
            starts[d] = (spatial_out[d] * in_dims[d + 1]) / out_dims[d + 1];
            ends[d] = ((spatial_out[d] + 1) * in_dims[d + 1])
                .div_ceil(out_dims[d + 1])
                .min(in_dims[d + 1]);
        }

        let k_dims: Vec<usize> = (0..N).map(|d| ends[d] - starts[d]).collect();
        let total_k: usize = k_dims.iter().product();
        let (mut sum, mut count) = (0.0, 0);

        for k_idx in 0..total_k {
            let k_coords = decode_index_simple(k_idx, &k_dims);
            let mut in_coords = vec![0; N + 2];
            in_coords[0] = batch;
            in_coords[N + 1] = channel;
            for d in 0..N {
                in_coords[d + 1] = starts[d] + k_coords[d];
            }

            sum += input.get_f32(&in_coords);
            count += 1;
        }
        *out_val = if count > 0 { sum / count as f32 } else { 0.0 };
    }
    out
}