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

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

pub fn run_avg_pool_backward<const N: usize>(
    grad_output: &HostData,
    opts: &AvgPoolOptions<N>,
    grad_input_dims: &[usize],
    grad_output_dims: &[usize],
    grad_input_strides: &[usize],
) -> Vec<f32> {
    let total: usize = grad_input_dims.iter().product();
    let mut grad_input = vec![0.0; total];

    if N != 2 {
        return grad_input;
    }

    let stride_h = opts.window.stride[0] as i32;
    let stride_w = opts.window.stride[1] as i32;
    let pad_h = opts.window.padding[0] as i32;
    let pad_w = opts.window.padding[1] as i32;
    let kernel_h = opts.window.kernel_size[0] as i32;
    let kernel_w = opts.window.kernel_size[1] as i32;

    let grad_h = grad_output_dims[1] as i32;
    let grad_w = grad_output_dims[2] as i32;

    for (i, grad_val) in grad_input.iter_mut().enumerate().take(total) {
        let coords = decode_index(i, grad_input_dims, grad_input_strides);
        let batch = coords[0] as i32;
        let ih = coords[1] as i32;
        let iw = coords[2] as i32;
        let channel = coords[3] as i32;

        let kms_h = kernel_h - stride_h;
        let kms_w = kernel_w - stride_w;

        let oh_start = ((ih + pad_h - kms_h).max(0)) / stride_h;
        let ow_start = ((iw + pad_w - kms_w).max(0)) / stride_w;
        let oh_end = ((kms_h.max(0) + oh_start).min(grad_h - 1)) + 1;
        let ow_end = ((kms_w.max(0) + ow_start).min(grad_w - 1)) + 1;

        let begin_h = ih + pad_h;
        let begin_w = iw + pad_w;

        let out_h = grad_input_dims[1] as i32;
        let out_w = grad_input_dims[2] as i32;
        let border_bottom = out_h + pad_h;
        let border_right = out_w + pad_w;

        let mut grad_acc = 0.0f32;

        for oh in oh_start..oh_end {
            let ih_start = (oh * stride_h).min(border_bottom).max(pad_h);
            let ih_end = (oh * stride_h + kernel_h).min(border_bottom);

            if begin_h >= ih_start && ih < ih_end {
                for ow in ow_start..ow_end {
                    let iw_start = (ow * stride_w).min(border_right).max(pad_w);
                    let iw_end = (ow * stride_w + kernel_w).min(border_right);

                    if begin_w >= iw_start && iw < iw_end {
                        let out_coords =
                            vec![batch as usize, oh as usize, ow as usize, channel as usize];
                        let grad_val = grad_output.get_f32(&out_coords);

                        if opts.count_include_pad {
                            let padded_h = kernel_h.min(out_h + 2 * pad_h - oh * stride_h);
                            let padded_w = kernel_w.min(out_w + 2 * pad_w - ow * stride_w);
                            grad_acc += grad_val / (padded_h * padded_w) as f32;
                        } else {
                            let ih_diff = (ih_end - ih_start) as f32;
                            let iw_diff = (iw_end - iw_start) as f32;
                            grad_acc += grad_val / (ih_diff * iw_diff);
                        }
                    }
                }
            }
        }

        *grad_val = grad_acc;
    }

    grad_input
}