333 lines
9.7 KiB
Rust
333 lines
9.7 KiB
Rust
//! Metal Device Management
|
|
//!
|
|
//! This module provides device discovery, initialization, and management
|
|
//! for Apple Metal GPUs.
|
|
|
|
use crate::error::{MetalError, Result};
|
|
use std::sync::Arc;
|
|
use tracing::{debug, info};
|
|
|
|
#[cfg(target_os = "macos")]
|
|
use objc2::rc::Retained;
|
|
#[cfg(target_os = "macos")]
|
|
use objc2::runtime::ProtocolObject;
|
|
#[cfg(target_os = "macos")]
|
|
use objc2_metal::{MTLCreateSystemDefaultDevice, MTLDevice, MTLCommandQueue, MTLCommandBuffer};
|
|
|
|
/// Information about a Metal device
|
|
#[derive(Debug, Clone)]
|
|
pub struct MetalDeviceInfo {
|
|
/// Human-readable device name
|
|
pub name: String,
|
|
/// Registry ID for unique identification
|
|
pub registry_id: u64,
|
|
/// Whether this is a headless (compute-only) device
|
|
pub is_headless: bool,
|
|
/// Whether this is a low-power device
|
|
pub is_low_power: bool,
|
|
/// Whether this device is removable (eGPU)
|
|
pub is_removable: bool,
|
|
/// Recommended maximum working set size in bytes
|
|
pub recommended_max_working_set_size: u64,
|
|
/// Maximum buffer length in bytes
|
|
pub max_buffer_length: usize,
|
|
/// Maximum threads per threadgroup
|
|
pub max_threads_per_threadgroup: usize,
|
|
}
|
|
|
|
/// Metal device wrapper providing safe access to Metal GPU
|
|
#[cfg(target_os = "macos")]
|
|
pub struct MetalDevice {
|
|
/// The underlying Metal device
|
|
device: Retained<ProtocolObject<dyn MTLDevice>>,
|
|
/// Default command queue for this device
|
|
command_queue: Retained<ProtocolObject<dyn MTLCommandQueue>>,
|
|
/// Cached device info
|
|
info: MetalDeviceInfo,
|
|
}
|
|
|
|
/// Stub implementation for non-macOS platforms
|
|
#[cfg(not(target_os = "macos"))]
|
|
pub struct MetalDevice {
|
|
_private: (),
|
|
}
|
|
|
|
#[cfg(target_os = "macos")]
|
|
use objc2_metal::MTLCopyAllDevices;
|
|
|
|
#[cfg(target_os = "macos")]
|
|
impl MetalDevice {
|
|
/// Check if Metal is available on this system
|
|
pub fn is_available() -> bool {
|
|
MTLCreateSystemDefaultDevice().is_some()
|
|
}
|
|
|
|
/// Get the number of available Metal devices
|
|
pub fn device_count() -> usize {
|
|
let devices = MTLCopyAllDevices();
|
|
if devices.is_empty() {
|
|
// Fallback: check if at least default device exists
|
|
usize::from(MTLCreateSystemDefaultDevice().is_some())
|
|
} else {
|
|
devices.len()
|
|
}
|
|
}
|
|
|
|
/// Get a Metal device by index
|
|
///
|
|
/// # Arguments
|
|
/// * `index` - The device index (0-based)
|
|
///
|
|
/// # Returns
|
|
/// The Metal device at the specified index, or an error if not found.
|
|
pub fn by_index(index: usize) -> Result<Self> {
|
|
debug!("Initializing Metal device at index {}", index);
|
|
|
|
let devices = MTLCopyAllDevices();
|
|
|
|
if devices.is_empty() {
|
|
// Fallback to default device if index is 0
|
|
if index == 0 {
|
|
return Self::system_default();
|
|
}
|
|
return Err(MetalError::device_init(format!(
|
|
"No Metal devices available (requested index {index})"
|
|
)));
|
|
}
|
|
|
|
if index >= devices.len() {
|
|
return Err(MetalError::device_init(format!(
|
|
"Metal device {} not found (only {} devices available)",
|
|
index, devices.len()
|
|
)));
|
|
}
|
|
|
|
// Use iterator to get device at index
|
|
let device = devices.iter().nth(index)
|
|
.ok_or_else(|| MetalError::device_init(format!(
|
|
"Failed to get Metal device at index {index}"
|
|
)))?
|
|
.clone();
|
|
|
|
let command_queue = device.newCommandQueue()
|
|
.ok_or_else(|| MetalError::device_init("Failed to create command queue"))?;
|
|
|
|
let info = Self::extract_device_info(&device);
|
|
|
|
info!("Metal device {} initialized: {} (eGPU: {})",
|
|
index, info.name, info.is_removable);
|
|
debug!("Device info: {:?}", info);
|
|
|
|
Ok(Self {
|
|
device,
|
|
command_queue,
|
|
info,
|
|
})
|
|
}
|
|
|
|
/// Get the system default Metal device
|
|
pub fn system_default() -> Result<Self> {
|
|
debug!("Initializing system default Metal device");
|
|
|
|
let device = MTLCreateSystemDefaultDevice()
|
|
.ok_or(MetalError::NoDevice)?;
|
|
|
|
let command_queue = device.newCommandQueue()
|
|
.ok_or_else(|| MetalError::device_init("Failed to create command queue"))?;
|
|
|
|
let info = Self::extract_device_info(&device);
|
|
|
|
info!("Metal device initialized: {}", info.name);
|
|
debug!("Device info: {:?}", info);
|
|
|
|
Ok(Self {
|
|
device,
|
|
command_queue,
|
|
info,
|
|
})
|
|
}
|
|
|
|
/// Check if this device is an eGPU (external GPU)
|
|
pub fn is_egpu(&self) -> bool {
|
|
self.info.is_removable
|
|
}
|
|
|
|
/// Extract device information
|
|
fn extract_device_info(device: &ProtocolObject<dyn MTLDevice>) -> MetalDeviceInfo {
|
|
MetalDeviceInfo {
|
|
name: device.name().to_string(),
|
|
registry_id: device.registryID(),
|
|
is_headless: device.isHeadless(),
|
|
is_low_power: device.isLowPower(),
|
|
is_removable: device.isRemovable(),
|
|
recommended_max_working_set_size: device.recommendedMaxWorkingSetSize(),
|
|
max_buffer_length: device.maxBufferLength(),
|
|
max_threads_per_threadgroup: 1024, // Metal maximum
|
|
}
|
|
}
|
|
|
|
/// Get the device name
|
|
pub fn name(&self) -> &str {
|
|
&self.info.name
|
|
}
|
|
|
|
/// Get device information
|
|
pub fn info(&self) -> &MetalDeviceInfo {
|
|
&self.info
|
|
}
|
|
|
|
/// Get the underlying MTLDevice
|
|
pub fn mtl_device(&self) -> &ProtocolObject<dyn MTLDevice> {
|
|
&self.device
|
|
}
|
|
|
|
/// Get the retained MTLDevice
|
|
pub fn mtl_device_retained(&self) -> &Retained<ProtocolObject<dyn MTLDevice>> {
|
|
&self.device
|
|
}
|
|
|
|
/// Get the command queue
|
|
pub fn command_queue(&self) -> &ProtocolObject<dyn MTLCommandQueue> {
|
|
&self.command_queue
|
|
}
|
|
|
|
/// Get the retained command queue
|
|
pub fn command_queue_retained(&self) -> &Retained<ProtocolObject<dyn MTLCommandQueue>> {
|
|
&self.command_queue
|
|
}
|
|
|
|
/// Create a new command queue
|
|
pub fn new_command_queue(&self) -> Result<Retained<ProtocolObject<dyn MTLCommandQueue>>> {
|
|
self.device.newCommandQueue()
|
|
.ok_or_else(|| MetalError::device_init("Failed to create command queue"))
|
|
}
|
|
|
|
/// Get maximum threads per threadgroup
|
|
pub fn max_threads_per_threadgroup(&self) -> usize {
|
|
self.info.max_threads_per_threadgroup
|
|
}
|
|
|
|
/// Get maximum buffer length
|
|
pub fn max_buffer_length(&self) -> usize {
|
|
self.info.max_buffer_length
|
|
}
|
|
|
|
/// Check if unified memory is supported (always true on Apple Silicon)
|
|
pub fn supports_unified_memory(&self) -> bool {
|
|
// Apple Silicon uses unified memory architecture
|
|
true
|
|
}
|
|
|
|
/// Synchronize all pending operations
|
|
pub fn synchronize(&self) -> Result<()> {
|
|
// Create a command buffer and wait for it to complete
|
|
let cmd_buffer = self.command_queue.commandBuffer()
|
|
.ok_or_else(|| MetalError::CommandBuffer("Failed to create sync command buffer".into()))?;
|
|
cmd_buffer.commit();
|
|
cmd_buffer.waitUntilCompleted();
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[cfg(not(target_os = "macos"))]
|
|
impl MetalDevice {
|
|
/// Check if Metal is available on this system (always false on non-macOS)
|
|
pub fn is_available() -> bool {
|
|
false
|
|
}
|
|
|
|
/// Get the number of available Metal devices (always 0 on non-macOS)
|
|
pub fn device_count() -> usize {
|
|
0
|
|
}
|
|
|
|
/// Get a Metal device by index (not available on non-macOS)
|
|
pub fn by_index(_index: usize) -> Result<Self> {
|
|
Err(MetalError::NotAvailable)
|
|
}
|
|
|
|
/// Get the system default Metal device (not available on non-macOS)
|
|
pub fn system_default() -> Result<Self> {
|
|
Err(MetalError::NotAvailable)
|
|
}
|
|
|
|
/// Check if this device is an eGPU
|
|
pub fn is_egpu(&self) -> bool {
|
|
false
|
|
}
|
|
|
|
/// Get the device name
|
|
pub fn name(&self) -> &str {
|
|
"Metal not available"
|
|
}
|
|
|
|
/// Get device information
|
|
pub fn info(&self) -> MetalDeviceInfo {
|
|
MetalDeviceInfo {
|
|
name: "Metal not available".to_string(),
|
|
registry_id: 0,
|
|
is_headless: false,
|
|
is_low_power: false,
|
|
is_removable: false,
|
|
recommended_max_working_set_size: 0,
|
|
max_buffer_length: 0,
|
|
max_threads_per_threadgroup: 0,
|
|
}
|
|
}
|
|
|
|
/// Get maximum threads per threadgroup
|
|
pub fn max_threads_per_threadgroup(&self) -> usize {
|
|
0
|
|
}
|
|
|
|
/// Get maximum buffer length
|
|
pub fn max_buffer_length(&self) -> usize {
|
|
0
|
|
}
|
|
|
|
/// Check if unified memory is supported
|
|
pub fn supports_unified_memory(&self) -> bool {
|
|
false
|
|
}
|
|
|
|
/// Synchronize all pending operations
|
|
pub fn synchronize(&self) -> Result<()> {
|
|
Err(MetalError::NotAvailable)
|
|
}
|
|
}
|
|
|
|
impl std::fmt::Debug for MetalDevice {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("MetalDevice")
|
|
.field("name", &self.name())
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
/// Thread-safe wrapper for MetalDevice
|
|
pub type SharedMetalDevice = Arc<MetalDevice>;
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_metal_availability() {
|
|
let available = MetalDevice::is_available();
|
|
println!("Metal available: {}", available);
|
|
|
|
#[cfg(not(target_os = "macos"))]
|
|
assert!(!available);
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_device_creation() {
|
|
if MetalDevice::is_available() {
|
|
let device = MetalDevice::system_default().expect("Failed to create device");
|
|
println!("Device name: {}", device.name());
|
|
assert!(!device.name().is_empty());
|
|
}
|
|
}
|
|
}
|