Skip to main content

burn_compute/tune/
tune_benchmark.rs

1use burn_common::benchmark::Benchmark;
2
3use crate::channel::ComputeChannel;
4use crate::client::ComputeClient;
5use crate::server::ComputeServer;
6
7use super::AutotuneOperation;
8use alloc::boxed::Box;
9use alloc::string::{String, ToString};
10
11/// A benchmark that runs on server handles
12#[derive(new)]
13pub struct TuneBenchmark<S: ComputeServer, C> {
14    operation: Box<dyn AutotuneOperation>,
15    client: ComputeClient<S, C>,
16}
17
18impl<S: ComputeServer, C: ComputeChannel<S>> Benchmark for TuneBenchmark<S, C> {
19    type Args = Box<dyn AutotuneOperation>;
20
21    fn prepare(&self) -> Self::Args {
22        self.operation.clone()
23    }
24
25    fn num_samples(&self) -> usize {
26        10
27    }
28
29    fn execute(&self, operation: Self::Args) {
30        AutotuneOperation::execute(operation);
31    }
32
33    fn name(&self) -> String {
34        "autotune".to_string()
35    }
36
37    fn sync(&self) {
38        self.client.sync();
39    }
40}