cubek-convolution 0.3.0

CubeK: Convolution Kernels
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
103
104
105

//! Benchmark catalogue for `cubek-convolution`.
//!
//! Gated behind the `benchmarks` cargo feature. The top-level `benchmarks`
//! crate re-exports [`Category`] from here and aggregates it with the other
//! kernels' catalogues.

mod benchmark;
mod correctness;
pub mod depthwise;
pub mod direct;
mod problem;
mod strategy;

pub use benchmark::bench;
pub use correctness::Conv2dCorrectness;
pub use problem::{Conv2dProblem, problems};
pub use strategy::strategies;

use cubecl::prelude::*;
use cubek_test_utils::{CatalogEntry, CategoryWork, ComputeWork, RunSamples, client};

use cubek_matmul::definition::MatmulGlobalElems;

use crate::definition::Conv2dCost;

use crate::Strategy;

pub struct Category;

impl cubek_test_utils::Category for Category {
    type Problem = Conv2dProblem;
    type Strategy = Strategy;

    fn id(&self) -> &'static str {
        "conv2d"
    }

    fn label(&self) -> &'static str {
        "Conv2d"
    }

    fn problems(&self) -> Vec<CatalogEntry<Conv2dProblem>> {
        problems()
    }

    fn strategies(&self) -> Vec<CatalogEntry<Strategy>> {
        strategies()
    }

    fn bench(
        &self,
        strategy: &Strategy,
        problem: &Conv2dProblem,
        num_samples: usize,
    ) -> Result<RunSamples, String> {
        bench(strategy, problem, num_samples)
    }
    fn correctness(
        &self,
    ) -> Option<&dyn cubek_test_utils::Correctness<Problem = Conv2dProblem, Strategy = Strategy>>
    {
        Some(&Conv2dCorrectness)
    }

    /// Standard conv flop count: one multiply-add per (output element, kernel tap,
    /// input channel). Runs in `half::f16`, the precision `bench` fixes.
    fn work(&self, problem: &Conv2dProblem) -> Option<CategoryWork> {
        let dtype = half::f16::elem_type_native();

        let [n, c_in, h_in, w_in] = problem.input_shape;
        let [c_out, _, k_h, k_w] = problem.weight_shape;
        let [s_h, s_w] = problem.args.stride;
        let [p_h, p_w] = problem.args.padding;
        let [d_h, d_w] = problem.args.dilation;

        let cost = Conv2dCost {
            batch: n,
            channels_in: c_in,
            spatial_in: [h_in, w_in],
            channels_out: c_out,
            kernel: [k_h, k_w],
            spatial_out: [
                (h_in + 2 * p_h - d_h * (k_h - 1) - 1) / s_h + 1,
                (w_in + 2 * p_w - d_w * (k_w - 1) - 1) / s_w + 1,
            ],
            bias_elems: problem.bias_shape,
            elems: MatmulGlobalElems {
                lhs: dtype,
                rhs: dtype,
                out: dtype,
            },
        };

        let (bytes_read, bytes_written) = cost.traffic();

        Some(CategoryWork {
            compute: Some(ComputeWork {
                ops: cost.compute_ops(),
                key: cost.compute_key(&client()),
            }),
            bytes_read,
            bytes_written,
        })
    }
}