embed/
device.rs

1use crate::error::EmbedError;
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub enum EmbedDevice {
5    WebGpu,
6    DirectMl,
7    CoreMl,
8    Cpu,
9    #[cfg(debug_assertions)]
10    Broken,
11}
12
13impl EmbedDevice {
14    pub const fn compiled() -> Self {
15        if cfg!(target_os = "linux") {
16            EmbedDevice::WebGpu
17        } else if cfg!(target_os = "windows") {
18            EmbedDevice::DirectMl
19        } else {
20            EmbedDevice::CoreMl
21        }
22    }
23
24    pub const fn label(self) -> &'static str {
25        match self {
26            EmbedDevice::WebGpu => "Vulkan",
27            EmbedDevice::DirectMl => "DirectML",
28            EmbedDevice::CoreMl => "CoreML",
29            EmbedDevice::Cpu => "the processor",
30            #[cfg(debug_assertions)]
31            EmbedDevice::Broken => "a broken card (test door)",
32        }
33    }
34}
35
36impl TryFrom<probe::Lane> for EmbedDevice {
37    type Error = EmbedError;
38
39    fn try_from(lane: probe::Lane) -> Result<Self, Self::Error> {
40        match lane {
41            probe::Lane::WebGpu => Ok(EmbedDevice::WebGpu),
42            probe::Lane::DirectMl => Ok(EmbedDevice::DirectMl),
43            probe::Lane::CoreMl => Ok(EmbedDevice::CoreMl),
44            probe::Lane::Processor => Ok(EmbedDevice::Cpu),
45            lane => Err(EmbedError::GpuUnavailable {
46                detail: format!("{lane:?} is not an embedding execution provider"),
47            }),
48        }
49    }
50}
51
52pub(crate) fn provider(device: EmbedDevice) -> ort::ep::ExecutionProviderDispatch {
53    match device {
54        #[cfg(target_os = "linux")]
55        EmbedDevice::WebGpu => ort::ep::WebGPU::default().build().error_on_failure(),
56        #[cfg(target_os = "windows")]
57        EmbedDevice::DirectMl => ort::ep::DirectML::default().build().error_on_failure(),
58        #[cfg(target_os = "macos")]
59        EmbedDevice::CoreMl => ort::ep::CoreML::default().build().error_on_failure(),
60        EmbedDevice::Cpu => ort::ep::CPU::default().build().error_on_failure(),
61        #[allow(unreachable_patterns)]
62        _ => unreachable!("a platform GPU provider cannot be used on another platform"),
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69
70    #[test]
71    fn maps_every_ml_lane_to_its_execution_provider() {
72        assert_eq!(
73            EmbedDevice::try_from(probe::Lane::WebGpu),
74            Ok(EmbedDevice::WebGpu)
75        );
76        assert_eq!(
77            EmbedDevice::try_from(probe::Lane::DirectMl),
78            Ok(EmbedDevice::DirectMl)
79        );
80        assert_eq!(
81            EmbedDevice::try_from(probe::Lane::CoreMl),
82            Ok(EmbedDevice::CoreMl)
83        );
84        assert_eq!(
85            EmbedDevice::try_from(probe::Lane::Processor),
86            Ok(EmbedDevice::Cpu)
87        );
88    }
89
90    #[test]
91    fn labels_name_the_engine_the_build_drives() {
92        assert_eq!(EmbedDevice::WebGpu.label(), "Vulkan");
93        assert_eq!(EmbedDevice::DirectMl.label(), "DirectML");
94        assert_eq!(EmbedDevice::CoreMl.label(), "CoreML");
95        assert_eq!(EmbedDevice::Cpu.label(), "the processor");
96    }
97}
embed/device.rs

embed/
device.rs
