burn-jit 0.16.1

Generic backend that can be compiled just-in-time to any shader language target
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98

use burn_tensor::{
    ops::{TransactionOps, TransactionPrimitiveResult},
    DType, TensorData,
};

use crate::{element::BoolElement, FloatElement, IntElement, JitBackend, JitRuntime};

impl<R, F, I, BT> TransactionOps<Self> for JitBackend<R, F, I, BT>
where
    R: JitRuntime,
    F: FloatElement,
    I: IntElement,
    BT: BoolElement,
{
    fn tr_execute(
        transaction: burn_tensor::ops::TransactionPrimitive<Self>,
    ) -> impl std::future::Future<Output = burn_tensor::ops::TransactionPrimitiveResult> + 'static + Send
    {
        let mut bindings = Vec::new();
        let mut client = None;

        enum Kind {
            Float(usize, Vec<usize>, DType),
            Int(usize, Vec<usize>, DType),
            Bool(usize, Vec<usize>, DType),
        }

        let mut num_bindings = 0;

        let mut kinds = Vec::new();

        transaction.read_floats.into_iter().for_each(|t| {
            if client.is_none() {
                client = Some(t.client.clone());
            }

            kinds.push(Kind::Float(num_bindings, t.shape.into(), F::dtype()));
            num_bindings += 1;
            bindings.push(t.handle.binding())
        });
        transaction.read_ints.into_iter().for_each(|t| {
            if client.is_none() {
                client = Some(t.client.clone());
            }

            kinds.push(Kind::Int(num_bindings, t.shape.into(), I::dtype()));
            num_bindings += 1;
            bindings.push(t.handle.binding())
        });
        transaction.read_bools.into_iter().for_each(|t| {
            if client.is_none() {
                client = Some(t.client.clone());
            }

            kinds.push(Kind::Bool(num_bindings, t.shape.into(), BT::dtype()));
            num_bindings += 1;
            bindings.push(t.handle.binding())
        });

        let client = client.unwrap();

        async move {
            let mut data: Vec<Option<_>> = client
                .read_async(bindings)
                .await
                .into_iter()
                .map(Some)
                .collect::<Vec<Option<_>>>();

            let mut result = TransactionPrimitiveResult::default();

            for kind in kinds {
                match kind {
                    Kind::Float(index, shape, dtype) => {
                        let bytes = data.get_mut(index).unwrap().take().unwrap();
                        result
                            .read_floats
                            .push(TensorData::from_bytes(bytes, shape, dtype));
                    }
                    Kind::Int(index, shape, dtype) => {
                        let bytes = data.get_mut(index).unwrap().take().unwrap();
                        result
                            .read_ints
                            .push(TensorData::from_bytes(bytes, shape, dtype));
                    }
                    Kind::Bool(index, shape, dtype) => {
                        let bytes = data.get_mut(index).unwrap().take().unwrap();
                        result
                            .read_bools
                            .push(TensorData::from_bytes(bytes, shape, dtype));
                    }
                }
            }

            result
        }
    }
}