torsh_ffi/python/tensor/
device.rs

1use crate::error::{FfiError, FfiResult};
2
3/// Device types for tensor operations
4#[derive(Debug, Clone, PartialEq)]
5pub enum DeviceType {
6    /// CPU device
7    CPU,
8    /// CUDA GPU device with device index
9    CUDA(usize),
10    /// Metal GPU device (future support)
11    Metal(usize),
12    /// WebGPU device (future support)
13    WebGPU(usize),
14}
15
16impl DeviceType {
17    /// Get the device name as a string
18    pub fn name(&self) -> String {
19        match self {
20            DeviceType::CPU => "cpu".to_string(),
21            DeviceType::CUDA(idx) => format!("cuda:{}", idx),
22            DeviceType::Metal(idx) => format!("metal:{}", idx),
23            DeviceType::WebGPU(idx) => format!("webgpu:{}", idx),
24        }
25    }
26
27    /// Parse device string into DeviceType
28    pub fn from_string(device_str: &str) -> FfiResult<Self> {
29        if device_str == "cpu" {
30            Ok(DeviceType::CPU)
31        } else if device_str.starts_with("cuda:") {
32            let idx_str = &device_str[5..];
33            let idx = idx_str
34                .parse::<usize>()
35                .map_err(|_| FfiError::DeviceTransfer {
36                    message: format!("Invalid CUDA device index: {}", idx_str),
37                })?;
38            Ok(DeviceType::CUDA(idx))
39        } else if device_str.starts_with("metal:") {
40            let idx_str = &device_str[6..];
41            let idx = idx_str
42                .parse::<usize>()
43                .map_err(|_| FfiError::DeviceTransfer {
44                    message: format!("Invalid Metal device index: {}", idx_str),
45                })?;
46            Ok(DeviceType::Metal(idx))
47        } else if device_str.starts_with("webgpu:") {
48            let idx_str = &device_str[7..];
49            let idx = idx_str
50                .parse::<usize>()
51                .map_err(|_| FfiError::DeviceTransfer {
52                    message: format!("Invalid WebGPU device index: {}", idx_str),
53                })?;
54            Ok(DeviceType::WebGPU(idx))
55        } else {
56            Err(FfiError::DeviceTransfer {
57                message: format!("Unsupported device type: {}", device_str),
58            })
59        }
60    }
61
62    /// Check if this device is available
63    pub fn is_available(&self) -> bool {
64        match self {
65            DeviceType::CPU => true,
66            DeviceType::CUDA(idx) => {
67                // For now, we'll use a simple check - in a real implementation,
68                // this would check if CUDA is available and the device exists
69                *idx < device_count_cuda()
70            }
71            DeviceType::Metal(_) => false,  // Not implemented yet
72            DeviceType::WebGPU(_) => false, // Not implemented yet
73        }
74    }
75
76    /// Get device compute capability or properties
77    pub fn properties(&self) -> DeviceProperties {
78        match self {
79            DeviceType::CPU => DeviceProperties {
80                name: "CPU".to_string(),
81                memory_total: get_system_memory(),
82                memory_available: get_available_memory(),
83                compute_capability: "N/A".to_string(),
84                multi_processor_count: num_cpus::get(),
85                is_integrated: false,
86            },
87            DeviceType::CUDA(idx) => DeviceProperties {
88                name: format!("CUDA Device {}", idx),
89                memory_total: get_cuda_memory(*idx).unwrap_or(0),
90                memory_available: get_cuda_available_memory(*idx).unwrap_or(0),
91                compute_capability: get_cuda_compute_capability(*idx)
92                    .unwrap_or("Unknown".to_string()),
93                multi_processor_count: get_cuda_sm_count(*idx).unwrap_or(0),
94                is_integrated: false,
95            },
96            DeviceType::Metal(idx) => DeviceProperties {
97                name: format!("Metal Device {}", idx),
98                memory_total: 0,
99                memory_available: 0,
100                compute_capability: "Not implemented".to_string(),
101                multi_processor_count: 0,
102                is_integrated: true,
103            },
104            DeviceType::WebGPU(idx) => DeviceProperties {
105                name: format!("WebGPU Device {}", idx),
106                memory_total: 0,
107                memory_available: 0,
108                compute_capability: "Not implemented".to_string(),
109                multi_processor_count: 0,
110                is_integrated: false,
111            },
112        }
113    }
114}
115
116/// Device properties information
117#[derive(Debug, Clone)]
118pub struct DeviceProperties {
119    pub name: String,
120    pub memory_total: usize,
121    pub memory_available: usize,
122    pub compute_capability: String,
123    pub multi_processor_count: usize,
124    pub is_integrated: bool,
125}
126
127/// Device utility functions (simplified implementations)
128#[allow(dead_code)]
129fn device_count_cuda() -> usize {
130    // In a real implementation, this would query CUDA runtime
131    // For now, we'll assume 0 or 1 device based on availability
132    if std::env::var("CUDA_VISIBLE_DEVICES").is_ok() {
133        1
134    } else {
135        0
136    }
137}
138
139#[allow(dead_code)]
140fn get_system_memory() -> usize {
141    // Get total system memory (simplified)
142    8 * 1024 * 1024 * 1024 // 8GB default
143}
144
145#[allow(dead_code)]
146fn get_available_memory() -> usize {
147    // Get available system memory (simplified)
148    4 * 1024 * 1024 * 1024 // 4GB default
149}
150
151#[allow(dead_code)]
152fn get_cuda_memory(device_idx: usize) -> Option<usize> {
153    // In a real implementation, this would query CUDA device memory
154    if device_idx < device_count_cuda() {
155        Some(8 * 1024 * 1024 * 1024) // 8GB default
156    } else {
157        None
158    }
159}
160
161#[allow(dead_code)]
162fn get_cuda_available_memory(device_idx: usize) -> Option<usize> {
163    // In a real implementation, this would query available CUDA memory
164    if device_idx < device_count_cuda() {
165        Some(6 * 1024 * 1024 * 1024) // 6GB available
166    } else {
167        None
168    }
169}
170
171#[allow(dead_code)]
172fn get_cuda_compute_capability(device_idx: usize) -> Option<String> {
173    // In a real implementation, this would query CUDA compute capability
174    if device_idx < device_count_cuda() {
175        Some("7.5".to_string()) // Common compute capability
176    } else {
177        None
178    }
179}
180
181#[allow(dead_code)]
182fn get_cuda_sm_count(device_idx: usize) -> Option<usize> {
183    // In a real implementation, this would query CUDA SM count
184    if device_idx < device_count_cuda() {
185        Some(68) // Common SM count for modern GPUs
186    } else {
187        None
188    }
189}
torsh_ffi/python/tensor/device.rs

torsh_ffi/python/tensor/
device.rs
