diff --git a/nvml-wrapper/src/device.rs b/nvml-wrapper/src/device.rs index 591c640..8838262 100644 --- a/nvml-wrapper/src/device.rs +++ b/nvml-wrapper/src/device.rs @@ -742,6 +742,188 @@ impl<'nvml> Device<'nvml> { } } + /** + Checks simultaneously if confidential compute is enabled, if the device is in a production environment, + and if the device is accepting client requests. + # Errors + * `Uninitialized`, if the library has not been successfully initialized + * `NotSupported`, if this query is not supported by the device + * `InvalidArg`, if confidential compute state is invalid + */ + pub fn check_confidential_compute_status(&self) -> Result { + let cc_state_sym = nvml_sym(self.nvml.lib.nvmlSystemGetConfComputeState.as_ref())?; + let cc_gpus_ready_sym = nvml_sym( + self.nvml + .lib + .nvmlSystemGetConfComputeGpusReadyState + .as_ref(), + )?; + + unsafe { + let mut state: nvmlConfComputeSystemState_t = mem::zeroed(); + nvml_try(cc_state_sym(&mut state))?; + + let is_cc_enabled = state.ccFeature == NVML_CC_SYSTEM_FEATURE_ENABLED; + let is_prod_environment = state.environment == NVML_CC_SYSTEM_ENVIRONMENT_PROD; + + let mut cc_gpus_ready: std::os::raw::c_uint = 0; + nvml_try(cc_gpus_ready_sym(&mut cc_gpus_ready))?; + let is_accepting_client_requests = + cc_gpus_ready == NVML_CC_ACCEPTING_CLIENT_REQUESTS_TRUE; + + Ok(is_cc_enabled && is_prod_environment && is_accepting_client_requests) + } + } + + /** + Gets the confidential compute state for this `Device`. + # Errors + * `Uninitialized`, if the library has not been successfully initialized + * `InvalidArg`, if device is invalid or memory is NULL + * `NotSupported`, if this query is not supported by the device + */ + #[doc(alias = "nvmlDeviceGetConfComputeGpusReadyState")] + pub fn get_confidential_compute_state(&self) -> Result { + let sym = nvml_sym( + self.nvml + .lib + .nvmlSystemGetConfComputeGpusReadyState + .as_ref(), + )?; + + unsafe { + let mut is_accepting_work: u32 = 0; + nvml_try(sym(&mut is_accepting_work))?; + Ok(is_accepting_work == NVML_CC_ACCEPTING_CLIENT_REQUESTS_TRUE) + } + } + + /** + Sets the confidential compute state for this `Device`. + # Errors + * `Uninitialized`, if the library has not been successfully initialized + * `InvalidArg`, if device is invalid or memory is NULL + * `NotSupported`, if this query is not supported by the device + */ + #[doc(alias = "nvmlDeviceSetConfComputeState")] + pub fn set_confidential_compute_state(&self, is_accepting_work: bool) -> Result<(), NvmlError> { + let sym = nvml_sym( + self.nvml + .lib + .nvmlSystemSetConfComputeGpusReadyState + .as_ref(), + )?; + + unsafe { + nvml_try(sym(is_accepting_work as u32))?; + Ok(()) + } + } + + /** + Gets the confidential compute state for this `Device`. + # Errors + + * `Uninitialized`, if the library has not been successfully initialized + * `InvalidArg`, if device is invalid or counters is NULL + * `NotSupported`, if the device does not support this feature + * `GpuLost`, if the target GPU has fallen off the bus or is otherwise inaccessible + * `ArgumentVersionMismatch`, if the provided version is invalid/unsupported + * `Unknown`, on any unexpected error + */ + #[doc(alias = "nvmlDeviceSetConfComputeSettings")] + pub fn is_cc_enabled(&self) -> Result { + let sym = nvml_sym(self.nvml.lib.nvmlSystemGetConfComputeSettings.as_ref())?; + + unsafe { + let mut settings: nvmlSystemConfComputeSettings_t = mem::zeroed(); + nvml_try(sym(&mut settings))?; + Ok(settings.ccFeature == NVML_CC_SYSTEM_FEATURE_ENABLED) + } + } + + /** + Gets the confidential compute state for this `Device`. + # Errors + + * `Uninitialized`, if the library has not been successfully initialized + * `InvalidArg`, if device is invalid or counters is NULL + * `NotSupported`, if the device does not support this feature + * `GpuLost`, if the target GPU has fallen off the bus or is otherwise inaccessible + * `ArgumentVersionMismatch`, if the provided version is invalid/unsupported + * `Unknown`, on any unexpected error + */ + #[doc(alias = "nvmlSystemGetConfComputeSettings")] + pub fn is_multi_gpu_protected_pcie_enabled(&self) -> Result { + let sym = nvml_sym(self.nvml.lib.nvmlSystemGetConfComputeSettings.as_ref())?; + + unsafe { + let mut settings: nvmlSystemConfComputeSettings_t = mem::zeroed(); + nvml_try(sym(&mut settings))?; + Ok(settings.multiGpuMode == NVML_CC_SYSTEM_MULTIGPU_PROTECTED_PCIE) + } + } + + /** + Gets the confidential compute state for this `Device`. + # Errors + + * `Uninitialized`, if the library has not been successfully initialized + * `InvalidArg`, if device is invalid or counters is NULL + * `NotSupported`, if the device does not support this feature + * `GpuLost`, if the target GPU has fallen off the bus or is otherwise inaccessible + * `ArgumentVersionMismatch`, if the provided version is invalid/unsupported + * `Unknown`, on any unexpected error + */ + #[doc(alias = "nvmlSystemGetConfComputeSettings")] + pub fn is_cc_dev_mode_enabled(&self) -> Result { + let sym = nvml_sym(self.nvml.lib.nvmlSystemGetConfComputeSettings.as_ref())?; + + unsafe { + let mut settings: nvmlSystemConfComputeSettings_t = mem::zeroed(); + nvml_try(sym(&mut settings))?; + Ok(settings.devToolsMode == NVML_CC_SYSTEM_DEVTOOLS_MODE_ON) + } + } + + /** + Gets the confidential compute capabilities for this `Device`. + # Errors + * `Uninitialized`, if the library has not been successfully initialized + * `InvalidArg`, if device is invalid or memory is NULL + * `NotSupported`, if this query is not supported by the device + */ + pub fn get_confidential_compute_capabilities( + &self, + ) -> Result { + let sym = nvml_sym(self.nvml.lib.nvmlSystemGetConfComputeCapabilities.as_ref())?; + + unsafe { + let mut capabilities: nvmlConfComputeSystemCaps_t = mem::zeroed(); + nvml_try(sym(&mut capabilities))?; + + let cpu_caps = match capabilities.cpuCaps { + NVML_CC_SYSTEM_CPU_CAPS_NONE => ConfidentialComputeCpuCapabilities::None, + NVML_CC_SYSTEM_CPU_CAPS_AMD_SEV => ConfidentialComputeCpuCapabilities::AmdSev, + NVML_CC_SYSTEM_CPU_CAPS_INTEL_TDX => ConfidentialComputeCpuCapabilities::IntelTdx, + _ => return Err(NvmlError::Unknown), + }; + + let gpus_caps = match capabilities.gpusCaps { + NVML_CC_SYSTEM_GPUS_CC_CAPABLE => ConfidentialComputeGpuCapabilities::Capable, + NVML_CC_SYSTEM_GPUS_CC_NOT_CAPABLE => { + ConfidentialComputeGpuCapabilities::NotCapable + } + _ => return Err(NvmlError::Unknown), + }; + + Ok(ConfidentialComputeCapabilities { + cpu_caps, + gpus_caps, + }) + } + } + /** Fetches the confidential compute attestation report for this [`Device`]. diff --git a/nvml-wrapper/src/structs/device.rs b/nvml-wrapper/src/structs/device.rs index ba1121c..b23db41 100644 --- a/nvml-wrapper/src/structs/device.rs +++ b/nvml-wrapper/src/structs/device.rs @@ -4,6 +4,38 @@ use crate::enum_wrappers::device::OperationMode; #[cfg(feature = "serde")] use serde_derive::{Deserialize, Serialize}; +/// Returned from `Device.get_confidential_compute_capabilities()` +#[derive(Debug, Clone, Eq, PartialEq, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub struct ConfidentialComputeCapabilities { + /// The CPU capabilities. + pub cpu_caps: ConfidentialComputeCpuCapabilities, + /// The GPU capabilities. + pub gpus_caps: ConfidentialComputeGpuCapabilities, +} + +/// The possible CPU capabilities for confidential compute (either None, AMD SEV or Intel TDX) +#[derive(Debug, Clone, Eq, PartialEq, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub enum ConfidentialComputeCpuCapabilities { + /// No CPU capabilities. + None, + /// AMD SEV confidential compute capabilities. + AmdSev, + /// Intel TDX confidential compute capabilities. + IntelTdx, +} + +/// The possible GPU capabilities for confidential compute (either not capable or capable) +#[derive(Debug, Clone, Eq, PartialEq, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub enum ConfidentialComputeGpuCapabilities { + /// Capable. + Capable, + /// Not capable. + NotCapable, +} + /// Returned from `Device.confidential_compute_gpu_attestation_report_bytes()` #[derive(Debug, Clone, Eq, PartialEq, Hash)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]