ort_candle/
lib.rs

1use candle_core::DType;
2use ort_sys::OrtErrorCode;
3
4mod api;
5mod memory;
6mod session;
7mod tensor;
8
9pub(crate) use ort_sys::stub::Error;
10
11pub use self::api::api;
12
13pub(crate) struct Environment {}
14
15impl Environment {
16	pub fn new_sys() -> *mut ort_sys::OrtEnv {
17		(Box::leak(Box::new(Self {})) as *mut Environment).cast()
18	}
19
20	pub unsafe fn cast_from_sys<'e>(ptr: *const ort_sys::OrtEnv) -> &'e Environment {
21		unsafe { &*ptr.cast::<Environment>() }
22	}
23
24	pub unsafe fn consume_sys(ptr: *mut ort_sys::OrtEnv) -> Box<Environment> {
25		Box::from_raw(ptr.cast::<Environment>())
26	}
27}
28
29fn convert_sys_to_dtype(sys: ort_sys::ONNXTensorElementDataType) -> Result<DType, Error> {
30	match sys {
31		ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8 => Ok(DType::U8),
32		ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32 => Ok(DType::U32),
33		ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16 => Ok(DType::I16),
34		ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 => Ok(DType::I32),
35		ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 => Ok(DType::I64),
36		ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16 => Ok(DType::BF16),
37		ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16 => Ok(DType::F16),
38		ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT => Ok(DType::F32),
39		ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE => Ok(DType::F64),
40		ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E4M3FN => Ok(DType::F8E4M3),
41		_ => Err(Error::new(OrtErrorCode::ORT_FAIL, "Element type not supported by candle"))
42	}
43}
44
45fn convert_dtype_to_sys(dtype: DType) -> Result<ort_sys::ONNXTensorElementDataType, Error> {
46	match dtype {
47		DType::U8 => Ok(ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8),
48		DType::U32 => Ok(ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32),
49		DType::I16 => Ok(ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16),
50		DType::I32 => Ok(ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32),
51		DType::I64 => Ok(ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64),
52		DType::BF16 => Ok(ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16),
53		DType::F16 => Ok(ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16),
54		DType::F32 => Ok(ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT),
55		DType::F64 => Ok(ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE),
56		DType::F8E4M3 => Ok(ort_sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E4M3FN),
57		_ => Err(Error::new(OrtErrorCode::ORT_FAIL, "Element type not supported by ONNX Runtime"))
58	}
59}
ort_candle/lib.rs

ort_candle/
lib.rs
