st-zrt 0.2.0

Stellarrion st-zrt: ultra-fast, zero-overhead Rust runtime over onnxruntime.
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

//! Session initializers backed by caller-owned tensors.
use crate::element::TensorElement;
use crate::tensor::{RunInput, TensorBuffer};
use crate::{Error, Result, sys};
use std::ffi::{CString, c_char};

trait InitializerValue: Send + Sync {
    fn as_value_ptr(&self) -> *const sys::ValueHandle;
}

impl<T> InitializerValue for TensorBuffer<T>
where
    T: TensorElement + Send + Sync,
{
    #[inline]
    fn as_value_ptr(&self) -> *const sys::ValueHandle {
        RunInput::as_value_ptr(self)
    }
}

/// An initializer tensor owned by ZRT and kept alive by the session using it.
///
/// Use this when a model initializer/weight should come from external caller memory
/// instead of being copied from the model file. Construction is zero-copy with respect to
/// the tensor backing buffer: ORT sees the `TensorBuffer`'s storage directly.
pub struct OwnedInitializer {
    name: CString,
    value: Box<dyn InitializerValue>,
}

impl OwnedInitializer {
    pub fn tensor<T>(name: &str, value: TensorBuffer<T>) -> Result<Self>
    where
        T: TensorElement + Send + Sync + 'static,
    {
        Ok(Self {
            name: CString::new(name)
                .map_err(|_| Error::new(-1, "initializer name contains a NUL"))?,
            value: Box::new(value),
        })
    }

    #[inline]
    pub fn name(&self) -> &str {
        self.name
            .to_str()
            .expect("initializer names are constructed from Rust UTF-8")
    }

    #[inline]
    pub(crate) fn name_ptr(&self) -> *const c_char {
        self.name.as_ptr()
    }

    #[inline]
    pub(crate) fn value_ptr(&self) -> *const sys::ValueHandle {
        self.value.as_value_ptr()
    }
}

unsafe impl Send for OwnedInitializer {}
unsafe impl Sync for OwnedInitializer {}