tract-metal 0.23.0-dev.5

Tiny, no-nonsense, self contained, TensorFlow and ONNX inference
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

use metal::{CommandBuffer, ComputeCommandEncoder, ComputeCommandEncoderRef};
use std::ops::{Deref, DerefMut};

#[derive(Debug, Clone)]
pub struct TCommandBuffer {
    inner: CommandBuffer,
    encoder: ComputeCommandEncoder,
}

impl TCommandBuffer {
    pub fn new(command_buffer: CommandBuffer) -> Self {
        let encoder = command_buffer.new_compute_command_encoder().to_owned();

        TCommandBuffer { inner: command_buffer, encoder }
    }

    pub fn encoder(&self) -> &ComputeCommandEncoder {
        &self.encoder
    }

    pub fn encode<EncodeCallback>(&self, encode_cb: EncodeCallback)
    where
        EncodeCallback: Fn(&ComputeCommandEncoderRef),
    {
        encode_cb(&self.encoder);
    }
}

impl Deref for TCommandBuffer {
    type Target = CommandBuffer;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl DerefMut for TCommandBuffer {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}