burn_compute/channel/
mutex.rs

1use super::ComputeChannel;
2use crate::server::{ComputeServer, Handle};
3use alloc::sync::Arc;
4use alloc::vec::Vec;
5use burn_common::reader::Reader;
6use spin::Mutex;
7
8/// The MutexComputeChannel ensures thread-safety by locking the server
9/// on every operation
10#[derive(Debug)]
11pub struct MutexComputeChannel<Server> {
12    server: Arc<Mutex<Server>>,
13}
14
15impl<S> Clone for MutexComputeChannel<S> {
16    fn clone(&self) -> Self {
17        Self {
18            server: self.server.clone(),
19        }
20    }
21}
22impl<Server> MutexComputeChannel<Server>
23where
24    Server: ComputeServer,
25{
26    /// Create a new mutex compute channel.
27    pub fn new(server: Server) -> Self {
28        Self {
29            server: Arc::new(Mutex::new(server)),
30        }
31    }
32}
33
34impl<Server> ComputeChannel<Server> for MutexComputeChannel<Server>
35where
36    Server: ComputeServer,
37{
38    fn read(&self, handle: &Handle<Server>) -> Reader<Vec<u8>> {
39        self.server.lock().read(handle)
40    }
41
42    fn create(&self, data: &[u8]) -> Handle<Server> {
43        self.server.lock().create(data)
44    }
45
46    fn empty(&self, size: usize) -> Handle<Server> {
47        self.server.lock().empty(size)
48    }
49
50    fn execute(&self, kernel: Server::Kernel, handles: &[&Handle<Server>]) {
51        self.server.lock().execute(kernel, handles)
52    }
53
54    fn sync(&self) {
55        self.server.lock().sync()
56    }
57}
burn_compute/channel/mutex.rs

burn_compute/channel/
mutex.rs
