burn_router/
tensor.rs

1use core::sync::atomic::{AtomicU32, Ordering};
2
3use alloc::format;
4use alloc::{sync::Arc, vec::Vec};
5
6use super::RunnerClient;
7use burn_backend::{DType, Shape, TensorData, TensorMetadata, backend::ExecutionError};
8use burn_ir::{TensorId, TensorIr, TensorStatus};
9
10/// Tensor primitive for the [router backend](crate::BackendRouter).
11pub struct RouterTensor<C: RunnerClient> {
12    pub(crate) id: TensorId,
13    pub(crate) shape: Shape,
14    pub(crate) dtype: DType,
15    /// The client that has this tensor
16    pub client: C,
17    pub(crate) count: Arc<AtomicU32>,
18}
19
20impl<C: RunnerClient> TensorMetadata for RouterTensor<C> {
21    fn dtype(&self) -> DType {
22        self.dtype
23    }
24
25    fn shape(&self) -> Shape {
26        self.shape.clone()
27    }
28
29    fn rank(&self) -> usize {
30        self.shape.num_dims()
31    }
32}
33
34impl<C: RunnerClient> RouterTensor<C> {
35    /// Create a new router tensor.
36    pub fn new(id: TensorId, shape: Shape, dtype: DType, client: C) -> Self {
37        Self {
38            id,
39            shape,
40            dtype,
41            client,
42            count: Arc::new(AtomicU32::new(1)),
43        }
44    }
45
46    pub(crate) async fn into_data(self) -> Result<TensorData, ExecutionError> {
47        self.client.clone().read_tensor_async(self.into_ir()).await
48    }
49
50    /// Get the ir for this tensor
51    pub fn into_ir(mut self) -> TensorIr {
52        let count = self.count.load(Ordering::Relaxed);
53        let status = self.status(count);
54        let mut shape_out = Shape::from(Vec::<usize>::new());
55        core::mem::swap(&mut self.shape, &mut shape_out);
56
57        if let TensorStatus::ReadWrite = status {
58            // Avoids an unwanted drop on the same thread.
59            //
60            // Since `drop` is called after `into_ir`, we must not register a drop if the tensor
61            // was consumed with a `ReadWrite` status.
62            self.count.fetch_add(1, Ordering::Relaxed);
63        }
64
65        TensorIr {
66            status,
67            shape: shape_out,
68            id: self.id,
69            dtype: self.dtype,
70        }
71    }
72
73    pub(crate) fn to_ir_out(&self) -> TensorIr {
74        TensorIr {
75            status: TensorStatus::NotInit,
76            shape: self.shape.clone(),
77            id: self.id,
78            dtype: self.dtype,
79        }
80    }
81
82    pub(crate) fn status(&self, count: u32) -> TensorStatus {
83        if count <= 1 {
84            TensorStatus::ReadWrite
85        } else {
86            TensorStatus::ReadOnly
87        }
88    }
89}
90
91impl<C: RunnerClient> core::fmt::Debug for RouterTensor<C> {
92    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
93        f.write_str(
94            format!(
95                "{{ id: {:?}, shape: {:?}, dtype: {:?}, device: {:?} }}",
96                self.id,
97                self.shape,
98                self.dtype,
99                self.client.device().clone(),
100            )
101            .as_str(),
102        )
103    }
104}
105
106impl<C: RunnerClient> Clone for RouterTensor<C> {
107    fn clone(&self) -> Self {
108        self.count.fetch_add(1, Ordering::Relaxed);
109
110        Self {
111            id: self.id,
112            shape: self.shape.clone(),
113            client: self.client.clone(),
114            dtype: self.dtype,
115            count: self.count.clone(),
116        }
117    }
118}
119
120impl<C: RunnerClient> Drop for RouterTensor<C> {
121    fn drop(&mut self) {
122        let count = self.count.fetch_sub(1, Ordering::Relaxed);
123
124        match self.status(count) {
125            TensorStatus::ReadWrite => {
126                let id = self.id;
127                let mut shape = Shape::from(Vec::<usize>::new());
128                core::mem::swap(&mut shape, &mut self.shape);
129
130                let ir = TensorIr {
131                    id,
132                    shape,
133                    status: TensorStatus::ReadWrite,
134                    dtype: self.dtype,
135                };
136                self.client.register_op(burn_ir::OperationIr::Drop(ir));
137            }
138            TensorStatus::ReadOnly => {}
139            TensorStatus::NotInit => {}
140        }
141    }
142}
burn_router/tensor.rs

burn_router/
tensor.rs
