1use std::fmt;
4
5pub trait Device: Clone + Send + Sync + 'static {
7 fn name(&self) -> &'static str;
9
10 fn is_cpu(&self) -> bool;
12
13 fn is_cuda(&self) -> bool;
15}
16
17#[derive(Debug, Clone, Copy, Default)]
19pub struct Cpu;
20
21impl Device for Cpu {
22 fn name(&self) -> &'static str {
23 "cpu"
24 }
25
26 fn is_cpu(&self) -> bool {
27 true
28 }
29
30 fn is_cuda(&self) -> bool {
31 false
32 }
33}
34
35impl fmt::Display for Cpu {
36 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37 write!(f, "cpu")
38 }
39}
40
41#[derive(Debug, Clone, Copy)]
43pub struct Cuda {
44 pub device_id: i32,
45}
46
47impl Cuda {
48 pub fn new(device_id: i32) -> Self {
49 Cuda { device_id }
50 }
51}
52
53impl Default for Cuda {
54 fn default() -> Self {
55 Cuda { device_id: 0 }
56 }
57}
58
59impl Device for Cuda {
60 fn name(&self) -> &'static str {
61 "cuda"
62 }
63
64 fn is_cpu(&self) -> bool {
65 false
66 }
67
68 fn is_cuda(&self) -> bool {
69 true
70 }
71}
72
73impl fmt::Display for Cuda {
74 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75 write!(f, "cuda:{}", self.device_id)
76 }
77}