burn_compute/
client.rs

1use crate::{
2    channel::ComputeChannel,
3    server::{ComputeServer, Handle},
4    tune::{AutotuneOperationSet, Tuner},
5};
6use alloc::vec::Vec;
7use alloc::{boxed::Box, sync::Arc};
8use burn_common::reader::Reader;
9use burn_common::stub::RwLock;
10
11/// The ComputeClient is the entry point to require tasks from the ComputeServer.
12/// It should be obtained for a specific device via the Compute struct.
13#[derive(Debug)]
14pub struct ComputeClient<Server: ComputeServer, Channel> {
15    channel: Channel,
16    tuner: Arc<RwLock<Tuner<Server, Channel>>>,
17}
18
19impl<S, C> Clone for ComputeClient<S, C>
20where
21    S: ComputeServer,
22    C: ComputeChannel<S>,
23{
24    fn clone(&self) -> Self {
25        Self {
26            channel: self.channel.clone(),
27            tuner: self.tuner.clone(),
28        }
29    }
30}
31
32impl<Server, Channel> ComputeClient<Server, Channel>
33where
34    Server: ComputeServer,
35    Channel: ComputeChannel<Server>,
36{
37    /// Create a new client.
38    pub fn new(channel: Channel, tuner: Arc<RwLock<Tuner<Server, Channel>>>) -> Self {
39        Self { channel, tuner }
40    }
41
42    /// Given a handle, returns owned resource as bytes.
43    pub fn read(&self, handle: &Handle<Server>) -> Reader<Vec<u8>> {
44        self.channel.read(handle)
45    }
46
47    /// Given a resource, stores it and returns the resource handle.
48    pub fn create(&self, data: &[u8]) -> Handle<Server> {
49        self.channel.create(data)
50    }
51
52    /// Reserves `size` bytes in the storage, and returns a handle over them.
53    pub fn empty(&self, size: usize) -> Handle<Server> {
54        self.channel.empty(size)
55    }
56
57    /// Executes the `kernel` over the given `handles`.
58    pub fn execute(&self, kernel: Server::Kernel, handles: &[&Handle<Server>]) {
59        self.channel.execute(kernel, handles)
60    }
61
62    /// Wait for the completion of every task in the server.
63    pub fn sync(&self) {
64        self.channel.sync()
65    }
66
67    /// Executes the fastest kernel in the autotune operation, using (cached) runtime benchmarks
68    pub fn autotune_execute(
69        &self,
70        autotune_operation_set: Box<dyn AutotuneOperationSet<Server::AutotuneKey>>,
71    ) {
72        self.tuner
73            .write()
74            .unwrap()
75            .execute_autotune(autotune_operation_set, self);
76    }
77
78    /// Get the fastest kernel for the given autotune key if it exists.
79    pub fn autotune_result(&self, key: &Server::AutotuneKey) -> Option<usize> {
80        self.tuner.read().unwrap().autotune_fastest(key)
81    }
82}
burn_compute/client.rs

burn_compute/
client.rs
