ruda_tensor/ops/
transaction.rs

1use alloc::vec::Vec;
2use core::future::Future;
3
4use crate::tensor::{BoolTensor, FloatTensor, IntTensor, QuantizedTensor};
5use crate::{Backend, ExecutionError, TensorData, TensorPrimitive};
6
7pub use ruda_core::tensor::transaction::TransactionData as TransactionPrimitiveData;
8use ruda_core::tensor::transaction::ReadbackOrder as Order;
9
10#[derive(Default)]
11/// Contains all tensor primitives that are going to be read.
12pub struct TransactionPrimitive<B: Backend> {
13    /// Float tensors.
14    pub read_floats: Vec<FloatTensor<B>>,
15    /// Quantized tensors.
16    pub read_qfloats: Vec<QuantizedTensor<B>>,
17    /// Int tensors.
18    pub read_ints: Vec<IntTensor<B>>,
19    /// Bool tensors.
20    pub read_bools: Vec<BoolTensor<B>>,
21    orders: Vec<Order>,
22}
23
24/// Operations that are sync by nature and that can be batch together in transactions to improve
25/// compute utilization with efficient laziness.
26pub trait TransactionOps<B: Backend> {
27    /// Executes a [transaction](TransactionPrimitive) and return its
28    /// [data](TransactionPrimitiveData).
29    fn tr_execute(
30        transaction: TransactionPrimitive<B>,
31    ) -> impl Future<Output = Result<TransactionPrimitiveData, ExecutionError>> + Send {
32        async move {
33            let mut floats = Vec::new();
34            let mut qfloats = Vec::new();
35            let mut ints = Vec::new();
36            let mut bools = Vec::new();
37
38            for t in transaction.read_floats {
39                floats.push(B::float_into_data(t).await?);
40            }
41            for t in transaction.read_qfloats {
42                qfloats.push(B::q_into_data(t).await?);
43            }
44            for t in transaction.read_ints {
45                ints.push(B::int_into_data(t).await?);
46            }
47            for t in transaction.read_bools {
48                bools.push(B::bool_into_data(t).await?);
49            }
50
51            Ok(TransactionPrimitiveData {
52                read_floats: floats,
53                read_qfloats: qfloats,
54                read_ints: ints,
55                read_bools: bools,
56            })
57        }
58    }
59}
60
61impl<B: Backend> TransactionPrimitive<B> {
62    /// Creates a new transaction.
63    pub fn new(
64        read_floats: Vec<FloatTensor<B>>,
65        read_qfloats: Vec<QuantizedTensor<B>>,
66        read_ints: Vec<IntTensor<B>>,
67        read_bools: Vec<BoolTensor<B>>,
68    ) -> Self {
69        Self {
70            read_floats,
71            read_qfloats,
72            read_ints,
73            read_bools,
74            orders: Vec::default(),
75        }
76    }
77    /// Executes the transaction asynchronously and returns the [data](TensorData) in the same order
78    /// in which they were [registered](crate::tensor::TransactionOp::register_transaction).
79    pub async fn execute_async(mut self) -> Result<Vec<TensorData>, ExecutionError> {
80        let mut orders = Vec::new();
81        core::mem::swap(&mut orders, &mut self.orders);
82        let result = B::tr_execute(self).await?;
83
84        Ok(result.into_ordered(orders))
85    }
86
87    pub(crate) fn register_float(&mut self, tensor: TensorPrimitive<B>) {
88        match tensor {
89            TensorPrimitive::Float(tensor) => {
90                self.orders.push(Order::Float(self.read_floats.len()));
91                self.read_floats.push(tensor);
92            }
93            TensorPrimitive::QFloat(tensor) => {
94                self.orders.push(Order::QFloat(self.read_qfloats.len()));
95                self.read_qfloats.push(tensor);
96            }
97        }
98    }
99
100    pub(crate) fn register_int(&mut self, tensor: IntTensor<B>) {
101        self.orders.push(Order::Int(self.read_ints.len()));
102        self.read_ints.push(tensor);
103    }
104
105    pub(crate) fn register_bool(&mut self, tensor: BoolTensor<B>) {
106        self.orders.push(Order::Bool(self.read_bools.len()));
107        self.read_bools.push(tensor);
108    }
109}
ruda_tensor/ops/transaction.rs

ruda_tensor/ops/
transaction.rs
