ruda_tensor_dispatch/ops/
transaction.rs

1use super::*;
2use crate::{DispatchDevice, BackendTensor};
3use ruda_tensor::DeviceOps;
4
5#[derive(Default)]
6struct Reads {
7    floats: Vec<(usize, DispatchTensor)>,
8    quantized: Vec<(usize, DispatchTensor)>,
9    ints: Vec<(usize, DispatchTensor)>,
10    bools: Vec<(usize, DispatchTensor)>,
11}
12
13impl TransactionOps<Self> for Dispatch {
14    fn tr_execute(transaction: TransactionPrimitive<Self>) -> impl Future<Output = Result<TransactionPrimitiveData, ExecutionError>> + Send {
15        let mut groups: Vec<(DispatchDevice, Reads)> = Vec::new();
16        let sizes = [transaction.read_floats.len(), transaction.read_qfloats.len(), transaction.read_ints.len(), transaction.read_bools.len()];
17        macro_rules! group {
18            ($values:expr, $field:ident) => {
19                for (index, tensor) in $values.into_iter().enumerate() {
20                    let device = tensor.device().inner().clone();
21                    let position = groups.iter().position(|(candidate, _)| candidate == &device)
22                        .unwrap_or_else(|| { groups.push((device, Reads::default())); groups.len() - 1 });
23                    groups[position].1.$field.push((index, tensor));
24                }
25            };
26        }
27        group!(transaction.read_floats, floats);
28        group!(transaction.read_qfloats, quantized);
29        group!(transaction.read_ints, ints);
30        group!(transaction.read_bools, bools);
31        let pending = groups.into_iter().map(|(device, reads)| {
32            macro_rules! execute {
33                ($B:ty, $variant:ident, $device:expr) => {{
34                    let (float_indices, floats): (Vec<_>, Vec<_>) = reads.floats.into_iter()
35                        .map(|(i, t)| (i, unpack!(native, $variant, Float, t, AutodiffContext::Disabled))).unzip();
36                    let (quantized_indices, quantized): (Vec<_>, Vec<_>) = reads.quantized.into_iter()
37                        .map(|(i, t)| (i, unpack!(native, $variant, Quantized, t, AutodiffContext::Disabled))).unzip();
38                    let (int_indices, ints): (Vec<_>, Vec<_>) = reads.ints.into_iter()
39                        .map(|(i, t)| (i, unpack!(native, $variant, Int, t, AutodiffContext::Disabled))).unzip();
40                    let (bool_indices, bools): (Vec<_>, Vec<_>) = reads.bools.into_iter()
41                        .map(|(i, t)| (i, unpack!(native, $variant, Bool, t, AutodiffContext::Disabled))).unzip();
42                    let future = <$B>::tr_execute(TransactionPrimitive::new(floats, quantized, ints, bools));
43                    let future: core::pin::Pin<Box<dyn Future<Output = Result<TransactionPrimitiveData, ExecutionError>> + Send>> = Box::pin(future);
44                    ([float_indices, quantized_indices, int_indices, bool_indices], future)
45                }};
46            }
47            native_dispatch!(&device, execute)
48        }).collect::<Vec<_>>();
49        async move {
50            let mut output = sizes.map(|size| (0..size).map(|_| None).collect::<Vec<Option<TensorData>>>());
51            for (indices, future) in pending {
52                let data = future.await?;
53                for ((indices, values), destination) in indices.into_iter()
54                    .zip([data.read_floats, data.read_qfloats, data.read_ints, data.read_bools])
55                    .zip(output.iter_mut()) {
56                    assert_eq!(indices.len(), values.len(), "Backend returned an incomplete transaction");
57                    for (index, value) in indices.into_iter().zip(values) { destination[index] = Some(value); }
58                }
59            }
60            let [floats, quantized, ints, bools] = output.map(|values| values.into_iter()
61                .map(|value| value.expect("Missing transaction result")).collect());
62            Ok(TransactionPrimitiveData { read_floats: floats, read_qfloats: quantized, read_ints: ints, read_bools: bools })
63        }
64    }
65}
ruda_tensor_dispatch/ops/transaction.rs

ruda_tensor_dispatch/ops/
transaction.rs
