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
106
107
108
109
110
111
112
113

use cubecl::{
    benchmark::{Benchmark, ProfileDuration, TimingMethod},
    client::Client,
    future,
    prelude::*,
    std::tensor::TensorHandle,
    zspace::Shape,
};
use cubek_test_utils::{RunSamples, TestInput};

use crate::{ConvolutionArgs, DirectTensors, launch_direct};

use super::problem::{DirectProblem, DirectStrategy};

pub fn bench(
    strategy: &DirectStrategy,
    problem: &DirectProblem,
    num_samples: usize,
) -> Result<RunSamples, String> {
    let device = cubecl::test_device();
    let client = device.client();

    let bench = DirectBench {
        problem: *problem,
        strategy: *strategy,
        client,
        device,
        samples: num_samples,
    };

    let durations = bench
        .run(cubek_test_utils::timing_method(TimingMethod::System))
        .map_err(|e| format!("benchmark failed: {e}"))?
        .durations;

    Ok(RunSamples::new(durations))
}

struct DirectBench {
    problem: DirectProblem,
    #[allow(dead_code)]
    strategy: DirectStrategy,
    device: cubecl::Device,
    client: Client,
    samples: usize,
}

fn dtype() -> ElemType {
    f32::elem_type_native()
}

fn uniform(client: &Client, shape: [usize; 4], seed: u64) -> TensorHandle {
    TestInput::builder(client.clone(), Shape::new(shape))
        .dtype(dtype())
        .uniform(seed, 0.0, 1.0)
        .generate_without_host_data()
}

impl Benchmark for DirectBench {
    type Input = (TensorHandle, TensorHandle);
    type Output = ();

    fn prepare(&self) -> Self::Input {
        (
            uniform(&self.client, self.problem.in_shape(), 0),
            uniform(&self.client, self.problem.weight_shape(), 1),
        )
    }

    fn execute(&self, (input, weight): Self::Input) -> Result<(), String> {
        let problem = &self.problem;
        let out: TensorHandle =
            TensorHandle::empty(&self.client, problem.out_shape().to_vec(), dtype());

        let padding = problem.padding();
        let args = ConvolutionArgs::<2> {
            stride: [problem.stride; 2],
            padding: [padding; 2],
            dilation: [1; 2],
        };

        launch_direct::<2>(
            &self.client,
            DirectTensors {
                input: input.binding(),
                weight: weight.binding(),
                bias: None,
                out: out.binding(),
            },
            args,
            1,
            dtype(),
        )
        .map_err(|e| format!("{e:?}"))
    }

    fn num_samples(&self) -> usize {
        self.samples
    }

    fn name(&self) -> String {
        let client = self.device.client();
        format!("{}-direct-conv-{}", client.name(), dtype()).to_lowercase()
    }

    fn sync(&self) {
        future::block_on(self.client.sync()).unwrap()
    }

    fn profile(&self, args: Self::Input) -> Result<ProfileDuration, String> {
        cubek_test_utils::profile_launch(&self.client, "direct-conv-bench", || self.execute(args))
    }
}