burn_cubecl/ops/
transaction.rs

1use burn_backend::{
2    DType, TensorData,
3    backend::ExecutionError,
4    ops::{TransactionOps, TransactionPrimitive, TransactionPrimitiveData},
5};
6use burn_std::{Shape, Strides};
7use cubecl::server::{CopyDescriptor, Handle};
8
9use crate::{CubeBackend, CubeRuntime};
10
11impl<R: CubeRuntime> TransactionOps<Self> for CubeBackend<R> {
12    async fn tr_execute(
13        transaction: TransactionPrimitive<Self>,
14    ) -> Result<TransactionPrimitiveData, ExecutionError> {
15        let mut client = None;
16
17        enum Kind {
18            Float,
19            Int,
20            Bool,
21        }
22
23        #[derive(new)]
24        struct BindingData {
25            index: usize,
26            kind: Kind,
27            handle: Option<Handle>,
28            shape: Shape,
29            strides: Strides,
30            dtype: DType,
31        }
32
33        let mut num_bindings = 0;
34
35        let mut kinds = Vec::new();
36
37        for t in transaction.read_floats.into_iter() {
38            if client.is_none() {
39                client = Some(t.client.clone());
40            }
41
42            let t = crate::kernel::into_contiguous_aligned(t);
43            let binding = BindingData::new(
44                num_bindings,
45                Kind::Float,
46                Some(t.handle.clone()),
47                t.meta.shape.clone(),
48                t.meta.strides.clone(),
49                t.dtype,
50            );
51
52            kinds.push(binding);
53            num_bindings += 1;
54        }
55        for t in transaction.read_ints.into_iter() {
56            if client.is_none() {
57                client = Some(t.client.clone());
58            }
59
60            let t = crate::kernel::into_contiguous_aligned(t);
61            let binding = BindingData::new(
62                num_bindings,
63                Kind::Int,
64                Some(t.handle.clone()),
65                t.meta.shape.clone(),
66                t.meta.strides.clone(),
67                t.dtype,
68            );
69
70            kinds.push(binding);
71            num_bindings += 1;
72        }
73        for t in transaction.read_bools.into_iter() {
74            if client.is_none() {
75                client = Some(t.client.clone());
76            }
77
78            let t = crate::kernel::into_contiguous_aligned(t);
79            let binding = BindingData::new(
80                num_bindings,
81                Kind::Bool,
82                Some(t.handle.clone()),
83                t.meta.shape.clone(),
84                t.meta.strides.clone(),
85                t.dtype,
86            );
87
88            kinds.push(binding);
89            num_bindings += 1;
90        }
91
92        let client = client.unwrap();
93
94        let bindings = kinds
95            .iter_mut()
96            .map(|b| {
97                CopyDescriptor::new(
98                    b.handle.take().unwrap().binding(),
99                    b.shape.clone(),
100                    b.strides.clone(),
101                    b.dtype.size(),
102                )
103            })
104            .collect();
105
106        let mut data: Vec<Option<_>> = client
107            .read_tensor_async(bindings)
108            .await
109            .map_err(|err| ExecutionError::WithContext {
110                reason: format!("{err:?}"),
111            })?
112            .into_iter()
113            .map(Some)
114            .collect::<Vec<Option<_>>>();
115
116        let mut result = TransactionPrimitiveData::default();
117
118        for binding in kinds {
119            let bytes = data.get_mut(binding.index).unwrap().take().unwrap();
120            let t_data = TensorData::from_bytes(bytes, binding.shape, binding.dtype);
121
122            match binding.kind {
123                Kind::Float => {
124                    result.read_floats.push(t_data);
125                }
126                Kind::Int => {
127                    result.read_ints.push(t_data);
128                }
129                Kind::Bool => {
130                    result.read_bools.push(t_data);
131                }
132            }
133        }
134
135        Ok(result)
136    }
137}
burn_cubecl/ops/transaction.rs

burn_cubecl/ops/
transaction.rs
