burn_compute/channel/
cell.rs

1use super::ComputeChannel;
2use crate::server::{ComputeServer, Handle};
3use alloc::sync::Arc;
4use alloc::vec::Vec;
5use burn_common::reader::Reader;
6
7/// A channel using a [ref cell](core::cell::RefCell) to access the server with mutability.
8///
9/// # Important
10///
11/// Only use this channel if you don't use any threading in your application, otherwise it will
12/// panic or cause undefined behaviors.
13///
14/// This is mosly useful for `no-std` environments where threads aren't supported, otherwise prefer
15/// the [mutex](super::MutexComputeChannel) or the [mpsc](super::MpscComputeChannel) channels.
16#[derive(Debug)]
17pub struct RefCellComputeChannel<Server> {
18    server: Arc<core::cell::RefCell<Server>>,
19}
20
21impl<S> Clone for RefCellComputeChannel<S> {
22    fn clone(&self) -> Self {
23        Self {
24            server: self.server.clone(),
25        }
26    }
27}
28
29impl<Server> RefCellComputeChannel<Server>
30where
31    Server: ComputeServer,
32{
33    /// Create a new cell compute channel.
34    pub fn new(server: Server) -> Self {
35        Self {
36            server: Arc::new(core::cell::RefCell::new(server)),
37        }
38    }
39}
40
41impl<Server> ComputeChannel<Server> for RefCellComputeChannel<Server>
42where
43    Server: ComputeServer,
44{
45    fn read(&self, handle: &Handle<Server>) -> Reader<Vec<u8>> {
46        self.server.borrow_mut().read(handle)
47    }
48
49    fn create(&self, resource: &[u8]) -> Handle<Server> {
50        self.server.borrow_mut().create(resource)
51    }
52
53    fn empty(&self, size: usize) -> Handle<Server> {
54        self.server.borrow_mut().empty(size)
55    }
56
57    fn execute(&self, kernel_description: Server::Kernel, handles: &[&Handle<Server>]) {
58        self.server
59            .borrow_mut()
60            .execute(kernel_description, handles)
61    }
62
63    fn sync(&self) {
64        self.server.borrow_mut().sync()
65    }
66}
67
68/// This is unsafe, since no concurrency is supported by the `RefCell` channel.
69/// However using this channel should only be done in single threaded environments such as `no-std`.
70unsafe impl<Server: ComputeServer> Send for RefCellComputeChannel<Server> {}
71unsafe impl<Server: ComputeServer> Sync for RefCellComputeChannel<Server> {}
burn_compute/channel/cell.rs

burn_compute/channel/
cell.rs
