tfledge/
device.rs

1use core::ptr::null_mut;
2
3use crate::ffi::*;
4
5/// Represents the type of Coral device.
6#[derive(PartialEq, Eq)]
7pub enum CoralDeviceKind {
8    /// A Coral device connected via USB.
9    Usb,
10    /// A Coral device installed on a PCI slot.
11    Pci,
12}
13
14impl CoralDeviceKind {
15    /// Converts the FFI representation of device type to the `CoralDeviceKind` enum.
16    fn from(ffi: edgetpu_device_type) -> Self {
17        match ffi {
18            edgetpu_device_type::EDGETPU_APEX_USB => Self::Usb,
19            edgetpu_device_type::EDGETPU_APEX_PCI => Self::Pci,
20        }
21    }
22}
23
24/// Represents a Coral Edge TPU device.
25///
26/// This struct provides information about a detected Coral device, including its type and system path.
27/// You'll primarily use this struct to:
28///
29/// - Identify the type of Coral device connected (USB or PCI).
30/// - Retrieve the system path to the device, which can be used for further interactions with the device.
31///
32/// # Examples
33///
34/// ```
35/// use tfledge::{list_devices, CoralDeviceKind};
36///
37/// let devices = list_devices();
38///
39/// for device in devices {
40///     // Print the path of the device
41///     println!("Device path: {}", device.path());
42///     
43///     // Check the kind of device
44///     match device.kind() {
45///         CoralDeviceKind::Usb => println!("This is a USB Coral device."),
46///         CoralDeviceKind::Pci => println!("This is a PCI Coral device."),
47///     }
48/// }
49/// ```
50pub struct CoralDevice {
51    ptr: edgetpu_device,
52}
53
54impl CoralDevice {
55    /// Returns the type of the Coral device.
56    ///
57    /// # Examples
58    ///
59    /// See the example in [`CoralDevice`].
60    pub fn kind(&self) -> CoralDeviceKind {
61        CoralDeviceKind::from(self.ptr.type_)
62    }
63
64    /// Returns the system path of the Coral device.
65    ///
66    /// # Examples
67    ///
68    /// See the example in [`CoralDevice`].
69    pub fn path(&self) -> &str {
70        unsafe { core::ffi::CStr::from_ptr(self.ptr.path).to_str().unwrap() }
71    }
72
73    /// Creates a new `TfLiteDelegate` for this device.
74    ///
75    /// This method is used internally to set up the Coral device for use with a TensorFlow Lite interpreter.
76    pub(crate) fn create_delegate(&self) -> *mut TfLiteDelegate {
77        unsafe {
78            // Use default options for the delegate
79            let opts: *mut edgetpu_option = null_mut();
80
81            // Create the delegate using the FFI function
82            let ptr = edgetpu_create_delegate(self.ptr.type_, self.ptr.path, opts, 0);
83
84            // Panic if the delegate creation fails.
85            // In a real application, you'd want to handle this error gracefully.
86            if ptr.is_null() {
87                panic!("Failed to create Edge TPU delegate.");
88            }
89
90            ptr
91        }
92    }
93}
94
95/// An iterator over available Coral Edge TPU devices.
96///
97/// This struct is returned by the [`list_devices()`] function and lets you iterate over all
98/// detected Coral devices. You can use this iterator to find a specific device or to list
99/// all available devices.
100///
101/// # Examples
102///
103/// ```
104/// use tfledge::{list_devices, CoralDeviceKind};
105///
106/// // Iterate over each detected Coral device
107/// for device in list_devices() {
108///     println!("Found Coral device: {}", device.path());
109/// }
110/// ```
111pub struct CoralDeviceList {
112    ptr: *mut edgetpu_device,
113    ct: usize,
114    curr: usize,
115}
116
117impl Iterator for CoralDeviceList {
118    type Item = CoralDevice;
119
120    fn next(&mut self) -> Option<Self::Item> {
121        unsafe {
122            // Check if we're past the end of the device list
123            if self.curr >= self.ct {
124                return None;
125            }
126
127            // Get a reference to the current device from the raw pointer
128            let device = self.ptr.add(self.curr);
129
130            // Increment the counter to the next device
131            self.curr += 1;
132
133            // Construct and return a `CoralDevice` from the raw pointer.
134            Some(CoralDevice { ptr: *device })
135        }
136    }
137}
138
139impl Drop for CoralDeviceList {
140    fn drop(&mut self) {
141        // Free the memory allocated for the device list by the FFI
142        unsafe {
143            edgetpu_free_devices(self.ptr);
144        }
145    }
146}
147
148/// Lists all available Coral Edge TPU devices.
149///
150/// This function returns a [`CoralDeviceList`], which can be used to iterate over all detected
151/// Coral devices.
152///
153/// # Examples
154///
155/// ```
156/// use tfledge::list_devices;
157///
158/// for device in list_devices() {
159///     println!("Found Coral device: {}", device.path());
160/// }
161/// ```
162pub fn list_devices() -> CoralDeviceList {
163    unsafe {
164        let mut ct: usize = 0;
165
166        // Get a pointer to the list of devices and the number of devices
167        let ptr = edgetpu_list_devices(&mut ct);
168
169        // Construct and return the CoralDeviceList
170        CoralDeviceList { ptr, ct, curr: 0 }
171    }
172}
tfledge/device.rs

tfledge/
device.rs
